Source code for meld.runner.transform.rest2

#
# Copyright 2015 by Justin MacCallum, Alberto Perez, Ken Dill
# All rights reserved
#

"""
This module implements a transformer that implements REST2.
"""

import math
from typing import Dict, List, Tuple

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

from meld import interfaces
from meld.runner.transform import TransformerBase
from meld.system import (
    density,
    mapping,
    options,
    param_sampling,
    restraints,
    temperature,
)


[docs]class REST2Transformer(TransformerBase): """ An implementation of REST2 We use the updated version of Replica Exchange with Solute Scaling. Warnings: Currently, the REST2 implmentation in MELD has a limitation that any CMAP / AMAP potentials are not scaled. References: L. Wang, R.A. Friesner, B.J. Berne, Replica exchange with solute scaling: a more efficient version of replica exchange with solute tempering. """ scaler: temperature.REST2Scaler ramp: restraints.TimeRamp nb_force: mm.NonbondedForce dihedral_force: mm.PeriodicTorsionForce protein_nonbonded_params: Dict[int, Tuple[float, float, float]] protein_exception_params: Dict[int, Tuple[int, int, float, float, float]] protein_dihedrals: Dict[int, Tuple[int, int, int, int, int, float, float]]
[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.active = False self.protein_nonbonded_params = {} self.protein_exception_params = {} self.protein_dihedrals = {} if options.use_rest2: self.active = True if isinstance(options.rest2_scaler, temperature.REST2Scaler): self.scaler = options.rest2_scaler else: raise ValueError("Trying to use REST2 without a REST2Scaler") if builder_info["builder"] != "amber": raise ValueError("REST2 only works with Amber") if builder_info["solvation"] != "explicit": raise ValueError("Cannot use REST2 without explicit solvent")
[docs] def finalize( self, state: interfaces.IState, system: mm.System, topology: app.Topology ) -> None: if self.active: nonsolvent_atoms = self._find_nonsolvent_atoms(topology) self._find_nb_force(system) self._find_dihedral_force(system) self._gather_nonbonded_params(nonsolvent_atoms) self._gather_dihedral_params(nonsolvent_atoms, topology)
[docs] def update( self, state: interfaces.IState, simulation: app.Simulation, alpha: float, timestep: int, ) -> None: if self.active: scale = self.scaler(alpha) self._update_nonbonded(simulation, scale) self._update_dihedrals(simulation, scale)
def _find_nonsolvent_atoms(self, topology: app.Topology) -> List[int]: solvent_residue_names = ["WAT", "SOL", "H2O", "HOH"] nonsolvent_atoms = [] for atom in topology.atoms(): if not atom.residue.name in solvent_residue_names: nonsolvent_atoms.append(atom.index) return nonsolvent_atoms def _gather_nonbonded_params(self, nonsolvent_atoms: List[int]) -> None: # gather the nonbonded parameters for index in nonsolvent_atoms: self.protein_nonbonded_params[index] = self.nb_force.getParticleParameters( index ) # gather the exception parameters for param_index in range(self.nb_force.getNumExceptions()): params = self.nb_force.getExceptionParameters(param_index) if params[0] in nonsolvent_atoms and params[1] in nonsolvent_atoms: self.protein_exception_params[param_index] = params def _gather_dihedral_params( self, nonsolvent_atoms: List[int], topology: app.Topology ) -> None: bond_idxs = [sorted([i.index, j.index]) for i, j in topology.bonds()] for parm_index in range(self.dihedral_force.getNumTorsions()): params = self.dihedral_force.getTorsionParameters(parm_index) i, j, k, l, _, _, _ = params not_solvent = ( i in nonsolvent_atoms and j in nonsolvent_atoms and k in nonsolvent_atoms and l in nonsolvent_atoms ) not_improper = ( sorted([i, j]) in bond_idxs and sorted([j, k]) in bond_idxs and sorted([k, l]) in bond_idxs ) # only modify dihedrals involving non-solvent atoms # and those where sequential atoms are bonded (proper dihedrals) if not_solvent and not_improper: self.protein_dihedrals[parm_index] = params def _find_nb_force(self, system: mm.System) -> None: forces = [system.getForce(i) for i in range(system.getNumForces())] nb_forces = [f for f in forces if isinstance(f, mm.NonbondedForce)] if not nb_forces: raise RuntimeError("REST2 could not find NonbondedForce") if len(nb_forces) > 1: raise RuntimeError("REST2 found more than one NonbondedForce") self.nb_force = nb_forces[0] def _find_dihedral_force(self, system: mm.System) -> None: forces = [system.getForce(i) for i in range(system.getNumForces())] dihed_forces = [f for f in forces if isinstance(f, mm.PeriodicTorsionForce)] if not dihed_forces: raise RuntimeError("REST2 could not find PeriodicTorsionForce") if len(dihed_forces) > 1: raise RuntimeError("REST2 found more than one PeriodicTorsionForce") self.dihedral_force = dihed_forces[0] def _update_nonbonded(self, simulation: app.Simulation, scale: float) -> None: for index in self.protein_nonbonded_params: nb_params = self.protein_nonbonded_params[index] q, sigma, eps = nb_params self.nb_force.setParticleParameters( index, q * math.sqrt(scale), sigma, eps * scale ) for index in self.protein_exception_params: except_params = self.protein_exception_params[index] i, j, q, sigma, eps = except_params self.nb_force.setExceptionParameters( index, i, j, q * scale, sigma, eps * scale ) self.nb_force.updateParametersInContext(simulation.context) def _update_dihedrals(self, simulation: app.Simulation, scale: float) -> None: for index in self.protein_dihedrals: params = self.protein_dihedrals[index] i, j, k, l, mult, phi, fc = params self.dihedral_force.setTorsionParameters( index, i, j, k, l, mult, phi, fc * scale ) self.dihedral_force.updateParametersInContext(simulation.context)