Ejemplo n.º 1
0
def test_lambda_protocol():
    """

    Tests LambdaProtocol, ensures that it can be instantiated with defaults, and that it fails if disallowed functions are tried

    """

    # check that it's possible to instantiate a LambdaProtocol for all the default types
    for protocol in ['default', 'namd', 'quarters']:
        lp = LambdaProtocol(functions=protocol)

    # check that if we give an incomplete set of parameters it will add in the missing terms
    missing_functions = {'lambda_sterics_delete': lambda x: x}
    lp = LambdaProtocol(functions=missing_functions)
    assert (len(missing_functions) == 1)
    assert (len(lp.get_functions()) == 9)
Ejemplo n.º 2
0
def create_hss(pkl, suffix, selection, checkpoint_interval, n_states):
    with open(pkl, 'rb') as f:
        htf = pickle.load(f)
    lambda_protocol = LambdaProtocol(functions='default')
    reporter_file = pkl[:-3] + suffix + '.nc'
    reporter = MultiStateReporter(
        reporter_file,
        analysis_particle_indices=htf.hybrid_topology.select(selection),
        checkpoint_interval=checkpoint_interval)
    hss = HybridRepexSampler(mcmc_moves=mcmc.LangevinSplittingDynamicsMove(
        timestep=4.0 * unit.femtoseconds,
        collision_rate=5.0 / unit.picosecond,
        n_steps=250,
        reassign_velocities=False,
        n_restart_attempts=20,
        splitting="V R R R O R R R V",
        constraint_tolerance=1e-06),
                             hybrid_factory=htf,
                             online_analysis_interval=10)
    hss.setup(n_states=n_states,
              temperature=300 * unit.kelvin,
              storage_file=reporter,
              lambda_protocol=lambda_protocol,
              endstates=False)
    return hss, reporter
Ejemplo n.º 3
0
def create_hss(reporter_name,
               hybrid_factory,
               selection_string='all',
               checkpoint_interval=1,
               n_states=13):
    lambda_protocol = LambdaProtocol(functions='default')
    reporter = MultiStateReporter(
        reporter_name,
        analysis_particle_indices=hybrid_factory.hybrid_topology.select(
            selection_string),
        checkpoint_interval=checkpoint_interval)
    hss = HybridRepexSampler(mcmc_moves=mcmc.LangevinSplittingDynamicsMove(
        timestep=4.0 * unit.femtoseconds,
        collision_rate=5.0 / unit.picosecond,
        n_steps=250,
        reassign_velocities=False,
        n_restart_attempts=20,
        splitting="V R R R O R R R V",
        constraint_tolerance=1e-06),
                             hybrid_factory=hybrid_factory,
                             online_analysis_interval=10)
    hss.setup(n_states=n_states,
              temperature=300 * unit.kelvin,
              storage_file=reporter,
              lambda_protocol=lambda_protocol,
              endstates=False)
    return hss, reporter
Ejemplo n.º 4
0
def test_create_endstates():
    """
    test the creation of unsampled endstates
    """
    from pkg_resources import resource_filename
    smiles_filename = resource_filename("perses",
                                        os.path.join("data", "test.smi"))
    fe_setup = RelativeFEPSetup(ligand_input=smiles_filename,
                                old_ligand_index=0,
                                new_ligand_index=1,
                                forcefield_files=[],
                                small_molecule_forcefield='gaff-2.11',
                                phases=['vacuum'])

    hybrid_factory = HybridTopologyFactory(
        topology_proposal=fe_setup._vacuum_topology_proposal,
        current_positions=fe_setup._vacuum_positions_old,
        new_positions=fe_setup._vacuum_positions_new,
        neglected_new_angle_terms=fe_setup._vacuum_forward_neglected_angles,
        neglected_old_angle_terms=fe_setup._vacuum_reverse_neglected_angles,
        softcore_LJ_v2=True,
        interpolate_old_and_new_14s=False)

    zero_state_error, one_state_error = validate_endstate_energies(
        fe_setup._vacuum_topology_proposal,
        hybrid_factory,
        added_energy=fe_setup._vacuum_added_valence_energy,
        subtracted_energy=fe_setup._vacuum_subtracted_valence_energy,
        beta=beta,
        platform=openmm.Platform.getPlatformByName('Reference'),
        ENERGY_THRESHOLD=ENERGY_THRESHOLD)

    lambda_alchemical_state = RelativeAlchemicalState.from_system(
        hybrid_factory.hybrid_system)
    lambda_protocol = LambdaProtocol(functions='default')
    lambda_alchemical_state.set_alchemical_parameters(0.0, lambda_protocol)
    thermodynamic_state = CompoundThermodynamicState(
        ThermodynamicState(hybrid_factory.hybrid_system,
                           temperature=temperature),
        composable_states=[lambda_alchemical_state])
    zero_endstate = copy.deepcopy(thermodynamic_state)
    one_endstate = copy.deepcopy(thermodynamic_state)
    one_endstate.set_alchemical_parameters(1.0, lambda_protocol)
    new_endstates = create_endstates(zero_endstate, one_endstate)
