diffsol-la 0.3.0

use crate::context::broadcast_batch;
use crate::{error::LaError, linear_solver_error};
use crate::{Context, FaerContext};

use crate::{
    linear_solver::LinearSolver, FaerMat, FaerScalar, FaerVec, LinearOp, Matrix, MatrixCommon,
};

use faer::dyn_stack::{MemBuffer, MemStack};
use faer::linalg::lu::full_pivoting::{factor, solve};
use faer::perm::Perm;
use faer::reborrow::{Reborrow, ReborrowMut};
use faer::{Conj, Mat, Par};

/// The $LU$ factors of one batch, stored packed: the strict lower triangle of `lu` is $L$ (unit
/// diagonal), the upper triangle is $U$.  faer's triangular solves read only their own triangle
/// and substitute ones on $L$'s diagonal, so the same matrix serves as both operands and no
/// splitting copy is needed.
struct PackedLu<T: FaerScalar> {
    lu: Mat<T>,
    row: Perm<usize>,
    col: Perm<usize>,
}

impl<T: FaerScalar> PackedLu<T> {
    fn new(a: faer::MatRef<'_, T>, par: Par) -> Self {
        let (m, n) = a.shape();
        let mut lu = a.to_owned();
        let mut row_fwd = vec![0usize; m];
        let mut row_bwd = vec![0usize; m];
        let mut col_fwd = vec![0usize; n];
        let mut col_bwd = vec![0usize; n];
        factor::lu_in_place(
            lu.as_mut(),
            &mut row_fwd,
            &mut row_bwd,
            &mut col_fwd,
            &mut col_bwd,
            par,
            MemStack::new(&mut MemBuffer::new(
                factor::lu_in_place_scratch::<usize, T>(m, n, par, Default::default()),
            )),
            Default::default(),
        );
        Self {
            lu,
            row: Perm::new_checked(row_fwd.into_boxed_slice(), row_bwd.into_boxed_slice(), m),
            col: Perm::new_checked(col_fwd.into_boxed_slice(), col_bwd.into_boxed_slice(), n),
        }
    }

    fn solve_in_place(&self, rhs: faer::MatMut<'_, T>, par: Par, transpose: bool) {
        let n = self.lu.nrows();
        let k = rhs.ncols();
        if transpose {
            solve::solve_transpose_in_place_with_conj(
                self.lu.rb(),
                self.lu.rb(),
                self.row.rb(),
                self.col.rb(),
                Conj::No,
                rhs,
                par,
                MemStack::new(&mut MemBuffer::new(
                    solve::solve_transpose_in_place_scratch::<usize, T>(n, k, par),
                )),
            );
        } else {
            solve::solve_in_place_with_conj(
                self.lu.rb(),
                self.lu.rb(),
                self.row.rb(),
                self.col.rb(),
                Conj::No,
                rhs,
                par,
                MemStack::new(&mut MemBuffer::new(
                    solve::solve_in_place_scratch::<usize, T>(n, k, par),
                )),
            );
        }
    }
}

/// A [LinearSolver] that uses the LU decomposition in the [`faer`](https://github.com/sarah-ek/faer-rs) library to solve the linear system.
///
/// The low-level `faer::linalg::lu` API is used rather than `faer::linalg::solvers` so that the
/// parallelism from [`FaerContext::par`] is honoured; the wrapper types hardcode faer's
/// process-global setting instead.
pub struct LU<T>
where
    T: FaerScalar,
{
    lu: Vec<PackedLu<T>>,
    matrix: Option<FaerMat<T>>,
}

impl<T> Default for LU<T>
where
    T: FaerScalar,
{
    fn default() -> Self {
        Self {
            lu: Vec::new(),
            matrix: None,
        }
    }
}

impl<T: FaerScalar> LinearSolver<FaerMat<T>> for LU<T> {
    fn set_linearisation<C: LinearOp<T = T, V = FaerVec<T>, M = FaerMat<T>, C = FaerContext>>(
        &mut self,
        op: &C,
    ) {
        let matrix = self.matrix.as_mut().expect("Matrix not set");
        op.matrix_inplace(matrix);
        let nc = matrix.ncols();
        let par = matrix.context.par;
        self.lu = (0..matrix.context.nbatch())
            .map(|b| PackedLu::new(matrix.data.rb().subcols(b * nc, nc), par))
            .collect();
    }

    fn solve_in_place(&self, x: &mut FaerVec<T>) -> Result<(), LaError> {
        self.solve_in_place_impl(x, false)
    }

    fn solve_transpose_in_place(&self, x: &mut FaerVec<T>) -> Result<(), LaError> {
        self.solve_in_place_impl(x, true)
    }

    fn set_sparsity<C: LinearOp<T = T, V = FaerVec<T>, M = FaerMat<T>, C = FaerContext>>(
        &mut self,
        op: &C,
    ) {
        let ncols = op.ncols();
        let nrows = op.nrows();
        let matrix = C::M::new_from_sparsity(nrows, ncols, op.sparsity(), *op.context());
        self.matrix = Some(matrix);
    }
}

impl<T: FaerScalar> LU<T> {
    fn solve_in_place_impl(&self, x: &mut FaerVec<T>, transpose: bool) -> Result<(), LaError> {
        if self.lu.is_empty() {
            return Err(linear_solver_error!(LuNotInitialized));
        }
        x.context
            .assert_broadcastable_into(self.lu.len(), "lu_solve");
        let par = x.context.par;
        let nlu = self.lu.len();
        let nb = x.data.ncols();
        for batch in 0..nb {
            self.lu[broadcast_batch(batch, nlu, nb)].solve_in_place(
                x.data.rb_mut().col_mut(batch).as_mat_mut(),
                par,
                transpose,
            );
        }
        Ok(())
    }
}

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

    crate::linear_solver::generate_lu_tests!(
        f64,
        FaerMat<f64>,
        LU<f64>,
        FaerContext::with_nbatch(2)
    );
    crate::linear_solver::generate_lu_tests!(
        f32,
        FaerMat<f32>,
        LU<f32>,
        FaerContext::with_nbatch(2)
    );
}