kups.md.integrators ¶

@dataclass
class CSVRStep[State](Propagator[State]):
    r"""Canonical Sampling through Velocity Rescaling (CSVR) thermostat step.

    Implements the Bussi-Donadio-Parrinello algorithm for canonical sampling
    by stochastically rescaling velocities to maintain the target temperature.
    This produces correct canonical ensemble sampling unlike deterministic
    velocity rescaling (Berendsen thermostat).

    The scaling factor $\alpha^2$ is sampled from the conditional distribution:

    $$\alpha^2 \sim (K'/K) \text{ where } K' \text{ follows the target kinetic energy distribution}$$

    The algorithm uses:

    $$\alpha^2 = c_1 + c_2(R_1^2 + R_2) + 2R_1\sqrt{c_1 c_2}$$

    where:

    - $c_1 = e^{-\Delta t/\tau}$ — exponential decay factor
    - $c_2 = (1-c_1) \cdot K_{\text{target}}/(K_{\text{current}} \cdot N_{\text{dof}})$ — correction factor
    - $R_1 \sim \mathcal{N}(0,1)$ — Gaussian random variable
    - $R_2 \sim \chi^2(N_{\text{dof}}-1)$ — chi-squared random variable

    Type Parameters:
        State: Simulation state type

    Attributes:
        particles: Lens to get/set indexed particle data (momenta $\\mathbf{p}$, masses $m$)
        systems: View to extract system data (time step $\\Delta t$, temperature $T$,
            degrees of freedom $N_{\\text{dof}}$, thermostat time constant $\\tau$)

    References:
        Bussi, G., Donadio, D., & Parrinello, M. (2007). Canonical sampling
        through velocity rescaling. J. Chem. Phys., 126(1), 014101.
        DOI: 10.1063/1.2408420
    """

    particles: Lens[State, Table[ParticleId, IsCSVRParticleData]] = field(static=True)
    systems: View[State, Table[SystemId, IsCSVRParams]] = field(static=True)

    def __call__(self, key: Array, state: State) -> State:
        """Apply CSVR stochastic velocity rescaling.

        Args:
            key: JAX PRNG key for generating random noise
            state: Current simulation state

        Returns:
            Updated state with rescaled momenta matching target temperature distribution
        """
        # Extract parameters
        system = self.systems(state)
        particles = self.particles.get(state)
        # Δt: timestep [time]
        timestep = system.data.time_step
        # kT: thermal energy [energy]
        target_thermal_energy = system.data.temperature * BOLTZMANN_CONSTANT
        # τ: thermostat time constant [time]
        thermostat_timescale = system.data.thermostat_time_constant
        # N_dof: degrees of freedom [dimensionless]
        # TODO: Update once we have constraints that could limit the degrees of freedom
        degrees_of_freedom = particles.data.system.counts.data * 3 - 3

        # Compute current kinetic energy from particles
        per_particle_ke = particle_kinetic_energy(
            particles.data.momenta, particles.data.masses
        )
        # K: total kinetic energy per system [energy]
        kinetic_energy_current = particles.data.system.sum_over(per_particle_ke).data
        # K_target = N_dof·kT/2 [energy]
        kinetic_energy_target = degrees_of_freedom * target_thermal_energy / 2

        # Generate random numbers for scaling
        key1, key2 = jax.random.split(key)
        # R₁ ~ N(0,1) [dimensionless]
        gaussian_noise = jax.random.normal(key1, dtype=float)

        # R₂ ~ χ²(N_dof-1) [dimensionless]
        dof_minus_one = degrees_of_freedom - 1
        chi_squared_noise = jnp.where(
            dof_minus_one > 0,
            jax.random.chisquare(key2, df=dof_minus_one, dtype=float),
            0.0,
        )

        # CSVR scaling coefficients
        # c₁ = e^(-Δt/τ) [dimensionless]
        exponential_decay = jnp.exp(-timestep / thermostat_timescale)
        # c₂ = (1-c₁)·K_target/(K_current·N_dof) [dimensionless]
        correction_factor = (
            (1 - exponential_decay)
            * kinetic_energy_target
            / (kinetic_energy_current * degrees_of_freedom)
        )

        # α² = c₁ + c₂(R₁² + R₂) + 2R₁√(c₁c₂) [dimensionless]
        scaling_squared = (
            exponential_decay
            + correction_factor * (gaussian_noise**2 + chi_squared_noise)
            + 2 * gaussian_noise * jnp.sqrt(exponential_decay * correction_factor)
        )
        # α = √(α²), ensure non-negative [dimensionless]
        velocity_scale = jnp.sqrt(jnp.maximum(scaling_squared, 0.0))

        # Scale momenta by system
        scale_per_system = velocity_scale[particles.data.system.indices]
        new_momenta = particles.data.momenta * scale_per_system[..., None]

        assert new_momenta.shape == particles.data.momenta.shape
        return (
            self.particles.bind(state).focus(lambda x: x.data.momenta).set(new_momenta)
        )

