spintronics 0.3.2

//! Adaptive Heun (predictor-corrector) integrator for stochastic LLG.
//!
//! Builds on the existing fixed-step Heun (`dynamics/llg.rs:step_heun`) by adding
//! an embedded Euler-Heun error estimate (Heun is 2nd order, Euler is 1st order
//! — their difference approximates the local truncation error). A PI-style
//! controller adjusts `dt` to keep the estimated error below `atol + rtol·|m|`.
//!
//! For thermal noise (SLLG), the noise sample is held fixed within a step
//! (otherwise no convergence concept exists in the mean-square sense). On
//! step rejection the same noise sample is re-used with the smaller dt;
//! this preserves the Brownian increment statistics across retries.
//!
//! # Algorithm
//!
//! Given current state `m`, deterministic field `H_eff(m)`, frozen thermal
//! field `H_th`, and trial step `dt`:
//!
//! 1. **Predictor (Euler-Maruyama)**:
//!    `m_E = m + dt · f(m, H_eff(m) + H_th)`, then re-normalized.
//! 2. **Corrector (Heun)**:
//!    `m_H = m + (dt/2) · [f(m, H_eff(m) + H_th) + f(m_E, H_eff(m_E) + H_th)]`,
//!    then re-normalized.
//! 3. **Error estimate**: `err = |m_H - m_E|`.
//! 4. **Tolerance**: `tol = atol + rtol · max(|m|, |m_H|)`.
//! 5. **Accept** if `err ≤ tol`; otherwise reject and shrink `dt`.
//! 6. **Step-size update**: `dt_new = dt · clamp(safety · (tol/err)^(1/2),
//!    min_factor, max_factor)`.
//!
//! Because Heun has local truncation order 2 (global order 1 in the
//! mean-square sense for the multiplicative-noise SDE), the exponent
//! `1/(p+1) = 1/2` is appropriate.
//!
//! # References
//!
//! - Klauder, J. R. and Petersen, W. P., *Numerical integration of
//!   multiplicative-noise stochastic differential equations*, J. Stat.
//!   Phys. **39**, 53 (1985).
//! - García-Palacios, J. L. and Lázaro, F. J., *Langevin-dynamics study
//!   of the dynamical properties of small magnetic particles*,
//!   Phys. Rev. B **58**, 14937 (1998).
//! - Kloeden, P. E. and Platen, E., *Numerical Solution of Stochastic
//!   Differential Equations*, Springer (1992).
//! - Hairer, E. and Wanner, G., *Solving Ordinary Differential Equations II*,
//!   Springer, Chapter IV.8 (1996) — PI step-size control.

use crate::error::{invalid_param, Result};
use crate::stochastic::thermal::ThermalField;
use crate::vector3::Vector3;

/// Adaptive Heun integrator for the stochastic Landau-Lifshitz-Gilbert
/// equation.
///
/// Performs predictor-corrector steps with an embedded Euler-Heun error
/// estimate; the step is shrunk and retried (re-using the same thermal
/// noise sample) when the estimate exceeds `atol + rtol · |m|`.
#[derive(Debug, Clone)]
pub struct HeunAdaptive {
    /// Current (last-suggested) time step `[s]`.
    pub current_dt: f64,
    /// Absolute tolerance per Cartesian component.
    pub atol: f64,
    /// Relative tolerance (multiplied by `|m| ≈ 1`).
    pub rtol: f64,
    /// PI safety factor (typical: 0.9).
    pub safety: f64,
    /// Smallest allowed time step `[s]`.
    pub dt_min: f64,
    /// Largest allowed time step `[s]`.
    pub dt_max: f64,
    /// Maximum step-growth factor per step (typical: 5.0).
    pub max_factor: f64,
    /// Minimum step-shrink factor per step (typical: 0.2).
    pub min_factor: f64,
    /// Gyromagnetic ratio magnitude `[rad/(s·T)]`.
    pub gamma: f64,
    /// Gilbert damping parameter (dimensionless).
    pub alpha: f64,
    /// Maximum number of rejected retries per accepted step.
    pub max_retries: usize,
}

