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