Skip to main content

cjc_runtime/
sparse_eigen.rs

1//! Sparse Eigenvalue Solvers — Lanczos and Arnoldi
2//!
3//! Deterministic sparse eigenvalue extraction for symmetric (Lanczos) and
4//! general (Arnoldi) matrices. All operations use Kahan summation for
5//! order-invariant deterministic reductions.
6
7use crate::sparse::SparseCsr;
8use cjc_repro::kahan_sum_f64;
9
10// ---------------------------------------------------------------------------
11// Lanczos Algorithm (Symmetric Matrices)
12// ---------------------------------------------------------------------------
13
14/// Result of a Lanczos eigenvalue computation.
15#[derive(Debug, Clone)]
16pub struct LanczosResult {
17    /// Eigenvalues (sorted ascending).
18    pub eigenvalues: Vec<f64>,
19    /// Eigenvectors as rows (each row is an eigenvector).
20    pub eigenvectors: Vec<Vec<f64>>,
21}
22
23/// Compute the `k` largest eigenvalues of a symmetric sparse matrix using
24/// the Lanczos algorithm.
25///
26/// # Arguments
27/// * `matrix` - Symmetric sparse CSR matrix
28/// * `k` - Number of eigenvalues to compute (must be <= matrix dimension)
29/// * `max_iter` - Maximum Lanczos iterations (typically 2*k to 10*k)
30///
31/// # Returns
32/// `LanczosResult` with eigenvalues and corresponding eigenvectors.
33///
34/// # Determinism
35/// Uses Kahan summation for all dot products and norms.
36/// Result is deterministic for the same input.
37pub 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    // Lanczos vectors (orthonormal basis)
45    let mut v_prev = vec![0.0_f64; n];
46    // Start with uniform vector (deterministic, has components along all eigenvectors)
47    let inv_sqrt_n = 1.0 / (n as f64).sqrt();
48    let mut v_curr = vec![inv_sqrt_n; n];
49
50    // Tridiagonal matrix elements
51    let mut alpha = Vec::with_capacity(iters); // diagonal
52    let mut beta = Vec::with_capacity(iters);  // off-diagonal
53
54    for j in 0..iters {
55        // w = A * v_j
56        let w_raw = matrix.matvec(&v_curr).unwrap();
57        let mut w = w_raw;
58
59        // alpha_j = v_j^T * w
60        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        // w = w - alpha_j * v_j - beta_{j-1} * v_{j-1}
67        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        // beta_j = ||w||
73        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            // Invariant subspace found
78            break;
79        }
80        beta.push(beta_j);
81
82        // v_{j+1} = w / beta_j
83        v_prev = v_curr;
84        v_curr = w.iter().map(|&x| x / beta_j).collect();
85    }
86
87    // Solve tridiagonal eigenvalue problem using QR iteration
88    let (eigenvalues, _) = tridiagonal_qr(&alpha, &beta);
89
90    // Return the k largest eigenvalues (sorted ascending)
91    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![], // Eigenvector recovery deferred to library layer
100    }
101}
102
103// ---------------------------------------------------------------------------
104// Arnoldi Algorithm (General Matrices)
105// ---------------------------------------------------------------------------
106
107/// Result of an Arnoldi eigenvalue computation.
108#[derive(Debug, Clone)]
109pub struct ArnoldiResult {
110    /// Eigenvalue approximations (real parts).
111    pub eigenvalues_real: Vec<f64>,
112    /// Eigenvalue approximations (imaginary parts).
113    pub eigenvalues_imag: Vec<f64>,
114}
115
116/// Compute eigenvalue approximations of a general sparse matrix using
117/// the Arnoldi iteration to build an upper Hessenberg matrix, then
118/// extract eigenvalues from the Hessenberg form.
119///
120/// # Arguments
121/// * `matrix` - Sparse CSR matrix (square)
122/// * `k` - Number of Arnoldi iterations (determines subspace size)
123///
124/// # Returns
125/// `ArnoldiResult` with real and imaginary parts of eigenvalue approximations.
126pub 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    // Arnoldi vectors
132    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]; // Upper Hessenberg (m+1) x m
134
135    // q_0 = e_0 (deterministic starting vector)
136    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        // w = A * q_j
142        let mut w = matrix.matvec(&q[j]).unwrap();
143
144        // Modified Gram-Schmidt
145        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        // h_{j+1,j} = ||w||
156        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        // q_{j+1} = w / h_{j+1,j}
168        let q_next: Vec<f64> = w.iter().map(|&x| x / h_next).collect();
169        q.push(q_next);
170    }
171
172    // Extract eigenvalues from the m x m upper Hessenberg matrix
173    // For now, use the diagonal as approximations (Ritz values)
174    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
184// ---------------------------------------------------------------------------
185// Tridiagonal QR Iteration (Internal)
186// ---------------------------------------------------------------------------
187
188/// Solve a symmetric tridiagonal eigenvalue problem using implicit QR iteration.
189///
190/// # Arguments
191/// * `diag` - Diagonal elements (alpha)
192/// * `offdiag` - Off-diagonal elements (beta)
193///
194/// # Returns
195/// (eigenvalues, number of iterations used)
196fn 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    // Pad e to length n (e[n-1] unused but simplifies indexing)
208    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    // Implicit QL algorithm with Wilkinson shift (LAPACK tql1 approach)
216    for l_start in 0..n {
217        let mut iters_this = 0;
218        loop {
219            // Find the unreduced submatrix: smallest m >= l_start with e[m] negligible
220            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; // eigenvalue d[l_start] converged
230            }
231
232            iters_this += 1;
233            total_iters += 1;
234            if iters_this > max_iter {
235                break;
236            }
237
238            // Wilkinson shift
239            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            // Chase bulge from m-1 down to l_start
249            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                    // Deflation
256                    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// ---------------------------------------------------------------------------
282// Tests
283// ---------------------------------------------------------------------------
284
285#[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        // All eigenvalues of identity are 1.0
322        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        // Top 2 eigenvalues should be 4 and 5
332        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        // Arnoldi on identity should give eigenvalues close to 1.0
343        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        // 2x2 tridiagonal: [[3, 1], [1, 3]] -> eigenvalues 2 and 4
351        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}