@dataclass
class CellMomentumKick[State](Propagator[State]):
    r"""B$^h$ kick (Gao Algorithm lines 2 & 13).

    Updates the cell-momentum tensor by the pressure-deviation virial. In
    paper convention (Eq. 8d) the deterministic term is

    $$\text{kick}^{\text{paper}} = \det(h)\,(P_{\text{ins}} - P_0 I)\,h^{-T}
                                   - \chi k_B T\, h^{-T}.$$

    In kUPS convention $V = h^T$, so $h^{-T} = V^{-1}$, and $p^h_{\text{kUPS}}
    = (p^h_{\text{paper}})^T$. Transposing the kick and projecting to the
    lower triangle gives the kUPS-side update applied here:

    $$p^h \leftarrow p^h + \Delta t \cdot \text{tril}\!\left(
        \det(V)\,V^{-T}(P_{\text{ins}} - P_0 I) - \chi k_B T\, V^{-T}
    \right),$$

    The integration uses $\Delta t / 2$ — the algorithm calls this kick four
    times per step, twice at the start and twice at the end of the BAOAB
    palindrome, each at a half time-step. $P_{\text{ins}}$ is the symmetric
    pressure tensor from
    [instantaneous_pressure_tensor][kups.md.observables.instantaneous_pressure_tensor]
    and includes the lattice-gradient contribution. $\chi = 1$ for the
    3D lower-triangular cell parameterisation (Appendix A, $\nu/d-1$).
    """

    particles: Lens[State, Table[ParticleId, _BarostatParticleData]] = field(
        static=True
    )
    systems: Lens[State, Table[SystemId, IsBAOABNPTLangevinSystem]] = field(static=True)

    def __call__(self, key: Array, state: State) -> State:
        del key  # Deterministic step
        sys = self.systems.get(state)
        par = self.particles.get(state)
        params = sys.data.integrator_params

        V_inv = sys.data.cell.inverse_vectors  # lower-tri
        V_det = sys.data.cell.volume  # (n,)
        chi_kT = _GAO_CHI * params.temperature * BOLTZMANN_CONSTANT  # (n,)
        P_target = params.target_pressure  # (n,)
        P_ins = instantaneous_pressure_tensor(par, sys)  # (n, 3, 3) symmetric
        P_diff = P_ins - P_target[..., None, None] * jnp.eye(3)

        # Paper-side kick = det(h)·(P_diff @ V_inv) − χkT·V_inv.
        # kUPS-side kick = (paper kick)^T  (since kUPS p^h = paper p^h transposed).
        kick_paper = (
            V_det[..., None, None] * (P_diff @ V_inv) - chi_kT[..., None, None] * V_inv
        )
        kick_kups = jnp.tril(jnp.swapaxes(kick_paper, -1, -2))

        # Half time-step: this kick fires twice per BAOAB step on each side.
        dt_half = params.time_step[..., None, None] / 2
        new_cell_momentum = sys.data.cell_momentum + dt_half * kick_kups
        return self.systems.focus(lambda x: x.data.cell_momentum).set(
            state, new_cell_momentum
        )

