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