1#[cfg(feature = "complex")]
2use crate::algebra::bridge::BridgeScratch;
3use crate::algebra::parallel;
4use crate::algebra::prelude::*;
5use crate::error::KError;
6#[cfg(all(not(feature = "complex"), feature = "backend-faer"))]
7use crate::matrix::convert::csr_from_linop;
8#[cfg(feature = "backend-faer")]
9use crate::matrix::op::GenericCsrOp;
10use crate::matrix::op::LinOp;
11use crate::matrix::sparse::CsrMatrix;
12#[cfg(feature = "complex")]
13use crate::ops::kpc::KPreconditioner;
14#[cfg(feature = "complex")]
15use crate::preconditioner::Preconditioner as ObjPreconditioner;
16#[cfg(feature = "complex")]
17use crate::preconditioner::bridge::{
18 apply_pc_mut_s as bridge_apply_pc_mut_s, apply_pc_s as bridge_apply_pc_s,
19};
20use crate::preconditioner::stats::{PcIntrospect, PcStats};
21use crate::preconditioner::{LocalPreconditioner, PcDistributedSupport, PcSide, Preconditioner};
22#[cfg(feature = "backend-faer")]
23use faer::Mat;
24use std::sync::atomic::{AtomicPtr, AtomicU64, Ordering};
25
26#[derive(Clone, Copy, Debug, PartialEq, Eq)]
27pub enum JacobiDiagMode {
28 ZeroOnDefect,
29 FixDiagonal,
30 RowL1OnDefect,
31}
32
33pub struct Jacobi {
34 pub(crate) diag_inv: Vec<S>,
35 n: usize,
36 tiny_diag_threshold: R,
37 fix_diag_replacement: R,
38 diag_mode: JacobiDiagMode,
39 applies: AtomicU64,
40}
41impl Default for Jacobi {
42 fn default() -> Self {
43 Self::new()
44 }
45}
46
47impl Jacobi {
48 pub fn new() -> Self {
49 Self {
50 diag_inv: Vec::new(),
51 n: 0,
52 tiny_diag_threshold: 1e-14,
53 fix_diag_replacement: 1e-12,
54 diag_mode: JacobiDiagMode::ZeroOnDefect,
55 applies: AtomicU64::new(0),
56 }
57 }
58
59 pub fn with_diag_mode(mut self, mode: JacobiDiagMode) -> Self {
60 self.diag_mode = mode;
61 self
62 }
63
64 pub fn with_tiny_diag_threshold(mut self, threshold: R) -> Self {
65 self.tiny_diag_threshold = threshold;
66 self
67 }
68
69 pub fn with_fix_diag_replacement(mut self, replacement: R) -> Self {
70 self.fix_diag_replacement = replacement;
71 self
72 }
73
74 fn diag_inverse_for_row(&self, aii: S, row_l1: R) -> S {
75 if aii.abs() > self.tiny_diag_threshold {
76 return aii.inv();
77 }
78 match self.diag_mode {
79 JacobiDiagMode::ZeroOnDefect => S::zero(),
80 JacobiDiagMode::FixDiagonal => {
81 let replacement = self
82 .fix_diag_replacement
83 .max(self.tiny_diag_threshold)
84 .max(1e-30);
85 S::from_real(replacement).inv()
86 }
87 JacobiDiagMode::RowL1OnDefect => {
88 let scale = row_l1.max(self.tiny_diag_threshold).max(1e-30);
89 S::from_real(scale).inv()
90 }
91 }
92 }
93
94 fn fill_diag_from_csr(&mut self, csr: &CsrMatrix<S>) {
95 let n = csr.nrows().min(csr.ncols());
96 self.diag_inv.resize(n, S::zero());
97 for i in 0..n {
98 let rs = csr.row_ptr()[i];
99 let re = csr.row_ptr()[i + 1];
100 let mut aii = S::zero();
101 let mut row_l1 = 0.0;
102 for p in rs..re {
103 row_l1 += csr.values()[p].abs();
104 if csr.col_idx()[p] == i {
105 aii = csr.values()[p];
106 }
107 }
108 self.diag_inv[i] = self.diag_inverse_for_row(aii, row_l1);
109 }
110 self.n = n;
111 }
112
113 fn recompute(&mut self, pmat: &dyn LinOp<S = S>) -> Result<(), KError> {
114 if let Some(csr) = pmat.as_any().downcast_ref::<CsrMatrix<S>>() {
115 self.fill_diag_from_csr(csr);
116 return Ok(());
117 }
118 #[cfg(feature = "backend-faer")]
119 if let Some(generic) = pmat.as_any().downcast_ref::<GenericCsrOp<S>>() {
120 let csr = generic.matrix();
121 let n = csr.nrows.min(csr.ncols);
122 self.diag_inv.resize(n, S::zero());
123 for i in 0..n {
124 let rs = csr.rowptr[i];
125 let re = csr.rowptr[i + 1];
126 let mut aii = S::zero();
127 let mut row_l1 = 0.0;
128 for p in rs..re {
129 row_l1 += csr.values[p].abs();
130 if csr.colind[p] == i {
131 aii = csr.values[p];
132 }
133 }
134 self.diag_inv[i] = self.diag_inverse_for_row(aii, row_l1);
135 }
136 self.n = n;
137 return Ok(());
138 }
139 #[cfg(feature = "backend-faer")]
140 if let Some(d) = pmat.as_any().downcast_ref::<Mat<S>>() {
141 let n = d.nrows().min(d.ncols());
142 self.diag_inv.resize(n, S::zero());
143 for i in 0..n {
144 let aii = d[(i, i)];
145 let mut row_l1 = 0.0;
146 for j in 0..d.ncols() {
147 row_l1 += d[(i, j)].abs();
148 }
149 self.diag_inv[i] = self.diag_inverse_for_row(aii, row_l1);
150 }
151 self.n = n;
152 return Ok(());
153 }
154 #[cfg(all(not(feature = "complex"), feature = "backend-faer"))]
155 {
156 let csr = csr_from_linop(pmat, 0.0)?;
157 self.fill_diag_from_csr(&csr);
158 return Ok(());
159 }
160 Err(KError::InvalidInput("Jacobi needs Dense or CSR".into()))
161 }
162}
163impl Preconditioner for Jacobi {
164 fn dims(&self) -> (usize, usize) {
165 (self.n, self.n)
166 }
167
168 fn setup(&mut self, pmat: &dyn LinOp<S = S>) -> Result<(), KError> {
169 self.recompute(pmat)
170 }
171 fn supports_numeric_update(&self) -> bool {
172 true
173 }
174
175 fn update_numeric(&mut self, pmat: &dyn LinOp<S = S>) -> Result<(), KError> {
176 self.recompute(pmat)
177 }
178
179 fn required_format(&self) -> crate::matrix::format::OpFormat {
180 crate::matrix::format::OpFormat::Csr
181 }
182 fn apply(&self, _side: PcSide, r: &[S], z: &mut [S]) -> Result<(), KError> {
183 if r.len() != self.n || z.len() != self.n {
184 return Err(KError::InvalidInput(format!(
185 "Jacobi::apply dimension mismatch: n={}, r.len()={}, z.len()={}",
186 self.n,
187 r.len(),
188 z.len()
189 )));
190 }
191 let z_ptr = AtomicPtr::new(z.as_mut_ptr());
192 parallel::par_for_each_index(r.len(), move |i| unsafe {
193 let z_ptr = z_ptr.load(Ordering::Relaxed);
194 *z_ptr.add(i) = self.diag_inv[i] * r[i];
195 });
196 self.applies.fetch_add(1, Ordering::Relaxed);
197 Ok(())
198 }
199
200 fn distributed_support(&self) -> PcDistributedSupport {
201 PcDistributedSupport::Distributed
202 }
203}
204
205impl LocalPreconditioner for Jacobi {
206 fn dims(&self) -> (usize, usize) {
207 (self.n, self.n)
208 }
209
210 fn apply_local(&self, x: &[S], y: &mut [S]) -> Result<(), KError> {
211 if x.len() != self.n || y.len() != self.n {
212 return Err(KError::InvalidInput(format!(
213 "Jacobi::apply_local dimension mismatch: n={}, x.len()={}, y.len()={}",
214 self.n,
215 x.len(),
216 y.len()
217 )));
218 }
219
220 let y_ptr = AtomicPtr::new(y.as_mut_ptr());
221 parallel::par_for_each_index(x.len(), move |i| unsafe {
222 let y_ptr = y_ptr.load(Ordering::Relaxed);
223 *y_ptr.add(i) = self.diag_inv[i] * x[i];
224 });
225 self.applies.fetch_add(1, Ordering::Relaxed);
226 Ok(())
227 }
228}
229
230#[cfg(feature = "complex")]
231impl KPreconditioner for Jacobi {
232 type Scalar = S;
233
234 #[inline]
235 fn dims(&self) -> (usize, usize) {
236 (self.n, self.n)
237 }
238
239 fn apply_s(
240 &self,
241 side: PcSide,
242 x: &[S],
243 y: &mut [S],
244 scratch: &mut BridgeScratch,
245 ) -> Result<(), KError> {
246 bridge_apply_pc_s(self, side, x, y, scratch)
247 }
248
249 fn apply_mut_s(
250 &mut self,
251 side: PcSide,
252 x: &[S],
253 y: &mut [S],
254 scratch: &mut BridgeScratch,
255 ) -> Result<(), KError> {
256 bridge_apply_pc_mut_s(self, side, x, y, scratch)
257 }
258
259 fn on_restart_s(&mut self, outer_iter: usize, residual_norm: R) -> Result<(), KError> {
260 ObjPreconditioner::on_restart(self, outer_iter, residual_norm)
261 }
262}
263
264#[cfg(all(feature = "backend-faer", not(feature = "complex")))]
265impl crate::preconditioner::legacy::Preconditioner<Mat<f64>, Vec<f64>> for Jacobi {
266 fn setup(&mut self, a: &Mat<f64>) -> Result<(), KError> {
267 self.recompute(a)
268 }
269 fn apply(&self, side: PcSide, r: &Vec<f64>, z: &mut Vec<f64>) -> Result<(), KError> {
270 crate::preconditioner::Preconditioner::apply(self, side, r.as_slice(), z.as_mut_slice())
271 }
272}
273
274#[cfg(all(test, feature = "backend-faer", not(feature = "complex")))]
275mod tests {
276 use super::*;
277 use crate::algebra::bridge::BridgeScratch;
278 use crate::ops::kpc::KPreconditioner;
279
280 #[test]
281 fn apply_s_matches_real_path() {
282 let mut pc = Jacobi::new();
283 let dense = Mat::<f64>::from_fn(2, 2, |i, j| if i == j { [4.0, 9.0][i] } else { 0.0 });
284 pc.setup(&dense).expect("jacobi setup");
285
286 let rhs_real = [8.0, 18.0];
287 let mut out_real = [0.0; 2];
288 pc.apply(PcSide::Left, &rhs_real, &mut out_real)
289 .expect("jacobi apply real");
290
291 let rhs_s: Vec<S> = rhs_real.iter().copied().map(S::from_real).collect();
292 let mut out_s = vec![S::zero(); rhs_s.len()];
293 let mut scratch = BridgeScratch::default();
294 pc.apply_s(PcSide::Left, &rhs_s, &mut out_s, &mut scratch)
295 .expect("jacobi apply_s");
296
297 for (ys, &yr) in out_s.iter().zip(out_real.iter()) {
298 assert!((ys.real() - yr).abs() < 1e-12);
299 }
300 }
301
302 #[test]
303 fn fix_diagonal_mode_replaces_tiny_diagonal() {
304 let mut pc = Jacobi::new()
305 .with_diag_mode(JacobiDiagMode::FixDiagonal)
306 .with_tiny_diag_threshold(1e-10)
307 .with_fix_diag_replacement(2.0);
308 let dense = Mat::<f64>::from_fn(2, 2, |i, j| {
309 if i == 0 && j == 0 {
310 0.0
311 } else if i == j {
312 4.0
313 } else {
314 0.0
315 }
316 });
317 pc.setup(&dense).expect("jacobi setup");
318
319 let rhs_real = [8.0, 8.0];
320 let mut out_real = [0.0; 2];
321 pc.apply(PcSide::Left, &rhs_real, &mut out_real)
322 .expect("jacobi apply real");
323 assert!((out_real[0] - 4.0).abs() < 1e-12);
324 assert!((out_real[1] - 2.0).abs() < 1e-12);
325 }
326
327 #[test]
328 fn row_l1_mode_uses_row_sum_on_missing_diagonal() {
329 let mut pc = Jacobi::new()
330 .with_diag_mode(JacobiDiagMode::RowL1OnDefect)
331 .with_tiny_diag_threshold(1e-12);
332 let csr = CsrMatrix::from_csr(2, 2, vec![0, 1, 2], vec![1, 1], vec![3.0, 4.0]);
333 pc.setup(&csr).expect("jacobi setup");
334
335 let rhs_real = [6.0, 8.0];
336 let mut out_real = [0.0; 2];
337 pc.apply(PcSide::Left, &rhs_real, &mut out_real)
338 .expect("jacobi apply real");
339 assert!((out_real[0] - 2.0).abs() < 1e-12);
340 assert!((out_real[1] - 2.0).abs() < 1e-12);
341 }
342}
343
344impl PcIntrospect for Jacobi {
345 fn stats(&self) -> PcStats {
346 PcStats {
347 name: "Jacobi",
348 n: self.n,
349 build_ms: 0.0,
350 nnz_pc: self.n,
351 fill_ratio: 0.0,
352 applies: self.applies.load(Ordering::Relaxed),
353 complex_kernel: None,
354 setup_mode: None,
355 fallback_reason: None,
356 residual_reduction_per_time: None,
357 }
358 }
359}