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
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]