impl HeunAdaptive {
    /// Construct a new adaptive Heun integrator with the supplied initial
    /// step size and Gilbert damping.
    ///
    /// Default tolerances `atol = 1e-6`, `rtol = 1e-4`; default step bounds
    /// span 14 orders of magnitude (`dt_min = 1e-20 s`, `dt_max = 1e-6 s`),
    /// chosen to accommodate sub-femtosecond exchange dynamics through to
    /// millisecond relaxation. The gyromagnetic ratio defaults to
    /// `crate::constants::GAMMA`.
    ///
    /// # Errors
    /// Returns `InvalidParameter` if `initial_dt` is non-positive,
    /// non-finite, or if `alpha` is negative.
    pub fn new(initial_dt: f64, alpha: f64) -> Result<Self> {
        if !initial_dt.is_finite() || initial_dt <= 0.0 {
            return Err(invalid_param(
                "initial_dt",
                "must be a positive, finite number",
            ));
        }
        if !alpha.is_finite() || alpha < 0.0 {
            return Err(invalid_param(
                "alpha",
                "Gilbert damping must be non-negative and finite",
            ));
        }

        Ok(Self {
            current_dt: initial_dt,
            atol: 1.0e-6,
            rtol: 1.0e-4,
            safety: 0.9,
            dt_min: 1.0e-20,
            dt_max: 1.0e-6,
            max_factor: 5.0,
            min_factor: 0.2,
            gamma: crate::constants::GAMMA,
            alpha,
            max_retries: 20,
        })
    }

    /// Set absolute and relative tolerances.
    pub fn with_tolerances(mut self, atol: f64, rtol: f64) -> Self {
        self.atol = atol;
        self.rtol = rtol;
        self
    }

    /// Set the minimum and maximum allowed step sizes.
    pub fn with_dt_bounds(mut self, dt_min: f64, dt_max: f64) -> Self {
        self.dt_min = dt_min;
        self.dt_max = dt_max;
        self
    }

    /// Set the PI safety factor (typical 0.9).
    pub fn with_safety(mut self, safety: f64) -> Self {
        self.safety = safety;
        self
    }

    /// Set the maximum number of rejection retries per call to [`Self::step`].
    pub fn with_max_retries(mut self, n: usize) -> Self {
        self.max_retries = n;
        self
    }

    /// Set the gyromagnetic ratio magnitude (default: electron value).
    pub fn with_gamma(mut self, gamma: f64) -> Self {
        self.gamma = gamma;
        self
    }

    /// Perform one adaptive step of the stochastic LLG equation, with up to
    /// `max_retries` rejection retries.
    ///
    /// The thermal sample is generated once at the beginning of the step
    /// from `thermal.generate(self.current_dt)` and re-used for each
    /// retry — this is the *frozen-noise* convention required to preserve
    /// the mean-square accuracy of Heun under step-size changes.
    ///
    /// On success returns `(new_m, dt_used, error_estimate)`.
    ///
    /// # Errors
    /// Returns `NumericalError` if `dt_min` is reached without the error
    /// estimate falling below tolerance, or if `max_retries` is exceeded.
    pub fn step<F>(
        &mut self,
        m: Vector3<f64>,
        h_eff_fn: F,
        thermal: &mut ThermalField,
    ) -> Result<(Vector3<f64>, f64, f64)>
    where
        F: Fn(Vector3<f64>) -> Vector3<f64>,
    {
        // Freeze the noise sample at the start of the step; the thermal
        // field variance scales as 1/dt, so we rescale it to a "unit-time"
        // increment that can be re-used across rejections with different dt.
        let dt_for_noise = self.current_dt;
        let h_thermal_unit_scaled = if dt_for_noise > 0.0 {
            // Sample at the current dt; the FDT variance σ² ∝ 1/dt means
            // the raw thermal field already encodes the (1/dt) scaling.
            // To re-use the same Wiener increment ΔW across dt-rescales,
            // store H_th · sqrt(dt_for_noise) (the underlying Wiener
            // increment), then rescale to 1/sqrt(new_dt) on each retry.
            let h_raw = thermal.generate(dt_for_noise);
            h_raw * dt_for_noise.sqrt()
        } else {
            Vector3::zero()
        };

        let mut current_dt = self
            .current_dt
            .clamp(self.dt_min.max(0.0), self.dt_max.max(self.dt_min));

        for _ in 0..self.max_retries.max(1) {
            let h_thermal_step = if current_dt > 0.0 {
                h_thermal_unit_scaled * (1.0 / current_dt.sqrt())
            } else {
                Vector3::zero()
            };

            let (m_euler, m_heun) = Self::evolve_one_dt(
                m,
                &h_eff_fn,
                h_thermal_step,
                current_dt,
                self.gamma,
                self.alpha,
            );

            let err_vec = m_heun - m_euler;
            let err = err_vec.magnitude();
            let scale = m.magnitude().max(m_heun.magnitude()).max(1.0);
            let tol = self.atol + self.rtol * scale;

            // Step-size adjustment factor (Heun: p = 2 — embedded Euler is
            // p−1 = 1, so the exponent for the error reduction is 1/p = 1/2).
            let factor = if err > 0.0 {
                self.safety * (tol / err).sqrt()
            } else {
                self.max_factor
            };
            let factor = factor.clamp(self.min_factor, self.max_factor);

            if err <= tol {
                // Accept: schedule the next step at the suggested dt.
                let dt_next = (current_dt * factor).clamp(self.dt_min, self.dt_max);
                self.current_dt = dt_next;
                return Ok((m_heun, current_dt, err));
            }

            // Reject: shrink and retry with the *same* underlying noise.
            let new_dt = (current_dt * factor).clamp(self.dt_min, self.dt_max);
            if new_dt >= current_dt {
                // Cannot shrink further (already at dt_min) — accept best
                // effort to avoid stalling.
                self.current_dt = self.dt_min;
                return Ok((m_heun, current_dt, err));
            }
            current_dt = new_dt;
        }

        Err(crate::error::numerical_error(
            "HeunAdaptive: exceeded max_retries without meeting tolerance",
        ))
    }

