faer_precond/fsai/
mod.rs

1//! Factorized sparse approximate inverse (FSAI) preconditioner.
2//!
3//! FSAI builds an *explicit* approximation to `A^{-1}` for a Hermitian
4//! positive-definite `A`: a sparse lower-triangular factor `G ~= L^{-1}` (with
5//! `A = L L^H`) such that `M^{-1} = G^H G`. Applying the preconditioner is then
6//! two sparse matrix-vector products — no triangular solves — which, like the
7//! [polynomial](crate::poly) preconditioner, parallelises far better than the
8//! sequential sweeps of ILU/IC. Building `G` is also naturally parallel: each
9//! row is an independent small dense SPD solve.
10//!
11//! # When to use it
12//!
13//! Reach for FSAI on SPD systems when you want an approximate inverse whose
14//! *apply* is a matvec — many cores, accelerators, or repeated applies where the
15//! triangular-solve dependency chain of [`crate::Ic0`] is the bottleneck. The
16//! trade-off is a heavier build (one dense factorisation per row) and a quality
17//! that hinges on the chosen pattern.
18//!
19//! # Pattern
20//!
21//! The accuracy/cost knob is the prescribed lower-triangular pattern of `G`:
22//!
23//! - [`FsaiPattern::LowerOfA`] — the lower triangle of `A` (cheapest).
24//! - [`FsaiPattern::LowerOfPower`] — the lower triangle of `pattern(A^power)`,
25//!   denser and more accurate for larger `power`.
26//!
27//! Adaptive pattern selection is left as future work.
28//!
29//! # Storage
30//!
31//! `G` is stored CSC. Apply reads it through a
32//! [`faer::sparse::SparseColMatRef`] and its adjoint view, with
33//! a single work column block drawn from the caller's [`MemStack`] — no heap
34//! allocation.
35
36use core::fmt::Debug;
37
38use dyn_stack::{MemStack, StackReq};
39use faer::matrix_free::{BiLinOp, BiPrecond, LinOp, Precond};
40use faer::sparse::{SparseColMatRef, SymbolicSparseColMatRef};
41use faer::{MatMut, MatRef, Par};
42use faer_traits::{ComplexField, Index};
43
44mod apply;
45mod build;
46
47/// Prescribed lower-triangular pattern for the FSAI factor `G`.
48#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
49pub enum FsaiPattern {
50    /// Lower triangle of `A`.
51    #[default]
52    LowerOfA,
53    /// Lower triangle of `pattern(A^power)` (`power >= 1`; `1` equals
54    /// [`FsaiPattern::LowerOfA`]).
55    LowerOfPower { power: usize },
56}
57
58/// Error returned by FSAI construction.
59#[derive(Debug, Clone, PartialEq, Eq)]
60pub enum FsaiError {
61    /// The source matrix was not square.
62    NonSquareMatrix { nrows: usize, ncols: usize },
63    /// A `LowerOfPower` power of zero was requested.
64    InvalidPower,
65    /// The local block at row `row` was not positive definite.
66    NotPositiveDefinite { row: usize },
67}
68
69impl core::fmt::Display for FsaiError {
70    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
71        match self {
72            Self::NonSquareMatrix { nrows, ncols } => {
73                write!(f, "matrix must be square but is {nrows}x{ncols}")
74            }
75            Self::InvalidPower => f.write_str("FSAI pattern power must be at least 1"),
76            Self::NotPositiveDefinite { row } => {
77                write!(f, "local block at row {row} is not positive definite")
78            }
79        }
80    }
81}
82
83impl core::error::Error for FsaiError {}
84
85/// Factorized sparse approximate inverse: `M^{-1} = G^H G`.
86///
87/// Stores the lower-triangular factor `G` in CSC. Apply is two sparse matvecs
88/// with no triangular solves and no heap allocation. See the
89/// [module documentation](self) for guidance.
90#[derive(Debug, Clone)]
91pub struct Fsai<I, T> {
92    pub(crate) dim: usize,
93    pub(crate) g_col_ptr: Vec<I>,
94    pub(crate) g_row_idx: Vec<I>,
95    pub(crate) g_values: Vec<T>,
96}
97
98impl<I, T> Fsai<I, T> {
99    /// Dimension `n` of the preconditioner.
100    #[inline]
101    pub fn dim(&self) -> usize {
102        self.dim
103    }
104}
105
106impl<I: Index, T: ComplexField> Fsai<I, T> {
107    /// View over the lower-triangular factor `G`.
108    #[inline]
109    pub(crate) fn g_view(&self) -> SparseColMatRef<'_, I, T> {
110        let symbolic = unsafe {
111            SymbolicSparseColMatRef::<'_, I>::new_unchecked(
112                self.dim,
113                self.dim,
114                &self.g_col_ptr,
115                None,
116                &self.g_row_idx,
117            )
118        };
119        SparseColMatRef::new(symbolic, &self.g_values)
120    }
121}
122
123impl<I, T> LinOp<T> for Fsai<I, T>
124where
125    I: Index,
126    T: ComplexField + Debug + Sync,
127{
128    fn apply_scratch(&self, rhs_ncols: usize, _par: Par) -> StackReq {
129        apply::scratch(self, rhs_ncols)
130    }
131
132    fn nrows(&self) -> usize {
133        self.dim
134    }
135
136    fn ncols(&self) -> usize {
137        self.dim
138    }
139
140    fn apply(&self, out: MatMut<'_, T>, rhs: MatRef<'_, T>, par: Par, stack: &mut MemStack) {
141        apply::apply_out(self, out, rhs, false, par, stack);
142    }
143
144    fn conj_apply(&self, out: MatMut<'_, T>, rhs: MatRef<'_, T>, par: Par, stack: &mut MemStack) {
145        apply::apply_out(self, out, rhs, true, par, stack);
146    }
147}
148
149impl<I, T> Precond<T> for Fsai<I, T>
150where
151    I: Index,
152    T: ComplexField + Debug + Sync,
153{
154    fn apply_in_place_scratch(&self, rhs_ncols: usize, _par: Par) -> StackReq {
155        apply::scratch(self, rhs_ncols)
156    }
157
158    fn apply_in_place(&self, rhs: MatMut<'_, T>, par: Par, stack: &mut MemStack) {
159        apply::apply_inplace(self, rhs, false, par, stack);
160    }
161
162    fn conj_apply_in_place(&self, rhs: MatMut<'_, T>, par: Par, stack: &mut MemStack) {
163        apply::apply_inplace(self, rhs, true, par, stack);
164    }
165}
166
167impl<I, T> BiLinOp<T> for Fsai<I, T>
168where
169    I: Index,
170    T: ComplexField + Debug + Sync,
171{
172    fn transpose_apply_scratch(&self, rhs_ncols: usize, _par: Par) -> StackReq {
173        apply::scratch(self, rhs_ncols)
174    }
175
176    fn transpose_apply(
177        &self,
178        out: MatMut<'_, T>,
179        rhs: MatRef<'_, T>,
180        par: Par,
181        stack: &mut MemStack,
182    ) {
183        // M^{-T} = conj(M^{-1}).
184        apply::apply_out(self, out, rhs, true, par, stack);
185    }
186
187    fn adjoint_apply(
188        &self,
189        out: MatMut<'_, T>,
190        rhs: MatRef<'_, T>,
191        par: Par,
192        stack: &mut MemStack,
193    ) {
194        // M^{-H} = M^{-1} for Hermitian M.
195        apply::apply_out(self, out, rhs, false, par, stack);
196    }
197}
198
199impl<I, T> BiPrecond<T> for Fsai<I, T>
200where
201    I: Index,
202    T: ComplexField + Debug + Sync,
203{
204    fn transpose_apply_in_place_scratch(&self, rhs_ncols: usize, _par: Par) -> StackReq {
205        apply::scratch(self, rhs_ncols)
206    }
207
208    fn transpose_apply_in_place(&self, rhs: MatMut<'_, T>, par: Par, stack: &mut MemStack) {
209        apply::apply_inplace(self, rhs, true, par, stack);
210    }
211
212    fn adjoint_apply_in_place(&self, rhs: MatMut<'_, T>, par: Par, stack: &mut MemStack) {
213        apply::apply_inplace(self, rhs, false, par, stack);
214    }
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220    use core::mem::MaybeUninit;
221    use faer::sparse::{SparseColMat, Triplet};
222    use faer::{Mat, MatRef, mat};
223
224    fn with_stack(req: StackReq, f: impl FnOnce(&mut MemStack)) {
225        let nbytes = req.unaligned_bytes_required().max(1);
226        let mut buf = vec![MaybeUninit::<u8>::uninit(); nbytes].into_boxed_slice();
227        f(MemStack::new(&mut buf));
228    }
229
230    fn assert_close(lhs: MatRef<'_, f64>, rhs: MatRef<'_, f64>, tol: f64) {
231        assert_eq!(lhs.nrows(), rhs.nrows());
232        assert_eq!(lhs.ncols(), rhs.ncols());
233        for j in 0..lhs.ncols() {
234            for i in 0..lhs.nrows() {
235                let diff = (*lhs.get(i, j) - *rhs.get(i, j)).abs();
236                assert!(
237                    diff <= tol,
238                    "mismatch at ({i}, {j}): lhs={}, rhs={}, diff={diff}",
239                    *lhs.get(i, j),
240                    *rhs.get(i, j),
241                );
242            }
243        }
244    }
245
246    fn to_dense(a: &SparseColMat<usize, f64>) -> Mat<f64> {
247        let n = a.nrows();
248        let mut out = Mat::<f64>::zeros(n, a.ncols());
249        let a_ref = a.as_ref();
250        for j in 0..a.ncols() {
251            let rows = a_ref.symbolic().row_idx_of_col_raw(j);
252            let vals = a_ref.val_of_col(j);
253            for (r, v) in rows.iter().zip(vals.iter()) {
254                *out.as_mut().get_mut(*r, j) = *v;
255            }
256        }
257        out
258    }
259
260    fn tridiagonal(n: usize, diag: f64, off: f64) -> SparseColMat<usize, f64> {
261        let mut triplets = Vec::new();
262        for i in 0..n {
263            triplets.push(Triplet::new(i, i, diag));
264            if i > 0 {
265                triplets.push(Triplet::new(i, i - 1, off));
266                triplets.push(Triplet::new(i - 1, i, off));
267            }
268        }
269        SparseColMat::try_new_from_triplets(n, n, &triplets).unwrap()
270    }
271
272    fn laplacian_2d(grid: usize) -> SparseColMat<usize, f64> {
273        let n = grid * grid;
274        let mut triplets = Vec::new();
275        for gy in 0..grid {
276            for gx in 0..grid {
277                let idx = gy * grid + gx;
278                triplets.push(Triplet::new(idx, idx, 4.0));
279                if gx > 0 {
280                    triplets.push(Triplet::new(idx, idx - 1, -1.0));
281                }
282                if gx + 1 < grid {
283                    triplets.push(Triplet::new(idx, idx + 1, -1.0));
284                }
285                if gy > 0 {
286                    triplets.push(Triplet::new(idx, idx - grid, -1.0));
287                }
288                if gy + 1 < grid {
289                    triplets.push(Triplet::new(idx, idx + grid, -1.0));
290                }
291            }
292        }
293        SparseColMat::try_new_from_triplets(n, n, &triplets).unwrap()
294    }
295
296    fn apply_inplace(pc: &Fsai<usize, f64>, rhs: &mut Mat<f64>) {
297        with_stack(pc.apply_in_place_scratch(rhs.ncols(), Par::Seq), |stack| {
298            pc.apply_in_place(rhs.as_mut(), Par::Seq, stack);
299        });
300    }
301
302    fn residual_ratio(a: &SparseColMat<usize, f64>, pc: &Fsai<usize, f64>, b: &Mat<f64>) -> f64 {
303        let a_dense = to_dense(a);
304        let mut x = b.clone();
305        apply_inplace(pc, &mut x);
306        let residual = &a_dense * &x - b;
307        let b_norm: f64 = b.as_ref().col(0).iter().map(|v| v * v).sum::<f64>().sqrt();
308        let r_norm: f64 = residual
309            .as_ref()
310            .col(0)
311            .iter()
312            .map(|v| v * v)
313            .sum::<f64>()
314            .sqrt();
315        r_norm / b_norm
316    }
317
318    #[test]
319    fn diagonal_is_exact_inverse() {
320        let mut triplets = Vec::new();
321        for (i, &v) in [2.0, 4.0, 8.0].iter().enumerate() {
322            triplets.push(Triplet::new(i, i, v));
323        }
324        let a = SparseColMat::<usize, f64>::try_new_from_triplets(3, 3, &triplets).unwrap();
325        let pc = Fsai::try_new(a.as_ref(), FsaiPattern::LowerOfA).unwrap();
326        let mut x = mat![[2.0_f64], [8.0], [16.0]];
327        apply_inplace(&pc, &mut x);
328        let expected = mat![[1.0_f64], [2.0], [2.0]];
329        assert_close(x.as_ref(), expected.as_ref(), 1e-12);
330    }
331
332    #[test]
333    fn reduces_residual_on_laplacian() {
334        let a = laplacian_2d(8);
335        let n = a.nrows();
336        let pc = Fsai::try_new(a.as_ref(), FsaiPattern::LowerOfA).unwrap();
337        let b = Mat::<f64>::from_fn(n, 1, |i, _| (i % 7) as f64 - 3.0);
338        let ratio = residual_ratio(&a, &pc, &b);
339        assert!(ratio < 1.0, "FSAI should reduce the residual: {ratio}");
340    }
341
342    #[test]
343    fn denser_pattern_improves_accuracy() {
344        let a = laplacian_2d(8);
345        let n = a.nrows();
346        let b = Mat::<f64>::from_fn(n, 1, |i, _| (i % 7) as f64 - 3.0);
347        let p1 = Fsai::try_new(a.as_ref(), FsaiPattern::LowerOfA).unwrap();
348        let p2 = Fsai::try_new(a.as_ref(), FsaiPattern::LowerOfPower { power: 2 }).unwrap();
349        let r1 = residual_ratio(&a, &p1, &b);
350        let r2 = residual_ratio(&a, &p2, &b);
351        assert!(r2 < r1, "denser FSAI pattern should help: {r2} !< {r1}");
352    }
353
354    #[test]
355    fn symmetric_two_matvec_matches_dense() {
356        let a = tridiagonal(6, 4.0, -1.0);
357        let pc = Fsai::try_new(a.as_ref(), FsaiPattern::LowerOfA).unwrap();
358        // Forward apply == transpose apply for the real symmetric case.
359        let rhs = mat![[1.0_f64], [-2.0], [3.0], [0.5], [-1.0], [2.0]];
360        let mut fwd = rhs.clone();
361        apply_inplace(&pc, &mut fwd);
362        let mut tr = rhs.clone();
363        with_stack(pc.transpose_apply_in_place_scratch(1, Par::Seq), |stack| {
364            pc.transpose_apply_in_place(tr.as_mut(), Par::Seq, stack);
365        });
366        assert_close(fwd.as_ref(), tr.as_ref(), 1e-12);
367    }
368
369    #[test]
370    fn rejects_non_square() {
371        let mut triplets = Vec::new();
372        for i in 0..3 {
373            triplets.push(Triplet::new(i, i, 1.0));
374        }
375        let a = SparseColMat::<usize, f64>::try_new_from_triplets(3, 4, &triplets).unwrap();
376        assert_eq!(
377            Fsai::try_new(a.as_ref(), FsaiPattern::LowerOfA).unwrap_err(),
378            FsaiError::NonSquareMatrix { nrows: 3, ncols: 4 }
379        );
380    }
381}
faer_precond/fsai/mod.rs

faer_precond/fsai/
mod.rs
