diffsol_la/linear_solver/faer/
sparse_lu.rs

1use std::cell::RefCell;
2
3use crate::context::broadcast_batch;
4use crate::{
5    error::LaError, linear_solver::LinearSolver, linear_solver_error, scalar::IndexType, Context,
6    FaerContext, FaerScalar, FaerSparseMat, FaerVec, LinearOp, Matrix,
7};
8
9use faer::dyn_stack::{MemBuffer, MemStack};
10use faer::reborrow::{Reborrow, ReborrowMut};
11use faer::sparse::linalg::lu::{factorize_symbolic_lu, LuRef, NumericLu, SymbolicLu};
12use faer::Conj;
13
14/// Sparse $LU$ factorization state for a (possibly batched) [FaerSparseMat].
15///
16/// The symbolic factorization, one numeric factorization per batch and both scratch buffers are
17/// allocated once, in [`FaerSparseLU::set_sparsity`].  faer sizes a [`NumericLu`]'s own buffers on
18/// its first `factorize_numeric_lu` rather than on construction, so the first
19/// [`FaerSparseLU::set_linearisation`] after that still allocates inside faer; later ones reuse
20/// the capacity.
21struct SparseLu<T: FaerScalar> {
22    a: FaerSparseMat<T>,
23    symbolic: SymbolicLu<IndexType>,
24    numeric: Vec<NumericLu<IndexType, T>>,
25    factor_scratch: MemBuffer,
26    solve_scratch: RefCell<MemBuffer>,
27}
28
29impl<T: FaerScalar> SparseLu<T> {
30    fn new(a: FaerSparseMat<T>) -> Self {
31        let symbolic = factorize_symbolic_lu(a.data[0].symbolic(), Default::default())
32            .expect("Failed to create symbolic LU");
33        let par = a.context.par;
34        // the solve always runs one column at a time, so this size holds however many batches the
35        // right-hand side broadcasts into
36        let factor_scratch =
37            MemBuffer::new(symbolic.factorize_numeric_lu_scratch::<T>(par, Default::default()));
38        let solve_scratch =
39            RefCell::new(MemBuffer::new(symbolic.solve_in_place_scratch::<T>(1, par)));
40        Self {
41            numeric: (0..a.data.len()).map(|_| NumericLu::new()).collect(),
42            a,
43            symbolic,
44            factor_scratch,
45            solve_scratch,
46        }
47    }
48
49    fn nbatch(&self) -> usize {
50        self.numeric.len()
51    }
52
53    fn factor(&mut self) {
54        let par = self.a.context.par;
55        for (numeric, block) in self.numeric.iter_mut().zip(self.a.data.iter()) {
56            self.symbolic
57                .factorize_numeric_lu(
58                    numeric,
59                    block.rb(),
60                    par,
61                    MemStack::new(&mut self.factor_scratch),
62                    Default::default(),
63                )
64                .expect("Failed to factorise matrix");
65        }
66    }
67
68    fn solve_block(&self, batch: usize, rhs: faer::MatMut<'_, T>, par: faer::Par) {
69        LuRef::new_unchecked(&self.symbolic, &self.numeric[batch]).solve_in_place_with_conj(
70            Conj::No,
71            rhs,
72            par,
73            MemStack::new(&mut self.solve_scratch.borrow_mut()),
74        );
75    }
76}
77
78/// A [LinearSolver] that uses the LU decomposition in the [`faer`](https://github.com/sarah-ek/faer-rs) library to solve the linear system.
79///
80/// The low-level `faer::sparse::linalg::lu` API is used rather than
81/// `faer::sparse::linalg::solvers` both so that the parallelism from [`FaerContext::par`] is
82/// honoured -- the wrapper types hardcode faer's process-global setting -- and so that the
83/// factorizations and scratch space can be allocated once in [`Self::set_sparsity`] and reused.
84pub struct FaerSparseLU<T>
85where
86    T: FaerScalar,
87{
88    lu: Option<SparseLu<T>>,
89}
90
91impl<T> Default for FaerSparseLU<T>
92where
93    T: FaerScalar,
94{
95    fn default() -> Self {
96        Self { lu: None }
97    }
98}
99
100impl<T: FaerScalar> LinearSolver<FaerSparseMat<T>> for FaerSparseLU<T> {
101    fn set_linearisation<
102        C: LinearOp<T = T, V = FaerVec<T>, M = FaerSparseMat<T>, C = FaerContext>,
103    >(
104        &mut self,
105        op: &C,
106    ) {
107        let lu = self.lu.as_mut().expect("Matrix not set");
108        op.matrix_inplace(&mut lu.a);
109        lu.factor();
110    }
111
112    fn solve_in_place(&self, x: &mut FaerVec<T>) -> Result<(), LaError> {
113        let Some(lu) = self.lu.as_ref() else {
114            return Err(linear_solver_error!(LuNotInitialized));
115        };
116        x.context
117            .assert_broadcastable_into(lu.nbatch(), "sparse_lu_solve");
118        let par = x.context.par;
119        let nb = x.data.ncols();
120        for batch in 0..nb {
121            lu.solve_block(
122                broadcast_batch(batch, lu.nbatch(), nb),
123                x.data.rb_mut().col_mut(batch).as_mat_mut(),
124                par,
125            );
126        }
127        Ok(())
128    }
129
130    fn set_sparsity<C: LinearOp<T = T, V = FaerVec<T>, M = FaerSparseMat<T>, C = FaerContext>>(
131        &mut self,
132        op: &C,
133    ) {
134        let ncols = op.ncols();
135        let nrows = op.nrows();
136        let a = C::M::new_from_sparsity(nrows, ncols, op.sparsity(), *op.context());
137        self.lu = Some(SparseLu::new(a));
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144    use crate::{
145        linear_solver::tests::{diagonal_op, test_grouped_lu_solve, test_narrow_state_lu_solve},
146        Vector,
147    };
148
149    #[test]
150    fn test_sparse_lu() {
151        let mut s = FaerSparseLU::<f64>::default();
152        let op = diagonal_op::<FaerSparseMat<f64>>(2.0);
153        s.set_sparsity(&op);
154        s.set_linearisation(&op);
155        let b = FaerVec::from_vec(vec![2.0, 4.0], Default::default());
156        let x = s.solve(&b).unwrap();
157        x.assert_eq_st(
158            &FaerVec::from_vec(vec![1.0, 2.0], Default::default()),
159            1e-10,
160        );
161    }
162
163    #[test]
164    fn test_grouped_sparse_lu() {
165        test_grouped_lu_solve::<FaerSparseMat<f64>, FaerSparseLU<f64>>(FaerContext::with_nbatch(2));
166    }
167
168    #[test]
169    #[should_panic(expected = "incompatible nbatch")]
170    fn test_narrow_state_sparse_lu() {
171        test_narrow_state_lu_solve::<FaerSparseMat<f64>, FaerSparseLU<f64>>(
172            FaerContext::with_nbatch(2),
173        );
174    }
175}
diffsol_la/linear_solver/faer/sparse_lu.rs

diffsol_la/linear_solver/faer/
sparse_lu.rs
