Coverage for openxps/dynamical_variable.py: 99%
75 statements
« prev ^ index » next coverage.py v7.11.3, created at 2025-11-13 22:08 +0000
1"""
2.. module:: openxps.dynamical_variable
3 :platform: Linux, Windows, macOS
4 :synopsis: Dynamical variables for XPS simulations
6.. moduleauthor:: Charlles Abreu <craabreu@gmail.com>
8"""
10import typing as t
11from dataclasses import dataclass
13import cvpack
14import openmm as mm
15import typing_extensions as te
16from cvpack.serialization import Serializable
17from cvpack.units import Quantity
18from openmm import unit as mmunit
20from .bounds import Bounds, NoBounds, PeriodicBounds
21from .utils import preprocess_args
24@dataclass(frozen=True)
25class DynamicalVariable(Serializable):
26 """
27 Dynamical variable for extended phase-space simulations with OpenMM.
29 Parameters
30 ----------
31 name
32 The name of the context parameter to be turned into a dynamical variable.
33 unit
34 The unity of measurement of this dynamical variable. If it does not have a
35 unit, use ``openmm.unit.dimensionless``.
36 mass
37 The mass assigned to this dynamical variable, whose unit of measurement
38 must be compatible with ``dalton*(nanometer/unit)**2``, where ``unit`` is the
39 dynamical variable's own unit (see ``unit`` above).
40 bounds
41 The boundary condition to be applied to this dynamical variable. It must
42 be an instance of :class:`PeriodicBounds`, :class:`ReflectiveBounds`, or
43 :class:`NoBounds` (for unbounded variables). If it is not
44 :class:`~openxps.NoBounds`, its unit of measurement must be compatible with
45 the dynamical variable's own unit.
47 Example
48 -------
49 >>> import openxps as xps
50 >>> import yaml
51 >>> from openmm import unit
52 >>> dv = xps.DynamicalVariable(
53 ... "psi",
54 ... unit.radian,
55 ... 3 * unit.dalton*(unit.nanometer/unit.radian)**2,
56 ... xps.PeriodicBounds(-180, 180, unit.degree)
57 ... )
58 >>> dv
59 DynamicalVariable(name='psi', unit=rad, mass=3 nm**2 Da/(rad**2), bounds=...)
60 >>> dv.bounds
61 PeriodicBounds(lower=-3.14159..., upper=3.14159..., unit=rad)
62 >>> assert yaml.safe_load(yaml.safe_dump(dv)) == dv
63 """
65 name: str
66 unit: mmunit.Unit
67 mass: mmunit.Quantity
68 bounds: Bounds
70 def __post_init__(self) -> None:
71 if not mmunit.is_unit(self.unit):
72 raise ValueError("The unit must be a valid OpenMM unit.")
73 if not mmunit.is_quantity(self.mass):
74 raise TypeError("Mass must be have units of measurement.")
75 mass_unit = mmunit.dalton * (mmunit.nanometer / self.unit) ** 2
76 if not self.mass.unit.is_compatible(mass_unit):
77 raise TypeError(f"Mass units must be compatible with {mass_unit}.")
78 object.__setattr__(self, "mass", Quantity(self.mass.in_units_of(mass_unit)))
80 if isinstance(self.bounds, NoBounds):
81 return
83 if isinstance(self.bounds, Bounds):
84 if not self.bounds.unit.is_compatible(self.unit):
85 raise ValueError("Provided bounds have incompatible units.")
86 object.__setattr__(self, "bounds", self.bounds.convert(self.unit))
87 else:
88 raise TypeError("The bounds must be an instance of Bounds.")
90 def __getstate__(self) -> dict[str, t.Any]:
91 return {
92 "name": self.name,
93 "unit": self.unit,
94 "mass": self.mass,
95 "bounds": self.bounds,
96 }
98 def __setstate__(self, keywords: dict[str, t.Any]) -> None:
99 self.__init__(**keywords)
101 def in_md_units(self) -> "DynamicalVariable":
102 """
103 Return the dynamical variable in the MD unit system.
104 """
105 unit = self.unit.in_unit_system(mmunit.md_unit_system)
106 mass_unit = mmunit.dalton * (mmunit.nanometer / self.unit) ** 2
107 return DynamicalVariable(
108 self.name,
109 unit,
110 self.mass.in_units_of(mass_unit),
111 self.bounds.in_md_units(),
112 )
114 def isPeriodic(self) -> bool:
115 """
116 Returns whether this dynamical variable is periodic.
118 Returns
119 -------
120 bool
121 ``True`` if this dynamical variable is periodic, ``False`` otherwise.
123 Example
124 -------
125 >>> import openxps as xps
126 >>> from openmm import unit
127 >>> dv = xps.DynamicalVariable(
128 ... "psi",
129 ... unit.radian,
130 ... 3 * unit.dalton*(unit.nanometer/unit.radian)**2,
131 ... xps.PeriodicBounds(-180, 180, unit.degree)
132 ... )
133 >>> dv.isPeriodic()
134 True
135 """
136 return isinstance(self.bounds, PeriodicBounds)
138 def createCollectiveVariable(
139 self, index: int, name: t.Optional[str] = None
140 ) -> cvpack.OpenMMForceWrapper:
141 """
142 Returns a :CVPack:`OpenMMForceWrapper` object associating this extra degree of
143 freedom with the x coordinate of a particle in an OpenMM system.
145 Parameters
146 ----------
147 index
148 The index of the particle whose x coordinate will be associated with this
149 dynamical variable.
150 name
151 The name of the context parameter to be used in the OpenMM system. If
152 ``None``, the name of the dynamical variable will be used.
154 Returns
155 -------
156 cvpack.OpenMMForceWrapper
157 The collective variable object representing this dynamical variable.
159 Example
160 -------
161 >>> import openxps as xps
162 >>> from openmm import unit
163 >>> dv = xps.DynamicalVariable(
164 ... "psi",
165 ... unit.radian,
166 ... 3 * unit.dalton*(unit.nanometer/unit.radian)**2,
167 ... xps.PeriodicBounds(-180, 180, unit.degree)
168 ... )
169 >>> cv = dv.createCollectiveVariable(0)
170 """
171 bounds = self.bounds
172 if bounds is None:
173 force = mm.CustomExternalForce("x")
174 else:
175 force = mm.CustomExternalForce(bounds.leptonExpression("x"))
176 bounds = bounds.asQuantity() if self.isPeriodic() else None
177 force.addParticle(index, [])
178 return cvpack.OpenMMForceWrapper(
179 force,
180 self.unit,
181 bounds,
182 self.name if name is None else name,
183 )
185 def _distanceToCV(self, other: cvpack.CollectiveVariable) -> str:
186 name = other.getName()
187 diff = f"{name}-{self.name}"
188 if other.getPeriodicBounds() is None and not self.isPeriodic():
189 return f"({diff})"
190 if self.isPeriodic() and other.getPeriodicBounds() == self.bounds.asQuantity():
191 period = self.bounds.period
192 return f"({diff}-{period}*floor(0.5+({diff})/{period}))"
193 raise ValueError("Incompatible boundary conditions.")
195 def _distanceToDV(self, other: te.Self) -> str:
196 diff = f"{other.name}-{self.name}"
197 if not (self.isPeriodic() or other.isPeriodic()):
198 return f"({diff})"
199 if other.isPeriodic() and self.isPeriodic():
200 period = self.bounds.period
201 return f"({diff}-{period}*floor(0.5+({diff})/{period}))"
202 raise ValueError("Incompatible boundary conditions.")
204 def distanceTo(self, other: cvpack.CollectiveVariable | te.Self) -> str:
205 """
206 Returns a Lepton expression representing the distance between this extra degree
207 of freedom and another variable.
209 Parameters
210 ----------
211 other
212 The other variable to which the distance will be calculated.
214 Returns
215 -------
216 str
217 A string representing the distance between the two variables.
219 Example
220 -------
221 >>> import openxps as xps
222 >>> import pytest
223 >>> from openmm import unit
224 >>> dv = xps.DynamicalVariable(
225 ... "psi0",
226 ... unit.radian,
227 ... 3 * unit.dalton*(unit.nanometer/unit.radian)**2,
228 ... xps.PeriodicBounds(-180, 180, unit.degree)
229 ... )
230 >>> psi = cvpack.Torsion(6, 8, 14, 16, name="psi")
231 >>> dv.distanceTo(psi)
232 '(psi-psi0-6.28318...*floor(0.5+(psi-psi0)/6.28318...))'
233 >>>
235 """
236 if isinstance(other, cvpack.CollectiveVariable):
237 return self._distanceToCV(other)
238 if type(other) is type(self):
239 return self._distanceToDV(other)
240 raise TypeError(f"Method distanceTo not implemented for type {type(other)}.")
243DynamicalVariable.__init__ = preprocess_args(DynamicalVariable.__init__)
245DynamicalVariable.registerTag("!openxps.DynamicalVariable")