oxiflow 0.6.0

//! # Module `solver::methods::crank_nicolson`
//!
//! Crank-Nicolson integrator — semi-implicit, 2nd order (issue #43).
//!
//! ## Algorithm
//!
//! $$u^{n+1} = u^n + \frac{\Delta t}{2}\left[f(u^n, t^n) + f(u^{n+1}, t^{n+1})\right]$$
//!
//! A thin wrapper around the shared generalised theta method
//! ([`super::implicit::theta_method_step`], θ=0.5) — see that module's
//! docs for the frozen-Jacobian v1 scope and the path to a future
//! iterated Newton solver (DD-033). Identical structure to
//! [`BackwardEulerSolver`](super::backward_euler::BackwardEulerSolver) —
//! only `theta` differs.
//!
//! ## Stability
//!
//! A-stable but not L-stable: oscillatory (non-decaying) error modes can
//! persist for very stiff problems, unlike Backward Euler's strong
//! damping. Preferred over Backward Euler when 2nd-order accuracy matters
//! more than damping every mode aggressively.
//!
//! ## Scope at J4a
//!
//! Same restrictions as `BackwardEulerSolver`: single-domain,
//! `StepControl::Fixed` only, BC interaction with the frozen Jacobian not
//! yet covered by a dedicated test (see [`super::implicit`]).

use crate::context::error::OxiflowError;
use crate::context::value::ContextValue;
use crate::context::ContextCalculator;
use crate::solver::linear::{LinearSolver, NalgebraDenseSolver};
use crate::solver::methods::implicit::theta_method_step;
#[cfg(feature = "sparse")]
use crate::solver::methods::implicit::theta_method_step_adaptive;
use crate::solver::methods::SteppableSolver;
use crate::solver::scenario::{Domain, Scenario};
#[cfg(feature = "sparse")]
use crate::solver::sparse::SparseLinearSolver;
use crate::solver::{SimulationResult, Solver, SolverConfiguration};

/// Crank-Nicolson solver — semi-implicit, 2nd order.
pub struct CrankNicolsonSolver {
    linear_solver: Box<dyn LinearSolver>,
    /// Sparse backend (DD-043) — see
    /// [`BackwardEulerSolver`](super::backward_euler::BackwardEulerSolver)'s
    /// fields for the full rationale, identical here.
    #[cfg(feature = "sparse")]
    sparse_solver: Option<Box<dyn SparseLinearSolver>>,
    #[cfg(feature = "sparse")]
    sparse_threshold: usize,
    #[cfg(feature = "sparse")]
    jacobian_bandwidth: Option<usize>,
}

impl Default for CrankNicolsonSolver {
    fn default() -> Self {
        Self {
            linear_solver: Box::new(NalgebraDenseSolver),
            #[cfg(feature = "sparse")]
            sparse_solver: None,
            #[cfg(feature = "sparse")]
            sparse_threshold: 100,
            #[cfg(feature = "sparse")]
            jacobian_bandwidth: None,
        }
    }
}

impl CrankNicolsonSolver {
    /// Creates a solver using the default `nalgebra` dense backend.
    pub fn new() -> Self {
        Self::default()
    }

    /// Substitutes the linear solver backend (DD-013).
    pub fn with_linear_solver(mut self, linear_solver: Box<dyn LinearSolver>) -> Self {
        self.linear_solver = linear_solver;
        self
    }

    /// Configures the sparse backend (DD-043) — see
    /// [`BackwardEulerSolver::with_sparse_solver`](super::backward_euler::BackwardEulerSolver::with_sparse_solver).
    #[cfg(feature = "sparse")]
    pub fn with_sparse_solver(mut self, sparse_solver: Box<dyn SparseLinearSolver>) -> Self {
        self.sparse_solver = Some(sparse_solver);
        self
    }

    /// System size above which the sparse path is used (default 100).
    #[cfg(feature = "sparse")]
    pub fn with_sparse_threshold(mut self, sparse_threshold: usize) -> Self {
        self.sparse_threshold = sparse_threshold;
        self
    }

