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)