def test_parallel_tempering_pressure():
nc_filename = tempfile.mkdtemp() + "/out.nc"
T_min = 1.0 * unit.kelvin
T_max = 10.0 * unit.kelvin
n_temps = 3
pressure = 1.0 * unit.atmospheres
ho = testsystems.HarmonicOscillator()
system = ho.system
positions = ho.positions
coordinates = [positions] * n_temps
mpicomm = dummympi.DummyMPIComm()
parameters = {"number_of_iterations": 4}
replica_exchange = ParallelTempering.create(system,
coordinates,
nc_filename,
T_min=T_min,
T_max=T_max,
n_temps=n_temps,
pressure=pressure,
mpicomm=mpicomm,
parameters=parameters)
eq(replica_exchange.n_replicas, n_temps)
replica_exchange.run()
def test_double_set_thermodynamic_states():
nc_filename = tempfile.mkdtemp() + "/out.nc"
T_min = 1.0 * unit.kelvin
T_max = 10.0 * unit.kelvin
n_temps = 3
ho = testsystems.HarmonicOscillator()
system = ho.system
positions = ho.positions
coordinates = [positions] * n_temps
mpicomm = dummympi.DummyMPIComm()
parameters = {"number_of_iterations": 3}
replica_exchange = ParallelTempering.create(system,
coordinates,
nc_filename,
T_min=T_min,
T_max=T_max,
n_temps=n_temps,
mpicomm=mpicomm,
parameters=parameters)
eq(replica_exchange.n_replicas, n_temps)
replica_exchange.run()
states = replica_exchange.thermodynamic_states
replica_exchange.database.thermodynamic_states = states
def test_parallel_tempering_explicit_temperature_input(mpicomm):
nc_filename = tempfile.mkdtemp() + "/out.nc"
T0 = 1.0 * unit.kelvin
temperatures = [T0, T0 * 2, T0 * 4]
n_temps = len(temperatures)
ho = testsystems.HarmonicOscillator()
system = ho.system
positions = ho.positions
coordinates = [positions] * n_temps
parameters = {"number_of_iterations": 100}
replica_exchange = ParallelTempering.create(system,
coordinates,
nc_filename,
temperatures=temperatures,
mpicomm=mpicomm,
parameters=parameters)
replica_exchange.run()
eq(replica_exchange.n_replicas, n_temps)
def test_check_positions():
nc_filename = tempfile.mkdtemp() + "/out.nc"
T_min = 1.0 * unit.kelvin
T_i = [T_min, T_min * 10., T_min * 100.]
n_replicas = len(T_i)
ho = testsystems.HarmonicOscillator()
system = ho.system
positions = ho.positions
states = [
ThermodynamicState(system=system, temperature=T_i[i])
for i in range(n_replicas)
]
coordinates = [positions] * n_replicas
mpicomm = dummympi.DummyMPIComm()
parameters = {"number_of_iterations": 10}
replica_exchange = ReplicaExchange.create(states,
coordinates,
nc_filename,
mpicomm=mpicomm,
parameters=parameters)
replica_exchange.run()
db = replica_exchange.database
db.check_energies()
def test_external_protocol_work_accumulation():
"""When `measure_protocol_work==True`, assert that global `protocol_work` is initialized to zero and
reaches a zero value after integrating a few dozen steps without perturbation.
By default (`measure_protocol_work=False`), assert that there is no global name for `protocol_work`."""
testsystem = testsystems.HarmonicOscillator()
system, topology = testsystem.system, testsystem.topology
temperature = 298.0 * unit.kelvin
integrator = integrators.ExternalPerturbationLangevinIntegrator(splitting="O V R V O", temperature=temperature)
context = openmm.Context(system, integrator)
context.setPositions(testsystem.positions)
context.setVelocitiesToTemperature(temperature)
# Check that initial step accumulates no protocol work
assert(integrator.get_protocol_work(dimensionless=True) == 0), "Protocol work should be 0 initially"
assert(integrator.get_protocol_work() / unit.kilojoules_per_mole == 0), "Protocol work should have units of energy"
integrator.step(1)
assert(integrator.get_protocol_work(dimensionless=True) == 0), "There should be no protocol work."
# Check that a single step accumulates protocol work
pe_1 = context.getState(getEnergy=True).getPotentialEnergy()
perturbed_K=99.0 * unit.kilocalories_per_mole / unit.angstroms**2
context.setParameter('testsystems_HarmonicOscillator_K', perturbed_K)
pe_2 = context.getState(getEnergy=True).getPotentialEnergy()
integrator.step(1)
assert (integrator.get_protocol_work(dimensionless=True) != 0), "There should be protocol work after perturbing."
assert (integrator.protocol_work == (pe_2 - pe_1)), "The potential energy difference should be equal to protocol work."
del context, integrator
integrator = integrators.VVVRIntegrator(temperature)
context = openmm.Context(system, integrator)
context.setPositions(testsystem.positions)
context.setVelocitiesToTemperature(temperature)
integrator.step(25)
del context, integrator
def test_parallel_tempering_save_and_load(mpicomm):
nc_filename = tempfile.mkdtemp() + "/out.nc"
T_min = 1.0 * unit.kelvin
T_max = 10.0 * unit.kelvin
n_temps = 3
ho = testsystems.HarmonicOscillator()
system = ho.system
positions = ho.positions
coordinates = [positions] * n_temps
parameters = {"number_of_iterations": 200}
replica_exchange = ParallelTempering.create(system,
coordinates,
nc_filename,
T_min=T_min,
T_max=T_max,
n_temps=n_temps,
mpicomm=mpicomm,
parameters=parameters)
replica_exchange.run()
replica_exchange.extend(100)
replica_exchange = resume(nc_filename, mpicomm=mpicomm)
eq(replica_exchange.iteration, 200)
replica_exchange.run()
def test_baoab_heat_accumulation():
"""When `measure_heat==True`, assert that global `heat` is initialized to zero and
reaches a nonzero value after integrating a few dozen steps."""
# test `measure_shadow_work=True` --> accumulation of a nonzero value in global `shadow_work`
testsystem = testsystems.HarmonicOscillator()
system, topology = testsystem.system, testsystem.topology
temperature = 298.0 * unit.kelvin
integrator = integrators.BAOABIntegrator(temperature, measure_heat=True)
context = openmm.Context(system, integrator)
context.setPositions(testsystem.positions)
context.setVelocitiesToTemperature(temperature)
assert(integrator.get_heat(dimensionless=True) == 0), "Heat should initially be zero."
assert(integrator.get_heat() / unit.kilojoules_per_mole == 0), "integrator.get_heat() should have units of energy."
assert(integrator.heat / unit.kilojoules_per_mole == 0), "integrator.heat should have units of energy."
integrator.step(25)
assert(integrator.get_heat(dimensionless=True) != 0), "integrator.get_heat() should be nonzero after dynamics"
integrator = integrators.VVVRIntegrator(temperature)
context = openmm.Context(system, integrator)
context.setPositions(testsystem.positions)
context.setVelocitiesToTemperature(temperature)
integrator.step(25)
del context, integrator
def test_repex_save_and_load():
nc_filename = tempfile.mkdtemp() + "/out.nc"
T_min = 1.0 * unit.kelvin
T_i = [T_min, T_min * 10., T_min * 100.]
n_replicas = len(T_i)
ho = testsystems.HarmonicOscillator()
system = ho.system
positions = ho.positions
states = [
ThermodynamicState(system=system, temperature=T_i[i])
for i in range(n_replicas)
]
coordinates = [positions] * n_replicas
parameters = {"number_of_iterations": 5}
rex = replica_exchange.ReplicaExchange.create(states,
coordinates,
nc_filename,
parameters=parameters)
rex.run()
rex.extend(5)
rex = resume(nc_filename)
rex.run()
eq(rex.__class__.__name__, "ReplicaExchange")
def test_parallel_tempering_save_and_load():
nc_filename = tempfile.mkdtemp() + "/out.nc"
T_min = 1.0 * unit.kelvin
T_max = 10.0 * unit.kelvin
n_temps = 3
ho = testsystems.HarmonicOscillator()
system = ho.system
positions = ho.positions
coordinates = [positions] * n_temps
parameters = {"number_of_iterations": 5}
rex = parallel_tempering.ParallelTempering.create(system,
coordinates,
nc_filename,
T_min=T_min,
T_max=T_max,
n_temps=n_temps,
parameters=parameters)
rex.run()
rex.extend(5)
rex = resume(nc_filename)
rex.run()
eq(rex.__class__.__name__, "ParallelTempering")
def test_harmonic_oscillators_save_and_load():
nc_filename = tempfile.mkdtemp() + "/out.nc"
T_min = 1.0 * unit.kelvin
T_i = [T_min, T_min * 10., T_min * 100.]
n_replicas = len(T_i)
ho = testsystems.HarmonicOscillator()
system = ho.system
positions = ho.positions
states = [
ThermodynamicState(system=system, temperature=T_i[i])
for i in range(n_replicas)
]
coordinates = [positions] * n_replicas
mpicomm = dummympi.DummyMPIComm()
parameters = {"number_of_iterations": 200}
replica_exchange = ReplicaExchange.create(states,
coordinates,
nc_filename,
mpicomm=mpicomm,
parameters=parameters)
replica_exchange.run()
replica_exchange.extend(100)
replica_exchange = resume(nc_filename, mpicomm=mpicomm)
eq(replica_exchange.iteration, 200)
replica_exchange.run()
def test_protocol_work_accumulation_harmonic_oscillator(self):
"""Testing protocol work accumulation for ExternalPerturbationLangevinIntegrator with HarmonicOscillator
"""
testsystem = testsystems.HarmonicOscillator()
parameter_name = 'testsystems_HarmonicOscillator_x0'
parameter_initial = 0.0 * unit.angstroms
parameter_final = 10.0 * unit.angstroms
for platform_name in ['Reference', 'CPU']:
self.compare_external_protocol_work_accumulation(testsystem, parameter_name, parameter_initial, parameter_final, platform_name=platform_name)
def test_temperature_getter_setter():
"""Test that temperature setter and getter modify integrator variables."""
temperature = 350 * unit.kelvin
test = testsystems.HarmonicOscillator()
custom_integrators = get_all_custom_integrators()
thermostated_integrators = dict(
get_all_custom_integrators(only_thermostated=True))
for integrator_name, integrator_class in custom_integrators:
# If this is not a ThermostatedIntegrator, the interface should not be added.
if integrator_name not in thermostated_integrators:
integrator = integrator_class()
# NoseHooverChainVelocityVerletIntegrator will print a severe warning here,
# because it is being initialized without a system. That's OK.
assert ThermostatedIntegrator.is_thermostated(integrator) is False
assert ThermostatedIntegrator.restore_interface(
integrator) is False
assert not hasattr(integrator, 'getTemperature')
continue
# Test original integrator.
check_integrator_temperature_getter_setter.description = (
'Test temperature setter and '
'getter of {}').format(integrator_name)
integrator = integrator_class(temperature=temperature)
# NoseHooverChainVelocityVerletIntegrator will print a severe warning here,
# because it is being initialized without a system. That's OK.
context = openmm.Context(test.system, integrator)
context.setPositions(test.positions)
# Integrator temperature is initialized correctly.
check_integrator_temperature(integrator, temperature, 1)
yield check_integrator_temperature_getter_setter, integrator
del context
# Test Context integrator wrapper.
check_integrator_temperature_getter_setter.description = (
'Test temperature wrapper '
'of {}').format(integrator_name)
integrator = integrator_class()
# NoseHooverChainVelocityVerletIntegrator will print a severe warning here,
# because it is being initialized without a system. That's OK.
context = openmm.Context(test.system, integrator)
context.setPositions(test.positions)
integrator = context.getIntegrator()
# Setter and getter should be added successfully.
assert ThermostatedIntegrator.is_thermostated(integrator) is True
assert ThermostatedIntegrator.restore_interface(integrator) is True
assert isinstance(integrator, integrator_class)
yield check_integrator_temperature_getter_setter, integrator
del context
def test_stabilities_harmonic_oscillator():
"""
Test integrators for stability over a short number of steps of a harmonic oscillator.
"""
test = testsystems.HarmonicOscillator()
for methodname in dir(integrators):
if re.match('.*Integrator$', methodname):
integrator = getattr(integrators, methodname)()
integrator.__doc__ = methodname
check_stability.description = "Testing %s for stability over a short number of integration steps of a harmonic oscillator." % methodname
yield check_stability, integrator, test
def test_no_work_accumulation_harmonic_oscillator(self):
"""Testing protocol work accumulation for ExternalPerturbationLangevinIntegrator with HarmonicOscillator
"""
testsystem = testsystems.HarmonicOscillator()
parameter_name = "testsystems_HarmonicOscillator_x0"
parameter_initial = 0.0 * unit.angstroms
parameter_final = 10.0 * unit.angstroms
for platform_name in ["Reference", "CPU"]:
self.run_without_work_accumulation(
testsystem,
parameter_name,
parameter_initial,
parameter_final,
platform_name=platform_name,
)
def test_repex_multiple_save_and_load():
nc_filename = tempfile.mkdtemp() + "/out.nc"
T_min = 1.0 * unit.kelvin
T_i = [T_min, T_min * 2.0, T_min * 2.0]
n_replicas = len(T_i)
ho = testsystems.HarmonicOscillator()
system = ho.system
positions = ho.positions
states = [
ThermodynamicState(system=system, temperature=T_i[i])
for i in range(n_replicas)
]
coordinates = [positions] * n_replicas
parameters = {"number_of_iterations": steps}
rex = replica_exchange.ReplicaExchange.create(states,
coordinates,
nc_filename,
parameters=parameters)
rex.run()
for repeat in range(repeats):
replica_states0 = rex.replica_states
Nij_proposed0 = rex.Nij_proposed
Nij_accepted0 = rex.Nij_accepted
sampler_states0 = rex.sampler_states
rex.extend(steps)
rex = resume(nc_filename)
replica_states1 = rex.replica_states
Nij_proposed1 = rex.Nij_proposed
Nij_accepted1 = rex.Nij_accepted
sampler_states1 = rex.sampler_states
eq(replica_states0, replica_states1)
eq(Nij_proposed0, Nij_proposed1)
eq(Nij_accepted0, Nij_accepted1)
rex.run()
def test_parallel_tempering_multiple_save_and_load():
nc_filename = tempfile.mkdtemp() + "/out.nc"
T_min = 1.0 * unit.kelvin
T_max = 2.0 * unit.kelvin
n_temps = 3
ho = testsystems.HarmonicOscillator()
system = ho.system
positions = ho.positions
coordinates = [positions] * n_temps
parameters = {"number_of_iterations": steps}
rex = parallel_tempering.ParallelTempering.create(system,
coordinates,
nc_filename,
T_min=T_min,
T_max=T_max,
n_temps=n_temps,
parameters=parameters)
rex.run()
for repeat in range(repeats):
replica_states0 = rex.replica_states
Nij_proposed0 = rex.Nij_proposed
Nij_accepted0 = rex.Nij_accepted
sampler_states0 = rex.sampler_states
rex.extend(steps)
rex = resume(nc_filename)
replica_states1 = rex.replica_states
Nij_proposed1 = rex.Nij_proposed
Nij_accepted1 = rex.Nij_accepted
sampler_states1 = rex.sampler_states
eq(replica_states0, replica_states1)
eq(Nij_proposed0, Nij_proposed1)
eq(Nij_accepted0, Nij_accepted1)
rex.run()
def test_parallel_tempering():
nc_filename = tempfile.mkdtemp() + "/out.nc"
T_min = 1.0 * unit.kelvin
T_max = 10.0 * unit.kelvin
n_temps = 3
ho = testsystems.HarmonicOscillator()
system = ho.system
positions = ho.positions
coordinates = [positions] * n_temps
mpicomm = dummympi.DummyMPIComm()
parameters = {"number_of_iterations": 1000}
replica_exchange = ParallelTempering.create(system,
coordinates,
nc_filename,
T_min=T_min,
T_max=T_max,
n_temps=n_temps,
mpicomm=mpicomm,
parameters=parameters)
eq(replica_exchange.n_replicas, n_temps)
replica_exchange.run()
u_permuted = replica_exchange.database.ncfile.variables["energies"][:]
s = replica_exchange.database.ncfile.variables["states"][:]
u = permute_energies(u_permuted, s)
states = replica_exchange.thermodynamic_states
u0 = np.array([[ho.reduced_potential_expectation(s0, s1) for s1 in states]
for s0 in states])
l0 = np.log(
u0
) # Compare on log scale because uncertainties are proportional to values
l1 = np.log(u.mean(0))
eq(l0, l1, decimal=1)
def test_harmonic_oscillators():
nc_filename = tempfile.mkdtemp() + "/out.nc"
T_min = 1.0 * unit.kelvin
T_i = [T_min, T_min * 10., T_min * 100.]
n_replicas = len(T_i)
ho = testsystems.HarmonicOscillator()
system = ho.system
positions = ho.positions
states = [
ThermodynamicState(system=system, temperature=T_i[i])
for i in range(n_replicas)
]
coordinates = [positions] * n_replicas
mpicomm = dummympi.DummyMPIComm()
parameters = {"number_of_iterations": 1000}
replica_exchange = ReplicaExchange.create(states,
coordinates,
nc_filename,
mpicomm=mpicomm,
parameters=parameters)
replica_exchange.run()
u_permuted = replica_exchange.database.ncfile.variables["energies"][:]
s = replica_exchange.database.ncfile.variables["states"][:]
u = permute_energies(u_permuted, s)
u0 = np.array([[ho.reduced_potential_expectation(s0, s1) for s1 in states]
for s0 in states])
l0 = np.log(
u0
) # Compare on log scale because uncertainties are proportional to values
l1 = np.log(u.mean(0))
eq(l0, l1, decimal=1)
def run_alchemical_langevin_integrator(nsteps=0,
splitting="O { V R H R V } O"):
"""Check that the AlchemicalLangevinSplittingIntegrator reproduces the analytical free energy difference for a harmonic oscillator deformation, using BAR.
Up to 6*sigma is tolerated for error.
The total work (protocol work + shadow work) is used.
"""
#max deviation from the calculated free energy
NSIGMA_MAX = 6
n_iterations = 200 # number of forward and reverse protocols
# These are the alchemical functions that will be used to control the system
temperature = 298.0 * unit.kelvin
sigma = 1.0 * unit.angstrom # stddev of harmonic oscillator
kT = kB * temperature # thermal energy
beta = 1.0 / kT # inverse thermal energy
K = kT / sigma**2 # spring constant corresponding to sigma
mass = 39.948 * unit.amu
period = unit.sqrt(mass / K) # period of harmonic oscillator
timestep = period / 20.0
collision_rate = 1.0 / period
dF_analytical = 1.0
parameters = dict()
parameters['testsystems_HarmonicOscillator_x0'] = (0 * sigma, 2 * sigma)
parameters['testsystems_HarmonicOscillator_U0'] = (0 * kT, 1 * kT)
alchemical_functions = {
'forward': {
name: '(1-lambda)*%f + lambda*%f' %
(value[0].value_in_unit_system(unit.md_unit_system),
value[1].value_in_unit_system(unit.md_unit_system))
for (name, value) in parameters.items()
},
'reverse': {
name: '(1-lambda)*%f + lambda*%f' %
(value[1].value_in_unit_system(unit.md_unit_system),
value[0].value_in_unit_system(unit.md_unit_system))
for (name, value) in parameters.items()
},
}
# Create harmonic oscillator testsystem
testsystem = testsystems.HarmonicOscillator(K=K, mass=mass)
system = testsystem.system
positions = testsystem.positions
# Get equilibrium samples from initial and final states
burn_in = 5 * 20 # 5 periods
thinning = 5 * 20 # 5 periods
# Collect forward and reverse work values
directions = ['forward', 'reverse']
work = {
direction: np.zeros([n_iterations], np.float64)
for direction in directions
}
platform = openmm.Platform.getPlatformByName("Reference")
for direction in directions:
positions = testsystem.positions
# Create equilibrium and nonequilibrium integrators
equilibrium_integrator = GHMCIntegrator(temperature=temperature,
collision_rate=collision_rate,
timestep=timestep)
nonequilibrium_integrator = AlchemicalNonequilibriumLangevinIntegrator(
temperature=temperature,
collision_rate=collision_rate,
timestep=timestep,
alchemical_functions=alchemical_functions[direction],
splitting=splitting,
nsteps_neq=nsteps,
measure_shadow_work=True)
# Create compound integrator
compound_integrator = openmm.CompoundIntegrator()
compound_integrator.addIntegrator(equilibrium_integrator)
compound_integrator.addIntegrator(nonequilibrium_integrator)
# Create Context
context = openmm.Context(system, compound_integrator, platform)
context.setPositions(positions)
# Collect work samples
for iteration in range(n_iterations):
#
# Generate equilibrium sample
#
compound_integrator.setCurrentIntegrator(0)
equilibrium_integrator.reset()
compound_integrator.step(thinning)
#
# Generate nonequilibrium work sample
#
compound_integrator.setCurrentIntegrator(1)
nonequilibrium_integrator.reset()
# Check initial conditions after reset
current_lambda = nonequilibrium_integrator.getGlobalVariableByName(
'lambda')
assert current_lambda == 0.0, 'initial lambda should be 0.0 (was %f)' % current_lambda
current_step = nonequilibrium_integrator.getGlobalVariableByName(
'step')
assert current_step == 0.0, 'initial step should be 0 (was %f)' % current_step
compound_integrator.step(max(
1, nsteps)) # need to execute at least one step
work[direction][
iteration] = nonequilibrium_integrator.get_total_work(
dimensionless=True)
# Check final conditions before reset
current_lambda = nonequilibrium_integrator.getGlobalVariableByName(
'lambda')
assert current_lambda == 1.0, 'final lambda should be 1.0 (was %f) for splitting %s' % (
current_lambda, splitting)
current_step = nonequilibrium_integrator.getGlobalVariableByName(
'step')
assert int(current_step) == max(
1, nsteps
), 'final step should be %d (was %f) for splitting %s' % (max(
1, nsteps), current_step, splitting)
nonequilibrium_integrator.reset()
# Clean up
del context
del compound_integrator
dF, ddF = pymbar.BAR(work['forward'], work['reverse'])
nsigma = np.abs(dF - dF_analytical) / ddF
print(
"analytical DeltaF: {:12.4f}, DeltaF: {:12.4f}, dDeltaF: {:12.4f}, nsigma: {:12.1f}"
.format(dF_analytical, dF, ddF, nsigma))
if nsigma > NSIGMA_MAX:
raise Exception(
"The free energy difference for the nonequilibrium switching for splitting '%s' and %d steps is not zero within statistical error."
% (splitting, nsteps))
def test_periodic_langevin_integrator(splitting="H V R O R V H",
ncycles=40,
nsteps_neq=1000,
nsteps_eq=1000,
write_trajectory=False):
"""
Test PeriodicNonequilibriumIntegrator
Parameters
----------
integrator_flavor : openmmtools.integrator.PeriodicNonequilibriumIntegrator (or subclass)
integrator to run
ncycles : int, optional, default=40
number of cycles
nsteps_neq : int, optional, default=1000
number of forward/backward annealing steps
nsteps_eq : int, optional, default=1000
number of equilibration steps to run at endstates before annealing
write_trajectory : bool, optional, default=True
If True, will generate a PDB file that contains the harmonic oscillator trajectory
"""
#max deviation from the calculated free energy
NSIGMA_MAX = 6
# These are the alchemical functions that will be used to control the system
temperature = 298.0 * unit.kelvin
sigma = 1.0 * unit.angstrom # stddev of harmonic oscillator
kT = kB * temperature # thermal energy
beta = 1.0 / kT # inverse thermal energy
K = kT / sigma**2 # spring constant corresponding to sigma
mass = 39.948 * unit.amu
period = unit.sqrt(mass / K) # period of harmonic oscillator
timestep = period / 20.0
collision_rate = 1.0 / period
dF_analytical = 5.0
parameters = dict()
displacement = 10 * sigma
parameters['testsystems_HarmonicOscillator_x0'] = (0 * sigma, displacement)
parameters['testsystems_HarmonicOscillator_U0'] = (0 * kT, 5 * kT)
integrator_kwargs = {
'temperature': temperature,
'collision_rate': collision_rate,
'timestep': timestep,
'measure_shadow_work': False,
'measure_heat': False
}
alchemical_functions = {
name: '(1-lambda)*%f + lambda*%f' %
(value[0].value_in_unit_system(unit.md_unit_system),
value[1].value_in_unit_system(unit.md_unit_system))
for (name, value) in parameters.items()
}
# Create harmonic oscillator testsystem
testsystem = testsystems.HarmonicOscillator(K=K, mass=mass)
system = testsystem.system
positions = testsystem.positions
topology = testsystem.topology
# Create integrator
from openmmtools.integrators import PeriodicNonequilibriumIntegrator
integrator = PeriodicNonequilibriumIntegrator(
alchemical_functions=alchemical_functions,
splitting=splitting,
nsteps_eq=nsteps_eq,
nsteps_neq=nsteps_neq,
**integrator_kwargs)
platform = openmm.Platform.getPlatformByName("Reference")
context = openmm.Context(system, integrator, platform)
context.setPositions(positions)
nsteps_per_cycle = nsteps_eq + nsteps_neq + nsteps_eq + nsteps_neq
assert integrator.getGlobalVariableByName(
"n_steps_per_cycle") == nsteps_per_cycle
if write_trajectory:
from simtk.openmm.app import PDBFile
filename = 'neq-trajectory.pdb'
print(f'Writing trajectory to {filename}')
with open(filename, 'wt') as outfile:
# Write reference
import copy
pos1 = copy.deepcopy(positions)
pos2 = copy.deepcopy(positions)
pos2[0, 0] += displacement
PDBFile.writeModel(topology, pos1, outfile)
PDBFile.writeModel(topology, pos2, outfile)
interval = 10
PDBFile.writeModel(topology, positions, outfile, modelIndex=0)
for step in range(0, 2 * nsteps_per_cycle, interval):
integrator.step(interval)
positions = context.getState(getPositions=True).getPositions(
asNumpy=True)
PDBFile.writeModel(topology,
positions,
outfile,
modelIndex=step)
PDBFile.writeModel(topology, pos1, outfile)
PDBFile.writeModel(topology, pos2, outfile)
# Reset the integrator
integrator.reset()
step = 0
for cycle in range(2):
# eq (0)
for i in range(nsteps_eq):
integrator.step(1)
step += 1
assert integrator.getGlobalVariableByName("step") == (
step % nsteps_per_cycle)
assert np.isclose(integrator.getGlobalVariableByName("lambda"),
0.0)
# neq (0 -> 1)
for i in range(nsteps_neq):
integrator.step(1)
step += 1
assert integrator.getGlobalVariableByName("step") == (
step % nsteps_per_cycle)
assert np.isclose(
integrator.getGlobalVariableByName("lambda"),
(i + 1) / nsteps_neq
), f'{step} {integrator.getGlobalVariableByName("lambda")}'
# eq (1)
for i in range(nsteps_eq):
integrator.step(1)
step += 1
assert integrator.getGlobalVariableByName("step") == (
step % nsteps_per_cycle)
assert np.isclose(integrator.getGlobalVariableByName("lambda"),
1.0)
# neq (1 -> 0)
for i in range(nsteps_neq):
integrator.step(1)
step += 1
assert integrator.getGlobalVariableByName("step") == (
step % nsteps_per_cycle)
assert np.isclose(integrator.getGlobalVariableByName("lambda"),
1 - (i + 1) / nsteps_neq)
assert np.isclose(integrator.getGlobalVariableByName("lambda"), 0.0)
# Reset the integrator
integrator.reset()
forward_works, reverse_works = list(), list()
for _ in range(ncycles):
# Equilibrium (lambda = 0)
integrator.step(nsteps_eq)
# Forward (0 -> 1)
initial_work = integrator.get_protocol_work(dimensionless=True)
integrator.step(nsteps_neq)
final_work = integrator.get_protocol_work(dimensionless=True)
forward_work = final_work - initial_work
forward_works.append(forward_work)
# Equilibrium (lambda = 1)
integrator.step(nsteps_eq)
# Reverse work (1 -> 0)
initial_work = integrator.get_protocol_work(dimensionless=True)
integrator.step(nsteps_neq)
final_work = integrator.get_protocol_work(dimensionless=True)
reverse_work = final_work - initial_work
reverse_works.append(reverse_work)
print(np.array(forward_works).std())
print(np.array(reverse_works).std())
dF, ddF = pymbar.BAR(np.array(forward_works), np.array(reverse_works))
nsigma = np.abs(dF - dF_analytical) / ddF
assert np.isclose(integrator.getGlobalVariableByName("lambda"), 0.0)
print(
"analytical DeltaF: {:12.4f}, DeltaF: {:12.4f}, dDeltaF: {:12.4f}, nsigma: {:12.1f}"
.format(dF_analytical, dF, ddF, nsigma))
if nsigma > NSIGMA_MAX:
raise Exception(
f"The free energy difference for the nonequilibrium switching for splitting {splitting} is not zero within statistical error."
)
# Clean up
del context
del integrator
import simtk.openmm as openmm
import simtk.unit as units
import openmmtools.testsystems
from openmmtools import testsystems
from pymbar import timeseries
from functools import partial
from openmmmcmc.mcmc import HMCMove, GHMCMove, LangevinDynamicsMove, MonteCarloBarostatMove
import logging
# Test various combinations of systems and MCMC schemes
analytical_testsystems = [
("HarmonicOscillator", testsystems.HarmonicOscillator(), [ GHMCMove(timestep=10.0*units.femtoseconds,nsteps=100) ]),
("HarmonicOscillator", testsystems.HarmonicOscillator(), { GHMCMove(timestep=10.0*units.femtoseconds,nsteps=100) : 0.5, HMCMove(timestep=10*units.femtosecond, nsteps=10) : 0.5 }),
("HarmonicOscillatorArray", testsystems.HarmonicOscillatorArray(N=4), [ LangevinDynamicsMove(timestep=10.0*units.femtoseconds,nsteps=100) ]),
("IdealGas", testsystems.IdealGas(nparticles=216), [ HMCMove(timestep=10*units.femtosecond, nsteps=10), MonteCarloBarostatMove() ])
]
NSIGMA_CUTOFF = 6.0 # cutoff for significance testing
debug = False # set to True only for manual debugging of this nose test
def test_minimizer_all_testsystems():
#testsystem_classes = testsystems.TestSystem.__subclasses__()
testsystem_classes = [ testsystems.AlanineDipeptideVacuum ]
for testsystem_class in testsystem_classes:
class_name = testsystem_class.__name__
def run_alchemical_langevin_integrator(nsteps=0, splitting="O { V R H R V } O"):
"""Check that the AlchemicalLangevinSplittingIntegrator reproduces the analytical free energy difference for a harmonic oscillator deformation, using BAR.
Up to 6*sigma is tolerated for error.
The total work (protocol work + shadow work) is used.
"""
#max deviation from the calculated free energy
NSIGMA_MAX = 6
n_iterations = 100 # number of forward and reverse protocols
# These are the alchemical functions that will be used to control the system
temperature = 298.0 * unit.kelvin
sigma = 1.0 * unit.angstrom # stddev of harmonic oscillator
kT = kB * temperature # thermal energy
beta = 1.0 / kT # inverse thermal energy
K = kT / sigma**2 # spring constant corresponding to sigma
mass = 39.948 * unit.amu
period = unit.sqrt(mass/K) # period of harmonic oscillator
timestep = period / 20.0
collision_rate = 1.0 / period
dF_analytical = 1.0
parameters = dict()
parameters['testsystems_HarmonicOscillator_x0'] = (0 * sigma, 2 * sigma)
parameters['testsystems_HarmonicOscillator_U0'] = (0 * kT, 1 * kT)
forward_functions = { name : '(1-lambda)*%f + lambda*%f' % (value[0].value_in_unit_system(unit.md_unit_system), value[1].value_in_unit_system(unit.md_unit_system)) for (name, value) in parameters.items() }
reverse_functions = { name : '(1-lambda)*%f + lambda*%f' % (value[1].value_in_unit_system(unit.md_unit_system), value[0].value_in_unit_system(unit.md_unit_system)) for (name, value) in parameters.items() }
# Create harmonic oscillator testsystem
testsystem = testsystems.HarmonicOscillator(K=K, mass=mass)
system = testsystem.system
positions = testsystem.positions
# Get equilibrium samples from initial and final states
burn_in = 5 * 20 # 5 periods
thinning = 5 * 20 # 5 periods
# Collect forward and reverse work values
w_f = np.zeros([n_iterations], np.float64)
w_r = np.zeros([n_iterations], np.float64)
platform = openmm.Platform.getPlatformByName("Reference")
for direction in ['forward', 'reverse']:
positions = testsystem.positions
for iteration in range(n_iterations):
# Generate equilibrium sample
equilibrium_integrator = GHMCIntegrator(temperature=temperature, collision_rate=collision_rate, timestep=timestep)
equilibrium_context = openmm.Context(system, equilibrium_integrator, platform)
for (name, value) in parameters.items():
if direction == 'forward':
equilibrium_context.setParameter(name, value[0].value_in_unit_system(unit.md_unit_system))
else:
equilibrium_context.setParameter(name, value[1].value_in_unit_system(unit.md_unit_system))
equilibrium_context.setPositions(positions)
equilibrium_integrator.step(thinning)
positions = equilibrium_context.getState(getPositions=True).getPositions(asNumpy=True)
del equilibrium_context, equilibrium_integrator
# Generate nonequilibrium work sample
if direction == 'forward':
alchemical_functions = forward_functions
else:
alchemical_functions = reverse_functions
nonequilibrium_integrator = AlchemicalNonequilibriumLangevinIntegrator(temperature=temperature, collision_rate=collision_rate, timestep=timestep,
alchemical_functions=alchemical_functions, splitting=splitting, nsteps_neq=nsteps,
measure_shadow_work=True)
nonequilibrium_context = openmm.Context(system, nonequilibrium_integrator, platform)
nonequilibrium_context.setPositions(positions)
if nsteps == 0:
nonequilibrium_integrator.step(1) # need to execute at least one step
else:
nonequilibrium_integrator.step(nsteps)
if direction == 'forward':
w_f[iteration] = nonequilibrium_integrator.get_total_work(dimensionless=True)
else:
w_r[iteration] = nonequilibrium_integrator.get_total_work(dimensionless=True)
del nonequilibrium_context, nonequilibrium_integrator
dF, ddF = pymbar.BAR(w_f, w_r)
nsigma = np.abs(dF - dF_analytical) / ddF
print("analytical DeltaF: {:12.4f}, DeltaF: {:12.4f}, dDeltaF: {:12.4f}, nsigma: {:12.1f}".format(dF_analytical, dF, ddF, nsigma))
if nsigma > NSIGMA_MAX:
raise Exception("The free energy difference for the nonequilibrium switching for splitting '%s' and %d steps is not zero within statistical error." % (splitting, nsteps))
def notest_hamiltonian_exchange(mpicomm=None, verbose=True):
"""
Test that free energies and average potential energies of a 3D harmonic oscillator are correctly computed when running HamiltonianExchange.
TODO
* Integrate with test_replica_exchange.
* Test with different combinations of input parameters.
"""
if verbose and ((not mpicomm) or (mpicomm.rank==0)): sys.stdout.write("Testing Hamiltonian exchange facility with harmonic oscillators: ")
# Create test system of harmonic oscillators
testsystem = testsystems.HarmonicOscillatorArray()
[system, coordinates] = [testsystem.system, testsystem.positions]
# Define mass of carbon atom.
mass = 12.0 * units.amu
# Define thermodynamic states.
sigmas = [0.2, 0.3, 0.4] * units.angstroms # standard deviations: beta K = 1/sigma^2 so K = 1/(beta sigma^2)
temperature = 300.0 * units.kelvin # temperatures
seed_positions = list()
analytical_results = list()
f_i_analytical = list() # dimensionless free energies
u_i_analytical = list() # reduced potential
systems = list() # Systems list for HamiltonianExchange
for sigma in sigmas:
# Compute corresponding spring constant.
kB = units.BOLTZMANN_CONSTANT_kB * units.AVOGADRO_CONSTANT_NA
kT = kB * temperature # thermal energy
beta = 1.0 / kT # inverse temperature
K = 1.0 / (beta * sigma**2)
# Create harmonic oscillator system.
testsystem = testsystems.HarmonicOscillator(K=K, mass=mass, mm=openmm)
[system, positions] = [testsystem.system, testsystem.positions]
# Append to systems list.
systems.append(system)
# Append positions.
seed_positions.append(positions)
# Store analytical results.
results = computeHarmonicOscillatorExpectations(K, mass, temperature)
analytical_results.append(results)
f_i_analytical.append(results['f'])
reduced_potential = results['potential']['mean'] / kT
u_i_analytical.append(reduced_potential)
# DEBUG
print("")
print(seed_positions)
print(analytical_results)
print(u_i_analytical)
print(f_i_analytical)
print("")
# 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("Storing data in temporary file: %s" % str(store_filename))
# Create reference thermodynamic state.
reference_state = ThermodynamicState(systems[0], temperature=temperature)
# Create and configure simulation object.
simulation = HamiltonianExchange(store_filename, mpicomm=mpicomm)
simulation.create(reference_state, systems, seed_positions)
simulation.platform = openmm.Platform.getPlatformByName('Reference')
simulation.number_of_iterations = 200
simulation.timestep = 2.0 * units.femtoseconds
simulation.nsteps_per_iteration = 500
simulation.collision_rate = 9.2 / units.picosecond
simulation.verbose = False
simulation.show_mixing_statistics = False
# Run simulation.
utils.config_root_logger(True)
simulation.run() # run the simulation
utils.config_root_logger(False)
# Stop here if not root node.
if mpicomm and (mpicomm.rank != 0): return
# Analyze simulation to compute free energies.
analysis = simulation.analyze()
# TODO: Check if deviations exceed tolerance.
Delta_f_ij = analysis['Delta_f_ij']
dDelta_f_ij = analysis['dDelta_f_ij']
error = 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)
if verbose: print("PASSED.")
return
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
from functools import partial
import nose
from pymbar import timeseries
from openmmtools import testsystems
from openmmtools.states import SamplerState, ThermodynamicState
from openmmtools.mcmc import *
# =============================================================================
# GLOBAL TEST CONSTANTS
# =============================================================================
# Test various combinations of systems and MCMC schemes
analytical_testsystems = [
("HarmonicOscillator", testsystems.HarmonicOscillator(),
GHMCMove(timestep=10.0 * unit.femtoseconds, n_steps=100)),
("HarmonicOscillator", testsystems.HarmonicOscillator(),
WeightedMove([(GHMCMove(timestep=10.0 * unit.femtoseconds,
n_steps=100), 0.5),
(HMCMove(timestep=10 * unit.femtosecond,
n_steps=10), 0.5)])),
("HarmonicOscillatorArray", testsystems.HarmonicOscillatorArray(N=4),
LangevinDynamicsMove(timestep=10.0 * unit.femtoseconds, n_steps=100)),
("IdealGas", testsystems.IdealGas(nparticles=216),
SequenceMove([
HMCMove(timestep=10 * unit.femtosecond, n_steps=10),
MonteCarloBarostatMove()
]))
]
