cjc_runtime/
sparse_eigen.rs1use crate::sparse::SparseCsr;
8use cjc_repro::kahan_sum_f64;
9
10#[derive(Debug, Clone)]
16pub struct LanczosResult {
17 pub eigenvalues: Vec<f64>,
19 pub eigenvectors: Vec<Vec<f64>>,
21}
22
23pub fn lanczos_eigsh(matrix: &SparseCsr, k: usize, max_iter: usize) -> LanczosResult {
38 let n = matrix.nrows;
39 assert_eq!(n, matrix.ncols, "lanczos_eigsh: matrix must be square");
40 assert!(k > 0 && k <= n, "lanczos_eigsh: k must be in 1..=n");
41
42 let iters = max_iter.min(n);
43
44 let mut v_prev = vec![0.0_f64; n];
46 let inv_sqrt_n = 1.0 / (n as f64).sqrt();
48 let mut v_curr = vec![inv_sqrt_n; n];
49
50 let mut alpha = Vec::with_capacity(iters); let mut beta = Vec::with_capacity(iters); for j in 0..iters {
55 let w_raw = matrix.matvec(&v_curr).unwrap();
57 let mut w = w_raw;
58
59 let dot_products: Vec<f64> = v_curr.iter().zip(w.iter())
61 .map(|(&a, &b)| a * b)
62 .collect();
63 let alpha_j = kahan_sum_f64(&dot_products);
64 alpha.push(alpha_j);
65
66 let beta_prev = if j > 0 { beta[j - 1] } else { 0.0 };
68 for i in 0..n {
69 w[i] -= alpha_j * v_curr[i] + beta_prev * v_prev[i];
70 }
71
72 let norm_sq: Vec<f64> = w.iter().map(|&x| x * x).collect();
74 let beta_j = kahan_sum_f64(&norm_sq).sqrt();
75
76 if beta_j < 1e-14 {
77 break;
79 }
80 beta.push(beta_j);
81
82 v_prev = v_curr;
84 v_curr = w.iter().map(|&x| x / beta_j).collect();
85 }
86
87 let (eigenvalues, _) = tridiagonal_qr(&alpha, &beta);
89
90 let mut sorted_evals: Vec<f64> = eigenvalues;
92 sorted_evals.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
93
94 let k_actual = k.min(sorted_evals.len());
95 let top_k: Vec<f64> = sorted_evals[sorted_evals.len() - k_actual..].to_vec();
96
97 LanczosResult {
98 eigenvalues: top_k,
99 eigenvectors: vec![], }
101}
102
103#[derive(Debug, Clone)]
109pub struct ArnoldiResult {
110 pub eigenvalues_real: Vec<f64>,
112 pub eigenvalues_imag: Vec<f64>,
114}
115
116pub fn arnoldi_eigs(matrix: &SparseCsr, k: usize) -> ArnoldiResult {
127 let n = matrix.nrows;
128 assert_eq!(n, matrix.ncols, "arnoldi_eigs: matrix must be square");
129 let m = k.min(n);
130
131 let mut q: Vec<Vec<f64>> = Vec::with_capacity(m + 1);
133 let mut h: Vec<Vec<f64>> = vec![vec![0.0; m]; m + 1]; let mut q0 = vec![0.0_f64; n];
137 q0[0] = 1.0;
138 q.push(q0);
139
140 for j in 0..m {
141 let mut w = matrix.matvec(&q[j]).unwrap();
143
144 for i in 0..=j {
146 let dots: Vec<f64> = q[i].iter().zip(w.iter())
147 .map(|(&a, &b)| a * b)
148 .collect();
149 h[i][j] = kahan_sum_f64(&dots);
150 for l in 0..n {
151 w[l] -= h[i][j] * q[i][l];
152 }
153 }
154
155 let norm_sq: Vec<f64> = w.iter().map(|&x| x * x).collect();
157 let h_next = kahan_sum_f64(&norm_sq).sqrt();
158
159 if j + 1 < m + 1 {
160 h[j + 1][j] = h_next;
161 }
162
163 if h_next < 1e-14 {
164 break;
165 }
166
167 let q_next: Vec<f64> = w.iter().map(|&x| x / h_next).collect();
169 q.push(q_next);
170 }
171
172 let actual_m = (q.len() - 1).min(m);
175 let eigenvalues_real: Vec<f64> = (0..actual_m).map(|i| h[i][i]).collect();
176 let eigenvalues_imag = vec![0.0_f64; actual_m];
177
178 ArnoldiResult {
179 eigenvalues_real,
180 eigenvalues_imag,
181 }
182}
183
184fn tridiagonal_qr(diag: &[f64], offdiag: &[f64]) -> (Vec<f64>, usize) {
197 let n = diag.len();
198 if n == 0 {
199 return (vec![], 0);
200 }
201 if n == 1 {
202 return (vec![diag[0]], 0);
203 }
204
205 let mut d = diag.to_vec();
206 let mut e = offdiag.to_vec();
207 while e.len() < n {
209 e.push(0.0);
210 }
211
212 let max_iter = 30 * n;
213 let mut total_iters = 0;
214
215 for l_start in 0..n {
217 let mut iters_this = 0;
218 loop {
219 let mut m = l_start;
221 while m < n - 1 {
222 let tst = d[m].abs() + d[m + 1].abs();
223 if e[m].abs() <= 1e-15 * tst {
224 break;
225 }
226 m += 1;
227 }
228 if m == l_start {
229 break; }
231
232 iters_this += 1;
233 total_iters += 1;
234 if iters_this > max_iter {
235 break;
236 }
237
238 let mut g = (d[l_start + 1] - d[l_start]) / (2.0 * e[l_start]);
240 let r = g.hypot(1.0);
241 g = d[m] - d[l_start] + e[l_start] / (g + r.copysign(g));
242
243 let mut s = 1.0;
244 let mut c = 1.0;
245 let mut p = 0.0;
246 let mut early_break = false;
247
248 for i in (l_start..m).rev() {
250 let f = s * e[i];
251 let b = c * e[i];
252 let r = f.hypot(g);
253 e[i + 1] = r;
254 if r.abs() < 1e-30 {
255 d[i + 1] -= p;
257 e[m] = 0.0;
258 early_break = true;
259 break;
260 }
261 s = f / r;
262 c = g / r;
263 g = d[i + 1] - p;
264 let r2 = (d[i] - g) * s + 2.0 * c * b;
265 p = s * r2;
266 d[i + 1] = g + p;
267 g = c * r2 - b;
268 }
269
270 if !early_break {
271 d[l_start] -= p;
272 e[l_start] = g;
273 e[m] = 0.0;
274 }
275 }
276 }
277
278 (d, total_iters)
279}
280
281#[cfg(test)]
286mod tests {
287 use super::*;
288 use crate::sparse::SparseCoo;
289
290 fn identity_csr(n: usize) -> SparseCsr {
291 let mut values = Vec::new();
292 let mut row_indices = Vec::new();
293 let mut col_indices = Vec::new();
294 for i in 0..n {
295 values.push(1.0);
296 row_indices.push(i);
297 col_indices.push(i);
298 }
299 let coo = SparseCoo::new(values, row_indices, col_indices, n, n);
300 SparseCsr::from_coo(&coo)
301 }
302
303 fn diagonal_csr(diag: &[f64]) -> SparseCsr {
304 let n = diag.len();
305 let mut values = Vec::new();
306 let mut row_indices = Vec::new();
307 let mut col_indices = Vec::new();
308 for (i, &d) in diag.iter().enumerate() {
309 values.push(d);
310 row_indices.push(i);
311 col_indices.push(i);
312 }
313 let coo = SparseCoo::new(values, row_indices, col_indices, n, n);
314 SparseCsr::from_coo(&coo)
315 }
316
317 #[test]
318 fn test_lanczos_identity() {
319 let mat = identity_csr(5);
320 let result = lanczos_eigsh(&mat, 3, 20);
321 for &ev in &result.eigenvalues {
323 assert!((ev - 1.0).abs() < 1e-10, "expected 1.0, got {ev}");
324 }
325 }
326
327 #[test]
328 fn test_lanczos_diagonal() {
329 let mat = diagonal_csr(&[1.0, 2.0, 3.0, 4.0, 5.0]);
330 let result = lanczos_eigsh(&mat, 2, 50);
331 assert!(result.eigenvalues.len() >= 2);
333 let top = &result.eigenvalues;
334 assert!((top[top.len() - 1] - 5.0).abs() < 1e-8,
335 "expected 5.0, got {}", top[top.len() - 1]);
336 }
337
338 #[test]
339 fn test_arnoldi_identity() {
340 let mat = identity_csr(4);
341 let result = arnoldi_eigs(&mat, 4);
342 for &ev in &result.eigenvalues_real {
344 assert!((ev - 1.0).abs() < 1e-10, "expected ~1.0, got {ev}");
345 }
346 }
347
348 #[test]
349 fn test_tridiagonal_qr_simple() {
350 let (evals, _) = tridiagonal_qr(&[3.0, 3.0], &[1.0]);
352 let mut sorted = evals.clone();
353 sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
354 assert!((sorted[0] - 2.0).abs() < 1e-10, "expected 2.0, got {}", sorted[0]);
355 assert!((sorted[1] - 4.0).abs() < 1e-10, "expected 4.0, got {}", sorted[1]);
356 }
357
358 #[test]
359 fn test_lanczos_deterministic() {
360 let mat = diagonal_csr(&[1.0, 3.0, 5.0, 7.0, 9.0]);
361 let r1 = lanczos_eigsh(&mat, 3, 30);
362 let r2 = lanczos_eigsh(&mat, 3, 30);
363 assert_eq!(r1.eigenvalues.len(), r2.eigenvalues.len());
364 for (a, b) in r1.eigenvalues.iter().zip(r2.eigenvalues.iter()) {
365 assert_eq!(a.to_bits(), b.to_bits(), "non-deterministic eigenvalue");
366 }
367 }
368}