def test_replica_exchange(mpicomm=None, verbose=True):
"""
Test that free energies and avergae potential energies of a 3D harmonic oscillator are correctly computed.
TODO
* Test ParallelTempering and HamiltonianExchange subclasses as well.
* Test with different combinations of input parameters.
"""
if verbose and ((not mpicomm) or (mpicomm.rank==0)): print "Testing replica exchange facility with harmonic oscillators: ",
# Create test system of harmonic oscillators
import testsystems
[system, coordinates] = testsystems.HarmonicOscillatorArray()
# Define mass of carbon atom.
mass = 12.0 * units.amu
# Define thermodynamic states.
from thermodynamics import ThermodynamicState
states = list() # thermodynamic states
sigmas = [0.2, 0.3, 0.4] * units.angstroms # standard deviations: beta K = 1/sigma^2 so K = 1/(beta sigma^2)
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 (sigma, temperature) in zip(sigmas, temperatures):
# 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.
[system, positions] = testsystems.HarmonicOscillator(K=K, mass=mass, mm=openmm)
# 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'])
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.
from repex import ReplicaExchange
simulation = ReplicaExchange(states, seed_positions, store_filename, mpicomm=mpicomm) # initialize the replica-exchange simulation
simulation.number_of_iterations = 1000 # set the simulation to only run 2 iterations
simulation.timestep = 2.0 * units.femtoseconds # set the timestep for integration
simulation.nsteps_per_iteration = 500 # run 500 timesteps per iteration
simulation.collision_rate = 0.001 / units.picosecond # DEBUG: Use a low collision rate
simulation.platform = openmm.Platform.getPlatformByName('Reference') # use reference platform
simulation.verbose = True # DEBUG
# Run simulation.
simulation.run() # run the simulation
# 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
import tempfile
file = tempfile.NamedTemporaryFile() # use a temporary file for testing -- you will want to keep this file, since it stores output and checkpoint data
# Select platform: one of 'Reference' (CPU-only), 'Cuda' (NVIDIA Cuda), or 'OpenCL' (for OS X 10.6 with OpenCL OpenMM compiled)
platform = simtk.openmm.Platform.getPlatformByName("OpenCL")
platform = simtk.openmm.Platform.getPlatformByName("Cuda")
# Set up device to bind to.
print "Selecting MPI communicator and selecting a GPU device..."
from mpi4py import MPI # MPI wrapper
hostname = os.uname()[1]
ngpus = 6 # number of GPUs per system
comm = MPI.COMM_WORLD # MPI communicator
deviceid = comm.rank % ngpus # select a unique GPU for this node assuming block allocation (not round-robin)
platform.setPropertyDefaultValue('CudaDeviceIndex', '%d' % deviceid) # select Cuda device index
platform.setPropertyDefaultValue('OpenCLDeviceIndex', '%d' % deviceid) # select OpenCL device index
print "node '%s' deviceid %d / %d, MPI rank %d / %d" % (hostname, deviceid, ngpus, comm.rank, comm.size)
# Make sure random number generators have unique seeds.
seed = numpy.random.randint(sys.maxint - comm.size) + comm.rank
numpy.random.seed(seed)
# Set up replica exchange simulation.
simulation = ReplicaExchange(states, coordinates, file.name, mpicomm=comm) # initialize the replica-exchange simulation
simulation.verbose = True
simulation.number_of_iterations = 100 # set the simulation to only run 2 iterations
simulation.timestep = 2.0 * units.femtoseconds # set the timestep for integration
simulation.nsteps_per_iteration = 500 # run 500 timesteps per iteration
simulation.platform = platform
simulation.run() # run the simulation