diffsol 0.17.0

use crate::{Matrix, NonLinearOpJacobian};
use diffsol_la::LinearSolver as LaLinearSolver;
use num_traits::Zero;

pub use diffsol_la::{FaerLU, FaerSparseLU, NalgebraLU, NalgebraNativeLU};

#[cfg(feature = "suitesparse")]
pub use diffsol_la::KLU;

#[cfg(feature = "cuda")]
pub use diffsol_la::CudaLU;

#[cfg(feature = "cuda-oxide")]
pub use diffsol_la::OxideLU;

/// Re-export of the [`diffsol_la`] linear-solver backends' module paths, so that
/// existing `crate::linear_solver::<backend>` paths keep resolving.
pub use diffsol_la::linear_solver::{faer, nalgebra};

#[cfg(feature = "suitesparse")]
pub use diffsol_la::linear_solver::suitesparse;

#[cfg(feature = "cuda")]
pub use diffsol_la::linear_solver::cuda;

#[cfg(feature = "cuda-oxide")]
pub use diffsol_la::linear_solver::cuda_oxide;

/// A borrowing adapter that presents a [NonLinearOpJacobian] (evaluated at a
/// fixed state `x` and time `t`) as a [`diffsol_la::LinearOp`].
///
/// This is the bridge that allows the time-aware, operator-based
/// [LinearSolver] trait in `diffsol` to be implemented on top of the
/// time-unaware [`diffsol_la::LinearSolver`] backends.
pub struct LinearisedRef<'a, C: NonLinearOpJacobian> {
    op: &'a C,
    x: Option<&'a C::V>,
    t: C::T,
}

impl<'a, C: NonLinearOpJacobian> LinearisedRef<'a, C> {
    /// Create an adapter used only to query the sparsity pattern (no state set).
    pub fn sparsity_only(op: &'a C) -> Self {
        Self {
            op,
            x: None,
            t: C::T::zero(),
        }
    }

    /// Create an adapter that evaluates the Jacobian at `(x, t)`.
    pub fn at(op: &'a C, x: &'a C::V, t: C::T) -> Self {
        Self { op, x: Some(x), t }
    }
}

impl<C: NonLinearOpJacobian> diffsol_la::LinearOp for LinearisedRef<'_, C> {
    type T = C::T;
    type V = C::V;
    type M = C::M;
    type C = C::C;

    fn nrows(&self) -> crate::IndexType {
        self.op.nout()
    }

    fn ncols(&self) -> crate::IndexType {
        self.op.nstates()
    }

    fn context(&self) -> &Self::C {
        self.op.context()
    }

    fn matrix_inplace(&self, y: &mut Self::M) {
        let x = self.x.expect("LinearisedRef: state x not set");
        self.op.jacobian_inplace(x, self.t, y);
    }

    fn sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
        self.op.jacobian_sparsity()
    }
}

/// A solver for the linear problem `Ax = b`, where `A` is a linear operator that is obtained by taking the linearisation of a nonlinear operator `C`.
///
/// This trait extends the time-unaware [`diffsol_la::LinearSolver`] with time-aware,
/// operator-based setup methods that linearise a [NonLinearOpJacobian] at a given state
/// and time (bridged via [LinearisedRef]). The actual solving is provided by the
/// [`diffsol_la::LinearSolver`] supertrait ([`diffsol_la::LinearSolver::solve`] and
/// [`diffsol_la::LinearSolver::solve_in_place`]).
pub trait LinearSolver<M: Matrix>: LaLinearSolver<M> {
    // sets the point at which the linearisation of the operator is evaluated
    // the operator is assumed to have the same sparsity as that given to [Self::set_problem]
    fn set_linearisation<C: NonLinearOpJacobian<V = M::V, T = M::T, M = M, C = M::C>>(
        &mut self,
        op: &C,
        x: &M::V,
        t: M::T,
    );

    /// Set the problem to be solved, any previous problem is discarded.
    /// Any internal state of the solver is reset.
    /// This function will normally set the sparsity pattern of the matrix to be solved.
    fn set_problem<C: NonLinearOpJacobian<V = M::V, T = M::T, M = M, C = M::C>>(&mut self, op: &C);
}

