physics_in_parallel 4.1.1-alpha

/*!
Thermostats for canonical massive-particle state.

Purpose:
Thermostats modify particle velocities so a simulated particle system exchanges
energy with an implicit heat bath. This module is particle-specific: it mutates
`ATTR_V`, reads `ATTR_M_INV`, honors `ATTR_ALIVE` through `ParticleSelection`,
and always skips particles marked rigid by `ATTR_RIGID`.

Langevin convention:
`LangevinThermostat` applies the exact Ornstein-Uhlenbeck velocity update

`v_next = exp(-gamma * dt) * v_old + sqrt(tau_target * m_inv * (1 - exp(-2 * gamma * dt))) * z`

where `z` is a standard normal random value generated independently for each
particle and velocity component. The random stream is deterministic for a fixed
`seed`, `step_counter`, particle index, and component traversal order.
*/

use core::fmt;

use crate::threading::for_each_chunk_mut;

use crate::engines::soa::phys_obj::{AttrsError, PhysObj};
use crate::models::particles::attrs::{ATTR_M_INV, ATTR_V, ParticleSelection};
use crate::models::particles::state::{
    BorrowedMasks, ParticleStateError, mask_labels, validate_scalar_shape,
};
use crate::rng::{ResolvedRng, RngError};
use crate::space::discrete::square_lattice::random::IndexedRng;

const DOMAIN_LANGEVIN_NORMAL: u64 = 0xc7c7_d252_9b53_6071;

/// Errors returned by thermostat modules.
#[derive(Debug, Clone, PartialEq)]
pub enum ThermostatError {
    /// Parameter must be finite and within expected range.
    InvalidParam {
        field: &'static str,
        value: f64,
    },
    /// Time-step must be finite and strictly positive.
    InvalidDt {
        dt: f64,
    },
    State(ParticleStateError),
    Rng(RngError),
}

impl From<AttrsError> for ThermostatError {
    #[inline]
    fn from(value: AttrsError) -> Self {
        Self::State(value.into())
    }
}

impl From<ParticleStateError> for ThermostatError {
    fn from(value: ParticleStateError) -> Self {
        Self::State(value)
    }
}

impl fmt::Display for ThermostatError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InvalidParam { field, value } => write!(
                f,
                "thermostat parameter `{field}` must be finite and non-negative; got {value}"
            ),
            Self::InvalidDt { dt } => {
                write!(
                    f,
                    "thermostat time step must be finite and positive; got {dt}"
                )
            }
            Self::State(error) => write!(f, "invalid particle state: {error}"),
            Self::Rng(error) => write!(f, "thermostat RNG error: {error}"),
        }
    }
}

impl std::error::Error for ThermostatError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            Self::State(error) => Some(error),
            Self::Rng(error) => Some(error),
            _ => None,
        }
    }
}

/// Generic thermostat contract for particle state.
pub trait Thermostat: Send {
    /// Applies one thermostat step with time step `dt`.
    fn apply(&mut self, objects: &mut PhysObj, dt: f64) -> Result<(), ThermostatError>;
}

#[derive(Debug, Clone)]
pub struct LangevinThermostat {
    tau_target: f64,
    gamma: f64,
    rng: IndexedRng,
    step_counter: u64,
    selection: ParticleSelection,
}

impl LangevinThermostat {
    /// Constructs a Langevin thermostat.
    ///
    /// `tau_target` is the temperature-like energy scale used in the noise
    /// amplitude. `gamma` is the friction rate. Both must be finite and
    /// non-negative. `selection` controls whether dead particles are skipped;
    /// rigid particles are always skipped.
    pub fn new(
        tau_target: f64,
        gamma: f64,
        rng: ResolvedRng,
        selection: ParticleSelection,
    ) -> Result<Self, ThermostatError> {
        validate_nonnegative("tau_target", tau_target)?;
        validate_nonnegative("gamma", gamma)?;
        let rng = IndexedRng::new(rng).map_err(ThermostatError::Rng)?;

        Ok(Self {
            tau_target,
            gamma,
            rng,
            step_counter: 0,
            selection,
        })
    }

