Coverage for ufedmm/ufedmm.py: 73%
520 statements
« prev ^ index » next coverage.py v7.4.1, created at 2024-02-02 04:20 +0000
1"""
2.. module:: ufedmm
3 :platform: Unix, Windows
4 :synopsis: Unified Free Energy Dynamics with OpenMM
6.. moduleauthor:: Charlles Abreu <abreu@eq.ufrj.br>
7"""
9import functools
10from copy import deepcopy
12import numpy as np
13import openmm
14from openmm import app, unit
16import ufedmm
19def _standardized(quantity):
20 """Returns the numerical value of a quantity in a unit of measurement compatible
21 with the Molecular Dynamics unit system (mass in Da, distance in nm, time in ps,
22 temperature in K, energy in kJ/mol, angle in rad)."""
23 if unit.is_quantity(quantity):
24 return quantity.value_in_unit_system(unit.md_unit_system)
25 else:
26 return quantity
29def _add_fake_particles(force, num_particles):
30 if isinstance(force, openmm.CustomCVForce):
31 for i in range(force.getNumCollectiveVariables()):
32 _add_fake_particles(force.getCollectiveVariable(i), num_particles)
33 elif isinstance(force, openmm.NonbondedForce):
34 for i in range(num_particles - force.getNumParticles()):
35 force.addParticle(0.0, 1.0, 0.0)
36 elif isinstance(force, openmm.CustomNonbondedForce):
37 for i in range(num_particles - force.getNumParticles()):
38 force.addParticle([0.0] * force.getNumPerParticleParameters())
39 elif isinstance(force, openmm.CustomGBForce):
40 if force.getNumParticles() < num_particles:
41 parameter_list = list(
42 map(force.getParticleParameters, range(force.getNumParticles()))
43 )
44 force.addPerParticleParameter("isreal")
45 for index in range(force.getNumEnergyTerms()):
46 expression, type = force.getEnergyTermParameters(index)
47 energy = (
48 "isreal*E_GB"
49 if type == force.SingleParticle
50 else "isreal1*isreal2*E_GB"
51 )
52 force.setEnergyTermParameters(
53 index, f"{energy}; E_GB={expression}", type
54 )
55 for index, parameters in enumerate(parameter_list):
56 force.setParticleParameters(index, parameters + (1.0,))
57 for i in range(num_particles - force.getNumParticles()):
58 force.addParticle(parameter_list[0] + (0.0,))
59 elif isinstance(force, openmm.GBSAOBCForce):
60 raise RuntimeError("GBSAOBCForce not supported")
61 elif hasattr(force, "getReferencePositions"):
62 refpos = force.getReferencePositions()._value
63 if len(refpos) < num_particles:
64 for _ in range(num_particles - len(refpos)):
65 refpos.append(openmm.Vec3(0, 0, 0))
66 force.setReferencePositions(refpos)
69class CollectiveVariable(object):
70 """A function of the particle coordinates, evaluated by means of an :OpenMM:`Force`
71 object.
73 Quoting the manual entry for :OpenMM:`CustomCVForce`:
75 "Each collective variable is defined by a Force object. The Force's potential
76 energy is computed, and that becomes the value of the variable. This provides
77 enormous flexibility in defining collective variables, especially by using
78 custom forces. Anything that can be computed as a potential function can also
79 be used as a collective variable."
81 Parameters
82 ----------
83 id : str
84 A valid identifier string for this collective variable.
85 force : openmm.Force
86 An :OpenMM:`Force` object whose energy function is used to evaluate this
87 collective variable.
89 Example
90 -------
91 >>> import openmm
92 >>> import ufedmm
93 >>> from openmm import unit
94 >>> cv = ufedmm.CollectiveVariable('psi', openmm.CustomTorsionForce('theta'))
95 >>> cv.force.addTorsion(0, 1, 2, 3, [])
96 0
97 """
99 def __init__(self, id, force):
100 if not id.isidentifier():
101 raise ValueError("Parameter id must be a valid variable identifier")
102 self.id = id
103 self.force = force
105 def _create_context(self, system, positions):
106 system_copy = deepcopy(system)
107 for force in system_copy.getForces():
108 force.setForceGroup(0)
109 force_copy = deepcopy(self.force)
110 force_copy.setForceGroup(1)
111 _add_fake_particles(force_copy, system.getNumParticles())
112 system_copy.addForce(force_copy)
113 platform = openmm.Platform.getPlatformByName("Reference")
114 context = openmm.Context(system_copy, openmm.CustomIntegrator(0), platform)
115 context.setPositions(positions)
116 return context
118 def evaluate(self, system, positions, cv_unit=None):
119 """Computes the value of the collective variable for a given system and a given
120 set of particle coordinates.
122 Parameters
123 ----------
124 system : openmm.System
125 The system for which the collective variable will be evaluated.
126 positions : list of openmm.Vec3
127 A list whose size equals the number of particles in the system and which
128 contains the coordinates of these particles.
130 Keyword Args
131 ------------
132 cv_unit : unit.Unit, default=None
133 The unity of measurement of the collective variable. If this is `None`,
134 then a numerical value is returned based on the OpenMM default units.
136 Returns
137 -------
138 float or unit.Quantity
140 Example
141 -------
142 >>> import ufedmm
143 >>> from openmm import unit
144 >>> model = ufedmm.AlanineDipeptideModel()
145 >>> model.phi.evaluate(model.system, model.positions)
146 3.141592653589793
147 >>> model.psi.evaluate(model.system, model.positions)
148 3.141592653589793
149 """
150 context = self._create_context(system, positions)
151 energy = context.getState(getEnergy=True, groups={1}).getPotentialEnergy()
152 value = energy.value_in_unit(unit.kilojoules_per_mole)
153 if cv_unit is not None:
154 value *= cv_unit / _standardized(1 * cv_unit)
155 return value
157 def effective_mass(self, system, positions, cv_unit=None):
158 """Computes the effective mass of the collective variable for a given system and
159 a given set of particle coordinates.
161 The effective mass of a collective variable :math:`q(\\mathbf{r})` is defined as
162 :cite:`Cuendet_2014`:
164 .. math::
165 m_\\mathrm{eff} = \\left(
166 \\sum_{j=1}^N \\frac{1}{m_j} \\left\\|
167 \\frac{dq}{d\\mathbf{r}_j}
168 \\right\\|^2
169 \\right)^{-1}
171 Parameters
172 ----------
173 system : openmm.System
174 The system for which the collective variable will be evaluated.
175 positions : list of openmm.Vec3
176 A list whose size equals the number of particles in the system and which
177 contains the coordinates of these particles.
179 Keyword Args
180 ------------
181 cv_unit : unit.Unit, default=None
182 The unity of measurement of the collective variable. If this is `None`,
183 then a numerical value is returned based on the OpenMM default units.
185 Returns
186 -------
187 float or unit.Quantity
189 Example
190 -------
191 >>> import ufedmm
192 >>> from openmm import unit
193 >>> model = ufedmm.AlanineDipeptideModel()
194 >>> model.phi.effective_mass(model.system, model.positions)
195 0.04795887...
196 >>> model.psi.effective_mass(model.system, model.positions)
197 0.05115582...
198 """
199 context = self._create_context(system, positions)
200 forces = _standardized(
201 context.getState(getForces=True, groups={1}).getForces(asNumpy=True)
202 )
203 denom = sum(
204 f.dot(f) / _standardized(system.getParticleMass(i))
205 for i, f in enumerate(forces)
206 )
207 effective_mass = 1.0 / float(denom)
208 if cv_unit is not None:
209 factor = _standardized(1 * cv_unit) ** 2
210 effective_mass *= factor * unit.dalton * (unit.nanometers / cv_unit) ** 2
211 return effective_mass
214class DynamicalVariable(object):
215 """An extended phase-space variable, whose dynamics is coupled to that of one of
216 more collective variables of a system.
218 The coupling occurs in the form of a potential energy term involving this dynamical
219 variable and its associated collective variables.
221 The default potential is a harmonic driving of the type:
223 .. math::
224 V(s, \\mathbf r) = \\frac{\\kappa}{2} [s - q(\\mathbf r)]^2
226 where :math:`s` is the new dynamical variable, :math:`q(\\mathbf r)` is its
227 associated collective variable, and :math:`kappa` is a force constant.
229 Parameters
230 ----------
231 id : str
232 A valid identifier string for this dynamical variable.
233 min_value : float or unit.Quantity
234 The minimum allowable value for this dynamical variable.
235 max_value : float or unit.Quantity
236 The maximum allowable value for this dynamical variable.
237 mass : float or unit.Quantity
238 The mass assigned to this dynamical variable, whose unit of measurement must
239 be compatible with `unit.dalton*(unit.nanometers/X)**2`, where `X` is the
240 unit of measurement of the dynamical variable itself.
241 temperature : float or unit.Quantity
242 The temperature of the heat bath attached to this variable.
243 colvars : :class:`~ufedmm.ufedmm.CollectiveVariable` or list thereof
244 Either a single colective variable or a list.
245 potential : float or unit.Quantity or str
246 Either the value of the force constant of a harmonic driving potential or an
247 algebraic expression giving the energy of the system as a function of this
248 dynamical variable and its associated collective variable. Such expression
249 can also contain a set of global parameters, whose values must be passed as
250 keyword arguments (see below).
252 Keyword Args
253 ------------
254 periodic : bool, default=True
255 Whether the collective variable is periodic with period
256 `L=max_value-min_value`.
257 sigma : float or unit.Quantity, default=None
258 The standard deviation. If this is `None`, then no bias will be considered.
259 grid_size : int, default=None
260 The grid size. If this is `None` and `sigma` is finite, then a convenient
261 value will be automatically chosen.
262 **parameters
263 Names and values of global parameters present in the algebraic expression
264 defined as `potential` (see above).
266 Example
267 -------
268 >>> import openmm
269 >>> import ufedmm
270 >>> from openmm import unit
271 >>> cv = ufedmm.CollectiveVariable('psi', openmm.CustomTorsionForce('theta'))
272 >>> cv.force.addTorsion(0, 1, 2, 3, [])
273 0
274 >>> mass = 50*unit.dalton*(unit.nanometer/unit.radians)**2
275 >>> K = 1000*unit.kilojoules_per_mole/unit.radians**2
276 >>> Ts = 1500*unit.kelvin
277 >>> ufedmm.DynamicalVariable(
278 ... 's_psi', -180*unit.degrees, 180*unit.degrees, mass, Ts, cv, K
279 ... )
280 <s_psi in [-3.141592653589793, 3.141592653589793], periodic, m=50, T=1500>
281 """
283 def __init__(
284 self,
285 id,
286 min_value,
287 max_value,
288 mass,
289 temperature,
290 colvars,
291 potential,
292 periodic=True,
293 sigma=None,
294 grid_size=None,
295 **parameters,
296 ):
297 self.id = id
298 self.min_value = _standardized(min_value)
299 self.max_value = _standardized(max_value)
300 self._range = self.max_value - self.min_value
301 self.mass = _standardized(mass)
302 self.temperature = _standardized(temperature)
304 self.colvars = colvars if isinstance(colvars, (list, tuple)) else [colvars]
306 if isinstance(potential, str):
307 self.potential = potential
308 self.parameters = {
309 key: _standardized(value) for key, value in parameters.items()
310 }
311 else:
312 cv_id = self.colvars[0].id
313 if periodic:
314 self.potential = f"0.5*K_{cv_id}*min(d{cv_id},{self._range}-d{cv_id})^2"
315 self.potential += f"; d{cv_id}=abs({cv_id}-{self.id})"
316 else:
317 self.potential = f"0.5*K_{cv_id}*({cv_id}-{self.id})^2"
318 self.parameters = {f"K_{cv_id}": _standardized(potential)}
320 self.periodic = periodic
322 if sigma is None or sigma == 0.0:
323 self.sigma = self.grid_size = None
324 else:
325 self.sigma = _standardized(sigma)
326 self._scaled_variance = (self.sigma / self._range) ** 2
327 if grid_size is None:
328 self.grid_size = int(np.ceil(5 * self._range / self.sigma)) + 1
329 else:
330 self.grid_size = grid_size
332 self.force = openmm.CustomExternalForce(self._get_energy_function())
333 self.force.addGlobalParameter("Lx", 0.0)
334 self.force.addParticle(0, [])
336 def __repr__(self):
337 properties = f"m={self.mass}, T={self.temperature}"
338 status = "periodic" if self.periodic else "non-periodic"
339 return (
340 f"<{self.id} in "
341 f"[{self.min_value}, {self.max_value}], {status}, {properties}>"
342 )
344 def __getstate__(self):
345 return dict(
346 id=self.id,
347 min_value=self.min_value,
348 max_value=self.max_value,
349 mass=self.mass,
350 temperature=self.temperature,
351 colvars=self.colvars,
352 potential=self.potential,
353 periodic=self.periodic,
354 sigma=self.sigma,
355 grid_size=self.grid_size,
356 **self.parameters,
357 )
359 def __setstate__(self, kw):
360 self.__init__(**kw)
362 def _particle_mass(self, Lx):
363 length = Lx if self.periodic else Lx / 2
364 return self.mass * (self._range / length) ** 2
366 def _particle_position(self, value, Lx, y=0):
367 length = Lx if self.periodic else Lx / 2
368 return (
369 openmm.Vec3(length * (value - self.min_value) / self._range, y, 0)
370 * unit.nanometer
371 )
373 def _get_energy_function(self, index=""):
374 """Returns the algebraic expression that transforms the x coordinate of a
375 particle into this dynamical variables.
377 Keyword Args
378 ------------
379 index : str or int, default=''
380 An index for the particle in question, if needed.
382 Returns
383 -------
384 str
385 """
386 if self.periodic:
387 energy = f"{self.min_value}+{self._range}*(x{index}/Lx-floor(x{index}/Lx))"
388 else:
389 energy = f"{self.min_value}+{2*self._range}*min(pos{index},1-pos{index})"
390 energy += f";pos{index}=x{index}/Lx-floor(x{index}/Lx)"
391 return energy
393 def evaluate(self, x, Lx):
394 """Computes the value of this dynamical variable for a given x coordinate and a
395 given length of the simulation box in the x direction.
397 Parameters
398 ----------
399 x : float or unit.Quantity
400 The x coordinate.
401 Lx : float or unit.Quantity
402 The length of the simulation box in the x direction.
404 Returns
405 -------
406 float
407 """
408 pos = x / Lx - np.floor(x / Lx)
409 if self.periodic:
410 return self.min_value + self._range * pos
411 else:
412 return self.min_value + 2 * self._range * min(pos, 1 - pos)
415def _get_energy_function(variables):
416 energies = [v.potential.split(";", 1) for v in variables]
417 energy_terms = [energy[0] for energy in energies]
418 definitions = [energy[1] for energy in energies if len(energy) == 2]
419 expression = ";".join(["+".join(energy_terms)] + list(definitions))
420 return expression
423def _get_parameters(variables):
424 parameters = {}
425 for v in variables:
426 parameters.update(v.parameters)
427 return parameters
430class PeriodicTask(object):
431 def __init__(self, frequency):
432 self.frequency = frequency
434 def initialize(self, simulation, force_group):
435 pass
437 def update(self, simulation, steps):
438 pass
440 def describeNextReport(self, simulation):
441 steps = self.frequency - simulation.context.getStepCount() % self.frequency
442 return (steps, True, False, False, False, False)
444 def report(self, simulation, state):
445 pass
447 def saveCheckpoint(self, file):
448 pass
450 def loadCheckpoint(self, file):
451 pass
454class _Metadynamics(PeriodicTask):
455 """Extended-space Metadynamics.
457 Parameters
458 ----------
459 variables : list of :class:`DynamicalVariable`
460 A list of extended-space dynamical variables to which the metadynamics bias
461 must be applied. In fact, dynamical variables with `sigma = None` will not
462 be considered.
463 height : float or unit.Quantity
464 The height of the Gaussian potential hills to be deposited. If the
465 `bias_factor` keyword is defined (see below), then this is the unscaled
466 height.
467 frequency : int
468 The frequency of Gaussian hill deposition.
470 Keyword Args
471 ------------
472 bias_factor : float, default=None
473 Scales the height of the hills added to the bias potential. If it is `None`,
474 then the hills will have a constant height over time. For a bias factor to
475 be applicable, all bias variables must be at the same temperature T. The
476 extended-space dynamical variables are sampled as if the effective
477 temperature were T*bias_factor.
478 buffer_size : int, default=100
479 The buffer size.
480 grid_expansion : int, default=20
481 The number of extra grid points to be used in periodic directions of
482 multidimensional tabulated functions. This aims at avoiding boundary
483 discontinuity artifacts.
484 enforce_gridless : bool, default=False
485 Enforce gridless metadynamics even for 1D to 3D problems.
486 """
488 def __init__(
489 self,
490 variables,
491 height,
492 frequency,
493 bias_factor=None,
494 buffer_size=100,
495 grid_expansion=20,
496 enforce_gridless=False,
497 ):
498 super().__init__(frequency)
499 self.bias_indices = [i for i, v in enumerate(variables) if v.sigma is not None]
500 self.bias_variables = [variables[i] for i in self.bias_indices]
501 self.initial_height = self.height = _standardized(height)
502 self.bias_factor = bias_factor
503 if bias_factor is not None:
504 temperature = self.bias_variables[0].temperature
505 if any(v.temperature != temperature for v in self.bias_variables):
506 raise ValueError(
507 "Well-tempered metadynamics requires same temperature "
508 "for all variables."
509 )
510 self.delta_kT = _standardized(
511 (bias_factor - 1) * unit.MOLAR_GAS_CONSTANT_R * temperature
512 )
513 self.buffer_size = buffer_size
514 self.grid_expansion = grid_expansion
515 self._use_grid = len(self.bias_variables) < 4 and not enforce_gridless
516 self.force = (
517 self._interpolation_grid_force() if self._use_grid else self._hills_force()
518 )
520 def _interpolation_grid_force(self):
521 self._widths = []
522 self._bounds = []
523 self._extra_points = []
524 full_periodic = self._full_periodic = all(
525 v.periodic for v in self.bias_variables
526 )
527 for v in self.bias_variables:
528 extra_points = (
529 min(self.grid_expansion, v.grid_size - 1)
530 if v.periodic and not full_periodic
531 else 0
532 )
533 extra_range = extra_points * v._range / (v.grid_size - 1)
534 self._widths.append(v.grid_size + 2 * extra_points)
535 self._bounds += [v.min_value - extra_range, v.max_value + extra_range]
536 self._extra_points.append(extra_points)
537 self._bias = np.zeros(np.prod(self._widths))
538 num_bias_variables = len(self.bias_variables)
539 if num_bias_variables == 1:
540 self._table = openmm.Continuous1DFunction(
541 self._bias, *self._bounds, full_periodic
542 )
543 elif num_bias_variables == 2:
544 self._table = openmm.Continuous2DFunction(
545 *self._widths, self._bias, *self._bounds, full_periodic
546 )
547 else:
548 self._table = openmm.Continuous3DFunction(
549 *self._widths, self._bias, *self._bounds, full_periodic
550 )
551 expression = (
552 f'metad_bias_scale*bias({",".join(v.id for v in self.bias_variables)})'
553 )
554 for i, v in enumerate(self.bias_variables):
555 expression += f";{v.id}={v._get_energy_function(i+1)}"
556 force = openmm.CustomCVForce(expression)
557 for i in range(num_bias_variables):
558 x = openmm.CustomExternalForce("x")
559 x.addParticle(0, [])
560 force.addCollectiveVariable(f"x{i+1}", x)
561 force.addTabulatedFunction("bias", self._table)
562 force.addGlobalParameter("Lx", 0)
563 force.addGlobalParameter("metad_bias_scale", 1)
564 force.addEnergyParameterDerivative("metad_bias_scale")
565 return force
567 def _hills_force(self):
568 num_bias_variables = len(self.bias_variables)
569 centers = [f"center{i+1}" for i in range(num_bias_variables)]
570 exponents = []
571 for v, center in zip(self.bias_variables, centers):
572 if v.periodic: # von Mises
573 factor = 2 * np.pi / v._range
574 exponents.append(
575 f"{1.0/(factor*v.sigma)**2}*(cos({factor}*({v.id}-{center}))-1)"
576 )
577 else: # Gauss
578 exponents.append(f"({-0.5/v.sigma**2})*({v.id}-{center})^2")
579 expression = f'height*exp({"+".join(exponents)})'
580 for i, v in enumerate(self.bias_variables):
581 expression += f";{v.id}={v._get_energy_function(i+1)}"
582 force = openmm.CustomCompoundBondForce(num_bias_variables, expression)
583 force.addPerBondParameter("height")
584 for center in centers:
585 force.addPerBondParameter(center)
586 force.addGlobalParameter("Lx", 0)
587 return force
589 def _add_buffer(self, simulation):
590 size = min(self.buffer_size, self._total_hills - self.force.getNumBonds())
591 for i in range(size):
592 self.force.addBond(self.particles, [0] * (len(self.bias_variables) + 1))
593 simulation.context.reinitialize(preserveState=True)
595 def add_bias(self, simulation, bias):
596 if self._use_grid:
597 self._bias += bias.ravel()
598 if len(self.bias_variables) == 1:
599 self._table.setFunctionParameters(self._bias, *self._bounds)
600 else:
601 self._table.setFunctionParameters(
602 *self._widths, self._bias, *self._bounds
603 )
604 self.force.updateParametersInContext(simulation.context)
606 def initialize(self, simulation):
607 context = simulation.context
608 np = context.getSystem().getNumParticles() - len(context.variables)
609 self.particles = [np + index for index in self.bias_indices]
610 self.Lx = context.getParameter("Lx")
611 if self._use_grid:
612 for i, particle in enumerate(self.particles):
613 self.force.getCollectiveVariable(i).setParticleParameters(
614 0, particle, []
615 )
616 else:
617 self._num_hills = 0
618 self._total_hills = 0
619 simulation.system.addForce(self.force)
620 context.reinitialize(preserveState=True)
622 def update(self, simulation, steps):
623 if not self._use_grid:
624 steps_until_next_report = (
625 self.frequency - simulation.currentStep % self.frequency
626 )
627 if steps_until_next_report > steps:
628 required_hills = 0
629 else:
630 required_hills = (steps - steps_until_next_report) // self.frequency + 1
631 self._total_hills = self._num_hills + required_hills
632 self._add_buffer(simulation)
634 def report(self, simulation, state):
635 centers = state.getDynamicalVariables()
636 if self.bias_factor is None:
637 self.height = self.initial_height
638 else:
639 hills_state = simulation.context.getState(getParameterDerivatives=True)
640 energy = hills_state.getEnergyParameterDerivatives()["metad_bias_scale"]
641 self.height = self.initial_height * np.exp(-energy / self.delta_kT)
642 if self._use_grid:
643 hills = []
644 for i, v in enumerate(self.bias_variables):
645 x = (centers[i] - v.min_value) / v._range
646 dist = np.linspace(0, 1, num=v.grid_size) - x
647 if v.periodic: # von Mises
648 exponents = (np.cos(2 * np.pi * dist) - 1) / (
649 4 * np.pi * np.pi * v._scaled_variance
650 )
651 exponents[0] = exponents[-1] = 0.5 * (exponents[0] + exponents[-1])
652 else: # Gauss
653 exponents = -0.5 * dist * dist / v._scaled_variance
654 hills.append(self.height * np.exp(exponents))
655 ndim = len(self.bias_variables)
656 bias = (
657 hills[0]
658 if ndim == 1
659 else functools.reduce(np.multiply.outer, reversed(hills))
660 )
661 if not self._full_periodic:
662 for axis, v in enumerate(self.bias_variables):
663 if v.periodic:
664 n = self._extra_points[axis] + 1
665 begin = tuple(
666 slice(1, n) if i == axis else slice(None)
667 for i in range(ndim)
668 )
669 end = tuple(
670 slice(-n, -1) if i == axis else slice(None)
671 for i in range(ndim)
672 )
673 bias = np.concatenate((bias[end], bias, bias[begin]), axis=axis)
674 self.add_bias(simulation, bias)
675 else:
676 if self._num_hills == self.force.getNumBonds():
677 self._add_buffer(simulation)
678 self.force.setBondParameters(
679 self._num_hills, self.particles, [self.height] + centers
680 )
681 self._num_hills += 1
682 self.force.updateParametersInContext(simulation.context)
684 def saveCheckpoint(self, file):
685 if self._use_grid:
686 self._bias.tofile(file)
687 else:
688 np.array([self._num_hills], dtype=int).tofile(file)
689 parameter_list = []
690 for index in range(self._num_hills):
691 _, parameters = self.force.getBondParameters(index)
692 parameter_list.append(parameters)
693 np.array(parameter_list).tofile(file)
694 np.array([self.height]).tofile(file)
696 def loadCheckpoint(self, file, context):
697 if self._use_grid:
698 self._bias = np.fromfile(file, count=len(self._bias))
699 if len(self.bias_variables) == 1:
700 self._table.setFunctionParameters(self._bias, *self._bounds)
701 else:
702 self._table.setFunctionParameters(
703 *self._widths, self._bias, *self._bounds
704 )
705 self.force.updateParametersInContext(context)
706 else:
707 npars = len(self.bias_variables) + 1
708 nhills = self._num_hills = np.fromfile(file, dtype=int, count=1)[0]
709 parameter_array = np.fromfile(file, count=nhills * npars).reshape(
710 (nhills, npars)
711 )
712 for index in range(nhills - self.force.getNumBonds()):
713 self.force.addBond(self.particles, [0] * npars)
714 self._total_hills = self.force.getNumBonds()
715 for index, parameters in enumerate(parameter_array):
716 self.force.setBondParameters(index, self.particles, parameters)
717 for index in range(self._total_hills - nhills):
718 self.force.setBondParameters(index, self.particles, [0] * npars)
719 context.reinitialize(preserveState=True)
720 self.height = np.fromfile(file, count=1)[0]
723class ExtendedSpaceState(openmm.State):
724 """An extension of the :OpenMM:`State` class."""
726 def __init__(self, variables, state):
727 self.__class__ = type(
728 state.__class__.__name__, (self.__class__, state.__class__), {}
729 )
730 self.__dict__ = state.__dict__
731 self._variables = variables
732 a, _, _ = self.getPeriodicBoxVectors()
733 self._Lx = a.x
735 def _split(self, vector, asNumpy=False):
736 num = (vector.shape[0] if asNumpy else len(vector)) - len(self._variables)
737 particles_contribution = vector[:num, :] if asNumpy else vector[:num]
738 variables_contribution = (
739 vector[num:, 0] if asNumpy else [vec.x for vec in vector[num:]]
740 )
741 return particles_contribution, variables_contribution
743 def getPositions(self, asNumpy=False, extended=False, split=True):
744 """Gets the positions of all physical particles and optionally also gets the
745 positions of the extra particles from which the extended-space dynamical
746 variables are computed.
748 Keyword Args
749 ------------
750 asNumpy : bool, default=False
751 Whether to return Numpy arrays instead of lists of openmm.Vec3.
752 extended : bool, default=False
753 Whether to include the positions of the extra particles from which the
754 extended-space dynamical variables are computed.
755 split : bool, default=True
756 Whether to return the positions of the physical particles and the
757 positions of the extra particles separately.
759 Returns
760 -------
761 list(openmm.Vec3) or numpy.ndarray
762 If `extended=False` or `split=False`.
763 list(openmm.Vec3) or numpy.ndarray, list(float) or numpy.ndarray
764 If `extended=True` and `split=True`.
766 Raises
767 ------
768 Exception
769 If positions were not requested in the ``context.getState()`` call.
770 """
771 all_positions = super().getPositions(asNumpy)
772 if extended and not split:
773 return all_positions
774 positions, xvars = self._split(all_positions, asNumpy)
775 return (positions, xvars) if extended else positions
777 def getVelocities(self, asNumpy=False, extended=False):
778 """Gets the velocities of all physical particles and optionally also gets the
779 velocities of the extra particles associated to the extended- space dynamical
780 variables.
782 Keyword Args
783 ------------
784 asNumpy : bool, default=False
785 Whether to return Numpy arrays instead of lists of openmm.Vec3.
786 extended : bool, default=False
787 Whether to include the velocities of the extra particles.
789 Returns
790 -------
791 list(openmm.Vec3)
792 If `asNumpy=False` and `extended=False`.
793 numpy.ndarray
794 If `asNumpy=True` and `extended=False`.
795 list(openmm.Vec3), list(float)
796 If `asNumpy=False` and `extended=True`.
797 numpy.ndarray, numpy.ndarray
798 If `asNumpy=True` and `extended=True`.
800 Raises
801 ------
802 Exception
803 If velocities were not requested in the ``context.getState()`` call.
804 """
805 velocities = super().getVelocities(asNumpy)
806 if not extended:
807 velocities, _ = self._split(velocities, asNumpy)
808 return velocities
810 def getDynamicalVariables(self):
811 """Gets the values of the extended-space dynamical variables.
813 Returns
814 -------
815 list(float)
817 Raises
818 ------
819 Exception
820 If positions were not requested in the ``context.getState()`` call.
822 Example
823 -------
824 >>> import ufedmm
825 >>> import numpy as np
826 >>> model = ufedmm.AlanineDipeptideModel()
827 >>> integrator = ufedmm.CustomIntegrator(300, 0.001)
828 >>> args = [-np.pi, np.pi, 50, 1500, model.phi, 1000]
829 >>> s_phi = ufedmm.DynamicalVariable('s_phi', *args)
830 >>> s_psi = ufedmm.DynamicalVariable('s_psi', *args)
831 >>> context = ufedmm.ExtendedSpaceContext(
832 ... [s_phi, s_psi], model.system, integrator
833 ... )
834 >>> context.setPositions(model.positions)
835 >>> context.getState(getPositions=True).getDynamicalVariables()
836 [-3.141592653589793, -3.141592653589793]
837 """
838 _, xvars = self._split(super().getPositions())
839 return [v.evaluate(x, self._Lx) for v, x in zip(self._variables, xvars)]
842class ExtendedSpaceContext(openmm.Context):
843 """An extension of the :OpenMM:`Context` class.
845 Parameters
846 ----------
847 variables : list(DynamicalVariable)
848 The dynamical variables to be added to the system's phase space.
849 system : openmm.System
850 The :OpenMM:`System` which will be simulated
851 integrator : openmm.Integrator
852 The :OpenMM:`Integrator` which will be used to simulate the
853 :OpenMM:`System`.
854 platform : openmm.Platform
855 The :OpenMM:`Platform` to use for calculations.
856 properties : dict(str: str)
857 A set of values for platform-specific properties. Keys are the property
858 names.
860 Attributes
861 ----------
862 variables : list(DynamicalVariable)
863 The dynamical variables added to the system's phase space.
864 driving_forces : list[openmm.CustomCVForce]
865 A list of :OpenMM:`CustomCVForce` objects responsible for evaluating the
866 potential energy terms which couples the extra dynamical variables to their
867 associated collective variables.
868 """
870 def __init__(self, variables, system, *args, **kwargs):
871 def get_driving_force(variables):
872 driving_force = openmm.CustomCVForce(_get_energy_function(variables))
873 for name, value in _get_parameters(variables).items():
874 driving_force.addGlobalParameter(name, value)
875 for var in variables:
876 driving_force.addCollectiveVariable(var.id, var.force)
877 for colvar in var.colvars:
878 driving_force.addCollectiveVariable(colvar.id, colvar.force)
879 return driving_force
881 driving_force_variables = [[]]
882 added_colvars = 0
883 for var in variables:
884 colvars_to_add = len(var.colvars) + 1
885 if added_colvars + colvars_to_add > 32:
886 driving_force_variables.append([])
887 added_colvars = 0
888 driving_force_variables[-1].append(deepcopy(var))
889 added_colvars += colvars_to_add
890 driving_forces = [get_driving_force(vars) for vars in driving_force_variables]
891 box_length_x = system.getDefaultPeriodicBoxVectors()[0].x
892 num_particles = system.getNumParticles()
893 collective_variables = []
894 for dfvars in driving_force_variables:
895 for var in dfvars:
896 system.addParticle(var._particle_mass(box_length_x))
897 var.force.setParticleParameters(0, num_particles, [])
898 num_particles += 1
899 collective_variables += var.colvars
900 for force in system.getForces():
901 _add_fake_particles(force, num_particles)
902 for cv in collective_variables:
903 _add_fake_particles(cv.force, num_particles)
904 for driving_force in driving_forces:
905 system.addForce(driving_force)
906 super().__init__(system, *args, **kwargs)
907 self.setParameter("Lx", box_length_x)
908 self.variables = variables
909 self.driving_forces = driving_forces
910 if len(driving_forces) == 1:
911 self.driving_force = driving_forces[0] # for backwards compatibility
913 def getState(self, **kwargs):
914 """Returns a :class:`ExtendedSpaceState` object.
916 .. _getState: http://docs.openmm.org/latest/api-python/generated/openmm.openmm.Context.html#openmm.openmm.Context.getState # noqa: E501
918 Keyword Args
919 ------------
920 **kwargs
921 See getState_.
922 """
923 return ExtendedSpaceState(self.variables, super().getState(**kwargs))
925 def setPeriodicBoxVectors(self, a, b, c):
926 """Set the vectors defining the axes of the periodic box.
928 .. warning::
929 Only orthorhombic boxes are allowed.
931 Parameters
932 ----------
933 a : openmm.Vec3
934 The vector defining the first edge of the periodic box.
935 b : openmm.Vec3
936 The vector defining the second edge of the periodic box.
937 c : openmm.Vec3
938 The vector defining the third edge of the periodic box.
939 """
940 a = openmm.Vec3(*map(_standardized, a))
941 b = openmm.Vec3(*map(_standardized, b))
942 c = openmm.Vec3(*map(_standardized, c))
943 if not (a.y == a.z == b.x == b.z == c.x == c.y == 0.0):
944 raise ValueError("Only orthorhombic boxes are allowed")
945 self.setParameter("Lx", a.x)
946 super().setPeriodicBoxVectors(a, b, c)
947 system = self.getSystem()
948 ntotal = system.getNumParticles()
949 nvars = len(self.variables)
950 for i, v in enumerate(self.variables):
951 system.setParticleMass(ntotal - nvars + i, v._particle_mass(a.x))
952 self.reinitialize(preserveState=True)
954 def setPositions(self, positions, extended_positions=None):
955 """Sets the positions of all particles and extended-space variables in this
956 context. If the latter are not provided, then suitable values are automatically
957 determined from the particle positions.
959 Parameters
960 ----------
961 positions : list of openmm.Vec3
962 The positions of physical particles.
964 Keyword Args
965 ------------
966 extended_positions : list of float or unit.Quantity
967 The positions of extended-space particles.
968 """
969 a, b, _ = self.getState().getPeriodicBoxVectors()
970 nvars = len(self.variables)
971 particle_positions = _standardized(positions)
972 if extended_positions is None:
973 extra_positions = [
974 openmm.Vec3(0, b.y * (i + 1) / (nvars + 2), 0) for i in range(nvars)
975 ]
976 minisystem = openmm.System()
977 expression = _get_energy_function(self.variables)
978 for i, v in enumerate(self.variables):
979 expression += f"; {v.id}={v._get_energy_function(index=i+1)}"
980 force = openmm.CustomCompoundBondForce(nvars, expression)
981 force.addBond(range(nvars), [])
982 for name, value in _get_parameters(self.variables).items():
983 force.addGlobalParameter(name, value)
984 force.addGlobalParameter("Lx", a.x)
985 for v in self.variables:
986 minisystem.addParticle(v._particle_mass(a.x))
987 for cv in v.colvars:
988 value = cv.evaluate(
989 self.getSystem(), particle_positions + extra_positions
990 )
991 force.addGlobalParameter(cv.id, value)
992 minisystem.addForce(force)
993 minicontext = openmm.Context(
994 minisystem,
995 openmm.CustomIntegrator(0),
996 openmm.Platform.getPlatformByName("Reference"),
997 )
998 minicontext.setPositions(extra_positions)
999 openmm.LocalEnergyMinimizer.minimize(
1000 minicontext, 1 * unit.kilojoules_per_mole / unit.nanometers, 0
1001 )
1002 ministate = minicontext.getState(getPositions=True)
1003 extra_positions = ministate.getPositions().value_in_unit(unit.nanometers)
1004 else:
1005 extra_positions = [
1006 openmm.Vec3(x, b.y * (i + 1) / (nvars + 2), 0)
1007 for i, x in enumerate(extended_positions)
1008 ]
1009 super().setPositions(particle_positions + extra_positions)
1011 def setVelocitiesToTemperature(self, temperature, randomSeed=None):
1012 """Sets the velocities of all particles in the system to random values chosen
1013 from a Maxwell-Boltzmann distribution at a given temperature.
1015 .. warning ::
1016 The velocities of the extended-space variables are set to values consistent
1017 with the distribution at its own specified temperature.
1019 Parameters
1020 ----------
1021 temperature : float or unit.Quantity
1022 The temperature of the system.
1024 Keyword Args
1025 ------------
1026 randomSeed : int, default=None
1027 A seed for the random number generator.
1028 """
1029 system = self.getSystem()
1030 Ntotal = system.getNumParticles()
1031 Natoms = Ntotal - len(self.variables)
1032 m = np.array([system.getParticleMass(i) / unit.dalton for i in range(Ntotal)])
1033 T = np.array(
1034 [_standardized(temperature)] * Natoms
1035 + [v.temperature for v in self.variables]
1036 )
1037 sigma = np.sqrt(_standardized(unit.MOLAR_GAS_CONSTANT_R) * T / m)
1038 random_state = np.random.RandomState(randomSeed)
1039 velocities = sigma[:, np.newaxis] * random_state.normal(0, 1, (Ntotal, 3))
1040 super().setVelocities(velocities)
1043class ExtendedSpaceSimulation(app.Simulation):
1044 """A simulation involving extended phase-space variables.
1046 Parameters
1047 ----------
1048 variables : list of DynamicalVariable
1049 The dynamical variables to be added to the system's phase space.
1050 topology : openmm.app.Topology
1051 A Topology describing the system to be simulated.
1052 system : openmm.System
1053 The :OpenMM:`System` object to be simulated.
1054 integrator : openmm.Integrator
1055 The OpenMM Integrator to use for simulating the system dynamics.
1057 Keyword Args
1058 ------------
1059 platform : openmm.Platform, default=None
1060 If not None, the :OpenMM:`Platform` to use.
1061 platformProperties : dict, default=None
1062 If not None, a set of platform-specific properties to pass to the Context's
1063 constructor.
1064 """
1066 def __init__(
1067 self,
1068 variables,
1069 topology,
1070 system,
1071 integrator,
1072 platform=None,
1073 platformProperties=None,
1074 ):
1075 self.variables = variables
1076 self._periodic_tasks = []
1078 for force in system.getForces():
1079 cls = force.__class__.__name__
1080 if cls == "CMMotionRemover" or cls.endswith("Barostat"):
1081 raise Exception(
1082 "UFED: system cannot contain CMMotionRemover nor any Barostat"
1083 )
1085 box_vectors = topology.getPeriodicBoxVectors()
1086 if box_vectors is None:
1087 raise Exception("UFED: system must be confined in a simulation box")
1089 self.topology = topology
1090 self.system = system
1091 self.integrator = integrator
1092 if openmm.__version__ < "7.7":
1093 self.currentStep = 0
1094 self.reporters = []
1095 self._usesPBC = True
1096 if platform is None:
1097 self.context = ExtendedSpaceContext(variables, system, integrator)
1098 elif platformProperties is None:
1099 self.context = ExtendedSpaceContext(variables, system, integrator, platform)
1100 else:
1101 self.context = ExtendedSpaceContext(
1102 variables, system, integrator, platform, platformProperties
1103 )
1105 def add_periodic_task(self, task, force_group=0):
1106 """Adds a task to be executed periodically along this simulation.
1108 Parameters
1109 ----------
1110 task : PeriodicTask
1111 A :class:`~ufedmm.ufedmm.PeriodicTask` object.
1112 """
1113 task.initialize(self)
1114 self._periodic_tasks.append(task)
1116 def step(self, steps):
1117 """Executed a given number of simulation steps.
1119 Parameters
1120 ----------
1121 steps : int
1122 The number of steps to be executed.
1123 """
1124 if isinstance(self.integrator, ufedmm.AbstractMiddleRespaIntegrator):
1125 if (
1126 self.integrator._num_rattles == 0
1127 and self.system.getNumConstraints() > 0
1128 ):
1129 raise RuntimeError("Integrator cannot handle constraints")
1130 if self._periodic_tasks:
1131 for task in self._periodic_tasks:
1132 task.update(self, steps)
1133 self.reporters = self._periodic_tasks + self.reporters
1134 self._simulate(endStep=self.currentStep + steps)
1135 self.reporters = self.reporters[len(self._periodic_tasks) :]
1136 else:
1137 self._simulate(endStep=self.currentStep + steps)
1139 def saveCheckpoint(self, file):
1140 if isinstance(file, str):
1141 with open(file, "wb") as f:
1142 self.saveCheckpoint(f)
1143 else:
1144 for task in self._periodic_tasks:
1145 task.saveCheckpoint(file)
1146 file.write(self.context.createCheckpoint())
1148 def loadCheckpoint(self, file):
1149 if isinstance(file, str):
1150 with open(file, "rb") as f:
1151 self.loadCheckpoint(f)
1152 else:
1153 for task in self._periodic_tasks:
1154 task.loadCheckpoint(file, self.context)
1155 self.context.loadCheckpoint(file.read())
1158class UnifiedFreeEnergyDynamics(object):
1159 """A Unified Free-Energy Dynamics (UFED) setup.
1161 Parameters
1162 ----------
1163 variables : list of DynamicalVariable
1164 The variables.
1165 temperature : float or unit.Quantity
1166 The temperature.
1168 Keyword Args
1169 ------------
1170 height : float or unit.Quantity, default=None
1171 The height.
1172 frequency : int, default=None
1173 The frequency.
1174 bias_factor : float, default=None
1175 Scales the height of the hills added to the metadynamics bias potential. If
1176 it is `None`, then the hills will have a constant height over time. For a
1177 bias factor to be applicable, all bias variables must be at the same
1178 temperature T. The extended-space dynamical variables are sampled as if the
1179 effective temperature were T*bias_factor.
1180 grid_expansion : int, default=20
1181 The grid expansion.
1182 enforce_gridless : bool, default=False
1183 If this is `True`, gridless metadynamics is enforced even for 1D to 3D
1184 problems.
1185 buffer_size : int, default=100
1186 The buffer size.
1188 Example
1189 -------
1190 >>> import ufedmm
1191 >>> from openmm import unit
1192 >>> model = ufedmm.AlanineDipeptideModel(water='tip3p')
1193 >>> mass = 50*unit.dalton*(unit.nanometer/unit.radians)**2
1194 >>> Ks = 1000*unit.kilojoules_per_mole/unit.radians**2
1195 >>> Ts = 1500*unit.kelvin
1196 >>> limit = 180*unit.degrees
1197 >>> s_phi = ufedmm.DynamicalVariable(
1198 ... 's_phi', -limit, limit, mass, Ts, model.phi, Ks
1199 ... )
1200 >>> s_psi = ufedmm.DynamicalVariable(
1201 ... 's_psi', -limit, limit, mass, Ts, model.psi, Ks
1202 ... )
1203 >>> ufedmm.UnifiedFreeEnergyDynamics([s_phi, s_psi], 300*unit.kelvin)
1204 <variables=[s_phi, s_psi], temperature=300, height=None, frequency=None>
1205 """
1207 def __init__(
1208 self,
1209 variables,
1210 temperature,
1211 height=None,
1212 frequency=None,
1213 bias_factor=None,
1214 grid_expansion=20,
1215 enforce_gridless=False,
1216 buffer_size=100,
1217 ):
1218 self.variables = variables
1219 self.temperature = _standardized(temperature)
1220 self.height = _standardized(height)
1221 self.frequency = frequency
1222 self.bias_factor = bias_factor
1223 self.grid_expansion = grid_expansion
1224 self.enforce_gridless = enforce_gridless
1225 self.buffer_size = buffer_size
1227 dimensions = sum(v.sigma is not None for v in variables)
1228 self._metadynamics = not (
1229 dimensions == 0 or height is None or frequency is None
1230 )
1232 def __repr__(self):
1233 properties = (
1234 f"temperature={self.temperature}, "
1235 f"height={self.height}, "
1236 f"frequency={self.frequency}"
1237 )
1238 return f'<variables=[{", ".join(v.id for v in self.variables)}], {properties}>'
1240 def __getstate__(self):
1241 return dict(
1242 variables=self.variables,
1243 temperature=self.temperature,
1244 height=self.height,
1245 frequency=self.frequency,
1246 bias_factor=self.bias_factor,
1247 grid_expansion=self.grid_expansion,
1248 enforce_gridless=self.enforce_gridless,
1249 buffer_size=self.buffer_size,
1250 )
1252 def __setstate__(self, kw):
1253 self.__init__(**kw)
1255 def simulation(
1256 self, topology, system, integrator, platform=None, platformProperties=None
1257 ):
1258 """Returns a ExtendedSpaceSimulation object.
1260 .. warning::
1261 If the temperature of any driving parameter is different from the
1262 particle-system temperature, then the passed integrator must be a
1263 :OpenMM:`CustomIntegrator` object containing a per-dof variable `kT` whose
1264 content is the Boltzmann constant times the temperature associated to each
1265 degree of freedom. This is true for all integrators available in
1266 :mod:`ufedmm.integrators`, which are subclasses of
1267 :class:`ufedmm.integrators.CustomIntegrator`.
1269 Parameters
1270 ----------
1271 topology : openmm.app.Topology
1272 The topology.
1273 system : openmm.System
1274 The system.
1275 integrator :
1276 The integrator.
1278 Keyword Args
1279 ------------
1280 platform : openmm.Platform, default=None
1281 The platform.
1282 platformProperties : dict, default=None
1283 The platform properties.
1285 Example
1286 -------
1287 >>> import ufedmm
1288 >>> from openmm import unit
1289 >>> model = ufedmm.AlanineDipeptideModel(water='tip3p')
1290 >>> mass = 50*unit.dalton*(unit.nanometer/unit.radians)**2
1291 >>> Ks = 1000*unit.kilojoules_per_mole/unit.radians**2
1292 >>> Ts = 1500*unit.kelvin
1293 >>> dt = 2*unit.femtoseconds
1294 >>> gamma = 10/unit.picoseconds
1295 >>> limit = 180*unit.degrees
1296 >>> s_phi = ufedmm.DynamicalVariable(
1297 ... 's_phi', -limit, limit, mass, Ts, model.phi, Ks
1298 ... )
1299 >>> s_psi = ufedmm.DynamicalVariable(
1300 ... 's_psi', -limit, limit, mass, Ts, model.psi, Ks
1301 ... )
1302 >>> ufed = ufedmm.UnifiedFreeEnergyDynamics([s_phi, s_psi], 300*unit.kelvin)
1303 >>> integrator = ufedmm.GeodesicLangevinIntegrator(
1304 ... 300*unit.kelvin, gamma, dt
1305 ... )
1306 >>> simulation = ufed.simulation(model.topology, model.system, integrator)
1307 """
1308 simulation = ExtendedSpaceSimulation(
1309 self.variables,
1310 topology,
1311 system,
1312 integrator,
1313 platform,
1314 platformProperties,
1315 )
1317 if self._metadynamics:
1318 simulation.add_periodic_task(
1319 _Metadynamics(
1320 self.variables,
1321 self.height,
1322 self.frequency,
1323 bias_factor=self.bias_factor,
1324 buffer_size=self.buffer_size,
1325 grid_expansion=self.grid_expansion,
1326 enforce_gridless=self.enforce_gridless,
1327 ),
1328 )
1330 if any(v.temperature != self.temperature for v in self.variables):
1331 ntotal = system.getNumParticles()
1332 nvars = len(self.variables)
1333 simulation.context.setPositions(
1334 [openmm.Vec3(0, 0, 0)] * (ntotal - nvars), [0] * nvars
1335 )
1336 vartemps = [v.temperature for v in self.variables]
1337 if isinstance(integrator, ufedmm.integrators.CustomIntegrator):
1338 integrator.update_temperatures(self.temperature, vartemps)
1339 else:
1340 temperatures = [self.temperature] * (ntotal - len(vartemps)) + vartemps
1341 kT = [
1342 unit.MOLAR_GAS_CONSTANT_R * T * openmm.Vec3(1, 1, 1)
1343 for T in temperatures
1344 ]
1345 integrator.setPerDofVariableByName("kT", kT)
1347 return simulation