Source code for meld.runner.transform.restraints.cartesian

#
# All rights reserved
#

"""
This module implements transformers that add cartesian restraints
"""

import logging
from typing import List

import openmm as mm  # type: ignore
from openmm import app  # type: ignore

from meld import interfaces
from meld.runner import transform
from meld.runner.transform.restraints.util import _delete_from_always_active
from meld.system import density, mapping, options, param_sampling, restraints

logger = logging.getLogger(__name__)

FORCE_GROUP = 5


[docs]class CartesianRestraintTransformer(transform.TransformerBase): """ Transformer to handle Cartesian restraints """ force: mm.CustomExternalForce
[docs] def __init__( self, param_manager: param_sampling.ParameterManager, mapper: mapping.PeakMapManager, density_manager: density.DensityManager, builder_info: dict, options: options.RunOptions, always_active_restraints: List[restraints.Restraint], selectively_active_restraints: List[restraints.SelectivelyActiveCollection], ) -> None: self.use_pbc = builder_info.get("solvation", "implicit") == "explicit" self.restraints = [ r for r in always_active_restraints if isinstance(r, restraints.CartesianRestraint) ] _delete_from_always_active(self.restraints, always_active_restraints) if self.restraints: self.active = True else: self.active = False
[docs] def add_interactions( self, state: interfaces.IState, system: mm.System, topology: app.Topology ) -> mm.System: if self.active: if self.use_pbc: cartesian_force = mm.CustomExternalForce( "0.5 * cart_force_const * r_eff^2;" "r_eff = max(0.0, r - cart_delta);" "r = periodicdistance(x, y, z, cart_x, cart_y, cart_z)" ) else: cartesian_force = mm.CustomExternalForce( "0.5 * cart_force_const * r_eff^2;" "r_eff = max(0.0, r - cart_delta);" "r = sqrt(dx*dx + dy*dy + dz*dz);" "dx = x - cart_x;" "dy = y - cart_y;" "dz = z - cart_z;" ) cartesian_force.addPerParticleParameter("cart_x") cartesian_force.addPerParticleParameter("cart_y") cartesian_force.addPerParticleParameter("cart_z") cartesian_force.addPerParticleParameter("cart_delta") cartesian_force.addPerParticleParameter("cart_force_const") # add the atoms for r in self.restraints: weight = r.force_const cartesian_force.addParticle( r.atom_index, [r.x, r.y, r.z, r.delta, weight] ) cartesian_force.setForceGroup(FORCE_GROUP) system.addForce(cartesian_force) self.force = cartesian_force return system
[docs] def update( self, state: interfaces.IState, simulation: app.Simulation, alpha: float, timestep: int, ) -> None: if self.active: for index, r in enumerate(self.restraints): weight = r.force_const * r.scaler(alpha) * r.ramp(timestep) self.force.setParticleParameters( index, r.atom_index, [r.x, r.y, r.z, r.delta, weight] ) self.force.updateParametersInContext(simulation.context)
[docs]class YZCartesianTransformer(transform.TransformerBase): """ Transformer to handle YZCartesian restraints """ force: mm.CustomExternalForce
[docs] def __init__( self, param_manager: param_sampling.ParameterManager, mapper: mapping.PeakMapManager, density_manager: density.DensityManager, builder_info: dict, options: options.RunOptions, always_active_restraints: List[restraints.Restraint], selectively_active_restraints: List[restraints.SelectivelyActiveCollection], ) -> None: self.use_pbc = builder_info.get("solvation", "implicit") == "explicit" self.restraints = [ r for r in always_active_restraints if isinstance(r, restraints.YZCartesianRestraint) ] _delete_from_always_active(self.restraints, always_active_restraints) if self.restraints: self.active = True else: self.active = False
[docs] def add_interactions( self, state: interfaces.IState, system: mm.System, topology: app.Topology ) -> mm.System: if self.active: # create the confinement force if self.use_pbc: cartesian_force = mm.CustomExternalForce( "0.5 * cart_force_const * r_eff^2;" "r_eff = max(0.0, r - cart_delta);" "r = periodicdistance(0, y, z, 0, cart_y, cart_z);" ) else: cartesian_force = mm.CustomExternalForce( "0.5 * cart_force_const * r_eff^2;" "r_eff = max(0.0, r - cart_delta);" "r = sqrt(r2);" "r2 = dy*dy + dz*dz;" "dy = y - cart_y;" "dz = z - cart_z;" ) cartesian_force.addPerParticleParameter("cart_y") cartesian_force.addPerParticleParameter("cart_z") cartesian_force.addPerParticleParameter("cart_delta") cartesian_force.addPerParticleParameter("cart_force_const") # add the atoms for r in self.restraints: weight = r.force_const cartesian_force.addParticle(r.atom_index, [r.y, r.z, r.delta, weight]) cartesian_force.setForceGroup(FORCE_GROUP) system.addForce(cartesian_force) self.force = cartesian_force return system
[docs] def update( self, state: interfaces.IState, simulation: app.Simulation, alpha: float, timestep: int, ) -> None: if self.active: for index, r in enumerate(self.restraints): weight = r.force_const * r.scaler(alpha) * r.ramp(timestep) self.force.setParticleParameters( index, r.atom_index, [r.y, r.z, r.delta, weight] ) self.force.updateParametersInContext(simulation.context)