Source code for meld.system.patchers.rdc_integrator

"""
Add RDC alignments to the system
"""

import math
import openmm as mm  # type: ignore
from openmm import unit as u


GAS_CONSTANT = 8.314e-3


[docs]class CustomRDCIntegrator: _temperature: float _stepsize: int _num_alignments: int _alignment_mass: float _wrapped_integrator: mm.CustomIntegrator
[docs] def __init__( self, temperature, friction_coefficient, step_size, num_alignments, alignment_mass=1e4, ): self._temperature = self._handle_temp(temperature) if isinstance(friction_coefficient, u.Quantity): friction_coefficient = friction_coefficient.value_in_unit( (1 / u.picoseconds).unit ) assert friction_coefficient > 0 self._friction_coefficient = friction_coefficient if isinstance(step_size, u.Quantity): step_size = step_size.value_in_unit(u.picosecond) assert step_size > 0 assert num_alignments > 0 assert alignment_mass > 0 if isinstance(step_size, u.Quantity): step_size = step_size.value_in_unit(u.picosecond) self._step_size = step_size self._num_alignments = num_alignments self._alignment_mass = alignment_mass self._wrapped_integrator = self._create_custom_integrator()
def getTemperature(self): return self._temperature * u.kelvin def setTemperature(self, temp): self._temperature = self._handle_temp(temp) self._wrapped_integrator.setGlobalVariableByName( "kT", self._temperature * GAS_CONSTANT ) @property def num_alignments(self): return self._num_alignments def step(self, steps): self._wrapped_integrator.step(steps) def _create_custom_integrator(self): # # Setup variables # integrator = mm.CustomIntegrator(self._step_size) integrator.addGlobalVariable( "a", math.exp(-self._friction_coefficient * self._step_size) ) integrator.addGlobalVariable( "b", math.sqrt(1 - math.exp(-2 * self._friction_coefficient * self._step_size)), ) integrator.addGlobalVariable("kT", GAS_CONSTANT * self._temperature) # Add in the alignment tensor components for i in range(self._num_alignments): for j in range(5): integrator.addGlobalVariable(f"rdc_{i}_s{j + 1}", 0.0) integrator.addGlobalVariable(f"rdc_{i}_s{j + 1}_vel", 0.0) integrator.addPerDofVariable("x1", 0) # # Steps during integration # integrator.addUpdateContextState() integrator.addComputePerDof("v", "v + dt*f/m") for i in range(self._num_alignments): for j in range(5): integrator.addComputeGlobal( f"rdc_{i}_s{j + 1}_vel", f"rdc_{i}_s{j + 1}_vel - dt * deriv(energy, rdc_{i}_s{j + 1} / {self._alignment_mass}", ) integrator.addConstrainVelocities() integrator.addComputePerDof("x", "x + 0.5*dt*v") for i in range(self._num_alignments): for j in range(5): integrator.addComputeGlobal( f"rdc_{i}_s{j + 1}", f"rdc_{i}_s{j + 1} + 0.5*dt*rdc_{i}_s{j + 1}_vel", ) integrator.addComputePerDof("v", "a*v + b*sqrt(kT/m)*gaussian") for i in range(self._num_alignments): for j in range(5): integrator.addComputeGlobal( f"rdc_{i}_s{j + 1}_vel", f"a*rdc_{i}_s{j + 1}_vel + b*sqrt(kT/{self._alignment_mass})*gaussian", ) integrator.addComputePerDof("x", "x + 0.5*dt*v") for i in range(self._num_alignments): for j in range(5): integrator.addComputeGlobal( f"rdc_{i}_s{j + 1}", f"rdc_{i}_s{j + 1} + 0.5*dt*rdc_{i}_s{j + 1}_vel", ) integrator.addComputePerDof("x1", "x") integrator.addConstrainPositions() integrator.addComputePerDof("v", "v + (x-x1)/dt") return integrator def _handle_temp(self, temp): if isinstance(temp, u.Quantity): temp = temp.value_in_unit(u.kelvin) assert temp > 0 return temp