diffsol-la 0.3.0

use crate::{error::LaError, LinearOp, Matrix};

#[cfg(feature = "nalgebra")]
pub mod nalgebra;

#[cfg(feature = "faer")]
pub mod faer;

#[cfg(feature = "suitesparse")]
pub mod suitesparse;

#[cfg(feature = "cuda")]
#[allow(deprecated)]
pub mod cuda;

#[cfg(feature = "cuda-oxide")]
pub mod cuda_oxide;

pub use faer::lu::LU as FaerLU;
pub use nalgebra::lu::ReusableLU as NalgebraLU;

/// A solver for the linear problem `Ax = b`, where `A` is a [LinearOp].
pub trait LinearSolver<M: Matrix>: Default {
    /// Set the point at which the linear operator `A` is evaluated and factorise it.
    ///
    /// The operator is assumed to have the same sparsity as that given to
    /// [Self::set_sparsity].
    fn set_linearisation<C: LinearOp<V = M::V, T = M::T, M = M, C = M::C>>(&mut self, op: &C);

    /// Set the sparsity of 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_sparsity<C: LinearOp<V = M::V, T = M::T, M = M, C = M::C>>(&mut self, op: &C);

    /// Solve the problem `Ax = b` and return the solution `x`.
    ///
    /// Panics if [Self::set_linearisation] has not been called previously.
    fn solve(&self, b: &M::V) -> Result<M::V, LaError> {
        let mut b = b.clone();
        self.solve_in_place(&mut b)?;
        Ok(b)
    }

    fn solve_in_place(&self, b: &mut M::V) -> Result<(), LaError>;

    /// Solve the transposed problem `A^T x = b` in place, reusing the factorisation of `A`.
    ///
    /// Returns an error if [Self::set_linearisation] has not been called previously.
    fn solve_transpose_in_place(&self, b: &mut M::V) -> Result<(), LaError>;
}

/// The shared [LinearSolver] tests for solver `$LS` on matrix `$M`, with `$ctx2` a context of
/// two batches.
#[cfg(test)]
macro_rules! generate_lu_tests {
    ($suffix:ident, $M:ty, $LS:ty, $ctx2:expr) => {
        paste::paste! {
            #[test]
            fn [<test_lu_ $suffix>]() {
                $crate::linear_solver::tests::test_diagonal_lu_solve::<$M, $LS>();
            }
            #[test]
            fn [<test_grouped_lu_ $suffix>]() {
                $crate::linear_solver::tests::test_grouped_lu_solve::<$M, $LS>($ctx2);
            }
            #[test]
            fn [<test_transpose_lu_ $suffix>]() {
                $crate::linear_solver::tests::test_transpose_lu_solve::<$M, $LS>($ctx2);
            }
            #[test]
            #[should_panic(expected = "incompatible nbatch")]
            fn [<test_narrow_state_lu_ $suffix>]() {
                $crate::linear_solver::tests::test_narrow_state_lu_solve::<$M, $LS>($ctx2);
            }
        }
    };
}
#[cfg(test)]
pub(crate) use generate_lu_tests;

#[cfg(test)]
pub(crate) mod tests {
    use super::LinearSolver;
    use crate::{Context, IndexType, LinearOp, Matrix, Vector};

    /// A simple diagonal [LinearOp] used for testing the linear solvers.
    pub struct DiagonalOp<M: Matrix> {
        pub matrix: M,
    }

    /// Create a 2x2 diagonal operator `A = diag(value, value)`.
    pub fn diagonal_op<M: Matrix>(value: f64) -> DiagonalOp<M> {
        use num_traits::FromPrimitive;
        let v = M::T::from_f64(value).unwrap();
        let diag = M::V::from_vec(vec![v, v], Default::default());
        DiagonalOp {
            matrix: M::from_diagonal(&diag),
        }
    }

    /// Create an `nstates`-square diagonal operator `A = diag(value, ..., value)`.
    #[cfg_attr(not(feature = "cuda-oxide"), allow(dead_code))]
    pub fn diagonal_op_n<M: Matrix>(nstates: IndexType, value: f64, ctx: M::C) -> DiagonalOp<M> {
        use num_traits::FromPrimitive;
        let v = M::T::from_f64(value).unwrap();
        let diag = M::V::from_element(nstates, v, ctx);
        DiagonalOp {
            matrix: M::from_diagonal(&diag),
        }
    }

    /// Create a batched 2x2 diagonal operator with `A_b = diag(values[b], values[b])`.
    pub fn batched_diagonal_op<M: Matrix>(values: &[f64], ctx: M::C) -> DiagonalOp<M> {
        use num_traits::FromPrimitive;
        assert_eq!(values.len(), ctx.nbatch());
        let diag = M::V::from_vec(
            values
                .iter()
                .flat_map(|&v| [M::T::from_f64(v).unwrap(); 2])
                .collect(),
            ctx,
        );
        DiagonalOp {
            matrix: M::from_diagonal(&diag),
        }
    }

