Coverage for cvpack/units/units.py: 95%
143 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 00:10 +0000
1"""
2.. module:: units
3 :platform: Linux, MacOS, Windows
4 :synopsis: Units of measurement for CVPack.
6.. classauthor:: Charlles Abreu <craabreu@gmail.com>
8"""
10import ast
11import functools
12import operator
13import typing as t
14from numbers import Real
16import numpy as np
17import openmm
18from openmm import unit as mmunit
20from ..serialization import Serializable
22ScalarQuantity = mmunit.Quantity | Real
23VectorQuantity = mmunit.Quantity | np.ndarray | openmm.Vec3
24MatrixQuantity = (
25 mmunit.Quantity | np.ndarray | t.Sequence[openmm.Vec3] | t.Sequence[np.ndarray]
26)
29def _to_builtin(value: t.Any) -> t.Any:
30 """Convert numpy and OpenMM containers to plain Python lists and scalars."""
31 if isinstance(value, np.ndarray | np.generic):
32 return value.tolist()
33 if isinstance(value, openmm.Vec3):
34 return [_to_builtin(item) for item in value]
35 if type(value) in (list, tuple):
36 return type(value)(_to_builtin(item) for item in value)
37 return value
40_UNIT_OPERATORS = {
41 ast.Mult: operator.mul,
42 ast.Div: operator.truediv,
43 ast.Pow: operator.pow,
44}
47@functools.cache
48def _unit_names() -> dict[t.Any, str]:
49 """Map each OpenMM base or scaled unit to an ``openmm.unit`` attribute name."""
50 names = {}
51 for attribute in dir(mmunit):
52 value = getattr(mmunit, attribute)
53 if isinstance(value, mmunit.Unit):
54 components = list(value.iter_base_or_scaled_units())
55 if len(components) == 1 and components[0][1] == 1:
56 component = components[0][0]
57 # Prefer the attribute named like the unit itself (e.g. kilojoule)
58 if component not in names or attribute == component.name:
59 names[component] = attribute
60 return names
63def _evaluate_unit_node(node: ast.AST) -> t.Any: # noqa: PLR0911
64 """Evaluate a parsed unit expression, allowing only unit names and arithmetic."""
65 if isinstance(node, ast.Expression):
66 return _evaluate_unit_node(node.body)
67 if isinstance(node, ast.Name):
68 value = getattr(mmunit, node.id, None)
69 if isinstance(value, mmunit.Unit):
70 return value
71 raise ValueError(f"Unknown unit: {node.id}")
72 if isinstance(node, ast.Constant) and type(node.value) in (int, float):
73 return node.value
74 if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub | ast.UAdd):
75 operand = _evaluate_unit_node(node.operand)
76 if isinstance(operand, int | float):
77 return -operand if isinstance(node.op, ast.USub) else operand
78 if isinstance(node, ast.BinOp) and type(node.op) in _UNIT_OPERATORS:
79 left = _evaluate_unit_node(node.left)
80 right = _evaluate_unit_node(node.right)
81 if isinstance(node.op, ast.Pow) and not (
82 isinstance(left, mmunit.Unit) and isinstance(right, int | float)
83 ):
84 raise ValueError("Only units can be raised to numeric powers")
85 return _UNIT_OPERATORS[type(node.op)](left, right)
86 raise ValueError(f"Invalid unit expression: {ast.unparse(node)}")
89def _parse_unit(expression: str) -> mmunit.Unit:
90 """Safely parse a unit expression such as ``kilojoule/(mole*radian**2)``."""
91 text = f"1{expression}" if expression.startswith("/") else expression
92 try:
93 tree = ast.parse(text, mode="eval")
94 except SyntaxError as error:
95 raise ValueError(f"Invalid unit expression: {expression}") from error
96 result = _evaluate_unit_node(tree)
97 if mmunit.is_quantity(result) and result.value_in_unit(result.unit) == 1:
98 result = result.unit
99 if not isinstance(result, mmunit.Unit):
100 raise ValueError(f"Invalid unit expression: {expression}")
101 return result
104class Unit(mmunit.Unit, Serializable):
105 """
106 Extension of the OpenMM Unit class to allow serialization and deserialization.
108 Parameters
109 ----------
110 data
111 The data to be used to create the unit.
112 """
114 def __init__(self, data: str | mmunit.Unit | dict) -> None:
115 if isinstance(data, str):
116 data = _parse_unit(data)
117 if isinstance(data, mmunit.Unit):
118 data = dict(data.iter_base_or_scaled_units())
119 super().__init__(data)
121 def __repr__(self) -> str:
122 return self.get_symbol()
124 def __getstate__(self) -> dict[str, str]:
125 names = _unit_names()
126 factors = []
127 for component, exponent in self.iter_base_or_scaled_units():
128 if component not in names:
129 return {"data": str(self)}
130 power = int(exponent) if float(exponent).is_integer() else exponent
131 name = names[component]
132 factors.append(name if power == 1 else f"{name}**{power}")
133 return {"data": "*".join(factors) or "dimensionless"}
135 def __setstate__(self, keywords: dict[str, str]) -> None:
136 self.__init__(keywords["data"])
139Unit.registerTag("!cvpack.Unit")
142class Quantity(mmunit.Quantity, Serializable):
143 """
144 Extension of the OpenMM Quantity class to allow serialization and deserialization.
145 """
147 def __init__(self, *args: t.Any) -> None:
148 if len(args) == 1 and mmunit.is_quantity(args[0]):
149 super().__init__(args[0].value_in_unit(args[0].unit), Unit(args[0].unit))
150 else:
151 super().__init__(*args)
153 def __repr__(self):
154 return str(self)
156 def __getstate__(self) -> dict[str, t.Any]:
157 return {
158 "value": _to_builtin(self._value),
159 "unit": Unit(self.unit).__getstate__()["data"],
160 }
162 def __setstate__(self, keywords: dict[str, t.Any]) -> None:
163 self.__init__(keywords["value"], Unit(keywords["unit"]))
165 @property
166 def value(self) -> t.Any:
167 """The value of the quantity."""
168 return self._value
170 def in_md_units(self) -> t.Any:
171 """The value of the quantity in MD units."""
172 return self.value_in_unit_system(mmunit.md_unit_system)
175Quantity.registerTag("!cvpack.Quantity")
178_MD_BASE_UNITS = {
179 "mass": "dalton",
180 "length": "nanometer",
181 "time": "picosecond",
182 "charge": "elementary_charge",
183 "temperature": "kelvin",
184 "amount": "mole",
185 "angle": "radian",
186}
189def _md_exponents(unit: mmunit.Unit) -> dict[str, float]:
190 """Exponents of a unit in the base units of OpenMM's MD unit system."""
191 exponents = {}
192 for base, exponent in unit.iter_all_base_units():
193 name = _MD_BASE_UNITS[base.dimension.name]
194 exponents[name] = exponents.get(name, 0) + exponent
195 if "dalton" in exponents:
196 # openmm.unit defines the dalton as gram/mole
197 exponents["mole"] = exponents.get("mole", 0) + exponents["dalton"]
198 return {name: value for name, value in exponents.items() if value != 0}
201def _md_unit(exponents: dict[str, float]) -> mmunit.Unit:
202 """Build a unit from MD exponents, writing Da nm^2/ps^2 as kJ/mol if possible."""
203 exponents = dict(exponents)
204 result = mmunit.dimensionless
205 energy = exponents.get("dalton", 0)
206 if energy != 0 and exponents.get("picosecond", 0) == -2 * energy:
207 # The time dimension comes only from energies, so express them in kJ/mol
208 result = result * mmunit.kilojoule_per_mole**energy
209 exponents.pop("dalton")
210 exponents.pop("picosecond")
211 exponents["nanometer"] = exponents.get("nanometer", 0) - 2 * energy
212 for name, exponent in exponents.items():
213 if exponent != 0:
214 result = result * getattr(mmunit, name) ** exponent
215 return result
218def in_md_units(
219 quantity: ScalarQuantity | VectorQuantity | MatrixQuantity,
220) -> Quantity:
221 """
222 Return a quantity in the MD unit system (e.g. mass in Da, distance in
223 nm, time in ps, temperature in K, energy in kJ/mol, angle in rad).
225 As in kJ/mol, amounts of substance follow OpenMM's per-mole convention: a quantity
226 defined per particle, such as a pressure, is multiplied by Avogadro's number
227 (e.g. 1 bar becomes 0.0602 kJ/(mol nm^3)). Energies are written in kJ/mol.
229 Parameters
230 ----------
231 quantity
232 The quantity to be converted.
234 Returns
235 -------
236 Quantity
237 The quantity in the MD unit system.
238 """
239 if not mmunit.is_quantity(quantity):
240 return quantity * Unit("dimensionless")
241 moles = _md_exponents(quantity.unit).get("mole", 0)
242 if moles != 0:
243 quantity = quantity * mmunit.AVOGADRO_CONSTANT_NA**moles
244 exponents = _md_exponents(quantity.unit)
245 exponents.pop("mole", None)
246 unit_in_md_system = _md_unit(exponents)
247 if 1 * quantity.unit / unit_in_md_system == 1:
248 # Keep the original spelling of a unit that is already in the MD system
249 return Quantity(quantity)
250 return Quantity(quantity.in_units_of(unit_in_md_system))
253def value_in_md_units(
254 quantity: ScalarQuantity | VectorQuantity | MatrixQuantity,
255) -> t.Any:
256 """
257 Return the value of a quantity in the MD unit system (e.g. mass in Da, distance in
258 nm, time in ps, temperature in K, energy in kJ/mol, angle in rad).
260 Parameters
261 ----------
262 quantity
263 The quantity to be converted.
265 Returns
266 -------
267 Any
268 The value of the quantity in the MD unit system.
270 """
271 if mmunit.is_quantity(quantity):
272 return in_md_units(quantity).value
273 return quantity