faer_precond/spai/
mod.rs

1//! Sparse approximate inverse (SPAI) preconditioner.
2//!
3//! SPAI builds an explicit sparse `M ~= A^{-1}` by minimising `||A M - I||_F`.
4//! That objective separates by column, so `M` is assembled from independent
5//! small least-squares problems — embarrassingly parallel to build. Applying the
6//! preconditioner is then a single sparse matrix-vector product, with no
7//! triangular solves; the transpose and adjoint are matvecs against the
8//! transposed/adjoint views of `M`.
9//!
10//! Unlike [`crate::fsai::Fsai`], SPAI does **not** assume symmetry — it is the
11//! nonsymmetric approximate-inverse, the matvec-only counterpart to
12//! [`crate::Ilu0`] / [`crate::Ilutp`].
13//!
14//! # When to use it
15//!
16//! Reach for SPAI on nonsymmetric systems when matvecs are cheap and you want an
17//! approximate inverse that applies as a matvec — parallel hardware, or settings
18//! where the sequential triangular solve of an ILU does not scale. The build is
19//! heavier than an ILU (a dense least-squares solve per column), so it pays off
20//! when the same preconditioner is applied many times.
21//!
22//! # Pattern
23//!
24//! The accuracy/cost knob is the prescribed column pattern of `M`:
25//!
26//! - [`SpaiPattern::ColumnsOfA`] — each column of `M` takes the pattern of the
27//!   matching column of `A` (cheapest).
28//! - [`SpaiPattern::ColumnsOfPower`] — the pattern of `A^power`'s columns,
29//!   denser and more accurate for larger `power`.
30//!
31//! Adaptive (dynamic) pattern growth — the Grote–Huckle algorithm — is left as
32//! future work.
33//!
34//! # Storage
35//!
36//! `M` is stored CSC. Apply reads it through a
37//! [`faer::sparse::SparseColMatRef`] (and its transpose/adjoint
38//! views); the in-place path uses one work column block from the caller's
39//! [`MemStack`], so no heap allocation occurs during apply.
40
41use core::fmt::Debug;
42
43use dyn_stack::{MemStack, StackReq};
44use faer::matrix_free::{BiLinOp, BiPrecond, LinOp, Precond};
45use faer::sparse::{SparseColMatRef, SymbolicSparseColMatRef};
46use faer::{MatMut, MatRef, Par};
47use faer_traits::{ComplexField, Index};
48
49mod apply;
50mod build;
51
52/// Prescribed column pattern for the SPAI factor `M`.
53#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
54pub enum SpaiPattern {
55    /// Each column of `M` takes the pattern of the matching column of `A`.
56    #[default]
57    ColumnsOfA,
58    /// Columns take the pattern of `A^power` (`power >= 1`; `1` equals
59    /// [`SpaiPattern::ColumnsOfA`]).
60    ColumnsOfPower { power: usize },
61}
62
63/// Error returned by SPAI construction.
64#[derive(Debug, Clone, PartialEq, Eq)]
65pub enum SpaiError {
66    /// The source matrix was not square.
67    NonSquareMatrix { nrows: usize, ncols: usize },
68    /// A `ColumnsOfPower` power of zero was requested.
69    InvalidPower,
70}
71
72impl core::fmt::Display for SpaiError {
73    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
74        match self {
75            Self::NonSquareMatrix { nrows, ncols } => {
76                write!(f, "matrix must be square but is {nrows}x{ncols}")
77            }
78            Self::InvalidPower => f.write_str("SPAI pattern power must be at least 1"),
79        }
80    }
81}
82
83impl core::error::Error for SpaiError {}
84
85/// Sparse approximate inverse: `M^{-1} = M ~= A^{-1}`.
86///
87/// Stores `M` in CSC. Apply is a single sparse matvec with no triangular solves.
88/// See the [module documentation](self) for guidance.
89#[derive(Debug, Clone)]
90pub struct Spai<I, T> {
91    pub(crate) dim: usize,
92    pub(crate) m_col_ptr: Vec<I>,
93    pub(crate) m_row_idx: Vec<I>,
94    pub(crate) m_values: Vec<T>,
95}
96
97impl<I, T> Spai<I, T> {
98    /// Dimension `n` of the preconditioner.
99    #[inline]
100    pub fn dim(&self) -> usize {
101        self.dim
102    }
103}
104
105impl<I: Index, T: ComplexField> Spai<I, T> {
106    /// View over the approximate inverse `M`.
107    #[inline]
108    pub(crate) fn m_view(&self) -> SparseColMatRef<'_, I, T> {
109        let symbolic = unsafe {
110            SymbolicSparseColMatRef::<'_, I>::new_unchecked(
111                self.dim,
112                self.dim,
113                &self.m_col_ptr,
114                None,
115                &self.m_row_idx,
116            )
117        };
118        SparseColMatRef::new(symbolic, &self.m_values)
119    }
120}
121
122impl<I, T> LinOp<T> for Spai<I, T>
123where
124    I: Index,
125    T: ComplexField + Debug + Sync,
126{
127    fn apply_scratch(&self, _rhs_ncols: usize, _par: Par) -> StackReq {
128        StackReq::EMPTY
129    }
130
131    fn nrows(&self) -> usize {
132        self.dim
133    }
134
135    fn ncols(&self) -> usize {
136        self.dim
137    }
138
139    fn apply(&self, out: MatMut<'_, T>, rhs: MatRef<'_, T>, par: Par, _stack: &mut MemStack) {
140        apply::apply_out(self, false, false, out, rhs, par);
141    }
142
143    fn conj_apply(&self, out: MatMut<'_, T>, rhs: MatRef<'_, T>, par: Par, _stack: &mut MemStack) {
144        apply::apply_out(self, false, true, out, rhs, par);
145    }
146}
147
148impl<I, T> Precond<T> for Spai<I, T>
149where
150    I: Index,
151    T: ComplexField + Debug + Sync,
152{
153    fn apply_in_place_scratch(&self, rhs_ncols: usize, _par: Par) -> StackReq {
154        apply::inplace_scratch(self, rhs_ncols)
155    }
156
157    fn apply_in_place(&self, rhs: MatMut<'_, T>, par: Par, stack: &mut MemStack) {
158        apply::apply_inplace(self, false, false, rhs, par, stack);
159    }
160
161    fn conj_apply_in_place(&self, rhs: MatMut<'_, T>, par: Par, stack: &mut MemStack) {
162        apply::apply_inplace(self, false, true, rhs, par, stack);
163    }
164}
165
166impl<I, T> BiLinOp<T> for Spai<I, T>
167where
168    I: Index,
169    T: ComplexField + Debug + Sync,
170{
171    fn transpose_apply_scratch(&self, _rhs_ncols: usize, _par: Par) -> StackReq {
172        StackReq::EMPTY
173    }
174
175    fn transpose_apply(
176        &self,
177        out: MatMut<'_, T>,
178        rhs: MatRef<'_, T>,
179        par: Par,
180        _stack: &mut MemStack,
181    ) {
182        apply::apply_out(self, true, false, out, rhs, par);
183    }
184
185    fn adjoint_apply(
186        &self,
187        out: MatMut<'_, T>,
188        rhs: MatRef<'_, T>,
189        par: Par,
190        _stack: &mut MemStack,
191    ) {
192        apply::apply_out(self, true, true, out, rhs, par);
193    }
194}
195
196impl<I, T> BiPrecond<T> for Spai<I, T>
197where
198    I: Index,
199    T: ComplexField + Debug + Sync,
200{
201    fn transpose_apply_in_place_scratch(&self, rhs_ncols: usize, _par: Par) -> StackReq {
202        apply::inplace_scratch(self, rhs_ncols)
203    }
204
205    fn transpose_apply_in_place(&self, rhs: MatMut<'_, T>, par: Par, stack: &mut MemStack) {
206        apply::apply_inplace(self, true, false, rhs, par, stack);
207    }
208
209    fn adjoint_apply_in_place(&self, rhs: MatMut<'_, T>, par: Par, stack: &mut MemStack) {
210        apply::apply_inplace(self, true, true, rhs, par, stack);
211    }
212}
213
214#[cfg(test)]
215mod tests {
216    use super::*;
217    use core::mem::MaybeUninit;
218    use faer::sparse::{SparseColMat, Triplet};
219    use faer::{Mat, MatRef, mat};
220
221    fn with_stack(req: StackReq, f: impl FnOnce(&mut MemStack)) {
222        let nbytes = req.unaligned_bytes_required().max(1);
223        let mut buf = vec![MaybeUninit::<u8>::uninit(); nbytes].into_boxed_slice();
224        f(MemStack::new(&mut buf));
225    }
226
227    fn assert_close(lhs: MatRef<'_, f64>, rhs: MatRef<'_, f64>, tol: f64) {
228        assert_eq!(lhs.nrows(), rhs.nrows());
229        assert_eq!(lhs.ncols(), rhs.ncols());
230        for j in 0..lhs.ncols() {
231            for i in 0..lhs.nrows() {
232                let diff = (*lhs.get(i, j) - *rhs.get(i, j)).abs();
233                assert!(
234                    diff <= tol,
235                    "mismatch at ({i}, {j}): lhs={}, rhs={}, diff={diff}",
236                    *lhs.get(i, j),
237                    *rhs.get(i, j),
238                );
239            }
240        }
241    }
242
243    fn to_dense(a: &SparseColMat<usize, f64>) -> Mat<f64> {
244        let n = a.nrows();
245        let mut out = Mat::<f64>::zeros(n, a.ncols());
246        let a_ref = a.as_ref();
247        for j in 0..a.ncols() {
248            let rows = a_ref.symbolic().row_idx_of_col_raw(j);
249            let vals = a_ref.val_of_col(j);
250            for (r, v) in rows.iter().zip(vals.iter()) {
251                *out.as_mut().get_mut(*r, j) = *v;
252            }
253        }
254        out
255    }
256
257    fn tridiagonal(n: usize, diag: f64, sub: f64, sup: f64) -> SparseColMat<usize, f64> {
258        let mut triplets = Vec::new();
259        for i in 0..n {
260            triplets.push(Triplet::new(i, i, diag));
261            if i > 0 {
262                triplets.push(Triplet::new(i, i - 1, sub));
263                triplets.push(Triplet::new(i - 1, i, sup));
264            }
265        }
266        SparseColMat::try_new_from_triplets(n, n, &triplets).unwrap()
267    }
268
269    fn apply_inplace(pc: &Spai<usize, f64>, rhs: &mut Mat<f64>) {
270        with_stack(pc.apply_in_place_scratch(rhs.ncols(), Par::Seq), |stack| {
271            pc.apply_in_place(rhs.as_mut(), Par::Seq, stack);
272        });
273    }
274
275    fn residual_ratio(a: &SparseColMat<usize, f64>, pc: &Spai<usize, f64>, b: &Mat<f64>) -> f64 {
276        let a_dense = to_dense(a);
277        let mut x = b.clone();
278        apply_inplace(pc, &mut x);
279        let residual = &a_dense * &x - b;
280        let b_norm: f64 = b.as_ref().col(0).iter().map(|v| v * v).sum::<f64>().sqrt();
281        let r_norm: f64 = residual
282            .as_ref()
283            .col(0)
284            .iter()
285            .map(|v| v * v)
286            .sum::<f64>()
287            .sqrt();
288        r_norm / b_norm
289    }
290
291    #[test]
292    fn diagonal_is_exact_inverse() {
293        let mut triplets = Vec::new();
294        for (i, &v) in [2.0, 4.0, 8.0].iter().enumerate() {
295            triplets.push(Triplet::new(i, i, v));
296        }
297        let a = SparseColMat::<usize, f64>::try_new_from_triplets(3, 3, &triplets).unwrap();
298        let pc = Spai::try_new(a.as_ref(), SpaiPattern::ColumnsOfA).unwrap();
299        let mut x = mat![[2.0_f64], [8.0], [16.0]];
300        apply_inplace(&pc, &mut x);
301        let expected = mat![[1.0_f64], [2.0], [2.0]];
302        assert_close(x.as_ref(), expected.as_ref(), 1e-12);
303    }
304
305    #[test]
306    fn reduces_residual_on_nonsymmetric() {
307        let a = tridiagonal(12, 4.0, -2.0, -1.0);
308        let n = a.nrows();
309        let pc = Spai::try_new(a.as_ref(), SpaiPattern::ColumnsOfPower { power: 2 }).unwrap();
310        let b = Mat::<f64>::from_fn(n, 1, |i, _| (i % 7) as f64 - 3.0);
311        let ratio = residual_ratio(&a, &pc, &b);
312        assert!(ratio < 1.0, "SPAI should reduce the residual: {ratio}");
313    }
314
315    #[test]
316    fn transpose_differs_from_forward_on_nonsymmetric() {
317        let a = tridiagonal(8, 4.0, -2.0, -1.0);
318        let pc = Spai::try_new(a.as_ref(), SpaiPattern::ColumnsOfA).unwrap();
319        let rhs = Mat::<f64>::from_fn(8, 1, |i, _| (i % 5) as f64 - 2.0);
320
321        let mut fwd = rhs.clone();
322        apply_inplace(&pc, &mut fwd);
323        let mut tr = rhs.clone();
324        with_stack(pc.transpose_apply_in_place_scratch(1, Par::Seq), |stack| {
325            pc.transpose_apply_in_place(tr.as_mut(), Par::Seq, stack);
326        });
327        // For a nonsymmetric operator the two must differ.
328        let diff: f64 = (0..8)
329            .map(|i| (fwd.as_ref().get(i, 0) - tr.as_ref().get(i, 0)).abs())
330            .sum();
331        assert!(diff > 1e-8, "transpose apply should differ from forward");
332    }
333
334    #[test]
335    fn out_of_place_matches_in_place() {
336        let a = tridiagonal(10, 4.0, -2.0, -1.0);
337        let pc = Spai::try_new(a.as_ref(), SpaiPattern::ColumnsOfA).unwrap();
338        let rhs = Mat::<f64>::from_fn(10, 2, |i, j| ((i + 3 * j) % 7) as f64 - 3.0);
339        let mut out = Mat::<f64>::zeros(10, 2);
340        pc.apply(out.as_mut(), rhs.as_ref(), Par::Seq, MemStack::new(&mut []));
341        let mut inplace = rhs.clone();
342        apply_inplace(&pc, &mut inplace);
343        assert_close(out.as_ref(), inplace.as_ref(), 1e-12);
344    }
345
346    #[test]
347    fn rejects_non_square() {
348        let mut triplets = Vec::new();
349        for i in 0..3 {
350            triplets.push(Triplet::new(i, i, 1.0));
351        }
352        let a = SparseColMat::<usize, f64>::try_new_from_triplets(3, 4, &triplets).unwrap();
353        assert_eq!(
354            Spai::try_new(a.as_ref(), SpaiPattern::ColumnsOfA).unwrap_err(),
355            SpaiError::NonSquareMatrix { nrows: 3, ncols: 4 }
356        );
357    }
358}
faer_precond/spai/mod.rs

faer_precond/spai/
mod.rs
