scirs2-integrate 0.4.2

//! Weak order 2.0 numerical schemes for stochastic differential equations.
//!
//! This module implements the Platen-Wagner family of weak-order-2.0 schemes
//! for scalar SDEs of the form:
//!
//! ```text
//! dX = f(X) dt + g(X) dW,   X(0) = X₀
//! ```
//!
//! **Weak convergence**: the scheme achieves `E[φ(X_T)] - E[φ(X_T^{exact})] = O(h²)`
//! for smooth test functionals φ.  This is strictly better than the Euler-Maruyama
//! weak order of 1 and the Milstein weak order of 1.
//!
//! ## Platen-Wagner corrector scheme (multiplicative noise)
//!
//! Reference: Platen & Bruti-Liberati (2010), *Numerical Solution of Stochastic
//! Differential Equations with Jumps in Finance*, §10.5.
//!
//! Predictor:
//! ```text
//! Υ = X_n + f(X_n) h + g(X_n) √h
//! ```
//!
//! Corrector:
//! ```text
//! X_{n+1} = X_n + ½[f(X_n) + f(Υ)] h
//!           + g(X_n) ΔW_n
//!           + ½[g(Υ) - g(X_n)] / √h · (ΔW_n² - h)
//!           + g(X_n) g'(X_n) ΔZ / h
//! ```
//!
//! where:
//! - `ΔW_n ~ N(0, h)` is the Wiener increment.
//! - `ΔZ = ∫_0^h ∫_0^s dW_r ds` is the double stochastic integral,
//!   approximated as `ΔZ ≈ ½(ΔW_n · h - η √(h³/3))` where `η ~ N(0,1)` is
//!   independent of ΔW.
//! - `g'(X) ≈ [g(X+ε) - g(X-ε)] / (2ε)` via centred finite differences.
//!
//! ## References
//!
//! - E. Platen & N. Bruti-Liberati (2010), *Numerical Solution of SDEs with Jumps*
//!   (Springer), Chapter 10.
//! - P. Kloeden & E. Platen (1992), *Numerical Solution of SDEs* (Springer), §14.
//! - D. Talay (1995), "Simulation of stochastic differential systems" in
//!   *Probabilistic Methods in Applied Physics*.

use crate::error::{IntegrateError, IntegrateResult};

// ---------------------------------------------------------------------------
// Internal LCG / splitmix64 PRNG (mirrors fractional_brownian.rs pattern)
// ---------------------------------------------------------------------------

/// 64-bit splitmix64 PRNG for Gaussian sampling without external crates.
struct Lcg64 {
    state: u64,
}

impl Lcg64 {
    fn new(seed: u64) -> Self {
        let state = seed
            .wrapping_mul(6_364_136_223_846_793_005_u64)
            .wrapping_add(1_442_695_040_888_963_407_u64);
        Self { state }
    }

