Coverage for cvpack/utils.py: 96%
117 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 00:10 +0000
1"""
2.. class:: Utils
3 :platform: Linux, MacOS, Windows
4 :synopsis: Utility functions and classes for CVPack
6.. classauthor:: Charlles Abreu <craabreu@gmail.com>
8"""
10from __future__ import annotations
12import functools
13import inspect
14import typing as t
16import numpy as np
17import openmm
18from numpy import typing as npt
19from openmm import XmlSerializer
20from openmm import app as mmapp
21from openmm import unit as mmunit
23from .serialization import (
24 Serializable,
25 SerializableAtom,
26 SerializableForce,
27 SerializableResidue,
28)
29from .units import Quantity, Unit, value_in_md_units
32def evaluate_in_context(
33 forces: openmm.Force | t.Iterable[openmm.Force], context: openmm.Context
34) -> float | tuple[float, ...]:
35 """Evaluate the potential energies of OpenMM Forces in a given context.
37 Parameters
38 ----------
39 forces
40 The forces to be evaluated.
41 context
42 The context in which the forces will be evaluated.
44 Returns
45 -------
46 float | tuple[float, ...]
47 The potential energy of a single force, or a tuple with the potential
48 energies of multiple forces, in the given context.
49 """
50 is_single = isinstance(forces, openmm.Force)
51 forces = [forces] if is_single else list(forces)
52 if len(forces) > 32: # noqa: PLR2004
53 raise ValueError("At most 32 forces can be evaluated at once")
54 system = openmm.System()
55 for _ in range(context.getSystem().getNumParticles()):
56 system.addParticle(1.0)
57 for i, force in enumerate(forces):
58 force_copy = XmlSerializer.deserialize(XmlSerializer.serialize(force))
59 force_copy.setForceGroup(i)
60 system.addForce(force_copy)
61 state = context.getState(getPositions=True)
62 new_context = openmm.Context(
63 system, openmm.VerletIntegrator(1.0), context.getPlatform()
64 )
65 new_context.setPositions(state.getPositions())
66 new_context.setPeriodicBoxVectors(*state.getPeriodicBoxVectors())
67 parameters = context.getParameters()
68 for name in new_context.getParameters():
69 if name in parameters:
70 new_context.setParameter(name, parameters[name])
71 energies = []
72 for i in range(len(forces)):
73 state = new_context.getState(getEnergy=True, groups=1 << i)
74 energies.append(value_in_md_units(state.getPotentialEnergy()))
75 return energies[0] if is_single else tuple(energies)
78def convert_to_matrix(array: npt.ArrayLike) -> tuple[np.ndarray, int, int]:
79 """Convert a 1D or 2D array-like object to a 2D numpy array.
81 Parameters
82 ----------
83 array : array_like
84 The array to be converted.
86 Returns
87 -------
88 numpy.ndarray
89 The 2D numpy array.
90 int
91 The number of rows in the array.
92 int
93 The number of columns in the array.
94 """
95 array = np.atleast_2d(array)
96 numrows, numcols, *other_dimensions = array.shape
97 if other_dimensions:
98 raise ValueError("Array-like object cannot have more than two dimensions.")
99 return array, numrows, numcols
102def check_contiguous_residues(residues: t.Sequence) -> None:
103 """
104 Raise a ValueError unless the residues are consecutive members of one chain.
106 Parameters
107 ----------
108 residues
109 The residues to check, as OpenMM or serializable residues.
110 """
112 def chain_index(residue):
113 if hasattr(residue, "_chain_index"):
114 return residue._chain_index
115 return residue.chain.index
117 for previous, current in zip(residues[:-1], residues[1:]):
118 if current.index != previous.index + 1 or chain_index(current) != chain_index(
119 previous
120 ):
121 raise ValueError(
122 "The residues must be a contiguous sequence from a single chain, but "
123 f"{previous.name}{previous.id} is not followed by "
124 f"{current.name}{current.id}"
125 )
128def find_unused_force_group(system: openmm.System) -> int:
129 """
130 Return the lowest force group not used by any force in a system.
132 Parameters
133 ----------
134 system
135 The system to search for unused force groups.
137 Raises
138 ------
139 RuntimeError
140 If all force groups are already in use
141 """
142 used_groups = {force.getForceGroup() for force in system.getForces()}
143 for group in range(32):
144 if group not in used_groups:
145 return group
146 raise RuntimeError("All force groups are already in use.")
149def get_single_force_state(
150 force: openmm.Force,
151 context: openmm.Context,
152 allowReinitialization: bool = False,
153 **kwargs: bool,
154) -> openmm.State:
155 """
156 Get an OpenMM State containing the potential energy and/or force values computed
157 from a single force object.
159 Parameters
160 ----------
161 force
162 The force object from which the state should be extracted.
163 context
164 The context from which the state should be extracted.
165 allowReinitialization
166 If True, the force group of the given force will be temporarily changed to a
167 group that is not used by any other force in the system, if necessary.
169 Keyword Args
170 ------------
171 **kwargs
172 Additional keyword arguments to be passed to the `getState` method, except for
173 the `groups` argument.
175 Returns
176 -------
177 openmm.State
178 The state containing the requested values.
180 Raises
181 ------
182 ValueError
183 If this force is not present in the given context.
184 """
185 forces_and_groups = [
186 (f, f.getForceGroup()) for f in context.getSystem().getForces()
187 ]
188 if not any(f.this == force.this for f, _ in forces_and_groups):
189 raise RuntimeError("This force is not present in the given context.")
190 self_group = force.getForceGroup()
191 other_groups = {g for f, g in forces_and_groups if f.this != force.this}
192 if self_group not in other_groups:
193 return context.getState(groups=1 << self_group, **kwargs)
194 if not allowReinitialization:
195 raise ValueError("Context reinitialization required, but not allowed.")
196 new_group = find_unused_force_group(context.getSystem())
197 force.setForceGroup(new_group)
198 try:
199 context.reinitialize(preserveState=True)
200 return context.getState(groups=1 << new_group, **kwargs)
201 finally:
202 force.setForceGroup(self_group)
203 context.reinitialize(preserveState=True)
206def compute_effective_mass(
207 force: openmm.Force, context: openmm.Context, allowReinitialization: bool = False
208) -> float:
209 r"""
210 Compute the effective mass of an :OpenMM:`Force` at a given :OpenMM:`Context`.
212 Parameters
213 ----------
214 force
215 The force object from which the effective mass should be computed
216 context
217 The context at which the force's effective mass should be evaluated
218 allowReinitialization
219 If True, the force group of the given force will be temporarily changed to a
220 group that is not used by any other force in the system, if necessary.
222 Returns
223 -------
224 float
225 The effective mass of the force at the given context
226 """
227 state = get_single_force_state(
228 force, context, allowReinitialization, getForces=True
229 )
230 get_mass = functools.partial(
231 openmm._openmm.System_getParticleMass, context.getSystem()
232 )
233 force_vectors = state.getForces(asNumpy=True)._value
234 squared_forces = np.sum(np.square(force_vectors), axis=1)
235 nonzeros = np.nonzero(squared_forces)[0]
236 mass_values = np.fromiter(map(get_mass, nonzeros), dtype=np.float64)
237 movable = mass_values > 0 # OpenMM treats massless particles as fixed
238 if not np.any(movable):
239 return np.inf
240 return 1.0 / np.sum(squared_forces[nonzeros][movable] / mass_values[movable])
243def preprocess_args(func: t.Callable) -> t.Callable:
244 """
245 A decorator that converts instances of unserializable classes to their
246 serializable counterparts.
248 Parameters
249 ----------
250 func
251 The function to be decorated.
253 Returns
254 -------
255 The decorated function.
257 Example
258 -------
259 >>> from cvpack import units, utils
260 >>> from openmm import unit as mmunit
261 >>> @utils.preprocess_args
262 ... def function(data):
263 ... return data
264 >>> assert isinstance(function(mmunit.angstrom), units.Unit)
265 >>> assert isinstance(function(5 * mmunit.angstrom), units.Quantity)
266 >>> seq = [mmunit.angstrom, mmunit.nanometer]
267 >>> assert isinstance(function(seq), list)
268 >>> assert all(isinstance(item, units.Unit) for item in function(seq))
269 >>> dct = {"length": 3 * mmunit.angstrom, "time": 2 * mmunit.picosecond}
270 >>> assert isinstance(function(dct), dict)
271 >>> assert all(isinstance(item, units.Quantity) for item in function(dct).values())
272 """
273 signature = inspect.signature(func)
275 def convert(data: t.Any) -> t.Any: # noqa: PLR0911
276 if isinstance(data, np.integer):
277 return int(data)
278 if isinstance(data, np.floating):
279 return float(data)
280 if isinstance(data, mmunit.Quantity):
281 return Quantity(data)
282 if isinstance(data, mmunit.Unit):
283 return Unit(data)
284 if isinstance(data, mmapp.Atom):
285 return SerializableAtom(data)
286 if isinstance(data, mmapp.Residue):
287 return SerializableResidue(data)
288 if isinstance(data, openmm.Force) and not isinstance(data, Serializable):
289 return SerializableForce(data)
290 if isinstance(data, np.ndarray):
291 return convert(data.tolist())
292 if isinstance(data, t.Sequence) and not isinstance(data, str):
293 items = [convert(item) for item in data]
294 return tuple(items) if type(data) is tuple else items
295 if isinstance(data, dict):
296 return type(data)((key, convert(value)) for key, value in data.items())
297 return data
299 @functools.wraps(func)
300 def wrapper(*args, **kwargs):
301 bound = signature.bind(*args, **kwargs)
302 for name, data in bound.arguments.items():
303 bound.arguments[name] = convert(data)
304 return func(*bound.args, **bound.kwargs)
306 return wrapper