Coverage for cvpack/atomic_function.py: 93%
105 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 00:10 +0000
1"""
2.. class:: AtomicFunction
3 :platform: Linux, MacOS, Windows
4 :synopsis: A generic function of the coordinates of a group of atoms
6.. classauthor:: Charlles Abreu <craabreu@gmail.com>
8"""
10from __future__ import annotations
12import numpy as np
13import openmm
14from numpy.typing import ArrayLike
15from openmm import unit as mmunit
17from .base_custom_function import BaseCustomFunction
18from .units import ScalarQuantity, VectorQuantity
21class AtomicFunction(BaseCustomFunction, openmm.CustomCompoundBondForce):
22 r"""
23 A generic function of the coordinates of :math:`m` groups of :math:`n` atoms:
25 .. math::
27 f({\bf r}) = \sum_{i=1}^m F\left(
28 {\bf r}_{i,1}, {\bf r}_{i,2}, \dots, {\bf r}_{i,n}
29 \right)
31 where :math:`F` is a user-defined function and :math:`{\bf r}_{i,j}` is the
32 coordinate of the :math:`j`-th atom of the :math:`i`-th group.
34 The function :math:`F` is defined as a string and can be any expression supported
35 by :OpenMM:`CustomCompoundBondForce`. If the expression contains named parameters,
36 the value of each parameter can be passed in one of three ways:
38 #. By a semicolon-separated definition in the function string, such as described
39 in the :OpenMM:`CustomCompoundBondForce` documentation. In this case, the
40 parameter value will be the same for all groups of atoms.
42 #. By a 1D array or list of length :math:`m` passed as a keyword argument to
43 the :class:`AtomicFunction` constructor. In this case, each value will be
44 assigned to a different group of atoms.
46 #. By a scalar passed as a keyword argument to the :class:`AtomicFunction`
47 constructor. In this case, the parameter will apply to all atom groups and will
48 become available for on-the-fly modification during a simulation via the
49 ``setParameter`` method of an :OpenMM:`Context` object. **Warning**: other
50 collective variables or :OpenMM:`Force` objects in the same system will share
51 the same values of equal-named parameters.
53 Parameters
54 ----------
55 function
56 The function to be evaluated. It must be a valid
57 :OpenMM:`CustomCompoundBondForce` expression.
58 groups
59 The indices of the atoms in each group, passed as a 2D array-like object of
60 shape `(m, n)`, where `m` is the number of groups and `n` is the number of
61 atoms per group. If a 1D object is passed, it is assumed that `m` is 1 and
62 `n` is the length of the object.
63 unit
64 The unit of measurement of the collective variable. It must be compatible
65 with the MD unit system (mass in `daltons`, distance in `nanometers`, time
66 in `picoseconds`, temperature in `kelvin`, energy in `kilojoules_per_mol`,
67 angle in `radians`). If the collective variable does not have a unit, use
68 `unit.dimensionless`.
69 periodicBounds
70 The lower and upper bounds of the collective variable if it is periodic, or
71 ``None`` if it is not.
72 pbc
73 Whether to use periodic boundary conditions when computing atomic distances.
74 This should be ``True`` for periodic systems.
75 name
76 The name of the collective variable.
78 Keyword Args
79 ------------
80 **parameters
81 The named parameters of the function. Each parameter can be a 1D array-like
82 object or a scalar. In the former case, the array must have length :math:`m`
83 and each entry will be assigned to a different group of atoms. In the latter
84 case, it will define a global :OpenMM:`Context` parameter.
86 Raises
87 ------
88 ValueError
89 If the groups are not specified as a 1D or 2D array-like object.
90 ValueError
91 If the unit of the collective variable is not compatible with the MD unit
92 system.
94 Example
95 -------
96 >>> import cvpack
97 >>> import openmm
98 >>> import numpy as np
99 >>> from openmm import unit
100 >>> from openmmtools import testsystems
101 >>> model = testsystems.AlanineDipeptideVacuum()
102 >>> angle1 = cvpack.Angle(0, 5, 10)
103 >>> angle2 = cvpack.Angle(10, 15, 20)
104 >>> colvar = cvpack.AtomicFunction(
105 ... "(kappa/2)*(angle(p1, p2, p3) - theta0)^2",
106 ... unit.kilojoules_per_mole,
107 ... [[0, 5, 10], [10, 15, 20]],
108 ... kappa = 1000 * unit.kilojoules_per_mole/unit.radian**2,
109 ... theta0 = [np.pi/2, np.pi/3] * unit.radian,
110 ... )
111 >>> for cv in [angle1, angle2, colvar]:
112 ... cv.addToSystem(model.system)
113 >>> integrator = openmm.VerletIntegrator(0)
114 >>> platform = openmm.Platform.getPlatformByName('Reference')
115 >>> context = openmm.Context(model.system, integrator, platform)
116 >>> context.setPositions(model.positions)
117 >>> theta1 = angle1.getValue(context) / openmm.unit.radian
118 >>> theta2 = angle2.getValue(context) / openmm.unit.radian
119 >>> 500*((theta1 - np.pi/2)**2 + (theta2 - np.pi/3)**2)
120 429.479...
121 >>> colvar.getValue(context)
122 429.479... kJ/mol
123 >>> context.setParameter(
124 ... "kappa",
125 ... 2000 * unit.kilojoules_per_mole/unit.radian**2
126 ... )
127 >>> colvar.getValue(context)
128 858.958... kJ/mol
129 """
131 def __init__(
132 self,
133 function: str,
134 unit: mmunit.Unit,
135 groups: ArrayLike,
136 periodicBounds: VectorQuantity | None = None,
137 pbc: bool = False,
138 name: str = "atomic_function",
139 **parameters: ScalarQuantity | VectorQuantity,
140 ) -> None:
141 groups = np.atleast_2d(groups)
142 num_groups, atoms_per_group, *other_dimensions = groups.shape
143 if other_dimensions:
144 raise ValueError("Array `groups` cannot have more than 2 dimensions")
145 super().__init__(atoms_per_group, function)
146 overalls, perbonds = self._extractParameters(num_groups, **parameters)
147 self._addParameters(overalls, perbonds, groups, pbc)
148 groups = [[int(atom) for atom in group] for group in groups]
149 self._registerCV(
150 name,
151 unit,
152 function=function,
153 unit=unit,
154 groups=groups,
155 periodicBounds=periodicBounds,
156 pbc=pbc,
157 **overalls,
158 **perbonds,
159 )
160 if periodicBounds is not None:
161 self._registerPeriodicBounds(*periodicBounds)
163 @classmethod
164 def _fromCustomForce( # noqa: PLR0912
165 cls,
166 force: openmm.CustomAngleForce
167 | openmm.CustomBondForce
168 | openmm.CustomCompoundBondForce
169 | openmm.CustomExternalForce
170 | openmm.CustomTorsionForce,
171 unit: mmunit.Unit,
172 periodicBounds: VectorQuantity | None = None,
173 pbc: bool = False,
174 ) -> AtomicFunction:
175 """
176 Create an :class:`AtomicFunction` from an object of :openmm:`CustomBondForce`,
177 :openmm:`CustomAngleForce`, :openmm:`CustomTorsionForce`, or
178 :openmm:`CustomCompoundBondForce` class.
179 """
180 if isinstance(force, openmm.CustomExternalForce):
181 if "periodicdistance" in force.getEnergyFunction():
182 raise ValueError(
183 "Converting a CustomExternalForce that uses periodicdistance is "
184 "not supported"
185 )
186 number, item, definition = 1, "Particle", "; x=x1; y=y1; z=z1"
187 elif isinstance(force, openmm.CustomBondForce):
188 number, item, definition = 2, "Bond", "; r=distance(p1, p2)"
189 elif isinstance(force, openmm.CustomAngleForce):
190 number, item, definition = 3, "Angle", "; theta=angle(p1, p2, p3)"
191 elif isinstance(force, openmm.CustomTorsionForce):
192 number, item, definition = 4, "Torsion", "; theta=dihedral(p1, p2, p3, p4)"
193 elif isinstance(force, openmm.CustomCompoundBondForce):
194 number, item, definition = force.getNumParticlesPerBond(), "Bond", ""
195 else:
196 raise TypeError(f"Unsupported force type: {type(force).__name__}")
197 function = force.getEnergyFunction() + definition
198 parameters = {}
199 for i in range(force.getNumGlobalParameters()):
200 name = force.getGlobalParameterName(i)
201 value = force.getGlobalParameterDefaultValue(i)
202 parameters[name] = value
203 per_item_parameter_names = []
204 for i in range(getattr(force, f"getNumPer{item}Parameters")()):
205 per_item_parameter_names.append(
206 getattr(force, f"getPer{item}ParameterName")(i)
207 )
208 for name in per_item_parameter_names:
209 parameters[name] = []
210 atoms = []
211 for i in range(getattr(force, f"getNum{item}s")()):
212 if isinstance(force, openmm.CustomCompoundBondForce):
213 indices, per_item_parameters = force.getBondParameters(i)
214 else:
215 *indices, per_item_parameters = getattr(
216 force,
217 f"get{item}Parameters",
218 )(i)
219 atoms.append(indices)
220 for name, value in zip(per_item_parameter_names, per_item_parameters):
221 parameters[name].append(value)
222 atoms = np.asarray(atoms).reshape(-1, number)
223 return cls(function, unit, atoms, periodicBounds, pbc, **parameters)
225 @classmethod
226 def _fromHarmonicBondForce(
227 cls,
228 force: openmm.HarmonicBondForce,
229 unit: mmunit.Unit,
230 pbc: bool = False,
231 ) -> AtomicFunction:
232 """
233 Create an :class:`AtomicFunction` from a :OpenMM:`HarmonicBondForce`.
234 """
235 parameters = {"r0": [], "k": []}
236 atoms = []
237 for i in range(force.getNumBonds()):
238 *indices, length, k = force.getBondParameters(i)
239 atoms.append(indices)
240 parameters["r0"].append(length)
241 parameters["k"].append(k)
242 return cls(
243 "(k/2)*(distance(p1, p2)-r0)^2",
244 unit,
245 atoms,
246 None,
247 pbc,
248 **parameters,
249 )
251 @classmethod
252 def _fromHarmonicAngleForce(
253 cls,
254 force: openmm.HarmonicAngleForce,
255 unit: mmunit.Unit,
256 pbc: bool = False,
257 ) -> AtomicFunction:
258 """
259 Create an :class:`AtomicFunction` from a :OpenMM:`HarmonicAngleForce`.
260 """
261 parameters = {"theta0": [], "k": []}
262 atoms = []
263 for i in range(force.getNumAngles()):
264 *indices, angle, k = force.getAngleParameters(i)
265 atoms.append(indices)
266 parameters["theta0"].append(angle)
267 parameters["k"].append(k)
268 return cls(
269 "(k/2)*(angle(p1, p2, p3)-theta0)^2",
270 unit,
271 atoms,
272 None,
273 pbc,
274 **parameters,
275 )
277 @classmethod
278 def _fromPeriodicTorsionForce(
279 cls,
280 force: openmm.PeriodicTorsionForce,
281 unit: mmunit.Unit,
282 pbc: bool = False,
283 ) -> AtomicFunction:
284 """
285 Create an :class:`AtomicFunction` from a :OpenMM:`PeriodicTorsionForce`.
286 """
287 parameters = {"periodicity": [], "phase": [], "k": []}
288 atoms = []
289 for i in range(force.getNumTorsions()):
290 *indices, periodicity, phase, k = force.getTorsionParameters(i)
291 atoms.append(indices)
292 parameters["periodicity"].append(periodicity)
293 parameters["phase"].append(phase)
294 parameters["k"].append(k)
295 return cls(
296 "k*(1 + cos(periodicity*dihedral(p1, p2, p3, p4) - phase))",
297 unit,
298 atoms,
299 None,
300 pbc,
301 **parameters,
302 )
304 @classmethod
305 def fromOpenMMForce(
306 cls,
307 force: openmm.Force,
308 unit: mmunit.Unit,
309 periodicBounds: VectorQuantity | None = None,
310 pbc: bool | None = None,
311 ) -> AtomicFunction:
312 """
313 Create an :class:`AtomicFunction` from an :OpenMM:`Force`.
315 Parameters
316 ----------
317 force
318 The force to be converted
319 unit
320 The unit of measurement of the collective variable. It must be
321 compatible with the MD unit system (mass in `daltons`, distance in
322 `nanometers`, time in `picoseconds`, temperature in `kelvin`, energy in
323 `kilojoules_per_mol`, angle in `radians`). If the collective variable
324 does not have a unit, use `unit.dimensionless`.
325 periodicBounds
326 The lower and upper bounds of the collective variable if it is periodic,
327 or ``None`` if it is not. This parameter is considered only if `force` is
328 a custom :OpenMM:`Force`.
329 pbc
330 Whether to use periodic boundary conditions when computing atomic
331 distances. If ``None``, use the setting of `force`.
333 Raises
334 ------
335 TypeError
336 If the force is not convertible to an :class:`AtomicFunction`
337 ValueError
338 If the force uses tabulated functions, or if it is a
339 :OpenMM:`CustomExternalForce` that uses ``periodicdistance``
341 Examples
342 --------
343 >>> import cvpack
344 >>> import numpy as np
345 >>> import openmm
346 >>> from openmm import unit
347 >>> from openmm import app
348 >>> from openmmtools import testsystems
349 >>> model = testsystems.LysozymeImplicit()
350 >>> residues = [r for r in model.topology.residues() if 59 <= r.index <= 79]
351 >>> helix_content = cvpack.HelixTorsionContent(residues)
352 >>> model.system.addForce(helix_content)
353 6
354 >>> num_atoms = model.system.getNumParticles()
355 >>> mean_x = openmm.CustomExternalForce("x/num_atoms")
356 >>> mean_x.addGlobalParameter("num_atoms", num_atoms)
357 0
358 >>> for i in range(num_atoms):
359 ... _ = mean_x.addParticle(i, [])
360 >>> model.system.addForce(mean_x)
361 7
362 >>> forces = {force.getName(): force for force in model.system.getForces()}
363 >>> copies = {
364 ... name: cvpack.AtomicFunction.fromOpenMMForce(
365 ... force, unit.kilojoules_per_mole
366 ... )
367 ... for name, force in forces.items()
368 ... if name in [
369 ... "HarmonicBondForce",
370 ... "HarmonicAngleForce",
371 ... "PeriodicTorsionForce",
372 ... "CustomExternalForce"
373 ... ]
374 ... }
375 >>> copies["HelixTorsionContent"] = cvpack.AtomicFunction.fromOpenMMForce(
376 ... helix_content, unit.dimensionless
377 ... )
378 >>> indices = {}
379 >>> for index, (name, force) in enumerate(copies.items(), start=1):
380 ... _ = model.system.addForce(force)
381 ... force.setForceGroup(index)
382 ... indices[name] = index
383 >>> platform = openmm.Platform.getPlatformByName('Reference')
384 >>> integrator = openmm.VerletIntegrator(0)
385 >>> context = openmm.Context(model.system, integrator, platform)
386 >>> context.setPositions(model.positions)
387 >>> for name in copies:
388 ... state = context.getState(getEnergy=True, groups={indices[name]})
389 ... value = state.getPotentialEnergy() / unit.kilojoules_per_mole
390 ... copy_value = copies[name].getValue(context)
391 ... print(f"{name}: original={value:.6f}, copy={copy_value}")
392 HarmonicBondForce: original=2094.312..., copy=2094.312... kJ/mol
393 HarmonicAngleForce: original=3239.795..., copy=3239.795... kJ/mol
394 PeriodicTorsionForce: original=4226.05..., copy=4226.05... kJ/mol
395 CustomExternalForce: original=5.02155..., copy=5.02155... kJ/mol
396 HelixTorsionContent: original=17.4528..., copy=17.4528... dimensionless
397 """
398 if not isinstance(force, openmm.Force):
399 raise TypeError(f"Force {force} is not convertible to an AtomicFunction")
400 if pbc is None:
401 pbc = force.usesPeriodicBoundaryConditions()
402 if getattr(force, "getNumTabulatedFunctions", lambda: 0)() > 0:
403 raise ValueError("Forces with tabulated functions are not supported")
404 if isinstance(
405 force,
406 (
407 openmm.CustomAngleForce,
408 openmm.CustomBondForce,
409 openmm.CustomCompoundBondForce,
410 openmm.CustomExternalForce,
411 openmm.CustomTorsionForce,
412 ),
413 ):
414 return cls._fromCustomForce(force, unit, periodicBounds, pbc)
415 if isinstance(force, openmm.HarmonicBondForce):
416 return cls._fromHarmonicBondForce(force, unit, pbc)
417 if isinstance(force, openmm.HarmonicAngleForce):
418 return cls._fromHarmonicAngleForce(force, unit, pbc)
419 if isinstance(force, openmm.PeriodicTorsionForce):
420 return cls._fromPeriodicTorsionForce(force, unit, pbc)
421 raise TypeError(f"Force {force} is not convertible to an AtomicFunction")
424AtomicFunction.registerTag("!cvpack.AtomicFunction")