Skip to main content
`Ctrl`+`K`

#

# Source code for ase.md.nose_hoover_chain

    # fmt: off

    from __future__ import annotations

    import numpy as np
    from scipy.special import exprel

    import ase.units
    from ase import Atoms
    from ase.md.md import MolecularDynamics

    # Coefficients for the fourth-order Suzuki-Yoshida integration scheme
    # Ref: H. Yoshida, Phys. Lett. A 150, 5-7, 262-268 (1990).
    #      https://doi.org/10.1016/0375-9601(90)90092-3
    FOURTH_ORDER_COEFFS = [
        1 / (2 - 2 ** (1 / 3)),
        -(2 ** (1 / 3)) / (2 - 2 ** (1 / 3)),
        1 / (2 - 2 ** (1 / 3)),
    ]

    class NoseHooverChainNVT(MolecularDynamics):
        """Isothermal molecular dynamics with Nose-Hoover chain.

        This implementation is based on the Nose-Hoover chain equations and
        the Liouville-operator derived integrator for non-Hamiltonian systems [1-3].

        - [1] G. J. Martyna, M. L. Klein, and M. E. Tuckerman, J. Chem. Phys. 97,
              2635 (1992). https://doi.org/10.1063/1.463940
        - [2] M. E. Tuckerman, J. Alejandre, R. López-Rendón, A. L. Jochim,
              and G. J. Martyna, J. Phys. A: Math. Gen. 39, 5629 (2006).
              https://doi.org/10.1088/0305-4470/39/19/S18
        - [3] M. E. Tuckerman, Statistical Mechanics: Theory and Molecular
              Simulation, Oxford University Press (2010).

        While the algorithm and notation for the thermostat are largely adapted
        from Ref. [4], the core equations are detailed in the implementation
        note available at
        https://github.com/lan496/lan496.github.io/blob/main/notes/nose_hoover_chain/main.pdf.

        - [4] M. E. Tuckerman, Statistical Mechanics: Theory and Molecular
              Simulation, 2nd ed. (Oxford University Press, 2009).
        """

        def __init__(
            self,
            atoms: Atoms,
            timestep: float,
            temperature_K: float,
            tdamp: float,
            tchain: int = 3,
            tloop: int = 1,
            **kwargs,
        ):
            """
            Parameters
            ----------
            atoms: ase.Atoms
                The atoms object.
            timestep: float
                The time step in ASE time units.
            temperature_K: float
                The target temperature in K.
            tdamp: float
                The characteristic time scale for the thermostat in ASE time units.
                Typically, it is set to 100 times of `timestep`.
            tchain: int
                The number of thermostat variables in the Nose-Hoover chain.
            tloop: int
                The number of sub-steps in thermostat integration.
            **kwargs : dict, optional
                Additional arguments passed to :class:~ase.md.md.MolecularDynamics
                base class.
            """
            super().__init__(
                atoms=atoms,
                timestep=timestep,
                **kwargs,
            )
            assert self.masses.shape == (len(self.atoms), 1)

            num_atoms = self.atoms.get_global_number_of_atoms()
            self._thermostat = NoseHooverChainThermostat(
                num_atoms_global=num_atoms,
                masses=self.masses,
                temperature_K=temperature_K,
                tdamp=tdamp,
                tchain=tchain,
                tloop=tloop,
            )

            # The following variables are updated during self.step()
            self._q = self.atoms.get_positions()
            self._p = self.atoms.get_momenta()

        def step(self) -> None:
            dt2 = self.dt / 2
            self._p = self._thermostat.integrate_nhc(self._p, dt2)
            self._integrate_p(dt2)
            self._integrate_q(self.dt)
            self._integrate_p(dt2)
            self._p = self._thermostat.integrate_nhc(self._p, dt2)

            self._update_atoms()

        def get_conserved_energy(self) -> float:
            """Return the conserved energy-like quantity.

            This method is mainly used for testing.
            """
            conserved_energy = (
                self.atoms.get_potential_energy(force_consistent=True)
                + self.atoms.get_kinetic_energy()
                + self._thermostat.get_thermostat_energy()
            )
            return float(conserved_energy)

        def _update_atoms(self) -> None:
            self.atoms.set_positions(self._q)
            self.atoms.set_momenta(self._p)

        def _get_forces(self) -> np.ndarray:
            self._update_atoms()
            return self.atoms.get_forces(md=True)

        def _integrate_q(self, delta: float) -> None:
            """Integrate exp(i * L_1 * delta)"""
            self._q += delta * self._p / self.masses

        def _integrate_p(self, delta: float) -> None:
            """Integrate exp(i * L_2 * delta)"""
            forces = self._get_forces()
            self._p += delta * forces

    class NoseHooverChainThermostat:
        """Nose-Hoover chain style thermostats.

        See `NoseHooverChainNVT` for the references.
        """
        def __init__(
            self,
            num_atoms_global: int,
            masses: np.ndarray,
            temperature_K: float,
            tdamp: float,
            tchain: int = 3,
            tloop: int = 1,
        ):
            """See `NoseHooverChainNVT` for the parameters."""
            self._num_atoms_global = num_atoms_global
            self._masses = masses  # (len(atoms), 1)
            self._tdamp = tdamp
            self._tchain = tchain
            self._tloop = tloop

            self._kT = ase.units.kB * temperature_K

            assert tchain >= 1
            self._Q = np.zeros(tchain)
            self._Q[0] = 3 * self._num_atoms_global * self._kT * tdamp**2
            self._Q[1:] = self._kT * tdamp**2

            # The following variables are updated during self.step()
            self._eta = np.zeros(self._tchain)
            self._p_eta = np.zeros(self._tchain)

        def get_thermostat_energy(self) -> float:
            """Return energy-like contribution from the thermostat variables."""
            energy = (
                3 * self._num_atoms_global * self._kT * self._eta[0]
                + self._kT * np.sum(self._eta[1:])
                + np.sum(0.5 * self._p_eta**2 / self._Q)
            )
            return float(energy)

        def integrate_nhc(self, p: np.ndarray, delta: float) -> np.ndarray:
            """Integrate exp(i * L_NHC * delta) and update momenta `p`."""
            for _ in range(self._tloop):
                for coeff in FOURTH_ORDER_COEFFS:
                    p = self._integrate_nhc_loop(
                        p, coeff * delta / self._tloop
                    )

            return p

        def _integrate_p_eta_j(self, p: np.ndarray, j: int,
                               delta2: float, delta4: float) -> None:
            if j < self._tchain - 1:
                self._p_eta[j] *= np.exp(
                    -delta4 * self._p_eta[j + 1] / self._Q[j + 1]
                )

            if j == 0:
                g_j = np.sum(p**2 / self._masses) \
                    - 3 * self._num_atoms_global * self._kT
            else:
                g_j = self._p_eta[j - 1] ** 2 / self._Q[j - 1] - self._kT
            self._p_eta[j] += delta2 * g_j

            if j < self._tchain - 1:
                self._p_eta[j] *= np.exp(
                    -delta4 * self._p_eta[j + 1] / self._Q[j + 1]
                )

        def _integrate_eta(self, delta: float) -> None:
            self._eta += delta * self._p_eta / self._Q

        def _integrate_nhc_p(self, p: np.ndarray, delta: float) -> None:
            p *= np.exp(-delta * self._p_eta[0] / self._Q[0])

        def _integrate_nhc_loop(self, p: np.ndarray, delta: float) -> np.ndarray:
            delta2 = delta / 2
            delta4 = delta / 4

            for j in range(self._tchain):
                self._integrate_p_eta_j(p, self._tchain - j - 1, delta2, delta4)
            self._integrate_eta(delta)
            self._integrate_nhc_p(p, delta)
            for j in range(self._tchain):
                self._integrate_p_eta_j(p, j, delta2, delta4)

            return p

    class IsotropicMTKNPT(MolecularDynamics):
        """Isothermal-isobaric molecular dynamics with isotropic volume fluctuations
        by Martyna-Tobias-Klein (MTK) method [1].

        See also `NoseHooverChainNVT` for the references.

        - [1] G. J. Martyna, D. J. Tobias, and M. L. Klein, J. Chem. Phys. 101,
              4177-4189 (1994). https://doi.org/10.1063/1.467468
        """
        def __init__(
            self,
            atoms: Atoms,
            timestep: float,
            temperature_K: float,
            pressure_au: float,
            tdamp: float,
            pdamp: float,
            tchain: int = 3,
            pchain: int = 3,
            tloop: int = 1,
            ploop: int = 1,
            **kwargs,
        ):
            """
            Parameters
            ----------
            atoms: ase.Atoms
                The atoms object.
            timestep: float
                The time step in ASE time units.
            temperature_K: float
                The target temperature in K.
            pressure_au: float
                The external pressure in eV/Ang^3.
            tdamp: float
                The characteristic time scale for the thermostat in ASE time units.
                Typically, it is set to 100 times of `timestep`.
            pdamp: float
                The characteristic time scale for the barostat in ASE time units.
                Typically, it is set to 1000 times of `timestep`.
            tchain: int
                The number of thermostat variables in the Nose-Hoover thermostat.
            pchain: int
                The number of barostat variables in the MTK barostat.
            tloop: int
                The number of sub-steps in thermostat integration.
            ploop: int
                The number of sub-steps in barostat integration.
            **kwargs : dict, optional
                Additional arguments passed to :class:~ase.md.md.MolecularDynamics
                base class.
            """
            super().__init__(
                atoms=atoms,
                timestep=timestep,
                **kwargs,
            )
            assert self.masses.shape == (len(self.atoms), 1)

            if len(atoms.constraints) > 0:
                raise NotImplementedError(
                    "Current implementation does not support constraints"
                )

            self._num_atoms_global = self.atoms.get_global_number_of_atoms()
            self._thermostat = NoseHooverChainThermostat(
                num_atoms_global=self._num_atoms_global,
                masses=self.masses,
                temperature_K=temperature_K,
                tdamp=tdamp,
                tchain=tchain,
                tloop=tloop,
            )
            self._barostat = IsotropicMTKBarostat(
                num_atoms_global=self._num_atoms_global,
                temperature_K=temperature_K,
                pdamp=pdamp,
                pchain=pchain,
                ploop=ploop,
            )

            self._temperature_K = temperature_K
            self._pressure_au = pressure_au

            self._kT = ase.units.kB * self._temperature_K
            self._volume0 = self.atoms.get_volume()
            self._cell0 = np.array(self.atoms.get_cell())

            # The following variables are updated during self.step()
            self._q = self.atoms.get_positions()  # positions
            self._p = self.atoms.get_momenta()  # momenta
            self._eps = 0.0  # volume
            self._p_eps = 0.0  # volume momenta

        def step(self) -> None:
            dt2 = self.dt / 2

            self._p_eps = self._barostat.integrate_nhc_baro(self._p_eps, dt2)
            self._p = self._thermostat.integrate_nhc(self._p, dt2)
            self._integrate_p_cell(dt2)
            self._integrate_p(dt2)
            self._integrate_q(self.dt)
            self._integrate_q_cell(self.dt)
            self._integrate_p(dt2)
            self._integrate_p_cell(dt2)
            self._p = self._thermostat.integrate_nhc(self._p, dt2)
            self._p_eps = self._barostat.integrate_nhc_baro(self._p_eps, dt2)

            self._update_atoms()

        def get_conserved_energy(self) -> float:
            """Return the conserved energy-like quantity.

            This method is mainly used for testing.
            """
            conserved_energy = (
                self.atoms.get_potential_energy(force_consistent=True)
                + self.atoms.get_kinetic_energy()
                + self._thermostat.get_thermostat_energy()
                + self._barostat.get_barostat_energy()
                + self._p_eps * self._p_eps / (2 * self._barostat.W)
                + self._pressure_au * self._get_volume()
            )
            return float(conserved_energy)

        def _update_atoms(self) -> None:
            self.atoms.set_positions(self._q)
            self.atoms.set_momenta(self._p)
            cell = self._cell0 * np.exp(self._eps)
            # Never set scale_atoms=True
            self.atoms.set_cell(cell, scale_atoms=False)

        def _get_volume(self) -> float:
            return self._volume0 * np.exp(3 * self._eps)

        def _get_forces(self) -> np.ndarray:
            self._update_atoms()
            return self.atoms.get_forces(md=True)

        def _get_pressure(self) -> np.ndarray:
            self._update_atoms()
            stress = self.atoms.get_stress(voigt=False, include_ideal_gas=True)
            pressure = -np.trace(stress) / 3
            return pressure

        def _integrate_q(self, delta: float) -> None:
            """Integrate exp(i * L_1 * delta)"""
            x = delta * self._p_eps / self._barostat.W
            self._q = (
                self._q * np.exp(x)
                + self._p * delta / self.masses * exprel(x)
            )

        def _integrate_p(self, delta: float) -> None:
            """Integrate exp(i * L_2 * delta)"""
            x = (1 + 1 / self._num_atoms_global) * self._p_eps * delta \
                    / self._barostat.W
            forces = self._get_forces()
            self._p = self._p * np.exp(-x) + delta * forces * exprel(-x)

        def _integrate_q_cell(self, delta: float) -> None:
            """Integrate exp(i * L_(epsilon, 1) * delta)"""
            self._eps += delta * self._p_eps / self._barostat.W

        def _integrate_p_cell(self, delta: float) -> None:
            """Integrate exp(i * L_(epsilon, 2) * delta)"""
            pressure = self._get_pressure()
            volume = self._get_volume()
            G = (
                3 * volume * (pressure - self._pressure_au)
                + np.sum(self._p**2 / self.masses) / self._num_atoms_global
            )
            self._p_eps += delta * G

    class IsotropicMTKBarostat:
        """MTK barostat for isotropic volume fluctuations.

        See `IsotropicMTKNPT` for the references.
        """
        def __init__(
            self,
            num_atoms_global: int,
            temperature_K: float,
            pdamp: float,
            pchain: int = 3,
            ploop: int = 1,
        ):
            self._num_atoms_global = num_atoms_global
            self._pdamp = pdamp
            self._pchain = pchain
            self._ploop = ploop

            self._kT = ase.units.kB * temperature_K

            self._W = (3 * self._num_atoms_global + 3) * self._kT * self._pdamp**2

            assert pchain >= 1
            self._R = np.zeros(self._pchain)
            self._R[0] = 9 * self._kT * self._pdamp**2
            self._R[1:] = self._kT * self._pdamp**2

            self._xi = np.zeros(self._pchain)  # barostat coordinates
            self._p_xi = np.zeros(self._pchain)

        @property
        def W(self) -> float:
            """Virtual mass for barostat momenta `p_xi`."""
            return self._W

        def get_barostat_energy(self) -> float:
            """Return energy-like contribution from the barostat variables."""
            energy = (
                + np.sum(0.5 * self._p_xi**2 / self._R)
                + self._kT * np.sum(self._xi)
            )
            return float(energy)

        def integrate_nhc_baro(self, p_eps: float, delta: float) -> float:
            """Integrate exp(i * L_NHC-baro * delta)"""
            for _ in range(self._ploop):
                for coeff in FOURTH_ORDER_COEFFS:
                    p_eps = self._integrate_nhc_baro_loop(
                        p_eps, coeff * delta / self._ploop
                    )
            return p_eps

        def _integrate_nhc_baro_loop(self, p_eps: float, delta: float) -> float:
            delta2 = delta / 2
            delta4 = delta / 4

            for j in range(self._pchain):
                self._integrate_p_xi_j(p_eps, self._pchain - j - 1, delta2, delta4)
            self._integrate_xi(delta)
            p_eps = self._integrate_nhc_p_eps(p_eps, delta)
            for j in range(self._pchain):
                self._integrate_p_xi_j(p_eps, j, delta2, delta4)

            return p_eps

        def _integrate_p_xi_j(self, p_eps: float, j: int,
                              delta2: float, delta4: float) -> None:
            if j < self._pchain - 1:
                self._p_xi[j] *= np.exp(
                    -delta4 * self._p_xi[j + 1] / self._R[j + 1]
                )

            if j == 0:
                g_j = p_eps ** 2 / self._W - self._kT
            else:
                g_j = self._p_xi[j - 1] ** 2 / self._R[j - 1] - self._kT
            self._p_xi[j] += delta2 * g_j

            if j < self._pchain - 1:
                self._p_xi[j] *= np.exp(
                    -delta4 * self._p_xi[j + 1] / self._R[j + 1]
                )

        def _integrate_xi(self, delta: float) -> None:
            self._xi += delta * self._p_xi / self._R

        def _integrate_nhc_p_eps(self, p_eps: float, delta: float) -> float:
            p_eps_new = p_eps * float(
                np.exp(-delta * self._p_xi[0] / self._R[0])
            )
            return p_eps_new

    class MTKNPT(MolecularDynamics):
        """Isothermal-isobaric molecular dynamics with volume-and-cell fluctuations
        by Martyna-Tobias-Klein (MTK) method [1].

        See also `NoseHooverChainNVT` for the references.

        - [1] G. J. Martyna, D. J. Tobias, and M. L. Klein, J. Chem. Phys. 101,
              4177-4189 (1994). https://doi.org/10.1063/1.467468
        """
        def __init__(
            self,
            atoms: Atoms,
            timestep: float,
            temperature_K: float,
            pressure_au: float,
            tdamp: float,
            pdamp: float,
            tchain: int = 3,
            pchain: int = 3,
            tloop: int = 1,
            ploop: int = 1,
            **kwargs,
        ):
            """
            Parameters
            ----------
            atoms: ase.Atoms
                The atoms object.
            timestep: float
                The time step in ASE time units.
            temperature_K: float
                The target temperature in K.
            pressure_au: float
                The external pressure in eV/Ang^3.
            tdamp: float
                The characteristic time scale for the thermostat in ASE time units.
                Typically, it is set to 100 times of `timestep`.
            pdamp: float
                The characteristic time scale for the barostat in ASE time units.
                Typically, it is set to 1000 times of `timestep`.
            tchain: int
                The number of thermostat variables in the Nose-Hoover thermostat.
            pchain: int
                The number of barostat variables in the MTK barostat.
            tloop: int
                The number of sub-steps in thermostat integration.
            ploop: int
                The number of sub-steps in barostat integration.
            **kwargs : dict, optional
                Additional arguments passed to :class:~ase.md.md.MolecularDynamics
                base class.
            """
            super().__init__(
                atoms=atoms,
                timestep=timestep,
                **kwargs,
            )
            assert self.masses.shape == (len(self.atoms), 1)

            if len(atoms.constraints) > 0:
                raise NotImplementedError(
                    "Current implementation does not support constraints"
                )

            self._num_atoms_global = self.atoms.get_global_number_of_atoms()
            self._thermostat = NoseHooverChainThermostat(
                num_atoms_global=self._num_atoms_global,
                masses=self.masses,
                temperature_K=temperature_K,
                tdamp=tdamp,
                tchain=tchain,
                tloop=tloop,
            )
            self._barostat = MTKBarostat(
                num_atoms_global=self._num_atoms_global,
                temperature_K=temperature_K,
                pdamp=pdamp,
                pchain=pchain,
                ploop=ploop,
            )

            self._temperature_K = temperature_K
            self._pressure_au = pressure_au

            self._kT = ase.units.kB * self._temperature_K

            # The following variables are updated during self.step()
            self._q = self.atoms.get_positions()  # positions
            self._p = self.atoms.get_momenta()  # momenta
            self._h = np.array(self.atoms.get_cell())  # cell
            self._p_g = np.zeros((3, 3))  # cell momenta

        def step(self) -> None:
            dt2 = self.dt / 2

            self._p_g = self._barostat.integrate_nhc_baro(self._p_g, dt2)
            self._p = self._thermostat.integrate_nhc(self._p, dt2)
            self._integrate_p_cell(dt2)
            self._integrate_p(dt2)
            self._integrate_q(self.dt)
            self._integrate_q_cell(self.dt)
            self._integrate_p(dt2)
            self._integrate_p_cell(dt2)
            self._p = self._thermostat.integrate_nhc(self._p, dt2)
            self._p_g = self._barostat.integrate_nhc_baro(self._p_g, dt2)

            self._update_atoms()

        def get_conserved_energy(self) -> float:
            conserved_energy = (
                self.atoms.get_total_energy()
                + self._thermostat.get_thermostat_energy()
                + self._barostat.get_barostat_energy()
                + np.trace(self._p_g.T @ self._p_g) / (2 * self._barostat.W)
                + self._pressure_au * self._get_volume()
            )
            return float(conserved_energy)

        def _update_atoms(self) -> None:
            self.atoms.set_positions(self._q)
            self.atoms.set_momenta(self._p)
            self.atoms.set_cell(self._h, scale_atoms=False)

        def _get_volume(self) -> float:
            return np.abs(np.linalg.det(self._h))

        def _get_forces(self) -> np.ndarray:
            self._update_atoms()
            return self.atoms.get_forces(md=True)

        def _get_stress(self) -> np.ndarray:
            self._update_atoms()
            stress = self.atoms.get_stress(voigt=False, include_ideal_gas=True)
            return -stress

        def _integrate_q(self, delta: float) -> None:
            """Integrate exp(i * L_1 * delta)"""
            # eigvals: (3-eigvec), U: (3-xyz, 3-eigvec)
            eigvals, U = np.linalg.eigh(self._p_g)
            x = self._q @ U  # (num_atoms, 3-eigvec)
            y = self._p @ U  # (num_atoms, 3-eigvec)
            sol = (
                x * np.exp(eigvals * delta / self._barostat.W)[None, :]
                + delta * y / self.masses * exprel(
                    eigvals * delta / self._barostat.W
                )[None, :]
            )  # (num_atoms, 3-eigvec)
            self._q = sol @ U.T

        def _integrate_p(self, delta: float) -> None:
            """Integrate exp(i * L_2 * delta)"""
            forces = self._get_forces()  # (num_atoms, 3-xyz)

            # eigvals: (3-eigvec), U: (3-xyz, 3-eigvec)
            eigvals, U = np.linalg.eigh(self._p_g)
            kappas = eigvals \
                + np.trace(self._p_g) / (3 * self._num_atoms_global)  # (3-eigvec)
            y = self._p @ U  # (num_atoms, 3-eigvec)
            sol = (
                y * np.exp(-kappas * delta / self._barostat.W)[None, :]
                + delta * (forces @ U) * exprel(
                    -kappas * delta / self._barostat.W
                )[None, :]
            )  # (num_atoms, 3-eigvec)
            self._p = sol @ U.T

        def _integrate_q_cell(self, delta: float) -> None:
            """Integrate exp(i * L_(g, 1) * delta)"""
            # U @ np.diag(eigvals) @ U.T = self._p_g
            # eigvals: (3-eigvec), U: (3-xyz, 3-eigvec)
            eigvals, U = np.linalg.eigh(self._p_g)
            n = self._h @ U  # (3-axis, 3-eigvec)
            sol = n * np.exp(
                eigvals * delta / self._barostat.W
            )[None, :]  # (3-axis, 3-eigvec)
            self._h = sol @ U.T

        def _integrate_p_cell(self, delta: float) -> None:
            """Integrate exp(i * L_(g, 2) * delta)"""
            stress = self._get_stress()
            G = (
                self._get_volume() * (stress - self._pressure_au * np.eye(3))
                + np.sum(self._p**2 / self.masses) / (3 * self._num_atoms_global)
                    * np.eye(3)
            )
            self._p_g += delta * G

    class MTKBarostat:
        """MTK barostat for volume-and-cell fluctuations.

        See `MTKNPT` for the references.
        """
        def __init__(
            self,
            num_atoms_global: int,
            temperature_K: float,
            pdamp: float,
            pchain: int = 3,
            ploop: int = 1,
        ):
            self._num_atoms_global = num_atoms_global
            self._pdamp = pdamp
            self._pchain = pchain
            self._ploop = ploop

            self._kT = ase.units.kB * temperature_K

            self._W = (self._num_atoms_global + 1) * self._kT * self._pdamp**2

            assert pchain >= 1
            self._R = np.zeros(self._pchain)
            cell_dof = 9  # TODO:
            self._R[0] = cell_dof * self._kT * self._pdamp**2
            self._R[1:] = self._kT * self._pdamp**2

            self._xi = np.zeros(self._pchain)  # barostat coordinates
            self._p_xi = np.zeros(self._pchain)

        @property
        def W(self) -> float:
            return self._W

        def get_barostat_energy(self) -> float:
            energy = (
                np.sum(self._p_xi**2 / self._R) / 2
                + 9 * self._kT * self._xi[0]
                + self._kT * np.sum(self._xi[1:])
            )
            return float(energy)

        def integrate_nhc_baro(self, p_g: np.ndarray, delta: float) -> np.ndarray:
            """Integrate exp(i * L_NHC-baro * delta)"""
            for _ in range(self._ploop):
                for coeff in FOURTH_ORDER_COEFFS:
                    p_g = self._integrate_nhc_baro_loop(
                        p_g, coeff * delta / self._ploop
                    )
            return p_g

        def _integrate_nhc_baro_loop(
            self, p_g: np.ndarray, delta: float
        ) -> np.ndarray:
            delta2 = delta / 2
            delta4 = delta / 4

            for j in range(self._pchain):
                self._integrate_p_xi_j(p_g, self._pchain - j - 1, delta2, delta4)
            self._integrate_xi(delta)
            self._integrate_nhc_p_eps(p_g, delta)
            for j in range(self._pchain):
                self._integrate_p_xi_j(p_g, j, delta2, delta4)

            return p_g

        def _integrate_p_xi_j(
            self, p_g: np.ndarray, j: int, delta2: float, delta4: float
        ) -> None:
            if j < self._pchain - 1:
                self._p_xi[j] *= np.exp(
                    -delta4 * self._p_xi[j + 1] / self._R[j + 1]
                )

            if j == 0:
                # TODO: do we need to substitute 9 with cell_dof?
                g_j = np.trace(p_g.T @ p_g) / self._W - 9 * self._kT
            else:
                g_j = self._p_xi[j - 1] ** 2 / self._R[j - 1] - self._kT
            self._p_xi[j] += delta2 * g_j

            if j < self._pchain - 1:
                self._p_xi[j] *= np.exp(
                    -delta4 * self._p_xi[j + 1] / self._R[j + 1]
                )

        def _integrate_xi(self, delta: float) -> None:
            for j in range(self._pchain):
                self._xi[j] += delta * self._p_xi[j] / self._R[j]

        def _integrate_nhc_p_eps(self, p_g: np.ndarray, delta: float) -> None:
            p_g *= np.exp(-delta * self._p_xi[0] / self._R[0])