Skip to main content

lau_diffusion_agents/
heat_kernel.rs

1//! Heat kernel on Riemannian manifolds, Gaussian diffusion, fundamental solutions.
2
3use nalgebra::{DVector, DMatrix};
4
5pub fn heat_kernel_euclidean(dim: usize, t: f64, x: &DVector<f64>, y: &DVector<f64>) -> f64 {
6    assert!(t > 0.0, "Time must be positive");
7    let diff = x - y;
8    let dist_sq = diff.iter().map(|v| v * v).sum::<f64>();
9    let norm = (4.0 * std::f64::consts::PI * t).powf(-(dim as f64) / 2.0);
10    norm * (-dist_sq / (4.0 * t)).exp()
11}
12
13pub fn heat_kernel_circle(t: f64, x: f64, y: f64, r: f64, n_terms: usize) -> f64 {
14    assert!(t > 0.0);
15    let mut sum = 0.0;
16    let circumference = 2.0 * std::f64::consts::PI * r;
17    for k in -(n_terms as i32)..=(n_terms as i32) {
18        let shift = k as f64 * circumference;
19        let d = x - y + shift;
20        sum += (-d * d / (4.0 * t)).exp();
21    }
22    sum / (4.0 * std::f64::consts::PI * t).sqrt()
23}
24
25pub fn heat_kernel_sphere(t: f64, theta1: f64, phi1: f64, theta2: f64, phi2: f64, r: f64, max_l: usize) -> f64 {
26    assert!(t > 0.0);
27    let cos_gamma = theta1.cos() * theta2.cos() + theta1.sin() * theta2.sin() * (phi1 - phi2).cos();
28    let cos_gamma = cos_gamma.clamp(-1.0, 1.0);
29    let mut sum = 0.0;
30    for l in 0..=max_l {
31        let cl = (2.0 * l as f64 + 1.0) / (4.0 * std::f64::consts::PI);
32        let pl = legendre_polynomial(l, cos_gamma);
33        sum += cl * pl * (-(l as f64) * (l as f64 + 1.0) * t / (r * r)).exp();
34    }
35    sum / (r * r)
36}
37
38pub fn legendre_polynomial(l: usize, x: f64) -> f64 {
39    match l {
40        0 => 1.0,
41        1 => x,
42        _ => {
43            let mut p_prev = 1.0;
44            let mut p_curr = x;
45            for n in 1..l {
46                let p_next = ((2 * n + 1) as f64 * x * p_curr - n as f64 * p_prev) / (n + 1) as f64;
47                p_prev = p_curr;
48                p_curr = p_next;
49            }
50            p_curr
51        }
52    }
53}
54
55pub fn gaussian_kernel(x: &DVector<f64>, mu: &DVector<f64>, sigma_inv: &DMatrix<f64>, det_sigma: f64) -> f64 {
56    let diff = x - mu;
57    let dim = x.nrows() as f64;
58    let exponent: f64 = -0.5 * (sigma_inv * &diff).dot(&diff);
59    let norm = 1.0 / ((2.0 * std::f64::consts::PI).powf(dim / 2.0) * det_sigma.sqrt());
60    norm * exponent.exp()
61}
62
63pub fn heat_equation_1d(u0: &[f64], alpha: f64, dx: f64, dt: f64, n_steps: usize) -> Vec<f64> {
64    let n = u0.len();
65    let r = alpha * dt / (dx * dx);
66    assert!(r < 0.5, "CFL condition violated: r={} >= 0.5", r);
67    let mut u = u0.to_vec();
68    let mut u_new = u.clone();
69
70    for _ in 0..n_steps {
71        for i in 1..n - 1 {
72            u_new[i] = u[i] + r * (u[i + 1] - 2.0 * u[i] + u[i - 1]);
73        }
74        u_new[0] = u_new[1];
75        u_new[n - 1] = u_new[n - 2];
76        std::mem::swap(&mut u, &mut u_new);
77    }
78    u
79}
80
81pub fn graph_heat_kernel(laplacian: &DMatrix<f64>, t: f64, n_terms: usize) -> DMatrix<f64> {
82    let n = laplacian.nrows();
83    let mut result = DMatrix::identity(n, n);
84    let mut power = laplacian.clone();
85    let mut factorial = 1.0;
86
87    for k in 1..=n_terms {
88        factorial *= k as f64;
89        let coeff = (-t).powi(k as i32) / factorial;
90        result += &power.scale(coeff);
91        power = &power * laplacian;
92    }
93    result
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    use approx::assert_relative_eq;
100
101    #[test]
102    fn test_heat_kernel_euclidean_point() {
103        let x = DVector::from_vec(vec![0.0]);
104        let y = DVector::from_vec(vec![0.0]);
105        let k = heat_kernel_euclidean(1, 1.0, &x, &y);
106        assert_relative_eq!(k, 1.0 / (4.0 * std::f64::consts::PI).sqrt(), epsilon = 1e-10);
107    }
108
109    #[test]
110    fn test_heat_kernel_euclidean_2d() {
111        let x = DVector::from_vec(vec![0.0, 0.0]);
112        let y = DVector::from_vec(vec![0.0, 0.0]);
113        let k = heat_kernel_euclidean(2, 1.0, &x, &y);
114        assert_relative_eq!(k, 1.0 / (4.0 * std::f64::consts::PI), epsilon = 1e-10);
115    }
116
117    #[test]
118    fn test_heat_kernel_symmetry() {
119        let x = DVector::from_vec(vec![1.0, 2.0]);
120        let y = DVector::from_vec(vec![3.0, 4.0]);
121        let k1 = heat_kernel_euclidean(2, 1.0, &x, &y);
122        let k2 = heat_kernel_euclidean(2, 1.0, &y, &x);
123        assert_relative_eq!(k1, k2, epsilon = 1e-12);
124    }
125
126    #[test]
127    fn test_heat_kernel_decays_with_distance() {
128        let origin = DVector::from_vec(vec![0.0]);
129        let near = DVector::from_vec(vec![1.0]);
130        let far = DVector::from_vec(vec![5.0]);
131        let k_near = heat_kernel_euclidean(1, 1.0, &origin, &near);
132        let k_far = heat_kernel_euclidean(1, 1.0, &origin, &far);
133        assert!(k_near > k_far);
134    }
135
136    #[test]
137    fn test_legendre_polynomials() {
138        assert_relative_eq!(legendre_polynomial(0, 0.5), 1.0);
139        assert_relative_eq!(legendre_polynomial(1, 0.5), 0.5);
140        assert_relative_eq!(legendre_polynomial(2, 0.5), -0.125, epsilon = 1e-10);
141    }
142
143    #[test]
144    fn test_heat_equation_1d_decay() {
145        let n = 50;
146        let u0: Vec<f64> = (0..n).map(|i| {
147            let x = i as f64 / (n - 1) as f64;
148            (-((x - 0.5) * 10.0).powi(2)).exp()
149        }).collect();
150        let result = heat_equation_1d(&u0, 0.1, 0.02, 0.0001, 100);
151        let max_val = result.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
152        let initial_max = u0.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
153        assert!(max_val < initial_max);
154    }
155
156    #[test]
157    fn test_heat_equation_1d_preserves_mass() {
158        let n = 20;
159        let u0: Vec<f64> = (0..n).map(|i| {
160            let x = i as f64 / (n - 1) as f64;
161            (-((x - 0.5) * 8.0).powi(2)).exp()
162        }).collect();
163        let initial_mass: f64 = u0.iter().sum();
164        let result = heat_equation_1d(&u0, 0.01, 0.1, 0.001, 10);
165        let final_mass: f64 = result.iter().sum();
166        assert_relative_eq!(initial_mass, final_mass, epsilon = 0.5);
167    }
168
169    #[test]
170    fn test_gaussian_kernel_normalized() {
171        let mu = DVector::from_vec(vec![0.0, 0.0]);
172        let sigma = DMatrix::identity(2, 2);
173        let sigma_inv = sigma.clone();
174        let mut sum = 0.0;
175        let step = 0.5;
176        for x in (-40..=40).map(|i| i as f64 * step) {
177            for y in (-40..=40).map(|i| i as f64 * step) {
178                let p = DVector::from_vec(vec![x, y]);
179                sum += gaussian_kernel(&p, &mu, &sigma_inv, 1.0) * step * step;
180            }
181        }
182        assert_relative_eq!(sum, 1.0, epsilon = 0.1);
183    }
184
185    #[test]
186    fn test_heat_kernel_circle_positive() {
187        let k = heat_kernel_circle(0.1, 0.0, 1.0, 1.0, 5);
188        assert!(k > 0.0);
189    }
190
191    #[test]
192    fn test_heat_kernel_sphere_positive() {
193        let k = heat_kernel_sphere(0.1, 0.0, 0.0, 0.5, 0.5, 1.0, 5);
194        assert!(k > 0.0);
195    }
196}