/// Any [`diffsol_la::LinearSolver`] backend automatically implements the
/// time-aware, operator-based [LinearSolver] trait via the [LinearisedRef] bridge.
impl<M: Matrix, LS: LaLinearSolver<M>> LinearSolver<M> for LS {
    fn set_linearisation<C: NonLinearOpJacobian<V = M::V, T = M::T, M = M, C = M::C>>(
        &mut self,
        op: &C,
        x: &M::V,
        t: M::T,
    ) {
        LaLinearSolver::set_linearisation(self, &LinearisedRef::at(op, x, t));
    }

    fn set_problem<C: NonLinearOpJacobian<V = M::V, T = M::T, M = M, C = M::C>>(&mut self, op: &C) {
        LaLinearSolver::set_sparsity(self, &LinearisedRef::sparsity_only(op));
    }
}

pub struct LinearSolveSolution<V> {
    pub x: V,
    pub b: V,
}

impl<V> LinearSolveSolution<V> {
    pub fn new(b: V, x: V) -> Self {
        Self { x, b }
    }
}

#[cfg(test)]
pub mod tests {
    use crate::{
        matrix::dense_nalgebra_serial::NalgebraMat,
        op::{closure::Closure, ParameterisedOp},
        scalar::scale,
        vector::VectorRef,
        Context, FaerMat, LinearSolver, Matrix, NonLinearOpJacobian, Vector,
    };
    use num_traits::{FromPrimitive, One, Zero};

    use super::LinearSolveSolution;
    use super::{FaerLU, NalgebraLU, NalgebraNativeLU};

    /// `y = diag(p) * x`, one lane at a time.
    fn diagonal_mul<M: Matrix>(x: &[M::T], p: &[M::T], _t: M::T, y: &mut [M::T]) {
        for ((y, p), x) in y.iter_mut().zip(p.iter()).zip(x.iter()) {
            *y = *p * *x;
        }
    }