@dataclass
class CellPositionStep[State](Propagator[State]):
    r"""A$^h$ drift (Gao Algorithm lines 4 & 9): cell-only drift.

    Implements the paper's :math:`\mathcal{F}^h_K = \sum_{\alpha\beta}
    (p^h_{\alpha\beta}/M_{\alpha\beta})\,\partial/\partial h_{\alpha\beta}`,
    which moves only $h$:

    $$V \leftarrow V + \frac{\Delta t}{2}\,p^h / M.$$

    **Atomic positions are not modified here.** The convective response of
    particle positions to cell motion ($\dot{h} h^{-1} r$, paper Eq. 8a) is
    handled in :class:`CoupledPositionStep` via the analytic ODE solver;
    rescaling positions in this step too would double-count the convective
    term and break the symplectic Trotter splitting.

    The new cell is reconstructed as a :class:`TriclinicFrame` from the
    drifted lower-triangular vectors, which forces a fresh ``volume`` cache.
    """

    systems: Lens[State, Table[SystemId, IsBAOABNPTLangevinSystem]] = field(static=True)

    def __call__(self, key: Array, state: State) -> State:
        del key  # Deterministic step
        sys = self.systems.get(state)
        params = sys.data.integrator_params
        V_old = sys.data.cell.vectors  # lower-tri
        V_dot = _cell_velocity(sys.data.cell_momentum, params.barostat_mass)
        # Half time-step (BAOAB: drifts twice per side, each at Δt/2).
        # V_old, V_dot both lower-tri ⇒ V_new lower-tri.
        V_new = V_old + (params.time_step[..., None, None] / 2) * V_dot
        # Reconstruct TriclinicFrame so cached volume refreshes.
        return self.systems.focus(lambda x: x.data.cell.frame).apply(
            state, lambda frame: frame.from_matrix(V_new)
        )

@dataclass
class CellStochasticStep[State](Propagator[State]):
    r"""O$^h$ Ornstein–Uhlenbeck thermostat on the cell-momentum tensor
    (Gao Algorithm line 6).

    Exact OU solution per lower-triangular component:

    $$p^h_{\alpha\beta} \leftarrow e^{-\gamma_{\alpha\beta}\Delta t}\,p^h_{\alpha\beta}
        + \sqrt{M_{\alpha\beta}\,k_B T\,(1 - e^{-2\gamma_{\alpha\beta}\Delta t})}\,R_{\alpha\beta},$$

    with $R_{\alpha\beta} \sim \mathcal{N}(0, 1)$. Noise is masked to the
    lower triangle so strict-upper components of $p^h$ remain zero.
    """

    systems: Lens[State, Table[SystemId, IsBAOABNPTLangevinSystem]] = field(static=True)

    def __call__(self, key: Array, state: State) -> State:
        sys = self.systems.get(state)
        params = sys.data.integrator_params
        gamma = params.barostat_friction
        M = params.barostat_mass
        kT = params.temperature * BOLTZMANN_CONSTANT
        dt = params.time_step
        damp = jnp.exp(-gamma * dt[..., None, None])
        noise_amp = jnp.sqrt(M * kT[..., None, None] * (1.0 - damp**2))
        R = jnp.tril(jax.random.normal(key, M.shape, dtype=M.dtype))
        new_cell_momentum = jnp.tril(damp * sys.data.cell_momentum + noise_amp * R)
        return self.systems.focus(lambda x: x.data.cell_momentum).set(
            state, new_cell_momentum
        )

