diffsol_la/linear_solver/faer/
lu.rs

1use crate::context::broadcast_batch;
2use crate::{error::LaError, linear_solver_error};
3use crate::{Context, FaerContext};
4
5use crate::{
6    linear_solver::LinearSolver, FaerMat, FaerScalar, FaerVec, LinearOp, Matrix, MatrixCommon,
7};
8
9use faer::dyn_stack::{MemBuffer, MemStack};
10use faer::linalg::lu::full_pivoting::{factor, solve};
11use faer::perm::Perm;
12use faer::reborrow::{Reborrow, ReborrowMut};
13use faer::{Conj, Mat, Par};
14
15/// The $LU$ factors of one batch, stored packed: the strict lower triangle of `lu` is $L$ (unit
16/// diagonal), the upper triangle is $U$.  faer's triangular solves read only their own triangle
17/// and substitute ones on $L$'s diagonal, so the same matrix serves as both operands and no
18/// splitting copy is needed.
19struct PackedLu<T: FaerScalar> {
20    lu: Mat<T>,
21    row: Perm<usize>,
22    col: Perm<usize>,
23}
24
25impl<T: FaerScalar> PackedLu<T> {
26    fn new(a: faer::MatRef<'_, T>, par: Par) -> Self {
27        let (m, n) = a.shape();
28        let mut lu = a.to_owned();
29        let mut row_fwd = vec![0usize; m];
30        let mut row_bwd = vec![0usize; m];
31        let mut col_fwd = vec![0usize; n];
32        let mut col_bwd = vec![0usize; n];
33        factor::lu_in_place(
34            lu.as_mut(),
35            &mut row_fwd,
36            &mut row_bwd,
37            &mut col_fwd,
38            &mut col_bwd,
39            par,
40            MemStack::new(&mut MemBuffer::new(
41                factor::lu_in_place_scratch::<usize, T>(m, n, par, Default::default()),
42            )),
43            Default::default(),
44        );
45        Self {
46            lu,
47            row: Perm::new_checked(row_fwd.into_boxed_slice(), row_bwd.into_boxed_slice(), m),
48            col: Perm::new_checked(col_fwd.into_boxed_slice(), col_bwd.into_boxed_slice(), n),
49        }
50    }
51
52    fn solve_in_place(&self, rhs: faer::MatMut<'_, T>, par: Par) {
53        let n = self.lu.nrows();
54        let k = rhs.ncols();
55        solve::solve_in_place_with_conj(
56            self.lu.rb(),
57            self.lu.rb(),
58            self.row.rb(),
59            self.col.rb(),
60            Conj::No,
61            rhs,
62            par,
63            MemStack::new(&mut MemBuffer::new(
64                solve::solve_in_place_scratch::<usize, T>(n, k, par),
65            )),
66        );
67    }
68}
69
70/// A [LinearSolver] that uses the LU decomposition in the [`faer`](https://github.com/sarah-ek/faer-rs) library to solve the linear system.
71///
72/// The low-level `faer::linalg::lu` API is used rather than `faer::linalg::solvers` so that the
73/// parallelism from [`FaerContext::par`] is honoured; the wrapper types hardcode faer's
74/// process-global setting instead.
75pub struct LU<T>
76where
77    T: FaerScalar,
78{
79    lu: Vec<PackedLu<T>>,
80    matrix: Option<FaerMat<T>>,
81}
82
83impl<T> Default for LU<T>
84where
85    T: FaerScalar,
86{
87    fn default() -> Self {
88        Self {
89            lu: Vec::new(),
90            matrix: None,
91        }
92    }
93}
94
95impl<T: FaerScalar> LinearSolver<FaerMat<T>> for LU<T> {
96    fn set_linearisation<C: LinearOp<T = T, V = FaerVec<T>, M = FaerMat<T>, C = FaerContext>>(
97        &mut self,
98        op: &C,
99    ) {
100        let matrix = self.matrix.as_mut().expect("Matrix not set");
101        op.matrix_inplace(matrix);
102        let nc = matrix.ncols();
103        let par = matrix.context.par;
104        self.lu = (0..matrix.context.nbatch())
105            .map(|b| PackedLu::new(matrix.data.rb().subcols(b * nc, nc), par))
106            .collect();
107    }
108
109    fn solve_in_place(&self, x: &mut FaerVec<T>) -> Result<(), LaError> {
110        if self.lu.is_empty() {
111            return Err(linear_solver_error!(LuNotInitialized));
112        }
113        x.context
114            .assert_broadcastable_into(self.lu.len(), "lu_solve");
115        let par = x.context.par;
116        let nlu = self.lu.len();
117        let nb = x.data.ncols();
118        for batch in 0..nb {
119            self.lu[broadcast_batch(batch, nlu, nb)]
120                .solve_in_place(x.data.rb_mut().col_mut(batch).as_mat_mut(), par);
121        }
122        Ok(())
123    }
124
125    fn set_sparsity<C: LinearOp<T = T, V = FaerVec<T>, M = FaerMat<T>, C = FaerContext>>(
126        &mut self,
127        op: &C,
128    ) {
129        let ncols = op.ncols();
130        let nrows = op.nrows();
131        let matrix = C::M::new_from_sparsity(nrows, ncols, op.sparsity(), *op.context());
132        self.matrix = Some(matrix);
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139    use crate::{
140        linear_solver::tests::{diagonal_op, test_grouped_lu_solve, test_narrow_state_lu_solve},
141        Vector,
142    };
143
144    #[test]
145    fn test_lu() {
146        let mut s = LU::<f64>::default();
147        let op = diagonal_op::<FaerMat<f64>>(2.0);
148        s.set_sparsity(&op);
149        s.set_linearisation(&op);
150        let b = FaerVec::from_vec(vec![2.0, 4.0], Default::default());
151        let x = s.solve(&b).unwrap();
152        x.assert_eq_st(
153            &FaerVec::from_vec(vec![1.0, 2.0], Default::default()),
154            1e-10,
155        );
156    }
157
158    #[test]
159    fn test_grouped_lu() {
160        test_grouped_lu_solve::<FaerMat<f64>, LU<f64>>(FaerContext::with_nbatch(2));
161    }
162
163    #[test]
164    #[should_panic(expected = "incompatible nbatch")]
165    fn test_narrow_state_lu() {
166        test_narrow_state_lu_solve::<FaerMat<f64>, LU<f64>>(FaerContext::with_nbatch(2));
167    }
168}
diffsol_la/linear_solver/faer/lu.rs

diffsol_la/linear_solver/faer/
lu.rs
