|
|
|
|
|
import numpy as np |
|
|
import openmm.unit as unit |
|
|
from abc import abstractmethod, ABC |
|
|
from scipy.constants import physical_constants |
|
|
|
|
|
|
|
|
class BaseDynamics(ABC): |
|
|
def __init__(self, args, state): |
|
|
super().__init__() |
|
|
self.start_file = f"./data/{args.molecule}/{state}.pdb" |
|
|
self.temperature = args.temperature * unit.kelvin |
|
|
self.friction = args.friction / unit.femtoseconds |
|
|
self.timestep = args.timestep * unit.femtoseconds |
|
|
self.pdb, self.integrator, self.simulation, self.external_force = self.setup() |
|
|
self.get_md_info() |
|
|
self.simulation.minimizeEnergy() |
|
|
self.position = self.report()[0] |
|
|
|
|
|
@abstractmethod |
|
|
def setup(self): |
|
|
pass |
|
|
|
|
|
def get_md_info(self): |
|
|
self.num_particles = self.simulation.system.getNumParticles() |
|
|
m = np.array([ |
|
|
self.simulation.system.getParticleMass(i).value_in_unit(unit.dalton) |
|
|
for i in range(self.num_particles) |
|
|
]) |
|
|
self.heavy_atoms = m > 1.1 |
|
|
m = unit.Quantity(m, unit.dalton) |
|
|
unadjusted_variance = ( |
|
|
2 * self.timestep * self.friction * unit.BOLTZMANN_CONSTANT_kB * self.temperature / m[:, None] |
|
|
) |
|
|
std_SI_units = ( |
|
|
1 / physical_constants["unified atomic mass unit"][0] * |
|
|
unadjusted_variance.value_in_unit(unit.joule / unit.dalton) |
|
|
) |
|
|
self.std = unit.Quantity( |
|
|
np.sqrt(std_SI_units), unit.meter / unit.second |
|
|
).value_in_unit(unit.nanometer / unit.femtosecond) |
|
|
self.m = m.value_in_unit(unit.dalton) |
|
|
|
|
|
def step(self, forces): |
|
|
for i in range(forces.shape[0]): |
|
|
self.external_force.setParticleParameters(i, i, forces[i]) |
|
|
self.external_force.updateParametersInContext(self.simulation.context) |
|
|
self.simulation.step(1) |
|
|
|
|
|
def report(self): |
|
|
state = self.simulation.context.getState(getPositions=True, getForces=True) |
|
|
positions = state.getPositions().value_in_unit(unit.nanometer) |
|
|
forces = state.getForces().value_in_unit( |
|
|
unit.dalton * unit.nanometer / unit.femtosecond ** 2 |
|
|
) |
|
|
return positions, forces |
|
|
|
|
|
def reset(self): |
|
|
for i in range(len(self.position)): |
|
|
self.external_force.setParticleParameters(i, i, [0, 0, 0]) |
|
|
self.external_force.updateParametersInContext(self.simulation.context) |
|
|
self.simulation.context.setPositions(self.position) |
|
|
self.simulation.context.setVelocitiesToTemperature(self.temperature) |
|
|
|
|
|
def set_temperature(self, temperature): |
|
|
self.integrator.setTemperature(temperature * unit.kelvin) |
|
|
|
|
|
def energy_function(self, positions): |
|
|
forces, potentials = [], [] |
|
|
for i in range(len(positions)): |
|
|
self.simulation.context.setPositions(positions[i]) |
|
|
state = self.simulation.context.getState(getForces=True, getEnergy=True) |
|
|
force = state.getForces().value_in_unit( |
|
|
unit.dalton * unit.nanometer / unit.femtosecond ** 2 |
|
|
) |
|
|
potential = state.getPotentialEnergy().value_in_unit( |
|
|
unit.kilojoules / unit.mole |
|
|
) |
|
|
forces.append(force) |
|
|
potentials.append(potential) |
|
|
return np.array(forces), np.array(potentials) |
|
|
|