    /// Perform a single Heun predictor-corrector step at the supplied dt
    /// using a frozen thermal field. Returns both the Euler predictor (for
    /// error estimation) and the Heun corrector result. Both are
    /// re-normalised so that `|m|` stays on the unit sphere.
    pub fn evolve_one_dt<F>(
        m: Vector3<f64>,
        h_eff_fn: &F,
        h_thermal: Vector3<f64>,
        dt: f64,
        gamma: f64,
        alpha: f64,
    ) -> (Vector3<f64>, Vector3<f64>)
    where
        F: Fn(Vector3<f64>) -> Vector3<f64>,
    {
        let h0 = h_eff_fn(m) + h_thermal;
        let k1 = sllg_rhs(m, h0, gamma, alpha);

        let m_euler = (m + k1 * dt).normalize();

        // Heun corrector: evaluate the slope at the predicted point with the
        // *same* thermal sample (frozen-noise convention).
        let h1 = h_eff_fn(m_euler) + h_thermal;
        let k2 = sllg_rhs(m_euler, h1, gamma, alpha);

        let m_heun = (m + (k1 + k2) * (0.5 * dt)).normalize();
        (m_euler, m_heun)
    }
}

/// Explicit form of the LLG right-hand side.
///
/// `f(m, h) = -γ/(1+α²) · [m × h + α · m × (m × h)]`
#[inline]
fn sllg_rhs(m: Vector3<f64>, h: Vector3<f64>, gamma: f64, alpha: f64) -> Vector3<f64> {
    let m_cross_h = m.cross(&h);
    let damping = m.cross(&m_cross_h) * alpha;
    (m_cross_h + damping) * (-gamma / (1.0 + alpha * alpha))
}

