Coverage for cvpack/shortest_distance.py: 100%
47 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 00:10 +0000
1"""
2.. class:: ShortestDistance
3 :platform: Linux, MacOS, Windows
4 :synopsis: A smooth approximation for the shortest distance between two atom groups
6.. classauthor:: Charlles Abreu <craabreu@gmail.com>
8"""
10import typing as t
12import numpy as np
13import openmm
14from openmm import unit as mmunit
16from .collective_variable import CollectiveVariable
17from .units import Quantity
20class ShortestDistance(CollectiveVariable, openmm.CustomCVForce):
21 r"""
22 A smooth approximation for the shortest distance between two atom groups:
24 .. math::
25 r_{\rm min}({\bf r}) = \frac{S_{rw}({\bf r})}{S_w({\bf r})}
27 with
29 .. math::
31 S_{rw}({\bf r}) = r_c e^{-\beta} +
32 \sum_{i \in {\bf g}_1} \sum_{j \in {\bf g}_2 \atop r_{ij} < r_c}
33 r_{ij} \left(e^{-\beta \frac{r_{ij}}{r_c}} - e^{-\beta}\right)
35 and
37 .. math::
39 S_w({\bf r}) = e^{-\beta} +
40 \sum_{i \in {\bf g}_1} \sum_{j \in {\bf g}_2 \atop r_{ij} < r_c}
41 \left(e^{-\beta \frac{r_{ij}}{r_c}} - e^{-\beta}\right)
43 where :math:`r_{ij} = \|{\bf r}_j - {\bf r}_i\|` is the distance between atoms
44 :math:`i` and :math:`j`, :math:`{\bf g}_1` and :math:`{\bf g}_2` are the sets of
45 indices of the atoms in the first and second groups, respectively, :math:`r_c` is
46 the cutoff distance, and :math:`\beta` is a parameter that controls the degree of
47 approximation.
49 The larger the value of :math:`\beta`, the closer the approximation to the true
50 shortest distance. However, an excessively large value may lead to numerical
51 instability.
53 The terms outside the summations guarantee that the shortest distance is
54 well-defined even when there are no atom pairs within the cutoff distance. In this
55 case, the collective variable evaluates to the cutoff distance. Shifting each weight
56 by :math:`e^{-\beta}` makes the contribution of a pair vanish at the cutoff
57 distance, so the collective variable is continuous when a pair crosses it.
59 .. note::
61 Atoms are allowed to be in both groups. In this case, terms for which
62 :math:`i = j` are ignored, and each pair of distinct atoms contributes a
63 single term, even if both atoms are in both groups.
65 .. note::
67 GPU platforms accumulate forces in 64-bit fixed point, which cannot represent
68 the forces of :math:`S_{rw}` and :math:`S_w` when the weights span many orders
69 of magnitude. To keep them in range, the collective variable is evaluated in
70 several windows. Each window multiplies all weights by a different constant,
71 which leaves :math:`r_{\rm min}` unchanged, and the window used is the one
72 whose :math:`S_w` is the largest not exceeding :math:`10^6`. The number of
73 windows, and thus the cost, grows with :math:`\beta`: one window for
74 :math:`\beta \leq 13`, five for the default :math:`\beta = 100`, and at most
75 16. For :math:`\beta > 344`, 16 windows are not enough, and forces computed on
76 GPU platforms lose precision.
78 Parameters
79 ----------
80 group1
81 The indices of the atoms in the first group.
82 group2
83 The indices of the atoms in the second group.
84 numAtoms
85 The total number of atoms in the system, including those that are not in any
86 of the groups. This is necessary for the correct initialization of the
87 collective variable.
88 beta
89 The parameter that controls the degree of approximation.
90 cutoffDistance
91 The cutoff distance for evaluating atom pairs.
92 pbc
93 Whether to use periodic boundary conditions in distance calculations. This
94 should be ``True`` for periodic systems.
95 name
96 The name of the collective variable.
98 Example
99 -------
100 >>> import cvpack
101 >>> import numpy as np
102 >>> from openmm import unit
103 >>> from openmmtools import testsystems
104 >>> model = testsystems.HostGuestVacuum()
105 >>> group1, group2 = [], []
106 >>> for residue in model.topology.residues():
107 ... group = group1 if residue.name == "B2" else group2
108 ... group.extend(atom.index for atom in residue.atoms())
109 >>> coords = model.positions.value_in_unit(mmunit.nanometers)
110 >>> print(np.linalg.norm(
111 ... coords[group1, None, :] - coords[None, group2, :],
112 ... axis=-1,
113 ... ).min())
114 0.1757...
115 >>> num_atoms = model.system.getNumParticles()
116 >>> min_dist = cvpack.ShortestDistance(group1, group2, num_atoms)
117 >>> min_dist.addToSystem(model.system)
118 >>> platform = openmm.Platform.getPlatformByName("Reference")
119 >>> integrator = openmm.VerletIntegrator(1.0 * mmunit.femtoseconds)
120 >>> context = openmm.Context(model.system, integrator, platform)
121 >>> context.setPositions(model.positions)
122 >>> min_dist.getValue(context)
123 0.1785... nm
124 """
126 WINDOW_EFOLDS = 22
127 MAX_WEIGHT_SUM = 1e6
128 MAX_WINDOWS = 16
130 def __init__(
131 self,
132 group1: t.Sequence[int],
133 group2: t.Sequence[int],
134 numAtoms: int,
135 beta: float = 100,
136 cutoffDistance: mmunit.Quantity = Quantity(1 * mmunit.nanometers),
137 pbc: bool = False,
138 name: str = "shortest_distance",
139 ) -> None:
140 rc = cutoffDistance
141 if mmunit.is_quantity(rc):
142 rc = rc.value_in_unit(mmunit.nanometers)
143 # Window k multiplies all weights by exp(beta*a_k/rc), which leaves S_rw/S_w
144 # unchanged. The offsets a_k decrease from rc in steps of efolds*rc/beta until
145 # the weights of the last window are at most MAX_WEIGHT_SUM. A CustomCVForce
146 # holds at most 32 inner forces, so large beta values need wider steps.
147 log_max = np.log(self.MAX_WEIGHT_SUM)
148 num_steps = max(0, int(np.ceil((beta - log_max) / self.WINDOW_EFOLDS)))
149 num_windows = min(1 + num_steps, self.MAX_WINDOWS)
150 efolds = self.WINDOW_EFOLDS
151 if num_windows < 1 + num_steps:
152 efolds = (beta - log_max) / (num_windows - 1)
153 definitions = []
154 expression = None
155 super().__init__("")
156 for k in reversed(range(num_windows)):
157 offset = rc - k * efolds * rc / beta
158 shift = np.exp(-k * efolds)
159 # The clamp only acts on windows whose weight sum exceeds MAX_WEIGHT_SUM,
160 # which are never selected, and prevents overflows in them. Factoring out
161 # the shift keeps constants that underflow in single precision out of the
162 # GPU kernels.
163 exponent = f"min({2 * log_max}, {beta / rc}*({offset}-r))"
164 weight = f"exp({exponent})*(1-exp({beta / rc}*(r-{rc})))"
165 for cv, function in [(f"wrsum{k}", f"{weight}*r"), (f"wsum{k}", weight)]:
166 force = openmm.CustomNonbondedForce(function)
167 force.setNonbondedMethod(
168 force.CutoffPeriodic if pbc else force.CutoffNonPeriodic
169 )
170 force.setCutoffDistance(cutoffDistance)
171 force.setUseSwitchingFunction(False)
172 force.setUseLongRangeCorrection(False)
173 for _ in range(numAtoms):
174 force.addParticle([])
175 force.addInteractionGroup(group1, group2)
176 self.addCollectiveVariable(cv, force)
177 definitions.append(f"value{k}=({rc * shift}+wrsum{k})/sum{k}")
178 definitions.append(f"sum{k}={shift}+wsum{k}")
179 if expression is None:
180 expression = f"value{k}"
181 else:
182 condition = f"step({self.MAX_WEIGHT_SUM}-sum{k})"
183 expression = f"select({condition}, value{k}, {expression})"
184 self.setEnergyFunction("; ".join([expression, *definitions]))
185 self._registerCV(
186 name,
187 mmunit.nanometers,
188 group1=group1,
189 group2=group2,
190 numAtoms=numAtoms,
191 beta=beta,
192 cutoffDistance=rc,
193 pbc=pbc,
194 )
197ShortestDistance.registerTag("!cvpack.ShortestDistance")