Coverage for cvpack/collective_variable.py: 100%
59 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 00:10 +0000
1"""
2.. class:: CollectiveVariable
3 :platform: Linux, MacOS, Windows
4 :synopsis: An abstract class with common attributes and methods for all CVs
6.. moduleauthor:: Charlles Abreu <craabreu@gmail.com>
8"""
10from __future__ import annotations
12import typing as t
14import openmm
15import yaml
16from openmm import _openmm as mmswig
17from openmm import unit as mmunit
19from .serialization import Serializable
20from .units import Quantity, ScalarQuantity, Unit
21from .utils import (
22 compute_effective_mass,
23 find_unused_force_group,
24 get_single_force_state,
25 preprocess_args,
26)
29class CollectiveVariable(Serializable):
30 r"""
31 An abstract class with common attributes and methods for all CVs.
32 """
34 _unit: Unit = Unit("dimensionless")
35 _mass_unit: Unit = Unit("dalton")
36 _args: dict[str, t.Any] = {}
37 _periodic_bounds: Quantity | None = None
39 def __getstate__(self) -> dict[str, t.Any]:
40 # The force group is not serialized: addToSystem assigns it
41 return {**self._args, "name": self.getName()}
43 def __setstate__(self, keywords: dict[str, t.Any]) -> None:
44 self.__init__(**keywords)
46 def __copy__(self) -> CollectiveVariable:
47 return yaml.safe_load(yaml.safe_dump(self))
49 def __deepcopy__(self, _) -> CollectiveVariable:
50 return yaml.safe_load(yaml.safe_dump(self))
52 @preprocess_args
53 def _registerCV(
54 self,
55 name: str,
56 cvUnit: Unit,
57 **kwargs: t.Any,
58 ) -> None:
59 """
60 Register the newly created CollectiveVariable subclass instance.
62 This method must always be called from Subclass.__init__.
64 Parameters
65 ----------
66 name
67 The name of this collective variable.
68 cvUnit
69 The unit of measurement of this collective variable. It must be a unit
70 in the MD unit system (mass in Da, distance in nm, time in ps,
71 temperature in K, energy in kJ/mol, angle in rad).
72 kwargs
73 The keyword arguments needed to construct this collective variable
75 Raises
76 ------
77 ValueError
78 If the unit is not compatible with the MD unit system
79 """
80 if (1 * cvUnit).value_in_unit_system(mmunit.md_unit_system) != 1:
81 raise ValueError(
82 f"Unit {cvUnit} is not compatible with the MD unit system."
83 )
84 mmswig.Force_setName(self, name)
85 self._unit = cvUnit
86 self._mass_unit = Unit(mmunit.dalton * (mmunit.nanometers / self._unit) ** 2)
87 self._args = {"name": name}
88 self._args.update(kwargs)
90 def _registerPeriodicBounds(
91 self, lower: ScalarQuantity, upper: ScalarQuantity
92 ) -> None:
93 """
94 Register the periodic bounds of this collective variable.
96 This method must be called from Subclass.__init__ if the collective variable is
97 periodic.
99 Parameters
100 ----------
101 lower
102 The lower bound of this collective variable
103 upper
104 The upper bound of this collective variable
105 """
106 if mmunit.is_quantity(lower):
107 lower = lower.value_in_unit(self.getUnit())
108 if mmunit.is_quantity(upper):
109 upper = upper.value_in_unit(self.getUnit())
110 self._periodic_bounds = Quantity((lower, upper), self.getUnit())
112 def _setUnusedForceGroup(self, system: openmm.System) -> int:
113 """
114 Set the force group of this collective variable to the lowest force group not
115 used by any force in an :OpenMM:`System`.
117 .. note::
119 Evaluating a collective variable (see :meth:`getValue`) or computing its
120 effective mass (see :meth:`getEffectiveMass`) is more efficient when the
121 collective variable is the only force in its own force group.
123 Parameters
124 ----------
125 system
126 The system to search for unused force groups
128 Returns
129 -------
130 int
131 The force group assigned to this collective variable
133 Raises
134 ------
135 RuntimeError
136 If all force groups are already in use
137 """
138 new_group = find_unused_force_group(system)
139 mmswig.Force_setForceGroup(self, new_group)
140 return new_group
142 def getName(self) -> str:
143 """
144 Get the name of this collective variable.
145 """
146 return mmswig.Force_getName(self)
148 def getUnit(self) -> Unit:
149 """
150 Get the unit of measurement of this collective variable.
151 """
152 return self._unit
154 def getMassUnit(self) -> Unit:
155 """
156 Get the unit of measurement of the effective mass of this collective variable.
157 """
158 return self._mass_unit
160 def getPeriodicBounds(self) -> Quantity | None:
161 """
162 Get the periodic bounds of this collective variable.
164 Returns
165 -------
166 Union[Quantity, None]
167 The periodic bounds of this collective variable or None if it is not
168 periodic
169 """
170 return self._periodic_bounds
172 def addToSystem(
173 self, system: openmm.System, setUnusedForceGroup: bool = True
174 ) -> None:
175 """
176 Add this collective variable to an :OpenMM:`System`.
178 Parameters
179 ----------
180 system
181 The system to which this collective variable should be added
182 setUnusedForceGroup
183 If True, the force group of this collective variable will be set to the
184 first available force group in the system
185 """
186 if setUnusedForceGroup:
187 self._setUnusedForceGroup(system)
188 system.addForce(self)
190 def getValue(
191 self, context: openmm.Context, allowReinitialization: bool = False
192 ) -> Quantity:
193 """
194 Evaluate this collective variable at a given :OpenMM:`Context`.
196 If this collective variable shares the force group with other forces, then its
197 evaluation requires reinitializing the :OpenMM:`Context` twice at each call.
198 This is inefficient and should be avoided. To allow this behavior, the
199 ``allowReinitialization`` parameter must be set to True.
201 .. note::
203 By adding this collective variable to the system using the
204 :meth:`addToSystem` method, the force group of this collective variable
205 is set to an available force group in the system by default.
207 Parameters
208 ----------
209 context
210 The context at which this collective variable should be evaluated.
211 allowReinitialization
212 If True, the context will be reinitialized if necessary.
214 Returns
215 -------
216 Quantity
217 The value of this collective variable at the given context.
219 Example
220 -------
221 In this example, we compute the values of the backbone dihedral angles and
222 the radius of gyration of an alanine dipeptide molecule in water:
224 >>> import cvpack
225 >>> import openmm
226 >>> from openmmtools import testsystems
227 >>> model = testsystems.AlanineDipeptideExplicit()
228 >>> top = model.mdtraj_topology
229 >>> backbone_atoms = top.select("not water and name N CA C")
230 >>> phi = cvpack.Torsion(*backbone_atoms[0:4])
231 >>> psi = cvpack.Torsion(*backbone_atoms[1:5])
232 >>> radius_of_gyration = cvpack.RadiusOfGyration(
233 ... top.select("not water")
234 ... )
235 >>> for cv in [phi, psi, radius_of_gyration]:
236 ... cv.addToSystem(model.system)
237 >>> context = openmm.Context(
238 ... model.system, openmm.VerletIntegrator(0)
239 ... )
240 >>> context.setPositions(model.positions)
241 >>> phi.getValue(context)
242 -3.1415... rad
243 >>> psi.getValue(context)
244 3.1415... rad
245 >>> radius_of_gyration.getValue(context)
246 0.29514... nm
247 """
248 state = get_single_force_state(
249 self,
250 context,
251 allowReinitialization,
252 getEnergy=True,
253 )
254 value = state.getPotentialEnergy().value_in_unit(mmunit.kilojoules_per_mole)
255 return Quantity(value, self.getUnit())
257 def getEffectiveMass(
258 self, context: openmm.Context, allowReinitialization: bool = False
259 ) -> Quantity:
260 r"""
261 Compute the effective mass of this collective variable at a given
262 :OpenMM:`Context`.
264 The effective mass of a collective variable :math:`q({\bf r})` is defined as
265 :cite:`Chipot_2007`:
267 .. math::
269 m_\mathrm{eff}({\bf r}) = \left(
270 \sum_{i=1}^N \frac{1}{m_i} \left\| \frac{dq}{d{\bf r}_i} \right\|^2
271 \right)^{-1}
273 If this collective variable shares the force group with other forces, then
274 evaluating its effective mass requires reinitializing the :OpenMM:`Context`
275 twice at each call. This is inefficient and should be avoided. To allow this
276 behavior, the ``allowReinitialization`` parameter must be set to True.
278 .. note::
280 By adding this collective variable to the system using the
281 :meth:`addToSystem` method, the force group of this collective variable
282 is set to an available force group in the system by default.
284 Parameters
285 ----------
286 context
287 The context at which this collective variable's effective mass should be
288 evaluated.
289 allowReinitialization
290 If True, the context will be reinitialized if necessary.
292 Returns
293 -------
294 Quantity
295 The effective mass of this collective variable at the given context.
297 Example
298 -------
299 In this example, we compute the effective masses of the backbone dihedral
300 angles and the radius of gyration of an alanine dipeptide molecule in water:
302 >>> import cvpack
303 >>> import openmm
304 >>> from openmmtools import testsystems
305 >>> model = testsystems.AlanineDipeptideExplicit()
306 >>> top = model.mdtraj_topology
307 >>> backbone_atoms = top.select("not water and name N CA C")
308 >>> phi = cvpack.Torsion(*backbone_atoms[0:4])
309 >>> psi = cvpack.Torsion(*backbone_atoms[1:5])
310 >>> radius_of_gyration = cvpack.RadiusOfGyration(
311 ... top.select("not water")
312 ... )
313 >>> for cv in [phi, psi, radius_of_gyration]:
314 ... cv.addToSystem(model.system)
315 >>> context = openmm.Context(
316 ... model.system, openmm.VerletIntegrator(0)
317 ... )
318 >>> context.setPositions(model.positions)
319 >>> phi.getEffectiveMass(context)
320 0.04795... nm**2 Da/(rad**2)
321 >>> psi.getEffectiveMass(context)
322 0.05119... nm**2 Da/(rad**2)
323 >>> radius_of_gyration.getEffectiveMass(context)
324 30.946... Da
325 """
326 return Quantity(
327 compute_effective_mass(self, context, allowReinitialization),
328 self._mass_unit,
329 )