// =========================================================================
// Tests
// =========================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use crate::constants::{GAMMA, KB};

    fn close(a: f64, b: f64, tol: f64) -> bool {
        (a - b).abs() < tol
    }

    // -- Test 1: constructor & defaults ---------------------------------

    #[test]
    fn test_construct_and_defaults() {
        let h = HeunAdaptive::new(1.0e-13, 0.01).expect("construct ok");
        assert!(close(h.current_dt, 1.0e-13, 1e-25));
        assert!(close(h.atol, 1.0e-6, 1e-20));
        assert!(close(h.rtol, 1.0e-4, 1e-20));
        assert!(close(h.safety, 0.9, 1e-15));
        assert!(close(h.alpha, 0.01, 1e-15));
        assert!(close(h.gamma, GAMMA, 1.0));
        assert_eq!(h.max_retries, 20);
    }

    #[test]
    fn test_construct_rejects_bad_args() {
        assert!(HeunAdaptive::new(-1.0, 0.01).is_err());
        assert!(HeunAdaptive::new(f64::NAN, 0.01).is_err());
        assert!(HeunAdaptive::new(0.0, 0.01).is_err());
        assert!(HeunAdaptive::new(1.0e-13, -0.01).is_err());
        assert!(HeunAdaptive::new(1.0e-13, f64::INFINITY).is_err());
    }

    // -- Test 2: builder setters -----------------------------------------

    #[test]
    fn test_tolerance_setters() {
        let h = HeunAdaptive::new(1.0e-13, 0.01)
            .expect("ok")
            .with_tolerances(1.0e-8, 1.0e-6)
            .with_dt_bounds(1.0e-18, 1.0e-9)
            .with_safety(0.85)
            .with_max_retries(50)
            .with_gamma(0.5 * GAMMA);
        assert!(close(h.atol, 1.0e-8, 1e-20));
        assert!(close(h.rtol, 1.0e-6, 1e-20));
        assert!(close(h.dt_min, 1.0e-18, 1e-25));
        assert!(close(h.dt_max, 1.0e-9, 1e-25));
        assert!(close(h.safety, 0.85, 1e-15));
        assert_eq!(h.max_retries, 50);
        assert!(close(h.gamma, 0.5 * GAMMA, 1.0));
    }

    // -- Test 3: T = 0 limit recovers fixed-step Heun ---------------------

    #[test]
    fn test_zero_temperature_matches_deterministic_heun() {
        use crate::dynamics::llg::LlgSolver;

        let alpha = 0.05_f64;
        let dt = 1.0e-13_f64;
        let m0 = Vector3::new(1.0, 0.1, 0.05).normalize();
        let h_field = Vector3::new(0.0, 0.0, 1.0e5);

        // Deterministic reference: spintronics::dynamics::llg::step_heun.
        let solver = LlgSolver::new(alpha, dt);
        let m_ref = solver.step_heun(m0, |_m| h_field);

        // Adaptive integrator with very tight tolerances and T = 0.
        let mut integ = HeunAdaptive::new(dt, alpha)
            .expect("ok")
            .with_tolerances(1.0e-12, 1.0e-12)
            .with_dt_bounds(dt, dt); // pin dt
        let mut thermal = ThermalField::new(0.0, 1.0e-24, 1.0e6, alpha);

        let (m_new, dt_used, _err) = integ.step(m0, |_m| h_field, &mut thermal).expect("step ok");
        assert!(close(dt_used, dt, 1e-25));
        assert!(
            (m_new - m_ref).magnitude() < 1e-12,
            "adaptive Heun at T=0 must reproduce fixed-step Heun: \
             |m_new - m_ref| = {:.3e}",
            (m_new - m_ref).magnitude()
        );
    }

    // -- Test 4: error scales with dt^2 for smooth deterministic problem -

    #[test]
    fn test_error_scales_with_step_squared() {
        let alpha = 0.02_f64;
        let m0 = Vector3::new(1.0, 0.0, 0.0);
        // Field in "T-like" units (matches existing `step_heun` tests).
        let h_field = Vector3::new(0.0, 0.0, 1.0);

        let mut thermal = ThermalField::new(0.0, 1.0e-24, 1.0e6, alpha);

        // The embedded error E_emb = m_H − m_E scales asymptotically as
        // O(dt²) for fixed analytic trajectory; pick dt small enough that
        // we are in the asymptotic regime, but large enough that the
        // result is not dominated by round-off.
        let mut errs = Vec::new();
        for &dt in &[5.0e-13_f64, 1.0e-12_f64] {
            let mut integ = HeunAdaptive::new(dt, alpha)
                .expect("ok")
                .with_tolerances(1.0e-30, 1.0e-30) // never accept on first try
                .with_dt_bounds(dt, dt)
                .with_max_retries(1);
            let (_, _, err) = integ
                .step(m0, |_m| h_field, &mut thermal)
                .expect("step still returns best-effort");
            errs.push(err);
        }
        // For 2nd-order method, err(2dt) / err(dt) ≈ 4 in the asymptotic
        // regime. Allow a generous window because the predictor uses a
        // first-order step which competes with corrector.
        let ratio = errs[1] / errs[0].max(1.0e-300);
        assert!(
            ratio > 2.5 && ratio < 8.0,
            "error-vs-dt ratio = {:.3} (expected ~4)",
            ratio
        );
    }

    // -- Test 5: |m| stays on the unit sphere ----------------------------

    #[test]
    fn test_norm_preserved_after_step() {
        let alpha = 0.01_f64;
        let mut integ = HeunAdaptive::new(1.0e-13, alpha).expect("ok");
        let mut thermal = ThermalField::new(300.0, 1.0e-24, 1.0e6, alpha);
        let mut m = Vector3::new(1.0, 0.0, 0.0);
        let h_field = Vector3::new(0.0, 0.0, 1.0e5);

        for _ in 0..50 {
            let (m_new, _, _) = integ.step(m, |_m| h_field, &mut thermal).expect("step ok");
            m = m_new;
            assert!(
                (m.magnitude() - 1.0).abs() < 1e-10,
                "|m| drifted from unit sphere: {}",
                m.magnitude()
            );
        }
    }

    // -- Test 6: thermal averaging — ⟨|m|⟩ stays ≈ 1 (we re-normalize) ---

    #[test]
    fn test_thermal_norm_stability_at_300k() {
        // With normalization after every step, |m| ≈ 1 always; this test
        // checks that the integrator does not blow up over a long run with
        // strong thermal noise.
        let alpha = 0.1_f64;
        let mut integ = HeunAdaptive::new(1.0e-14, alpha)
            .expect("ok")
            .with_tolerances(1.0e-4, 1.0e-3);
        let mut thermal = ThermalField::new(300.0, 1.0e-26, 5.0e5, alpha);
        let mut m = Vector3::new(0.0, 0.0, 1.0);
        let h_field = Vector3::new(0.0, 0.0, 1.0e4);

        let mut mean_norm_sq = 0.0_f64;
        let n_steps = 500;
        for _ in 0..n_steps {
            let (m_new, _, _) = integ.step(m, |_m| h_field, &mut thermal).expect("step ok");
            m = m_new;
            mean_norm_sq += m.magnitude_squared();
            assert!(m.magnitude().is_finite());
        }
        mean_norm_sq /= n_steps as f64;
        assert!(
            (mean_norm_sq - 1.0).abs() < 1e-8,
            "⟨|m|²⟩ = {} should be ≈ 1 (we re-normalize)",
            mean_norm_sq
        );
    }

    // -- Test 7: dt grows when error is small ---------------------------

    #[test]
    fn test_dt_grows_when_error_small() {
        let alpha = 0.01_f64;
        let dt0 = 1.0e-16_f64;
        let mut integ = HeunAdaptive::new(dt0, alpha)
            .expect("ok")
            .with_tolerances(1.0e-3, 1.0e-3) // loose
            .with_dt_bounds(1.0e-20, 1.0e-9);
        let mut thermal = ThermalField::new(0.0, 1.0e-24, 1.0e6, alpha);
        let m = Vector3::new(1.0, 0.0, 0.0);
        let h_field = Vector3::new(0.0, 0.0, 1.0e4);

        let (_, dt_used, _) = integ.step(m, |_m| h_field, &mut thermal).expect("step ok");
        assert!(close(dt_used, dt0, 1e-25));
        // After the step, the controller should have grown dt towards
        // dt_max (limited by max_factor = 5).
        assert!(
            integ.current_dt > dt0,
            "dt did not grow (was {}, now {})",
            dt0,
            integ.current_dt
        );
    }

    // -- Test 8: dt shrinks when error is large --------------------------

    #[test]
    fn test_dt_shrinks_when_error_large() {
        let alpha = 0.5_f64; // heavy damping → fast dynamics
        let dt0 = 1.0e-9_f64; // way too large
        let mut integ = HeunAdaptive::new(dt0, alpha)
            .expect("ok")
            .with_tolerances(1.0e-12, 1.0e-12)
            .with_dt_bounds(1.0e-20, 1.0e-9)
            .with_max_retries(30);
        let mut thermal = ThermalField::new(0.0, 1.0e-24, 1.0e6, alpha);
        let m = Vector3::new(1.0, 0.0, 0.0);
        let h_field = Vector3::new(0.0, 0.0, 1.0e6);

        let result = integ.step(m, |_m| h_field, &mut thermal);
        // Either it rejected enough to converge (and dt was reduced) or it
        // returns NumericalError after exhausting retries.
        match result {
            Ok((_, dt_used, _)) => {
                assert!(
                    dt_used < dt0,
                    "dt should have shrunk from {}, got {}",
                    dt0,
                    dt_used
                );
            },
            Err(_) => {
                // Failure after exhausting retries is also acceptable
                // evidence that the controller tried to shrink.
                assert!(integ.current_dt <= dt0);
            },
        }
    }

    // -- Test 9: very small alpha (near-undamped precession) ------------

    #[test]
    fn test_small_alpha_precession() {
        let alpha = 1.0e-5_f64;
        let dt = 1.0e-13_f64;
        let m0 = Vector3::new(1.0, 0.0, 0.0);
        let h_field = Vector3::new(0.0, 0.0, 1.0e4);

        let mut integ = HeunAdaptive::new(dt, alpha)
            .expect("ok")
            .with_dt_bounds(dt, dt);
        let mut thermal = ThermalField::new(0.0, 1.0e-24, 1.0e6, alpha);

        let mut m = m0;
        // Larmor period ω = γ B with B = μ₀ H ≈ 4π·10⁻⁷ · 10⁴ ≈ 1.26·10⁻²;
        // T = 2π/(γ B) ≈ 2π / (1.76e11 · 1.26e-2) ≈ 2.84e-9 s.
        // Step a small fraction of that and check the z-component is
        // preserved (precession axis).
        for _ in 0..100 {
            let (m_new, _, _) = integ.step(m, |_m| h_field, &mut thermal).expect("step ok");
            m = m_new;
        }
        // With α ≈ 0, m_z should stay near 0 (rotation in xy plane); but
        // most importantly the orbit must not blow up.
        assert!(m.magnitude_squared().is_finite());
        assert!((m.magnitude() - 1.0).abs() < 1e-8);
    }

    // -- Test 10: rejection loop terminates within a few iterations -----

    #[test]
    fn test_reject_retry_terminates() {
        let alpha = 0.05_f64;
        let dt0 = 1.0e-10_f64; // too large for tight tol
        let mut integ = HeunAdaptive::new(dt0, alpha)
            .expect("ok")
            .with_tolerances(1.0e-10, 1.0e-10)
            .with_dt_bounds(1.0e-18, 1.0e-9)
            .with_max_retries(10);
        let mut thermal = ThermalField::new(0.0, 1.0e-24, 1.0e6, alpha);
        let m = Vector3::new(1.0, 0.0, 0.0);
        let h_field = Vector3::new(0.0, 0.0, 1.0e5);

        // Should either converge after a few shrinks or report a clear error.
        let _ = integ.step(m, |_m| h_field, &mut thermal);
        // current_dt should not be NaN regardless of outcome
        assert!(integ.current_dt.is_finite());
        assert!(integ.current_dt > 0.0);
    }

    // -- Bonus: SLLG RHS sanity (m·dm/dt = 0) ---------------------------

    #[test]
    fn test_sllg_rhs_orthogonality() {
        let m = Vector3::new(0.7, 0.1, -0.5).normalize();
        let h = Vector3::new(1.0e4, -2.0e3, 5.0e2);
        let dm = sllg_rhs(m, h, GAMMA, 0.05);
        // dm/dt is perpendicular to m → m · dm/dt = 0. Use a *relative*
        // tolerance because the absolute magnitudes are O(γ · |h|) ~ 1e15.
        let rel = m.dot(&dm).abs() / dm.magnitude().max(1.0);
        assert!(
            rel < 1.0e-12,
            "dm/dt not perpendicular to m (relative): {:.3e}",
            rel
        );
    }

    // Confirm we don't accidentally rely on KB elsewhere
    #[test]
    fn test_uses_kb_for_thermal() {
        // Smoke test: KB is used by ThermalField, which we invoke; just
        // assert KB is the expected SI value.
        assert!(close(KB, 1.380_649e-23, 1.0e-30));
    }
}