#!/usr/bin/env python
# =============================================================================
# MODULE DOCSTRING
# =============================================================================
"""
Custom OpenMM Forces classes and utilities.
"""
# =============================================================================
# GLOBAL IMPORTS
# =============================================================================
import abc
import collections
import copy
import inspect
import logging
import math
import re
import scipy
import numpy as np
try:
import openmm
from openmm import unit
except ImportError: # OpenMM < 7.6
from simtk import openmm, unit
from openmmtools import utils
from openmmtools.constants import ONE_4PI_EPS0, STANDARD_STATE_VOLUME
logger = logging.getLogger(__name__)
# =============================================================================
# UTILITY FUNCTIONS
# =============================================================================
[docs]
class MultipleForcesError(Exception):
"""Error raised when multiple forces of the same class are found."""
pass
[docs]
class NoForceFoundError(Exception):
"""Error raised when no forces matching the given criteria are found."""
pass
[docs]
def iterate_forces(system):
"""Iterate over and restore the Python interface of the forces in the system."""
for force in system.getForces():
utils.RestorableOpenMMObject.restore_interface(force)
yield force
# Yield empty generator if the system has no forces.
return
[docs]
def find_forces(system, force_type, only_one=False, include_subclasses=False):
"""Find all the ``Force`` object of a given type in an OpenMM system.
Parameters
----------
system : openmm.System
The system to search.
force_type : str, or type
The class of the force to search, or a regular expression that
is used to match its name. Note that ``re.match()`` is used in
this case, not ``re.search()``. The ``iter_subclasses`` argument
must be False when this is a string.
only_one : bool
If True, an exception is raised when multiple forces of the same
type are found in the system, and only a single force is returned.
include_subclasses : bool, optional
If True, all forces inheriting from ``force_type`` are returned
as well (default is False). This can't be enabled if `force_type``
is not a class.
Returns
-------
forces : OrderedDict or tuple
If ``only_one`` is False, a dictionary force_index: force is returned
with all the forces matching the criteria. Otherwise,, a single pair
``(force_idx, force)`` is returned.
Raises
------
NoForceFoundError
If ``only_one`` is True and no forces matching the criteria are found.
MultipleForcesError
If ``only_one`` is True and multiple forces matching the criteria
are found
Examples
--------
The ``only_one`` flag can be used to retrieve a single force.
>>> from openmmtools import testsystems
>>> system = testsystems.TolueneVacuum().system
>>> force_index, force = find_forces(system, openmm.NonbondedForce, only_one=True)
>>> force.__class__.__name__
'NonbondedForce'
It is possible to search for force subclasses.
>>> class MyHarmonicForce(utils.RestorableOpenMMObject, openmm.CustomBondForce):
... pass
>>> force_idx = system.addForce(openmm.CustomBondForce('0.0'))
>>> force_idx = system.addForce(MyHarmonicForce('0.0'))
>>> forces = find_forces(system, openmm.CustomBondForce, include_subclasses=True)
>>> [(force_idx, force.__class__.__name__) for force_idx, force in forces.items()]
[(5, 'CustomBondForce'), (6, 'MyHarmonicForce')]
A regular expression can be used instead of a class.
>>> forces = find_forces(system, 'HarmonicAngleForce')
>>> [(force_idx, force.__class__.__name__) for force_idx, force in forces.items()]
[(1, 'HarmonicAngleForce')]
>>> forces = find_forces(system, '.*Harmonic.*')
>>> [(force_idx, force.__class__.__name__) for force_idx, force in forces.items()]
[(0, 'HarmonicBondForce'), (1, 'HarmonicAngleForce'), (6, 'MyHarmonicForce')]
"""
# Handle force_type argument when it's not a class.
re_pattern = None
if not inspect.isclass(force_type):
re_pattern = re.compile(force_type)
# Find all forces matching the force_type.
forces = {}
for force_idx, force in enumerate(iterate_forces(system)):
# Check force name.
if re_pattern is not None:
if re_pattern.match(force.__class__.__name__):
forces[force_idx] = force
# Check if the force class matches the requirements.
elif type(force) is force_type or (include_subclasses and isinstance(force, force_type)):
forces[force_idx] = force
# Second pass to find all subclasses of the matching forces.
if include_subclasses and re_pattern is not None:
matched_force_classes = [force.__class__ for force in forces.values()]
for force_idx, force in enumerate(iterate_forces(system)):
if force_idx in forces:
continue
for matched_force_class in matched_force_classes:
if isinstance(force, matched_force_class):
forces[force_idx] = force
# Reorder forces by index.
forces = collections.OrderedDict(sorted(forces.items()))
# Handle only_one.
if only_one is True:
if len(forces) == 0:
raise NoForceFoundError('No force of type {} could be found.'.format(force_type))
if len(forces) > 1:
raise MultipleForcesError('Found multiple forces of type {}'.format(force_type))
return forces.popitem(last=False)
return forces
def _compute_sphere_volume(radius):
"""Compute the volume of a square well restraint."""
return 4.0 / 3 * np.pi * radius**3
def _compute_harmonic_volume(radius, spring_constant, beta):
"""Compute the volume of an harmonic potential from 0 to radius.
Parameters
----------
radius : openmm.unit.Quantity
The upper limit on the distance (units of length).
spring_constant : openmm.unit.Quantity
The spring constant of the harmonic potential (units of
energy/mole/length^2).
beta : openmm.unit.Quantity
Thermodynamic beta (units of mole/energy).
Returns
-------
volume : openmm.unit.Quantity
The volume of the harmonic potential (units of length^3).
"""
# Turn everything to consistent dimension-less units.
length_unit = unit.nanometers
energy_unit = unit.kilojoules_per_mole
radius /= length_unit
beta *= energy_unit
spring_constant /= energy_unit/length_unit**2
bk = beta * spring_constant
bk_2 = bk / 2
bkr2_2 = bk_2 * radius**2
volume = math.sqrt(math.pi/2) * math.erf(math.sqrt(bkr2_2)) / bk**(3.0/2)
volume -= math.exp(-bkr2_2) * radius / bk
return 4 * math.pi * volume * length_unit**3
def _compute_harmonic_radius(spring_constant, potential_energy):
"""Find the radius at which the harmonic potential is energy.
Parameters
----------
spring_constant : openmm.unit.Quantity
The spring constant of the harmonic potential (units of
energy/mole/length^2).
potential_energy : openmm.unit.Quantity
The energy of the harmonic restraint (units of energy/mole).
Returns
-------
radius : openmm.unit.Quantity
The radius at which the harmonic potential is energy.
"""
length_unit = unit.nanometers
spring_constant *= length_unit**2
return math.sqrt(2 * potential_energy / spring_constant) * length_unit
# =============================================================================
# GENERIC CLASSES FOR RADIALLY SYMMETRIC RECEPTOR-LIGAND RESTRAINTS
# =============================================================================
[docs]
class RadiallySymmetricRestraintForce(utils.RestorableOpenMMObject):
"""Base class for radially-symmetric restraint force.
Provide facility functions to compute the standard state correction
of a receptor-ligand restraint.
To create a subclass, implement the properties :func:`restrained_atom_indices1`
and :func:`restrained_atom_indices2` (with their setters) that return
the indices of the restrained atoms.
You will also have to implement :func:`_create_bond`, which should add
the bond using the correct function/signature.
Optionally, you can implement :func:`distance_at_energy` if an
analytical expression for distance(potential_energy) exists.
If you subclass this, and plan on adding additional global parameters,
you need to invoke this class ``super().__init__`` first as the
``controlling_parameter_name`` must be the first global variable.
Parameters
----------
restraint_parameters : OrderedDict
An ordered dictionary containing the bond parameters in the form
parameter_name: parameter_value. The order is important to make
sure that parameters can be retrieved from the bond force with
the correct force index.
restrained_atom_indices1 : iterable of int
The indices of the first group of atoms to restrain.
restrained_atom_indices2 : iterable of int
The indices of the second group of atoms to restrain.
controlling_parameter_name : str
The name of the global parameter controlling the energy function.
*args, **kwargs
Parameters to pass to the super constructor.
Attributes
----------
controlling_parameter_name
"""
[docs]
def __init__(self, restraint_parameters, restrained_atom_indices1,
restrained_atom_indices2, controlling_parameter_name,
*args, **kwargs):
super(RadiallySymmetricRestraintForce, self).__init__(*args, **kwargs)
# Unzip bond parameters names and values from dict.
assert len(restraint_parameters) == 1 or isinstance(restraint_parameters, collections.OrderedDict)
parameter_names, parameter_values = zip(*restraint_parameters.items())
# Let the subclass initialize its bond.
self._create_bond(parameter_values, restrained_atom_indices1, restrained_atom_indices2)
# Add parameters. First global parameter is _restorable_force__class_hash
# from the RestorableOpenMMObject class.
err_msg = ('The force should have a single global parameter at this point. '
'This is likely because the subclass called addGlobalParameter '
'before calling super().__init__')
assert self.getNumGlobalParameters() == 1, err_msg
self.addGlobalParameter(controlling_parameter_name, 1.0)
for parameter in parameter_names:
self.addPerBondParameter(parameter)
# -------------------------------------------------------------------------
# Abstract methods.
# -------------------------------------------------------------------------
@abc.abstractmethod
def _create_bond(self, bond_parameter_values, restrained_atom_indices1,
restrained_atom_indices2):
"""Create the bond modelling the restraint.
Parameters
----------
bond_parameter_values : list of floats
The list of the parameter values of the bond.
restrained_atom_indices1 : list of int
The indices of the first group of atoms to restrain.
restrained_atom_indices2 : list of int
The indices of the second group of atoms to restrain.
"""
pass
# -------------------------------------------------------------------------
# Properties.
# -------------------------------------------------------------------------
@abc.abstractproperty
def restrained_atom_indices1(self):
"""list: The indices of the first group of restrained atoms."""
pass
@abc.abstractproperty
def restrained_atom_indices2(self):
"""list: The indices of the first group of restrained atoms."""
pass
@property
def restraint_parameters(self):
"""OrderedDict: The restraint parameters in dictionary form."""
parameter_values = self.getBondParameters(0)[-1]
restraint_parameters = [(self.getPerBondParameterName(parameter_idx), parameter_value)
for parameter_idx, parameter_value in enumerate(parameter_values)]
return collections.OrderedDict(restraint_parameters)
@property
def controlling_parameter_name(self):
"""str: The name of the global parameter controlling the energy function (read-only)."""
return self.getGlobalParameterName(1)
def distance_at_energy(self, potential_energy):
"""Compute the distance at which the potential energy is ``potential_energy``.
Parameters
----------
potential_energy : openmm.unit.Quantity
The potential energy of the restraint (units of energy/mole).
Returns
-------
distance : openmm.unit.Quantity
The distance at which the potential energy is ``potential_energy``
(units of length).
"""
raise NotImplementedError()
# -------------------------------------------------------------------------
# Methods to compute the standard state correction.
# -------------------------------------------------------------------------
def compute_standard_state_correction(self, thermodynamic_state, square_well=False,
radius_cutoff=None, energy_cutoff=None,
max_volume=None):
"""Return the standard state correction of the restraint.
The standard state correction is computed as
- log(V_standard / V_restraint)
where V_standard is the volume at standard state concentration and
V_restraint is the restraint volume. V_restraint is bounded by the
volume of the periodic box.
The ``square_well`` parameter, can be used to re-compute the standard
state correction when removing the bias introduced by the restraint.
Parameters
----------
thermodynamic_state : states.ThermodynamicState
The thermodynamic state at which to compute the standard state
correction.
square_well : bool, optional
If True, this computes the standard state correction assuming
the restraint to obey a square well potential. The energy
cutoff is still applied to the original energy potential.
radius_cutoff : openmm.unit.Quantity, optional
The maximum distance achievable by the restraint (units
compatible with nanometers). This is equivalent to placing
a hard wall potential at this distance.
energy_cutoff : float, optional
The maximum potential energy achievable by the restraint in kT.
This is equivalent to placing a hard wall potential at a
distance such that ``potential_energy(distance) == energy_cutoff``.
max_volume : openmm.unit.Quantity or 'system', optional
The volume of the periodic box (units compatible with nanometer**3).
This must be provided the thermodynamic state is in NPT. If the
string 'system' is passed, the maximum volume is computed from
the system box vectors (this has no effect if the system is not
periodic).
Returns
-------
correction : float
The unit-less standard state correction in kT at the given
thermodynamic state.
Raises
------
TypeError
If the thermodynamic state is in the NPT ensemble, and
``max_volume`` is not provided, or if the system is non-periodic
and no cutoff is given.
"""
# Determine restraint bound volume.
is_npt = thermodynamic_state.pressure is not None
if max_volume == 'system':
# ThermodynamicState.volume is None in the NPT ensemble.
# max_volume will still be None if the system is not periodic.
max_volume = thermodynamic_state.get_volume(ignore_ensemble=True)
elif max_volume is None and not is_npt:
max_volume = thermodynamic_state.volume
elif max_volume is None:
raise TypeError('max_volume must be provided with NPT ensemble')
# Non periodic systems reweighted to a square-well restraint must always have a cutoff.
if (not thermodynamic_state.is_periodic and square_well is True and
radius_cutoff is None and energy_cutoff is None and max_volume is None):
raise TypeError('One between radius_cutoff, energy_cutoff, or max_volume '
'must be provided when reweighting non-periodic thermodynamic '
'states to a square-well restraint.')
# If we evaluate the square well potential with no cutoffs,
# just use the volume of the periodic box.
if square_well is True and energy_cutoff is None and radius_cutoff is None:
restraint_volume = max_volume
# If we evaluate the square well potential with no energy cutoff,
# this can easily be solved analytically.
elif square_well is True and radius_cutoff is not None:
restraint_volume = _compute_sphere_volume(radius_cutoff)
# Use numerical integration.
else:
restraint_volume = self._compute_restraint_volume(
thermodynamic_state, square_well, radius_cutoff, energy_cutoff)
# Bound the restraint volume to the periodic box volume.
if max_volume is not None and restraint_volume > max_volume:
debug_msg = 'Limiting the restraint volume to {} nm^3 (original was {} nm^3)'
logger.debug(debug_msg.format(max_volume / unit.nanometers**3,
restraint_volume / unit.nanometers**3))
restraint_volume = max_volume
return -math.log(STANDARD_STATE_VOLUME / restraint_volume)
def _compute_restraint_volume(self, thermodynamic_state, square_well,
radius_cutoff, energy_cutoff):
"""Compute the volume of the restraint.
This function is called by ``compute_standard_state_correction()`` when
the standard state correction depends on the restraint potential.
Parameters
----------
thermodynamic_state : states.ThermodynamicState
The thermodynamic state at which to compute the standard state
correction.
square_well : bool, optional
If True, this computes the standard state correction assuming
the restraint to obey a square well potential. The energy
cutoff is still applied to the original energy potential.
radius_cutoff : openmm.unit.Quantity, optional
The maximum distance achievable by the restraint (units
compatible with nanometers). This is equivalent to placing
a hard wall potential at this distance.
energy_cutoff : float, optional
The maximum potential energy achievable by the restraint in kT.
This is equivalent to placing a hard wall potential at a
distance such that ``potential_energy(distance) == energy_cutoff``.
Returns
-------
restraint_volume : openmm.unit.Quantity
The volume of the restraint (units of length^3).
"""
# By default, use numerical integration.
return self._integrate_restraint_volume(thermodynamic_state, square_well,
radius_cutoff, energy_cutoff)
def _integrate_restraint_volume(self, thermodynamic_state, square_well,
radius_cutoff, energy_cutoff):
"""Compute the restraint volume through numerical integration.
Parameters
----------
thermodynamic_state : states.ThermodynamicState
The thermodynamic state at which to compute the standard state
correction.
square_well : bool, optional
If True, this computes the standard state correction assuming
the restraint to obey a square well potential. The energy
cutoff is still applied to the original energy potential.
radius_cutoff : openmm.unit.Quantity, optional
The maximum distance achievable by the restraint (units
compatible with nanometers). This is equivalent to placing
a hard wall potential at this distance.
energy_cutoff : float, optional
The maximum potential energy achievable by the restraint in kT.
This is equivalent to placing a hard wall potential at a
distance such that ``potential_energy(distance) == energy_cutoff``.
Returns
-------
restraint_volume : openmm.unit.Quantity
The volume of the restraint (units of length^3).
"""
distance_unit = unit.nanometer
# Create a System object containing two particles
# connected by the restraint force.
system = openmm.System()
system.addParticle(1.0 * unit.amu)
system.addParticle(1.0 * unit.amu)
force = copy.deepcopy(self)
force.restrained_atom_indices1 = [0]
force.restrained_atom_indices2 = [1]
# Disable the PBC for this approximation of the analytical solution.
force.setUsesPeriodicBoundaryConditions(False)
system.addForce(force)
# Create a Reference context to evaluate energies on the CPU.
integrator = openmm.VerletIntegrator(1.0 * unit.femtoseconds)
platform = openmm.Platform.getPlatformByName('Reference')
context = openmm.Context(system, integrator, platform)
# Set default positions.
positions = unit.Quantity(np.zeros([2,3]), distance_unit)
context.setPositions(positions)
# Create a function to compute integrand as a function of interparticle separation.
beta = thermodynamic_state.beta
def restraint_potential_func(r):
"""Return the potential energy in kT from the distance in nanometers."""
positions[1, 0] = r * distance_unit
context.setPositions(positions)
state = context.getState(getEnergy=True)
return beta * state.getPotentialEnergy()
def integrand(r):
"""
Parameters
----------
r : float
Inter-particle separation in nanometers
Returns
-------
dI : float
Contribution to the integral (in nm^2).
"""
potential = restraint_potential_func(r)
# If above the energy cutoff, this doesn't contribute to the integral.
if energy_cutoff is not None and potential > energy_cutoff:
return 0.0
# Check if we're reweighting to a square well potential.
if square_well:
potential = 0.0
dI = 4.0 * math.pi * r**2 * math.exp(-potential)
return dI
# Determine integration limits.
r_min, r_max, analytical_volume = self._determine_integral_limits(
thermodynamic_state, radius_cutoff, energy_cutoff, restraint_potential_func)
# Integrate restraint volume.
restraint_volume, restraint_volume_error = scipy.integrate.quad(
lambda r: integrand(r), r_min / distance_unit, r_max / distance_unit)
restraint_volume = restraint_volume * distance_unit**3 + analytical_volume
logger.debug("restraint_volume = {} nm^3".format(restraint_volume / distance_unit**3))
return restraint_volume
def _determine_integral_limits(self, thermodynamic_state, radius_cutoff,
energy_cutoff, potential_energy_func):
"""Determine integration limits for the standard state correction calculation.
This is called by ``_integrate_restraint_volume()`` to determine
the limits for numerical integration. This is important if we have
a cutoff as the points evaluated by scipy.integrate.quad are adaptively
chosen, and the hard wall can create numerical problems.
If part of the potential energy function can be computed analytically
you can reduce the integration interval and return a non-zero constant
to be added to the result of the integration.
Parameters
----------
thermodynamic_state : states.ThermodynamicState
The thermodynamic state at which to compute the standard state
correction.
square_well : bool, optional
If True, this computes the standard state correction assuming
the restraint to obey a square well potential. The energy
cutoff is still applied to the original energy potential.
radius_cutoff : openmm.unit.Quantity, optional
The maximum distance achievable by the restraint (units
compatible with nanometers). This is equivalent to placing
a hard wall potential at this distance.
energy_cutoff : float, optional
The maximum potential energy achievable by the restraint in kT.
This is equivalent to placing a hard wall potential at a
distance such that ``potential_energy(distance) == energy_cutoff``.
Returns
-------
r_min : openmm.unit.Quantity
The lower limit for numerical integration.
r_max : openmm.unit.Quantity
The upper limit for numerical integration.
analytical_volume : openmm.unit.Quantity
Volume excluded from the numerical integration that has been
computed analytically. This will be summed to the volume
computed through numerical integration.
"""
distance_unit = unit.nanometers
# The lower limit is always 0. Find the upper limit.
r_min = 0.0 * distance_unit
r_max = float('inf')
analytical_volume = 0.0 * distance_unit**3
if radius_cutoff is not None:
r_max = min(r_max, radius_cutoff / distance_unit)
if energy_cutoff is not None:
# First check if an analytical solution is available.
try:
energy_cutoff_distance = self.distance_at_energy(energy_cutoff*thermodynamic_state.kT)
except NotImplementedError:
# Find the first distance that exceeds the cutoff.
potential = 0.0
energy_cutoff_distance = 0.0 # In nanometers.
while potential <= energy_cutoff and energy_cutoff_distance < r_max:
energy_cutoff_distance += 0.1 # 1 Angstrom.
potential = potential_energy_func(energy_cutoff_distance)
r_max = min(r_max, energy_cutoff_distance)
# Handle the case where there are no distance or energy cutoff.
if r_max == float('inf'):
# For periodic systems, take thrice the maximum dimension of the system.
if thermodynamic_state.is_periodic:
box_vectors = thermodynamic_state.default_box_vectors
max_dimension = np.max(unit.Quantity(box_vectors) / distance_unit)
r_max = 3.0 * max_dimension
else:
r_max = 100.0 # distance_unit
r_max *= distance_unit
return r_min, r_max, analytical_volume
[docs]
class RadiallySymmetricCentroidRestraintForce(RadiallySymmetricRestraintForce,
openmm.CustomCentroidBondForce):
"""Base class for radially-symmetric restraints between the centroids of two groups of atoms.
The restraint is applied between the centers of mass of two groups
of atoms. The restraint strength is controlled by a global context
parameter whose name is passed on construction through the optional
argument ``controlling_parameter_name``.
With OpenCL, only on 64bit platforms are supported.
Parameters
----------
energy_function : str
The energy function to pass to ``CustomCentroidBondForce``. The
name of the controlling global parameter will be prepended to
this expression.
restraint_parameters : OrderedDict
An ordered dictionary containing the bond parameters in the form
parameter_name: parameter_value. The order is important to make
sure that parameters can be retrieved from the bond force with
the correct force index.
restrained_atom_indices1 : iterable of int
The indices of the first group of atoms to restrain.
restrained_atom_indices2 : iterable of int
The indices of the second group of atoms to restrain.
controlling_parameter_name : str, optional
The name of the global parameter controlling the energy function.
The default value is 'lambda_restraints'.
Attributes
----------
restraint_parameters
restrained_atom_indices1
restrained_atom_indices2
controlling_parameter_name
"""
[docs]
def __init__(self, energy_function, restraint_parameters,
restrained_atom_indices1, restrained_atom_indices2,
controlling_parameter_name='lambda_restraints'):
# Initialize CustomCentroidBondForce.
energy_function = controlling_parameter_name + ' * (' + energy_function + ')'
custom_centroid_bond_force_args = [2, energy_function]
super(RadiallySymmetricCentroidRestraintForce, self).__init__(
restraint_parameters, restrained_atom_indices1, restrained_atom_indices2,
controlling_parameter_name, *custom_centroid_bond_force_args)
@property
def restrained_atom_indices1(self):
"""The indices of the first group of restrained atoms."""
restrained_atom_indices1, weights_group1 = self.getGroupParameters(0)
return list(restrained_atom_indices1)
@restrained_atom_indices1.setter
def restrained_atom_indices1(self, atom_indices):
self.setGroupParameters(0, atom_indices)
@property
def restrained_atom_indices2(self):
"""The indices of the first group of restrained atoms."""
restrained_atom_indices2, weights_group2 = self.getGroupParameters(1)
return list(restrained_atom_indices2)
@restrained_atom_indices2.setter
def restrained_atom_indices2(self, atom_indices):
self.setGroupParameters(1, atom_indices)
def _create_bond(self, bond_parameter_values, restrained_atom_indices1,
restrained_atom_indices2):
"""Create the bond modelling the restraint."""
self.addGroup(restrained_atom_indices1)
self.addGroup(restrained_atom_indices2)
self.addBond([0, 1], bond_parameter_values)
[docs]
class RadiallySymmetricBondRestraintForce(RadiallySymmetricRestraintForce,
openmm.CustomBondForce):
"""Base class for radially-symmetric restraints between two atoms.
This is a version of ``RadiallySymmetricCentroidRestraintForce`` that can
be used with OpenCL 32-bit platforms. It supports atom groups with only a
single atom.
"""
[docs]
def __init__(self, energy_function, restraint_parameters,
restrained_atom_index1, restrained_atom_index2,
controlling_parameter_name='lambda_restraints'):
# Initialize CustomBondForce.
energy_function = energy_function.replace('distance(g1,g2)', 'r')
energy_function = controlling_parameter_name + ' * (' + energy_function + ')'
super(RadiallySymmetricBondRestraintForce, self).__init__(
restraint_parameters, [restrained_atom_index1], [restrained_atom_index2],
controlling_parameter_name, energy_function)
# -------------------------------------------------------------------------
# Public properties.
# -------------------------------------------------------------------------
@property
def restrained_atom_indices1(self):
"""The indices of the first group of restrained atoms."""
atom1, atom2, parameters = self.getBondParameters(0)
return [atom1]
@restrained_atom_indices1.setter
def restrained_atom_indices1(self, atom_indices):
assert len(atom_indices) == 1
atom1, atom2, parameters = self.getBondParameters(0)
self.setBondParameters(0, atom_indices[0], atom2, parameters)
@property
def restrained_atom_indices2(self):
"""The indices of the first group of restrained atoms."""
atom1, atom2, parameters = self.getBondParameters(0)
return [atom2]
@restrained_atom_indices2.setter
def restrained_atom_indices2(self, atom_indices):
assert len(atom_indices) == 1
atom1, atom2, parameters = self.getBondParameters(0)
self.setBondParameters(0, atom1, atom_indices[0], parameters)
def _create_bond(self, bond_parameter_values, restrained_atom_indices1, restrained_atom_indices2):
"""Create the bond modelling the restraint."""
self.addBond(restrained_atom_indices1[0], restrained_atom_indices2[0], bond_parameter_values)
# =============================================================================
# HARMONIC RESTRAINTS
# =============================================================================
[docs]
class HarmonicRestraintForceMixIn(object):
"""A mix-in providing the interface for harmonic restraints."""
[docs]
def __init__(self, spring_constant, *args, **kwargs):
energy_function = '(K/2)*distance(g1,g2)^2'
restraint_parameters = collections.OrderedDict([('K', spring_constant)])
super(HarmonicRestraintForceMixIn, self).__init__(energy_function, restraint_parameters,
*args, **kwargs)
@property
def spring_constant(self):
"""unit.openmm.Quantity: The spring constant K (units of energy/mole/distance^2)."""
# This works for both CustomBondForce and CustomCentroidBondForce.
parameters = self.getBondParameters(0)[-1]
return parameters[0] * unit.kilojoule_per_mole/unit.nanometers**2
def distance_at_energy(self, potential_energy):
"""Compute the distance at which the potential energy is ``potential_energy``.
Parameters
----------
potential_energy : openmm.unit.Quantity
The potential energy of the restraint (units of energy/mole).
Returns
-------
distance : openmm.unit.Quantity
The distance at which the potential energy is ``potential_energy``
(units of length).
"""
return _compute_harmonic_radius(self.spring_constant, potential_energy)
def _compute_restraint_volume(self, thermodynamic_state, square_well,
radius_cutoff, energy_cutoff):
"""Compute the restraint volume analytically."""
# If there is not a cutoff, integrate up to 100kT
if energy_cutoff is None:
energy_cutoff = 100.0 # kT
radius = self.distance_at_energy(energy_cutoff * thermodynamic_state.kT)
if radius_cutoff is not None:
radius = min(radius, radius_cutoff)
if square_well:
return _compute_sphere_volume(radius)
return _compute_harmonic_volume(radius, self.spring_constant,
thermodynamic_state.beta)
[docs]
class HarmonicRestraintForce(HarmonicRestraintForceMixIn,
RadiallySymmetricCentroidRestraintForce):
"""Impose a single harmonic restraint between the centroids of two groups of atoms.
This can be used to prevent the ligand from drifting too far from the
protein in implicit solvent calculations or to keep the ligand close
to the binding pocket in the decoupled states to increase mixing.
The restraint is applied between the centroids of two groups of atoms
that belong to the receptor and the ligand respectively. The centroids
are determined by a mass-weighted average of the group particles positions.
The energy expression of the restraint is given by
``E = controlling_parameter * (K/2)*r^2``
where `K` is the spring constant, `r` is the distance between the
two group centroids, and `controlling_parameter` is a scale factor that
can be used to control the strength of the restraint.
With OpenCL, only on 64bit platforms are supported.
Parameters
----------
spring_constant : openmm.unit.Quantity
The spring constant K (see energy expression above) in units
compatible with joule/nanometer**2/mole.
restrained_atom_indices1 : iterable of int
The indices of the first group of atoms to restrain.
restrained_atom_indices2 : iterable of int
The indices of the second group of atoms to restrain.
controlling_parameter_name : str, optional
The name of the global parameter controlling the energy function.
The default value is 'lambda_restraints'.
Attributes
----------
spring_constant
restrained_atom_indices1
restrained_atom_indices2
restraint_parameters
controlling_parameter_name
"""
# All the methods are provided by the mix-ins.
pass
[docs]
class HarmonicRestraintBondForce(HarmonicRestraintForceMixIn,
RadiallySymmetricBondRestraintForce):
"""Impose a single harmonic restraint between two atoms.
This is a version of ``HarmonicRestraintForce`` that can be used with
OpenCL 32-bit platforms. It supports atom groups with only a single atom.
Parameters
----------
spring_constant : openmm.unit.Quantity
The spring constant K (see energy expression above) in units
compatible with joule/nanometer**2/mole.
restrained_atom_index1 : int
The index of the first atom to restrain.
restrained_atom_index2 : int
The index of the second atom to restrain.
controlling_parameter_name : str, optional
The name of the global parameter controlling the energy function.
The default value is 'lambda_restraints'.
Attributes
----------
spring_constant
restrained_atom_indices1
restrained_atom_indices2
restraint_parameters
controlling_parameter_name
"""
# All the methods are provided by the mix-ins.
pass
# =============================================================================
# FLAT-BOTTOM RESTRAINTS
# =============================================================================
[docs]
class FlatBottomRestraintForceMixIn(object):
"""A mix-in providing the interface for flat-bottom restraints."""
[docs]
def __init__(self, spring_constant, well_radius, *args, **kwargs):
energy_function = 'step(distance(g1,g2)-r0) * (K/2)*(distance(g1,g2)-r0)^2'
restraint_parameters = collections.OrderedDict([
('K', spring_constant),
('r0', well_radius)
])
super(FlatBottomRestraintForceMixIn, self).__init__(energy_function, restraint_parameters,
*args, **kwargs)
@property
def spring_constant(self):
"""unit.openmm.Quantity: The spring constant K (units of energy/mole/length^2)."""
# This works for both CustomBondForce and CustomCentroidBondForce.
parameters = self.getBondParameters(0)[-1]
return parameters[0] * unit.kilojoule_per_mole/unit.nanometers**2
@property
def well_radius(self):
"""unit.openmm.Quantity: The distance at which the harmonic restraint is imposed (units of length)."""
# This works for both CustomBondForce and CustomCentroidBondForce.
parameters = self.getBondParameters(0)[-1]
return parameters[1] * unit.nanometers
def distance_at_energy(self, potential_energy):
"""Compute the distance at which the potential energy is ``potential_energy``.
Parameters
----------
potential_energy : openmm.unit.Quantity
The potential energy of the restraint (units of energy/mole).
Returns
-------
distance : openmm.unit.Quantity
The distance at which the potential energy is ``potential_energy``
(units of length).
"""
if potential_energy == 0.0*unit.kilojoules_per_mole:
raise ValueError('Cannot compute the distance at this potential energy.')
harmonic_radius = _compute_harmonic_radius(self.spring_constant, potential_energy)
return self.well_radius + harmonic_radius
def _compute_restraint_volume(self, thermodynamic_state, square_well,
radius_cutoff, energy_cutoff):
"""Compute the restraint volume analytically."""
# Check if we are using square well and we can avoid numerical integration.
if square_well:
_, r_max, _ = self._determine_integral_limits(
thermodynamic_state, radius_cutoff, energy_cutoff)
return _compute_sphere_volume(r_max)
return self._integrate_restraint_volume(thermodynamic_state, square_well,
radius_cutoff, energy_cutoff)
def _determine_integral_limits(self, thermodynamic_state, radius_cutoff,
energy_cutoff, potential_energy_func=None):
# If there is not a cutoff, integrate up to 100kT.
if energy_cutoff is None:
energy_cutoff = 100.0 # kT
energy_cutoff = energy_cutoff * thermodynamic_state.kT
r_max = _compute_harmonic_radius(self.spring_constant, energy_cutoff)
r_max += self.well_radius
if radius_cutoff is not None:
r_max = min(r_max, radius_cutoff)
# Compute the volume from the flat-bottom part of the potential.
r_min = min(r_max, self.well_radius)
analytical_volume = _compute_sphere_volume(r_min)
return r_min, r_max, analytical_volume
[docs]
class FlatBottomRestraintForce(FlatBottomRestraintForceMixIn,
RadiallySymmetricCentroidRestraintForce):
"""A restraint between the centroids of two groups of atoms using a flat potential well with harmonic walls.
An alternative choice to receptor-ligand restraints that uses a flat
potential inside most of the protein volume with harmonic restraining
walls outside of this. It can be used to prevent the ligand from
drifting too far from protein in implicit solvent calculations while
still exploring the surface of the protein for putative binding sites.
The restraint is applied between the centroids of two groups of atoms
that belong to the receptor and the ligand respectively. The centroids
are determined by a mass-weighted average of the group particles positions.
More precisely, the energy expression of the restraint is given by
``E = controlling_parameter * step(r-r0) * (K/2)*(r-r0)^2``
where ``K`` is the spring constant, ``r`` is the distance between the
restrained atoms, ``r0`` is another parameter defining the distance
at which the restraint is imposed, and ``controlling_parameter``
is a scale factor that can be used to control the strength of the
restraint.
With OpenCL, only on 64bit platforms are supported.
Parameters
----------
spring_constant : openmm.unit.Quantity
The spring constant K (see energy expression above) in units
compatible with joule/nanometer**2/mole.
well_radius : openmm.unit.Quantity
The distance r0 (see energy expression above) at which the harmonic
restraint is imposed in units of distance.
restrained_atom_indices1 : iterable of int
The indices of the first group of atoms to restrain.
restrained_atom_indices2 : iterable of int
The indices of the second group of atoms to restrain.
controlling_parameter_name : str, optional
The name of the global parameter controlling the energy function.
The default value is 'lambda_restraints'.
Attributes
----------
spring_constant
well_radius
restrained_atom_indices1
restrained_atom_indices2
restraint_parameters
controlling_parameter_name
"""
# All the methods are provided by the mix-ins.
pass
[docs]
class FlatBottomRestraintBondForce(FlatBottomRestraintForceMixIn,
RadiallySymmetricBondRestraintForce):
"""A restraint between two atoms using a flat potential well with harmonic walls.
This is a version of ``FlatBottomRestraintForce`` that can be used with
OpenCL 32-bit platforms. It supports atom groups with only a single atom.
Parameters
----------
spring_constant : openmm.unit.Quantity
The spring constant K (see energy expression above) in units
compatible with joule/nanometer**2/mole.
well_radius : openmm.unit.Quantity
The distance r0 (see energy expression above) at which the harmonic
restraint is imposed in units of distance.
restrained_atom_index1 : int
The index of the first group of atoms to restrain.
restrained_atom_index2 : int
The index of the second group of atoms to restrain.
controlling_parameter_name : str, optional
The name of the global parameter controlling the energy function.
The default value is 'lambda_restraints'.
Attributes
----------
spring_constant
well_radius
restrained_atom_indices1
restrained_atom_indices2
restraint_parameters
controlling_parameter_name
"""
# All the methods are provided by the mix-ins.
pass
# =============================================================================
# REACTION FIELD
# =============================================================================
[docs]
class UnshiftedReactionFieldForce(openmm.CustomNonbondedForce):
"""A force modelling switched reaction-field electrostatics.
Contrarily to a normal `NonbondedForce` with `CutoffPeriodic` nonbonded
method, this force sets the `c_rf` to 0.0 and uses a switching function
to avoid forces discontinuities at the cutoff distance.
Parameters
----------
cutoff_distance : openmm.unit.Quantity, default 15*angstroms
The cutoff distance (units of distance).
switch_width : openmm.unit.Quantity, default 1.0*angstrom
Switch width for electrostatics (units of distance).
reaction_field_dielectric : float
The dielectric constant used for the solvent.
"""
[docs]
def __init__(self, cutoff_distance=15*unit.angstroms, switch_width=1.0*unit.angstrom,
reaction_field_dielectric=78.3):
k_rf = cutoff_distance**(-3) * (reaction_field_dielectric - 1.0) / (2.0*reaction_field_dielectric + 1.0)
# Energy expression omits c_rf constant term.
energy_expression = "ONE_4PI_EPS0*chargeprod*(r^(-1) + k_rf*r^2);"
energy_expression += "chargeprod = charge1*charge2;"
energy_expression += "k_rf = {:f};".format(k_rf.value_in_unit_system(unit.md_unit_system))
energy_expression += "ONE_4PI_EPS0 = {:f};".format(ONE_4PI_EPS0) # already in OpenMM units
# Create CustomNonbondedForce.
super(UnshiftedReactionFieldForce, self).__init__(energy_expression)
# Add parameters.
self.addPerParticleParameter("charge")
# Configure force.
self.setNonbondedMethod(openmm.CustomNonbondedForce.CutoffPeriodic)
self.setCutoffDistance(cutoff_distance)
self.setUseLongRangeCorrection(False)
if switch_width is not None:
self.setUseSwitchingFunction(True)
self.setSwitchingDistance(cutoff_distance - switch_width)
else: # Truncated
self.setUseSwitchingFunction(False)
@classmethod
def from_nonbonded_force(cls, nonbonded_force, switch_width=1.0*unit.angstrom):
"""Copy constructor from an OpenMM `NonbondedForce`.
The returned force has same cutoff distance and dielectric, and
its particles have the same charges. Exclusions corresponding to
`nonbonded_force` exceptions are also added.
.. warning
This only creates the force object. The electrostatics in
`nonbonded_force` remains unmodified. Use the function
`replace_reaction_field` to correctly convert a system to
use an unshifted reaction field potential.
Parameters
----------
nonbonded_force : openmm.NonbondedForce
The nonbonded force to copy.
switch_width : openmm.unit.Quantity
Switch width for electrostatics (units of distance).
Returns
-------
reaction_field_force : UnshiftedReactionFieldForce
The reaction field force with copied particles.
"""
# OpenMM gives unitless values.
cutoff_distance = nonbonded_force.getCutoffDistance()
reaction_field_dielectric = nonbonded_force.getReactionFieldDielectric()
reaction_field_force = cls(cutoff_distance, switch_width, reaction_field_dielectric)
# Set particle charges.
for particle_index in range(nonbonded_force.getNumParticles()):
charge, sigma, epsilon = nonbonded_force.getParticleParameters(particle_index)
reaction_field_force.addParticle([charge])
# Add exclusions to CustomNonbondedForce.
for exception_index in range(nonbonded_force.getNumExceptions()):
iatom, jatom, chargeprod, sigma, epsilon = nonbonded_force.getExceptionParameters(exception_index)
reaction_field_force.addExclusion(iatom, jatom)
return reaction_field_force
@classmethod
def from_system(cls, system, switch_width=1.0*unit.angstrom):
"""Copy constructor from the first OpenMM `NonbondedForce` in `system`.
If multiple `NonbondedForce`s are found, an exception is raised.
.. warning
This only creates the force object. The electrostatics in
`nonbonded_force` remains unmodified. Use the function
`replace_reaction_field` to correctly convert a system to
use an unshifted reaction field potential.
Parameters
----------
system : openmm.System
The system containing the nonbonded force to copy.
switch_width : openmm.unit.Quantity
Switch width for electrostatics (units of distance).
Returns
-------
reaction_field_force : UnshiftedReactionFieldForce
The reaction field force.
See Also
--------
UnshiftedReactionField.from_nonbonded_force
"""
force_idx, nonbonded_force = find_forces(system, openmm.NonbondedForce, only_one=True)
return cls.from_nonbonded_force(nonbonded_force, switch_width)
class SwitchedReactionFieldForce(openmm.CustomNonbondedForce):
"""A force modelling switched reaction-field electrostatics.
Parameters
----------
cutoff_distance : openmm.unit.Quantity, default 15*angstroms
The cutoff distance (units of distance).
switch_width : openmm.unit.Quantity, default 1.0*angstrom
Switch width for electrostatics (units of distance).
reaction_field_dielectric : float
The dielectric constant used for the solvent.
"""
def __init__(self, cutoff_distance=15*unit.angstroms, switch_width=1.0*unit.angstrom,
reaction_field_dielectric=78.3):
k_rf = cutoff_distance**(-3) * (reaction_field_dielectric - 1.0) / (2.0*reaction_field_dielectric + 1.0)
c_rf = cutoff_distance**(-1) * (3*reaction_field_dielectric) / (2.0*reaction_field_dielectric + 1.0)
# Energy expression omits c_rf constant term.
energy_expression = "ONE_4PI_EPS0*chargeprod*(r^(-1) + k_rf*r^2 - c_rf);"
energy_expression += "chargeprod = charge1*charge2;"
energy_expression += "k_rf = {:f};".format(k_rf.value_in_unit_system(unit.md_unit_system))
energy_expression += "c_rf = {:f};".format(c_rf.value_in_unit_system(unit.md_unit_system))
energy_expression += "ONE_4PI_EPS0 = {:f};".format(ONE_4PI_EPS0) # already in OpenMM units
# Create CustomNonbondedForce.
super(SwitchedReactionFieldForce, self).__init__(energy_expression)
# Add parameters.
self.addPerParticleParameter("charge")
# Configure force.
self.setNonbondedMethod(openmm.CustomNonbondedForce.CutoffPeriodic)
self.setCutoffDistance(cutoff_distance)
self.setUseLongRangeCorrection(False)
if switch_width is not None:
self.setUseSwitchingFunction(True)
self.setSwitchingDistance(cutoff_distance - switch_width)
else: # Truncated
self.setUseSwitchingFunction(False)
@classmethod
def from_nonbonded_force(cls, nonbonded_force, switch_width=1.0*unit.angstrom):
"""Copy constructor from an OpenMM `NonbondedForce`.
The returned force has same cutoff distance and dielectric, and
its particles have the same charges. Exclusions corresponding to
`nonbonded_force` exceptions are also added.
.. warning
This only creates the force object. The electrostatics in
`nonbonded_force` remains unmodified. Use the function
`replace_reaction_field` to correctly convert a system to
use an unshifted reaction field potential.
Parameters
----------
nonbonded_force : openmm.NonbondedForce
The nonbonded force to copy.
switch_width : openmm.unit.Quantity
Switch width for electrostatics (units of distance).
Returns
-------
reaction_field_force : UnshiftedReactionFieldForce
The reaction field force with copied particles.
"""
# OpenMM gives unitless values.
cutoff_distance = nonbonded_force.getCutoffDistance()
reaction_field_dielectric = nonbonded_force.getReactionFieldDielectric()
reaction_field_force = cls(cutoff_distance, switch_width, reaction_field_dielectric)
# Set particle charges.
for particle_index in range(nonbonded_force.getNumParticles()):
charge, sigma, epsilon = nonbonded_force.getParticleParameters(particle_index)
reaction_field_force.addParticle([charge])
# Add exclusions to CustomNonbondedForce.
for exception_index in range(nonbonded_force.getNumExceptions()):
iatom, jatom, chargeprod, sigma, epsilon = nonbonded_force.getExceptionParameters(exception_index)
reaction_field_force.addExclusion(iatom, jatom)
return reaction_field_force
@classmethod
def from_system(cls, system, switch_width=1.0*unit.angstrom):
"""Copy constructor from the first OpenMM `NonbondedForce` in `system`.
If multiple `NonbondedForce`s are found, an exception is raised.
.. warning
This only creates the force object. The electrostatics in
`nonbonded_force` remains unmodified. Use the function
`replace_reaction_field` to correctly convert a system to
use an unshifted reaction field potential.
Parameters
----------
system : openmm.System
The system containing the nonbonded force to copy.
switch_width : openmm.unit.Quantity
Switch width for electrostatics (units of distance).
Returns
-------
reaction_field_force : UnshiftedReactionFieldForce
The reaction field force.
See Also
--------
UnshiftedReactionField.from_nonbonded_force
"""
force_idx, nonbonded_force = find_forces(system, openmm.NonbondedForce, only_one=True)
return cls.from_nonbonded_force(nonbonded_force, switch_width)
if __name__ == '__main__':
import doctest
doctest.testmod()