Ejemplo n.º 5
0
def create_langevin_integrator(htf, constraint_tol):
    """
    create lambda alchemical states, thermodynamic states, sampler states, integrator, and return context, thermostate, sampler_state, integrator
    """
    fast_lambda_alchemical_state = RelativeAlchemicalState.from_system(
        htf.hybrid_system)
    fast_lambda_alchemical_state.set_alchemical_parameters(
        0.0, LambdaProtocol(functions='default'))

    fast_thermodynamic_state = CompoundThermodynamicState(
        ThermodynamicState(htf.hybrid_system, temperature=temperature),
        composable_states=[fast_lambda_alchemical_state])

    fast_sampler_state = SamplerState(
        positions=htf._hybrid_positions,
        box_vectors=htf.hybrid_system.getDefaultPeriodicBoxVectors())

    integrator_1 = integrators.LangevinIntegrator(
        temperature=temperature,
        timestep=4.0 * unit.femtoseconds,
        splitting='V R O R V',
        measure_shadow_work=False,
        measure_heat=False,
        constraint_tolerance=constraint_tol,
        collision_rate=5.0 / unit.picoseconds)
    #     mcmc_moves=mcmc.LangevinSplittingDynamicsMove(timestep = 4.0 * unit.femtoseconds,
    #                                                              collision_rate=5.0 / unit.picosecond,
    #                                                              n_steps=1,
    #                                                              reassign_velocities=False,
    #                                                              n_restart_attempts=20,
    #                                                              splitting="V R R R O R R R V",
    #                                                              constraint_tolerance=constraint_tol)

    #print(integrator_1.getConstraintTolerance())

    fast_context, fast_integrator = cache.global_context_cache.get_context(
        fast_thermodynamic_state, integrator_1)

    fast_sampler_state.apply_to_context(fast_context)

    return fast_context, fast_thermodynamic_state, fast_sampler_state, fast_integrator
Ejemplo n.º 6
0
def test_create_endstates():
    """
    test the creation of unsampled endstates
    """
    fe_setup = RelativeFEPSetup(ligand_input=f"{os.getcwd()}/test.smi",
                                old_ligand_index=0,
                                new_ligand_index=1,
                                forcefield_files=['gaff.xml'],
                                phases=['vacuum'])

    hybrid_factory = HybridTopologyFactory(
        topology_proposal=fe_setup._vacuum_topology_proposal,
        current_positions=fe_setup._vacuum_positions_old,
        new_positions=fe_setup._vacuum_positions_new,
        neglected_new_angle_terms=fe_setup._vacuum_forward_neglected_angles,
        neglected_old_angle_terms=fe_setup._vacuum_reverse_neglected_angles,
        softcore_LJ_v2=True,
        interpolate_old_and_new_14s=False)

    zero_state_error, one_state_error = validate_endstate_energies(
        fe_setup._vacuum_topology_proposal,
        hybrid_factory,
        added_energy=fe_setup._vacuum_added_valence_energy,
        subtracted_energy=fe_setup._vacuum_subtracted_valence_energy,
        beta=beta,
        ENERGY_THRESHOLD=ENERGY_THRESHOLD)

    lambda_alchemical_state = RelativeAlchemicalState.from_system(
        hybrid_factory.hybrid_system)
    lambda_protocol = LambdaProtocol(functions='default')
    lambda_alchemical_state.set_alchemical_parameters(0.0, lambda_protocol)
    thermodynamic_state = CompoundThermodynamicState(
        ThermodynamicState(hybrid_factory.hybrid_system,
                           temperature=temperature),
        composable_states=[lambda_alchemical_state])
    zero_endstate = copy.deepcopy(thermodynamic_state)
    one_endstate = copy.deepcopy(thermodynamic_state)
    one_endstate.set_alchemical_parameters(1.0, lambda_protocol)
    new_endstates = create_endstates(zero_endstate, one_endstate)
