def test_position_output():
"""
Test that the hybrid returns the correct positions for the new and old systems after construction
"""
from perses.annihilation.new_relative import HybridTopologyFactory
import numpy as np
#generate topology proposal
topology_proposal, old_positions, new_positions = utils.generate_vacuum_topology_proposal()
factory = HybridTopologyFactory(topology_proposal, old_positions, new_positions)
old_positions_factory = factory.old_positions(factory.hybrid_positions)
new_positions_factory = factory.new_positions(factory.hybrid_positions)
assert np.all(np.isclose(old_positions.in_units_of(unit.nanometers), old_positions_factory.in_units_of(unit.nanometers)))
assert np.all(np.isclose(new_positions.in_units_of(unit.nanometers), new_positions_factory.in_units_of(unit.nanometers)))
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_position_output():
"""
Test that the hybrid returns the correct positions for the new and old systems after construction
"""
from perses.annihilation.new_relative import HybridTopologyFactory
import numpy as np
#generate topology proposal
topology_proposal, old_positions, new_positions = generate_topology_proposal(
)
factory = HybridTopologyFactory(topology_proposal, old_positions,
new_positions)
old_positions_factory = factory.old_positions(factory.hybrid_positions)
new_positions_factory = factory.new_positions(factory.hybrid_positions)
assert np.all(
np.isclose(old_positions.in_units_of(unit.nanometers),
old_positions_factory.in_units_of(unit.nanometers)))
assert np.all(
np.isclose(new_positions.in_units_of(unit.nanometers),
new_positions_factory.in_units_of(unit.nanometers)))
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 make_alchemical_system(self, topology_proposal, current_positions,
new_positions):
"""
Generate an alchemically-modified system at the correct atoms
based on the topology proposal. This method generates a hybrid system using the new
HybridTopologyFactory. It memoizes so that calling multiple times (within a recent time period)
will immediately return a cached object.
Arguments
---------
topology_proposal : perses.rjmc.TopologyProposal
Unmodified real system corresponding to appropriate leg of transformation.
current_positions : np.ndarray of float
Positions of "old" system
new_positions : np.ndarray of float
Positions of "new" system atoms
Returns
-------
hybrid_factory : perses.annihilation.new_relative.HybridTopologyFactory
a factory object containing the hybrid system
"""
try:
hybrid_factory = self._hybrid_cache[topology_proposal]
#If we've retrieved the factory from the cache, update it to include the relevant positions
hybrid_factory._old_positions = current_positions
hybrid_factory._new_positions = new_positions
hybrid_factory._compute_hybrid_positions()
except KeyError:
try:
hybrid_factory = HybridTopologyFactory(topology_proposal,
current_positions,
new_positions)
self._hybrid_cache[topology_proposal] = hybrid_factory
except:
hybrid_factory = None
return hybrid_factory
def check_alchemical_hybrid_elimination_bar(topology_proposal, old_positions, new_positions, ncmc_nsteps=50, n_iterations=50, NSIGMA_MAX=6.0, geometry=False):
"""
Check that the hybrid topology, where both endpoints are identical, returns a free energy within NSIGMA_MAX of 0.
Parameters
----------
topology_proposal
positions
ncmc_nsteps
NSIGMA_MAX
Returns
-------
"""
#make the hybrid topology factory:
factory = HybridTopologyFactory(topology_proposal, old_positions, new_positions)
platform = openmm.Platform.getPlatformByName("CUDA")
hybrid_system = factory.hybrid_system
hybrid_topology = factory.hybrid_topology
initial_hybrid_positions = factory.hybrid_positions
#alchemical functions
functions = {
'lambda_sterics' : '2*lambda * step(0.5 - lambda) + (1.0 - step(0.5 - lambda))',
'lambda_electrostatics' : '2*(lambda - 0.5) * step(lambda - 0.5)',
'lambda_bonds' : 'lambda',
'lambda_angles' : 'lambda',
'lambda_torsions' : 'lambda'
}
w_f = np.zeros(n_iterations)
w_r = np.zeros(n_iterations)
#make the alchemical integrators:
forward_integrator = NCMCGHMCAlchemicalIntegrator(temperature, hybrid_system, functions, nsteps=ncmc_nsteps, direction='insert')
forward_context = openmm.Context(hybrid_system, forward_integrator, platform)
print("Minimizing for forward protocol...")
forward_context.setPositions(initial_hybrid_positions)
for parm in functions.keys():
forward_context.setParameter(parm, 0.0)
openmm.LocalEnergyMinimizer.minimize(forward_context, maxIterations=10)
initial_state = forward_context.getState(getPositions=True, getEnergy=True)
print("The initial energy after minimization is %s" % str(initial_state.getPotentialEnergy()))
initial_forward_positions = initial_state.getPositions(asNumpy=True)
equil_positions = simulate_hybrid(hybrid_system,functions, 0.0, initial_forward_positions)
print("Beginning forward protocols")
#first, do forward protocol (lambda=0 -> 1)
with progressbar.ProgressBar(max_value=n_iterations) as bar:
for i in range(n_iterations):
equil_positions = simulate_hybrid(hybrid_system, functions, 0.0, equil_positions)
forward_context.setPositions(equil_positions)
forward_integrator.step(ncmc_nsteps)
w_f[i] = -1.0 * forward_integrator.getLogAcceptanceProbability(forward_context)
bar.update(i)
del forward_context, forward_integrator
reverse_integrator = NCMCGHMCAlchemicalIntegrator(temperature, hybrid_system, functions, nsteps=ncmc_nsteps, direction='delete')
print("Minimizing for reverse protocol...")
reverse_context = openmm.Context(hybrid_system, reverse_integrator, platform)
reverse_context.setPositions(initial_hybrid_positions)
for parm in functions.keys():
reverse_context.setParameter(parm, 1.0)
openmm.LocalEnergyMinimizer.minimize(reverse_context, maxIterations=10)
initial_state = reverse_context.getState(getPositions=True, getEnergy=True)
print("The initial energy after minimization is %s" % str(initial_state.getPotentialEnergy()))
initial_reverse_positions = initial_state.getPositions(asNumpy=True)
equil_positions = simulate_hybrid(hybrid_system,functions, 1.0, initial_reverse_positions, nsteps=1000)
#now, reverse protocol
print("Beginning reverse protocols...")
with progressbar.ProgressBar(max_value=n_iterations) as bar:
for i in range(n_iterations):
equil_positions = simulate_hybrid(hybrid_system,functions, 1.0, equil_positions)
reverse_context.setPositions(equil_positions)
reverse_integrator.step(ncmc_nsteps)
w_r[i] = -1.0 * reverse_integrator.getLogAcceptanceProbability(reverse_context)
bar.update(i)
del reverse_context, reverse_integrator
from pymbar import BAR
[df, ddf] = BAR(w_f, w_r)
print("df = %12.6f +- %12.5f kT" % (df, ddf))
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 perses.rjmc.topology_proposal import SmallMoleculeSetProposalEngine, TopologyProposal
from perses.annihilation.new_relative import HybridTopologyFactory
import simtk.openmm as openmm
from perses.tests.utils import createOEMolFromIUPAC, createSystemFromIUPAC
mol_name = "naphthalene"
ref_mol_name = "benzene"
mol = createOEMolFromIUPAC(mol_name)
m, system, positions, topology = createSystemFromIUPAC(mol_name)
refmol = createOEMolFromIUPAC(ref_mol_name)
#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
integrator = openmm.VerletIntegrator(1)
platform = openmm.Platform.getPlatformByName("Reference")
context = openmm.Context(alchemical_system, integrator, platform)
context.setPositions(alchemical_positions)
functions = {
'lambda_sterics':
'2*lambda * step(0.5 - lambda) + (1.0 - step(0.5 - lambda))',
'lambda_electrostatics': '2*(lambda - 0.5) * step(lambda - 0.5)',
'lambda_bonds': 'lambda',
'lambda_angles': 'lambda',
'lambda_torsions': 'lambda'
}
#set all to zero
for parm in functions.keys():
context.setParameter(parm, 0.0)
initial_energy = context.getState(getEnergy=True).getPotentialEnergy()
#set all to one
for parm in functions.keys():
context.setParameter(parm, 1.0)
final_energy = context.getState(getEnergy=True).getPotentialEnergy()
if np.abs(final_energy -
initial_energy) > 1.0e-6 * unit.kilojoule_per_mole:
raise Exception(
"The energy at the endpoints was not equal for molecule %s" %
mol_name)
def run_hybrid_endpoint_overlap(topology_proposal, current_positions,
new_positions):
"""
Test that the variance of the perturbation from lambda={0,1} to the corresponding nonalchemical endpoint is not
too large.
Parameters
----------
topology_proposal : perses.rjmc.TopologyProposal
TopologyProposal object describing the transformation
current_positions : np.array, unit-bearing
Positions of the initial system
new_positions : np.array, unit-bearing
Positions of the new system
Returns
-------
hybrid_endpoint_results : list
list of [df, ddf, N_eff] for 1 and 0
"""
#create the hybrid system:
hybrid_factory = HybridTopologyFactory(topology_proposal,
current_positions,
new_positions,
use_dispersion_correction=True)
#get the relevant thermodynamic states:
nonalchemical_zero_thermodynamic_state, nonalchemical_one_thermodynamic_state, lambda_zero_thermodynamic_state, lambda_one_thermodynamic_state = generate_thermodynamic_states(
hybrid_factory.hybrid_system, topology_proposal)
nonalchemical_thermodynamic_states = [
nonalchemical_zero_thermodynamic_state,
nonalchemical_one_thermodynamic_state
]
alchemical_thermodynamic_states = [
lambda_zero_thermodynamic_state, lambda_one_thermodynamic_state
]
#create an MCMCMove, BAOAB with default parameters
mc_move = mcmc.LangevinDynamicsMove()
initial_sampler_state = SamplerState(
hybrid_factory.hybrid_positions,
box_vectors=hybrid_factory.hybrid_system.getDefaultPeriodicBoxVectors(
))
hybrid_endpoint_results = []
for lambda_state in (0, 1):
result = run_endpoint_perturbation(
alchemical_thermodynamic_states[lambda_state],
nonalchemical_thermodynamic_states[lambda_state],
initial_sampler_state,
mc_move,
100,
hybrid_factory,
lambda_index=lambda_state)
print(result)
hybrid_endpoint_results.append(result)
return hybrid_endpoint_results
def check_hybrid_round_trip_elimination(topology_proposal,
positions,
ncmc_nsteps=50,
NSIGMA_MAX=6.0):
"""
Test the hybrid system by switching between lambda = 1 and lambda = 0, then using BAR to compute the free energy
difference. As the test is designed so that both endpoints are the same, the free energy difference should be zero.
Parameters
----------
topology_proposal : TopologyProposal
The topology proposal to test.
This must be a null transformation, where topology_proposal.old_system == topology_proposal.new_system
ncmc_steps : int, optional, default=50
Number of NCMC switching steps, or 0 for instantaneous switching.
NSIGMA_MAX : float, optional, default=6.0
"""
functions = {
'lambda_sterics': 'lambda',
'lambda_electrostatics': 'lambda',
'lambda_bonds': 'lambda',
'lambda_angles': 'lambda',
'lambda_torsions': 'lambda'
}
# Initialize engine
from perses.annihilation import NCMCGHMCAlchemicalIntegrator
from perses.annihilation.new_relative import HybridTopologyFactory
#The current and "proposed" positions are the same, since the molecule is not changed.
factory = HybridTopologyFactory(topology_proposal, positions, positions)
forward_integrator = NCMCGHMCAlchemicalIntegrator(temperature,
factory.hybrid_system,
functions,
nsteps=ncmc_nsteps,
direction='insert')
reverse_integrator = NCMCGHMCAlchemicalIntegrator(temperature,
factory.hybrid_system,
functions,
nsteps=ncmc_nsteps,
direction='delete')
platform = openmm.Platform.getPlatformByName("Reference")
forward_context = openmm.Context(factory.hybrid_system, forward_integrator,
platform)
reverse_context = openmm.Context(factory.hybrid_system, reverse_integrator,
platform)
# Make sure that old system and new system are identical.
if not (topology_proposal.old_system == topology_proposal.new_system):
raise Exception(
"topology_proposal must be a null transformation for this test (old_system == new_system)"
)
for (k, v) in topology_proposal.new_to_old_atom_map.items():
if k != v:
raise Exception(
"topology_proposal must be a null transformation for this test (retailed atoms must map onto themselves)"
)
nequil = 5 # number of equilibration iterations
niterations = 50 # number of round-trip switching trials
logP_work_n_f = np.zeros([niterations], np.float64)
for iteration in range(nequil):
positions = simulate_hybrid(factory.hybrid_system, functions, 0.0,
factory.hybrid_positions)
#do forward switching:
for iteration in range(niterations):
# Equilibrate
positions = simulate_hybrid(factory.hybrid_system, functions, 0.0,
factory.hybrid_positions)
# Check that positions are not NaN
if (np.any(np.isnan(positions / unit.angstroms))):
raise Exception("Positions became NaN during equilibration")
# Hybrid NCMC
forward_integrator.reset()
forward_context.setPositions(positions)
forward_integrator.step(ncmc_nsteps)
logP_work = forward_integrator.getTotalWork(forward_context)
# Check that positions are not NaN
if (np.any(np.isnan(positions / unit.angstroms))):
raise Exception("Positions became NaN on Hybrid NCMC switch")
# Store log probability associated with work
logP_work_n_f[iteration] = logP_work
logP_work_n_r = np.zeros([niterations], np.float64)
for iteration in range(nequil):
positions = simulate_hybrid(factory.hybrid_system, functions, 1.0,
factory.hybrid_positions)
#do forward switching:
for iteration in range(niterations):
# Equilibrate
positions = simulate_hybrid(factory.hybrid_system, functions, 1.0,
factory.hybrid_positions)
# Check that positions are not NaN
if (np.any(np.isnan(positions / unit.angstroms))):
raise Exception("Positions became NaN during equilibration")
# Hybrid NCMC
reverse_integrator.reset()
reverse_context.setPositions(positions)
reverse_integrator.step(ncmc_nsteps)
logP_work = reverse_integrator.getTotalWork(forward_context)
# Check that positions are not NaN
if (np.any(np.isnan(positions / unit.angstroms))):
raise Exception("Positions became NaN on Hybrid NCMC switch")
# Store log probability associated with work
logP_work_n_r[iteration] = logP_work
work_f = -logP_work_n_f
work_r = -logP_work_n_r
from pymbar import BAR
[df, ddf] = BAR(work_f, work_r)
print("df = %12.6f +- %12.5f kT" % (df, ddf))
if (abs(df) > NSIGMA_MAX * ddf):
msg = 'Delta F (%d steps switching) = %f +- %f kT; should be within %f sigma of 0\n' % (
ncmc_nsteps, df, ddf, NSIGMA_MAX)
msg += 'logP_work_n:\n'
msg += str(work_f) + '\n'
msg += str(work_r) + '\n'
raise Exception(msg)
def check_alchemical_hybrid_elimination_bar(topology_proposal,
positions,
ncmc_nsteps=50,
NSIGMA_MAX=6.0,
geometry=False):
"""
Check that the hybrid topology, where both endpoints are identical, returns a free energy within NSIGMA_MAX of 0.
Parameters
----------
topology_proposal
positions
ncmc_nsteps
NSIGMA_MAX
Returns
-------
"""
from perses.annihilation import NCMCGHMCAlchemicalIntegrator
from perses.annihilation.new_relative import HybridTopologyFactory
#make the hybrid topology factory:
factory = HybridTopologyFactory(topology_proposal, positions, positions)
platform = openmm.Platform.getPlatformByName("Reference")
hybrid_system = factory.hybrid_system
hybrid_topology = factory.hybrid_topology
initial_hybrid_positions = factory.hybrid_positions
n_iterations = 50 #number of times to do NCMC protocol
#alchemical functions
functions = {
'lambda_sterics':
'2*lambda * step(0.5 - lambda) + (1.0 - step(0.5 - lambda))',
'lambda_electrostatics': '2*(lambda - 0.5) * step(lambda - 0.5)',
'lambda_bonds': 'lambda',
'lambda_angles': 'lambda',
'lambda_torsions': 'lambda'
}
w_f = np.zeros(n_iterations)
w_r = np.zeros(n_iterations)
#make the alchemical integrators:
#forward_integrator = NCMCGHMCAlchemicalIntegrator(temperature, hybrid_system, functions, nsteps=ncmc_nsteps, direction='insert')
#reverse_integrator = NCMCGHMCAlchemicalIntegrator(temperature, hybrid_system, functions, nsteps=ncmc_nsteps, direction='delete')
#first, do forward protocol (lambda=0 -> 1)
for i in range(n_iterations):
forward_integrator = NCMCGHMCAlchemicalIntegrator(temperature,
hybrid_system,
functions,
nsteps=ncmc_nsteps,
direction='insert')
equil_positions = simulate_hybrid(hybrid_system, functions, 0.0,
initial_hybrid_positions)
context = openmm.Context(hybrid_system, forward_integrator, platform)
context.setPositions(equil_positions)
forward_integrator.step(ncmc_nsteps)
w_f[i] = -1.0 * forward_integrator.getLogAcceptanceProbability(context)
print(i)
del context, forward_integrator
#now, reverse protocol
for i in range(n_iterations):
reverse_integrator = NCMCGHMCAlchemicalIntegrator(temperature,
hybrid_system,
functions,
nsteps=ncmc_nsteps,
direction='delete')
equil_positions = simulate_hybrid(hybrid_system, functions, 1.0,
initial_hybrid_positions)
context = openmm.Context(hybrid_system, reverse_integrator, platform)
context.setPositions(equil_positions)
reverse_integrator.step(ncmc_nsteps)
w_r[i] = -1.0 * reverse_integrator.getLogAcceptanceProbability(context)
print(i)
del context, reverse_integrator
from pymbar import BAR
[df, ddf] = BAR(w_f, w_r)
print("df = %12.6f +- %12.5f kT" % (df, ddf))
if (abs(df) > NSIGMA_MAX * ddf):
msg = 'Delta F (%d steps switching) = %f +- %f kT; should be within %f sigma of 0\n' % (
ncmc_nsteps, df, ddf, NSIGMA_MAX)
msg += 'logP_work_n:\n'
msg += str(w_f) + '\n'
msg += str(w_r) + '\n'
raise Exception(msg)
def __init__(self,
topology_proposal,
pos_old,
new_positions,
use_dispersion_correction=False,
forward_functions=None,
ncmc_nsteps=100,
nsteps_per_iteration=1,
concurrency=4,
platform_name="OpenCL",
temperature=300.0 * unit.kelvin,
trajectory_directory=None,
trajectory_prefix=None):
#construct the hybrid topology factory object
self._factory = HybridTopologyFactory(
topology_proposal,
pos_old,
new_positions,
use_dispersion_correction=use_dispersion_correction)
#use default functions if none specified
if forward_functions == None:
self._forward_functions = self.default_forward_functions
else:
self._forward_functions = forward_functions
#reverse functions to get a symmetric protocol
self._reverse_functions = {
param: param_formula.replace("lambda", "(1-lambda)")
for param, param_formula in self._forward_functions.items()
}
#set up some class attributes
self._hybrid_system = self._factory.hybrid_system
self._initial_hybrid_positions = self._factory.hybrid_positions
self._concurrency = concurrency
self._ncmc_nsteps = ncmc_nsteps
self._nsteps_per_iteration = nsteps_per_iteration
self._trajectory_prefix = trajectory_prefix
self._trajectory_directory = trajectory_directory
self._zero_endpoint_n_atoms = topology_proposal.n_atoms_old
self._one_endpoint_n_atoms = topology_proposal.n_atoms_new
#initialize lists for results
self._forward_nonequilibrium_trajectories = []
self._reverse_nonequilibrium_trajectories = []
self._forward_nonequilibrium_cumulative_works = []
self._reverse_nonequilibrium_cumulative_works = []
self._forward_nonequilibrium_results = []
self._reverse_nonequilibrium_results = []
self._forward_total_work = []
self._reverse_total_work = []
self._lambda_zero_reduced_potentials = []
self._lambda_one_reduced_potentials = []
self._nonalchemical_zero_endpt_reduced_potentials = []
self._nonalchemical_one_endpt_reduced_potentials = []
self._nonalchemical_zero_results = []
self._nonalchemical_one_results = []
#Set the number of times that the nonequilbrium move will have to be run in order to complete a protocol:
if self._ncmc_nsteps % self._nsteps_per_iteration != 0:
logging.warning(
"The number of ncmc steps is not divisible by the number of steps per iteration. You may not have a full protocol."
)
self._n_iterations_per_call = self._ncmc_nsteps // self._nsteps_per_iteration
#create the thermodynamic state
lambda_zero_alchemical_state = alchemy.AlchemicalState.from_system(
self._hybrid_system)
lambda_one_alchemical_state = copy.deepcopy(
lambda_zero_alchemical_state)
#ensure their states are set appropriately
lambda_zero_alchemical_state.set_alchemical_parameters(0.0)
lambda_one_alchemical_state.set_alchemical_parameters(0.0)
#create the base thermodynamic state with the hybrid system
self._thermodynamic_state = ThermodynamicState(self._hybrid_system,
temperature=temperature)
#Create thermodynamic states for the nonalchemical endpoints
self._nonalchemical_zero_thermodynamic_state = ThermodynamicState(
topology_proposal.old_system, temperature=temperature)
self._nonalchemical_one_thermodynamic_state = ThermodynamicState(
topology_proposal.new_system, temperature=temperature)
#Now create the compound states with different alchemical states
self._lambda_zero_thermodynamic_state = CompoundThermodynamicState(
self._thermodynamic_state,
composable_states=[lambda_zero_alchemical_state])
self._lambda_one_thermodynamic_state = CompoundThermodynamicState(
self._thermodynamic_state,
composable_states=[lambda_one_alchemical_state])
#create the forward and reverse integrators
self._forward_integrator = AlchemicalNonequilibriumLangevinIntegrator(
alchemical_functions=self._forward_functions,
nsteps_neq=ncmc_nsteps,
temperature=temperature)
self._reverse_integrator = AlchemicalNonequilibriumLangevinIntegrator(
alchemical_functions=self._reverse_functions,
nsteps_neq=ncmc_nsteps,
temperature=temperature)
#create the forward and reverse MCMoves
self._forward_ne_mc_move = NonequilibriumSwitchingMove(
self._forward_integrator, self._nsteps_per_iteration)
self._reverse_ne_mc_move = NonequilibriumSwitchingMove(
self._reverse_integrator, self._nsteps_per_iteration)
#create the equilibrium MCMove
self._equilibrium_mc_move = mcmc.LangevinSplittingDynamicsMove()
#set the SamplerState for the lambda 0 and 1 equilibrium simulations
self._lambda_one_sampler_state = SamplerState(
self._initial_hybrid_positions,
box_vectors=self._hybrid_system.getDefaultPeriodicBoxVectors())
self._lambda_zero_sampler_state = copy.deepcopy(
self._lambda_one_sampler_state)
#initialize by minimizing
self.minimize()
#initialize the trajectories for the lambda 0 and 1 equilibrium simulations
a_0, b_0, c_0, alpha_0, beta_0, gamma_0 = mdtrajutils.unitcell.box_vectors_to_lengths_and_angles(
*self._lambda_zero_sampler_state.box_vectors)
a_1, b_1, c_1, alpha_1, beta_1, gamma_1 = mdtrajutils.unitcell.box_vectors_to_lengths_and_angles(
*self._lambda_one_sampler_state.box_vectors)
self._lambda_zero_traj = md.Trajectory(
np.array(self._lambda_zero_sampler_state.positions),
self._factory.hybrid_topology,
unitcell_lengths=[a_0, b_0, c_0],
unitcell_angles=[alpha_0, beta_0, gamma_0])
self._lambda_one_traj = md.Trajectory(
np.array(self._lambda_one_sampler_state.positions),
self._factory.hybrid_topology,
unitcell_lengths=[a_1, b_1, c_1],
unitcell_angles=[alpha_1, beta_1, gamma_1])
class NonequilibriumSwitchingFEP(object):
"""
This class manages Nonequilibrium switching based relative free energy calculations, carried out on a distributed computing framework.
"""
default_forward_functions = {
'lambda_sterics': 'lambda',
'lambda_electrostatics': 'lambda',
'lambda_bonds': 'lambda',
'lambda_angles': 'lambda',
'lambda_torsions': 'lambda'
}
def __init__(self,
topology_proposal,
pos_old,
new_positions,
use_dispersion_correction=False,
forward_functions=None,
ncmc_nsteps=100,
nsteps_per_iteration=1,
concurrency=4,
platform_name="OpenCL",
temperature=300.0 * unit.kelvin,
trajectory_directory=None,
trajectory_prefix=None):
#construct the hybrid topology factory object
self._factory = HybridTopologyFactory(
topology_proposal,
pos_old,
new_positions,
use_dispersion_correction=use_dispersion_correction)
#use default functions if none specified
if forward_functions == None:
self._forward_functions = self.default_forward_functions
else:
self._forward_functions = forward_functions
#reverse functions to get a symmetric protocol
self._reverse_functions = {
param: param_formula.replace("lambda", "(1-lambda)")
for param, param_formula in self._forward_functions.items()
}
#set up some class attributes
self._hybrid_system = self._factory.hybrid_system
self._initial_hybrid_positions = self._factory.hybrid_positions
self._concurrency = concurrency
self._ncmc_nsteps = ncmc_nsteps
self._nsteps_per_iteration = nsteps_per_iteration
self._trajectory_prefix = trajectory_prefix
self._trajectory_directory = trajectory_directory
self._zero_endpoint_n_atoms = topology_proposal.n_atoms_old
self._one_endpoint_n_atoms = topology_proposal.n_atoms_new
#initialize lists for results
self._forward_nonequilibrium_trajectories = []
self._reverse_nonequilibrium_trajectories = []
self._forward_nonequilibrium_cumulative_works = []
self._reverse_nonequilibrium_cumulative_works = []
self._forward_nonequilibrium_results = []
self._reverse_nonequilibrium_results = []
self._forward_total_work = []
self._reverse_total_work = []
self._lambda_zero_reduced_potentials = []
self._lambda_one_reduced_potentials = []
self._nonalchemical_zero_endpt_reduced_potentials = []
self._nonalchemical_one_endpt_reduced_potentials = []
self._nonalchemical_zero_results = []
self._nonalchemical_one_results = []
#Set the number of times that the nonequilbrium move will have to be run in order to complete a protocol:
if self._ncmc_nsteps % self._nsteps_per_iteration != 0:
logging.warning(
"The number of ncmc steps is not divisible by the number of steps per iteration. You may not have a full protocol."
)
self._n_iterations_per_call = self._ncmc_nsteps // self._nsteps_per_iteration
#create the thermodynamic state
lambda_zero_alchemical_state = alchemy.AlchemicalState.from_system(
self._hybrid_system)
lambda_one_alchemical_state = copy.deepcopy(
lambda_zero_alchemical_state)
#ensure their states are set appropriately
lambda_zero_alchemical_state.set_alchemical_parameters(0.0)
lambda_one_alchemical_state.set_alchemical_parameters(0.0)
#create the base thermodynamic state with the hybrid system
self._thermodynamic_state = ThermodynamicState(self._hybrid_system,
temperature=temperature)
#Create thermodynamic states for the nonalchemical endpoints
self._nonalchemical_zero_thermodynamic_state = ThermodynamicState(
topology_proposal.old_system, temperature=temperature)
self._nonalchemical_one_thermodynamic_state = ThermodynamicState(
topology_proposal.new_system, temperature=temperature)
#Now create the compound states with different alchemical states
self._lambda_zero_thermodynamic_state = CompoundThermodynamicState(
self._thermodynamic_state,
composable_states=[lambda_zero_alchemical_state])
self._lambda_one_thermodynamic_state = CompoundThermodynamicState(
self._thermodynamic_state,
composable_states=[lambda_one_alchemical_state])
#create the forward and reverse integrators
self._forward_integrator = AlchemicalNonequilibriumLangevinIntegrator(
alchemical_functions=self._forward_functions,
nsteps_neq=ncmc_nsteps,
temperature=temperature)
self._reverse_integrator = AlchemicalNonequilibriumLangevinIntegrator(
alchemical_functions=self._reverse_functions,
nsteps_neq=ncmc_nsteps,
temperature=temperature)
#create the forward and reverse MCMoves
self._forward_ne_mc_move = NonequilibriumSwitchingMove(
self._forward_integrator, self._nsteps_per_iteration)
self._reverse_ne_mc_move = NonequilibriumSwitchingMove(
self._reverse_integrator, self._nsteps_per_iteration)
#create the equilibrium MCMove
self._equilibrium_mc_move = mcmc.LangevinSplittingDynamicsMove()
#set the SamplerState for the lambda 0 and 1 equilibrium simulations
self._lambda_one_sampler_state = SamplerState(
self._initial_hybrid_positions,
box_vectors=self._hybrid_system.getDefaultPeriodicBoxVectors())
self._lambda_zero_sampler_state = copy.deepcopy(
self._lambda_one_sampler_state)
#initialize by minimizing
self.minimize()
#initialize the trajectories for the lambda 0 and 1 equilibrium simulations
a_0, b_0, c_0, alpha_0, beta_0, gamma_0 = mdtrajutils.unitcell.box_vectors_to_lengths_and_angles(
*self._lambda_zero_sampler_state.box_vectors)
a_1, b_1, c_1, alpha_1, beta_1, gamma_1 = mdtrajutils.unitcell.box_vectors_to_lengths_and_angles(
*self._lambda_one_sampler_state.box_vectors)
self._lambda_zero_traj = md.Trajectory(
np.array(self._lambda_zero_sampler_state.positions),
self._factory.hybrid_topology,
unitcell_lengths=[a_0, b_0, c_0],
unitcell_angles=[alpha_0, beta_0, gamma_0])
self._lambda_one_traj = md.Trajectory(
np.array(self._lambda_one_sampler_state.positions),
self._factory.hybrid_topology,
unitcell_lengths=[a_1, b_1, c_1],
unitcell_angles=[alpha_1, beta_1, gamma_1])
def minimize(self, max_steps=50):
"""
Minimize both end states. This method updates the _sampler_state attributes for each lambda
Parameters
----------
max_steps : int, default 50
max number of steps for openmm minimizer.
"""
#Asynchronously invoke the tasks
minimized_lambda_zero_result = feptasks.minimize.delay(
self._lambda_zero_thermodynamic_state,
self._lambda_zero_sampler_state,
self._equilibrium_mc_move,
max_iterations=max_steps)
minimized_lambda_one_result = feptasks.minimize.delay(
self._lambda_one_thermodynamic_state,
self._lambda_one_sampler_state,
self._equilibrium_mc_move,
max_iterations=max_steps)
#now synchronously retrieve the results and save the sampler states.
self._lambda_zero_sampler_state = minimized_lambda_zero_result.get()
self._lambda_one_sampler_state = minimized_lambda_one_result.get()
def run(self, n_iterations=5, concurrency=1):
"""
Run one iteration of the nonequilibrium switching free energy calculations. This entails:
- 1 iteration of equilibrium at lambda=0 and lambda=1
- concurrency (parameter) many nonequilibrium trajectories in both forward and reverse
(e.g., if concurrency is 5, then 5 forward and 5 reverse protocols will be run)
- 1 iteration of equilibrium at lambda=0 and lambda=1
Parameters
----------
n_iterations : int, optional, default 5
The number of times to run the entire sequence described above
concurrency: int, default 1
The number of concurrent nonequilibrium protocols to run; note that with greater than one,
error estimation may be more complicated.
"""
for i in range(n_iterations):
self._run_equilibrium()
self._run_nonequilibrium(concurrency=concurrency,
n_iterations=self._n_iterations_per_call)
self._run_equilibrium()
if self._trajectory_directory:
self._write_equilibrium_trajectories(self._trajectory_directory,
self._trajectory_prefix)
def _run_equilibrium(self, n_iterations=1):
"""
Run one iteration of equilibrium at lambda=1 and lambda=0, and replace the current equilibrium sampler states
with the results of the equilibrium calculation, as well as extend the current equilibrium trajectories.
Parameters
----------
n_iterations : int, default 1
How many times to run the n_steps of equilibrium
"""
#run equilibrium for lambda=0 and lambda=1
lambda_zero_result = feptasks.run_equilibrium.delay(
self._lambda_zero_thermodynamic_state,
self._lambda_zero_sampler_state, self._equilibrium_mc_move,
self._factory.hybrid_topology, n_iterations)
lambda_one_result = feptasks.run_equilibrium.delay(
self._lambda_one_thermodynamic_state,
self._lambda_one_sampler_state, self._equilibrium_mc_move,
self._factory.hybrid_topology, n_iterations)
#retrieve the results of the calculation
self._lambda_zero_sampler_state, traj_zero_result, lambda_zero_reduced_potential = lambda_zero_result.get(
)
self._lambda_one_sampler_state, traj_one_result, lambda_one_reduced_potential = lambda_one_result.get(
)
#append the potential energies of the final frame of the trajectories
self._lambda_zero_reduced_potentials.append(
lambda_zero_reduced_potential)
self._lambda_one_reduced_potentials.append(
lambda_one_reduced_potential)
#Now create SamplerStates to generate the data for endpoint perturbations:
final_hybrid_positions_zero = self._lambda_zero_sampler_state.positions
final_hybrid_positions_one = self._lambda_one_sampler_state.positions
positions_zero = self._factory.old_positions(
final_hybrid_positions_zero)
positions_one = self._factory.new_positions(final_hybrid_positions_one)
#Create sampler states for each of these:
sampler_state_zero = SamplerState(
positions_zero,
box_vectors=self._lambda_zero_sampler_state.box_vectors)
sampler_state_one = SamplerState(
positions_one,
box_vectors=self._lambda_one_sampler_state.box_vectors)
#launch a task to compute the reduced potentials at these endpoints
self._nonalchemical_zero_results.append(
feptasks.compute_reduced_potential.delay(
self._nonalchemical_zero_thermodynamic_state,
sampler_state_zero))
self._nonalchemical_one_results.append(
feptasks.compute_reduced_potential.delay(
self._nonalchemical_one_thermodynamic_state,
sampler_state_one))
#join the trajectories to the reference trajectories, if the object exists,
#otherwise, simply create it
if self._lambda_zero_traj:
self._lambda_zero_traj = self._lambda_zero_traj.join(
traj_zero_result, check_topology=False)
else:
self._lambda_zero_traj = traj_zero_result
if self._lambda_one_traj:
self._lambda_one_traj = self._lambda_one_traj.join(
traj_one_result, check_topology=False)
else:
self._lambda_one_traj = traj_one_result
def _run_nonequilibrium(self, concurrency=1, n_iterations=1):
"""
Run concurrency-many nonequilibrium protocols in both directions. This method stores the result object, but does
not retrieve the results. Note that n_iterations is important, since in order to perform an entire switching trajectory
(from 0 to 1 or vice versa), we require that n_steps*n_iterations = protocol length
Parameters
----------
concurrency : int, default 1
The number of protocols to run in each direction simultaneously
n_iterations : int, default 1
The number of times to have the NE move applied. Note that as above if n_steps*n_iterations!=ncmc_nsteps,
the protocol will not be run properly.
"""
#set up the group object that will be used to compute the nonequilibrium results.
forward_protocol_group = celery.group(
feptasks.run_protocol.s(
self._lambda_zero_thermodynamic_state,
self._lambda_zero_sampler_state, self._forward_ne_mc_move,
self._factory.hybrid_topology, n_iterations)
for i in range(concurrency))
reverse_protocol_group = celery.group(
feptasks.run_protocol.s(
self._lambda_one_thermodynamic_state,
self._lambda_one_sampler_state, self._reverse_ne_mc_move,
self._factory.hybrid_topology, n_iterations)
for i in range(concurrency))
#get the result objects:
self._forward_nonequilibrium_results.append(
forward_protocol_group.apply_async())
self._reverse_nonequilibrium_results.append(
reverse_protocol_group.apply_async())
def retrieve_nonequilibrium_results(self):
"""
Retrieve any pending results that were generated by computations from the run() call. Note that this will block
until all have completed. This method will update the list of trajectories as well as the nonequilibrium work values.
"""
for result in self._forward_nonequilibrium_results:
result_group = result.join()
for result in result_group:
traj, cum_work = result
#we can take the final element as the total work
self._forward_total_work.append(cum_work[-1])
#we'll append the cumulative work and the trajectory to the appropriate lists
self._forward_nonequilibrium_cumulative_works.append(cum_work)
self._forward_nonequilibrium_trajectories.append(traj)
for result in self._reverse_nonequilibrium_results:
result_group = result.join()
for result in result_group:
traj, cum_work = result
#we can take the final element as the total work
self._reverse_total_work.append(cum_work[-1])
#we'll append the cumulative work and the trajectory to the appropriate lists
self._reverse_nonequilibrium_cumulative_works.append(cum_work)
self._reverse_nonequilibrium_trajectories.append(traj)
def write_nonequilibrium_trajectories(self, directory, file_prefix):
"""
Write out an MDTraj h5 file for each nonequilibrium trajectory. The files will be placed in
[directory]/file_prefix-[forward, reverse]-index.h5. This method will ensure that all pending
results are collected.
Parameters
----------
directory : str
The directory in which to place the files
file_prefix : str
A prefix for the filenames
"""
self.retrieve_nonequilibrium_results()
#loop through the forward trajectories
for index, forward_trajectory in enumerate(
self._forward_nonequilibrium_trajectories):
#construct the name for this file
full_filename = os.path.join(
directory, file_prefix + "forward" + str(index) + ".h5")
#save the trajectory
forward_trajectory.save_hdf5(full_filename)
#repeat for the reverse trajectories:
for index, reverse_trajectory in enumerate(
self._reverse_nonequilibrium_trajectories):
#construct the name for this file
full_filename = os.path.join(
directory, file_prefix + "reverse" + str(index) + ".h5")
#save the trajectory
reverse_trajectory.save_hdf5(full_filename)
def _write_equilibrium_trajectories(self, directory, file_prefix):
"""
Write out an MDTraj h5 file for each nonequilibrium trajectory. The files will be placed in
[directory]/file_prefix-[lambda0, lambda1].h5.
Parameters
----------
directory : str
The directory in which to place the files
file_prefix : str
A prefix for the filenames
"""
lambda_zero_filename = os.path.join(
directory, file_prefix + "-" + "lambda0" + ".h5")
lambda_one_filename = os.path.join(
directory, file_prefix + "-" + "lambda1" + ".h5")
filenames = [lambda_zero_filename, lambda_one_filename]
trajs = [self._lambda_zero_traj, self._lambda_one_traj]
#open the existing file if it exists, and append. Otherwise create it
for filename, traj in zip(filenames, trajs):
if not os.path.exists(filename):
traj.save_hdf5(filename)
else:
written_traj = md.load(filename)
concatenated_traj = written_traj.join(traj)
concatenated_traj.save_hdf5(filename)
#delete the trajectories.
self._lambda_one_traj = None
self._lambda_zero_traj = None
def retrieve_nonalchemical_results(self):
"""
Call this to retrieve the reduced potential results for the nonalchemical endpoints
"""
for nonalchemical_result_zero, nonalchemical_result_one in zip(
self._nonalchemical_zero_results,
self._nonalchemical_one_results):
self._nonalchemical_zero_endpt_reduced_potentials.append(
nonalchemical_result_zero.get())
self._nonalchemical_one_endpt_reduced_potentials.append(
nonalchemical_result_one.get())
self._nonalchemical_zero_results = []
self._nonalchemcal_one_results = []
@property
def zero_endpoint_perturbation(self):
hybrid_reduced_potentials = np.array(
self._lambda_zero_reduced_potentials)
nonalchemical_reduced_potentials = np.array(
self._nonalchemical_zero_endpt_reduced_potentials)
[df, ddf] = pymbar.EXP(nonalchemical_reduced_potentials -
hybrid_reduced_potentials)
return [df, ddf]
@property
def one_endpoint_perturbation(self):
hybrid_reduced_potentials = np.array(
self._lambda_one_reduced_potentials)
nonalchemical_reduced_potentials = np.array(
self._lambda_zero_reduced_potentials)
[df, ddf] = pymbar.EXP(nonalchemical_reduced_potentials -
hybrid_reduced_potentials)
return [df, ddf]
@property
def lambda_zero_equilibrium_trajectory(self):
return self._lambda_zero_traj
@property
def lambda_one_equilibrium_trajectory(self):
return self._lambda_one_traj
@property
def forward_nonequilibrium_trajectories(self):
return self._forward_nonequilibrium_trajectories
@property
def reverse_nonequilibrium_trajectories(self):
return self._reverse_nonequilibrium_trajectories
@property
def forward_cumulative_works(self):
return self._forward_nonequilibrium_cumulative_works
@property
def reverse_cumulative_works(self):
return self._reverse_nonequilibrium_cumulative_works
@property
def current_free_energy_estimate(self):
[df, ddf] = pymbar.BAR(self._forward_total_work,
self._reverse_total_work)
return [df, ddf]