Canonical Sampling through Velocity Rescaling in OpenMM

csvr.py

from typing import Union

import numpy as np
import openmm

from openmm import unit


class CSVRIntegrator(openmm.CustomIntegrator):
    """
    Implements the Canonical Sampling through Velocity Rescaling (CSVR) integrator,
    also known as the Bussi-Donadio-Parrinello thermostat.

    Parameters
    ----------
    temperature
        The temperature.
    frictionCoeff
        The friction coefficient.
    stepSize
        The integration step size.
    num_dof
        The number of degrees of freedom in the system.
    splitting
        The splitting scheme. Valid letters are R, V, and O, where R stands for a
        position update, V for a velocity update, and O for a rescaling of the
        velocities. The default is 'VROR', which is the CSVR version of the
        Leapfrog Middle (LF-Middle) scheme.
    """

    def __init__(
        self,
        temperature: Union[float, unit.Quantity],
        frictionCoeff: Union[float, unit.Quantity],
        stepSize: int,
        num_dof: int,
        splitting: str = "VROR",
    ) -> None:
        super().__init__(stepSize)
        self._num_dof = num_dof
        self._rng = np.random.default_rng(None)
        self._add_global_variables(temperature, frictionCoeff)
        self.addUpdateContextState()
        for letter in splitting:
            n = splitting.count(letter)
            timestep = "dt" if n == 1 else f"{1/n}*dt"
            if letter == "V":
                self._add_boost(timestep)
            elif letter == "R":
                self._add_translation(timestep)
            elif letter == "O":
                self._add_rescaling(timestep)
            else:
                raise ValueError("Valid splitting scheme letters are R, V, and O")

    def _add_global_variables(
        self, temperature: Union[float, unit.Quantity], frictionCoeff: Union[float, unit.Quantity]
    ) -> None:
        self.addPerDofVariable("x1", 0)
        self.addGlobalVariable("sumRsq", 0)
        self.addGlobalVariable("mvv", 0)
        self.addGlobalVariable("kT", unit.MOLAR_GAS_CONSTANT_R * temperature)
        self.addGlobalVariable("friction", frictionCoeff)

    def _add_translation(self, timestep: str) -> None:
        self.addComputePerDof("x", f"x + {timestep} * v")
        self.addComputePerDof("x1", "x")
        self.addConstrainPositions()
        self.addComputePerDof("v", f"v + (x - x1) / ({timestep})")
        self.addConstrainVelocities()

    def _add_boost(self, timestep: str) -> None:
        self.addComputePerDof("v", f"v + {timestep} * f / m")
        self.addConstrainVelocities()

    def _add_rescaling(self, timestep: str) -> None:
        rescaling = (
            "v * sqrt(A + BC * (R1 ^ 2 + sumRsq) + 2 * sqrt(A * BC) * R1)"
            "; R1 = gaussian"
            "; BC = (1 - A) * kT / mvv"
            f"; A = exp(-friction * {timestep})"
        )
        self.addComputeSum("mvv", "m * v * v")
        self.addComputePerDof("v", rescaling)

    def _sums_of_squared_gaussians(self, num_steps: int) -> np.ndarray:
        sumRsq =  2.0 * self._rng.standard_gamma((self._num_dof - 1) // 2, num_steps)
        if self._num_dof % 2 == 0:
            sumRsq += self._rng.standard_normal(num_steps) ** 2
        return sumRsq

    def setRandomNumberSeed(self, seed: int) -> None:
        """
        This method overrides the :class:`openmm.CustomIntegrator` method to also set
        the seed of the random number generator used to pick numbers from the gamma
        distribution.

        Parameters
        ----------
        seed
            The seed to use for the random number generator.
        """
        self._rng = np.random.default_rng(seed + 2**31)
        super().setRandomNumberSeed(self._rng.integers(-(2**31), 2**31))

    def step(self, steps: int) -> None:
        """
        This method overrides the :class:`openmm.CustomIntegrator` method to include the
        efficient computation of the sum of squares of normally distributed random
        numbers.

        Parameters
        ----------
        steps
            The number of steps to take.
        """
        for sumRsq in self._sums_of_squared_gaussians(steps):
            self.setGlobalVariableByName("sumRsq", sumRsq)
            super().step(1)