@dataclass
class CoupledMomentumStep[State](Propagator[State]):
    r"""B kick with cell coupling (Gao Algorithm lines 3 & 11).

    Solves $\dot{p}_i = F_i - h^{-\top}\dot{h}^{\top} p_i$ exactly over the
    params time-step via :func:`solve_affine_ode`. In kUPS row-vector form
    the coupling becomes $\dot{p}_{\text{row}} = F_{\text{row}} - p_{\text{row}}\cdot M^{\top}$
    with $M = V^{-1}\dot{V}$ (lower-triangular); the equivalent column-vector
    ODE has $A = -M$.

    Half time-step: the BAOAB palindrome calls this kick four times per
    full step, each at $\Delta t / 2$.
    """

    particles: Lens[State, Table[ParticleId, _BarostatParticleData]] = field(
        static=True
    )
    systems: Lens[State, Table[SystemId, IsBAOABNPTLangevinSystem]] = field(static=True)

    def __call__(self, key: Array, state: State) -> State:
        del key
        sys = self.systems.get(state)
        par_lens = self.particles.bind(state)
        par = par_lens.get()
        params = sys.data.integrator_params
        V_dot = _cell_velocity(sys.data.cell_momentum, params.barostat_mass)
        M_pos = sys.data.cell.inverse_vectors @ V_dot  # lower-tri
        A_per_p = -M_pos[par.data.system.indices]  # column-form A
        dt_half_per_p = params.time_step[par.data.system.indices] / 2
        new_momenta = jax.vmap(solve_affine_ode)(
            A_per_p, par.data.forces, par.data.momenta, dt_half_per_p
        )
        return par_lens.focus(lambda p: p.data.momenta).set(new_momenta)

@dataclass
class CoupledPositionStep[State](Propagator[State]):
    r"""A drift with cell coupling (Gao Algorithm lines 5 & 8).

    Solves $\dot{r}_i = p_i/m_i + \dot{h}h^{-1} r_i$ exactly over the params
    time-step via :func:`solve_affine_ode`. In kUPS row-vector form the
    coupling becomes $\dot{r}_{\text{row}} = v_{\text{row}} + r_{\text{row}}\cdot M$
    with $M = V^{-1}\dot{V}$ (lower-triangular); the column-vector ODE has
    $A = M^{\top}$ (upper-triangular).

    Half time-step: the BAOAB palindrome calls this drift four times per
    full step, each at $\Delta t / 2$.
    """

    particles: Lens[State, Table[ParticleId, _BarostatParticleData]] = field(
        static=True
    )
    systems: Lens[State, Table[SystemId, IsBAOABNPTLangevinSystem]] = field(static=True)

    def __call__(self, key: Array, state: State) -> State:
        del key
        sys = self.systems.get(state)
        par_lens = self.particles.bind(state)
        par = par_lens.get()
        params = sys.data.integrator_params
        V_dot = _cell_velocity(sys.data.cell_momentum, params.barostat_mass)
        M_pos = sys.data.cell.inverse_vectors @ V_dot  # lower-tri
        A_per_p = jnp.swapaxes(M_pos, -1, -2)[par.data.system.indices]  # upper-tri
        dt_half_per_p = params.time_step[par.data.system.indices] / 2
        v = par.data.momenta / par.data.masses[..., None]
        new_positions = jax.vmap(solve_affine_ode)(
            A_per_p, v, par.data.positions, dt_half_per_p
        )
        return par_lens.focus(lambda p: p.data.positions).set(new_positions)

@runtime_checkable
class Flow[State, PyTree](Protocol):
    """Protocol for position update flows with boundary conditions.

    A flow defines how positions evolve under velocity updates, potentially
    including boundary conditions like periodic wrapping or reflections.
    """

    def __call__(
        self, state: State, dt: Time, primal: PyTree, tangent: PyTree
    ) -> PyTree:
        """Apply flow to update positions.

        Args:
            state: Current simulation state.
            dt: Timestep $\\Delta t$ (units: time).
            primal: Position $\\mathbf{r}$ (units: length).
            tangent: Velocity $\\mathbf{v}$  (units: length/time).

        Returns:
            Updated position (units: length).
        """
        ...

@runtime_checkable
class IsBAOABLangevinParams(
    HasTimeStep, HasTemperature, HasFrictionCoefficient, Protocol
):
    r"""Integrator-params shape for :func:`make_baoab_langevin_step`."""

@runtime_checkable
class IsBAOABNPTLangevinParams(
    HasTimeStep,
    HasTemperature,
    HasFrictionCoefficient,
    HasTargetPressure,
    HasPressureCouplingTime,
    HasCompressibility,
    HasBarostatMass,
    HasBarostatFriction,
    Protocol,
):
    r"""Integrator-params shape for :func:`make_baoab_npt_langevin_step`."""