Ejemplo n.º 7
0
def run_setup(setup_options, serialize_systems=True, build_samplers=True):
    """
    Run the setup pipeline and return the relevant setup objects based on a yaml input file.
    Parameters
    ----------
    setup_options : dict
        result of loading yaml input file
    Returns
    -------
    setup_dict: dict
        {'topology_proposals': top_prop, 'hybrid_topology_factories': htf, 'hybrid_samplers': hss}
        - 'topology_proposals':
    """
    phases = setup_options['phases']
    known_phases = ['complex', 'solvent', 'vacuum']
    for phase in phases:
        assert (
            phase in known_phases
        ), f"Unknown phase, {phase} provided. run_setup() can be used with {known_phases}"

    if 'use_given_geometries' not in list(setup_options.keys()):
        use_given_geometries = False
    else:
        assert type(setup_options['use_given_geometries']) == type(True)
        use_given_geometries = setup_options['use_given_geometries']

    if 'complex' in phases:
        _logger.info(f"\tPulling receptor (as pdb or mol2)...")
        # We'll need the protein PDB file (without missing atoms)
        try:
            protein_pdb_filename = setup_options['protein_pdb']
            assert protein_pdb_filename is not None
            receptor_mol2 = None
        except KeyError:
            try:
                receptor_mol2 = setup_options['receptor_mol2']
                assert receptor_mol2 is not None
                protein_pdb_filename = None
            except KeyError as e:
                print(
                    "Either protein_pdb or receptor_mol2 must be specified if running a complex simulation"
                )
                raise e
    else:
        protein_pdb_filename = None
        receptor_mol2 = None

    # And a ligand file containing the pair of ligands between which we will transform
    ligand_file = setup_options['ligand_file']
    _logger.info(f"\tdetected ligand file: {ligand_file}")

    # get the indices of ligands out of the file:
    old_ligand_index = setup_options['old_ligand_index']
    new_ligand_index = setup_options['new_ligand_index']
    _logger.info(
        f"\told ligand index: {old_ligand_index}; new ligand index: {new_ligand_index}"
    )

    _logger.info(f"\tsetting up forcefield files...")
    forcefield_files = setup_options['forcefield_files']

    if "timestep" in setup_options:
        if isinstance(setup_options['timestep'], float):
            timestep = setup_options['timestep'] * unit.femtoseconds
        else:
            timestep = setup_options['timestep']
        _logger.info(f"\ttimestep: {timestep}.")
    else:
        timestep = 1.0 * unit.femtoseconds
        _logger.info(f"\tno timestep detected: setting default as 1.0fs.")

    if "neq_splitting" in setup_options:
        neq_splitting = setup_options['neq_splitting']
        _logger.info(f"\tneq_splitting: {neq_splitting}")

        try:
            eq_splitting = setup_options['eq_splitting']
            _logger.info(f"\teq_splitting: {eq_splitting}")
        except KeyError as e:
            print(
                "If you specify a nonequilibrium splitting string, you must also specify an equilibrium one."
            )
            raise e

    else:
        eq_splitting = "V R O R V"
        neq_splitting = "V R O R V"
        _logger.info(
            f"\tno splitting strings specified: defaulting to neq: {neq_splitting}, eq: {eq_splitting}."
        )

    if "measure_shadow_work" in setup_options:
        measure_shadow_work = setup_options['measure_shadow_work']
        _logger.info(f"\tmeasuring shadow work: {measure_shadow_work}.")
    else:
        measure_shadow_work = False
        _logger.info(
            f"\tno measure_shadow_work specified: defaulting to False.")
    if isinstance(setup_options['pressure'], float):
        pressure = setup_options['pressure'] * unit.atmosphere
    else:
        pressure = setup_options['pressure']
    if isinstance(setup_options['temperature'], float):
        temperature = setup_options['temperature'] * unit.kelvin
    else:
        temperature = setup_options['temperature']
    if isinstance(setup_options['solvent_padding'], float):
        solvent_padding_angstroms = setup_options[
            'solvent_padding'] * unit.angstrom
    else:
        solvent_padding_angstroms = setup_options['solvent_padding']
    if isinstance(setup_options['ionic_strength'], float):
        ionic_strength = setup_options['ionic_strength'] * unit.molar
    else:
        ionic_strength = setup_options['ionic_strength']
    _logger.info(f"\tsetting pressure: {pressure}.")
    _logger.info(f"\tsetting temperature: {temperature}.")
    _logger.info(f"\tsetting solvent padding: {solvent_padding_angstroms}A.")
    _logger.info(f"\tsetting ionic strength: {ionic_strength}M.")

    setup_pickle_file = setup_options[
        'save_setup_pickle_as'] if 'save_setup_pickle_as' in list(
            setup_options) else None
    _logger.info(f"\tsetup pickle file: {setup_pickle_file}")
    trajectory_directory = setup_options['trajectory_directory']
    _logger.info(f"\ttrajectory directory: {trajectory_directory}")
    try:
        atom_map_file = setup_options['atom_map']
        with open(atom_map_file, 'r') as f:
            atom_map = {
                int(x.split()[0]): int(x.split()[1])
                for x in f.readlines()
            }
        _logger.info(f"\tsucceeded parsing atom map.")
    except Exception:
        atom_map = None
        _logger.info(f"\tno atom map specified: default to None.")

    if 'topology_proposal' not in list(setup_options.keys(
    )) or setup_options['topology_proposal'] is None:
        _logger.info(
            f"\tno topology_proposal specified; proceeding to RelativeFEPSetup...\n\n\n"
        )
        if 'set_solvent_box_dims_to_complex' in list(setup_options.keys(
        )) and setup_options['set_solvent_box_dims_to_complex']:
            set_solvent_box_dims_to_complex = True
        else:
            set_solvent_box_dims_to_complex = False

        _logger.info(
            f'Box dimensions: {setup_options["complex_box_dimensions"]} and {setup_options["solvent_box_dimensions"]}'
        )
        fe_setup = RelativeFEPSetup(
            ligand_file,
            old_ligand_index,
            new_ligand_index,
            forcefield_files,
            phases=phases,
            protein_pdb_filename=protein_pdb_filename,
            receptor_mol2_filename=receptor_mol2,
            pressure=pressure,
            temperature=temperature,
            solvent_padding=solvent_padding_angstroms,
            spectator_filenames=setup_options['spectators'],
            map_strength=setup_options['map_strength'],
            atom_expr=setup_options['atom_expr'],
            bond_expr=setup_options['bond_expr'],
            atom_map=atom_map,
            neglect_angles=setup_options['neglect_angles'],
            anneal_14s=setup_options['anneal_1,4s'],
            small_molecule_forcefield=setup_options[
                'small_molecule_forcefield'],
            small_molecule_parameters_cache=setup_options[
                'small_molecule_parameters_cache'],
            trajectory_directory=trajectory_directory,
            trajectory_prefix=setup_options['trajectory_prefix'],
            nonbonded_method=setup_options['nonbonded_method'],
            complex_box_dimensions=setup_options['complex_box_dimensions'],
            solvent_box_dimensions=setup_options['solvent_box_dimensions'],
            ionic_strength=ionic_strength,
            remove_constraints=setup_options['remove_constraints'],
            use_given_geometries=use_given_geometries)

        _logger.info(f"\twriting pickle output...")
        if setup_pickle_file is not None:
            with open(
                    os.path.join(os.getcwd(), trajectory_directory,
                                 setup_pickle_file), 'wb') as f:
                try:
                    pickle.dump(fe_setup, f)
                    _logger.info(f"\tsuccessfully dumped pickle.")
                except Exception as e:
                    print(e)
                    print("\tUnable to save setup object as a pickle")

            _logger.info(
                f"\tsetup is complete.  Writing proposals and positions for each phase to top_prop dict..."
            )
        else:
            _logger.info(
                f"\tsetup is complete.  Omitted writing proposals and positions for each phase to top_prop dict..."
            )

        top_prop = dict()
        for phase in phases:
            top_prop[f'{phase}_topology_proposal'] = getattr(
                fe_setup, f'{phase}_topology_proposal')
            top_prop[f'{phase}_geometry_engine'] = getattr(
                fe_setup, f'_{phase}_geometry_engine')
            top_prop[f'{phase}_old_positions'] = getattr(
                fe_setup, f'{phase}_old_positions')
            top_prop[f'{phase}_new_positions'] = getattr(
                fe_setup, f'{phase}_new_positions')
            top_prop[f'{phase}_added_valence_energy'] = getattr(
                fe_setup, f'_{phase}_added_valence_energy')
            top_prop[f'{phase}_subtracted_valence_energy'] = getattr(
                fe_setup, f'_{phase}_subtracted_valence_energy')
            top_prop[f'{phase}_logp_proposal'] = getattr(
                fe_setup, f'_{phase}_logp_proposal')
            top_prop[f'{phase}_logp_reverse'] = getattr(
                fe_setup, f'_{phase}_logp_reverse')
            top_prop[f'{phase}_forward_neglected_angles'] = getattr(
                fe_setup, f'_{phase}_forward_neglected_angles')
            top_prop[f'{phase}_reverse_neglected_angles'] = getattr(
                fe_setup, f'_{phase}_reverse_neglected_angles')

        top_prop['ligand_oemol_old'] = fe_setup._ligand_oemol_old
        top_prop['ligand_oemol_new'] = fe_setup._ligand_oemol_new
        top_prop[
            'non_offset_new_to_old_atom_map'] = fe_setup.non_offset_new_to_old_atom_map
        _logger.info(f"\twriting atom_mapping.png")
        atom_map_outfile = os.path.join(os.getcwd(), trajectory_directory,
                                        'atom_mapping.png')

        if 'render_atom_map' in list(
                setup_options.keys()) and setup_options['render_atom_map']:
            render_atom_mapping(atom_map_outfile, fe_setup._ligand_oemol_old,
                                fe_setup._ligand_oemol_new,
                                fe_setup.non_offset_new_to_old_atom_map)

    else:
        _logger.info(f"\tloading topology proposal from yaml setup options...")
        top_prop = np.load(setup_options['topology_proposal']).item()

    n_steps_per_move_application = setup_options[
        'n_steps_per_move_application']
    _logger.info(
        f"\t steps per move application: {n_steps_per_move_application}")
    trajectory_directory = setup_options['trajectory_directory']

    trajectory_prefix = setup_options['trajectory_prefix']
    _logger.info(f"\ttrajectory prefix: {trajectory_prefix}")

    if 'atom_selection' in setup_options:
        atom_selection = setup_options['atom_selection']
        _logger.info(f"\tatom selection detected: {atom_selection}")
    else:
        _logger.info(f"\tno atom selection detected: default to all.")
        atom_selection = 'all'

    if setup_options['fe_type'] == 'neq':
        _logger.info(f"\tInstantiating nonequilibrium switching FEP")
        n_equilibrium_steps_per_iteration = setup_options[
            'n_equilibrium_steps_per_iteration']
        ncmc_save_interval = setup_options['ncmc_save_interval']
        write_ncmc_configuration = setup_options['write_ncmc_configuration']
        if setup_options['LSF']:
            _internal_parallelism = {
                'library': ('dask', 'LSF'),
                'num_processes': setup_options['processes']
            }
        else:
            _internal_parallelism = None

        ne_fep = dict()
        for phase in phases:
            _logger.info(f"\t\tphase: {phase}")
            hybrid_factory = HybridTopologyFactory(
                top_prop['%s_topology_proposal' % phase],
                top_prop['%s_old_positions' % phase],
                top_prop['%s_new_positions' % phase],
                neglected_new_angle_terms=top_prop[
                    f"{phase}_forward_neglected_angles"],
                neglected_old_angle_terms=top_prop[
                    f"{phase}_reverse_neglected_angles"],
                softcore_LJ_v2=setup_options['softcore_v2'],
                interpolate_old_and_new_14s=setup_options['anneal_1,4s'])

            if build_samplers:
                ne_fep[phase] = SequentialMonteCarlo(
                    factory=hybrid_factory,
                    lambda_protocol=setup_options['lambda_protocol'],
                    temperature=temperature,
                    trajectory_directory=trajectory_directory,
                    trajectory_prefix=f"{trajectory_prefix}_{phase}",
                    atom_selection=atom_selection,
                    timestep=timestep,
                    eq_splitting_string=eq_splitting,
                    neq_splitting_string=neq_splitting,
                    collision_rate=setup_options['ncmc_collision_rate_ps'],
                    ncmc_save_interval=ncmc_save_interval,
                    internal_parallelism=_internal_parallelism)

        print("Nonequilibrium switching driver class constructed")

        return {'topology_proposals': top_prop, 'ne_fep': ne_fep}

    else:
        _logger.info(f"\tno nonequilibrium detected.")
        htf = dict()
        hss = dict()
        _logger.info(f"\tcataloging HybridTopologyFactories...")

        for phase in phases:
            _logger.info(f"\t\tphase: {phase}:")
            #TODO write a SAMSFEP class that mirrors NonequilibriumSwitchingFEP
            _logger.info(
                f"\t\twriting HybridTopologyFactory for phase {phase}...")
            htf[phase] = HybridTopologyFactory(
                top_prop['%s_topology_proposal' % phase],
                top_prop['%s_old_positions' % phase],
                top_prop['%s_new_positions' % phase],
                neglected_new_angle_terms=top_prop[
                    f"{phase}_forward_neglected_angles"],
                neglected_old_angle_terms=top_prop[
                    f"{phase}_reverse_neglected_angles"],
                softcore_LJ_v2=setup_options['softcore_v2'],
                interpolate_old_and_new_14s=setup_options['anneal_1,4s'])

        for phase in phases:
            # Define necessary vars to check energy bookkeeping
            _top_prop = top_prop['%s_topology_proposal' % phase]
            _htf = htf[phase]
            _forward_added_valence_energy = top_prop['%s_added_valence_energy'
                                                     % phase]
            _reverse_subtracted_valence_energy = top_prop[
                '%s_subtracted_valence_energy' % phase]

            if not use_given_geometries:
                zero_state_error, one_state_error = validate_endstate_energies(
                    _top_prop,
                    _htf,
                    _forward_added_valence_energy,
                    _reverse_subtracted_valence_energy,
                    beta=1.0 / (kB * temperature),
                    ENERGY_THRESHOLD=ENERGY_THRESHOLD
                )  #, trajectory_directory=f'{xml_directory}{phase}')
                _logger.info(f"\t\terror in zero state: {zero_state_error}")
                _logger.info(f"\t\terror in one state: {one_state_error}")
            else:
                _logger.info(
                    f"'use_given_geometries' was passed to setup; skipping endstate validation"
                )

            #TODO expose more of these options in input
            if build_samplers:

                n_states = setup_options['n_states']
                _logger.info(f"\tn_states: {n_states}")
                if 'n_replicas' not in setup_options:
                    n_replicas = n_states
                else:
                    n_replicas = setup_options['n_replicas']

                checkpoint_interval = setup_options['checkpoint_interval']

                # generating lambda protocol
                lambda_protocol = LambdaProtocol(
                    functions=setup_options['protocol-type'])
                _logger.info(
                    f'Using lambda protocol : {setup_options["protocol-type"]}'
                )

                if atom_selection:
                    selection_indices = htf[phase].hybrid_topology.select(
                        atom_selection)
                else:
                    selection_indices = None

                storage_name = str(trajectory_directory) + '/' + str(
                    trajectory_prefix) + '-' + str(phase) + '.nc'
                _logger.info(f'\tstorage_name: {storage_name}')
                _logger.info(f'\tselection_indices {selection_indices}')
                _logger.info(f'\tcheckpoint interval {checkpoint_interval}')
                reporter = MultiStateReporter(
                    storage_name,
                    analysis_particle_indices=selection_indices,
                    checkpoint_interval=checkpoint_interval)

                if phase == 'vacuum':
                    endstates = False
                else:
                    endstates = True

                if setup_options['fe_type'] == 'fah':
                    _logger.info('SETUP FOR FAH DONE')
                    return {
                        'topology_proposals': top_prop,
                        'hybrid_topology_factories': htf
                    }

                if setup_options['fe_type'] == 'sams':
                    hss[phase] = HybridSAMSSampler(
                        mcmc_moves=mcmc.LangevinSplittingDynamicsMove(
                            timestep=timestep,
                            collision_rate=1.0 / unit.picosecond,
                            n_steps=n_steps_per_move_application,
                            reassign_velocities=False,
                            n_restart_attempts=20,
                            constraint_tolerance=1e-06),
                        hybrid_factory=htf[phase],
                        online_analysis_interval=setup_options['offline-freq'],
                        online_analysis_minimum_iterations=10,
                        flatness_criteria=setup_options['flatness-criteria'],
                        gamma0=setup_options['gamma0'])
                    hss[phase].setup(n_states=n_states,
                                     n_replicas=n_replicas,
                                     temperature=temperature,
                                     storage_file=reporter,
                                     lambda_protocol=lambda_protocol,
                                     endstates=endstates)
                elif setup_options['fe_type'] == 'repex':
                    hss[phase] = HybridRepexSampler(
                        mcmc_moves=mcmc.LangevinSplittingDynamicsMove(
                            timestep=timestep,
                            collision_rate=1.0 / unit.picosecond,
                            n_steps=n_steps_per_move_application,
                            reassign_velocities=False,
                            n_restart_attempts=20,
                            constraint_tolerance=1e-06),
                        hybrid_factory=htf[phase],
                        online_analysis_interval=setup_options['offline-freq'])
                    hss[phase].setup(n_states=n_states,
                                     temperature=temperature,
                                     storage_file=reporter,
                                     lambda_protocol=lambda_protocol,
                                     endstates=endstates)
            else:
                _logger.info(f"omitting sampler construction")

            if serialize_systems:
                # save the systems and the states
                pass

                _logger.info('WRITING OUT XML FILES')
                #old_thermodynamic_state, new_thermodynamic_state, hybrid_thermodynamic_state, _ = generate_endpoint_thermodynamic_states(htf[phase].hybrid_system, _top_prop)

                xml_directory = f'{setup_options["trajectory_directory"]}/xml/'
                if not os.path.exists(xml_directory):
                    os.makedirs(xml_directory)
                from perses.utils import data
                _logger.info('WRITING OUT XML FILES')
                _logger.info(f'Saving the hybrid, old and new system to disk')
                data.serialize(
                    htf[phase].hybrid_system,
                    f'{setup_options["trajectory_directory"]}/xml/{phase}-hybrid-system.gz'
                )
                data.serialize(
                    htf[phase]._old_system,
                    f'{setup_options["trajectory_directory"]}/xml/{phase}-old-system.gz'
                )
                data.serialize(
                    htf[phase]._new_system,
                    f'{setup_options["trajectory_directory"]}/xml/{phase}-new-system.gz'
                )

        return {
            'topology_proposals': top_prop,
            'hybrid_topology_factories': htf,
            'hybrid_samplers': hss
        }
