Coverage for cvpack/attraction_strength.py: 98%
61 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 00:10 +0000
1"""
2.. class:: AttractionStrength
3 :platform: Linux, MacOS, Windows
4 :synopsis: The strength of the attraction between two atom groups
6.. classauthor:: Charlles Abreu <craabreu@gmail.com>
8"""
10from __future__ import annotations
12import typing as t
14import numpy as np
15import openmm
16from openmm import unit as mmunit
18from .collective_variable import CollectiveVariable
19from .units import Quantity, ScalarQuantity
20from .utils import evaluate_in_context
22ONE_4PI_EPS0 = 138.93545764438198
25def _cutoff_distance(nonbondedForce: openmm.NonbondedForce) -> float:
26 """The cutoff distance of a NonbondedForce in nm, which the reaction field needs."""
27 if nonbondedForce.getNonbondedMethod() == openmm.NonbondedForce.NoCutoff:
28 raise ValueError("The NonbondedForce must use a cutoff distance")
29 return nonbondedForce.getCutoffDistance().value_in_unit(mmunit.nanometers)
32def _lennard_jones_switch(nonbondedForce: openmm.NonbondedForce, cutoff: float) -> str:
33 """
34 The factor and definitions that apply the switching function of a
35 NonbondedForce to a Lennard-Jones term, or an empty string if it uses none.
36 As in NonbondedForce, the switching function does not apply to electrostatics.
37 """
38 if not nonbondedForce.getUseSwitchingFunction():
39 return ""
40 rs = nonbondedForce.getSwitchingDistance().value_in_unit(mmunit.nanometers)
41 return (
42 f"*select(step(r - {rs}), 1 - t^3*(10 - t*(15 - 6*t)), 1)"
43 f"; t = (r - {rs})/{cutoff - rs}"
44 )
47class AttractionStrength(CollectiveVariable, openmm.CustomNonbondedForce):
48 r"""
49 The strength of the attraction between two atom groups:
51 .. math::
53 S_{\rm attr}({\bf r}) = S_{1,2}({\bf r}) = &-\frac{1}{E_{\rm ref}} \Bigg[
54 \sum_{i \in {\bf g}_1}
55 \sum_{\substack{j \in {\bf g}_2 \\ j \neq i}}
56 \epsilon_{ij} u_{\rm disp}\left(
57 \frac{\|{\bf r}_i - {\bf r}_j\|}{\sigma_{ij}}
58 \right) \\
59 &+\sum_{i \in {\bf g}_1}
60 \sum_{\substack{j \in {\bf g}_2 \\ q_jq_i < 0}}
61 \frac{q_i q_j}{4 \pi \varepsilon_0 r_{\rm c}} u_{\rm elec}\left(
62 \frac{\|{\bf r}_i - {\bf r}_j\|}{r_{\rm c}}
63 \right) \Bigg]
65 where :math:`{\bf g}_1` and :math:`{\bf g}_2` are the two atom groups,
66 :math:`r_{\rm c}` is the cutoff distance, :math:`\varepsilon_0` is the
67 permittivity of empty space, :math:`q_i` is the charge of atom :math:`i`, and
68 :math:`E_{\rm ref}` is a reference value (in energy units per mole).
69 The Lennard-Jones parameters are given by the Lorentz-Berthelot mixing rule, i.e.
70 :math:`\epsilon_{ij} = \sqrt{\epsilon_i \epsilon_j}`, and
71 :math:`\sigma_{ij} = (\sigma_i + \sigma_j)/2`.
73 Optionally, one can provide a third atom group :math:`{\bf g}_3` to contrast
74 the attraction strength between :math:`{\bf g}_1` and :math:`{\bf g}_2` with
75 that between :math:`{\bf g}_1` and :math:`{\bf g}_3`. One can also provide
76 a scaling factor :math:`\alpha` to balance the contributions of the two
77 interactions. In this case, the collective variable becomes
79 .. math::
81 S_{\rm attr}({\bf r}) = S_{1,2}({\bf r}) - \alpha S_{1,3}({\bf r})
83 .. note::
85 Groups :math:`{\bf g}_1` and :math:`{\bf g}_2` can overlap or even be the
86 same, in which case the collective variable will measure the strength of
87 the self-attraction of :math:`{\bf g}_1`. Each pair of distinct atoms
88 contributes a single term, even if both atoms are in both groups. On the
89 other hand, the contrast group :math:`{\bf g}_3` cannot overlap with either
90 :math:`{\bf g}_1` or :math:`{\bf g}_2`.
92 The function :math:`u_{\rm disp}(x)` is a Lennard-Jones-like reduced potential with
93 a highly softened repulsion part, defined as
95 .. math::
97 u_{\rm disp}(x) = 4\left(\frac{1}{y^2} - \frac{1}{y}\right),
98 \quad \text{where} \quad
99 y = |x^6 - 2| + 2
101 The function :math:`u_{\rm elec}(x)` provides a Coulomb-like decay with
102 reaction-field screening, defined as
104 .. math::
106 u_{\rm elec}(x) = \frac{1}{x} + \frac{x^2 - 3}{2}
108 The screening considers a perfect conductor as the surrounding medium
109 :cite:`Correa_2022`.
111 .. note::
113 Only attractive electrostatic interactions are considered (:math:`q_i q_j < 0`),
114 which gives :math:`S_{\rm attr}({\bf r})` a lower bound of zero. The upper
115 bound depends on the system details, the chosen groups of atoms, and the
116 adopted reference value.
118 The Lennard-Jones parameters, atomic charges, cutoff distance, boundary conditions,
119 as well as whether to use a switching function and its corresponding switching
120 distance, are taken from the :openmm:`NonbondedForce` object. As in
121 :openmm:`NonbondedForce`, the switching function applies only to the
122 Lennard-Jones term. Pairs with :math:`\sigma_{ij} = 0` have no Lennard-Jones term.
124 .. note::
125 Any non-exclusion exceptions involving an atom in :math:`{\bf g}_1` and an atom
126 in either :math:`{\bf g}_2` or :math:`{\bf g}_3` are turned into exclusions in
127 this collective variable.
129 Parameters
130 ----------
131 group1
132 The first atom group.
133 group2
134 The second atom group.
135 nonbondedForce
136 The :class:`openmm.NonbondedForce` object from which to collect the necessary
137 parameters.
138 contrastGroup
139 An optional third atom group to contrast the attraction strength between
140 :math:`{\bf g}_1` and :math:`{\bf g}_2` with that between :math:`{\bf g}_1`
141 and :math:`{\bf g}_3`.
142 reference
143 A reference value (in energy units per mole) to which the collective variable
144 should be normalized. One can also provide a :OpenMM:`Context` object from
145 which to extract a reference attraction strength. The extracted value will be
146 :math:`S_{1,2}({\bf r})` for :math:`E_{\rm ref} = 1`, regardless of whether
147 `contrastGroup` is provided or not.
148 contrastScaling
149 A scaling factor :math:`\alpha` to balance the contributions of the two
150 interactions. The default is 1.0.
151 name
152 The name of the collective variable.
154 Raises
155 ------
156 ValueError
157 If the contrast group overlaps with group1 or group2
158 ValueError
159 If contrastScaling is negative
160 ValueError
161 If the nonbonded force does not use a cutoff distance
163 Examples
164 --------
165 >>> import cvpack
166 >>> from openmm import unit
167 >>> from openmmtools import testsystems
168 >>> model = testsystems.HostGuestExplicit()
169 >>> guest = [a.index for a in model.topology.atoms() if a.residue.name == "B2"]
170 >>> host = [a.index for a in model.topology.atoms() if a.residue.name == "CUC"]
171 >>> forces = {f.getName(): f for f in model.system.getForces()}
172 >>> cv1 = cvpack.AttractionStrength(guest, host, forces["NonbondedForce"])
173 >>> cv1.addToSystem(model.system)
174 >>> platform = openmm.Platform.getPlatformByName("Reference")
175 >>> integrator = openmm.VerletIntegrator(1.0 * mmunit.femtoseconds)
176 >>> context = openmm.Context(model.system, integrator, platform)
177 >>> context.setPositions(model.positions)
178 >>> cv1.getValue(context)
179 4912.8... dimensionless
181 >>> water = [a.index for a in model.topology.atoms() if a.residue.name == "HOH"]
182 >>> cv2 = cvpack.AttractionStrength(guest, water, forces["NonbondedForce"])
183 >>> cv2.addToSystem(model.system)
184 >>> context.reinitialize(preserveState=True)
185 >>> cv2.getValue(context)
186 2090.2... dimensionless
188 >>> cv3 = cvpack.AttractionStrength(guest, host, forces["NonbondedForce"], water)
189 >>> cv3.addToSystem(model.system)
190 >>> context.reinitialize(preserveState=True)
191 >>> cv3.getValue(context)
192 2822.63... dimensionless
193 >>> print(cv1.getValue(context) - cv2.getValue(context))
194 2822.63... dimensionless
196 >>> cv4 = cvpack.AttractionStrength(
197 ... guest, host, forces["NonbondedForce"], water, contrastScaling=0.5
198 ... )
199 >>> cv4.addToSystem(model.system)
200 >>> context.reinitialize(preserveState=True)
201 >>> cv4.getValue(context)
202 3867.7... dimensionless
203 >>> 1 * cv1.getValue(context) - 0.5 * cv2.getValue(context)
204 3867.7...
205 """
207 def __init__(
208 self,
209 group1: t.Iterable[int],
210 group2: t.Iterable[int],
211 nonbondedForce: openmm.NonbondedForce,
212 contrastGroup: t.Iterable[int] | None = None,
213 reference: ScalarQuantity | openmm.Context = Quantity(
214 1.0 * mmunit.kilojoule_per_mole
215 ),
216 contrastScaling: float = 1.0,
217 name: str = "attraction_strength",
218 ) -> None:
219 group1 = list(group1)
220 group2 = list(group2)
221 contrasting = contrastGroup is not None
222 if contrasting:
223 contrastGroup = list(contrastGroup)
224 if set(contrastGroup) & (set(group1) | set(group2)):
225 raise ValueError(
226 "The contrast group must be disjoint from group1 and group2"
227 )
228 if contrastScaling < 0:
229 raise ValueError("contrastScaling must be non-negative")
230 cutoff = _cutoff_distance(nonbondedForce)
231 # Pairs with sigma = 0 have no Lennard-Jones term, and select() keeps their
232 # NaN derivatives out of the forces
233 expression = (
234 f"-(lj + coul){'*sign1*sign2' if contrasting else ''}/ref"
235 "; lj = select(sigma, 4*epsilon*(1/y^2 - 1/y), 0)"
236 f"{_lennard_jones_switch(nonbondedForce, cutoff)}"
237 f"; coul = {ONE_4PI_EPS0 / cutoff}*q1q2*(1/x + (x^2 - 3)/2)"
238 f"; x = r/{cutoff}"
239 "; y = abs((r/sigma)^6 - 2) + 2"
240 "; q1q2 = min(0, charge1*charge2)"
241 "; epsilon = sqrt(epsilon1*epsilon2)"
242 "; sigma = (sigma1 + sigma2)/2"
243 "; ref = 1"
244 )
245 super().__init__(expression)
246 self.setNonbondedMethod(
247 self.CutoffPeriodic
248 if nonbondedForce.usesPeriodicBoundaryConditions()
249 else self.CutoffNonPeriodic
250 )
251 self.setCutoffDistance(cutoff)
252 for parameter in ("charge", "sigma", "epsilon") + ("sign",) * contrasting:
253 self.addPerParticleParameter(parameter)
254 contrast_atoms = np.zeros(nonbondedForce.getNumParticles(), dtype=bool)
255 if contrasting:
256 contrast_atoms[contrastGroup] = True
257 for atom, in_group in enumerate(contrast_atoms):
258 charge, sigma, epsilon = nonbondedForce.getParticleParameters(atom)
259 if contrasting:
260 sign = -1 if in_group else 1
261 scale = contrastScaling if in_group else 1
262 self.addParticle([charge * scale, sigma, epsilon * scale**2, sign])
263 else:
264 self.addParticle([charge, sigma, epsilon])
265 for exception in range(nonbondedForce.getNumExceptions()):
266 i, j, *_ = nonbondedForce.getExceptionParameters(exception)
267 self.addExclusion(i, j)
268 self.setUseSwitchingFunction(False)
269 self.setUseLongRangeCorrection(False)
270 self.addInteractionGroup(group1, group2)
271 if isinstance(reference, openmm.Context):
272 reference = evaluate_in_context(self, reference)
273 elif mmunit.is_quantity(reference):
274 reference = reference.value_in_unit(mmunit.kilojoule_per_mole)
275 self.setEnergyFunction(expression.replace("ref = 1", f"ref = {reference}"))
276 if contrasting:
277 self.addInteractionGroup(group1, contrastGroup)
279 self._registerCV(
280 name,
281 mmunit.dimensionless,
282 group1=group1,
283 group2=group2,
284 nonbondedForce=nonbondedForce,
285 contrastGroup=contrastGroup,
286 reference=reference,
287 contrastScaling=contrastScaling,
288 )
291AttractionStrength.registerTag("!cvpack.AttractionStrength")