File size: 3,277 Bytes
7efee70 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
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)
|