Ejemplo n.º 8
0
    def initialize(self,
                   thermodynamic_state,
                   lambda_protocol='default',
                   timestep=1 * unit.femtoseconds,
                   collision_rate=1 / unit.picoseconds,
                   temperature=300 * unit.kelvin,
                   neq_splitting_string='V R O R V',
                   ncmc_save_interval=None,
                   topology=None,
                   subset_atoms=None,
                   measure_shadow_work=False,
                   integrator='langevin',
                   compute_endstate_correction=True):

        try:
            self.context_cache = cache.global_context_cache

            if measure_shadow_work:
                measure_heat = True
            else:
                measure_heat = False

            self.thermodynamic_state = thermodynamic_state
            if integrator == 'langevin':
                self.integrator = integrators.LangevinIntegrator(
                    temperature=temperature,
                    timestep=timestep,
                    splitting=neq_splitting_string,
                    measure_shadow_work=measure_shadow_work,
                    measure_heat=measure_heat,
                    constraint_tolerance=1e-6,
                    collision_rate=collision_rate)
            elif integrator == 'hmc':
                self.integrator = integrators.HMCIntegrator(
                    temperature=temperature, nsteps=2, timestep=timestep / 2)
            else:
                raise Exception(
                    f"integrator {integrator} is not supported. supported integrators include {self.supported_integrators}"
                )

            self.lambda_protocol_class = LambdaProtocol(
                functions=lambda_protocol)

            #create temperatures
            self.beta = 1.0 / (kB * temperature)
            self.temperature = temperature

            self.save_interval = ncmc_save_interval

            self.topology = topology
            self.subset_atoms = subset_atoms

            #if we have a trajectory, set up some ancillary variables:
            if self.topology is not None:
                n_atoms = self.topology.n_atoms
                self._trajectory_positions = []
                self._trajectory_box_lengths = []
                self._trajectory_box_angles = []

            self.compute_endstate_correction = compute_endstate_correction
            if self.compute_endstate_correction:
                self.thermodynamic_state.set_alchemical_parameters(
                    0.0, lambda_protocol=self.lambda_protocol_class)
                first_endstate = copy.deepcopy(self.thermodynamic_state)
                self.thermodynamic_state.set_alchemical_parameters(
                    1.0, lambda_protocol=self.lambda_protocol_class)
                last_endstate = copy.deepcopy(self.thermodynamic_state)
                endstates = create_endstates(first_endstate, last_endstate)
                self.endstates = {0.0: endstates[0], 1.0: endstates[1]}
            else:
                self.endstates = None

            #set a bool variable for pass or failure
            self.succeed = True
            return True
        except Exception as e:
            _logger.error(e)
            self.succeed = False
            return False
