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")]
13pub mod cuda;
14
15#[cfg(feature = "cuda-oxide")]
16pub mod cuda_oxide;
17
18pub use faer::lu::LU as FaerLU;
19pub use nalgebra::lu::LU as NalgebraNativeLU;
20pub use nalgebra::reusable_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
48#[cfg(test)]
49pub(crate) mod tests {
50    use super::LinearSolver;
51    use crate::{Context, IndexType, LinearOp, Matrix, Vector};
52
53    /// A simple diagonal [LinearOp] used for testing the linear solvers.
54    pub struct DiagonalOp<M: Matrix> {
55        matrix: M,
56    }
57
58    /// Create a 2x2 diagonal operator `A = diag(value, value)`.
59    pub fn diagonal_op<M: Matrix>(value: f64) -> DiagonalOp<M> {
60        use num_traits::FromPrimitive;
61        let v = M::T::from_f64(value).unwrap();
62        let diag = M::V::from_vec(vec![v, v], Default::default());
63        DiagonalOp {
64            matrix: M::from_diagonal(&diag),
65        }
66    }
67
68    /// Create an `nstates`-square diagonal operator `A = diag(value, ..., value)`.
69    #[cfg_attr(not(feature = "cuda-oxide"), allow(dead_code))]
70    pub fn diagonal_op_n<M: Matrix>(nstates: IndexType, value: f64, ctx: M::C) -> DiagonalOp<M> {
71        use num_traits::FromPrimitive;
72        let v = M::T::from_f64(value).unwrap();
73        let diag = M::V::from_element(nstates, v, ctx);
74        DiagonalOp {
75            matrix: M::from_diagonal(&diag),
76        }
77    }
78
79    /// Create a batched 2x2 diagonal operator with `A_b = diag(values[b], values[b])`.
80    pub fn batched_diagonal_op<M: Matrix>(values: &[f64], ctx: M::C) -> DiagonalOp<M> {
81        use num_traits::FromPrimitive;
82        assert_eq!(values.len(), ctx.nbatch());
83        let diag = M::V::from_vec(
84            values
85                .iter()
86                .flat_map(|&v| [M::T::from_f64(v).unwrap(); 2])
87                .collect(),
88            ctx,
89        );
90        DiagonalOp {
91            matrix: M::from_diagonal(&diag),
92        }
93    }
94
95    /// One factorization per group: `A` carries 2 batches and the right-hand side 4, so
96    /// batches 0 and 1 solve against `A_0` and batches 2 and 3 against `A_1`.  A cyclic
97    /// mapping would pair batch 1 with `A_1` and fail.
98    pub fn test_grouped_lu_solve<M: Matrix, LS: LinearSolver<M>>(ctx2: M::C) {
99        use num_traits::FromPrimitive;
100        let f = |x: f64| M::T::from_f64(x).unwrap();
101        let wide = ctx2.clone_with_nbatch(4).unwrap();
102        let op = batched_diagonal_op::<M>(&[2.0, 5.0], ctx2);
103        let mut s = LS::default();
104        s.set_sparsity(&op);
105        s.set_linearisation(&op);
106        let b = M::V::from_vec((1..=8).map(|i| f(i as f64)).collect(), wide.clone());
107        let x = s.solve(&b).unwrap();
108        let expected = M::V::from_vec(
109            vec![
110                f(0.5),
111                f(1.0),
112                f(1.5),
113                f(2.0),
114                f(1.0),
115                f(1.2),
116                f(1.4),
117                f(1.6),
118            ],
119            wide,
120        );
121        x.assert_eq_st(&expected, f(1e-10));
122    }
123
124    /// More factorizations than right-hand side batches would leave factorizations unused, so
125    /// the solve panics rather than silently dropping them.
126    pub fn test_narrow_state_lu_solve<M: Matrix, LS: LinearSolver<M>>(ctx2: M::C) {
127        let wide = ctx2.clone_with_nbatch(4).unwrap();
128        let op = batched_diagonal_op::<M>(&[2.0, 3.0, 4.0, 5.0], wide);
129        let mut s = LS::default();
130        s.set_sparsity(&op);
131        s.set_linearisation(&op);
132        let mut b = M::V::zeros(2, ctx2);
133        s.solve_in_place(&mut b).unwrap();
134    }
135
136    impl<M: Matrix> LinearOp for DiagonalOp<M> {
137        type T = M::T;
138        type V = M::V;
139        type M = M;
140        type C = M::C;
141
142        fn nrows(&self) -> IndexType {
143            self.matrix.nrows()
144        }
145        fn ncols(&self) -> IndexType {
146            self.matrix.ncols()
147        }
148        fn context(&self) -> &Self::C {
149            self.matrix.context()
150        }
151        fn matrix_inplace(&self, y: &mut Self::M) {
152            y.copy_from(&self.matrix);
153        }
154        fn sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
155            self.matrix.sparsity().map(|s| {
156                use crate::matrix::sparsity::MatrixSparsityRef;
157                s.to_owned()
158            })
159        }
160    }
161}
diffsol_la/linear_solver/mod.rs

diffsol_la/linear_solver/
mod.rs
