faer_precond/spai/
build.rs1use 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 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 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 *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
107fn 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}