    /// Restores a Langevin thermostat from resolved deterministic state.
    ///
    /// Unlike [`Self::new`], this constructor never fills missing RNG values
    /// from defaults or host entropy. `rng` must contain both a seed and the
    /// indexed SplitMix64 method previously returned by [`Self::resolved_rng`].
    pub fn from_state(
        tau_target: f64,
        gamma: f64,
        rng: ResolvedRng,
        step_counter: u64,
        selection: ParticleSelection,
    ) -> Result<Self, ThermostatError> {
        validate_nonnegative("tau_target", tau_target)?;
        validate_nonnegative("gamma", gamma)?;
        let rng = IndexedRng::new(rng).map_err(ThermostatError::Rng)?;

        Ok(Self {
            tau_target,
            gamma,
            rng,
            step_counter,
            selection,
        })
    }

    #[inline]
    pub fn tau_target(&self) -> f64 {
        self.tau_target
    }

    #[inline]
    pub fn gamma(&self) -> f64 {
        self.gamma
    }

    #[inline]
    pub fn resolved_rng(&self) -> ResolvedRng {
        self.rng.resolved_rng()
    }

    #[inline]
    pub fn step_counter(&self) -> u64 {
        self.step_counter
    }

    #[inline]
    pub fn selection(&self) -> ParticleSelection {
        self.selection
    }
}

impl Thermostat for LangevinThermostat {
    /// Validates all masses and noise amplitudes before mutation. On error,
    /// velocities and the replay counter are unchanged for either backend.
    fn apply(&mut self, objects: &mut PhysObj, dt: f64) -> Result<(), ThermostatError> {
        validate_dt(dt)?;

        let (dim, n) = {
            let v = objects.core.get::<f64>(ATTR_V)?;
            (v.dim(), v.num_vectors())
        };

        let flags = mask_labels(objects, self.selection);
        let ([v, mass], flags) = objects
            .core
            .get_mixed_mut::<f64, u8, 2, 2>([ATTR_V, ATTR_M_INV], flags)?;
        validate_scalar_shape(ATTR_M_INV, mass.dim(), mass.num_vectors(), n)?;
        let m_inv_values = mass.borrow_values();
        let masks = BorrowedMasks::new(flags, n)?;

        let c = (-self.gamma * dt).exp();
        let one_minus_c2 = (1.0 - c * c).max(0.0);
        let step = self.step_counter;
        let rng = self.rng;
        let tau_target = self.tau_target;

        // Validate every included mass and derived amplitude before any write.
        // Recomputing the cheap square root avoids a full-system sigma buffer.
        for (i, &mass) in m_inv_values.iter().enumerate() {
            if masks.should_skip(i) {
                continue;
            }
            if !mass.is_finite() || mass <= 0.0 {
                return Err(ThermostatError::InvalidParam {
                    field: ATTR_M_INV,
                    value: mass,
                });
            }
            let sigma = (tau_target * mass * one_minus_c2).sqrt();
            if !sigma.is_finite() {
                return Err(ThermostatError::InvalidParam {
                    field: "sigma",
                    value: sigma,
                });
            }
        }
        v.edit_values(|values| {
            for_each_chunk_mut(values, dim, |start, chunk| {
                for (row, values) in chunk.chunks_exact_mut(dim).enumerate() {
                    let i = start / dim + row;
                    if masks.should_skip(i) {
                        continue;
                    }
                    let sigma = (tau_target * m_inv_values[i] * one_minus_c2).sqrt();
                    for (component, value) in values.iter_mut().enumerate() {
                        let z = rng.standard_normal(
                            step,
                            DOMAIN_LANGEVIN_NORMAL,
                            i as u64,
                            component as u64,
                        );
                        *value = c * *value + sigma * z;
                    }
                }
            });
        });

        self.step_counter = self.step_counter.wrapping_add(1);
        Ok(())
    }
}

#[inline]
fn validate_nonnegative(field: &'static str, value: f64) -> Result<(), ThermostatError> {
    if !value.is_finite() || value < 0.0 {
        return Err(ThermostatError::InvalidParam { field, value });
    }
    Ok(())
}

#[inline]
fn validate_dt(dt: f64) -> Result<(), ThermostatError> {
    if !dt.is_finite() || dt <= 0.0 {
        return Err(ThermostatError::InvalidDt { dt });
    }
    Ok(())
}