    #[allow(clippy::type_complexity)]
    pub fn linear_problem<M: Matrix + 'static>() -> (
        Closure<
            M,
            impl Fn(&[M::T], &[M::T], M::T, &mut [M::T]),
            impl Fn(&[M::T], &[M::T], M::T, &[M::T], &mut [M::T]),
        >,
        M::V,
        M::T,
        M::V,
        Vec<LinearSolveSolution<M::V>>,
    ) {
        let ctx = M::C::default();
        // the operator is `diag(p) * x`, with the diagonal carried in the parameters so that a
        // batched context can give each lane its own
        let p = M::V::from_vec(
            vec![M::T::from_f64(2.0).unwrap(), M::T::from_f64(2.0).unwrap()],
            ctx.clone(),
        );
        let mut op = Closure::new(
            diagonal_mul::<M>,
            |_x: &[M::T], p: &[M::T], _t, v: &[M::T], y: &mut [M::T]| {
                diagonal_mul::<M>(v, p, _t, y)
            },
            2,
            2,
            p.len(),
            ctx.clone(),
        );
        op.calculate_sparsity(
            &M::V::from_element(2, M::T::one(), ctx.clone()),
            M::T::zero(),
            &p,
        );
        let rtol = M::T::from_f64(1e-6).unwrap();
        let atol = M::V::from_vec(
            vec![M::T::from_f64(1e-6).unwrap(), M::T::from_f64(1e-6).unwrap()],
            ctx.clone(),
        );
        let solns = vec![LinearSolveSolution::new(
            M::V::from_vec(
                vec![M::T::from_f64(2.0).unwrap(), M::T::from_f64(4.0).unwrap()],
                ctx.clone(),
            ),
            M::V::from_vec(vec![M::T::one(), M::T::from_f64(2.0).unwrap()], ctx.clone()),
        )];
        (op, p, rtol, atol, solns)
    }

    pub fn test_linear_solver<'a, C>(
        mut solver: impl LinearSolver<C::M>,
        op: C,
        rtol: C::T,
        atol: &'a C::V,
        solns: Vec<LinearSolveSolution<C::V>>,
    ) where
        C: NonLinearOpJacobian,
        for<'b> &'b C::V: VectorRef<C::V>,
    {
        solver.set_problem(&op);
        let x = C::V::zeros(op.nout(), op.context().clone());
        let t = C::T::zero();
        LinearSolver::set_linearisation(&mut solver, &op, &x, t);
        for soln in solns {
            let x = solver.solve(&soln.b).unwrap();
            let tol = { &soln.x * scale(rtol) + atol };
            x.assert_eq(&soln.x, &tol);
        }
    }

    #[allow(clippy::type_complexity)]
    pub fn linear_problem_batched<M: Matrix + 'static>(
        ctx: M::C,
    ) -> (
        Closure<
            M,
            impl Fn(&[M::T], &[M::T], M::T, &mut [M::T]),
            impl Fn(&[M::T], &[M::T], M::T, &[M::T], &mut [M::T]),
        >,
        M::V,
        M::T,
        M::V,
        Vec<LinearSolveSolution<M::V>>,
    ) {
        assert_eq!(ctx.nbatch(), 2);
        let two = M::T::from_f64(2.0).unwrap();
        let three = M::T::from_f64(3.0).unwrap();
        // one diagonal per lane, carried in the parameters
        let p = M::V::from_vec(vec![two, two, three, three], ctx.clone());
        let mut op = Closure::new(
            diagonal_mul::<M>,
            |_x: &[M::T], p: &[M::T], _t, v: &[M::T], y: &mut [M::T]| {
                diagonal_mul::<M>(v, p, _t, y)
            },
            2,
            2,
            p.len(),
            ctx.clone(),
        );
        op.calculate_sparsity(
            &M::V::from_element(2, M::T::one(), ctx.clone()),
            M::T::zero(),
            &p,
        );
        let rtol = M::T::from_f64(1e-6).unwrap();
        let atol_val = M::T::from_f64(1e-6).unwrap();
        let atol = M::V::from_vec(vec![atol_val; 4], ctx.clone());
        let one = M::T::one();
        let four = M::T::from_f64(4.0).unwrap();
        let six = M::T::from_f64(6.0).unwrap();
        let nine = M::T::from_f64(9.0).unwrap();
        let solns = vec![LinearSolveSolution::new(
            M::V::from_vec(vec![two, four, six, nine], ctx.clone()),
            M::V::from_vec(vec![one, two, two, three], ctx),
        )];
        (op, p, rtol, atol, solns)
    }

    #[test]
    fn test_lu_nalgebra() {
        let (op, p, rtol, atol, solns) = linear_problem::<NalgebraMat<f64>>();
        let op = ParameterisedOp::new(&op, &p);
        let s = NalgebraLU::default();
        test_linear_solver(s, op, rtol, &atol, solns);
    }

    #[test]
    fn test_lu_nalgebra_native() {
        let (op, p, rtol, atol, solns) = linear_problem::<NalgebraMat<f64>>();
        let op = ParameterisedOp::new(&op, &p);
        let s = NalgebraNativeLU::default();
        test_linear_solver(s, op, rtol, &atol, solns);
    }

    #[test]
    fn test_lu_faer() {
        let (op, p, rtol, atol, solns) = linear_problem::<FaerMat<f64>>();
        let op = ParameterisedOp::new(&op, &p);
        let s = FaerLU::default();
        test_linear_solver(s, op, rtol, &atol, solns);
    }

    #[test]
    fn test_sparse_lu_faer() {
        use crate::FaerSparseMat;
        let (op, p, rtol, atol, solns) = linear_problem::<FaerSparseMat<f64>>();
        let op = ParameterisedOp::new(&op, &p);
        let s = super::FaerSparseLU::default();
        test_linear_solver(s, op, rtol, &atol, solns);
    }

    #[cfg(feature = "suitesparse")]
    #[test]
    fn test_klu() {
        use crate::FaerSparseMat;
        let (op, p, rtol, atol, solns) = linear_problem::<FaerSparseMat<f64>>();
        let op = ParameterisedOp::new(&op, &p);
        let s = super::KLU::default();
        test_linear_solver(s, op, rtol, &atol, solns);
    }

    #[cfg(feature = "cuda")]
    #[test]
    fn test_lu_cuda() {
        use crate::{CudaLU, CudaMat};
        let (op, p, rtol, atol, solns) = linear_problem::<CudaMat<f64>>();
        let op = ParameterisedOp::new(&op, &p);
        let s = CudaLU::default();
        test_linear_solver(s, op, rtol, &atol, solns);
    }

    #[cfg(feature = "cuda")]
    #[test]
    fn test_lu_cuda_batched() {
        use crate::{CudaContext, CudaLU, CudaMat};
        let ctx = CudaContext::default().with_nbatch(2);
        let (op, p, rtol, atol, solns) = linear_problem_batched::<CudaMat<f64>>(ctx);
        let op = ParameterisedOp::new(&op, &p);
        let s = CudaLU::default();
        test_linear_solver(s, op, rtol, &atol, solns);
    }
}