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)
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
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
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)
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
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)
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 }
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
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)
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!")
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)
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)
def test_lambda_protocol_failure_ends(): bad_function = {'lambda_sterics_delete': lambda x: -x} lp = LambdaProtocol(functions=bad_function)