    /// Declares the Jacobian's half-bandwidth (e.g.
    /// `scheme.stencil_radius()`, DD-039).
    #[cfg(feature = "sparse")]
    pub fn with_jacobian_bandwidth(mut self, jacobian_bandwidth: usize) -> Self {
        self.jacobian_bandwidth = Some(jacobian_bandwidth);
        self
    }
}

impl Solver for CrankNicolsonSolver {
    fn solve(
        &self,
        scenario: &Scenario,
        config: &SolverConfiguration,
    ) -> Result<SimulationResult, OxiflowError> {
        self.solve_fixed_step(scenario, config)
    }
}

#[cfg(feature = "sparse")]
impl SteppableSolver for CrankNicolsonSolver {
    fn step(
        &self,
        domain: &Domain,
        chain: &[&dyn ContextCalculator],
        state: &mut ContextValue,
        _history: &[ContextValue],
        t: f64,
        dt: f64,
    ) -> Result<ContextValue, OxiflowError> {
        // history_depth() defaults to 0 -- Crank-Nicolson is a one-step
        // method, `_history` is always empty here and intentionally unused.
        //
        // See BackwardEulerSolver::step for why dispatch is gated on
        // `sparse_solver` alone.
        match &self.sparse_solver {
            Some(sparse_solver) => theta_method_step_adaptive(
                domain,
                chain,
                state,
                t,
                dt,
                0.5,
                self.linear_solver.as_ref(),
                sparse_solver.as_ref(),
                self.sparse_threshold,
                self.jacobian_bandwidth,
            ),
            None => theta_method_step(
                domain,
                chain,
                state,
                t,
                dt,
                0.5,
                self.linear_solver.as_ref(),
            ),
        }
    }
}

