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()