class Langevin():
_explorer = None
_initialized = False
_context = None
_integrator = None
_timestep = Quantity(value=2.0, unit=u.femtosecond)
_temperature = Quantity(value=298.0, unit=u.kelvin)
_collision_rate = Quantity(value=1.0, unit=1.0 / u.picosecond)
def __init__(self, explorer):
self._explorer = explorer
def _initialize(self):
system = self._explorer.context.getSystem()
platform = self._explorer.context.getPlatform()
properties = {}
if platform.getName() == 'CUDA':
properties['CudaPrecision'] = 'mixed'
self._integrator = LangevinIntegrator(self._temperature,
self._collision_rate,
self._timestep)
self._context = Context(system, self._integrator, platform, properties)
self._initialized = True
def set_parameters(self,
temperature=Quantity(value=298.0, unit=u.kelvin),
collision_rate=Quantity(value=1.0,
unit=1.0 / u.picosecond),
timestep=Quantity(value=2.0, unit=u.femtosecond)):
self._timestep = timestep
self._temperature = temperature
self._collision_rate = collision_rate
if self._initialized:
self._integrator.setFriction(
self._collision_rate.value_in_unit(u.picosecond**-1))
self._integrator.setTemperature(
self._temperature.value_in_unit(u.kelvin))
self._integrator.setStepSize(
self._timestep.value_in_unit(u.picoseconds))
else:
self._initialize()
def get_parameters(self):
parameters = {
'timestep': self._timestep,
'temperature': self._temperature,
'collision_rate': self._collision_rate
}
return parameters
def replicate_parameters(self, explorer):
timestep = explorer.md.langevin._timestep
temperature = explorer.md.langevin._temperature
collision_rate = explorer.md.langevin._collision_rate
self.set_paramters(temperature, collision_rate, timestep)
def _set_coordinates(self, coordinates):
self._context.setPositions(coordinates)
def _get_coordinates(self):
return self._context.getState(getPositions=True).getPositions(
asNumpy=True)
def _set_velocities(self, velocities):
self._context.setVelocities(velocities)
def _get_velocities(self):
return self._context.getState(getVelocities=True).getVelocities(
asNumpy=True)
def _coordinates_to_explorer(self):
self._explorer.set_coordinates(self._get_coordinates())
def _coordinates_from_explorer(self):
self._set_coordinates(self._explorer.get_coordinates())
def _velocities_to_explorer(self):
self._explorer.set_velocities(self._get_velocities())
def _velocities_from_explorer(self):
self._set_velocities(self._explorer.get_velocities())
def get_time(self):
return self._context.getState().getTime()
def run(self, steps=0):
if not self._initialized:
self._initialize()
self._coordinates_from_explorer()
self._velocities_from_explorer()
self._integrator.step(steps)
self._coordinates_to_explorer()
self._velocities_to_explorer()
def __call__(self, *args, **kwargs):
return self.run(*args, **kwargs)
class MCMCSampler(object):
"""
Markov chain Monte Carlo (MCMC) sampler.
This is a minimal functional implementation placeholder until we can replace this with MCMCSampler from `openmmmcmc`.
Properties
----------
positions : simtk.unit.Quantity of size [nparticles,3] with units compatible with nanometers
The current positions.
iteration : int
Iterations completed.
verbose : bool
If True, verbose output is printed
Examples
--------
>>> # Create a test system
>>> test = testsystems.AlanineDipeptideVacuum()
>>> # Create a sampler state.
>>> sampler_state = SamplerState(positions=test.positions)
>>> # Create a thermodynamic state.
>>> thermodynamic_state = ThermodynamicState(system=test.system, temperature=298.0*unit.kelvin)
>>> # Create an MCMC sampler
>>> sampler = MCMCSampler(thermodynamic_state, sampler_state)
>>> # Run the sampler
>>> sampler.verbose = False
>>> sampler.run()
"""
def __init__(self, thermodynamic_state=None, sampler_state=None, platform=None, ncfile=None):
"""
Create an MCMC sampler.
Parameters
----------
thermodynamic_state : ThermodynamicState
The thermodynamic state to simulate
sampler_state : SamplerState
The initial sampler state to simulate from.
platform : simtk.openmm.Platform, optional, default=None
If specified, this platform will be used
ncfile : netCDF4.Dataset, optional, default=None
NetCDF storage file.
"""
if thermodynamic_state is None:
raise Exception("'thermodynamic_state' must be specified")
if sampler_state is None:
raise Exception("'sampler_state' must be specified")
self.thermodynamic_state = thermodynamic_state
self.sampler_state = sampler_state
# Initialize
self.iteration = 0
# For GHMC / Langevin integrator
self.collision_rate = 1.0 / unit.picoseconds
self.timestep = 2.0 * unit.femtoseconds
self.nsteps = 500 # number of steps per update
self.verbose = True
self.platform = platform
# For writing PDB files
self.pdbfile = None
self.topology = None
self._timing = dict()
self._initializeNetCDF(ncfile)
self._initialized = False
def _initialize(self):
# Create an integrator
integrator_name = 'Langevin'
if integrator_name == 'GHMC':
from openmmtools.integrators import GHMCIntegrator
self.integrator = GHMCIntegrator(temperature=self.thermodynamic_state.temperature, collision_rate=self.collision_rate, timestep=self.timestep)
elif integrator_name == 'Langevin':
from simtk.openmm import LangevinIntegrator
self.integrator = LangevinIntegrator(self.thermodynamic_state.temperature, self.collision_rate, self.timestep)
else:
raise Exception("integrator_name '%s' not valid." % (integrator_name))
# Create a Context
if self.platform is not None:
self.context = openmm.Context(self.thermodynamic_state.system, self.integrator, self.platform)
else:
self.context = openmm.Context(self.thermodynamic_state.system, self.integrator)
self.thermodynamic_state.update_context(self.context, self.integrator)
self.sampler_state.update_context(self.context)
self.context.setVelocitiesToTemperature(self.thermodynamic_state.temperature)
self._initialized = True
def _initializeNetCDF(self, ncfile):
self.ncfile = ncfile
if self.ncfile == None:
return
natoms = self.thermodynamic_state.system.getNumParticles()
self.ncfile.createDimension('iterations', None)
self.ncfile.createDimension('atoms', natoms) # TODO: What do we do if dimension can change?
self.ncfile.createDimension('spatial', 3)
self.ncfile.createVariable('positions', 'f4', dimensions=('iterations', 'atoms', 'spatial'), zlib=True, chunksizes=(1,natoms,3))
self.ncfile.createVariable('box_vectors', 'f4', dimensions=('iterations', 'spatial', 'spatial'), zlib=True, chunksizes=(1,3,3))
self.ncfile.createVariable('potential', 'f8', dimensions=('iterations',), chunksizes=(1,))
self.ncfile.createVariable('sample_positions_time', 'f4', dimensions=('iterations',), chunksizes=(1,))
# Weight adaptation information
self.ncfile.createVariable('stage', 'i2', dimensions=('iterations',), chunksizes=(1,))
self.ncfile.createVariable('gamma', 'f8', dimensions=('iterations',), chunksizes=(1,))
def update(self):
"""
Update the sampler with one step of sampling.
"""
if not self._initialized:
self._initialize()
if self.verbose:
print("." * 80)
print("MCMC sampler iteration %d" % self.iteration)
initial_time = time.time()
# Reset statistics
if hasattr(self.integrator, 'setGlobalVariableByName'):
self.integrator.setGlobalVariableByName('naccept', 0)
# Take some steps
self.integrator.step(self.nsteps)
# Get new sampler state.
self.sampler_state = SamplerState.createFromContext(self.context)
# Report statistics
if hasattr(self.integrator, 'getGlobalVariableByName'):
naccept = self.integrator.getGlobalVariableByName('naccept')
fraction_accepted = float(naccept) / float(self.nsteps)
if self.verbose: print("Accepted %d / %d GHMC steps (%.2f%%)." % (naccept, self.nsteps, fraction_accepted * 100))
final_time = time.time()
elapsed_time = final_time - initial_time
self._timing['sample positions'] = elapsed_time
if self.verbose:
final_energy = self.context.getState(getEnergy=True).getPotentialEnergy() * self.thermodynamic_state.beta
print('Final energy is %12.3f kT' % (final_energy))
print('elapsed time %8.3f s' % elapsed_time)
if self.ncfile:
self.ncfile.variables['positions'][self.iteration,:,:] = self.sampler_state.positions[:,:] / unit.nanometers
for k in range(3):
self.ncfile.variables['box_vectors'][self.iteration,k,:] = self.sampler_state.box_vectors[k,:] / unit.nanometers
self.ncfile.variables['potential'][self.iteration] = self.thermodynamic_state.beta * self.context.getState(getEnergy=True).getPotentialEnergy()
self.ncfile.variables['sample_positions_time'][self.iteration] = elapsed_time
# Increment iteration count
self.iteration += 1
if self.verbose:
print("." * 80)
if self.pdbfile is not None:
print("Writing frame...")
from simtk.openmm.app import PDBFile
PDBFile.writeModel(self.topology, self.sampler_state.positions, self.pdbfile, self.iteration)
self.pdbfile.flush()
def run(self, niterations=1):
"""
Run the sampler for the specified number of iterations
Parameters
----------
niterations : int, optional, default=1
Number of iterations to run the sampler for.
"""
for iteration in range(niterations):
self.update()