Skip to main content

faer_precond/spai/
build.rs

1//! Construction of the SPAI approximate inverse `M`.
2//!
3//! SPAI minimises `||A M - I||_F`, which decouples into independent least-
4//! squares problems, one per column: `min ||A m_k - e_k||`. For a prescribed
5//! column pattern `J_k`, only the rows `I_k = union of pattern(A[:,j]) for j in
6//! J_k` are involved, so each solve is a small dense overdetermined system
7//! `A[I_k, J_k] m_k = e_k[I_k]`, handled by a QR least-squares solve.
8
9use faer::linalg::solvers::{Qr, SolveLstsq};
10use faer::sparse::SparseColMatRef;
11use faer::Mat;
12use faer_traits::math_utils::{copy, one};
13use faer_traits::{ComplexField, Index};
14
15use super::{Spai, SpaiError, SpaiPattern};
16
17impl<I: Index, T: ComplexField> Spai<I, T> {
18    /// Build a SPAI preconditioner for `A`.
19    ///
20    /// # Errors
21    ///
22    /// - [`SpaiError::NonSquareMatrix`] if `A` is not square.
23    /// - [`SpaiError::InvalidPower`] if a `ColumnsOfPower` power is zero.
24    pub fn try_new(a: SparseColMatRef<'_, I, T>, pattern: SpaiPattern) -> Result<Self, SpaiError> {
25        if a.nrows() != a.ncols() {
26            return Err(SpaiError::NonSquareMatrix {
27                nrows: a.nrows(),
28                ncols: a.ncols(),
29            });
30        }
31        let power = match pattern {
32            SpaiPattern::ColumnsOfA => 1,
33            SpaiPattern::ColumnsOfPower { power } => power,
34        };
35        if power == 0 {
36            return Err(SpaiError::InvalidPower);
37        }
38        let n = a.nrows();
39        let col_pats = column_patterns(a, power);
40
41        let mut m_col_ptr: Vec<I> = Vec::with_capacity(n + 1);
42        m_col_ptr.push(I::truncate(0));
43        let mut m_row_idx: Vec<I> = Vec::new();
44        let mut m_values: Vec<T> = Vec::new();
45
46        let mut marker = vec![usize::MAX; n];
47        let mut local = vec![0usize; n];
48
49        for k in 0..n {
50            let jk = &col_pats[k];
51
52            // Row set I_k = union of the patterns of A's columns in J_k.
53            let mut ik: Vec<usize> = Vec::new();
54            for &j in jk {
55                for raw in a.symbolic().row_idx_of_col_raw(j) {
56                    let r = raw.zx();
57                    if marker[r] != k {
58                        marker[r] = k;
59                        ik.push(r);
60                    }
61                }
62            }
63            ik.sort_unstable();
64            for (rr, &r) in ik.iter().enumerate() {
65                local[r] = rr;
66            }
67
68            let m_rows = ik.len();
69            let n_cols = jk.len();
70            let mut a_sub = Mat::<T>::zeros(m_rows, n_cols);
71            for (cc, &j) in jk.iter().enumerate() {
72                for (raw, val) in a
73                    .symbolic()
74                    .row_idx_of_col_raw(j)
75                    .iter()
76                    .zip(a.val_of_col(j).iter())
77                {
78                    let r = raw.zx();
79                    // every such r is in I_k by construction
80                    *a_sub.as_mut().get_mut(local[r], cc) = copy(val);
81                }
82            }
83
84            let mut e = Mat::<T>::zeros(m_rows, 1);
85            if marker[k] == k {
86                *e.as_mut().get_mut(local[k], 0) = one::<T>();
87            }
88
89            let m_k = Qr::new(a_sub.as_ref()).solve_lstsq(&e);
90
91            for (cc, &j) in jk.iter().enumerate() {
92                m_row_idx.push(I::truncate(j));
93                m_values.push(copy(m_k.as_ref().get(cc, 0)));
94            }
95            m_col_ptr.push(I::truncate(m_row_idx.len()));
96        }
97
98        Ok(Self {
99            dim: n,
100            m_col_ptr,
101            m_row_idx,
102            m_values,
103        })
104    }
105}
106
107/// Pattern of each column of `A^power` (each list sorted).
108fn column_patterns<I: Index, T: ComplexField>(
109    a: SparseColMatRef<'_, I, T>,
110    power: usize,
111) -> Vec<Vec<usize>> {
112    let n = a.ncols();
113    let mut p: Vec<Vec<usize>> = (0..n)
114        .map(|k| {
115            let mut v: Vec<usize> = a
116                .symbolic()
117                .row_idx_of_col_raw(k)
118                .iter()
119                .map(|r| r.zx())
120                .collect();
121            v.sort_unstable();
122            v
123        })
124        .collect();
125
126    let mut marker = vec![usize::MAX; n];
127    for _ in 1..power {
128        let mut np: Vec<Vec<usize>> = Vec::with_capacity(n);
129        for (k, col) in p.iter().enumerate() {
130            let mut rows = Vec::new();
131            for &m in col {
132                for raw in a.symbolic().row_idx_of_col_raw(m) {
133                    let r = raw.zx();
134                    if marker[r] != k {
135                        marker[r] = k;
136                        rows.push(r);
137                    }
138                }
139            }
140            rows.sort_unstable();
141            np.push(rows);
142        }
143        p = np;
144    }
145    p
146}