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 = RefCell::new(MemBuffer::new(
39            symbolic
40                .solve_in_place_scratch::<T>(1, par)
41                .or(symbolic.solve_transpose_in_place_scratch::<T>(1, par)),
42        ));
43        Self {
44            numeric: (0..a.data.len()).map(|_| NumericLu::new()).collect(),
45            a,
46            symbolic,
47            factor_scratch,
48            solve_scratch,
49        }
50    }
51
52    fn nbatch(&self) -> usize {
53        self.numeric.len()
54    }
55
56    fn factor(&mut self) {
57        let par = self.a.context.par;
58        for (numeric, block) in self.numeric.iter_mut().zip(self.a.data.iter()) {
59            self.symbolic
60                .factorize_numeric_lu(
61                    numeric,
62                    block.rb(),
63                    par,
64                    MemStack::new(&mut self.factor_scratch),
65                    Default::default(),
66                )
67                .expect("Failed to factorise matrix");
68        }
69    }
70
71    fn solve_block(&self, batch: usize, rhs: faer::MatMut<'_, T>, par: faer::Par, transpose: bool) {
72        let lu = LuRef::new_unchecked(&self.symbolic, &self.numeric[batch]);
73        let mut scratch = self.solve_scratch.borrow_mut();
74        let stack = MemStack::new(&mut scratch);
75        if transpose {
76            lu.solve_transpose_in_place_with_conj(Conj::No, rhs, par, stack);
77        } else {
78            lu.solve_in_place_with_conj(Conj::No, rhs, par, stack);
79        }
80    }
81}
82
83/// A [LinearSolver] that uses the LU decomposition in the [`faer`](https://github.com/sarah-ek/faer-rs) library to solve the linear system.
84///
85/// The low-level `faer::sparse::linalg::lu` API is used rather than
86/// `faer::sparse::linalg::solvers` both so that the parallelism from [`FaerContext::par`] is
87/// honoured -- the wrapper types hardcode faer's process-global setting -- and so that the
88/// factorizations and scratch space can be allocated once in [`Self::set_sparsity`] and reused.
89pub struct FaerSparseLU<T>
90where
91    T: FaerScalar,
92{
93    lu: Option<SparseLu<T>>,
94}
95
96impl<T> Default for FaerSparseLU<T>
97where
98    T: FaerScalar,
99{
100    fn default() -> Self {
101        Self { lu: None }
102    }
103}
104
105impl<T: FaerScalar> LinearSolver<FaerSparseMat<T>> for FaerSparseLU<T> {
106    fn set_linearisation<
107        C: LinearOp<T = T, V = FaerVec<T>, M = FaerSparseMat<T>, C = FaerContext>,
108    >(
109        &mut self,
110        op: &C,
111    ) {
112        let lu = self.lu.as_mut().expect("Matrix not set");
113        op.matrix_inplace(&mut lu.a);
114        lu.factor();
115    }
116
117    fn solve_in_place(&self, x: &mut FaerVec<T>) -> Result<(), LaError> {
118        self.solve_in_place_impl(x, false)
119    }
120
121    fn solve_transpose_in_place(&self, x: &mut FaerVec<T>) -> Result<(), LaError> {
122        self.solve_in_place_impl(x, true)
123    }
124
125    fn set_sparsity<C: LinearOp<T = T, V = FaerVec<T>, M = FaerSparseMat<T>, C = FaerContext>>(
126        &mut self,
127        op: &C,
128    ) {
129        let ncols = op.ncols();
130        let nrows = op.nrows();
131        let a = C::M::new_from_sparsity(nrows, ncols, op.sparsity(), *op.context());
132        self.lu = Some(SparseLu::new(a));
133    }
134}
135
136impl<T: FaerScalar> FaerSparseLU<T> {
137    fn solve_in_place_impl(&self, x: &mut FaerVec<T>, transpose: bool) -> Result<(), LaError> {
138        let Some(lu) = self.lu.as_ref() else {
139            return Err(linear_solver_error!(LuNotInitialized));
140        };
141        x.context
142            .assert_broadcastable_into(lu.nbatch(), "sparse_lu_solve");
143        let par = x.context.par;
144        let nb = x.data.ncols();
145        for batch in 0..nb {
146            lu.solve_block(
147                broadcast_batch(batch, lu.nbatch(), nb),
148                x.data.rb_mut().col_mut(batch).as_mat_mut(),
149                par,
150                transpose,
151            );
152        }
153        Ok(())
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160
161    crate::linear_solver::generate_lu_tests!(
162        f64,
163        FaerSparseMat<f64>,
164        FaerSparseLU<f64>,
165        FaerContext::with_nbatch(2)
166    );
167    crate::linear_solver::generate_lu_tests!(
168        f32,
169        FaerSparseMat<f32>,
170        FaerSparseLU<f32>,
171        FaerContext::with_nbatch(2)
172    );
173}
diffsol_la/linear_solver/faer/sparse_lu.rs

diffsol_la/linear_solver/faer/
sparse_lu.rs