Ejemplo n.º 9
0
    def setup(self,
              n_states,
              temperature,
              storage_file,
              minimisation_steps=100,
              n_replicas=None,
              lambda_schedule=None,
              lambda_protocol=LambdaProtocol(),
              endstates=True):

        from perses.dispersed import feptasks

        hybrid_system = self._factory.hybrid_system

        positions = self._factory.hybrid_positions
        lambda_zero_alchemical_state = RelativeAlchemicalState.from_system(
            hybrid_system)

        thermostate = ThermodynamicState(hybrid_system,
                                         temperature=temperature)
        compound_thermodynamic_state = CompoundThermodynamicState(
            thermostate, composable_states=[lambda_zero_alchemical_state])

        thermodynamic_state_list = []
        sampler_state_list = []

        context_cache = cache.ContextCache()

        if n_replicas is None:
            _logger.info(
                f'n_replicas not defined, setting to match n_states, {n_states}'
            )
            n_replicas = n_states
        elif n_replicas > n_states:
            _logger.warning(
                f'More sampler states: {n_replicas} requested greater than number of states: {n_states}. Setting n_replicas to n_states: {n_states}'
            )
            n_replicas = n_states

        # TODO this feels like it should be somewhere else... just not sure where. Maybe into lambda_protocol
        if lambda_schedule is None:
            lambda_schedule = np.linspace(0., 1., n_states)
        else:
            assert (
                len(lambda_schedule) == n_states
            ), 'length of lambda_schedule must match the number of states, n_states'
            assert (
                lambda_schedule[0] == 0.), 'lambda_schedule must start at 0.'
            assert (
                lambda_schedule[-1] == 1.), 'lambda_schedule must end at 1.'
            difference = np.diff(lambda_schedule)
            assert (all(i >= 0. for i in difference)
                    ), 'lambda_schedule must be monotonicly increasing'

        #starting with the initial positions generated py geometry.py
        sampler_state = SamplerState(
            positions,
            box_vectors=hybrid_system.getDefaultPeriodicBoxVectors())
        for lambda_val in lambda_schedule:
            compound_thermodynamic_state_copy = copy.deepcopy(
                compound_thermodynamic_state)
            compound_thermodynamic_state_copy.set_alchemical_parameters(
                lambda_val, lambda_protocol)
            thermodynamic_state_list.append(compound_thermodynamic_state_copy)

            # now generating a sampler_state for each thermodyanmic state, with relaxed positions
            context, context_integrator = context_cache.get_context(
                compound_thermodynamic_state_copy)
            feptasks.minimize(compound_thermodynamic_state_copy, sampler_state)
            sampler_state_list.append(copy.deepcopy(sampler_state))

        reporter = storage_file

        # making sure number of sampler states equals n_replicas
        if len(sampler_state_list) != n_replicas:
            # picking roughly evenly spaced sampler states
            # if n_replicas == 1, then it will pick the first in the list
            idx = np.round(
                np.linspace(0,
                            len(sampler_state_list) - 1,
                            n_replicas)).astype(int)
            sampler_state_list = [
                state for i, state in enumerate(sampler_state_list) if i in idx
            ]

        assert len(sampler_state_list) == n_replicas

        if endstates:
            # generating unsampled endstates
            _logger.info('Generating unsampled endstates.')
            unsampled_dispersion_endstates = create_endstates(
                copy.deepcopy(thermodynamic_state_list[0]),
                copy.deepcopy(thermodynamic_state_list[-1]))
            self.create(
                thermodynamic_states=thermodynamic_state_list,
                sampler_states=sampler_state_list,
                storage=reporter,
                unsampled_thermodynamic_states=unsampled_dispersion_endstates)
        else:
            self.create(thermodynamic_states=thermodynamic_state_list,
                        sampler_states=sampler_state_list,
                        storage=reporter)
