"""
OpenMM system extensions and tools
==================================
.. moduleauthor:: Benjamin Ye <GitHub: @bbye98>
This module contains implementations of common OpenMM system
transformations, like the slab correction method for pseudo-2D slab
systems, the method of image charges, and an applied potential
difference.
"""
import logging
from typing import Any, Union
import warnings
import mpmath
import numpy as np
import openmm
from openmm import app, unit
from scipy import special
from .unit import VACUUM_PERMITTIVITY
try:
from openmm_ic import ICLangevinIntegrator
FOUND_ICPLUGIN = True
except ImportError:
try:
from constvplugin import ConstVLangevinIntegrator as ICLangevinIntegrator
FOUND_ICPLUGIN = True
except ImportError:
FOUND_ICPLUGIN = False
[docs]
def register_particles(
system: openmm.System, topology: app.Topology,
N: int = 0, mass: Union[float, unit.Quantity] = 0.0, *,
chain: app.Chain = None, element: app.Element = None, name: str = "",
resname: str = "", nbforce: openmm.NonbondedForce = None,
charge: Union[float, unit.Quantity] = 0.0,
sigma: Union[float, unit.Quantity] = 0.0,
epsilon: Union[float, unit.Quantity] = 0.0,
cnbforces: dict[openmm.CustomNonbondedForce, tuple[Any]] = None
) -> None:
"""
Sequentially register particles of the same type to the simulation
system, topology, and nonbonded pair potentials.
Parameters
----------
system : `openmm.System`
OpenMM molecular system. Can be :code:`None` if particles are to
be registered to the topology only.
topology : `openmm.app.Topology`
Topological information about an OpenMM system.
N : `int`, default: :code:`0`
Number of atom(s). If not specified, no particles are added.
mass : `float` or `openmm.unit.Quantity`, default: :code:`0`
Molar mass. If not specified, particles are massless and will
not move.
**Reference unit**: :math:`\\mathrm{g/mol}`.
chain : `openmm.app.Chain`, keyword-only, optional
Chain that the atom(s) should be added to. If not provided,
a new chain is created for each atom.
element : `openmm.app.Element`, keyword-only, optional
Chemical element of the atom(s) to add.
name : `str`, keyword-only, optional
Name of the atom(s) to add.
resname : `str`, keyword-only, optional
Name of the residue(s) to add. If not specified, `name` is used
if available.
nbforce : `openmm.NonbondedForce`, keyword-only, optional
Standard OpenMM pair potential object implementing the nonbonded
Lennard-Jones and Coulomb potentials.
charge : `float` or `openmm.unit.Quantity`, keyword-only, \
default: :code`0`
Charge :math:`q` of the atom(s) for use in the Coulomb potential
in `nbforce`.
**Reference unit**: :math:`\\mathrm{e}`.
sigma : `float` or `openmm.unit.Quantity`, keyword-only, \
default: :code`0`
:math:`\\sigma` parameter of the Lennard-Jones potential in
`nbforce`.
**Reference unit**: :math:`\\mathrm{nm}`.
epsilon : `float` or `openmm.unit.Quantity`, keyword-only, \
default: :code`0`
:math:`\\epsilon` parameter of the Lennard-Jones potential in
`nbforce`.
**Reference unit**: :math:`\\mathrm{kJ/mol}`.
cnbforces : `dict`, keyword-only, optional
Custom pair potential objects implementing other non-standard
pair potentials and their corresponding per-particle parameters.
**Example**: :code:`{gauss: (0.3 * unit.nanometer,)}`, where
`gauss` is a custom Gaussian potential obtained using
:func:`mdhelper.openmm.pair.gauss()`.
"""
has_nbforce = nbforce is not None
has_system = system is not None
per_chain = chain is None
cnbforces = cnbforces or {}
for _ in range(N):
if has_system:
system.addParticle(mass)
if per_chain:
chain = topology.addChain()
residue = topology.addResidue(resname or name, chain)
topology.addAtom(name, element, residue)
if has_nbforce:
nbforce.addParticle(charge, sigma, epsilon)
for cnbforce, param in cnbforces.items():
cnbforce.addParticle(param)
[docs]
def add_slab_correction(
system: openmm.System, topology: app.Topology,
nbforce: Union[openmm.NonbondedForce, openmm.CustomNonbondedForce],
temp: Union[float, unit.Quantity], fric: Union[float, unit.Quantity],
dt: Union[float, unit.Quantity], axis: int = 2, *,
charge_index: int = 0, z_scale: float = 3, method: str = "force"
) -> openmm.Integrator:
r"""
Implements a slab correction so that efficient three-dimensional
Ewald methods can continue to be used to evaluate the electrostatics
for systems that are periodic in the :math:`x`- and
:math:`y`-directions but not the :math:`z`-direction. Effectively,
the system is treated as if it were periodic in the
:math:`z`-direction, but with empty volume added between the slabs
and the slab–slab interactions removed.
For electroneutral systems, the Yeh–Berkowitz correction [1]_ is
applied:
.. math::
\begin{gather*}
U^\mathrm{corr}=\frac{N_A}{2\varepsilon_0V}M_z^2\\
u_i^\mathrm{corr}=\frac{N_A}{2\varepsilon_0V}q_i\left(z_iM_z
-\frac{\sum_i q_iz_i^2}{2}\right)\\
f_{i,z}^\mathrm{corr}=-\frac{N_A}{\varepsilon_0V}q_iM_z
\end{gather*}
For systems with a net electric charge, the Ballenegger–Arnold–Cerdà
correction [2]_ is applied instead:
.. math::
\begin{gather*}
U^\mathrm{corr}=\frac{N_A}{2\varepsilon_0V}
\left(M_z^2-q_\mathrm{tot}\sum_i q_iz_i^2
-\frac{q_\mathrm{tot}^2L_z^2}{12}\right)\\
u_i^\mathrm{corr}=\frac{N_A}{2\varepsilon_0V}q_i
\left(z_iM_z-\frac{\sum_i q_iz_i^2+q_\mathrm{tot}z_i^2}{2}
-\frac{q_\mathrm{tot}L_z^2}{12}\right)\\
f_{i,z}^\mathrm{corr}=-\frac{N_A}{\varepsilon_0V}q_i
\left(M_z-q_\mathrm{tot}z_i\right)
\end{gather*}
Note that the the relative permittivity
:math:`\varepsilon_\mathrm{r}` does not appear in the equations
above because it is accounted for by scaling the particle charges in
the Coulomb potential.
Parameters
----------
system : `openmm.System`
OpenMM molecular system.
topology : `openmm.app.Topology`
Topological information about an OpenMM system.
nbforce : `openmm.NonbondedForce` or `openmm.CustomNonbondedForce`
Pair potential object containing particle charge information.
.. note::
It is assumed that the charge :math:`q` information is the
first per-particle parameter stored in `nbforce`. If not, the
index can be specified in `charge_index`.
temp : `float` or `openmm.unit.Quantity`
System temperature :math:`T`.
**Reference unit**: :math:`\mathrm{K}`.
fric : `float` or `openmm.unit.Quantity`
Friction coefficient :math:`\gamma` that couples the system to
the heat bath.
**Reference unit**: :math:`\mathrm{ps}^{-1}`.
dt : `float` or `openmm.unit.Quantity`
Integration step size :math:`\Delta t`.
**Reference unit**: :math:`\mathrm{ps}`.
axis : `int`, default: :code:`2`
Axis along which to apply the slab correction, with :math:`x`
being :code:`0`, :math:`y` being :code:`1`, and :math:`z` being
:code:`2`. The source code and outputs, if any, will refer to
this dimension as :code:`z` regardless of this value.
charge_index : `int`, keyword-only, default: :code:`0`
Index of charge :math:`q` information in the per-particle
parameters stored in `nbforce`.
z_scale : `float`, keyword-only, default: :code:`3`
Scaling factor for the dimension specified in `axis`.
method : `str`, keyword-only, default: :code:`"force"`
Slab correction methodology.
.. container::
**Valid values**:
* :code:`"force"`: Collective implementation via
:class:`openmm.CustomExternalForce` and
:class:`openmm.CustomCVForce` [3]_. This is
generally the most efficient.
* :code:`"integrator"`: Per-particle implementation via an
:class:`openmm.CustomIntegrator`.
Returns
-------
integrator : `openmm.Integrator` or `openmm.CustomIntegrator`
Integrator that simulates a system using Langevin dynamics, with
the LFMiddle discretization.
References
----------
.. [1] Yeh, I.-C.; Berkowitz, M. L. Ewald Summation for Systems with
Slab Geometry. *J. Chem. Phys.* **1999**, *111* (7), 3155–3162.
https://doi.org/10.1063/1.479595.
.. [2] Ballenegger, V.; Arnold, A.; Cerdà, J. J. Simulations of
Non-Neutral Slab Systems with Long-Range Electrostatic
Interactions in Two-Dimensional Periodic Boundary Conditions.
*J. Chem. Phys.* **2009**, *131* (9), 094107.
https://doi.org/10.1063/1.3216473.
.. [3] Son, C. Y.; Wang, Z.-G. Image-Charge Effects on Ion
Adsorption near Aqueous Interfaces. *Proc. Natl. Acad. Sci.
U.S.A.* **2021**, *118* (19), e2020615118.
https://doi.org/10.1073/pnas.2020615118.
"""
# Get system dimensions
dims = np.array(
topology.getUnitCellDimensions().value_in_unit(unit.nanometer)
) * unit.nanometer
pbv = system.getDefaultPeriodicBoxVectors()
# Scale system z-dimension by specified z-scale
if z_scale < 2:
wmsg = ("A z-scaling factor that is less than 2 may introduce "
"unwanted slab–slab interactions. The recommended "
"value is 3.")
warnings.warn(wmsg)
elif z_scale > 5:
wmsg = ("A z-scaling factor that is greater than 5 may "
"penalize performance. The recommended value is 3.")
warnings.warn(wmsg)
dims[axis] *= z_scale
pbv[axis] *= z_scale
# Set new system dimensions
topology.setUnitCellDimensions(dims)
system.setDefaultPeriodicBoxVectors(*pbv)
# Obtain particle charge information
if isinstance(nbforce.getParticleParameters(0)[charge_index], unit.Quantity):
qs = np.fromiter(
(nbforce.getParticleParameters(i)[charge_index]
.value_in_unit(unit.elementary_charge)
for i in range(nbforce.getNumParticles())),
dtype=float
)
else:
qs = np.fromiter(
(nbforce.getParticleParameters(i)[charge_index]
for i in range(nbforce.getNumParticles())),
dtype=float
)
neutral = qs.min() == qs.max()
if not neutral:
q_tot = qs.sum()
electroneutral = np.isclose(q_tot, 0)
# Calculate coefficient for slab correction
coef = unit.AVOGADRO_CONSTANT_NA / \
(2 * VACUUM_PERMITTIVITY * dims[0] * dims[1] * dims[2])
# Get letter representation of axis for formula
z = chr(120 + axis)
if neutral:
# Instantiate an integrator that simulates a system using
# Langevin dynamics, with the LFMiddle discretization
integrator = openmm.LangevinMiddleIntegrator(temp, fric, dt)
else:
if method == "integrator":
# Implement an integrator that simulates a system using Langevin
# dynamics, with the LFMiddle discretization
integrator = openmm.CustomIntegrator(dt)
integrator.addGlobalVariable("a", np.exp(-fric * dt))
integrator.addGlobalVariable("b", np.sqrt(1 - np.exp(-2 * fric * dt)))
integrator.addGlobalVariable(
"kT", unit.AVOGADRO_CONSTANT_NA * unit.BOLTZMANN_CONSTANT_kB * temp
)
integrator.addPerDofVariable("x1", 0)
integrator.addUpdateContextState()
integrator.addComputePerDof("v", "v+dt*f/m")
integrator.addConstrainVelocities()
integrator.addComputePerDof("x", "x+dt*v/2")
integrator.addComputePerDof("v", "a*v+b*sqrt(kT/m)*gaussian")
integrator.addComputePerDof("x", "x+dt*v/2")
integrator.addComputePerDof("x1", "x")
integrator.addConstrainPositions()
integrator.addComputePerDof("v", "v+(x-x1)/dt")
# Initialize per-degree-of-freedom variable q for charge
integrator.addPerDofVariable("q", 0)
# Add global dipole moment computation to integrator
integrator.addComputeSum("M_z", "q*x")
integrator.addComputeSum("M_zz", "q*x^2")
# Give particle charge information to integrator
q_vectors = np.zeros((len(qs), 3))
q_vectors[:, axis] = qs
integrator.setPerDofVariableByName("q", q_vectors)
# Implement per-particle slab correction
if electroneutral:
slab_corr = openmm.CustomExternalForce(
f"coef*q*({z}*M_z-M_zz/2)"
)
else:
slab_corr = openmm.CustomExternalForce(
f"coef*q*({z}*M_z-(M_zz+q_tot*{z}^2)/2-q_tot*dim_z^2/12)"
)
slab_corr.addGlobalParameter("dim_z", dims[axis])
slab_corr.addGlobalParameter("q_tot", q_tot)
slab_corr.addGlobalParameter("M_z", 0)
slab_corr.addGlobalParameter("M_zz", 0)
slab_corr.addGlobalParameter("coef", coef)
slab_corr.addPerParticleParameter("q")
# Register real particles to the slab correction
for i, q in enumerate(qs):
slab_corr.addParticle(i, (q,))
elif method == "force":
# Instantiate an integrator that simulates a system using
# Langevin dynamics, with the LFMiddle discretization
integrator = openmm.LangevinMiddleIntegrator(temp, fric, dt)
# Calculate instantaneous system dipole
M_z = openmm.CustomExternalForce(f"q*{z}")
M_z.addPerParticleParameter("q")
# Implement collective slab correction
if electroneutral:
slab_corr = openmm.CustomCVForce("coef*M_z^2")
else:
M_zz = openmm.CustomExternalForce(f"q*{z}^2")
M_zz.addPerParticleParameter("q")
slab_corr = openmm.CustomCVForce(
"coef*(M_z^2-q_tot*M_zz-q_tot^2*dim_z^2/12)"
)
slab_corr.addCollectiveVariable("M_zz", M_zz)
slab_corr.addGlobalParameter("dim_z", dims[axis])
slab_corr.addGlobalParameter("q_tot", q_tot)
slab_corr.addCollectiveVariable("M_z", M_z)
slab_corr.addGlobalParameter("coef", coef)
# Register real particles to the slab correction
for i, q in enumerate(qs):
M_z.addParticle(i, (q,))
if not electroneutral:
M_zz.addParticle(i, (q,))
# Register slab correction to the system
system.addForce(slab_corr)
return integrator
[docs]
def add_image_charges(
system: openmm.System, topology: app.Topology,
positions: Union[np.ndarray[float], unit.Quantity],
temp: Union[float, unit.Quantity], fric: Union[float, unit.Quantity],
dt: Union[float, unit.Quantity], *, gamma: float = -1,
n_cells: int = 2, nbforce: openmm.NonbondedForce = None,
cnbforces: dict[openmm.CustomNonbondedForce, dict[str, Any]] = None,
wall_indices: np.ndarray[int] = None, exclude: bool = False
) -> tuple[unit.Quantity, openmm.Integrator]:
r"""
Implements the method of image charges for dielectric boundaries.
For more information about the method, see Refs. [1]_, [2]_, and
[3]_ for perfectly conducting boundaries, and Refs. [4]_, [5]_, and
[6]_ otherwise.
.. note::
The boundaries must be in the :math:`xy`-plane and along the
:math:`z`-axis.
Parameters
----------
system : `openmm.System`
OpenMM molecular system.
topology : `openmm.app.Topology`
Topological information about an OpenMM system.
positions : `numpy.ndarray`
Positions of the :math:`N` particles in the real system.
**Shape**: :math:`(N,\,3)`.
**Reference unit**: :math:`\mathrm{nm}`.
temp : `float` or `openmm.unit.Quantity`
System temperature :math:`T`.
**Reference unit**: :math:`\mathrm{K}`.
fric : `float` or `openmm.unit.Quantity`
Friction coefficient :math:`\gamma` that couples the system to
the heat bath.
**Reference unit**: :math:`\mathrm{ps}^{-1}`.
dt : `float` or `openmm.unit.Quantity`
Integration step size :math:`\Delta t`.
**Reference unit**: :math:`\mathrm{ps}`.
gamma : `float`, keyword-only, default: :code:`-1`
Scaled image charge magnitude :math:`\gamma`.
n_cells : `int`, keyword-only, default: :code:`2`
Total number :math:`n_\mathrm{cell}` of cells, with :math:`1`
real and :math:`n_\mathrm{cell}-1` image cells.
.. tip::
Use :math:`n_\mathrm{cell}>2` to instantiate a
highly-confined simulation system with periodic dimensions
large enough to satisfy the constraint that the force field
cutoff be less than half the box size.
**Valid values**: Even integers greater than :math:`1` when
:math:`\gamma=-1`.
nbforce : `openmm.NonbondedForce`, keyword-only, optional
Standard OpenMM pair potential object implementing the nonbonded
Lennard-Jones and Coulomb potentials. For the image charges, the
charges :math:`q` have their signs flipped, and the LJ
parameters :math:`\sigma` and :math:`\epsilon` are both set to
:math:`0`.
cnbforces : `dict`, keyword-only, optional
Custom pair potential objects implementing other non-standard
pair potentials. The keys are the potentials and the values are
the corresponding per-particle parameters to use, which can be
provided as follows:
.. container::
* `None`: All per-particle parameters set to :math:`0`.
* `dict`: Information on how to transform the real particle
parameters into those for the image charges. The valid
(optional) key-value pairs are:
* :code:`"charge"`: Index (`int`) where the charge
:math:`q` information is stored.
* :code:`"zero"`: Index (`int`) or indices (`list`,
`slice`, or `numpy.ndarray`) of parameters to zero out.
* :code:`"replace"`: `dict` containing key-value pairs of
indices (`int`) of parameters to change and their
replacement value(s). If the value is an `int`, all
particles receive that value for the current pair
potential. If the value is another `dict`, the new
parameter value (value of inner `dict`) is determined by
the current parameter value (key of inner `dict`).
To prevent unintended behavior, ensure that each parameter index
for a custom pair potential is used only once across the keys
listed above. If one does appear more than once, the order
in which parameter values are updated follows that of the list
above.
See the Examples section below for a simple illustration of the
possibilities outlined above.
wall_indices : array-like, keyword-only, optional
Atom indices corresponding to wall particles. If not provided,
the wall particles are guessed using the system dimensions.
exclude : `bool`, keyword-only, default: :code:`False`
Specifies whether interactions between a wall particle and all
image wall particles should be excluded. If :code:`False`, only
the interaction between a wall particle and its own image is
removed.
.. tip::
If you want accurate forces acting on wall particles, set
:code:`exclude=True`.
Returns
-------
positions : `numpy.ndarray`
Positions of the :math:`N` particles in the real system and
their images.
**Shape**: :math:`(2N,\,3)`.
**Reference unit**: :math:`\mathrm{nm}`.
integrator : `openmm.Integrator`
Integrator that simulates a system using Langevin dynamics, with
the LFMiddle discretization.
Examples
--------
Prepare the `cnbforces` `dict` argument for a system with the
following custom pair potentials:
* :code:`pair_gauss`: Gaussian potential with per-particle
parameters `type` and `beta`.
* `type` is the particle type, and is used to look up a
tabulated function for parameter values. As such, the new
values will vary and depend on the old ones. For example, for
a system with :math:`2` types of particles, :math:`0` becomes
:math:`2` and :math:`1` becomes :math:`3`.
* `beta` is assumed to be constant, with the value for the real
particles differing from that for image charges. For this
example, the new value is :math:`\beta = 42` (arbitrary
units).
* :code:`pair_wca`: Weeks–Chander–Andersen potential with per-
particle parameters `sigma` and `epsilon`, both of which
should be set to :math:`0` to turn off this potential for
image charges.
* :code:`pair_coul_recip`: Smeared Coulomb potential with per-
particle parameter `charge`, which should be flipped for image
charges, and `dummy`, a value that should be set to :math:`0`.
.. code::
NEW_BETA_VALUE = 42
cnbforces = {
pair_gauss: {"replace": {0: {0: 2, 1: 3}, 1: NEW_BETA_VALUE}},
pair_wca: None,
pair_coul_recip: {"charge": 0, "zero": 1}
}
References
----------
.. [1] Hautman, J.; Halley, J. W.; Rhee, Y. ‐J. Molecular Dynamics
Simulation of Water Beween Two Ideal Classical Metal Walls.
*J. Chem. Phys.* **1989**, *91* (1), 467–472.
https://doi.org/10.1063/1.457481.
.. [2] Dwelle, K. A.; Willard, A. P. Constant Potential,
Electrochemically Active Boundary Conditions for Electrochemical
Simulation. *J. Phys. Chem. C* **2019**, *123* (39), 24095–24103.
https://doi.org/10.1021/acs.jpcc.9b06635.
.. [3] Son, C. Y.; Wang, Z.-G. Image-Charge Effects on Ion
Adsorption near Aqueous Interfaces. *Proc. Natl. Acad. Sci.
U.S.A.* **2021**, *118* (19), e2020615118.
https://doi.org/10.1073/pnas.2020615118.
.. [4] Dos Santos, A. P.; Girotto, M.; Levin, Y. Simulations of
Coulomb Systems with Slab Geometry Using an Efficient 3D Ewald
Summation Method. *The Journal of Chemical Physics* **2016**,
**144** (14), 144103. https://doi.org/10.1063/1.4945560.
.. [5] Dos Santos, A. P.; Girotto, M.; Levin, Y. Simulations of
Coulomb Systems Confined by Polarizable Surfaces Using Periodic
Green Functions. *The Journal of Chemical Physics* **2017**,
**147** (18), 184105. https://doi.org/10.1063/1.4997420.
.. [6] Son, C. Y.; Wang, Z.-G. Manuscript in preparation.
"""
if not FOUND_ICPLUGIN:
emsg = ("An integrator capable of simulating a system with "
"image charges was not found. As such, the method of "
"image charges is unavailable unless the openmm-ic "
"(https://github.com/bbye98/mdhelper/tree/main/lib/openmm-ic-plugin) "
"or the constvplugin (https://github.com/scychon/openmm_constV) "
"package is installed.")
raise ImportError(emsg)
if np.isclose(gamma, 0):
emsg = ("Use the slab correction, available via "
"mdhelper.openmm.system.slab_correction(), for gamma=0.")
raise ValueError(emsg)
if not np.isclose(gamma, -1) and n_cells != 2:
emsg = ("The method of image charges with gamma != -1 is only "
"implemented for n_cells=2.")
raise ValueError(emsg)
def _ic_beta(gamma: float, x: float) -> float:
r"""
Computes the :math:`\beta` value used in the higher-order term
correction for the method of image charges with
:math:`\gamma\neq\pm1`.
Parameters
----------
gamma : `float`
Scaled image charge magnitude :math:`\gamma`.
x : `float`
Scaled :math:`x`-coordinate of the ion.
Returns
-------
beta : `float`
Approximated :math:`\beta` value.
"""
if not 0 <= x <= 1:
raise ValueError("'x' must be between 0 and 1.")
if np.isclose(x, 0.5):
return float(2 * special.zeta(3, 1.5)
- 2 * gamma ** 4 * mpmath.lerchphi(gamma ** 2, 3, 1.5))
else:
return (
special.zeta(2, 2 - x) - special.zeta(2, 1 + x) - gamma ** 4 *
float(mpmath.lerchphi(gamma ** 2, 2, 2 - x)
- mpmath.lerchphi(gamma ** 2, 2, 1 + x))
) / (2 * x - 1)
# Get system information
dims = np.asarray(
topology.getUnitCellDimensions().value_in_unit(unit.nanometer)
) * unit.nanometer
pbv = system.getDefaultPeriodicBoxVectors()
N_real_atoms = positions.shape[0]
if isinstance(positions, unit.Quantity):
positions = positions.value_in_unit(unit.nanometer)
# Guess indices of left and right walls if not provided
if wall_indices is None:
wall_indices = np.concatenate(
(np.isclose(positions[:, 2], 0).nonzero()[0],
np.isclose(positions[:, 2],
dims[2].value_in_unit(unit.nanometer)).nonzero()[0])
)
# Find averaged beta value for image charges with gamma =/= +/-1
beta = (_ic_beta(gamma, 0) + _ic_beta(gamma, 0.5)) / 2
# Set up higher-order image charge and slab corrections
cv_E_corr = openmm.CustomExternalForce("q*(1-2*z/L)")
cv_E_corr.addGlobalParameter("L", dims[2]) # real system z-dimension such
# that 0 <= (x = z / L_z) <= 1
cv_E_corr.addPerParticleParameter("q")
cv_M_z = openmm.CustomExternalForce("q*z")
cv_M_z.addPerParticleParameter("q")
cv_M_zz = openmm.CustomExternalForce("q*z^2")
cv_M_zz.addPerParticleParameter("q")
# Obtain particle charge information and register charges to
# corrections
if nbforce is None:
charge_index = None
for force, params in cnbforces:
if "charge" in params:
charge_index = params["charge"]
break
if charge_index is None:
raise ValueError("No charge information provided.")
else:
force = nbforce
charge_index = 0
q_tot = 0
if isinstance(force.getParticleParameters(0)[charge_index], unit.Quantity):
for i in range(force.getNumParticles()):
q = (force.getParticleParameters(i)[charge_index]
.value_in_unit(unit.elementary_charge))
q_tot += q
if not np.isclose(q, 0):
cv_E_corr.addParticle(i, (q,))
cv_M_z.addParticle(i, (q,))
cv_M_zz.addParticle(i, (q,))
else:
for i in range(force.getNumParticles()):
q = force.getParticleParameters(i)[charge_index]
q_tot += q
if not np.isclose(q, 0):
cv_E_corr.addParticle(i, (q,))
cv_M_z.addParticle(i, (q,))
cv_M_zz.addParticle(i, (q,))
electroneutral = np.isclose(q_tot, 0)
# Update and set new system dimensions
dims[2] *= n_cells
topology.setUnitCellDimensions(dims)
pbv[2] *= n_cells
system.setDefaultPeriodicBoxVectors(*pbv)
logging.info(f"Increased z-dimension to {dims[2]}.")
# Determine correction energy expression
corr_energy = ""
corr = openmm.CustomCVForce(corr_energy)
if not np.isclose(beta, 0):
corr_energy += "coef1*E_corr*M_z"
corr.addCollectiveVariable("E_corr", cv_E_corr)
corr.addGlobalParameter(
"coef1",
(unit.AVOGADRO_CONSTANT_NA * gamma * beta
/ (4 * np.pi * VACUUM_PERMITTIVITY * dims[2] ** 2))
.in_units_of(unit.kilojoule_per_mole /
(unit.elementary_charge ** 2 * unit.nanometer))
)
if not np.isclose(gamma, -1):
corr_energy += "+coef2*M_z^2"
if not electroneutral:
if np.isclose(gamma, 1):
corr_energy += "-coef2*q_tot*M_z*L_z"
elif np.isclose(gamma, -1):
corr_energy += "+coef2*q_tot*(M_z*L_z-M_zz)"
else:
corr_energy += "-coef2*q_tot*M_zz"
corr.addGlobalParameter("q_tot", q_tot)
if "coef2" in corr_energy:
corr.addGlobalParameter(
"coef2",
(unit.AVOGADRO_CONSTANT_NA
/ (2 * VACUUM_PERMITTIVITY * dims[0] * dims[1] * dims[2]))
.in_units_of(unit.kilojoule_per_mole
/ (unit.elementary_charge * unit.nanometer) ** 2)
)
if "L_z" in corr_energy:
corr.addGlobalParameter("L_z", dims[2]) # periodic system z-dimension
if "M_z" in corr_energy:
corr.addCollectiveVariable("M_z", cv_M_z)
if "M_zz" in corr_energy:
corr.addCollectiveVariable("M_zz", cv_M_zz)
if corr_energy:
if corr_energy.startswith("+"):
corr_energy = corr_energy.lstrip("+")
corr.setEnergyFunction(corr_energy)
system.addForce(corr)
logging.info("Added higher-order image charge and/or slab "
"correction(s).")
# Mirror particle positions
if n_cells == 2:
positions = np.concatenate(
(positions, positions * np.array((1, 1, -1), dtype=int))
) * unit.nanometer
else:
positions = np.tile(positions, (n_cells, 1))
for cell_index in range(1, n_cells):
start = cell_index * N_real_atoms
stop = (cell_index + 1) * N_real_atoms
positions[start:stop, 2] = (
(1 - 2 * (cell_index % 2)) * positions[start:stop, 2]
- 2 * np.floor(cell_index / 2)
* dims[2].value_in_unit(unit.nanometer)
)
positions *= unit.nanometer
logging.info(f"Replicated {N_real_atoms:,} particles {n_cells - 1} "
"time(s) over the z-axis.")
# Instantiate an integrator that simulates a system using
# Langevin dynamics and updates the image charge positions
integrator = ICLangevinIntegrator(temp, fric, dt, n_cells)
# Register image charges to the system, topology, and force field
cnbforces = cnbforces or {}
N_real_chains = topology.getNumChains()
atoms = list(topology.atoms())
residues = list(topology.residues())
coefs = (1, gamma)
for c in range(1, n_cells):
coef = coefs[c % 2]
chains_ic = [topology.addChain() for _ in range(N_real_chains)]
residues_ic = [topology.addResidue(f"IC_{r.name}",
chains_ic[r.chain.index])
for r in residues]
for i, atom in enumerate(atoms):
system.addParticle(0)
topology.addAtom(f"IC_{atom.name}", atom.element,
residues_ic[atom.residue.index])
if nbforce is not None:
nbforce.addParticle(
0 if i in wall_indices
else coef * nbforce.getParticleParameters(i)[0],
0, 0
)
for force, kwargs in cnbforces.items():
params = np.array(force.getParticleParameters(i))
if kwargs is None:
params[:] = 0
else:
if "charge" in kwargs:
params[kwargs["charge"]] *= (0 if i in wall_indices
else coef)
if "zero" in kwargs:
params[kwargs["zero"]] = 0
if "replace" in kwargs:
for index, value in kwargs["replace"].items():
params[index] = (value[params[index]]
if isinstance(value, dict)
else value)
force.addParticle(params)
logging.info(f"Registered {system.getNumParticles() - N_real_atoms:,} "
"image particles to the force field.")
# Add existing particle exclusions to mirrored image charges
for i in range(nbforce.getNumExceptions()):
i1, i2, qq = nbforce.getExceptionParameters(i)[:3]
if i1 not in wall_indices and i2 not in wall_indices:
for c in range(1, n_cells):
nbforce.addException(c * N_real_atoms + i1,
c * N_real_atoms + i2, qq, 0, 0)
for force in cnbforces:
i1, i2 = force.getExclusionParticles(i)
force.addExclusion(c * N_real_atoms + i1,
c * N_real_atoms + i2)
logging.info("Mirrored excluded non-wall image particle–image "
"particle interactions.")
# Prevent wall particles from interacting with all image wall
# particles
if exclude:
for i in wall_indices:
for j in wall_indices:
for c in range(1, n_cells):
nbforce.addException(i, c * N_real_atoms + j, 0, 0, 0)
for force in cnbforces:
force.addExclusion(i, c * N_real_atoms + j)
# Prevent wall particles from interacting their images
else:
for i in wall_indices:
for c in range(1, n_cells):
nbforce.addException(i, c * N_real_atoms + i, 0, 0, 0)
for force in cnbforces:
force.addExclusion(i, c * N_real_atoms + i)
logging.info("Removed wall–image wall interactions.")
return positions, integrator
[docs]
def add_electric_field(
system: openmm.System, nbforce: openmm.NonbondedForce,
E: Union[float, unit.Quantity], *, axis: int = 2,
dielectric: float = 1, charge_index: int = 0,
atom_indices: Union[int, np.ndarray[int]] = None
) -> None:
r"""
Adds an electric field to all charged particles by adding a force
:math:`f_i=q_iE` in the axis specified in `axis`, where :math:`q_i`
is the per-particle charge and :math:`E` is the electric field.
.. hint::
The following schematic shows how directionality is handled:
.. code::
|-| (-) ---> |+|
|-| <-- E -- |+|
|-| <--- (+) |+|
With a positive potential difference
(:math:`\Delta V>0\;\mathrm{V}`), the electric field is negative
(:math:`E<0\;\mathrm{V/m}`) such that it is pointing from the
right (positive) plate to the left (negative) plate. If an ion
has a positive charge (:math:`q_i>0\;\mathrm{e}`), the force will
be negative, indicating that it will be drawn to the left plate,
and vice versa.
Parameters
----------
system : `openmm.System`
OpenMM molecular system.
nbforce : `openmm.NonbondedForce` or `openmm.CustomNonbondedForce`
Pair potential object containing particle charge information.
.. note::
It is assumed that the charge :math:`q` information is the
first per-particle parameter stored in `nbforce`. If not, the
index can be specified in `charge_index`.
E : `float` or `openmm.unit.Quantity`
Electric field :math:`E`.
**Reference unit**: :math:`\mathrm{kJ/(mol\cdot nm\cdot e)}`.
axis : `int`, keyword-only, default: :code:`2`
Axis along which the walls are placed. :code:`0`, :code:`1`, and
:code:`2` correspond to :math:`x`, :math:`y`, and :math:`z`,
respectively.
dielectric : `float`, keyword-only, default: :code:`1`
Relative permittivity :math:`\varepsilon_\mathrm{r}` of the
medium. Used to scale the particle charges by
:math:`\sqrt{\varepsilon_\mathrm{r}}` and recover the original
values.
charge_index : `int`, keyword-only, default: :code:`0`
Index of charge :math:`q` information in the per-particle
parameters stored in `nbforce`.
atom_indices : `int` or array-like, keyword-only, optional
Indices of atoms to apply the electric field to. By default,
the electric field is applied to all atoms, but this can be
computationally expensive when there are charged particles that
do not need to be included, such as image charges. If an `int`
is provided, all atoms up to that index are included.
"""
# Get letter representation of axis for formula
z = chr(120 + axis)
# Get indices of atoms that are affected by the electric field
if atom_indices is None:
atom_indices = range(nbforce.getNumParticles())
elif isinstance(atom_indices, int):
atom_indices = range(atom_indices)
# Create and register particles to the electric field
efield = openmm.CustomExternalForce(f"-q*E*{z}")
efield.addGlobalParameter("E", E)
efield.addPerParticleParameter("q")
for i in atom_indices:
q = nbforce.getParticleParameters(i)[charge_index]
if isinstance(q, unit.Quantity):
q = q.value_in_unit(unit.elementary_charge)
if not np.isclose(q, 0):
efield.addParticle(i, (q * np.sqrt(dielectric),))
system.addForce(efield)
[docs]
def estimate_pressure_tensor(
context: openmm.Context, dh: float = 1e-5, *, diag: bool = False
) -> np.ndarray[float]:
r"""
Computes the estimated pressure tensor using a central finite
difference for the virial contribution.
The pressure tensor is given by
.. math::
\mathbf{p}=\frac{1}{V}\left(
\sum_{i=1}^N m_i\mathbf{v}_i^\mathsf{T}\mathbf{v}_i
+\mathbf{h}^\mathsf{T}\frac{dU}{d\mathbf{h}}\right)
where :math:`V` is the volume, :math:`m_i` and :math:`\mathbf{v}_i`
are the mass and velocity vector of particle :math:`i`,
:math:`\mathbf{h}` is a :math:`3\times 3` matrix where the rows
contain the box vectors, and :math:`U` is the total pairwise
potential energy.
To evaluate :math:`dU/d\mathbf{h}`, the box vectors are perturbed by
`dh` in each of the six "directions" and the positions are updated
accordingly. Then, OpenMM reevaluates the potential
energy based on the new periodic box vectors and particle positions,
and the difference in potential energy is used in the central finite
difference formula to estimate the derivative.
Parameters
----------
context : `openmm.Context`
OpenMM simulation context.
dh : `float`, default: :code:`1e-5`
Finite difference step size.
diag : `bool`, keyword-only, default: :code:`False`
Determines whether only the values in the main diagonal of the
pressure tensor (i.e., the values that, when summed, gives the
system pressure) are calculated.
Returns
-------
pres : `numpy.ndarray`
Estimated pressure tensor (or diagonal components).
**Shape**: :math:`(3,)` or :math:`(3,\,3)`.
**Reference unit**: :math:`\mathrm{atm}`.
"""
try:
state = context.getState(getPositions=True, getVelocities=True,
getEnergy=True)
box = state.getPeriodicBoxVectors(asNumpy=True)
positions = state.getPositions(asNumpy=True)
velocities = state.getVelocities(asNumpy=True)
volume = box[0, 0] * box[1, 1] * box[2, 2]
except openmm.OpenMMException:
emsg = ("The simulation context must have information about "
"the particle positions and velocities.")
raise ValueError(emsg)
system = context.getSystem()
masses = np.fromiter(
(system.getParticleMass(i).value_in_unit(unit.dalton)
for i in range(system.getNumParticles())),
dtype=float
) * unit.dalton
if diag:
# Compute the ideal contribution
p_kinetic = (masses * velocities ** 2).sum(axis=0)
# Estimate the virial contribution with a central finite difference
p_virial = np.zeros(3) * unit.kilojoule_per_mole
for i in range(3):
box_ = box.copy()
box_[i, i] += dh
context.setPeriodicBoxVectors(*box_)
context.setPositions(
np.dot(positions,
np.divide(box_, box, out=np.zeros_like(box),
where=box.value_in_unit(unit.nanometer) != 0))
)
U_plus = context.getState(getEnergy=True).getPotentialEnergy()
box_ = box.copy()
box_[i, i] -= dh
context.setPeriodicBoxVectors(*box_)
context.setPositions(
np.dot(positions,
np.divide(box_, box, out=np.zeros_like(box),
where=box.value_in_unit(unit.nanometer) != 0))
)
U_minus = context.getState(getEnergy=True).getPotentialEnergy()
p_virial[i] = U_plus - U_minus
p_virial = (p_virial / (2 * dh)).in_units_of(p_kinetic.unit)
else:
# Compute the ideal contribution
p_kinetic = (masses * velocities * velocities[:, :, None]).sum(axis=0)
# Estimate the virial contribution with a central finite difference
p_virial = np.zeros((3, 3)) * unit.kilojoule_per_mole
for i in range(3):
for j in range(i + 1):
box_ = box.copy()
box_[i, j] += dh
context.setPeriodicBoxVectors(*box_)
context.setPositions(
np.dot(positions,
np.divide(box_, box, out=np.zeros_like(box),
where=box.value_in_unit(unit.nanometer) != 0))
)
U_plus = context.getState(getEnergy=True).getPotentialEnergy()
box_ = box.copy()
box_[i, j] -= dh
context.setPeriodicBoxVectors(*box_)
context.setPositions(
np.dot(positions,
np.divide(box_, box, out=np.zeros_like(box),
where=box.value_in_unit(unit.nanometer) != 0))
)
U_minus = context.getState(getEnergy=True).getPotentialEnergy()
p_virial[i, j] = U_plus - U_minus
p_virial = (p_virial / (2 * dh)).in_units_of(p_kinetic.unit)
p_virial = (
p_virial._value + np.tril(p_virial).T
- np.diag(np.diag(p_virial))
) * p_virial.unit
return (
(p_kinetic + p_virial) / (unit.AVOGADRO_CONSTANT_NA * volume)
).in_units_of(unit.atmosphere)