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
