Coverage for cvpack/collective_variable.py: 98%
58 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-09 16:14 +0000
1"""
2.. class:: CollectiveVariable
3 :platform: Linux, MacOS, Windows
4 :synopsis: An abstract class with common attributes and method for all CVs
6.. moduleauthor:: Charlles Abreu <craabreu@gmail.com>
8"""
10import typing as t
12import openmm
13import yaml
14from openmm import _openmm as mmswig
15from openmm import unit as mmunit
17from .serialization import Serializable
18from .units import Quantity, ScalarQuantity, Unit
19from .utils import compute_effective_mass, get_single_force_state, preprocess_args
22class CollectiveVariable(Serializable):
23 r"""
24 An abstract class with common attributes and method for all CVs.
25 """
27 _unit: Unit = Unit("dimensionless")
28 _mass_unit: Unit = Unit("dalton")
29 _args: t.Dict[str, t.Any] = {}
30 _periodic_bounds: t.Optional[Quantity] = None
32 def __getstate__(self) -> t.Dict[str, t.Any]:
33 return self._args
35 def __setstate__(self, keywords: t.Dict[str, t.Any]) -> None:
36 self.__init__(**keywords)
38 def __copy__(self) -> "CollectiveVariable":
39 return yaml.safe_load(yaml.safe_dump(self))
41 def __deepcopy__(self, _) -> "CollectiveVariable":
42 return yaml.safe_load(yaml.safe_dump(self))
44 @preprocess_args
45 def _registerCV(
46 self,
47 name: str,
48 cvUnit: Unit,
49 **kwargs: t.Any,
50 ) -> None:
51 """
52 Register the newly created CollectiveVariable subclass instance.
54 This method must always be called from Subclass.__init__.
56 Parameters
57 ----------
58 name
59 The name of this collective variable.
60 unit
61 The unit of measurement of this collective variable. It must be a unit
62 in the MD unit system (mass in Da, distance in nm, time in ps,
63 temperature in K, energy in kJ/mol, angle in rad).
64 kwargs
65 The keyword arguments needed to construct this collective variable
66 """
67 mmswig.Force_setName(self, name)
68 self._unit = cvUnit
69 self._mass_unit = Unit(mmunit.dalton * (mmunit.nanometers / self._unit) ** 2)
70 self._args = {"name": name}
71 self._args.update(kwargs)
73 def _registerPeriodicBounds(
74 self, lower: ScalarQuantity, upper: ScalarQuantity
75 ) -> None:
76 """
77 Register the periodic bounds of this collective variable.
79 This method must called from Subclass.__init__ if the collective variable is
80 periodic.
82 Parameters
83 ----------
84 lower
85 The lower bound of this collective variable
86 upper
87 The upper bound of this collective variable
88 """
89 if mmunit.is_quantity(lower):
90 lower = lower.value_in_unit(self.getUnit())
91 if mmunit.is_quantity(upper):
92 upper = upper.value_in_unit(self.getUnit())
93 self._periodic_bounds = Quantity((lower, upper), self.getUnit())
95 def _setUnusedForceGroup(self, system: openmm.System) -> None:
96 """
97 Set the force group of this collective variable to the one at a given position
98 in the ascending ordered list of unused force groups in an :OpenMM:`System`.
100 .. note::
102 Evaluating a collective variable (see :meth:`getValue`) or computing its
103 effective mass (see :meth:`getEffectiveMass`) is more efficient when the
104 collective variable is the only force in its own force group.
106 Parameters
107 ----------
108 system
109 The system to search for unused force groups
111 Raises
112 ------
113 RuntimeError
114 If all force groups are already in use
115 """
116 used_groups = {force.getForceGroup() for force in system.getForces()}
117 new_group = next(filter(lambda i: i not in used_groups, range(32)), None)
118 if new_group is None:
119 raise RuntimeError("All force groups are already in use.")
120 mmswig.Force_setForceGroup(self, new_group)
122 def getName(self) -> str: # pylint: disable=useless-parent-delegation
123 """
124 Get the name of this collective variable.
125 """
126 return mmswig.Force_getName(self)
128 def getUnit(self) -> Unit:
129 """
130 Get the unit of measurement of this collective variable.
131 """
132 return self._unit
134 def getMassUnit(self) -> Unit:
135 """
136 Get the unit of measurement of the effective mass of this collective variable.
137 """
138 return self._mass_unit
140 def getPeriodicBounds(self) -> t.Union[Quantity, None]:
141 """
142 Get the periodic bounds of this collective variable.
144 Returns
145 -------
146 Union[Quantity, None]
147 The periodic bounds of this collective variable or None if it is not
148 periodic
149 """
150 return self._periodic_bounds
152 def addToSystem(
153 self, system: openmm.System, setUnusedForceGroup: bool = True
154 ) -> None:
155 """
156 Add this collective variable to an :OpenMM:`System`.
158 Parameters
159 ----------
160 system
161 The system to which this collective variable should be added
162 setUnusedForceGroup
163 If True, the force group of this collective variable will be set to the
164 first available force group in the system
165 """
166 if setUnusedForceGroup:
167 self._setUnusedForceGroup(system)
168 system.addForce(self)
170 def getValue(
171 self, context: openmm.Context, allowReinitialization: bool = False
172 ) -> Quantity:
173 """
174 Evaluate this collective variable at a given :OpenMM:`Context`.
176 If this collective variable share the force group with other forces, then its
177 evaluation requires reinitializing the :OpenMM:`Context` twice at each call.
178 This is inefficient and should be avoided. To allow this behavior, the
179 ``allowReinitialization`` parameter must be set to True.
181 .. note::
183 By adding this collective variable to the system using the
184 :meth:`addToSystem` method, the force group of this collective variable
185 is set to an available force group in the system by default.
187 Parameters
188 ----------
189 context
190 The context at which this collective variable should be evaluated.
191 allowReinitialization
192 If True, the context will be reinitialized if necessary.
194 Returns
195 -------
196 Quantity
197 The value of this collective variable at the given context.
199 Example
200 -------
201 In this example, we compute the values of the backbone dihedral angles and
202 the radius of gyration of an alanine dipeptide molecule in water:
204 >>> import cvpack
205 >>> import openmm
206 >>> from openmmtools import testsystems
207 >>> model = testsystems.AlanineDipeptideExplicit()
208 >>> top = model.mdtraj_topology
209 >>> backbone_atoms = top.select("name N C CA and resid 1 2")
210 >>> phi = cvpack.Torsion(*backbone_atoms[0:4])
211 >>> psi = cvpack.Torsion(*backbone_atoms[1:5])
212 >>> radius_of_gyration = cvpack.RadiusOfGyration(
213 ... top.select("not water")
214 ... )
215 >>> for cv in [phi, psi, radius_of_gyration]:
216 ... cv.addToSystem(model.system)
217 >>> context = openmm.Context(
218 ... model.system, openmm.VerletIntegrator(0)
219 ... )
220 >>> context.setPositions(model.positions)
221 >>> phi.getValue(context)
222 3.1415... rad
223 >>> psi.getValue(context)
224 3.1415... rad
225 >>> radius_of_gyration.getValue(context)
226 0.29514... nm
227 """
228 state = get_single_force_state(
229 self,
230 context,
231 allowReinitialization,
232 getEnergy=True,
233 )
234 value = state.getPotentialEnergy().value_in_unit(mmunit.kilojoules_per_mole)
235 return Quantity(value, self.getUnit())
237 def getEffectiveMass(
238 self, context: openmm.Context, allowReinitialization: bool = False
239 ) -> Quantity:
240 r"""
241 Compute the effective mass of this collective variable at a given
242 :OpenMM:`Context`.
244 The effective mass of a collective variable :math:`q({\bf r})` is defined as
245 :cite:`Chipot_2007`:
247 .. math::
249 m_\mathrm{eff}({\bf r}) = \left(
250 \sum_{i=1}^N \frac{1}{m_i} \left\| \frac{dq}{d{\bf r}_i} \right\|^2
251 \right)^{-1}
253 If this collective variable share the force group with other forces, then
254 evaluating its effective mass requires reinitializing the :OpenMM:`Context`
255 twice at each call. This is inefficient and should be avoided. To allow this
256 behavior, the ``allowReinitialization`` parameter must be set to True.
258 .. note::
260 By adding this collective variable to the system using the
261 :meth:`addToSystem` method, the force group of this collective variable
262 is set to an available force group in the system by default.
264 Parameters
265 ----------
266 context
267 The context at which this collective variable's effective mass should be
268 evaluated.
269 allowReinitialization
270 If True, the context will be reinitialized if necessary.
272 Returns
273 -------
274 Quantity
275 The effective mass of this collective variable at the given context.
277 Example
278 -------
279 In this example, we compute the effective masses of the backbone dihedral
280 angles and the radius of gyration of an alanine dipeptide molecule in water:
282 >>> import cvpack
283 >>> import openmm
284 >>> from openmmtools import testsystems
285 >>> model = testsystems.AlanineDipeptideExplicit()
286 >>> top = model.mdtraj_topology
287 >>> backbone_atoms = top.select("name N C CA and resid 1 2")
288 >>> phi = cvpack.Torsion(*backbone_atoms[0:4])
289 >>> psi = cvpack.Torsion(*backbone_atoms[1:5])
290 >>> radius_of_gyration = cvpack.RadiusOfGyration(
291 ... top.select("not water")
292 ... )
293 >>> for cv in [phi, psi, radius_of_gyration]:
294 ... cv.addToSystem(model.system)
295 >>> context = openmm.Context(
296 ... model.system, openmm.VerletIntegrator(0)
297 ... )
298 >>> context.setPositions(model.positions)
299 >>> phi.getEffectiveMass(context)
300 0.05119... nm**2 Da/(rad**2)
301 >>> psi.getEffectiveMass(context)
302 0.05186... nm**2 Da/(rad**2)
303 >>> radius_of_gyration.getEffectiveMass(context)
304 30.946... Da
305 """
306 return Quantity(
307 compute_effective_mass(self, context, allowReinitialization),
308 self._mass_unit,
309 )