diffsol_la/linear_solver/
mod.rs

1use crate::{error::LaError, LinearOp, Matrix};
2
3#[cfg(feature = "nalgebra")]
4pub mod nalgebra;
5
6#[cfg(feature = "faer")]
7pub mod faer;
8
9#[cfg(feature = "suitesparse")]
10pub mod suitesparse;
11
12#[cfg(feature = "cuda")]
13#[allow(deprecated)]
14pub mod cuda;
15
16#[cfg(feature = "cuda-oxide")]
17pub mod cuda_oxide;
18
19pub use faer::lu::LU as FaerLU;
20pub use nalgebra::lu::ReusableLU as NalgebraLU;
21
22/// A solver for the linear problem `Ax = b`, where `A` is a [LinearOp].
23pub trait LinearSolver<M: Matrix>: Default {
24    /// Set the point at which the linear operator `A` is evaluated and factorise it.
25    ///
26    /// The operator is assumed to have the same sparsity as that given to
27    /// [Self::set_sparsity].
28    fn set_linearisation<C: LinearOp<V = M::V, T = M::T, M = M, C = M::C>>(&mut self, op: &C);
29
30    /// Set the sparsity of the problem to be solved, any previous problem is discarded.
31    ///
32    /// Any internal state of the solver is reset. This function will normally set
33    /// the sparsity pattern of the matrix to be solved.
34    fn set_sparsity<C: LinearOp<V = M::V, T = M::T, M = M, C = M::C>>(&mut self, op: &C);
35
36    /// Solve the problem `Ax = b` and return the solution `x`.
37    ///
38    /// Panics if [Self::set_linearisation] has not been called previously.
39    fn solve(&self, b: &M::V) -> Result<M::V, LaError> {
40        let mut b = b.clone();
41        self.solve_in_place(&mut b)?;
42        Ok(b)
43    }
44
45    fn solve_in_place(&self, b: &mut M::V) -> Result<(), LaError>;
46
47    /// Solve the transposed problem `A^T x = b` in place, reusing the factorisation of `A`.
48    ///
49    /// Returns an error if [Self::set_linearisation] has not been called previously.
50    fn solve_transpose_in_place(&self, b: &mut M::V) -> Result<(), LaError>;
51}
52
53/// The shared [LinearSolver] tests for solver `$LS` on matrix `$M`, with `$ctx2` a context of
54/// two batches.
55#[cfg(test)]
56macro_rules! generate_lu_tests {
57    ($suffix:ident, $M:ty, $LS:ty, $ctx2:expr) => {
58        paste::paste! {
59            #[test]
60            fn [<test_lu_ $suffix>]() {
61                $crate::linear_solver::tests::test_diagonal_lu_solve::<$M, $LS>();
62            }
63            #[test]
64            fn [<test_grouped_lu_ $suffix>]() {
65                $crate::linear_solver::tests::test_grouped_lu_solve::<$M, $LS>($ctx2);
66            }
67            #[test]
68            fn [<test_transpose_lu_ $suffix>]() {
69                $crate::linear_solver::tests::test_transpose_lu_solve::<$M, $LS>($ctx2);
70            }
71            #[test]
72            #[should_panic(expected = "incompatible nbatch")]
73            fn [<test_narrow_state_lu_ $suffix>]() {
74                $crate::linear_solver::tests::test_narrow_state_lu_solve::<$M, $LS>($ctx2);
75            }
76        }
77    };
78}
79#[cfg(test)]
80pub(crate) use generate_lu_tests;
81
82#[cfg(test)]
83pub(crate) mod tests {
84    use super::LinearSolver;
85    use crate::{Context, IndexType, LinearOp, Matrix, Vector};
86
87    /// A simple diagonal [LinearOp] used for testing the linear solvers.
88    pub struct DiagonalOp<M: Matrix> {
89        pub matrix: M,
90    }
91
92    /// Create a 2x2 diagonal operator `A = diag(value, value)`.
93    pub fn diagonal_op<M: Matrix>(value: f64) -> DiagonalOp<M> {
94        use num_traits::FromPrimitive;
95        let v = M::T::from_f64(value).unwrap();
96        let diag = M::V::from_vec(vec![v, v], Default::default());
97        DiagonalOp {
98            matrix: M::from_diagonal(&diag),
99        }
100    }
101
102    /// Create an `nstates`-square diagonal operator `A = diag(value, ..., value)`.
103    #[cfg_attr(not(feature = "cuda-oxide"), allow(dead_code))]
104    pub fn diagonal_op_n<M: Matrix>(nstates: IndexType, value: f64, ctx: M::C) -> DiagonalOp<M> {
105        use num_traits::FromPrimitive;
106        let v = M::T::from_f64(value).unwrap();
107        let diag = M::V::from_element(nstates, v, ctx);
108        DiagonalOp {
109            matrix: M::from_diagonal(&diag),
110        }
111    }
112
113    /// Create a batched 2x2 diagonal operator with `A_b = diag(values[b], values[b])`.
114    pub fn batched_diagonal_op<M: Matrix>(values: &[f64], ctx: M::C) -> DiagonalOp<M> {
115        use num_traits::FromPrimitive;
116        assert_eq!(values.len(), ctx.nbatch());
117        let diag = M::V::from_vec(
118            values
119                .iter()
120                .flat_map(|&v| [M::T::from_f64(v).unwrap(); 2])
121                .collect(),
122            ctx,
123        );
124        DiagonalOp {
125            matrix: M::from_diagonal(&diag),
126        }
127    }
128
129    /// Solve `diag(2, 2) x = (2, 4)`.
130    pub fn test_diagonal_lu_solve<M: Matrix, LS: LinearSolver<M>>() {
131        use num_traits::FromPrimitive;
132        let f = |x: f64| M::T::from_f64(x).unwrap();
133        let op = diagonal_op::<M>(2.0);
134        let mut s = LS::default();
135        s.set_sparsity(&op);
136        s.set_linearisation(&op);
137        let b = M::V::from_vec(vec![f(2.0), f(4.0)], Default::default());
138        let x = s.solve(&b).unwrap();
139        x.assert_eq_eps(
140            &M::V::from_vec(vec![f(1.0), f(2.0)], Default::default()),
141            100.0,
142        );
143    }
144
145    /// One factorization per group: `A` carries 2 batches and the right-hand side 4, so
146    /// batches 0 and 1 solve against `A_0` and batches 2 and 3 against `A_1`.  A cyclic
147    /// mapping would pair batch 1 with `A_1` and fail.
148    pub fn test_grouped_lu_solve<M: Matrix, LS: LinearSolver<M>>(ctx2: M::C) {
149        use num_traits::FromPrimitive;
150        let f = |x: f64| M::T::from_f64(x).unwrap();
151        let wide = ctx2.clone_with_nbatch(4).unwrap();
152        let op = batched_diagonal_op::<M>(&[2.0, 5.0], ctx2);
153        let mut s = LS::default();
154        s.set_sparsity(&op);
155        s.set_linearisation(&op);
156        let b = M::V::from_vec((1..=8).map(|i| f(i as f64)).collect(), wide.clone());
157        let x = s.solve(&b).unwrap();
158        let expected = M::V::from_vec(
159            vec![
160                f(0.5),
161                f(1.0),
162                f(1.5),
163                f(2.0),
164                f(1.0),
165                f(1.2),
166                f(1.4),
167                f(1.6),
168            ],
169            wide,
170        );
171        x.assert_eq_eps(&expected, 100.0);
172    }
173
174    /// Solve `A x = b` and `A^T x = b` with `A_0 = [[0, 1, 2], [3, 1, 0], [1, 0, 4]]` and
175    /// `A_1 = A_0^T`, both needing a pivot, against 4 right-hand side batches grouped as in
176    /// [test_grouped_lu_solve]. Every solution is `(1, 2, 3)`.
177    pub fn test_transpose_lu_solve<M: Matrix, LS: LinearSolver<M>>(ctx2: M::C) {
178        use num_traits::FromPrimitive;
179        let f = |x: f64| M::T::from_f64(x).unwrap();
180        let wide = ctx2.clone_with_nbatch(4).unwrap();
181        let indices = vec![(0, 1), (0, 2), (1, 0), (1, 1), (2, 0), (2, 2)];
182        let values = [1.0, 2.0, 3.0, 1.0, 1.0, 4.0, 3.0, 1.0, 1.0, 1.0, 2.0, 4.0];
183        let matrix =
184            M::try_from_triplets(3, 3, indices, values.iter().map(|&v| f(v)).collect(), ctx2)
185                .unwrap();
186        let op = DiagonalOp { matrix };
187        let mut s = LS::default();
188        s.set_sparsity(&op);
189        s.set_linearisation(&op);
190        let rhs = |a: [f64; 3], b: [f64; 3]| {
191            M::V::from_vec(
192                [a, a, b, b].iter().flatten().map(|&v| f(v)).collect(),
193                wide.clone(),
194            )
195        };
196        let expected = rhs([1.0, 2.0, 3.0], [1.0, 2.0, 3.0]);
197        let (a_x, at_x) = ([8.0, 5.0, 13.0], [9.0, 3.0, 14.0]);
198
199        let mut b = rhs(a_x, at_x);
200        s.solve_in_place(&mut b).unwrap();
201        b.assert_eq_eps(&expected, 100.0);
202
203        let mut b = rhs(at_x, a_x);
204        s.solve_transpose_in_place(&mut b).unwrap();
205        b.assert_eq_eps(&expected, 100.0);
206    }
207
208    /// More factorizations than right-hand side batches would leave factorizations unused, so
209    /// the solve panics rather than silently dropping them.
210    pub fn test_narrow_state_lu_solve<M: Matrix, LS: LinearSolver<M>>(ctx2: M::C) {
211        let wide = ctx2.clone_with_nbatch(4).unwrap();
212        let op = batched_diagonal_op::<M>(&[2.0, 3.0, 4.0, 5.0], wide);
213        let mut s = LS::default();
214        s.set_sparsity(&op);
215        s.set_linearisation(&op);
216        let mut b = M::V::zeros(2, ctx2);
217        s.solve_in_place(&mut b).unwrap();
218    }
219
220    impl<M: Matrix> LinearOp for DiagonalOp<M> {
221        type T = M::T;
222        type V = M::V;
223        type M = M;
224        type C = M::C;
225
226        fn nrows(&self) -> IndexType {
227            self.matrix.nrows()
228        }
229        fn ncols(&self) -> IndexType {
230            self.matrix.ncols()
231        }
232        fn context(&self) -> &Self::C {
233            self.matrix.context()
234        }
235        fn matrix_inplace(&self, y: &mut Self::M) {
236            y.copy_from(&self.matrix);
237        }
238        fn sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
239            self.matrix.sparsity().map(|s| {
240                use crate::matrix::sparsity::MatrixSparsityRef;
241                s.to_owned()
242            })
243        }
244    }
245}
diffsol_la/linear_solver/mod.rs

diffsol_la/linear_solver/
mod.rs
