def test_find_similar_sampler_states():
"""Test helper method AlchemicalPhase._find_similar_sampler_states."""
sampler_state1 = states.SamplerState(np.random.rand(100, 3))
sampler_state2 = states.SamplerState(np.random.rand(100, 3))
sampler_state3 = states.SamplerState(np.random.rand(100, 3))
sampler_states = [sampler_state1, sampler_state1, sampler_state2,
sampler_state1, sampler_state3, sampler_state2]
def test_find_similar_sampler_states():
"""Test helper method AlchemicalPhase._find_similar_sampler_states."""
sampler_state1 = states.SamplerState(np.random.rand(100, 3))
sampler_state2 = states.SamplerState(np.random.rand(100, 3))
sampler_state3 = states.SamplerState(np.random.rand(100, 3))
sampler_states = [sampler_state1, sampler_state1, sampler_state2,
sampler_state1, sampler_state3, sampler_state2]
similar_states = AlchemicalPhase._find_similar_sampler_states(sampler_states)
assert similar_states == {0: [1, 3], 2: [5], 4: []}
def setup_class(cls):
"""Shared test cases for the suite."""
temperature = 300 * unit.kelvin
# Default protocols for tests.
cls.protocol = dict(lambda_electrostatics=[1.0, 0.5, 0.0, 0.0, 0.0],
lambda_sterics=[1.0, 1.0, 1.0, 0.5, 0.0])
cls.restrained_protocol = dict(
lambda_electrostatics=[1.0, 1.0, 0.0, 0.0],
lambda_sterics=[1.0, 1.0, 1.0, 0.0],
lambda_restraints=[0.0, 1.0, 1.0, 1.0])
# Ligand-receptor in implicit solvent.
test_system = testsystems.HostGuestImplicit()
thermodynamic_state = states.ThermodynamicState(
test_system.system, temperature=temperature)
sampler_state = states.SamplerState(
positions=test_system.positions,
box_vectors=test_system.system.getDefaultPeriodicBoxVectors())
topography = Topography(test_system.topology,
ligand_atoms='resname B2')
cls.host_guest_implicit = ('Host-guest implicit', thermodynamic_state,
sampler_state, topography)
# Ligand-receptor in explicit solvent.
test_system = testsystems.HostGuestExplicit()
thermodynamic_state = states.ThermodynamicState(
test_system.system, temperature=temperature)
positions = test_system.positions
box_vectors = test_system.system.getDefaultPeriodicBoxVectors()
sampler_state = states.SamplerState(positions=positions,
box_vectors=box_vectors)
topography = Topography(test_system.topology,
ligand_atoms='resname B2')
cls.host_guest_explicit = ('Host-guest explicit', thermodynamic_state,
sampler_state, topography)
# Peptide solvated in explicit solvent.
test_system = testsystems.AlanineDipeptideExplicit()
thermodynamic_state = states.ThermodynamicState(
test_system.system, temperature=temperature)
positions = test_system.positions
box_vectors = test_system.system.getDefaultPeriodicBoxVectors()
sampler_state = states.SamplerState(positions=positions,
box_vectors=box_vectors)
topography = Topography(test_system.topology)
cls.alanine_explicit = ('Alanine dipeptide explicit',
thermodynamic_state, sampler_state, topography)
# All test cases
cls.all_test_cases = [
cls.host_guest_implicit, cls.host_guest_explicit,
cls.alanine_explicit
]
def test_minimize(self):
"""Test AlchemicalPhase minimization of positions in reference state."""
# Ligand-receptor in implicit solvent.
test_system = testsystems.AlanineDipeptideVacuum()
thermodynamic_state = states.ThermodynamicState(test_system.system,
temperature=300 *
unit.kelvin)
topography = Topography(test_system.topology)
# We create 3 different sampler states that will be distributed over
# replicas in a round-robin fashion.
displacement_vector = np.ones(3) * unit.nanometer
positions2 = test_system.positions + displacement_vector
positions3 = positions2 + displacement_vector
box_vectors = test_system.system.getDefaultPeriodicBoxVectors()
sampler_state1 = states.SamplerState(positions=test_system.positions,
box_vectors=box_vectors)
sampler_state2 = states.SamplerState(positions=positions2,
box_vectors=box_vectors)
sampler_state3 = states.SamplerState(positions=positions3,
box_vectors=box_vectors)
sampler_states = [sampler_state1, sampler_state2, sampler_state3]
with self.temporary_storage_path() as storage_path:
# Create alchemical phase.
alchemical_phase = AlchemicalPhase(ReplicaExchangeSampler())
alchemical_phase.create(thermodynamic_state, sampler_states,
topography, self.protocol, storage_path)
# Measure the average distance between positions. This should be
# maintained after minimization.
sampler_states = alchemical_phase._sampler.sampler_states
original_diffs = [
np.average(sampler_states[i].positions -
sampler_states[i + 1].positions)
for i in range(len(sampler_states) - 1)
]
# Minimize.
alchemical_phase.minimize()
# The minimized positions should be still more or less
# one displacement vector from each other.
sampler_states = alchemical_phase._sampler.sampler_states
new_diffs = [
np.average(sampler_states[i].positions -
sampler_states[i + 1].positions)
for i in range(len(sampler_states) - 1)
]
assert np.allclose(
original_diffs, new_diffs, rtol=0.1
), "original_diffs {} are not close to new_diffs {}".format(
original_diffs, new_diffs)
def setup_class(cls):
lysozyme = testsystems.LysozymeImplicit()
system, positions = lysozyme.system, lysozyme.positions
thermodynamic_state = states.ThermodynamicState(system, 300*unit.kelvin)
sampler_state = states.SamplerState(positions)
topography = Topography(lysozyme.topology, ligand_atoms='resname TMP')
cls.lysozyme_test_case = (thermodynamic_state, sampler_state, topography)
def restraint_selection_template(topography_ligand_atoms=None,
restrained_receptor_atoms=None,
restrained_ligand_atoms=None,
topography_regions=None):
"""The DSL atom selection works as expected."""
test_system = testsystems.HostGuestVacuum()
topography = Topography(test_system.topology, ligand_atoms=topography_ligand_atoms)
if topography_regions is not None:
for region, selection in topography_regions.items():
topography.add_region(region, selection)
sampler_state = states.SamplerState(positions=test_system.positions)
thermodynamic_state = states.ThermodynamicState(test_system.system,
temperature=300.0 * unit.kelvin)
# Initialize with DSL and without processing the string raises an error.
restraint = yank.restraints.Harmonic(spring_constant=2.0 * unit.kilojoule_per_mole / unit.nanometer ** 2,
restrained_receptor_atoms=restrained_receptor_atoms,
restrained_ligand_atoms=restrained_ligand_atoms)
with nose.tools.assert_raises(yank.restraints.RestraintParameterError):
restraint.restrain_state(thermodynamic_state)
# After parameter determination, the indices of the restrained atoms are correct.
restraint.determine_missing_parameters(thermodynamic_state, sampler_state, topography)
assert len(restraint.restrained_receptor_atoms) == 14
assert len(restraint.restrained_ligand_atoms) == 30
# The bond force is configured correctly.
restraint.restrain_state(thermodynamic_state)
system = thermodynamic_state.system
for force in system.getForces():
if isinstance(force, openmm.CustomCentroidBondForce):
assert force.getBondParameters(0)[0] == (0, 1)
assert len(force.getGroupParameters(0)[0]) == 14
assert len(force.getGroupParameters(1)[0]) == 30
assert isinstance(force, openmm.CustomCentroidBondForce) # We have found a force.
def traj_frame_to_sampler_state(traj: md.Trajectory, frame_number: int,
box_vectors):
xyz = traj.xyz[frame_number, :, :]
box_vectors = traj.openmm_boxes(frame_number)
sampler_state = states.SamplerState(
unit.Quantity(xyz, unit=unit.nanometers))
return sampler_state
def generate_example_waterbox_states(temperature=300.0 * unit.kelvin,
pressure=1.0 * unit.atmosphere):
"""
This is a convenience function to generate a CompoundThermodynamicState and SamplerState to use in other tests.
Here, we generate an alchemical water box
"""
#get the water box testsystem
water_ts = testsystems.AlchemicalWaterBox()
system = water_ts.system
positions = water_ts.positions
#construct the openmmtools objects for it
sampler_state = states.SamplerState(
positions, box_vectors=system.getDefaultPeriodicBoxVectors())
thermodynamic_state = states.ThermodynamicState(system,
temperature=temperature,
pressure=pressure)
#make an alchemical state
alchemical_state = alchemy.AlchemicalState.from_system(system)
alchemical_state.set_alchemical_parameters(0.0)
#make a compound thermodynamic state
cpd_thermodynamic_state = states.CompoundThermodynamicState(
thermodynamic_state, [alchemical_state])
return cpd_thermodynamic_state, sampler_state, water_ts.topology
def test_harmonic_standard_state():
"""
Test that the expected harmonic standard state correction is close to our approximation
Also ensures that PBC bonds are being computed and disabled correctly as expected
"""
LJ_fluid = testsystems.LennardJonesFluid()
# Create Harmonic restraint.
restraint = yank.restraints.create_restraint('Harmonic',
restrained_receptor_atoms=1)
# Determine other parameters.
ligand_atoms = [3, 4, 5]
topography = Topography(LJ_fluid.topology, ligand_atoms=ligand_atoms)
sampler_state = states.SamplerState(positions=LJ_fluid.positions)
thermodynamic_state = states.ThermodynamicState(system=LJ_fluid.system,
temperature=300.0 *
unit.kelvin)
restraint.determine_missing_parameters(thermodynamic_state, sampler_state,
topography)
spring_constant = restraint.spring_constant
# Compute standard-state volume for a single molecule in a box of size (1 L) / (avogadros number)
liter = 1000.0 * unit.centimeters**3 # one liter
box_volume = liter / (unit.AVOGADRO_CONSTANT_NA * unit.mole
) # standard state volume
analytical_shell_volume = (
2 * math.pi / (spring_constant * thermodynamic_state.beta))**(3.0 / 2)
analytical_standard_state_G = -math.log(
box_volume / analytical_shell_volume)
restraint_standard_state_G = restraint.get_standard_state_correction(
thermodynamic_state)
np.testing.assert_allclose(analytical_standard_state_G,
restraint_standard_state_G)
def run_rj_proposals(top_prop, configuration_traj, use_sterics, ncmc_nsteps, n_replicates, bond_softening_constant=1.0, angle_softening_constant=1.0):
ncmc_engine = NCMCEngine(nsteps=ncmc_nsteps, pressure=1.0*unit.atmosphere, bond_softening_constant=bond_softening_constant, angle_softening_constant=angle_softening_constant)
geometry_engine = FFAllAngleGeometryEngine(use_sterics=use_sterics, bond_softening_constant=bond_softening_constant, angle_softening_constant=angle_softening_constant)
initial_thermodynamic_state = states.ThermodynamicState(top_prop.old_system, temperature=temperature, pressure=1.0*unit.atmosphere)
final_thermodynamic_state = states.ThermodynamicState(top_prop.new_system, temperature=temperature, pressure=1.0*unit.atmosphere)
traj_indices = np.arange(0, configuration_traj.n_frames)
results = np.zeros([n_replicates, 4])
for i in tqdm.trange(n_replicates):
frame_index = np.random.choice(traj_indices)
initial_sampler_state = traj_frame_to_sampler_state(configuration_traj, frame_index)
initial_logP = - compute_reduced_potential(initial_thermodynamic_state, initial_sampler_state)
proposed_geometry, logP_geometry_forward = geometry_engine.propose(top_prop, initial_sampler_state.positions, beta)
proposed_sampler_state = states.SamplerState(proposed_geometry, box_vectors=initial_sampler_state.box_vectors)
final_old_sampler_state, final_sampler_state, logP_work, initial_hybrid_logP, final_hybrid_logP = ncmc_engine.integrate(top_prop, initial_sampler_state, proposed_sampler_state)
final_logP = - compute_reduced_potential(final_thermodynamic_state, final_sampler_state)
logP_reverse = geometry_engine.logp_reverse(top_prop, final_sampler_state.positions, final_old_sampler_state.positions, beta)
results[i, 0] = initial_hybrid_logP - initial_logP
results[i, 1] = logP_reverse - logP_geometry_forward
results[i, 2] = final_logP - final_hybrid_logP
results[i, 3] = logP_work
return results
def test_BoreschLike_standard_state_analytical():
"""
Perform some analytical tests of the Boresch standard state correction.
Also ensures that PBC is being handled correctly
"""
LJ_fluid = testsystems.LennardJonesFluid()
# Define receptor and ligand atoms
receptor_atoms = [0, 1, 2]
ligand_atoms = [3, 4, 5]
# Create restraint
K_r = 1.0 * unit.kilocalories_per_mole / unit.angstrom**2
r_0 = 0.0 * unit.angstrom
K_theta = 0.0 * unit.kilocalories_per_mole / unit.degrees**2
theta_0 = 30.0 * unit.degrees
topography = Topography(LJ_fluid.topology, ligand_atoms=ligand_atoms)
sampler_state = states.SamplerState(positions=LJ_fluid.positions)
thermodynamic_state = states.ThermodynamicState(system=LJ_fluid.system,
temperature=300.0 *
unit.kelvin)
for restraint_name in ['Boresch', 'PeriodicTorsionBoresch']:
restraint = yank.restraints.create_restraint(
'Boresch',
restrained_receptor_atoms=receptor_atoms,
restrained_ligand_atoms=ligand_atoms,
K_r=K_r,
r_aA0=r_0,
K_thetaA=K_theta,
theta_A0=theta_0,
K_thetaB=K_theta,
theta_B0=theta_0,
K_phiA=K_theta,
phi_A0=theta_0,
K_phiB=K_theta,
phi_B0=theta_0,
K_phiC=K_theta,
phi_C0=theta_0)
# Determine other parameters
restraint.determine_missing_parameters(thermodynamic_state,
sampler_state, topography)
# Compute standard-state volume for a single molecule in a box of size (1 L) / (avogadros number)
liter = 1000.0 * unit.centimeters**3 # one liter
box_volume = liter / (unit.AVOGADRO_CONSTANT_NA * unit.mole
) # standard state volume
analytical_shell_volume = (2 * math.pi /
(K_r * thermodynamic_state.beta))**(3.0 / 2)
analytical_standard_state_G = -math.log(
box_volume / analytical_shell_volume)
restraint_standard_state_G = restraint.get_standard_state_correction(
thermodynamic_state)
msg = 'Failed test for restraint {}'.format(restraint_name)
np.testing.assert_allclose(analytical_standard_state_G,
restraint_standard_state_G,
err_msg=msg)
def compare_energies(mol_name="naphthalene", ref_mol_name="benzene"):
"""
Make an atom map where the molecule at either lambda endpoint is identical, and check that the energies are also the same.
"""
from openmmtools import alchemy, states
from perses.rjmc.topology_proposal import SmallMoleculeSetProposalEngine, TopologyProposal
from perses.annihilation.relative import HybridTopologyFactory
import simtk.openmm as openmm
from perses.utils.openeye import createSystemFromIUPAC
from openmoltools.openeye import iupac_to_oemol, generate_conformers
mol = iupac_to_oemol(mol_name)
mol = generate_conformers(mol, max_confs=1)
m, system, positions, topology = createSystemFromIUPAC(mol_name)
refmol = iupac_to_oemol(ref_mol_name)
refmol = generate_conformers(refmol, max_confs=1)
#map one of the rings
atom_map = SmallMoleculeSetProposalEngine._get_mol_atom_map(mol, refmol)
#now use the mapped atoms to generate a new and old system with identical atoms mapped. This will result in the
#same molecule with the same positions for lambda=0 and 1, and ensures a contiguous atom map
effective_atom_map = {value: value for value in atom_map.values()}
#make a topology proposal with the appropriate data:
top_proposal = TopologyProposal(new_topology=topology,
new_system=system,
old_topology=topology,
old_system=system,
new_to_old_atom_map=effective_atom_map,
new_chemical_state_key="n1",
old_chemical_state_key='n2')
factory = HybridTopologyFactory(top_proposal, positions, positions)
alchemical_system = factory.hybrid_system
alchemical_positions = factory.hybrid_positions
platform = openmm.Platform.getPlatformByName("Reference")
_, _, alch_zero_state, alch_one_state = utils.generate_endpoint_thermodynamic_states(
alchemical_system, top_proposal)
rp_list = []
for state in [alch_zero_state, alch_one_state]:
integrator = openmm.VerletIntegrator(1)
context = state.create_context(integrator, platform)
samplerstate = states.SamplerState(
positions=alchemical_positions,
box_vectors=alchemical_system.getDefaultPeriodicBoxVectors())
samplerstate.apply_to_context(context)
rp = state.reduced_potential(context)
rp_list.append(rp)
del context, integrator
assert abs(rp_list[0] - rp_list[1]) < 1e-6
def test_partial_parametrization():
"""The automatic restraint parametrization doesn't overwrite user values."""
# Create states and identify ligand/receptor.
test_system = testsystems.HostGuestVacuum()
topography = Topography(test_system.topology, ligand_atoms='resname B2')
sampler_state = states.SamplerState(positions=test_system.positions)
thermodynamic_state = states.ThermodynamicState(test_system.system,
temperature=300.0*unit.kelvin)
# Test case: (restraint_type, constructor_kwargs)
boresch = dict(restrained_ligand_atoms=[130, 131, 136], K_r=1.0*unit.kilojoule_per_mole/unit.angstroms**2)
test_cases = [
('Harmonic', dict(spring_constant=2.0*unit.kilojoule_per_mole/unit.nanometer**2,
restrained_receptor_atoms=[5])),
('FlatBottom', dict(well_radius=1.0*unit.angstrom, restrained_ligand_atoms=[130])),
('Boresch', boresch),
('PeriodicTorsionBoresch', boresch),
]
if OpenMM73.dev_validate:
test_cases.append(('RMSD', dict(restrained_ligand_atoms=[130, 131, 136],
K_RMSD=1.0 * unit.kilojoule_per_mole / unit.angstroms ** 2)))
for restraint_type, kwargs in test_cases:
state = copy.deepcopy(thermodynamic_state)
restraint = yank.restraints.create_restraint(restraint_type, **kwargs)
# Test-precondition: The restraint has undefined parameters.
with nose.tools.assert_raises(yank.restraints.RestraintParameterError):
restraint.restrain_state(state)
# The automatic parametrization maintains user values.
restraint.determine_missing_parameters(state, sampler_state, topography)
for parameter_name, parameter_value in kwargs.items():
assert getattr(restraint, parameter_name) == parameter_value
# The rest of the parameters has been determined.
restraint.get_standard_state_correction(state)
# The force has been configured correctly.
restraint.restrain_state(state)
system = state.system
for force in system.getForces():
# RadiallySymmetricRestraint between two single atoms.
if isinstance(force, openmm.CustomBondForce):
particle1, particle2, _ = force.getBondParameters(0)
assert particle1 == restraint.restrained_receptor_atoms[0]
assert particle2 == restraint.restrained_ligand_atoms[0]
# Boresch restraint.
elif isinstance(force, openmm.CustomCompoundBondForce):
particles, _ = force.getBondParameters(0)
assert particles == tuple(restraint.restrained_receptor_atoms + restraint.restrained_ligand_atoms)
# RMSD restraint.
elif OpenMM73.dev_validate and isinstance(force, openmm.CustomCVForce):
rmsd_cv = force.getCollectiveVariable(0)
particles = rmsd_cv.getParticles()
assert particles == tuple(restraint.restrained_receptor_atoms + restraint.restrained_ligand_atoms)
def build_test_case():
"""Create a new ThermodynamicState, SamplerState and Topography."""
# Create a test system
t = HostGuestNoninteracting()
# Create states and topography encoding the info to determine the parameters.
topography = Topography(t.topology, ligand_atoms='resname B2')
sampler_state = states.SamplerState(positions=t.positions)
thermodynamic_state = states.ThermodynamicState(system=t.system, temperature=300.0*unit.kelvin)
return thermodynamic_state, sampler_state, topography
def compute_nonalchemical_perturbation(
equilibrium_result: EquilibriumResult,
hybrid_factory: HybridTopologyFactory,
nonalchemical_thermodynamic_state: states.ThermodynamicState,
lambda_state: int):
"""
Compute the perturbation of transforming the given hybrid equilibrium result into the system for the given nonalchemical_thermodynamic_state
Parameters
----------
equilibrium_result : EquilibriumResult
Result of the equilibrium simulation
hybrid_factory : HybridTopologyFactory
Hybrid factory necessary for getting the positions of the nonalchemical system
nonalchemical_thermodynamic_state : states.ThermodynamicState
ThermodynamicState of the nonalchemical system
lambda_state : int
Whether this is lambda 0 or 1
Returns
-------
work : float
perturbation in kT from the hybrid system to the nonalchemical one
"""
#get the objects we need to begin
hybrid_reduced_potential = equilibrium_result.reduced_potential
hybrid_sampler_state = equilibrium_result.sampler_state
hybrid_positions = hybrid_sampler_state.positions
#get the positions for the nonalchemical system
if lambda_state == 0:
nonalchemical_positions = hybrid_factory.old_positions(
hybrid_positions)
elif lambda_state == 1:
nonalchemical_positions = hybrid_factory.new_positions(
hybrid_positions)
else:
raise ValueError("lambda_state must be 0 or 1")
nonalchemical_sampler_state = states.SamplerState(
nonalchemical_positions, box_vectors=hybrid_sampler_state.box_vectors)
nonalchemical_reduced_potential = compute_reduced_potential(
nonalchemical_thermodynamic_state, nonalchemical_sampler_state)
return hybrid_reduced_potential - nonalchemical_reduced_potential
def test_restrain_atoms():
"""Check that the restrained molecule's centroid is in the origin."""
host_guest = testsystems.HostGuestExplicit()
topology = mdtraj.Topology.from_openmm(host_guest.topology)
sampler_state = states.SamplerState(positions=host_guest.positions)
thermodynamic_state = states.ThermodynamicState(
host_guest.system,
temperature=300 * unit.kelvin,
pressure=1.0 * unit.atmosphere)
# Restrain all the host carbon atoms.
restrained_atoms = [
atom.index for atom in topology.atoms
if atom.element.symbol is 'C' and atom.index <= 125
]
restrain_atoms(thermodynamic_state, sampler_state, restrained_atoms)
# Compute host center_of_geometry.
centroid = np.mean(sampler_state.positions[:126], axis=0)
assert np.allclose(centroid, np.zeros(3))
protocol[f'dr_dis{distance_index}'].append(
dunbrack_std[target, distance_index + 7])
print(protocol.keys())
constants = {
'temperature': temperature,
'pressure': pressure,
}
composable_state = MyComposableState.from_system(system)
thermo_states = states.create_thermodynamic_state_protocol(
system=system,
protocol=protocol,
constants=constants,
composable_states=[composable_state])
# assign sampler_state
sampler_state = states.SamplerState(
positions=pdb.positions, box_vectors=system.getDefaultPeriodicBoxVectors())
# Set up the context for mtd simulation
# at this step the CV and the system are separately passed to Metadynamics
from yank.multistate import SAMSSampler, MultiStateReporter
# TODO: You can use heavy hydrogens and 4 fs timesteps
move = mcmc.LangevinDynamicsMove(timestep=2.0 * unit.femtoseconds,
collision_rate=1.0 / unit.picosecond,
n_steps=1000,
reassign_velocities=False)
print("Done specifying integrator for simulation.")
simulation = SAMSSampler(mcmc_moves=move,
number_of_iterations=iteration,
online_analysis_interval=None,
gamma0=1.0,
flatness_threshold=0.2)
def _setup(self, topology, positions, box):
"""
Set up calculation.
"""
# Signal that system has now been set up
if self._setup_complete:
raise Exception(
"System has already been set up---cannot run again.")
self._setup_complete = True
# Compute thermal energy and inverse temperature
self.kT = kB * self.temperature
self.beta = 1.0 / self.kT
# Add a barostat
# TODO: Is this necessary, since ThermodynamicState handles this automatically? It may not correctly handle MonteCarloAnisotropicBarostat.
self._add_barostat()
# Create SamplerState for initial conditions
self.sampler_state = states.SamplerState(positions=positions,
box_vectors=self.box)
# Create reference thermodynamic state
self.reference_thermodynamic_state = states.ThermodynamicState(
system=self.system,
temperature=self.temperature,
pressure=self.pressure)
# Anneal ligand into binding site
if self.anneal_ligand:
self._anneal_ligand()
positions = self.sampler_state.positions
structure = pmd.openmm.load_topology(topology,
system=self.system,
xyz=positions)
# Create ThermodynamicStates for umbrella sampling along pore
self.thermodynamic_states = self._auto_create_thermodynamic_states(
structure, topology, self.reference_thermodynamic_state)
# Minimize initial thermodynamic state
# TODO: Select initial thermodynamic state based on which state has minimum energy
initial_state_index = 0
#self.sampler_state = self._minimize_sampler_state(self.thermodynamic_states[initial_state_index], self.sampler_state)
#pickle.dump(self.sampler_state, open('sampler_state.obj', 'wb'))
#pickle.dump(self.reference_thermodynamic_state, open('reference_thermodynamic_state.obj', 'wb'))
#pickle.dump(self.thermodynamic_states, open('thermodynamic_states.obj', 'wb'))
#pickle.dump(self.system, open('system.obj', 'wb'))
# Set up simulation
from yank.multistate import SAMSSampler, MultiStateReporter
# TODO: Change this to LangevinSplittingDynamicsMove
move = mcmc.LangevinDynamicsMove(timestep=self.timestep,
collision_rate=self.collision_rate,
n_steps=self.n_steps_per_iteration,
reassign_velocities=False)
self.simulation = SAMSSampler(
mcmc_moves=move,
number_of_iterations=self.n_iterations,
online_analysis_interval=None,
gamma0=self.gamma0,
flatness_threshold=self.flatness_threshold)
self.reporter = MultiStateReporter(
self.output_filename,
checkpoint_interval=self.checkpoint_interval,
analysis_particle_indices=self.analysis_particle_indices)
self.simulation.create(
thermodynamic_states=self.thermodynamic_states,
unsampled_thermodynamic_states=[
self.reference_thermodynamic_state
],
sampler_states=[self.sampler_state],
initial_thermodynamic_states=[initial_state_index],
storage=self.reporter)
def compute_nonalchemical_perturbation(
alchemical_thermodynamic_state: states.ThermodynamicState,
growth_thermodynamic_state: states.ThermodynamicState,
hybrid_sampler_state: states.SamplerState,
hybrid_factory: HybridTopologyFactory,
nonalchemical_thermodynamic_state: states.ThermodynamicState,
lambda_state: int) -> tuple:
"""
Compute the perturbation of transforming the given hybrid equilibrium result into the system for the given nonalchemical_thermodynamic_state
Parameters
----------
alchemical_thermodynamic_state: states.ThermodynamicState
alchemical thermostate
growth_thermodynamic_state : states.ThermodynamicState
hybrid_sampler_state: states.SamplerState
sampler state for the alchemical thermodynamic_state
hybrid_factory : HybridTopologyFactory
Hybrid factory necessary for getting the positions of the nonalchemical system
nonalchemical_thermodynamic_state : states.ThermodynamicState
ThermodynamicState of the nonalchemical system
lambda_state : int
Whether this is lambda 0 or 1
Returns
-------
valence_energy: float
reduced potential energy of the valence contribution of the alternate endstate
nonalchemical_reduced_potential : float
reduced potential energy of the nonalchemical endstate
hybrid_reduced_potential: float
reduced potential energy of the alchemical endstate
"""
#get the objects we need to begin
hybrid_reduced_potential = compute_reduced_potential(
alchemical_thermodynamic_state, hybrid_sampler_state)
hybrid_positions = hybrid_sampler_state.positions
#get the positions for the nonalchemical system
if lambda_state == 0:
nonalchemical_positions = hybrid_factory.old_positions(
hybrid_positions)
nonalchemical_alternate_positions = hybrid_factory.new_positions(
hybrid_positions)
elif lambda_state == 1:
nonalchemical_positions = hybrid_factory.new_positions(
hybrid_positions)
nonalchemical_alternate_positions = hybrid_factory.old_positions(
hybrid_positions)
else:
raise ValueError("lambda_state must be 0 or 1")
nonalchemical_sampler_state = states.SamplerState(
nonalchemical_positions, box_vectors=hybrid_sampler_state.box_vectors)
nonalchemical_alternate_sampler_state = states.SamplerState(
nonalchemical_alternate_positions,
box_vectors=hybrid_sampler_state.box_vectors)
nonalchemical_reduced_potential = compute_reduced_potential(
nonalchemical_thermodynamic_state, nonalchemical_sampler_state)
#now for the growth system (set at lambda 0 or 1) so we can get the valence energy
if growth_thermodynamic_state:
valence_energy = compute_reduced_potential(
growth_thermodynamic_state, nonalchemical_alternate_sampler_state)
else:
valence_energy = 0.0
#now, the corrected energy of the system (for dispersion correction) is the nonalchemical_reduced_potential + valence_energy
return (valence_energy, nonalchemical_reduced_potential,
hybrid_reduced_potential)
def main():
parser = argparse.ArgumentParser(
description=
'Compute a potential of mean force (PMF) for porin permeation.')
parser.add_argument('--index',
dest='index',
action='store',
type=int,
help='Index of ')
parser.add_argument('--output',
dest='output_filename',
action='store',
default='output.nc',
help='output netcdf filename (default: output.nc)')
args = parser.parse_args()
index = args.index
output_filename = args.output_filename
logger = logging.getLogger(__name__)
logging.root.setLevel(logging.DEBUG)
logging.basicConfig(level=logging.DEBUG)
yank.utils.config_root_logger(verbose=True, log_file_path=None)
# Configure ContextCache, platform and precision
from yank.experiment import ExperimentBuilder
platform = ExperimentBuilder._configure_platform('CUDA', 'mixed')
try:
openmmtools.cache.global_context_cache.platform = platform
except RuntimeError:
# The cache has been already used. Empty it before switching platform.
openmmtools.cache.global_context_cache.empty()
openmmtools.cache.global_context_cache.platform = platform
# Topology
pdbx = app.PDBxFile('mem_prot_md_system.pdbx')
# This system contains the CVforce with parameters different than zero
with open('openmm_system.xml', 'r') as infile:
openmm_system = XmlSerializer.deserialize(infile.read())
####### Indexes of configurations in trajectory ############################
configs = [
39, 141, 276, 406, 562, 668, 833, 1109, 1272, 1417, 1456, 1471, 1537,
1645, 1777, 1882
]
####### Indexes of states for series of replica exchange simulations #######
limits = [(0, 9), (10, 19), (20, 29), (30, 39), (40, 49), (50, 59),
(60, 69), (70, 79), (80, 89), (90, 99), (100, 109), (110, 119),
(120, 129), (130, 139), (140, 149), (150, 159)]
####### Reading positions from mdtraj trajectory ###########################
topology = md.Topology.from_openmm(pdbx.topology)
t = md.load('../../steered_md/comp7/forward/seed_0/steered_forward.nc',
top=topology)
positions = t.openmm_positions(configs[index])
thermodynamic_state_deserialized = states.ThermodynamicState(
system=openmm_system,
temperature=310 * unit.kelvin,
pressure=1.0 * unit.atmospheres)
sampler_state = states.SamplerState(
positions=positions,
box_vectors=t.unitcell_vectors[configs[index], :, :] * unit.nanometer)
logger.debug(type(sampler_state))
move = mcmc.LangevinDynamicsMove(timestep=2 * unit.femtosecond,
collision_rate=1.0 / unit.picoseconds,
n_steps=500,
reassign_velocities=False)
simulation = ReplicaExchangeSampler(mcmc_moves=move,
number_of_iterations=1)
analysis_particle_indices = topology.select(
'(protein and mass > 3.0) or (resname MER and mass > 3.0)')
reporter = MultiStateReporter(
output_filename,
checkpoint_interval=2000,
analysis_particle_indices=analysis_particle_indices)
first, last = limits[index]
# Initialize compound thermodynamic states
protocol = {
'lambda_restraints': [i / 159 for i in range(first, last + 1)],
'K_parallel': [
1250 * unit.kilojoules_per_mole / unit.nanometer**2
for i in range(first, last + 1)
],
'Kmax': [
500 * unit.kilojoules_per_mole / unit.nanometer**2
for i in range(first, last + 1)
],
'Kmin': [
500 * unit.kilojoules_per_mole / unit.nanometer**2
for i in range(first, last + 1)
]
}
my_composable_state = MyComposableState.from_system(openmm_system)
compound_states = states.create_thermodynamic_state_protocol(
thermodynamic_state_deserialized,
protocol=protocol,
composable_states=[my_composable_state])
simulation.create(thermodynamic_states=compound_states,
sampler_states=[sampler_state],
storage=reporter)
simulation.equilibrate(50, mcmc_moves=move)
simulation.run()
ts = simulation._thermodynamic_states[0]
context, _ = openmmtools.cache.global_context_cache.get_context(ts)
files_names = [
'state_{}_{}.log'.format(index, i) for i in range(first, last + 1)
]
files = []
for i, file in enumerate(files_names):
files.append(open(file, 'w'))
mpi.distribute(write_cv,
range(simulation.n_replicas),
context,
simulation,
files,
send_results_to=None)
for i in range(10000):
simulation.extend(n_iterations=2)
mpi.distribute(write_cv,
range(simulation.n_replicas),
context,
simulation,
files,
send_results_to=None)
topology = md.Topology.from_openmm(pdb.topology)
ligand_atoms = topology.select('(resname CM7)')
factory = AbsoluteAlchemicalFactory(
consistent_exceptions=False, disable_alchemical_dispersion_correction=True)
alchemical_region = AlchemicalRegion(alchemical_atoms=ligand_atoms)
alchemical_system = factory.create_alchemical_system(openmm_system,
alchemical_region)
thermodynamic_state_ref = states.ThermodynamicState(
system=alchemical_system,
temperature=310 * unit.kelvin,
pressure=1.0 * unit.atmospheres)
sampler_state = states.SamplerState(
positions=pdb.positions, box_vectors=pdb.topology.getPeriodicBoxVectors())
nstates = 160
#lambda_sterics= [1.00, 0.95, 0.90, 0.85, 0.80, 0.75, 0.70, 0.65, 0.59, 0.54, 0.47, 0.41, 0.36, 0.31, 0.27, 0.25, 0.225, 0.20, 0.18, 0.15, 0.13,
#0.11, 0.06, 0.03, 0.00]
#lambda_sterics =[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.00, 0.95, 0.90, 0.85, 0.80, 0.75, 0.70, 0.65, 0.59, 0.54, 0.47,
#0.41, 0.36, 0.31, 0.27, 0.25, 0.225, 0.20, 0.18, 0.15, 0.13, 0.11, 0.06, 0.03, 0.00]
#lambda_electrostatics=[1.0, 0.9, 0.79, 0.69, 0.58, 0.46, 0.32, 0.18, 0.05, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
#0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
lambda_electrostatics = [1.0, 0.95, 0.9, 0.85, 0.8]
lambda_sterics = [1.0, 1.0, 1.0, 1.0, 1.0]
lambda_restraints = [i / nstates for i in range(12, 112)]
K_parallel = [
6000 * unit.kilojoules_per_mole / unit.nanometer**2
for i in range(len(lambda_restraints))
def _setup(self):
"""
Set up calculation.
"""
# Signal that system has now been set up
if self._setup_complete:
raise Exception(
"System has already been set up---cannot run again.")
self._setup_complete = True
# Compute thermal energy and inverse temperature
self.kT = kB * self.temperature
self.beta = 1.0 / self.kT
# Create the system
self.system = self._create_system()
# Add a barostat
# TODO: Is this necessary, sicne ThermodynamicState handles this automatically? It may not correctly handle MonteCarloAnisotropicBarostat.
self._add_barostat()
# Restrain protein atoms in space
# TODO: Allow protein atom selection to be configured
selection = '((residue 342 and resname GLY) or (residue 97 and resname ASP) or (residue 184 and resname SER)) and (name CA)'
protein_atoms_to_restrain = self.mdtraj_topology.select(selection)
#self._restrain_protein(protein_atoms_to_restrain)
# Create SamplerState for initial conditions
self.sampler_state = states.SamplerState(
positions=self.grofile.positions,
box_vectors=self.grofile.getPeriodicBoxVectors())
# Create reference thermodynamic state
self.reference_thermodynamic_state = states.ThermodynamicState(
system=self.system,
temperature=self.temperature,
pressure=self.pressure)
# Anneal ligand into binding site
if self.anneal_ligand:
self._anneal_ligand()
# Create ThermodynamicStates for umbrella sampling along pore
self.thermodynamic_states = self._create_thermodynamic_states(
self.reference_thermodynamic_state)
# Minimize initial thermodynamic state
# TODO: Select initial thermodynamic state based on which state has minimum energy
initial_state_index = 0
self.sampler_state = self._minimize_sampler_state(
self.thermodynamic_states[initial_state_index], self.sampler_state)
# Set up simulation
from yank.multistate import SAMSSampler, MultiStateReporter
# TODO: Change this to LangevinSplittingDynamicsMove
move = mcmc.LangevinDynamicsMove(timestep=self.timestep,
collision_rate=self.collision_rate,
n_steps=self.n_steps_per_iteration,
reassign_velocities=False)
self.simulation = SAMSSampler(
mcmc_moves=move,
number_of_iterations=self.n_iterations,
online_analysis_interval=None,
gamma0=self.gamma0,
flatness_threshold=self.flatness_threshold)
self.reporter = MultiStateReporter(
self.output_filename,
checkpoint_interval=self.checkpoint_interval,
analysis_particle_indices=self.analysis_particle_indices)
self.simulation.create(
thermodynamic_states=self.thermodynamic_states,
unsampled_thermodynamic_states=[
self.reference_thermodynamic_state
],
sampler_states=[self.sampler_state],
initial_thermodynamic_states=[initial_state_index],
storage=self.reporter)
def validate_endstate_energies(topology_proposal,
htf,
added_energy,
subtracted_energy,
beta=1.0 / kT,
ENERGY_THRESHOLD=1e-6):
"""
Function to validate that the difference between the nonalchemical versus alchemical state at lambda = 0,1 is
equal to the difference in valence energy (forward and reverse).
Parameters
----------
topology_proposal : perses.topology_proposal.TopologyProposal object
top_proposal for relevant transformation
htf : perses.new_relative.HybridTopologyFactory object
hybrid top factory for setting alchemical hybrid states
added_energy : float
reduced added valence energy
subtracted_energy: float
reduced subtracted valence energy
Returns
-------
zero_state_energy_difference : float
reduced potential difference of the nonalchemical and alchemical lambda = 0 state (corrected for valence energy).
one_state_energy_difference : float
reduced potential difference of the nonalchemical and alchemical lambda = 1 state (corrected for valence energy).
"""
import copy
#import openmmtools.cache as cache
#context_cache = cache.global_context_cache
#create copies of old/new systems and set the dispersion correction
top_proposal = copy.deepcopy(topology_proposal)
top_proposal._old_system.getForce(3).setUseDispersionCorrection(False)
top_proposal._new_system.getForce(3).setUseDispersionCorrection(False)
#create copy of hybrid system, define old and new positions, and turn off dispersion correction
hybrid_system = copy.deepcopy(htf.hybrid_system)
hybrid_system_n_forces = hybrid_system.getNumForces()
for force_index in range(hybrid_system_n_forces):
forcename = hybrid_system.getForce(force_index).__class__.__name__
if forcename == 'NonbondedForce':
hybrid_system.getForce(force_index).setUseDispersionCorrection(
False)
old_positions, new_positions = htf._old_positions, htf._new_positions
#generate endpoint thermostates
nonalch_zero, nonalch_one, alch_zero, alch_one = generate_endpoint_thermodynamic_states(
hybrid_system, top_proposal)
# compute reduced energies
#for the nonalchemical systems...
attrib_list = [(nonalch_zero, old_positions,
top_proposal._old_system.getDefaultPeriodicBoxVectors()),
(alch_zero, htf._hybrid_positions,
hybrid_system.getDefaultPeriodicBoxVectors()),
(alch_one, htf._hybrid_positions,
hybrid_system.getDefaultPeriodicBoxVectors()),
(nonalch_one, new_positions,
top_proposal._new_system.getDefaultPeriodicBoxVectors())]
rp_list = []
platform = openmm.Platform.getPlatformByName('Reference')
for (state, pos, box_vectors) in attrib_list:
#print("\t\t\t{}".format(state))
integrator = openmm.VerletIntegrator(1.0 * unit.femtoseconds)
context = state.create_context(integrator, platform)
samplerstate = states.SamplerState(positions=pos,
box_vectors=box_vectors)
samplerstate.apply_to_context(context)
rp = state.reduced_potential(context)
rp_list.append(rp)
energy_comps = compute_potential_components(context)
for name, force in energy_comps:
print("\t\t\t{}: {}".format(name, force))
print(
f'added forces:{sum([energy*beta for name, energy in energy_comps])}'
)
print(f'rp: {rp}')
del context, integrator
#print(f"added_energy: {added_energy}; subtracted_energy: {subtracted_energy}")
nonalch_zero_rp, alch_zero_rp, alch_one_rp, nonalch_one_rp = rp_list[
0], rp_list[1], rp_list[2], rp_list[3]
assert abs(
nonalch_zero_rp - alch_zero_rp + added_energy
) < ENERGY_THRESHOLD, f"The zero state alchemical and nonalchemical energy absolute difference {abs(nonalch_zero_rp - alch_zero_rp + added_energy)} is greater than the threshold of {ENERGY_THRESHOLD}."
assert abs(
nonalch_one_rp - alch_one_rp + subtracted_energy
) < ENERGY_THRESHOLD, f"The one state alchemical and nonalchemical energy absolute difference {abs(nonalch_one_rp - alch_one_rp + subtracted_energy)} is greater than the threshold of {ENERGY_THRESHOLD}."
return abs(nonalch_zero_rp - alch_zero_rp +
added_energy), abs(nonalch_one_rp - alch_one_rp +
subtracted_energy)
def _anneal_ligand(self, structure, index):
"""
Anneal ligand interactions to clean up clashes.
Returns
-------
positions : unit.Quantity
Positions of all atoms after annealing the ligand
"""
reference_system = copy.deepcopy(self.phases[index].system)
protein = self.phases[index].mdtraj_top.select(
f'protein and backbone and type CA').tolist()
rmsd = openmm.RMSDForce(
self.phases[index].structure.positions,
protein + self.phases[index].lg1_idx + self.phases[index].lg2_idx)
energy_expression = 'step(dRMSD) * (K_RMSD/2)*dRMSD^2; dRMSD = (RMSD-RMSD0);'
restraint_force = openmm.CustomCVForce(energy_expression)
restraint_force.addCollectiveVariable('RMSD', rmsd)
restraint_force.addGlobalParameter(
'K_RMSD', 2 * unit.kilocalories_per_mole / (unit.angstroms**2))
restraint_force.addGlobalParameter('RMSD0', 2 * unit.angstroms)
reference_system.addForce(restraint_force)
alchemical_system = self._alchemically_modify_ligand(
reference_system, index)
from openmmtools.alchemy import AlchemicalState
alchemical_state_zero = mmtools.alchemy.AlchemicalState.from_system(
alchemical_system, parameters_name_suffix='zero')
alchemical_state_one = mmtools.alchemy.AlchemicalState.from_system(
alchemical_system, parameters_name_suffix='one')
thermodynamic_state = states.ThermodynamicState(
system=alchemical_system, temperature=300 * unit.kelvin)
composable_states = [alchemical_state_zero, alchemical_state_one]
compound_states = states.CompoundThermodynamicState(
thermodynamic_state, composable_states=composable_states)
sampler_state = states.SamplerState(
positions=structure.positions,
box_vectors=structure.topology.getPeriodicBoxVectors())
# Anneal
n_annealing_steps = 1000
integrator = openmm.LangevinIntegrator(300 * unit.kelvin,
90.0 / unit.picoseconds,
1.0 * unit.femtoseconds)
context, integrator = mmtools.cache.global_context_cache.get_context(
compound_states, integrator)
sampler_state.apply_to_context(context)
compound_states.lambda_sterics_one = 0.0
compound_states.lambda_electrostatics_one = 0.0
compound_states.apply_to_context(context)
print('Annealing sterics of ligand 1...')
for step in progressbar.progressbar(range(n_annealing_steps)):
compound_states.lambda_sterics_zero = float(step) / float(
n_annealing_steps)
compound_states.lambda_electrostatics_zero = 0.0
compound_states.apply_to_context(context)
integrator.step(1)
print('Annealing electrostatics of ligand 1...')
for step in progressbar.progressbar(range(n_annealing_steps)):
compound_states.lambda_sterics_zero = 1.0
compound_states.lambda_electrostatics_zero = float(step) / float(
n_annealing_steps)
compound_states.apply_to_context(context)
integrator.step(1)
compound_states.lambda_sterics_zero = 1.0
compound_states.lambda_electrostatics_zero = 1.0
compound_states.apply_to_context(context)
print('Annealing sterics of ligand 2...')
for step in progressbar.progressbar(range(n_annealing_steps)):
compound_states.lambda_sterics_one = float(step) / float(
n_annealing_steps)
compound_states.lambda_electrostatics_one = 0.0
compound_states.apply_to_context(context)
integrator.step(1)
print('Annealing electrostatics of ligand 2...')
for step in progressbar.progressbar(range(n_annealing_steps)):
compound_states.lambda_sterics_one = 1.0
compound_states.lambda_electrostatics_one = float(step) / float(
n_annealing_steps)
compound_states.apply_to_context(context)
integrator.step(1)
compound_states.apply_to_context(context)
sampler_state.update_from_context(context)
# Compute the final energy of the system.
final_energy = thermodynamic_state.reduced_potential(context)
print('final alchemical energy {:8.3f}kT'.format(final_energy))
return sampler_state.positions
def validate_endstate_energies(topology_proposal,
htf,
added_energy,
subtracted_energy,
beta=1.0 / kT,
ENERGY_THRESHOLD=1e-6,
platform=DEFAULT_PLATFORM,
trajectory_directory=None):
"""
Function to validate that the difference between the nonalchemical versus alchemical state at lambda = 0,1 is
equal to the difference in valence energy (forward and reverse).
Parameters
----------
topology_proposal : perses.topology_proposal.TopologyProposal object
top_proposal for relevant transformation
htf : perses.new_relative.HybridTopologyFactory object
hybrid top factory for setting alchemical hybrid states
added_energy : float
reduced added valence energy
subtracted_energy: float
reduced subtracted valence energy
Returns
-------
zero_state_energy_difference : float
reduced potential difference of the nonalchemical and alchemical lambda = 0 state (corrected for valence energy).
one_state_energy_difference : float
reduced potential difference of the nonalchemical and alchemical lambda = 1 state (corrected for valence energy).
"""
import copy
#import openmmtools.cache as cache
#context_cache = cache.global_context_cache
from perses.dispersed.utils import configure_platform
from perses.utils import data
platform = configure_platform(platform.getName(),
fallback_platform_name='Reference',
precision='double')
#create copies of old/new systems and set the dispersion correction
top_proposal = copy.deepcopy(topology_proposal)
forces = {
top_proposal._old_system.getForce(index).__class__.__name__:
top_proposal._old_system.getForce(index)
for index in range(top_proposal._old_system.getNumForces())
}
forces['NonbondedForce'].setUseDispersionCorrection(False)
forces = {
top_proposal._new_system.getForce(index).__class__.__name__:
top_proposal._new_system.getForce(index)
for index in range(top_proposal._new_system.getNumForces())
}
forces['NonbondedForce'].setUseDispersionCorrection(False)
#create copy of hybrid system, define old and new positions, and turn off dispersion correction
hybrid_system = copy.deepcopy(htf.hybrid_system)
hybrid_system_n_forces = hybrid_system.getNumForces()
for force_index in range(hybrid_system_n_forces):
forcename = hybrid_system.getForce(force_index).__class__.__name__
if forcename == 'NonbondedForce':
hybrid_system.getForce(force_index).setUseDispersionCorrection(
False)
old_positions, new_positions = htf._old_positions, htf._new_positions
#generate endpoint thermostates
nonalch_zero, nonalch_one, alch_zero, alch_one = generate_endpoint_thermodynamic_states(
hybrid_system, top_proposal)
# compute reduced energies
#for the nonalchemical systems...
attrib_list = [('real-old', nonalch_zero, old_positions,
top_proposal._old_system.getDefaultPeriodicBoxVectors()),
('hybrid-old', alch_zero, htf._hybrid_positions,
hybrid_system.getDefaultPeriodicBoxVectors()),
('hybrid-new', alch_one, htf._hybrid_positions,
hybrid_system.getDefaultPeriodicBoxVectors()),
('real-new', nonalch_one, new_positions,
top_proposal._new_system.getDefaultPeriodicBoxVectors())]
rp_list = []
for (state_name, state, pos, box_vectors) in attrib_list:
integrator = openmm.VerletIntegrator(1.0 * unit.femtoseconds)
context = state.create_context(integrator, platform)
samplerstate = states.SamplerState(positions=pos,
box_vectors=box_vectors)
samplerstate.apply_to_context(context)
rp = state.reduced_potential(context)
rp_list.append(rp)
energy_comps = compute_potential_components(context)
for name, force in energy_comps:
print("\t\t\t{}: {}".format(name, force))
_logger.debug(
f'added forces:{sum([energy for name, energy in energy_comps])}')
_logger.debug(f'rp: {rp}')
if trajectory_directory is not None:
_logger.info(
f'Saving {state_name} state xml to {trajectory_directory}/{state_name}-state.gz'
)
state = context.getState(getPositions=True,
getVelocities=True,
getForces=True,
getEnergy=True,
getParameters=True)
data.serialize(state,
f'{trajectory_directory}-{state_name}-state.gz')
del context, integrator
nonalch_zero_rp, alch_zero_rp, alch_one_rp, nonalch_one_rp = rp_list[
0], rp_list[1], rp_list[2], rp_list[3]
ratio = abs((nonalch_zero_rp - alch_zero_rp + added_energy) /
(nonalch_zero_rp + alch_zero_rp + added_energy))
assert ratio < ENERGY_THRESHOLD, f"The ratio in energy difference for the ZERO state is {ratio}.\n This is greater than the threshold of {ENERGY_THRESHOLD}.\n real-zero: {nonalch_zero_rp} \n alc-zero: {alch_zero_rp} \nadded-valence: {added_energy}"
ratio = abs((nonalch_one_rp - alch_one_rp + subtracted_energy) /
(nonalch_one_rp + alch_one_rp + subtracted_energy))
assert ratio < ENERGY_THRESHOLD, f"The ratio in energy difference for the ONE state is {ratio}.\n This is greater than the threshold of {ENERGY_THRESHOLD}.\n real-one: {nonalch_one_rp} \n alc-one: {alch_one_rp} \nsubtracted-valence: {subtracted_energy}"
return abs(nonalch_zero_rp - alch_zero_rp +
added_energy), abs(nonalch_one_rp - alch_one_rp +
subtracted_energy)
def setup_class(cls):
"""Shared test cases for the suite."""
temperature = 300 * unit.kelvin
# Default protocols for tests.
cls.protocol = dict(lambda_electrostatics=[1.0, 0.5, 0.0, 0.0, 0.0],
lambda_sterics=[1.0, 1.0, 1.0, 0.5, 0.0])
cls.restrained_protocol = dict(lambda_electrostatics=[1.0, 1.0, 0.0, 0.0],
lambda_sterics=[1.0, 1.0, 1.0, 0.0],
lambda_restraints=[0.0, 1.0, 1.0, 1.0])
# Ligand-receptor in implicit solvent.
test_system = testsystems.HostGuestImplicit()
thermodynamic_state = states.ThermodynamicState(test_system.system,
temperature=temperature)
sampler_state = states.SamplerState(positions=test_system.positions, box_vectors=test_system.system.getDefaultPeriodicBoxVectors())
topography = Topography(test_system.topology, ligand_atoms='resname B2')
cls.host_guest_implicit = ('Host-guest implicit', thermodynamic_state, sampler_state, topography)
# Ligand-receptor in explicit solvent.
test_system = testsystems.HostGuestExplicit()
thermodynamic_state = states.ThermodynamicState(test_system.system,
temperature=temperature)
positions = test_system.positions
box_vectors = test_system.system.getDefaultPeriodicBoxVectors()
sampler_state = states.SamplerState(positions=positions, box_vectors=box_vectors)
topography = Topography(test_system.topology, ligand_atoms='resname B2')
cls.host_guest_explicit = ('Host-guest explicit', thermodynamic_state, sampler_state, topography)
# Peptide solvated in explicit solvent.
test_system = testsystems.AlanineDipeptideExplicit()
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
top = mdtraj.Topology.from_openmm(testsystem.topology)
trj = mdtraj.Trajectory([testsystem.positions / unit.nanometers], top)
trj.save(stem + '.pdb')
# save system as .xml
serialized_system = mm.openmm.XmlSerializer.serialize(testsystem.system)
with open(stem + '.xml', 'w') as fp:
print(serialized_system, file=fp)
n_replicas = 3 # Number of temperature replicas.
T_min = 298.0 * unit.kelvin # Minimum temperature.
T_max = 600.0 * unit.kelvin # Maximum temperature.
reference_state = states.ThermodynamicState(system=testsystem.system,
temperature=T_min)
move = mcmc.GHMCMove(timestep=2.0 * unit.femtoseconds, n_steps=50)
sampler = ParallelTemperingSampler(mcmc_moves=move,
number_of_iterations=float('inf'),
online_analysis_interval=None)
storage_path = stem + '.nc'
reporter = MultiStateReporter(storage_path, checkpoint_interval=1)
sampler.create(reference_state,
states.SamplerState(testsystem.positions),
reporter,
min_temperature=T_min,
max_temperature=T_max,
n_temperatures=n_replicas)
sampler.run(n_iterations=10)
