| from abc import abstractmethod |
|
|
| from potentials.base import Potential |
| from potentials.md_utils import ForceReporter |
|
|
|
|
| class MoleculePotential(Potential): |
| """ |
| This is a baseclass for implementing molecular potentials using some framework. |
| """ |
|
|
| def __init__(self, start_file, index, save_file=None): |
| super().__init__() |
| self.start_file = start_file |
| self.index = index |
| self.save_file = save_file |
|
|
| self.pdb, self.simulation, self.external_force = self.setup() |
|
|
| self.reporter = ForceReporter(1) |
| self.simulation.reporters.append(self.reporter) |
|
|
| if self.index == 0: |
| self.simulation.step(1) |
| self.simulation.minimizeEnergy() |
| self.simulation.saveCheckpoint(self.get_position_file()) |
|
|
| if self.index > -1: |
| self.reset() |
|
|
| @abstractmethod |
| def setup(self): |
| """ |
| Initializes all Framework specific configurations |
| :return pdb: Description of the atom locations |
| :return simulation: Simulation environment that runs the MD |
| :return external_force: Reference to custom external force implementation to be updated by control |
| """ |
| pass |
|
|
| @abstractmethod |
| def get_position_file(self): |
| """ |
| Abstact method to point towards location of intermediate files. |
| :return path: Path to intermediate files |
| """ |
| pass |
|
|
| def drift(self, forces): |
| """ |
| This is the main workhorse. This function calculates the entire dynamics update. This includes the passive |
| dynamics as well as the controlled dynamics. This is done by offloading the computation step entirely to |
| OpenMM through the use of a "Custom External Force". |
| :param forces: Control force given by the policy |
| :return: New positions and velocity of the system. |
| """ |
| 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) |
|
|
| new_pos = self.reporter.latest_positions.copy() |
| new_forces = self.reporter.latest_forces.copy() |
|
|
| return new_pos, new_forces |
|
|
| def reset(self): |
| """ |
| Resets the MD simulation to the initial minimum energy conformation. |
| """ |
| self.simulation.loadCheckpoint(self.get_position_file()) |
| for i in range(len(self.pdb.positions)): |
| self.external_force.setParticleParameters(i, i, [0, 0, 0]) |
| self.external_force.updateParametersInContext(self.simulation.context) |
|
|
| self.simulation.step(1) |
| self.simulation.minimizeEnergy() |
|
|
| def latest_position(self): |
| """ |
| Returns the positions after the latest update step. |
| :return: Current coordinates of atoms |
| """ |
| return self.simulation.state.getPositions() |
|
|
