def minimize(thermodynamic_state: states.ThermodynamicState,
sampler_state: states.SamplerState,
max_iterations: int = 20) -> states.SamplerState:
"""
Minimize the given system and state, up to a maximum number of steps.
Parameters
----------
thermodynamic_state : openmmtools.states.ThermodynamicState
The state at which the system could be minimized
sampler_state : openmmtools.states.SamplerState
The starting state at which to minimize the system.
max_iterations : int, optional, default 20
The maximum number of minimization steps. Default is 20.
Returns
-------
sampler_state : openmmtools.states.SamplerState
The posititions and accompanying state following minimization
"""
mc_move = mcmc.LangevinSplittingDynamicsMove()
mcmc_sampler = mcmc.MCMCSampler(thermodynamic_state, sampler_state,
mc_move)
mcmc_sampler.minimize(max_iterations=max_iterations)
return mcmc_sampler.sampler_state
def minimize(thermodynamic_state, sampler_state, mc_move, max_iterations=20):
"""
Minimize the given system and state, up to a maximum number of steps.
Parameters
----------
thermodynamic_state : openmmtools.states.ThermodynamicState
The state at which the system could be minimized
sampler_state : openmmtools.states.SamplerState
The starting state at which to minimize the system.
mc_move : openmmtools.mcmc.MCMove
The move type. This is not directly relevant, but it will
determine whether a context can be reused. It is recommended that
the same move as the equilibrium protocol is used here.
max_iterations : int, optional, default 20
The maximum number of minimization steps. Default is 20.
Returns
-------
sampler_state : openmmtools.states.SamplerState
The posititions and accompanying state following minimization
"""
mcmc_sampler = mcmc.MCMCSampler(thermodynamic_state, sampler_state,
mc_move)
mcmc_sampler.minimize(max_iterations=max_iterations)
return mcmc_sampler.sampler_state
def run_endpoint_perturbation(lambda_thermodynamic_state,
nonalchemical_thermodynamic_state,
initial_hybrid_sampler_state,
mc_move,
n_iterations,
factory,
lambda_index=0,
print_work=False,
write_system=False,
write_state=False,
write_trajectories=False):
"""
Parameters
----------
lambda_thermodynamic_state : ThermodynamicState
The thermodynamic state corresponding to the hybrid system at a lambda endpoint
nonalchemical_thermodynamic_state : ThermodynamicState
The nonalchemical thermodynamic state for the relevant endpoint
initial_hybrid_sampler_state : SamplerState
Starting positions for the sampler. Must be compatible with lambda_thermodynamic_state
mc_move : MCMCMove
The MCMove that will be used for sampling at the lambda endpoint
n_iterations : int
The number of iterations
factory : HybridTopologyFactory
The hybrid topology factory
lambda_index : int, optional, default=0
The index, 0 or 1, at which to retrieve nonalchemical positions
print_work : bool, optional, default=False
If True, will print work values
write_system : bool, optional, default=False
If True, will write alchemical and nonalchemical System XML files
write_state : bool, optional, default=False
If True, write alchemical (hybrid) State XML files each iteration
write_trajectories : bool, optional, default=False
If True, will write trajectories
Returns
-------
df : float
Free energy difference between alchemical and nonalchemical systems, estimated with EXP
ddf : float
Standard deviation of estimate, corrected for correlation, from EXP estimator.
"""
import mdtraj as md
#run an initial minimization:
mcmc_sampler = mcmc.MCMCSampler(lambda_thermodynamic_state,
initial_hybrid_sampler_state, mc_move)
mcmc_sampler.minimize(max_iterations=20)
new_sampler_state = mcmc_sampler.sampler_state
if write_system:
with open(f'hybrid{lambda_index}-system.xml', 'w') as outfile:
outfile.write(
openmm.XmlSerializer.serialize(
lambda_thermodynamic_state.system))
with open(f'nonalchemical{lambda_index}-system.xml', 'w') as outfile:
outfile.write(
openmm.XmlSerializer.serialize(
nonalchemical_thermodynamic_state.system))
#initialize work array
w = np.zeros([n_iterations])
non_potential = np.zeros([n_iterations])
hybrid_potential = np.zeros([n_iterations])
#run n_iterations of the endpoint perturbation:
hybrid_trajectory = unit.Quantity(
np.zeros([
n_iterations,
lambda_thermodynamic_state.system.getNumParticles(), 3
]), unit.nanometers) # DEBUG
nonalchemical_trajectory = unit.Quantity(
np.zeros([
n_iterations,
nonalchemical_thermodynamic_state.system.getNumParticles(), 3
]), unit.nanometers) # DEBUG
for iteration in range(n_iterations):
# Generate a new sampler state for the hybrid system
mc_move.apply(lambda_thermodynamic_state, new_sampler_state)
# Compute the hybrid reduced potential at the new sampler state
hybrid_context, integrator = cache.global_context_cache.get_context(
lambda_thermodynamic_state)
new_sampler_state.apply_to_context(hybrid_context,
ignore_velocities=True)
hybrid_reduced_potential = lambda_thermodynamic_state.reduced_potential(
hybrid_context)
if write_state:
state = hybrid_context.getState(getPositions=True,
getParameters=True)
state_xml = openmm.XmlSerializer.serialize(state)
with open(f'state{iteration}_l{lambda_index}.xml', 'w') as outfile:
outfile.write(state_xml)
# Construct a sampler state for the nonalchemical system
if lambda_index == 0:
nonalchemical_positions = factory.old_positions(
new_sampler_state.positions)
elif lambda_index == 1:
nonalchemical_positions = factory.new_positions(
new_sampler_state.positions)
else:
raise ValueError(
"The lambda index needs to be either one or zero for this to be meaningful"
)
nonalchemical_sampler_state = SamplerState(
nonalchemical_positions, box_vectors=new_sampler_state.box_vectors)
if write_trajectories:
state = hybrid_context.getState(getPositions=True)
hybrid_trajectory[iteration, :, :] = state.getPositions(
asNumpy=True)
nonalchemical_trajectory[iteration, :, :] = nonalchemical_positions
# Compute the nonalchemical reduced potential
nonalchemical_context, integrator = cache.global_context_cache.get_context(
nonalchemical_thermodynamic_state)
nonalchemical_sampler_state.apply_to_context(nonalchemical_context,
ignore_velocities=True)
nonalchemical_reduced_potential = nonalchemical_thermodynamic_state.reduced_potential(
nonalchemical_context)
# Compute and store the work
w[iteration] = nonalchemical_reduced_potential - hybrid_reduced_potential
non_potential[iteration] = nonalchemical_reduced_potential
hybrid_potential[iteration] = hybrid_reduced_potential
if print_work:
print(
f'{iteration:8d} {hybrid_reduced_potential:8.3f} {nonalchemical_reduced_potential:8.3f} => {w[iteration]:8.3f}'
)
if write_trajectories:
if lambda_index == 0:
nonalchemical_mdtraj_topology = md.Topology.from_openmm(
factory._topology_proposal.old_topology)
elif lambda_index == 1:
nonalchemical_mdtraj_topology = md.Topology.from_openmm(
factory._topology_proposal.new_topology)
md.Trajectory(
hybrid_trajectory / unit.nanometers,
factory.hybrid_topology).save(f'hybrid{lambda_index}.pdb')
md.Trajectory(nonalchemical_trajectory / unit.nanometers,
nonalchemical_mdtraj_topology).save(
f'nonalchemical{lambda_index}.pdb')
# Analyze data and return results
[t0, g, Neff_max] = timeseries.detectEquilibration(w)
w_burned_in = w[t0:]
[df, ddf] = pymbar.EXP(w_burned_in)
ddf_corrected = ddf * np.sqrt(g)
results = [df, ddf_corrected, t0, Neff_max]
return results, non_potential, hybrid_potential
def __init__(self, molecules: List[str], output_filename: str, ncmc_switching_times: Dict[str, int], equilibrium_steps: Dict[str, int], timestep: unit.Quantity, initial_molecule: str=None, geometry_options: Dict=None):
self._molecules = [SmallMoleculeSetProposalEngine.canonicalize_smiles(molecule) for molecule in molecules]
environments = ['explicit', 'vacuum']
temperature = 298.15 * unit.kelvin
pressure = 1.0 * unit.atmospheres
constraints = app.HBonds
self._storage = NetCDFStorage(output_filename)
self._ncmc_switching_times = ncmc_switching_times
self._n_equilibrium_steps = equilibrium_steps
self._geometry_options = geometry_options
# Create a system generator for our desired forcefields.
from perses.rjmc.topology_proposal import SystemGenerator
system_generators = dict()
from pkg_resources import resource_filename
gaff_xml_filename = resource_filename('perses', 'data/gaff.xml')
barostat = openmm.MonteCarloBarostat(pressure, temperature)
system_generators['explicit'] = SystemGenerator([gaff_xml_filename, 'tip3p.xml'],
forcefield_kwargs={'nonbondedCutoff': 9.0 * unit.angstrom,
'implicitSolvent': None,
'constraints': constraints,
'ewaldErrorTolerance': 1e-5,
'hydrogenMass': 3.0*unit.amu}, periodic_forcefield_kwargs = {'nonbondedMethod': app.PME}
barostat=barostat)
system_generators['vacuum'] = SystemGenerator([gaff_xml_filename],
forcefield_kwargs={'implicitSolvent': None,
'constraints': constraints,
'hydrogenMass': 3.0*unit.amu}, nonperiodic_forcefield_kwargs = {'nonbondedMethod': app.NoCutoff})
#
# Create topologies and positions
#
topologies = dict()
positions = dict()
from openmoltools import forcefield_generators
forcefield = app.ForceField(gaff_xml_filename, 'tip3p.xml')
forcefield.registerTemplateGenerator(forcefield_generators.gaffTemplateGenerator)
# Create molecule in vacuum.
from perses.utils.openeye import extractPositionsFromOEMol
from openmoltools.openeye import smiles_to_oemol, generate_conformers
if initial_molecule:
smiles = initial_molecule
else:
smiles = np.random.choice(molecules)
molecule = smiles_to_oemol(smiles)
molecule = generate_conformers(molecule, max_confs=1)
topologies['vacuum'] = forcefield_generators.generateTopologyFromOEMol(molecule)
positions['vacuum'] = extractPositionsFromOEMol(molecule)
# Create molecule in solvent.
modeller = app.Modeller(topologies['vacuum'], positions['vacuum'])
modeller.addSolvent(forcefield, model='tip3p', padding=9.0 * unit.angstrom)
topologies['explicit'] = modeller.getTopology()
positions['explicit'] = modeller.getPositions()
# Set up the proposal engines.
proposal_metadata = {}
proposal_engines = dict()
for environment in environments:
proposal_engines[environment] = SmallMoleculeSetProposalEngine(self._molecules,
system_generators[environment])
# Generate systems
systems = dict()
for environment in environments:
systems[environment] = system_generators[environment].build_system(topologies[environment])
# Define thermodynamic state of interest.
thermodynamic_states = dict()
thermodynamic_states['explicit'] = states.ThermodynamicState(system=systems['explicit'],
temperature=temperature, pressure=pressure)
thermodynamic_states['vacuum'] = states.ThermodynamicState(system=systems['vacuum'], temperature=temperature)
# Create SAMS samplers
from perses.samplers.samplers import ExpandedEnsembleSampler, SAMSSampler
mcmc_samplers = dict()
exen_samplers = dict()
sams_samplers = dict()
for environment in environments:
storage = NetCDFStorageView(self._storage, envname=environment)
if self._geometry_options:
n_torsion_divisions = self._geometry_options['n_torsion_divsions'][environment]
use_sterics = self._geometry_options['use_sterics'][environment]
else:
n_torsion_divisions = 180
use_sterics = False
geometry_engine = geometry.FFAllAngleGeometryEngine(storage=storage, n_torsion_divisions=n_torsion_divisions, use_sterics=use_sterics)
move = mcmc.LangevinSplittingDynamicsMove(timestep=timestep, splitting="V R O R V",
n_restart_attempts=10)
chemical_state_key = proposal_engines[environment].compute_state_key(topologies[environment])
if environment == 'explicit':
sampler_state = states.SamplerState(positions=positions[environment],
box_vectors=systems[environment].getDefaultPeriodicBoxVectors())
else:
sampler_state = states.SamplerState(positions=positions[environment])
mcmc_samplers[environment] = mcmc.MCMCSampler(thermodynamic_states[environment], sampler_state, move)
exen_samplers[environment] = ExpandedEnsembleSampler(mcmc_samplers[environment], topologies[environment],
chemical_state_key, proposal_engines[environment],
geometry_engine,
options={'nsteps': self._ncmc_switching_times[environment]}, storage=storage, ncmc_write_interval=self._ncmc_switching_times[environment])
exen_samplers[environment].verbose = True
sams_samplers[environment] = SAMSSampler(exen_samplers[environment], storage=storage)
sams_samplers[environment].verbose = True
# Create test MultiTargetDesign sampler.
from perses.samplers.samplers import MultiTargetDesign
target_samplers = {sams_samplers['explicit']: 1.0, sams_samplers['vacuum']: -1.0}
designer = MultiTargetDesign(target_samplers, storage=self._storage)
# Store things.
self.molecules = molecules
self.environments = environments
self.topologies = topologies
self.positions = positions
self.system_generators = system_generators
self.proposal_engines = proposal_engines
self.thermodynamic_states = thermodynamic_states
self.mcmc_samplers = mcmc_samplers
self.exen_samplers = exen_samplers
self.sams_samplers = sams_samplers
self.designer = designer
def run_endpoint_perturbation(lambda_thermodynamic_state,
nonalchemical_thermodynamic_state,
initial_hybrid_sampler_state,
mc_move,
n_iterations,
factory,
lambda_index=0):
"""
Parameters
----------
lambda_thermodynamic_state : ThermodynamicState
The thermodynamic state corresponding to the hybrid system at a lambda endpoint
nonalchemical_thermodynamic_state : ThermodynamicState
The nonalchemical thermodynamic state for the relevant endpoint
initial_hybrid_sampler_state : SamplerState
Starting positions for the sampler. Must be compatible with lambda_thermodynamic_state
mc_move : MCMCMove
The MCMove that will be used for sampling at the lambda endpoint
n_iterations : int
The number of iterations
factory : HybridTopologyFactory
The hybrid topology factory
lambda_index : int, optional default 0
The index, 0 or 1, at which to retrieve nonalchemical positions
Returns
-------
df : float
Free energy difference between alchemical and nonalchemical systems, estimated with EXP
ddf : float
Standard deviation of estimate, corrected for correlation, from EXP estimator.
"""
#run an initial minimization:
mcmc_sampler = mcmc.MCMCSampler(lambda_thermodynamic_state,
initial_hybrid_sampler_state, mc_move)
mcmc_sampler.minimize(max_iterations=20)
new_sampler_state = mcmc_sampler.sampler_state
#initialize work array
w = np.zeros([n_iterations])
#run n_iterations of the endpoint perturbation:
for iteration in range(n_iterations):
mc_move.apply(lambda_thermodynamic_state, new_sampler_state)
#compute the reduced potential at the new state
hybrid_context, integrator = cache.global_context_cache.get_context(
lambda_thermodynamic_state)
new_sampler_state.apply_to_context(hybrid_context,
ignore_velocities=True)
hybrid_reduced_potential = lambda_thermodynamic_state.reduced_potential(
hybrid_context)
#generate a sampler state for the nonalchemical system
if lambda_index == 0:
nonalchemical_positions = factory.old_positions(
new_sampler_state.positions)
elif lambda_index == 1:
nonalchemical_positions = factory.new_positions(
new_sampler_state.positions)
else:
raise ValueError(
"The lambda index needs to be either one or zero for this to be meaningful"
)
nonalchemical_sampler_state = SamplerState(
nonalchemical_positions, box_vectors=new_sampler_state.box_vectors)
#compute the reduced potential at the nonalchemical system as well:
nonalchemical_context, integrator = cache.global_context_cache.get_context(
nonalchemical_thermodynamic_state)
nonalchemical_sampler_state.apply_to_context(nonalchemical_context,
ignore_velocities=True)
nonalchemical_reduced_potential = nonalchemical_thermodynamic_state.reduced_potential(
nonalchemical_context)
w[iteration] = nonalchemical_reduced_potential - hybrid_reduced_potential
[t0, g, Neff_max] = timeseries.detectEquilibration(w)
print(Neff_max)
w_burned_in = w[t0:]
[df, ddf] = pymbar.EXP(w_burned_in)
ddf_corrected = ddf * np.sqrt(g)
return [df, ddf_corrected, Neff_max]