Skip to main content
`Ctrl`+`K`

#

# Source code for ase.md.bussi

    # fmt: off

    """Bussi NVT dynamics class."""

    import math

    import numpy as np

    from ase import units
    from ase.md.verlet import VelocityVerlet

    class Bussi(VelocityVerlet):
        """Bussi stochastic velocity rescaling (NVT) molecular dynamics.

        Based on the paper from Bussi et al., J. Chem. Phys. 126, 014101 (2007)
        (also available from https://arxiv.org/abs/0803.4060).
        """

        def __init__(
            self,
            atoms,
            timestep,
            temperature_K,
            taut,
            rng=None,
            **kwargs,
        ):
            """
            Parameters
            ----------
            atoms : Atoms
                The atoms object.
            timestep : float
                The time step in ASE time units.
            temperature_K : float
                The desired temperature, in Kelvin.
            taut : float
                Time constant for Bussi temperature coupling in ASE time units.
            rng : RNG object, optional
                Random number generator, by default numpy.random.
            **kwargs : dict, optional
                Additional arguments are passed to
                :class:~ase.md.md.MolecularDynamics base class.
            """
            super().__init__(atoms, timestep, **kwargs)

            self.temp = temperature_K * units.kB
            self.taut = taut
            if rng is None:
                self.rng = np.random
            else:
                self.rng = rng

            self.ndof = self.atoms.get_number_of_degrees_of_freedom()

            self.target_kinetic_energy = 0.5 * self.temp * self.ndof

            if np.isclose(
                self.atoms.get_kinetic_energy(), 0.0, rtol=0, atol=1e-12
            ):
                raise ValueError(
                    "Initial kinetic energy is zero. "
                    "Please set the initial velocities before running Bussi NVT."
                )

            self._exp_term = math.exp(-self.dt / self.taut)
            self._masses = self.atoms.get_masses()[:, np.newaxis]

            self.transferred_energy = 0.0

        def scale_velocities(self):
            """Do the NVT Bussi stochastic velocity scaling."""
            kinetic_energy = self.atoms.get_kinetic_energy()
            alpha = self.calculate_alpha(kinetic_energy)

            momenta = self.atoms.get_momenta()
            self.atoms.set_momenta(alpha * momenta)

            self.transferred_energy += (alpha**2 - 1.0) * kinetic_energy

        def calculate_alpha(self, kinetic_energy):
            """Calculate the scaling factor alpha using equation (A7)
            from the Bussi paper."""

            energy_scaling_term = (
                (1 - self._exp_term)
                * self.target_kinetic_energy
                / kinetic_energy
                / self.ndof
            )

            # R1 in Eq. (A7)
            noisearray = self.rng.standard_normal(size=(1,))
            # ASE mpi interfaces can only broadcast arrays, not scalars
            self.comm.broadcast(noisearray, 0)
            normal_noise = noisearray[0]

            # \sum_{i=2}^{Nf} R_i^2 in Eq. (A7)
            # 2 * standard_gamma(n / 2) is equal to chisquare(n)
            sum_of_noises = 2.0 * self.rng.standard_gamma(0.5 * (self.ndof - 1))

            return math.sqrt(
                self._exp_term
                + energy_scaling_term * (sum_of_noises + normal_noise**2)
                + 2
                * normal_noise
                * math.sqrt(self._exp_term * energy_scaling_term)
            )

        def step(self, forces=None):
            """Move one timestep forward using Bussi NVT molecular dynamics."""
            self.scale_velocities()
            return super().step(forces)