def test_non_alchemical_protocol(self):
"""Any property of the ThermodynamicSystem can be specified in the protocol."""
name, thermodynamic_state, sampler_state, topography = self.host_guest_implicit
protocol = {
'lambda_sterics': [0.0, 0.5, 1.0],
'temperature': [300, 320, 300] * unit.kelvin,
'update_alchemical_charges': [True, True, False]
}
alchemical_phase = AlchemicalPhase(sampler=ReplicaExchange())
with self.temporary_storage_path() as storage_path:
alchemical_phase.create(thermodynamic_state, sampler_state, topography,
protocol, storage_path, restraint=yank.restraints.Harmonic())
self.check_protocol(alchemical_phase, protocol)
# If temperatures of the end states is different, an error is raised.
protocol['temperature'][-1] = 330 * unit.kelvin
alchemical_phase = AlchemicalPhase(sampler=ReplicaExchange())
with nose.tools.assert_raises(ValueError):
alchemical_phase.create(thermodynamic_state, sampler_state, topography,
protocol, 'not_created.nc', restraint=yank.restraints.Harmonic())
def test_illegal_protocol(self):
"""An error is raised when the protocol parameters have a different number of states."""
name, thermodynamic_state, sampler_state, topography = self.host_guest_implicit
protocol = {
'lambda_sterics': [0.0, 1.0],
'lambda_electrostatics': [1.0]
}
alchemical_phase = AlchemicalPhase(sampler=ReplicaExchange())
restraint = yank.restraints.Harmonic()
with nose.tools.assert_raises(ValueError):
alchemical_phase.create(thermodynamic_state, sampler_state, topography,
protocol, 'not_created.nc', restraint=restraint)
def test_illegal_restraint(self):
"""Raise an error when restraint is handled incorrectly."""
# An error is raised with ligand-receptor systems in implicit without restraint.
test_name, thermodynamic_state, sampler_state, topography = self.host_guest_implicit
alchemical_phase = AlchemicalPhase(sampler=ReplicaExchange())
with nose.tools.assert_raises(ValueError):
alchemical_phase.create(thermodynamic_state, sampler_state, topography,
self.protocol, 'not_created.nc')
# An error is raised when trying to apply restraint to non ligand-receptor systems.
restraint = yank.restraints.Harmonic()
test_name, thermodynamic_state, sampler_state, topography = self.alanine_explicit
with nose.tools.assert_raises(RuntimeError):
alchemical_phase.create(thermodynamic_state, sampler_state, topography,
self.protocol, 'not_created.nc', restraint=restraint)
def test_from_storage(self):
"""When resuming, the AlchemicalPhase recover the correct sampler."""
_, thermodynamic_state, sampler_state, topography = self.host_guest_implicit
restraint = yank.restraints.Harmonic()
with self.temporary_storage_path() as storage_path:
alchemical_phase = AlchemicalPhase(ReplicaExchange())
alchemical_phase.create(thermodynamic_state, sampler_state, topography,
self.protocol, storage_path, restraint=restraint)
# Delete old alchemical phase to close storage file.
del alchemical_phase
# Resume, the sampler has the correct class.
alchemical_phase = AlchemicalPhase.from_storage(storage_path)
assert isinstance(alchemical_phase._sampler, ReplicaExchange)
def test_create(self):
"""Alchemical state correctly creates the simulation object."""
available_restraints = list(yank.restraints.available_restraint_classes().values())
for test_index, test_case in enumerate(self.all_test_cases):
test_name, thermodynamic_state, sampler_state, topography = test_case
# Add random restraint if this is ligand-receptor system in implicit solvent.
if len(topography.ligand_atoms) > 0:
restraint_cls = available_restraints[np.random.randint(0, len(available_restraints))]
restraint = restraint_cls()
protocol = self.restrained_protocol
test_name += ' with restraint {}'.format(restraint_cls.__name__)
else:
restraint = None
protocol = self.protocol
# Add either automatic of fixed correction cutoff.
if test_index % 2 == 0:
correction_cutoff = 12 * unit.angstroms
else:
correction_cutoff = 'auto'
# Replace the reaction field of the reference system to compare
# also cutoff and switch width for the electrostatics.
reference_system = thermodynamic_state.system
mmtools.forcefactories.replace_reaction_field(reference_system, return_copy=False)
alchemical_phase = AlchemicalPhase(sampler=ReplicaExchange())
with self.temporary_storage_path() as storage_path:
alchemical_phase.create(thermodynamic_state, sampler_state, topography,
protocol, storage_path, restraint=restraint,
anisotropic_dispersion_cutoff=correction_cutoff)
yield prepare_yield(self.check_protocol, test_name, alchemical_phase, protocol)
yield prepare_yield(self.check_standard_state_correction, test_name, alchemical_phase,
topography)
yield prepare_yield(self.check_expanded_states, test_name, alchemical_phase,
protocol, correction_cutoff, reference_system)
# Free memory.
del alchemical_phase
def test_randomize_ligand(self):
"""Test method AlchemicalPhase.randomize_ligand."""
_, thermodynamic_state, sampler_state, topography = self.host_guest_implicit
restraint = yank.restraints.Harmonic()
ligand_atoms, receptor_atoms = topography.ligand_atoms, topography.receptor_atoms
ligand_positions = sampler_state.positions[ligand_atoms]
receptor_positions = sampler_state.positions[receptor_atoms]
with self.temporary_storage_path() as storage_path:
alchemical_phase = AlchemicalPhase(ReplicaExchange())
alchemical_phase.create(thermodynamic_state, sampler_state, topography,
self.protocol, storage_path, restraint=restraint)
# Randomize ligand positions.
alchemical_phase.randomize_ligand()
# The new sampler states have the same receptor positions
# but different ligand positions.
for sampler_state in alchemical_phase._sampler.sampler_states:
assert np.allclose(sampler_state.positions[receptor_atoms], receptor_positions)
assert not np.allclose(sampler_state.positions[ligand_atoms], ligand_positions)
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
sampler_state1 = states.SamplerState(test_system.positions)
sampler_state2 = states.SamplerState(positions2)
sampler_state3 = states.SamplerState(positions3)
sampler_states = [sampler_state1, sampler_state2, sampler_state3]
with self.temporary_storage_path() as storage_path:
# Create alchemical phase.
alchemical_phase = AlchemicalPhase(ReplicaExchange())
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)
def test_unknown_parameters():
"""Test ReplicaExchange raises exception on wrong initialization."""
ReplicaExchange(store_filename='test', wrong_parameter=False)
def test_parameters():
"""Test ReplicaExchange parameters initialization."""
repex = ReplicaExchange(store_filename='test', nsteps_per_iteration=1e6)
assert repex.nsteps_per_iteration == 1000000
assert repex.collision_rate == repex.default_parameters['collision_rate']
def test_replica_exchange(mpicomm=None, verbose=True):
"""
Test that free energies and average potential energies of a 3D harmonic oscillator are correctly computed by parallel tempering.
TODO
* Test ParallelTempering and HamiltonianExchange subclasses as well.
* Test with different combinations of input parameters.
"""
if verbose and ((not mpicomm) or (mpicomm.rank==0)): sys.stdout.write("Testing replica exchange facility with harmonic oscillators: ")
# Define mass of carbon atom.
mass = 12.0 * units.amu
# Define thermodynamic states.
states = list() # thermodynamic states
Ks = [500.00, 400.0, 300.0] * units.kilocalories_per_mole / units.angstroms**2 # spring constants
temperatures = [300.0, 350.0, 400.0] * units.kelvin # temperatures
seed_positions = list()
analytical_results = list()
f_i_analytical = list() # dimensionless free energies
u_i_analytical = list() # reduced potential
for (K, temperature) in zip(Ks, temperatures):
# Create harmonic oscillator system.
testsystem = testsystems.HarmonicOscillator(K=K, mass=mass, mm=openmm)
[system, positions] = [testsystem.system, testsystem.positions]
# Create thermodynamic state.
state = ThermodynamicState(system=system, temperature=temperature)
# Append thermodynamic state and positions.
states.append(state)
seed_positions.append(positions)
# Store analytical results.
results = computeHarmonicOscillatorExpectations(K, mass, temperature)
analytical_results.append(results)
f_i_analytical.append(results['f'])
kT = kB * temperature # thermal energy
reduced_potential = results['potential']['mean'] / kT
u_i_analytical.append(reduced_potential)
# Compute analytical Delta_f_ij
nstates = len(f_i_analytical)
f_i_analytical = numpy.array(f_i_analytical)
u_i_analytical = numpy.array(u_i_analytical)
s_i_analytical = u_i_analytical - f_i_analytical
Delta_f_ij_analytical = numpy.zeros([nstates,nstates], numpy.float64)
Delta_u_ij_analytical = numpy.zeros([nstates,nstates], numpy.float64)
Delta_s_ij_analytical = numpy.zeros([nstates,nstates], numpy.float64)
for i in range(nstates):
for j in range(nstates):
Delta_f_ij_analytical[i,j] = f_i_analytical[j] - f_i_analytical[i]
Delta_u_ij_analytical[i,j] = u_i_analytical[j] - u_i_analytical[i]
Delta_s_ij_analytical[i,j] = s_i_analytical[j] - s_i_analytical[i]
# Define file for temporary storage.
import tempfile # use a temporary file
file = tempfile.NamedTemporaryFile(delete=False)
store_filename = file.name
#print("node %d : Storing data in temporary file: %s" % (mpicomm.rank, str(store_filename))) # DEBUG
# Create and configure simulation object.
simulation = ReplicaExchange(store_filename, mpicomm=mpicomm)
simulation.create(states, seed_positions)
simulation.platform = openmm.Platform.getPlatformByName('Reference')
simulation.minimize = False
simulation.number_of_iterations = 200
simulation.nsteps_per_iteration = 500
simulation.timestep = 2.0 * units.femtoseconds
simulation.collision_rate = 20.0 / units.picosecond
simulation.verbose = False
simulation.show_mixing_statistics = False
simulation.online_analysis = True
# Run simulation.
utils.config_root_logger(False)
simulation.run() # run the simulation
utils.config_root_logger(True)
# Stop here if not root node.
if mpicomm and (mpicomm.rank != 0): return
# Retrieve extant analysis object.
online_analysis = simulation.analysis
# Analyze simulation to compute free energies.
analysis = simulation.analyze()
# Check if online analysis is close to final analysis.
error = numpy.abs(online_analysis['Delta_f_ij'] - analysis['Delta_f_ij'])
derror = (online_analysis['dDelta_f_ij']**2 + analysis['dDelta_f_ij']**2)
indices = numpy.where(derror > 0.0)
nsigma = numpy.zeros([nstates,nstates], numpy.float32)
nsigma[indices] = error[indices] / derror[indices]
MAX_SIGMA = 6.0 # maximum allowed number of standard errors
if numpy.any(nsigma > MAX_SIGMA):
print("Delta_f_ij from online analysis")
print(online_analysis['Delta_f_ij'])
print("Delta_f_ij from final analysis")
print(analysis['Delta_f_ij'])
print("error")
print(error)
print("derror")
print(derror)
print("nsigma")
print(nsigma)
raise Exception("Dimensionless free energy differences between online and final analysis exceeds MAX_SIGMA of %.1f" % MAX_SIGMA)
# TODO: Check if deviations exceed tolerance.
Delta_f_ij = analysis['Delta_f_ij']
dDelta_f_ij = analysis['dDelta_f_ij']
error = numpy.abs(Delta_f_ij - Delta_f_ij_analytical)
indices = numpy.where(dDelta_f_ij > 0.0)
nsigma = numpy.zeros([nstates,nstates], numpy.float32)
nsigma[indices] = error[indices] / dDelta_f_ij[indices]
MAX_SIGMA = 6.0 # maximum allowed number of standard errors
if numpy.any(nsigma > MAX_SIGMA):
print("Delta_f_ij")
print(Delta_f_ij)
print("Delta_f_ij_analytical")
print(Delta_f_ij_analytical)
print("error")
print(error)
print("stderr")
print(dDelta_f_ij)
print("nsigma")
print(nsigma)
raise Exception("Dimensionless free energy difference exceeds MAX_SIGMA of %.1f" % MAX_SIGMA)
error = analysis['Delta_u_ij'] - Delta_u_ij_analytical
nsigma = numpy.zeros([nstates,nstates], numpy.float32)
nsigma[indices] = error[indices] / dDelta_f_ij[indices]
if numpy.any(nsigma > MAX_SIGMA):
print("Delta_u_ij")
print(analysis['Delta_u_ij'])
print("Delta_u_ij_analytical")
print(Delta_u_ij_analytical)
print("error")
print(error)
print("nsigma")
print(nsigma)
raise Exception("Dimensionless potential energy difference exceeds MAX_SIGMA of %.1f" % MAX_SIGMA)
# Clean up.
del simulation
if verbose: print("PASSED.")
return
def setup_class(cls):
"""Shared test cases and variables."""
n_states = 3
n_steps = 5
checkpoint_interval = 2
# Test case with host guest in vacuum at 3 different positions and alchemical parameters.
# -----------------------------------------------------------------------------------------
hostguest_test = testsystems.HostGuestVacuum()
factory = mmtools.alchemy.AbsoluteAlchemicalFactory()
alchemical_region = mmtools.alchemy.AlchemicalRegion(
alchemical_atoms=range(126, 156))
hostguest_alchemical = factory.create_alchemical_system(
hostguest_test.system, alchemical_region)
# Translate the sampler states to be different one from each other.
hostguest_sampler_states = [
mmtools.states.SamplerState(hostguest_test.positions +
10 * i * unit.nanometers)
for i in range(n_states)
]
# Create the three basic thermodynamic states.
hostguest_thermodynamic_states = [
mmtools.states.ThermodynamicState(hostguest_alchemical,
300 * unit.kelvin)
for i in range(n_states)
]
# Create alchemical states at different parameter values.
alchemical_states = [
mmtools.alchemy.AlchemicalState.from_system(hostguest_alchemical)
for _ in range(n_states)
]
for i, alchemical_state in enumerate(alchemical_states):
alchemical_state.set_alchemical_parameters(
float(i) / (n_states - 1))
# Create compound states.
hostguest_compound_states = list()
for i in range(n_states):
hostguest_compound_states.append(
mmtools.states.CompoundThermodynamicState(
thermodynamic_state=hostguest_thermodynamic_states[i],
composable_states=[alchemical_states[i]]))
# Unsampled states.
nonalchemical_state = mmtools.states.ThermodynamicState(
hostguest_test.system, 300 * unit.kelvin)
hostguest_unsampled_states = [
copy.deepcopy(nonalchemical_state),
copy.deepcopy(nonalchemical_state)
]
cls.hostguest_test = (hostguest_compound_states,
hostguest_sampler_states,
hostguest_unsampled_states)
# Run a quick simulation
thermodynamic_states, sampler_states, unsampled_states = copy.deepcopy(
cls.hostguest_test)
n_states = len(thermodynamic_states)
# Remove one sampler state to verify distribution over states.
sampler_states = sampler_states[:-1]
# Prepare metadata for analysis.
reference_state = mmtools.states.ThermodynamicState(
hostguest_test.system, 300 * unit.kelvin)
topography = Topography(hostguest_test.topology,
ligand_atoms=range(126, 156))
metadata = {
'standard_state_correction': 4.0,
'reference_state': mmtools.utils.serialize(reference_state),
'topography': mmtools.utils.serialize(topography)
}
analysis_atoms = topography.ligand_atoms
# Create simulation and storage file.
cls.tmp_dir = tempfile.mkdtemp()
storage_path = os.path.join(cls.tmp_dir, 'test_analyze.nc')
move = mmtools.mcmc.LangevinDynamicsMove(n_steps=1)
cls.repex = ReplicaExchange(mcmc_moves=move,
number_of_iterations=n_steps)
cls.reporter = Reporter(storage_path,
checkpoint_interval=checkpoint_interval,
analysis_particle_indices=analysis_atoms)
cls.repex.create(thermodynamic_states,
sampler_states,
storage=cls.reporter,
unsampled_thermodynamic_states=unsampled_states,
metadata=metadata)
# run some iterations
cls.n_states = n_states
cls.n_steps = n_steps
cls.checkpoint_interval = checkpoint_interval
cls.analysis_atoms = analysis_atoms
cls.repex.run(cls.n_steps - 1) # Initial config
cls.repex_name = "RepexAnalyzer"