#[cfg(not(feature = "sparse"))]
impl SteppableSolver for CrankNicolsonSolver {
    fn step(
        &self,
        domain: &Domain,
        chain: &[&dyn ContextCalculator],
        state: &mut ContextValue,
        _history: &[ContextValue],
        t: f64,
        dt: f64,
    ) -> Result<ContextValue, OxiflowError> {
        // history_depth() defaults to 0 -- Crank-Nicolson is a one-step
        // method, `_history` is always empty here and intentionally unused.
        theta_method_step(
            domain,
            chain,
            state,
            t,
            dt,
            0.5,
            self.linear_solver.as_ref(),
        )
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::context::compute::ComputeContext;
    use crate::context::variable::ContextVariable;
    use crate::mesh::{Mesh, UniformGrid1D};
    use crate::model::traits::{PhysicalModel, RequiresContext};
    use crate::solver::config::{
        IntegratorKind, SolverConfiguration, StepControl, TimeConfiguration,
    };
    use nalgebra::{DMatrix, DVector};
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::sync::Arc;

    #[derive(Debug)]
    struct ExponentialDecay {
        lambda: f64,
    }

    impl RequiresContext for ExponentialDecay {
        fn required_variables(&self) -> Vec<ContextVariable> {
            vec![]
        }
    }

    impl PhysicalModel for ExponentialDecay {
        fn compute_physics(
            &self,
            state: &ContextValue,
            _ctx: &ComputeContext,
        ) -> Result<ContextValue, OxiflowError> {
            let u = state.as_scalar_field()?;
            Ok(ContextValue::ScalarField(u.map(|v| -self.lambda * v)))
        }

        fn initial_state(&self, mesh: &dyn Mesh) -> ContextValue {
            ContextValue::ScalarField(DVector::from_element(mesh.n_dof(), 1.0))
        }

        fn name(&self) -> &str {
            "exponential_decay"
        }
    }

    #[derive(Debug)]
    struct ZeroDerivative;

    impl RequiresContext for ZeroDerivative {
        fn required_variables(&self) -> Vec<ContextVariable> {
            vec![]
        }
    }

    impl PhysicalModel for ZeroDerivative {
        fn compute_physics(
            &self,
            state: &ContextValue,
            _ctx: &ComputeContext,
        ) -> Result<ContextValue, OxiflowError> {
            let u = state.as_scalar_field()?;
            Ok(ContextValue::ScalarField(DVector::zeros(u.len())))
        }

        fn initial_state(&self, mesh: &dyn Mesh) -> ContextValue {
            ContextValue::ScalarField(DVector::from_element(mesh.n_dof(), 2.5))
        }

        fn name(&self) -> &str {
            "zero_derivative"
        }
    }

    /// Delegates to `NalgebraDenseSolver` but counts calls.
    #[derive(Debug)]
    struct CountingLinearSolver {
        calls: Arc<AtomicUsize>,
    }

    impl LinearSolver for CountingLinearSolver {
        fn solve(&self, a: &DMatrix<f64>, b: &DVector<f64>) -> Result<DVector<f64>, OxiflowError> {
            self.calls.fetch_add(1, Ordering::SeqCst);
            NalgebraDenseSolver.solve(a, b)
        }
    }

    fn make_config(t_end: f64, dt: f64) -> SolverConfiguration {
        SolverConfiguration::new(
            TimeConfiguration::new(t_end, StepControl::Fixed { dt }),
            IntegratorKind::CrankNicolson,
        )
    }

    fn make_mesh(n: usize) -> Box<dyn Mesh> {
        Box::new(UniformGrid1D::new(n, 0.0, 1.0).unwrap())
    }

    #[test]
    fn zero_derivative_field_stays_constant() {
        let scenario = Scenario::single(Box::new(ZeroDerivative), make_mesh(5));
        let config = make_config(1.0, 0.1);
        let result = CrankNicolsonSolver::new()
            .solve(&scenario, &config)
            .unwrap();
        for state in &result.states {
            let field = state.as_scalar_field().unwrap();
            for v in field.iter() {
                assert!((v - 2.5).abs() < 1e-12);
            }
        }
    }

    #[test]
    fn exponential_decay_matches_analytical_over_many_steps() {
        // u^{n} = u^0 * [(1 - lambda*dt/2) / (1 + lambda*dt/2)]^n
        let lambda = 2.0;
        let dt = 0.1;
        let n_steps = 20;
        let scenario = Scenario::single(Box::new(ExponentialDecay { lambda }), make_mesh(2));
        let config = make_config(n_steps as f64 * dt, dt);
        let result = CrankNicolsonSolver::new()
            .solve(&scenario, &config)
            .unwrap();

        let ratio = (1.0 - lambda * dt / 2.0) / (1.0 + lambda * dt / 2.0);
        let expected = ratio.powi(n_steps);
        let final_field = result.states.last().unwrap().as_scalar_field().unwrap();
        for v in final_field.iter() {
            assert!((v - expected).abs() < 1e-9, "got {v}, expected {expected}");
        }
    }

    #[test]
    fn result_times_match_expected_steps() {
        let scenario = Scenario::single(Box::new(ZeroDerivative), make_mesh(3));
        let config = make_config(0.5, 0.1);
        let result = CrankNicolsonSolver::new()
            .solve(&scenario, &config)
            .unwrap();
        assert_eq!(result.states.len(), result.times.len());
        assert!((result.times[0] - 0.0).abs() < 1e-12);
        assert!((result.t_final().unwrap() - 0.5).abs() < 1e-9);
    }

    #[test]
    fn n_steps_is_correct() {
        let scenario = Scenario::single(Box::new(ZeroDerivative), make_mesh(2));
        let config = make_config(1.0, 0.25);
        let result = CrankNicolsonSolver::new()
            .solve(&scenario, &config)
            .unwrap();
        assert_eq!(result.n_steps, 4);
    }

    #[test]
    fn save_every_reduces_stored_states() {
        let scenario = Scenario::single(Box::new(ZeroDerivative), make_mesh(2));
        let config = SolverConfiguration::new(
            TimeConfiguration::new(1.0, StepControl::Fixed { dt: 0.1 }).saving_every(5),
            IntegratorKind::CrankNicolson,
        );
        let result = CrankNicolsonSolver::new()
            .solve(&scenario, &config)
            .unwrap();
        assert_eq!(result.states.len(), 3);
    }

    #[test]
    fn step_matches_one_iteration_of_solve() {
        let scenario = Scenario::single(Box::new(ExponentialDecay { lambda: 0.7 }), make_mesh(3));
        let config = make_config(0.1, 0.1);

        let solver = CrankNicolsonSolver::new();
        let via_solve = solver.solve(&scenario, &config).unwrap();
        let final_via_solve = via_solve.states.last().unwrap().as_scalar_field().unwrap();

        let domain = scenario.single_domain().unwrap();
        let requirements = scenario.context_requirements();
        let chain =
            crate::solver::chain::build_calculator_chain(&requirements, &config.calculators)
                .unwrap();
        let mut u = domain.model.initial_state(domain.mesh.as_ref());
        let next = solver.step(domain, &chain, &mut u, &[], 0.0, 0.1).unwrap();
        let final_via_step = next.as_scalar_field().unwrap();

        assert_eq!(final_via_solve.len(), final_via_step.len());
        for i in 0..final_via_solve.len() {
            assert!((final_via_solve[i] - final_via_step[i]).abs() < 1e-15);
        }
    }

    #[test]
    fn stable_for_very_stiff_problem_over_many_steps() {
        // Crank-Nicolson is A-stable but not L-stable -- damping is
        // weaker than Backward Euler's for very stiff modes (oscillatory
        // decay is possible), but the solution must still stay bounded.
        let lambda = 1.0e4;
        let dt = 0.1;
        let scenario = Scenario::single(Box::new(ExponentialDecay { lambda }), make_mesh(2));
        let config = make_config(2.0, dt);
        let result = CrankNicolsonSolver::new()
            .solve(&scenario, &config)
            .unwrap();

        for state in &result.states {
            let field = state.as_scalar_field().unwrap();
            for v in field.iter() {
                assert!(v.is_finite(), "value diverged: {v}");
                assert!(v.abs() <= 1.0, "expected bounded solution, got {v}");
            }
        }
    }

    #[test]
    fn with_linear_solver_substitutes_backend() {
        let calls = Arc::new(AtomicUsize::new(0));
        let solver =
            CrankNicolsonSolver::new().with_linear_solver(Box::new(CountingLinearSolver {
                calls: calls.clone(),
            }));

        let scenario = Scenario::single(Box::new(ExponentialDecay { lambda: 1.0 }), make_mesh(2));
        let config = make_config(0.5, 0.1);

        solver.solve(&scenario, &config).unwrap();

        assert_eq!(calls.load(Ordering::SeqCst), 5);
    }

    #[test]
    fn negative_dt_returns_error() {
        let scenario = Scenario::single(Box::new(ZeroDerivative), make_mesh(2));
        let config = make_config(1.0, -0.1);
        assert!(CrankNicolsonSolver::new()
            .solve(&scenario, &config)
            .is_err());
    }

    #[test]
    fn t_end_before_t_start_returns_error() {
        let scenario = Scenario::single(Box::new(ZeroDerivative), make_mesh(2)).with_t_start(5.0);
        let config = make_config(1.0, 0.1);
        assert!(CrankNicolsonSolver::new()
            .solve(&scenario, &config)
            .is_err());
    }

    #[test]
    fn missing_calculator_returns_error() {
        #[derive(Debug)]
        struct NeedsExternal;
        impl RequiresContext for NeedsExternal {
            fn required_variables(&self) -> Vec<ContextVariable> {
                vec![ContextVariable::External {
                    name: "missing".into(),
                }]
            }
        }
        impl PhysicalModel for NeedsExternal {
            fn compute_physics(
                &self,
                s: &ContextValue,
                _: &ComputeContext,
            ) -> Result<ContextValue, OxiflowError> {
                Ok(s.clone())
            }
            fn initial_state(&self, mesh: &dyn Mesh) -> ContextValue {
                ContextValue::ScalarField(DVector::from_element(mesh.n_dof(), 0.0))
            }
            fn name(&self) -> &str {
                "needs_external"
            }
        }

        let scenario = Scenario::single(Box::new(NeedsExternal), make_mesh(2));
        let config = make_config(1.0, 0.1);
        let err = CrankNicolsonSolver::new()
            .solve(&scenario, &config)
            .unwrap_err();
        assert!(matches!(err, OxiflowError::MissingCalculator(_)));
    }
}