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, transpose: bool) {
53        let n = self.lu.nrows();
54        let k = rhs.ncols();
55        if transpose {
56            solve::solve_transpose_in_place_with_conj(
57                self.lu.rb(),
58                self.lu.rb(),
59                self.row.rb(),
60                self.col.rb(),
61                Conj::No,
62                rhs,
63                par,
64                MemStack::new(&mut MemBuffer::new(
65                    solve::solve_transpose_in_place_scratch::<usize, T>(n, k, par),
66                )),
67            );
68        } else {
69            solve::solve_in_place_with_conj(
70                self.lu.rb(),
71                self.lu.rb(),
72                self.row.rb(),
73                self.col.rb(),
74                Conj::No,
75                rhs,
76                par,
77                MemStack::new(&mut MemBuffer::new(
78                    solve::solve_in_place_scratch::<usize, T>(n, k, par),
79                )),
80            );
81        }
82    }
83}
84
85/// A [LinearSolver] that uses the LU decomposition in the [`faer`](https://github.com/sarah-ek/faer-rs) library to solve the linear system.
86///
87/// The low-level `faer::linalg::lu` API is used rather than `faer::linalg::solvers` so that the
88/// parallelism from [`FaerContext::par`] is honoured; the wrapper types hardcode faer's
89/// process-global setting instead.
90pub struct LU<T>
91where
92    T: FaerScalar,
93{
94    lu: Vec<PackedLu<T>>,
95    matrix: Option<FaerMat<T>>,
96}
97
98impl<T> Default for LU<T>
99where
100    T: FaerScalar,
101{
102    fn default() -> Self {
103        Self {
104            lu: Vec::new(),
105            matrix: None,
106        }
107    }
108}
109
110impl<T: FaerScalar> LinearSolver<FaerMat<T>> for LU<T> {
111    fn set_linearisation<C: LinearOp<T = T, V = FaerVec<T>, M = FaerMat<T>, C = FaerContext>>(
112        &mut self,
113        op: &C,
114    ) {
115        let matrix = self.matrix.as_mut().expect("Matrix not set");
116        op.matrix_inplace(matrix);
117        let nc = matrix.ncols();
118        let par = matrix.context.par;
119        self.lu = (0..matrix.context.nbatch())
120            .map(|b| PackedLu::new(matrix.data.rb().subcols(b * nc, nc), par))
121            .collect();
122    }
123
124    fn solve_in_place(&self, x: &mut FaerVec<T>) -> Result<(), LaError> {
125        self.solve_in_place_impl(x, false)
126    }
127
128    fn solve_transpose_in_place(&self, x: &mut FaerVec<T>) -> Result<(), LaError> {
129        self.solve_in_place_impl(x, true)
130    }
131
132    fn set_sparsity<C: LinearOp<T = T, V = FaerVec<T>, M = FaerMat<T>, C = FaerContext>>(
133        &mut self,
134        op: &C,
135    ) {
136        let ncols = op.ncols();
137        let nrows = op.nrows();
138        let matrix = C::M::new_from_sparsity(nrows, ncols, op.sparsity(), *op.context());
139        self.matrix = Some(matrix);
140    }
141}
142
143impl<T: FaerScalar> LU<T> {
144    fn solve_in_place_impl(&self, x: &mut FaerVec<T>, transpose: bool) -> Result<(), LaError> {
145        if self.lu.is_empty() {
146            return Err(linear_solver_error!(LuNotInitialized));
147        }
148        x.context
149            .assert_broadcastable_into(self.lu.len(), "lu_solve");
150        let par = x.context.par;
151        let nlu = self.lu.len();
152        let nb = x.data.ncols();
153        for batch in 0..nb {
154            self.lu[broadcast_batch(batch, nlu, nb)].solve_in_place(
155                x.data.rb_mut().col_mut(batch).as_mat_mut(),
156                par,
157                transpose,
158            );
159        }
160        Ok(())
161    }
162}
163
164#[cfg(test)]
165mod tests {
166    use super::*;
167
168    crate::linear_solver::generate_lu_tests!(
169        f64,
170        FaerMat<f64>,
171        LU<f64>,
172        FaerContext::with_nbatch(2)
173    );
174    crate::linear_solver::generate_lu_tests!(
175        f32,
176        FaerMat<f32>,
177        LU<f32>,
178        FaerContext::with_nbatch(2)
179    );
180}
diffsol_la/linear_solver/faer/lu.rs

diffsol_la/linear_solver/faer/
lu.rs
