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 Explorer():
topology = None
context = None
pbc = False
n_atoms = 0
n_dof = 0
md = None
quench = None
move = None
def __init__(self, topology=None, system=None, pbc=False, platform='CUDA'):
from .md import MD
from .quench import Quench
from .move import Move
from .distance import Distance
from .acceptance import Acceptance
if topology is None:
raise ValueError('topology is needed')
if system is None:
raise ValueError('system is needed')
integrator = LangevinIntegrator(0 * u.kelvin, 1.0 / u.picoseconds,
2.0 * u.femtoseconds)
#integrator.setConstraintTolerance(0.00001)
if platform == 'CUDA':
platform = Platform.getPlatformByName('CUDA')
properties = {'CudaPrecision': 'mixed'}
elif platform == 'CPU':
platform = Platform.getPlatformByName('CPU')
properties = {}
self.topology = topology
self.context = Context(system, integrator, platform, properties)
self.n_atoms = msm.get(self.context, target='system', n_atoms=True)
self.n_dof = 0
for i in range(system.getNumParticles()):
if system.getParticleMass(i) > 0 * u.dalton:
self.n_dof += 3
for i in range(system.getNumConstraints()):
p1, p2, distance = system.getConstraintParameters(i)
if system.getParticleMass(
p1) > 0 * u.dalton or system.getParticleMass(
p2) > 0 * u.dalton:
self.n_dof -= 1
if any(
type(system.getForce(i)) == CMMotionRemover
for i in range(system.getNumForces())):
self.n_dof -= 3
self.pbc = pbc
if self.pbc:
raise NotImplementedError
self.md = MD(self)
self.quench = Quench(self)
self.move = Move(self)
self.distance = Distance(self)
self.acceptance = Acceptance(self)
def _copy(self):
topology = self.topology
coordinates = self.get_coordinates()
system = self.context.getSystem()
platform = self.context.getPlatform().getName()
pbc = self.pbc
tmp_explorer = Explorer(topology, system, pbc, platform)
tmp_explorer.set_coordinates(coordinates)
for ii, jj in vars(tmp_explorer.md).items():
if not ii.startswith('_'):
if jj._initialized:
jj.replicate_parameters(self)
for ii, jj in vars(tmp_explorer.quench).items():
if not ii.startswith('_'):
if jj._initialized:
jj.replicate_parameters(self)
for ii, jj in vars(tmp_explorer.move).items():
if not ii.startswith('_'):
if jj._initialized:
jj.replicate_parameters(self)
return tmp_explorer
def replicate(self, times=1):
from copy import deepcopy
if times == 1:
output = self._copy()
else:
output = [self._copy() for ii in range(times)]
return output
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 set_velocities_to_temperature(self, temperature):
self.context.setVelocitiesToTemperature(temperature)
def get_velocities(self):
return self.context.getState(getVelocities=True).getVelocities(
asNumpy=True)
def get_temperature(self):
return (2 * self.context.getState(getEnergy=True).getKineticEnergy() /
(self.n_dof * u.MOLAR_GAS_CONSTANT_R)).in_units_of(u.kelvin)
def get_potential_energy(self):
energy = self.context.getState(getEnergy=True).getPotentialEnergy()
return energy
def get_potential_energy_gradient(self):
gradient = -self.context.getState(getForces=True).getForces(
asNumpy=True)
gradient = gradient.ravel() * gradient.unit
return gradient
def get_potential_energy_hessian(self,
mass_weighted=False,
symmetric=True):
"""OpenMM single frame hessian evaluation
Since OpenMM doesnot provide a Hessian evaluation method, we used finite difference on forces
from: https://leeping.github.io/forcebalance/doc/html/api/openmmio_8py_source.html
Returns
-------
hessian: np.array with shape 3N x 3N, N = number of "real" atoms
The result hessian matrix.
The row indices are fx0, fy0, fz0, fx1, fy1, ...
The column indices are x0, y0, z0, x1, y1, ..
The unit is kilojoule / (nanometer^2 * mole * dalton) => 10^24 s^-2
"""
n_dof = self.n_atoms * 3
pos = self.get_coordinates()
hessian = np.empty(
(n_dof, n_dof),
dtype=float) * u.kilojoules_per_mole / (u.nanometers**2)
# finite difference step size
diff = 0.0001 * u.nanometer
coef = 1.0 / (2.0 * diff) # 1/2h
for i in range(self.n_atoms):
# loop over the x, y, z coordinates
for j in range(3):
# plus perturbation
pos[i][j] += diff
self.set_coordinates(pos)
grad_plus = self.get_potential_energy_gradient()
# minus perturbation
pos[i][j] -= 2 * diff
self.set_coordinates(pos)
grad_minus = self.get_potential_energy_gradient()
# set the perturbation back to zero
pos[i][j] += diff
# fill one row of the hessian matrix
hessian[i * 3 + j] = (grad_plus - grad_minus) * coef
if mass_weighted:
mass = np.array([
self.context.getSystem().getParticleMass(k).value_in_unit(
u.dalton) for k in range(self.n_atoms)
]) * u.dalton
mass_weight = 1.0 / np.sqrt(mass) * (mass.unit**-0.5)
mass_weight = np.repeat(mass_weight, 3) * mass_weight.unit
hessian = np.multiply(
hessian, mass_weight) * hessian.unit * mass_weight.unit
hessian = np.multiply(
hessian,
mass_weight[:, np.newaxis]) * hessian.unit * mass_weight.unit
# make hessian symmetric by averaging upper right and lower left
if symmetric:
hessian += hessian.T * hessian.unit
hessian *= 0.5
# recover the original position
self.set_coordinates(pos)
return hessian