Skip to main content

lau_diffusion_agents/
spectral_diffusion.rs

1//! Spectral decomposition of diffusion operators, eigenfunction expansion.
2
3use nalgebra::{DMatrix, DVector};
4use serde::{Serialize, Deserialize};
5
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct SpectralDecomposition {
8    pub eigenvalues: Vec<f64>,
9    pub eigenvectors: DMatrix<f64>,
10}
11
12impl SpectralDecomposition {
13    pub fn decompose(matrix: &DMatrix<f64>) -> Result<Self, String> {
14        let n = matrix.nrows();
15        if n != matrix.ncols() { return Err("Matrix must be square".into()); }
16        if n == 0 { return Ok(Self { eigenvalues: vec![], eigenvectors: DMatrix::zeros(0, 0) }); }
17        let eig = matrix.clone().symmetric_eigen();
18        Ok(Self {
19            eigenvalues: eig.eigenvalues.iter().copied().collect(),
20            eigenvectors: eig.eigenvectors,
21        })
22    }
23
24    pub fn reconstruct(&self) -> DMatrix<f64> {
25        let n = self.eigenvalues.len();
26        if n == 0 { return DMatrix::zeros(0, 0); }
27        let mut result = DMatrix::zeros(n, n);
28        for k in 0..n {
29            let v = self.eigenvectors.column(k);
30            result += &v.scale(self.eigenvalues[k]) * &v.transpose();
31        }
32        result
33    }
34
35    pub fn low_rank(&self, k: usize) -> DMatrix<f64> {
36        let n = self.eigenvalues.len();
37        if n == 0 { return DMatrix::zeros(0, 0); }
38        let mut indices: Vec<usize> = (0..n).collect();
39        indices.sort_by(|&a, &b| self.eigenvalues[b].abs().partial_cmp(&self.eigenvalues[a].abs()).unwrap());
40        let mut result = DMatrix::zeros(n, n);
41        for &idx in indices.iter().take(k) {
42            let v = self.eigenvectors.column(idx);
43            result += &v.scale(self.eigenvalues[idx]) * &v.transpose();
44        }
45        result
46    }
47}
48
49pub fn heat_kernel_spectral(eigenvalues: &[f64], eigenvectors: &DMatrix<f64>, t: f64, i: usize, j: usize, n_terms: usize) -> f64 {
50    eigenvalues.iter().take(n_terms)
51        .zip(eigenvectors.column_iter().take(n_terms))
52        .map(|(lambda, phi)| (-lambda * t).exp() * phi[i] * phi[j])
53        .sum()
54}
55
56pub fn spectral_diffusion_solve(eigenvalues: &[f64], eigenvectors: &DMatrix<f64>, initial: &DVector<f64>, t: f64) -> DVector<f64> {
57    let n = eigenvalues.len();
58    let mut result = DVector::zeros(n);
59    for k in 0..n {
60        let phi_k = eigenvectors.column(k);
61        let c_k = phi_k.dot(initial);
62        result += phi_k.scale(c_k * (-eigenvalues[k] * t).exp());
63    }
64    result
65}
66
67pub fn spectral_gap(eigenvalues: &[f64]) -> Option<f64> {
68    let mut sorted = eigenvalues.to_vec();
69    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap());
70    sorted.iter().find(|&&e| e > 1e-10).copied()
71}
72
73pub fn diffusion_kernel_matrix(eigenvalues: &[f64], eigenvectors: &DMatrix<f64>, t: f64) -> DMatrix<f64> {
74    let n = eigenvalues.len();
75    let mut result = DMatrix::zeros(n, n);
76    for k in 0..n {
77        let phi_k = eigenvectors.column(k);
78        result += &phi_k.scale((-eigenvalues[k] * t).exp()) * &phi_k.transpose();
79    }
80    result
81}
82
83pub fn laplacian_1d_matrix(n: usize, dx: f64) -> DMatrix<f64> {
84    let mut l = DMatrix::zeros(n, n);
85    let inv_dx2 = 1.0 / (dx * dx);
86    for i in 0..n {
87        l[(i, i)] = -2.0 * inv_dx2;
88        if i > 0 { l[(i, i - 1)] = 1.0 * inv_dx2; }
89        if i < n - 1 { l[(i, i + 1)] = 1.0 * inv_dx2; }
90    }
91    l
92}
93
94pub fn graph_laplacian(adjacency: &DMatrix<f64>) -> DMatrix<f64> {
95    let n = adjacency.nrows();
96    let degree: DVector<f64> = DVector::from_iterator(n, (0..n).map(|i| adjacency.row(i).sum()));
97    DMatrix::from_diagonal(&degree) - adjacency
98}
99
100pub fn normalized_graph_laplacian(adjacency: &DMatrix<f64>) -> DMatrix<f64> {
101    let n = adjacency.nrows();
102    let degree: DVector<f64> = DVector::from_iterator(n, (0..n).map(|i| adjacency.row(i).sum()));
103    let d_inv_sqrt: DVector<f64> = degree.map(|d| if d > 1e-15 { 1.0 / d.sqrt() } else { 0.0 });
104    let d_matrix = DMatrix::from_diagonal(&d_inv_sqrt);
105    DMatrix::identity(n, n) - &d_matrix * adjacency * &d_matrix
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111    use approx::assert_relative_eq;
112
113    #[test]
114    fn test_laplacian_1d() {
115        let l = laplacian_1d_matrix(3, 1.0);
116        assert_relative_eq!(l[(0, 0)], -2.0);
117        assert_relative_eq!(l[(0, 1)], 1.0);
118        assert_relative_eq!(l[(1, 0)], 1.0);
119    }
120
121    #[test]
122    fn test_graph_laplacian() {
123        let adj = DMatrix::from_row_slice(3, 3, &[0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0]);
124        let l = graph_laplacian(&adj);
125        assert_relative_eq!(l[(0, 0)], 2.0);
126        assert_relative_eq!(l[(0, 1)], -1.0);
127    }
128
129    #[test]
130    fn test_normalized_graph_laplacian() {
131        let adj = DMatrix::from_row_slice(3, 3, &[0.0, 1.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0]);
132        let l_norm = normalized_graph_laplacian(&adj);
133        assert_relative_eq!(l_norm[(0, 0)], 1.0, epsilon = 1e-10);
134    }
135
136    #[test]
137    fn test_spectral_decomposition() {
138        let l = laplacian_1d_matrix(5, 1.0);
139        let decomp = SpectralDecomposition::decompose(&l).unwrap();
140        assert_eq!(decomp.eigenvalues.len(), 5);
141        assert!(decomp.eigenvalues.iter().all(|&e| e <= 1e-6));
142    }
143
144    #[test]
145    fn test_spectral_reconstruction() {
146        let a = DMatrix::from_row_slice(3, 3, &[2.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 2.0]);
147        let decomp = SpectralDecomposition::decompose(&a).unwrap();
148        let reconstructed = decomp.reconstruct();
149        for i in 0..3 {
150            for j in 0..3 {
151                assert_relative_eq!(reconstructed[(i, j)], a[(i, j)], epsilon = 1e-8);
152            }
153        }
154    }
155
156    #[test]
157    fn test_low_rank_approximation() {
158        let a = DMatrix::from_row_slice(4, 4, &[
159            4.0, 1.0, 0.0, 0.0, 1.0, 3.0, 1.0, 0.0,
160            0.0, 1.0, 2.0, 1.0, 0.0, 0.0, 1.0, 1.0,
161        ]);
162        let decomp = SpectralDecomposition::decompose(&a).unwrap();
163        let lr = decomp.low_rank(2);
164        assert_eq!(lr.nrows(), 4);
165    }
166
167    #[test]
168    fn test_heat_kernel_spectral_decay() {
169        let l = laplacian_1d_matrix(5, 1.0);
170        let decomp = SpectralDecomposition::decompose(&l).unwrap();
171        let k1 = heat_kernel_spectral(&decomp.eigenvalues, &decomp.eigenvectors, 0.1, 0, 0, 5);
172        let k10 = heat_kernel_spectral(&decomp.eigenvalues, &decomp.eigenvectors, 1.0, 0, 0, 5);
173        // For negative semidefinite Laplacian, eigenvalues <= 0, so exp(-lambda*t) decays as t grows
174        // At i=j, both should be positive
175        assert!(k1 > 0.0, "k1={}", k1);
176        assert!(k10 > 0.0, "k10={}", k10);
177    }
178
179    #[test]
180    fn test_spectral_diffusion_solve() {
181        let l = laplacian_1d_matrix(5, 1.0);
182        let decomp = SpectralDecomposition::decompose(&l).unwrap();
183        let u0 = DVector::from_vec(vec![1.0, 0.0, 0.0, 0.0, 0.0]);
184        let ut = spectral_diffusion_solve(&decomp.eigenvalues, &decomp.eigenvectors, &u0, 0.1);
185        assert_eq!(ut.len(), 5);
186    }
187
188    #[test]
189    fn test_spectral_gap() {
190        let eigenvalues = vec![0.0, 1.0, 3.0, 5.0];
191        let gap = spectral_gap(&eigenvalues).unwrap();
192        assert_relative_eq!(gap, 1.0);
193    }
194
195    #[test]
196    fn test_diffusion_kernel_matrix() {
197        let l = laplacian_1d_matrix(3, 1.0);
198        let decomp = SpectralDecomposition::decompose(&l).unwrap();
199        let k = diffusion_kernel_matrix(&decomp.eigenvalues, &decomp.eigenvectors, 0.1);
200        assert_eq!(k.nrows(), 3);
201        assert_eq!(k.ncols(), 3);
202    }
203}