Coverage for openxps/integrators/csvr.py: 97%
70 statements
« prev ^ index » next coverage.py v7.11.3, created at 2025-11-13 22:08 +0000
1"""
2.. module:: openxps.integrators.csvr
3 :platform: Linux, MacOS, Windows
4 :synopsis: Canonical Sampling through Velocity Rescaling (CSVR) integrator.
6.. classauthor:: Charlles Abreu <craabreu@gmail.com>
8"""
10import numpy as np
11import openmm as mm
12from openmm import unit as mmunit
14from openxps.utils import preprocess_args
16from .utils import IntegratorMixin, add_property
19@add_property("temperature", mmunit.kelvin)
20@add_property("friction coefficient", 1 / mmunit.picosecond)
21class CSVRIntegrator(IntegratorMixin, mm.CustomIntegrator):
22 """The Canonical Sampling through Velocity Rescaling integrator :cite:`Bussi2007`.
24 The Canonical Sampling through Velocity Rescaling (CSVR) thermostat scales all
25 velocities by a single factor stochastically determined in order to preserve the
26 canonical distribution at the specified temperature.
28 .. warning::
29 When used outside OpenXPS, the :meth:`registerWithSystem` method must be called
30 with `isExtension=False` before starting a simulation.
32 Parameters
33 ----------
34 temperature
35 The temperature.
36 frictionCoeff
37 The friction coefficient.
38 stepSize
39 The integration step size.
40 forceFirst
41 If True, the integrator will apply a force-first scheme rather than a
42 symmetric operator splitting scheme.
44 Example
45 -------
46 >>> import openxps as xps
47 >>> import openmm as mm
48 >>> from openmm import unit
49 >>> from openmmtools import testsystems
51 Symmetric scheme (default)
53 >>> integrator = xps.integrators.CSVRIntegrator(
54 ... temperature=300 * unit.kelvin,
55 ... frictionCoeff=10 / unit.picosecond,
56 ... stepSize=2 * unit.femtoseconds,
57 ... )
58 >>> integrator
59 Per-dof variables:
60 x1
61 Global variables:
62 sumRsq = 0.0
63 mvv = 0.0
64 kT = 2.49433...
65 friction = 10.0
66 scaling_factor = 1.0
67 Computation steps:
68 0: allow forces to update the context state
69 1: v <- v + 0.5*dt*f/m
70 2: constrain velocities
71 3: x <- x + 0.5*dt*v
72 4: x1 <- x
73 5: constrain positions
74 6: v <- v + (x - x1)/(0.5*dt)
75 7: constrain velocities
76 8: mvv <- sum(m*v*v)
77 9: scaling_factor <- sqrt(A + BC*(R1^2 + sumRsq) + 2*sqrt(A*BC)*R1); R1 = ...
78 10: v <- v*scaling_factor
79 11: x <- x + 0.5*dt*v
80 12: x1 <- x
81 13: constrain positions
82 14: v <- v + (x - x1)/(0.5*dt)
83 15: constrain velocities
84 16: v <- v + 0.5*dt*f/m
85 17: constrain velocities
87 Force-first scheme
89 >>> integrator_ff = xps.integrators.CSVRIntegrator(
90 ... temperature=300 * unit.kelvin,
91 ... frictionCoeff=10 / unit.picosecond,
92 ... stepSize=2 * unit.femtoseconds,
93 ... forceFirst=True
94 ... )
95 >>> try:
96 ... integrator_ff.getNumDegreesOfFreedom()
97 ... except ValueError as e:
98 ... print(e)
99 The number of degrees of freedom has not been determined...
100 >>> model = testsystems.AlanineDipeptideVacuum()
101 >>> integrator_ff.registerWithSystem(model.system, False)
102 >>> integrator_ff.getNumDegreesOfFreedom()
103 51
104 >>> integrator_ff.setRandomNumberSeed(1234)
105 >>> platform = mm.Platform.getPlatformByName("Reference")
106 >>> context = mm.Context(model.system, integrator_ff, platform)
107 >>> context.setPositions(model.positions)
108 >>> context.setVelocitiesToTemperature(300 * unit.kelvin, 4321)
109 >>> integrator_ff.step(10)
110 """
112 @preprocess_args
113 def __init__(
114 self,
115 temperature: mmunit.Quantity,
116 frictionCoeff: mmunit.Quantity,
117 stepSize: mmunit.Quantity,
118 forceFirst: bool = False,
119 ) -> None:
120 super().__init__(stepSize)
121 self._init_temperature(temperature)
122 self._init_friction_coefficient(frictionCoeff)
123 self._forceFirst = forceFirst
124 self._num_dof = None
125 self._rng = np.random.default_rng(None)
126 self._add_variables()
127 self.addUpdateContextState()
128 self._add_boost(1 if forceFirst else 0.5)
129 self._add_translation(0.5)
130 self._add_rescaling(1)
131 self._add_translation(0.5)
132 if not forceFirst:
133 self._add_boost(0.5)
135 def _add_variables(self) -> None:
136 self.addPerDofVariable("x1", 0)
137 self.addGlobalVariable("sumRsq", 0)
138 self.addGlobalVariable("mvv", 0)
139 self.addGlobalVariable("kT", 0)
140 self.addGlobalVariable("friction", 0)
141 self.addGlobalVariable("scaling_factor", 1)
142 self._update_global_variables()
144 def _update_global_variables(self) -> None:
145 kt = mmunit.MOLAR_GAS_CONSTANT_R * self.getTemperature()
146 friction = self.getFrictionCoefficient()
147 self.setGlobalVariableByName("kT", kt)
148 self.setGlobalVariableByName("friction", friction)
150 def _add_translation(self, fraction: float) -> None:
151 self.addComputePerDof("x", f"x + {fraction}*dt*v")
152 self.addComputePerDof("x1", "x")
153 self.addConstrainPositions()
154 self.addComputePerDof("v", f"v + (x - x1)/({fraction}*dt)")
155 self.addConstrainVelocities()
157 def _add_boost(self, fraction: float) -> None:
158 self.addComputePerDof("v", f"v + {fraction}*dt*f/m")
159 self.addConstrainVelocities()
161 def _add_rescaling(self, fraction: float) -> None:
162 self.addComputeSum("mvv", "m*v*v")
163 self.addComputeGlobal(
164 "scaling_factor",
165 "sqrt(A + BC*(R1^2 + sumRsq) + 2*sqrt(A*BC)*R1)"
166 "; R1 = gaussian"
167 "; BC = (1 - A)*kT/mvv"
168 f"; A = exp(-{fraction}*dt*friction)",
169 )
170 self.addComputePerDof("v", "v*scaling_factor")
172 def _sums_of_squared_gaussians(self, num_steps: int) -> np.ndarray:
173 sumRsq = 2.0 * self._rng.standard_gamma((self._num_dof - 1) // 2, num_steps)
174 if self._num_dof % 2 == 0:
175 sumRsq += self._rng.standard_normal(num_steps) ** 2
176 return sumRsq
178 def getNumDegreesOfFreedom(self) -> int:
179 """Get the number of degrees of freedom in the system.
181 Raises
182 ------
183 ValueError
184 If registerWithSystem has not been called first.
185 """
186 if self._num_dof is None:
187 raise ValueError(
188 "The number of degrees of freedom has not been determined.\n"
189 "Call the `registerWithSystem` method first."
190 )
191 return self._num_dof
193 def registerWithSystem(self, system: mm.System, isExtension: bool) -> None:
194 if isExtension:
195 self._num_dof = system.getNumParticles()
196 else:
197 self._num_dof = IntegratorMixin._countDegreesOfFreedom(system)
199 def setRandomNumberSeed(self, seed: int) -> None:
200 """
201 This method overrides the :class:`openmm.CustomIntegrator` method to also set
202 the seed of the random number generator used to pick numbers from the gamma
203 distribution.
204 Parameters
205 ----------
206 seed
207 The seed to use for the random number generator.
208 """
209 self._rng = np.random.default_rng(seed + 2**31)
210 super().setRandomNumberSeed(self._rng.integers(-(2**31), 2**31))
212 def step(self, steps: int) -> None:
213 """
214 This method overrides the :class:`openmm.CustomIntegrator` method to include the
215 efficient computation of the sum of squares of normally distributed random
216 numbers.
217 Parameters
218 ----------
219 steps
220 The number of steps to take.
221 """
222 for sumRsq in self._sums_of_squared_gaussians(steps):
223 self.setGlobalVariableByName("sumRsq", sumRsq)
224 super().step(1)