Ejemplo n.º 10
0
    def __init__(self,
                 thermodynamic_state,
                 sampler_state,
                 nsteps,
                 direction,
                 splitting='V R O R V',
                 temperature=300 * unit.kelvin,
                 collision_rate=np.inf / unit.picoseconds,
                 timestep=1.0 * unit.femtosecond,
                 work_save_interval=None,
                 top=None,
                 subset_atoms=None,
                 save_configuration=False,
                 lambda_protocol='default',
                 measure_shadow_work=False,
                 label=None,
                 trajectory_filename=None):

        _logger.debug(f"Initializing Particle...")
        start = time.time()
        self._timers = {}  #instantiate timer
        self.label = [label]

        self.context_cache = cache.global_context_cache

        if measure_shadow_work:
            measure_heat = True
        else:
            measure_heat = False

        assert direction == 'forward' or direction == 'reverse', f"The direction of the annealing protocol ({direction}) is invalid; must be specified as 'forward' or 'reverse'"

        self._direction = direction

        #define the lambda schedule (linear)
        if self._direction == 'forward':
            self.start_lambda = 0.0
            self.end_lambda = 1.0
        elif self._direction == 'reverse':
            self.start_lambda = 1.0
            self.end_lambda = 0.0
        else:
            raise Error(f"direction must be 'forward' or 'reverse'")

        #create lambda protocol
        self._nsteps = nsteps
        if self._nsteps is None:
            self.trailblaze = True
            self._work_save_interval = None  # this is allowed to be None, in which case, the save_config method will never be called
            #likewise, if work work save interval is longer than the trailblazed lambda protocol, the save_configuration method will never be called
            self.current_lambda = self.start_lambda
            self.importance_samples = 1  #including the zero state
            #if we opt to trailblaze, we don't define a self.lambdas linsapce
            #instead, we define a self.current lambda and set it to the value of the start lambda
        else:
            self.trailblaze = False
            self.lambdas = np.linspace(self.start_lambda, self.end_lambda,
                                       self._nsteps)
            self.current_index = int(0)
            self.current_lambda = self.lambdas[self.current_index]
            if work_save_interval is None:
                self._work_save_interval = None
            else:
                self._work_save_interval = work_save_interval
                #check that the work write interval is a factor of the number of steps, so we don't accidentally record the
                #work before the end of the protocol as the end
                if self._nsteps % self._work_save_interval != 0:
                    raise ValueError(
                        "The work writing interval must be a factor of the total number of steps"
                    )

        #create sampling objects
        self.sampler_state = sampler_state
        self.thermodynamic_state = thermodynamic_state
        _logger.debug(f"thermodynamic state: {self.thermodynamic_state}")
        self.integrator = integrators.LangevinIntegrator(
            temperature=temperature,
            timestep=timestep,
            splitting=splitting,
            measure_shadow_work=measure_shadow_work,
            measure_heat=measure_heat,
            constraint_tolerance=1e-6,
            collision_rate=collision_rate)
        #platform = openmm.Platform.getPlatformByName(platform_name)
        self.context = openmm.Context(self.thermodynamic_state.system,
                                      self.integrator)
        #self.context, self.integrator = self.context_cache.get_context(self.thermodynamic_state, integrator)
        _logger.debug(f"context: {self.context}")
        self.lambda_protocol_class = LambdaProtocol(functions=lambda_protocol)
        self.thermodynamic_state.set_alchemical_parameters(
            self.start_lambda, lambda_protocol=self.lambda_protocol_class)
        self.thermodynamic_state.apply_to_context(self.context)
        self.sampler_state.apply_to_context(self.context,
                                            ignore_velocities=True)
        self.context.setVelocitiesToTemperature(
            self.thermodynamic_state.temperature)  #randomize velocities @ temp
        self.integrator.step(1)
        self.sampler_state.update_from_context(self.context)

        #create temperatures
        self._beta = 1.0 / (kB * temperature)
        self._temperature = temperature

        init_state = self.context.getState(getEnergy=True)
        self.initial_energy = self._beta * (init_state.getPotentialEnergy() +
                                            init_state.getKineticEnergy())

        self._save_configuration = save_configuration
        self._measure_shadow_work = measure_shadow_work
        if self._save_configuration:
            if trajectory_filename is None:
                raise Exception(
                    f"cannot save configuration when trajectory_filename is None"
                )
            else:
                self._trajectory_filename = trajectory_filename

        #use the number of step moves plus one, since the first is always zero
        self._cumulative_work = [0.0]
        self._shadow_work = 0.0
        self._heat = 0.0

        self._topology = top
        self._subset_atoms = subset_atoms
        self._trajectory = None

        #if we have a trajectory, set up some ancillary variables:
        if self._topology is not None:
            n_atoms = self._topology.n_atoms
            self._trajectory_positions = []
            self._trajectory_box_lengths = []
            self._trajectory_box_angles = []
        else:
            self._save_configuration = False

        self._timers['instantiate'] = time.time() - start
        self._timers['protocol'] = []
        self._timers['save'] = []

        #set a bool variable for pass or failure
        self.succeed = True
        self.failures = []
        _logger.debug(f"Initialization complete!")