@runtime_checkable
class IsBAOABNPTLangevinSystem(
    IsBAOABNPTLangevinSystemBase,
    HasIntegratorParams[IsBAOABNPTLangevinParams],
    Protocol,
):
    r"""NPT Langevin MD system row: :class:`IsBAOABNPTLangevinSystemBase` plus integrator params."""

@runtime_checkable
class IsBAOABNPTLangevinSystemBase(HasCell[Periodic3D], HasCellMomentum, Protocol):
    r"""Params-free NPT Langevin MD system row: periodic cell, cell-momentum tensor,
    and a cell-gradient leaf for the virial."""

    @property
    def cell_gradients(self) -> Cell[Periodic3D]: ...

@runtime_checkable
class IsCSVRNPTParams(
    HasTimeStep,
    HasTemperature,
    HasThermostatTimeConstant,
    HasTargetPressure,
    HasPressureCouplingTime,
    HasCompressibility,
    HasMinimumScaleFactor,
    Protocol,
):
    r"""Integrator-params shape for :func:`make_csvr_npt_step`."""

@runtime_checkable
class IsCSVRParams(
    HasTimeStep,
    HasTemperature,
    HasThermostatTimeConstant,
    Protocol,
):
    r"""Integrator-params shape for :func:`make_csvr_step`."""

@runtime_checkable
class IsMDSystem[P](HasCell[Periodic3D], HasIntegratorParams[P], Protocol):
    r"""Protocol for an MD system row: a periodic cell plus bundled integrator parameters."""

@runtime_checkable
class IsMDSystemNPT(IsMDSystemNPTBase, HasIntegratorParams[IsCSVRNPTParams], Protocol):
    r"""NPT MD system row: :class:`IsMDSystemNPTBase` plus bundled integrator parameters."""

@runtime_checkable
class IsMDSystemNPTBase(HasCell[Periodic3D], Protocol):
    r"""Params-free NPT MD system row: a periodic cell plus the barostat cell-gradient leaf."""

    @property
    def cell_gradients(self) -> Cell[Periodic3D]: ...

@runtime_checkable
class IsVerletParams(HasTimeStep, Protocol):
    r"""Integrator-params shape for :func:`make_velocity_verlet_step`."""

@dataclass
class MomentumStep[State](Propagator[State]):
    """Update momenta using forces according to Newton's second law.

    Implements the 'B' operator in splitting schemes, applying forces to
    update particle momenta. This is the dynamical update step that couples
    to the potential energy landscape.

    The momentum update follows:

    $$\\mathbf{p}(t+\\Delta t) = \\mathbf{p}(t) + \\mathbf{F}(t) \\cdot \\Delta t$$

    where $\\mathbf{F} = -\\nabla U$ is the force derived from potential energy $U$.

    Type Parameters:
        State: Simulation state type

    Attributes:
        particles: Lens to get/set indexed particle data (momenta $\\mathbf{p}$, forces $\\mathbf{F}$)
        systems: View to extract system data with time step $\\Delta t$
    """

    particles: Lens[State, Table[ParticleId, IsMomentumStepData]] = field(static=True)
    systems: View[State, Table[SystemId, HasTimeStep]] = field(static=True)

    def __call__(self, key: Array, state: State) -> State:
        """Apply momentum update step.

        Args:
            key: JAX PRNG key (unused in this deterministic step).
            state: Current simulation state.

        Returns:
            Updated state with new momenta.
        """
        del key  # Deterministic step
        # Extract current state
        particle_lens = self.particles.bind(state)
        particles = particle_lens.get()
        sys = self.systems(state)[particles.data.system]

        new_momenta = (
            particles.data.momenta + particles.data.forces * sys.time_step[..., None]
        )
        assert new_momenta.shape == particles.data.momenta.shape
        return particle_lens.focus(lambda x: x.data.momenta).set(new_momenta)