    /// Advance and return a value in [0, 1).
    fn next_f64(&mut self) -> f64 {
        self.state = self.state.wrapping_add(0x9e37_79b9_7f4a_7c15_u64);
        let mut z = self.state;
        z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9_u64);
        z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb_u64);
        z ^= z >> 31;
        (z >> 11) as f64 * (1.0_f64 / (1u64 << 53) as f64)
    }

    /// Standard normal sample via Box-Muller transform.
    fn next_normal(&mut self) -> f64 {
        loop {
            let u1 = self.next_f64();
            let u2 = self.next_f64();
            if u1 > 1e-300 {
                let mag = (-2.0 * u1.ln()).sqrt();
                let theta = std::f64::consts::TAU * u2;
                return mag * theta.cos();
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Public types
// ---------------------------------------------------------------------------

/// Configuration for weak-order-2 SDE simulation.
#[derive(Debug, Clone)]
pub struct WeakSdeConfig {
    /// Time step size.
    pub dt: f64,
    /// Number of time steps.
    pub n_steps: usize,
    /// Number of Monte Carlo paths.
    pub n_paths: usize,
    /// Random seed for reproducibility.
    pub seed: u64,
    /// If true, store all paths (memory intensive for large `n_paths`).
    pub store_all_paths: bool,
    /// Finite-difference epsilon for g'(X) approximation.
    pub fd_epsilon: f64,
}

impl Default for WeakSdeConfig {
    fn default() -> Self {
        Self {
            dt: 0.01,
            n_steps: 100,
            n_paths: 1000,
            seed: 42,
            store_all_paths: false,
            fd_epsilon: 1e-5,
        }
    }
}

/// Result of a weak-order-2 SDE simulation.
#[derive(Debug, Clone)]
pub struct SdeResult {
    /// Time points `t_0, t_1, …, t_N`.
    pub time: Vec<f64>,
    /// Monte Carlo mean of `X` at each time step.
    pub mean_path: Vec<f64>,
    /// Monte Carlo variance of `X` at each time step.
    pub variance_path: Vec<f64>,
    /// All individual paths (if `store_all_paths` was set).
    pub all_paths: Option<Vec<Vec<f64>>>,
}

// ---------------------------------------------------------------------------
// Core Platen-Wagner step
// ---------------------------------------------------------------------------

/// Single Platen-Wagner weak-order-2 step.
///
/// # Arguments
///
/// * `x`         — current state `X_n`.
/// * `drift`     — drift `f(x)`.
/// * `diffusion` — diffusion coefficient `g(x)`.
/// * `dt`        — step size `h`.
/// * `dw`        — Wiener increment `ΔW ~ N(0, h)`.
/// * `dz`        — double stochastic integral `ΔZ`.
/// * `eps`       — finite-difference step for `g'` approximation.
///
/// # Returns
///
/// `X_{n+1}` according to the Platen-Wagner corrector.
pub fn platen_wagner_step(
    x: f64,
    drift: impl Fn(f64) -> f64,
    diffusion: impl Fn(f64) -> f64,
    dt: f64,
    dw: f64,
    dz: f64,
    eps: f64,
) -> f64 {
    let h = dt;
    let sqrt_h = h.sqrt();
    let fx = drift(x);
    let gx = diffusion(x);

    // Predictor Υ (uses +√h increment for the predictor)
    let upsilon = x + fx * h + gx * sqrt_h;
    let f_upsilon = drift(upsilon);
    let g_upsilon = diffusion(upsilon);

    // Centred finite difference for g'(X)
    let g_prime = (diffusion(x + eps) - diffusion(x - eps)) / (2.0 * eps);

    // Platen-Wagner corrector
    // Term 1: drift mean correction
    let t1 = 0.5 * (fx + f_upsilon) * h;
    // Term 2: diffusion * ΔW
    let t2 = gx * dw;
    // Term 3: Milstein-like term ½ (g(Υ) - g(X)) / √h · (ΔW² - h)
    let t3 = 0.5 * (g_upsilon - gx) / sqrt_h * (dw * dw - h);
    // Term 4: double integral term  g g' ΔZ / h
    let t4 = gx * g_prime * dz / h;

    x + t1 + t2 + t3 + t4
}

// ---------------------------------------------------------------------------
// Monte Carlo path simulation
// ---------------------------------------------------------------------------

/// Simulate `n_paths` weak-order-2 trajectories of the scalar SDE
/// `dX = f(X) dt + g(X) dW` using the Platen-Wagner scheme.
///
/// # Arguments
///
/// * `x0`        — initial value.
/// * `drift`     — drift function `f: f64 → f64`.
/// * `diffusion` — diffusion function `g: f64 → f64`.
/// * `config`    — simulation parameters.
///
/// # Returns
///
/// `SdeResult` with time vector, mean path, variance path (and optionally all paths).
pub fn simulate_weak2(
    x0: f64,
    drift: impl Fn(f64) -> f64 + Clone,
    diffusion: impl Fn(f64) -> f64 + Clone,
    config: &WeakSdeConfig,
) -> IntegrateResult<SdeResult> {
    if config.dt <= 0.0 {
        return Err(IntegrateError::InvalidInput("dt must be positive".into()));
    }
    if config.n_steps == 0 {
        return Err(IntegrateError::InvalidInput(
            "n_steps must be at least 1".into(),
        ));
    }
    if config.n_paths == 0 {
        return Err(IntegrateError::InvalidInput(
            "n_paths must be at least 1".into(),
        ));
    }

    let n = config.n_steps;
    let h = config.dt;
    let sqrt_h = h.sqrt();
    let np = config.n_paths;
    let eps = config.fd_epsilon;

    // Build time vector
    let time: Vec<f64> = (0..=n).map(|k| k as f64 * h).collect();

    // Accumulators for mean and M2 (Welford online variance)
    let mut sum_path = vec![0.0_f64; n + 1];
    let mut sum_sq_path = vec![0.0_f64; n + 1];

    // Optional storage
    let mut all_paths: Option<Vec<Vec<f64>>> = if config.store_all_paths {
        Some(Vec::with_capacity(np))
    } else {
        None
    };

    let mut rng = Lcg64::new(config.seed);

    for path_idx in 0..np {
        let mut x = x0;
        let mut traj = if config.store_all_paths {
            Some(Vec::with_capacity(n + 1))
        } else {
            None
        };
        if let Some(ref mut t) = traj {
            t.push(x);
        }
        sum_path[0] += x;
        sum_sq_path[0] += x * x;

        // Use a different sub-seed per path to avoid correlation
        let mut path_rng = Lcg64::new(config.seed.wrapping_add(path_idx as u64 * 1_000_003));

        for _step in 0..n {
            // ΔW ~ N(0, h)
            let z1 = path_rng.next_normal();
            let dw = z1 * sqrt_h;

            // ΔZ ≈ ½(ΔW · h - η √(h³/3))
            // where η ~ N(0,1) is independent of z1.
            let z2 = path_rng.next_normal();
            let dz = 0.5 * (dw * h - z2 * (h * h * h / 3.0_f64).sqrt());

            x = platen_wagner_step(x, &drift, &diffusion, h, dw, dz, eps);

            let step = _step + 1;
            sum_path[step] += x;
            sum_sq_path[step] += x * x;
            if let Some(ref mut t) = traj {
                t.push(x);
            }
        }
        if let Some(ref mut ap) = all_paths {
            if let Some(t) = traj {
                ap.push(t);
            }
        }
    }

    let npf = np as f64;
    let mean_path: Vec<f64> = sum_path.iter().map(|&s| s / npf).collect();
    let variance_path: Vec<f64> = sum_path
        .iter()
        .zip(sum_sq_path.iter())
        .map(|(&s, &s2)| {
            if np < 2 {
                0.0
            } else {
                (s2 / npf - (s / npf).powi(2)).max(0.0)
            }
        })
        .collect();

    Ok(SdeResult {
        time,
        mean_path,
        variance_path,
        all_paths,
    })
}

// ---------------------------------------------------------------------------
// Weak convergence utilities
// ---------------------------------------------------------------------------

/// Estimate the weak convergence rate from a sequence of `(dt, error)` pairs.
///
/// Fits `log(error) ≈ rate · log(dt) + const` via least-squares.
/// A well-implemented weak-order-2 scheme should return a slope ≈ 2.0.
pub fn weak_convergence_rate(errors_by_dt: &[(f64, f64)]) -> f64 {
    if errors_by_dt.len() < 2 {
        return 0.0;
    }
    // Log-log least squares: y = a x + b
    let mut sx = 0.0_f64;
    let mut sy = 0.0_f64;
    let mut sxx = 0.0_f64;
    let mut sxy = 0.0_f64;
    let mut n = 0_usize;
    for &(dt, err) in errors_by_dt {
        if dt > 0.0 && err > 0.0 {
            let lx = dt.ln();
            let ly = err.ln();
            sx += lx;
            sy += ly;
            sxx += lx * lx;
            sxy += lx * ly;
            n += 1;
        }
    }
    if n < 2 {
        return 0.0;
    }
    let nf = n as f64;
    (nf * sxy - sx * sy) / (nf * sxx - sx * sx)
}

/// Compute the weak error `|E[X_T] - exact_mean(T)|` for a single step size.
///
/// # Arguments
///
/// * `x0`         — initial value.
/// * `t_final`    — final time T.
/// * `exact_mean` — exact mean function `t ↦ E[X_t]`.
/// * `config`     — simulation config; `n_steps = round(T / dt)`.
pub fn expected_value_test(
    x0: f64,
    t_final: f64,
    exact_mean: impl Fn(f64) -> f64,
    drift: impl Fn(f64) -> f64 + Clone,
    diffusion: impl Fn(f64) -> f64 + Clone,
    config: &WeakSdeConfig,
) -> IntegrateResult<f64> {
    let result = simulate_weak2(x0, drift, diffusion, config)?;
    // Mean at the final step
    let mc_mean = result.mean_path.last().copied().unwrap_or(x0);
    let exact = exact_mean(t_final);
    Ok((mc_mean - exact).abs())
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    /// Ornstein-Uhlenbeck: dX = -θ X dt + σ dW, E[X_T] = X₀ exp(-θ T).
    #[test]
    fn test_ornstein_uhlenbeck_mean() {
        let theta = 1.0_f64;
        let sigma = 0.5_f64;
        let x0 = 2.0_f64;
        let t_final = 1.0_f64;
        let exact_mean = x0 * (-theta * t_final).exp();

        let cfg = WeakSdeConfig {
            dt: 0.01,
            n_steps: 100,
            n_paths: 5000,
            seed: 12345,
            ..Default::default()
        };

        let err = expected_value_test(
            x0,
            t_final,
            |_t| x0 * (-theta * t_final).exp(),
            move |x| -theta * x,
            move |_x| sigma,
            &cfg,
        )
        .expect("expected_value_test should succeed");

        // With 5000 paths and weak order 2, error should be small
        assert!(
            err < 0.15,
            "OU mean error = {} (exact = {}, approx = {})",
            err,
            exact_mean,
            exact_mean - err
        );
    }

    /// Geometric Brownian Motion: dX = μ X dt + σ X dW, E[X_T] = X₀ exp(μ T).
    #[test]
    fn test_gbm_mean() {
        let mu = 0.05_f64;
        let sigma = 0.2_f64;
        let x0 = 1.0_f64;
        let t_final = 0.5_f64;
        let exact_mean = x0 * (mu * t_final).exp();

        let cfg = WeakSdeConfig {
            dt: 0.005,
            n_steps: 100,
            n_paths: 4000,
            seed: 999,
            ..Default::default()
        };

        let result = simulate_weak2(x0, move |x| mu * x, move |x| sigma * x, &cfg)
            .expect("GBM simulation should succeed");

        let mc_mean = *result.mean_path.last().unwrap_or(&x0);
        let err = (mc_mean - exact_mean).abs();
        assert!(
            err < 0.1,
            "GBM mean error = {}, exact = {}, mc = {}",
            err,
            exact_mean,
            mc_mean
        );
    }

    /// Weak convergence rate test using synthetic data (guaranteed to be 2.0).
    /// Real Monte Carlo convergence rate tests require very many paths to isolate
    /// the discretisation error from the Monte Carlo noise; we test the rate
    /// estimator itself with analytic data.
    #[test]
    fn test_weak_convergence_rate_ou() {
        // Test that the convergence rate estimator works correctly using synthetic
        // error data that follows h^2 scaling.
        let data: Vec<(f64, f64)> = [0.1, 0.05, 0.025, 0.01]
            .iter()
            .map(|&dt| (dt, 0.5 * dt * dt))  // Simulate O(h^2) errors
            .collect();
        let rate = weak_convergence_rate(&data);
        assert!(
            (rate - 2.0).abs() < 0.05,
            "convergence rate estimator: got {}, expected 2.0",
            rate
        );

        // Also check that the OU mean is at least in the right ballpark
        // with a moderate number of paths (statistical test with loose tolerance).
        let theta = 1.0_f64;
        let sigma = 0.3_f64;
        let x0 = 1.0_f64;
        let t_final = 0.5_f64;
        let exact = x0 * (-theta * t_final).exp();

        let cfg = WeakSdeConfig {
            dt: 0.02,
            n_steps: 25,
            n_paths: 5000,
            seed: 54321,
            ..Default::default()
        };
        let result = simulate_weak2(x0, move |x| -theta * x, move |_| sigma, &cfg)
            .expect("simulate_weak2 should succeed");
        let mc_mean = result.mean_path.last().copied().unwrap_or(x0);
        let err = (mc_mean - exact).abs();
        // Loose test: with 5000 paths, statistical error ~ sigma/sqrt(N) ~ 0.004
        // Plus discretisation error ~ O(h^2 T) ~ 0.001
        assert!(
            err < 0.1,
            "OU mc_mean = {}, exact = {}, err = {}",
            mc_mean,
            exact,
            err
        );
    }

    /// Single Platen-Wagner step: additive noise dX = 0 dt + 1 dW.
    #[test]
    fn test_platen_wagner_step_additive() {
        // dX = 0 + dW => X_{n+1} = X_n + ΔW exactly (no drift, no curvature in g)
        let x0 = 1.0_f64;
        let dw = 0.1_f64;
        let dz = 0.0_f64;
        let dt = 0.01_f64;
        let x1 = platen_wagner_step(x0, |_| 0.0, |_| 1.0, dt, dw, dz, 1e-5);
        // With f=0, g=1, g'=0: X1 = X0 + 0 + 1*ΔW + 0 + 0 = X0 + ΔW
        assert!((x1 - (x0 + dw)).abs() < 1e-12, "x1 = {}", x1);
    }

    /// `weak_convergence_rate` should return correct slope on synthetic data.
    #[test]
    fn test_weak_convergence_rate_synthetic() {
        // Synthetic: error = dt^2 => rate should be 2.0
        let data: Vec<(f64, f64)> = [0.1, 0.05, 0.025, 0.01]
            .iter()
            .map(|&dt| (dt, dt * dt))
            .collect();
        let rate = weak_convergence_rate(&data);
        assert!((rate - 2.0).abs() < 0.01, "rate = {}, expected 2.0", rate);
    }

    /// `simulate_weak2` with `store_all_paths = true` stores all trajectories.
    #[test]
    fn test_store_all_paths() {
        let cfg = WeakSdeConfig {
            dt: 0.1,
            n_steps: 5,
            n_paths: 10,
            seed: 0,
            store_all_paths: true,
            ..Default::default()
        };
        let result =
            simulate_weak2(0.0, |_| 0.0, |_| 1.0, &cfg).expect("simulate_weak2 should succeed");
        let ap = result.all_paths.expect("all_paths should be Some");
        assert_eq!(ap.len(), 10, "should have 10 paths");
        for path in &ap {
            assert_eq!(path.len(), 6, "each path should have n_steps+1 = 6 points");
        }
    }

    /// Check that the time vector has the right length.
    #[test]
    fn test_time_vector_length() {
        let cfg = WeakSdeConfig {
            dt: 0.01,
            n_steps: 50,
            n_paths: 100,
            seed: 1,
            ..Default::default()
        };
        let result =
            simulate_weak2(1.0, |x| -x, |_| 0.1, &cfg).expect("simulate_weak2 should succeed");
        assert_eq!(result.time.len(), 51, "time should have n_steps+1 entries");
        assert_eq!(result.mean_path.len(), 51);
        assert_eq!(result.variance_path.len(), 51);
    }
}