    /// Solve `diag(2, 2) x = (2, 4)`.
    pub fn test_diagonal_lu_solve<M: Matrix, LS: LinearSolver<M>>() {
        use num_traits::FromPrimitive;
        let f = |x: f64| M::T::from_f64(x).unwrap();
        let op = diagonal_op::<M>(2.0);
        let mut s = LS::default();
        s.set_sparsity(&op);
        s.set_linearisation(&op);
        let b = M::V::from_vec(vec![f(2.0), f(4.0)], Default::default());
        let x = s.solve(&b).unwrap();
        x.assert_eq_eps(
            &M::V::from_vec(vec![f(1.0), f(2.0)], Default::default()),
            100.0,
        );
    }

    /// One factorization per group: `A` carries 2 batches and the right-hand side 4, so
    /// batches 0 and 1 solve against `A_0` and batches 2 and 3 against `A_1`.  A cyclic
    /// mapping would pair batch 1 with `A_1` and fail.
    pub fn test_grouped_lu_solve<M: Matrix, LS: LinearSolver<M>>(ctx2: M::C) {
        use num_traits::FromPrimitive;
        let f = |x: f64| M::T::from_f64(x).unwrap();
        let wide = ctx2.clone_with_nbatch(4).unwrap();
        let op = batched_diagonal_op::<M>(&[2.0, 5.0], ctx2);
        let mut s = LS::default();
        s.set_sparsity(&op);
        s.set_linearisation(&op);
        let b = M::V::from_vec((1..=8).map(|i| f(i as f64)).collect(), wide.clone());
        let x = s.solve(&b).unwrap();
        let expected = M::V::from_vec(
            vec![
                f(0.5),
                f(1.0),
                f(1.5),
                f(2.0),
                f(1.0),
                f(1.2),
                f(1.4),
                f(1.6),
            ],
            wide,
        );
        x.assert_eq_eps(&expected, 100.0);
    }

    /// Solve `A x = b` and `A^T x = b` with `A_0 = [[0, 1, 2], [3, 1, 0], [1, 0, 4]]` and
    /// `A_1 = A_0^T`, both needing a pivot, against 4 right-hand side batches grouped as in
    /// [test_grouped_lu_solve]. Every solution is `(1, 2, 3)`.
    pub fn test_transpose_lu_solve<M: Matrix, LS: LinearSolver<M>>(ctx2: M::C) {
        use num_traits::FromPrimitive;
        let f = |x: f64| M::T::from_f64(x).unwrap();
        let wide = ctx2.clone_with_nbatch(4).unwrap();
        let indices = vec![(0, 1), (0, 2), (1, 0), (1, 1), (2, 0), (2, 2)];
        let values = [1.0, 2.0, 3.0, 1.0, 1.0, 4.0, 3.0, 1.0, 1.0, 1.0, 2.0, 4.0];
        let matrix =
            M::try_from_triplets(3, 3, indices, values.iter().map(|&v| f(v)).collect(), ctx2)
                .unwrap();
        let op = DiagonalOp { matrix };
        let mut s = LS::default();
        s.set_sparsity(&op);
        s.set_linearisation(&op);
        let rhs = |a: [f64; 3], b: [f64; 3]| {
            M::V::from_vec(
                [a, a, b, b].iter().flatten().map(|&v| f(v)).collect(),
                wide.clone(),
            )
        };
        let expected = rhs([1.0, 2.0, 3.0], [1.0, 2.0, 3.0]);
        let (a_x, at_x) = ([8.0, 5.0, 13.0], [9.0, 3.0, 14.0]);

        let mut b = rhs(a_x, at_x);
        s.solve_in_place(&mut b).unwrap();
        b.assert_eq_eps(&expected, 100.0);

        let mut b = rhs(at_x, a_x);
        s.solve_transpose_in_place(&mut b).unwrap();
        b.assert_eq_eps(&expected, 100.0);
    }

    /// More factorizations than right-hand side batches would leave factorizations unused, so
    /// the solve panics rather than silently dropping them.
    pub fn test_narrow_state_lu_solve<M: Matrix, LS: LinearSolver<M>>(ctx2: M::C) {
        let wide = ctx2.clone_with_nbatch(4).unwrap();
        let op = batched_diagonal_op::<M>(&[2.0, 3.0, 4.0, 5.0], wide);
        let mut s = LS::default();
        s.set_sparsity(&op);
        s.set_linearisation(&op);
        let mut b = M::V::zeros(2, ctx2);
        s.solve_in_place(&mut b).unwrap();
    }

    impl<M: Matrix> LinearOp for DiagonalOp<M> {
        type T = M::T;
        type V = M::V;
        type M = M;
        type C = M::C;

        fn nrows(&self) -> IndexType {
            self.matrix.nrows()
        }
        fn ncols(&self) -> IndexType {
            self.matrix.ncols()
        }
        fn context(&self) -> &Self::C {
            self.matrix.context()
        }
        fn matrix_inplace(&self, y: &mut Self::M) {
            y.copy_from(&self.matrix);
        }
        fn sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
            self.matrix.sparsity().map(|s| {
                use crate::matrix::sparsity::MatrixSparsityRef;
                s.to_owned()
            })
        }
    }
}