@dataclass
class PositionStep[State](Propagator[State]):
    """Update positions using velocities in molecular dynamics.

    Implements the 'A' operator in splitting schemes, propagating positions
    forward in time using the current velocities. This is the kinematic update
    step in velocity Verlet and related integrators.

    The position update follows:

    $$\\mathbf{r}(t+\\Delta t) = \\mathbf{r}(t) + \\mathbf{v}(t) \\cdot \\Delta t$$

    where $\\mathbf{v} = \\mathbf{p}/m$ is the velocity derived from momentum.

    Type Parameters:
        State: Simulation state type

    Attributes:
        particles: Lens to get/set indexed particle data (momenta $\\mathbf{p}$, positions $\\mathbf{r}$, masses $m$)
        systems: View to extract system data with time step $\\Delta t$
        flow: Flow operator defining how positions evolve (handles boundary conditions)
    """

    particles: Lens[State, Table[ParticleId, _PositionStepData]] = field(static=True)
    systems: View[State, Table[SystemId, HasTimeStep]] = field(static=True)
    flow: Flow[State, Array] = field(static=True)

    def __call__(self, key: Array, state: State) -> State:
        """Apply position update step.

        Args:
            key: JAX PRNG key (unused in this deterministic step).
            state: Current simulation state.

        Returns:
            Updated state with new positions.
        """
        del key  # Deterministic step
        # Extract current state
        particle_lens = self.particles.bind(state)
        particles = particle_lens.get()
        sys = self.systems(state)[particles.data.system]
        # Update particles: r_new = r + (p/m)·Δt
        velocity = particles.data.momenta / particles.data.masses[..., None]
        new_positions = self.flow(
            state, sys.time_step, particles.data.positions, velocity
        )
        assert new_positions.shape == particles.data.positions.shape
        return particle_lens.focus(lambda x: x.data.positions).set(new_positions)

