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 = jax.ops.segment_sum(
            per_particle_ke,
            particles.data.system.indices,
            particles.data.system.num_labels,
        )
        # 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)
        )

@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 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(IsMDSystem[IsCSVRNPTParams], Protocol):
    r"""NPT MD system row: adds the cell-gradient leaf needed for the barostat."""

    @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)

@dataclass
class StochasticStep[State](Propagator[State]):
    """Langevin thermostat stochastic step with exact Ornstein-Uhlenbeck solution.

    Implements the 'O' operator in the BAOAB splitting scheme. This step
    exactly solves the Ornstein-Uhlenbeck stochastic differential equation:

    $$d\\mathbf{p} = -\\gamma\\mathbf{p}\\,dt + \\sqrt{2\\gamma k_B T m}\\,dW$$

    The exact solution for timestep $\\Delta t$ is:

    $$\\mathbf{p}(t+\\Delta t) = e^{-\\gamma\\Delta t} \\mathbf{p}(t) + \\sqrt{k_B T(1-e^{-2\\gamma\\Delta t})} \\sqrt{m}\\,\\eta$$

    where $\\eta \\sim \\mathcal{N}(0,1)$ is Gaussian white noise. This preserves the correct
    Maxwell-Boltzmann distribution at temperature $T$.

    Type Parameters:
        State: Simulation state type

    Attributes:
        particles: Lens to get/set indexed particle data (momenta $\\mathbf{p}$, masses $m$)
        system: View to extract system data (time step $\\Delta t$, temperature $T$,
            friction coefficient $\\gamma$)

    References:
        Leimkuhler, B., & Matthews, C. (2013). Rational construction of
        stochastic numerical methods for molecular sampling.
        Appl. Math. Res. Express, 2013(1), 34-56.
        DOI: 10.1093/amrx/abs010
    """

    particles: Lens[State, Table[ParticleId, IsStochasticParticleData]] = field(
        static=True
    )
    system: View[State, Table[SystemId, IsBAOABLangevinParams]] = field(static=True)

    def __call__(self, key: Array, state: State) -> State:
        """Apply stochastic Ornstein-Uhlenbeck thermostat step.

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

        Returns:
            Updated state with thermostated momenta
        """
        # Extract current state
        particle_lens = self.particles.bind(state)
        particles = particle_lens.get()
        sys = self.system(state)[particles.data.system]
        # kT: thermal energy [energy]
        thermal_energy_per_particle = sys.temperature * BOLTZMANN_CONSTANT
        # Ornstein-Uhlenbeck coefficients
        # c₁ = e^(-γΔt) [dimensionless]
        damping_factor = jax.numpy.exp(-sys.friction_coefficient * sys.time_step)
        # c₂ = √(kT(1-e^(-2γΔt))) [√energy]
        noise_amplitude = jax.numpy.sqrt(
            thermal_energy_per_particle * (1 - damping_factor**2)
        )

        # η ~ N(0,1) [dimensionless]
        noise = sample_like(jax.random.normal, key, particles.data.momenta)

        # Exact OU solution: p_new = c₁·p + c₂·√m·η
        new_momenta = (
            damping_factor[..., None] * particles.data.momenta
            + (noise_amplitude * jnp.sqrt(particles.data.masses))[..., None] * noise
        )
        constrained_momenta = remove_center_of_mass_momentum(
            new_momenta, particles.data.masses, particles.data.system
        )
        # Deterministic gamma=0 or dt=0 steps should remain no-ops. For
        # active Langevin thermostats, keep the sampled ensemble in the
        # zero-total-momentum subspace used by the MD temperature analysis.
        should_project = (sys.friction_coefficient > 0) & (sys.time_step > 0)
        new_momenta = jnp.where(
            should_project[..., None], constrained_momenta, new_momenta
        )

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

@dataclass
class WrapFlow[State, PyTree](Flow[State, PyTree]):
    """Flow that applies the cell's wrap to updated positions.

    After the base flow updates positions, applies the cell's ``wrap`` method.
    On periodic axes this folds positions back into the box (minimum image
    convention); on non-periodic axes ``wrap`` is the identity and positions
    pass through unchanged.

    Type Parameters:
        State: Simulation state type
        PyTree: JAX PyTree type for positions

    Attributes:
        cell: View to extract the [Cell][kups.core.cell.Cell] from state
        flow: Underlying flow operator (typically [euclidean_flow][kups.md.integrators.euclidean_flow])

    Example:
        ```python
        from kups.md.integrators import WrapFlow, euclidean_flow

        wrap_flow = WrapFlow(
            cell=lambda s: s.cell,
            flow=euclidean_flow
        )
        ```
    """

    cell: View[State, Cell] = field(static=True)
    flow: Flow[State, PyTree] = field(static=True)

    def __call__(
        self, state: State, dt: Time, primal: PyTree, tangent: PyTree
    ) -> PyTree:
        return tree_map(self.cell(state).wrap, self.flow(state, dt, primal, tangent))