Ejemplo n.º 11
0
    def setup(self,
              n_states,
              temperature,
              storage_file,
              minimisation_steps=100,
              lambda_schedule=None,
              lambda_protocol=LambdaProtocol(),
              endstates=True):

        from perses.dispersed import feptasks

        hybrid_system = self._factory.hybrid_system

        positions = self._factory.hybrid_positions
        lambda_zero_alchemical_state = RelativeAlchemicalState.from_system(
            hybrid_system)

        thermostate = ThermodynamicState(hybrid_system,
                                         temperature=temperature)
        compound_thermodynamic_state = CompoundThermodynamicState(
            thermostate, composable_states=[lambda_zero_alchemical_state])

        thermodynamic_state_list = []
        sampler_state_list = []

        context_cache = cache.ContextCache()

        if lambda_schedule is None:
            lambda_schedule = np.linspace(0., 1., n_states)
        else:
            assert (
                len(lambda_schedule) == n_states
            ), 'length of lambda_schedule must match the number of states, n_states'
            assert (
                lambda_schedule[0] == 0.), 'lambda_schedule must start at 0.'
            assert (
                lambda_schedule[-1] == 1.), 'lambda_schedule must end at 1.'
            difference = np.diff(lambda_schedule)
            assert (all(i >= 0. for i in difference)
                    ), 'lambda_schedule must be monotonicly increasing'

        #starting with the initial positions generated py geometry.py
        sampler_state = SamplerState(
            positions,
            box_vectors=hybrid_system.getDefaultPeriodicBoxVectors())
        for lambda_val in lambda_schedule:
            compound_thermodynamic_state_copy = copy.deepcopy(
                compound_thermodynamic_state)
            compound_thermodynamic_state_copy.set_alchemical_parameters(
                lambda_val, lambda_protocol)
            thermodynamic_state_list.append(compound_thermodynamic_state_copy)

            # now generating a sampler_state for each thermodyanmic state, with relaxed positions
            context, context_integrator = context_cache.get_context(
                compound_thermodynamic_state_copy)
            feptasks.minimize(compound_thermodynamic_state_copy, sampler_state)
            sampler_state_list.append(copy.deepcopy(sampler_state))

        reporter = storage_file

        if endstates:
            # generating unsampled endstates
            logger.info('Generating unsampled endstates.')
            unsampled_dispersion_endstates = create_endstates(
                copy.deepcopy(thermodynamic_state_list[0]),
                copy.deepcopy(thermodynamic_state_list[-1]))
            self.create(
                thermodynamic_states=thermodynamic_state_list,
                sampler_states=sampler_state_list,
                storage=reporter,
                unsampled_thermodynamic_states=unsampled_dispersion_endstates)
        else:
            self.create(thermodynamic_states=thermodynamic_state_list,
                        sampler_states=sampler_state_list,
                        storage=reporter)
Ejemplo n.º 12
0
def test_lambda_protocol_naked_charges():
    naked_charge_functions = {
        'lambda_sterics_insert': lambda x: 0.0 if x < 0.5 else 2.0 * (x - 0.5),
        'lambda_electrostatics_insert': lambda x: 2.0 * x if x < 0.5 else 1.0
    }
    lp = LambdaProtocol(functions=naked_charge_functions)
Ejemplo n.º 13
0
def test_lambda_protocol_failure_ends():
    bad_function = {'lambda_sterics_delete': lambda x: -x}
    lp = LambdaProtocol(functions=bad_function)