@dataclass
class StochasticCellRescalingStep[State](Propagator[State]):
    """Stochastic cell rescaling barostat for NPT ensemble sampling.

    Implements the isotropic stochastic cell rescaling algorithm (Bernetti & Bussi, 2020)
    that correctly samples the NPT ensemble. This first-order barostat includes a
    stochastic term to ensure proper volume fluctuations, unlike the Berendsen
    barostat which artificially suppresses fluctuations.

    The algorithm scales both the simulation box and particle positions by a
    factor $\\mu$ determined by:

    $$\\mu \\approx 1 + \\frac{\\Delta t}{\\tau_P} \\beta (P - P_0) + \\sqrt{\\frac{2k_B T \\beta \\Delta t}{\\tau_P V}} \\, R$$

    where:

    - $\\tau_P$ = pressure coupling time constant
    - $P$ = instantaneous pressure
    - $P_0$ = target pressure
    - $\\beta$ = isothermal compressibility
    - $k_B T$ = thermal energy
    - $V$ = box volume
    - $R \\sim \\mathcal{N}(0,1)$ = Gaussian random noise

    The scaling is applied to both box and positions:

    $$\\mathbf{L}_{\\text{new}} = \\mu \\mathbf{L}, \\quad \\mathbf{r}_{\\text{new}} = \\mu \\mathbf{r}$$

    **Important:** The [Cell][kups.core.cell.Cell] must be reconstructed after
    scaling to ensure the cached volume is recomputed correctly.

    Type Parameters:
        State: Simulation state type

    Attributes:
        particles: Lens to get/set indexed particle data (positions $\\mathbf{r}$, momenta $\\mathbf{p}$, masses $m$)
        systems: Lens to get/set system data (lattice vectors $\\mathbf{L}$, stress tensor $\\mathbf{W}$,
            time step $\\Delta t$, temperature $T$, target pressure $P_0$,
            barostat time constant $\\tau_P$, compressibility $\\beta$, minimum scale factor $\\mu_{\\text{min}}$)

    References:
        Bernetti, M., & Bussi, G. (2020). Pressure control using stochastic
        cell rescaling. J. Chem. Phys., 153(11), 114107.
        DOI: 10.1063/5.0020514
    """

    particles: Lens[State, Table[ParticleId, _BarostatParticleData]] = field(
        static=True
    )
    systems: Lens[State, Table[SystemId, IsMDSystemNPT]] = field(static=True)

    def __call__(self, key: Array, state: State) -> State:
        """Apply stochastic cell rescaling for pressure control.

        Scales the simulation box and particle positions by a factor determined
        from pressure deviation and stochastic fluctuations. The Cell is
        reconstructed to ensure cached volume is updated correctly.

        Args:
            key: JAX PRNG key for generating volume fluctuation noise
            state: Current simulation state

        Returns:
            Updated state with rescaled box and positions matching NPT ensemble
        """
        # Extract parameters
        systems = self.systems.get(state)
        params = systems.data.integrator_params
        # Δt: timestep [time]
        timestep = params.time_step
        # kT: thermal energy [energy]
        thermal_energy = params.temperature * BOLTZMANN_CONSTANT
        # P₀: target pressure [pressure]
        target_pressure = params.target_pressure
        # τP: barostat time constant [time]
        barostat_timescale = params.pressure_coupling_time
        # β: isothermal compressibility [1/pressure]
        compressibility = params.compressibility

        # Get current state
        # Cell with lattice vectors L
        cell = systems.data.cell
        # V: volume [length³]
        volume = cell.volume
        # Compute kinetic energy from particles
        particles = self.particles.bind(state).get()
        per_particle_ke = particle_kinetic_energy(
            particles.data.momenta, particles.data.masses
        )
        # K: total kinetic energy per system [energy]
        kinetic_energy = jax.ops.segment_sum(
            per_particle_ke,
            particles.data.system.indices,
            particles.data.system.num_labels,
        )

        # Full Cauchy stress via virial theorem:
        # σ = -(1/V)(Σ ∂U/∂r_i ⊗ r_i + h^T · ∂U/∂h)
        cauchy_stress = stress_via_virial_theorem(particles, systems).data
        # P = 2K/(dV) + Tr(σ)/d
        current_pressure = instantaneous_pressure(kinetic_energy, cauchy_stress, volume)

        # Stochastic cell rescaling (Bernetti & Bussi 2020)
        # Linearized form for small timesteps:
        # μ ≈ 1 + (Δt/τP)·β·(P - P₀) + √(2kT·β·Δt/(τP·V))·R
        # where R ~ N(0,1)

        # Stochastic cell rescaling (Bernetti & Bussi 2020, Eq. in reference impl)
        # dε = -β/τP·Δt·(P₀ - P) + √(2kT·β·Δt/(τP·V))·R
        # where dε = d(ln V) is the log-volume change. The LINEAR scaling
        # factor for lattice vectors is exp(dε/3) = (V_new/V)^(1/3).

        pressure_deviation = current_pressure - target_pressure
        # dε: log-volume change [dimensionless]
        depsilon_det = (
            (timestep / barostat_timescale) * compressibility * pressure_deviation
        )
        random_noise = jax.random.normal(key, dtype=volume.dtype)
        depsilon_stoch = (
            jnp.sqrt(
                2.0
                * thermal_energy
                * compressibility
                * timestep
                / (barostat_timescale * volume)
            )
            * random_noise
        )

        depsilon = depsilon_det + depsilon_stoch
        # Linear scaling: exp(dε/3) — cube root of volume scaling
        scaling_factor = jnp.exp(depsilon / 3.0)

        # Safety clamp to prevent extreme scaling
        # μ ∈ [μ_min, μ_max]
        min_scaling = params.minimum_scale_factor
        max_scaling = 1.0 / min_scaling
        scaling_factor = jnp.clip(scaling_factor, min_scaling, max_scaling)

        # Scale cell: L_new = μ·L
        # CRITICAL: Must reconstruct Cell to recompute cached volume
        # L_new = μ·L [length]
        new_cell = cell * scaling_factor
        state = self.systems.focus(lambda x: x.data.cell).set(state, new_cell)

        # Scale positions: r_new = μ·r
        particle_lens = self.particles.bind(state)
        particles = particle_lens.get()
        # μ_i: scaling factor per system [dimensionless]
        scaling_per_system = scaling_factor[particles.data.system.indices]

        new_positions = particles.data.positions * scaling_per_system[..., None]
        assert new_positions.shape == particles.data.positions.shape
        return particle_lens.focus(lambda p: p.data.positions).set(new_positions)