diffsol_la/linear_solver/nalgebra/
lu.rs

1use nalgebra::Dyn;
2
3use crate::context::broadcast_batch;
4use crate::{
5    error::LaError, linear_solver_error, matrix::dense_nalgebra_serial::NalgebraMat, Context,
6    LinearOp, LinearSolver, Matrix, NalgebraContext, NalgebraScalar, NalgebraVec,
7};
8
9/// A [LinearSolver] that uses the LU decomposition in the [`nalgebra` library](https://nalgebra.org/) to solve the linear system.
10///
11/// Exported as [NalgebraNativeLU](crate::NalgebraNativeLU): re-allocates its factorization
12/// (a matrix clone plus a fresh permutation sequence) on every
13/// [`set_linearisation`](LinearSolver::set_linearisation) call, unlike the default
14/// [NalgebraLU](crate::NalgebraLU), which reuses its workspace in place. Kept around for
15/// comparison/benchmarking.
16#[derive(Clone)]
17pub struct LU<T>
18where
19    T: NalgebraScalar,
20{
21    matrix: Option<NalgebraMat<T>>,
22    lu: Vec<nalgebra::LU<T, Dyn, Dyn>>,
23}
24
25impl<T> Default for LU<T>
26where
27    T: NalgebraScalar,
28{
29    fn default() -> Self {
30        Self {
31            lu: Vec::new(),
32            matrix: None,
33        }
34    }
35}
36
37impl<T: NalgebraScalar> LinearSolver<NalgebraMat<T>> for LU<T> {
38    fn solve_in_place(&self, state: &mut NalgebraVec<T>) -> Result<(), LaError> {
39        if self.lu.is_empty() {
40            return Err(linear_solver_error!(LuNotInitialized));
41        }
42        state
43            .context
44            .assert_broadcastable_into(self.lu.len(), "lu_solve");
45        if state.context.nbatch() == 1 {
46            if self.lu[0].solve_mut(&mut state.data) {
47                return Ok(());
48            }
49            return Err(linear_solver_error!(LuSolveFailed));
50        }
51        let nb = state.context.nbatch();
52        for batch in 0..nb {
53            let mut state_batch = state.data.column_mut(batch);
54            if !self.lu[broadcast_batch(batch, self.lu.len(), nb)].solve_mut(&mut state_batch) {
55                return Err(linear_solver_error!(LuSolveFailed));
56            }
57        }
58        Ok(())
59    }
60
61    fn set_linearisation<
62        C: LinearOp<T = T, V = NalgebraVec<T>, M = NalgebraMat<T>, C = NalgebraContext>,
63    >(
64        &mut self,
65        op: &C,
66    ) {
67        let matrix = self.matrix.as_mut().expect("Matrix not set");
68        op.matrix_inplace(matrix);
69        if matrix.context.nbatch() == 1 {
70            self.lu = vec![matrix.data.clone().lu()];
71            return;
72        }
73        let ncols = matrix.data.ncols() / matrix.context.nbatch();
74        self.lu = (0..matrix.context.nbatch())
75            .map(|batch| matrix.data.columns(batch * ncols, ncols).into_owned().lu())
76            .collect();
77    }
78
79    fn set_sparsity<
80        C: LinearOp<T = T, V = NalgebraVec<T>, M = NalgebraMat<T>, C = NalgebraContext>,
81    >(
82        &mut self,
83        op: &C,
84    ) {
85        let ncols = op.ncols();
86        let nrows = op.nrows();
87        let matrix = C::M::new_from_sparsity(nrows, ncols, op.sparsity(), *op.context());
88        self.matrix = Some(matrix);
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95    use crate::{
96        linear_solver::tests::{diagonal_op, test_grouped_lu_solve, test_narrow_state_lu_solve},
97        Vector,
98    };
99
100    #[test]
101    fn test_lu() {
102        let mut s = LU::<f64>::default();
103        let op = diagonal_op::<NalgebraMat<f64>>(2.0);
104        s.set_sparsity(&op);
105        s.set_linearisation(&op);
106        let b = NalgebraVec::from_vec(vec![2.0, 4.0], Default::default());
107        let x = s.solve(&b).unwrap();
108        x.assert_eq_st(
109            &NalgebraVec::from_vec(vec![1.0, 2.0], Default::default()),
110            1e-10,
111        );
112    }
113    #[test]
114    fn test_grouped_lu() {
115        test_grouped_lu_solve::<NalgebraMat<f64>, LU<f64>>(NalgebraContext::with_nbatch(2));
116    }
117
118    #[test]
119    #[should_panic(expected = "incompatible nbatch")]
120    fn test_narrow_state_lu() {
121        test_narrow_state_lu_solve::<NalgebraMat<f64>, LU<f64>>(NalgebraContext::with_nbatch(2));
122    }
123}
diffsol_la/linear_solver/nalgebra/lu.rs

diffsol_la/linear_solver/nalgebra/
lu.rs
