1use nalgebra::{Matrix3, Vector3};
4use serde::{Serialize, Deserialize};
5
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct DiffusionTensor {
8 pub tensor: Matrix3<f64>,
9}
10
11impl DiffusionTensor {
12 pub fn isotropic(scalar: f64) -> Self {
13 Self { tensor: Matrix3::identity().scale(scalar) }
14 }
15
16 pub fn from_eigen(eigenvalues: &Vector3<f64>, eigenvectors: &Matrix3<f64>) -> Self {
17 let d = Matrix3::from_diagonal(eigenvalues);
18 Self { tensor: eigenvectors * d * eigenvectors.transpose() }
19 }
20
21 pub fn axis_aligned(dx: f64, dy: f64, dz: f64) -> Self {
22 Self { tensor: Matrix3::from_diagonal(&Vector3::new(dx, dy, dz)) }
23 }
24
25 pub fn fractional_anisotropy(&self) -> f64 {
26 let eig = self.tensor.symmetric_eigen();
27 let vals = eig.eigenvalues;
28 let mean = (vals[0] + vals[1] + vals[2]) / 3.0;
29 let numerator = (vals[0] - mean).powi(2) + (vals[1] - mean).powi(2) + (vals[2] - mean).powi(2);
30 let denominator = vals[0].powi(2) + vals[1].powi(2) + vals[2].powi(2);
31 if denominator < 1e-15 { return 0.0; }
32 (3.0 / 2.0 * numerator / denominator).sqrt()
33 }
34
35 pub fn mean_diffusivity(&self) -> f64 {
36 (self.tensor[(0, 0)] + self.tensor[(1, 1)] + self.tensor[(2, 2)]) / 3.0
37 }
38
39 pub fn radial_diffusivity(&self) -> f64 {
40 let eig = self.tensor.symmetric_eigen();
41 let vals = eig.eigenvalues;
42 let max_idx = if vals[0] >= vals[1] && vals[0] >= vals[2] { 0 }
43 else if vals[1] >= vals[2] { 1 } else { 2 };
44 let mut sum = 0.0;
45 let mut count = 0;
46 for i in 0..3 { if i != max_idx { sum += vals[i]; count += 1; } }
47 if count > 0 { sum / count as f64 } else { 0.0 }
48 }
49
50 pub fn apply(&self, gradient: &Vector3<f64>) -> Vector3<f64> {
51 &self.tensor * gradient
52 }
53
54 pub fn trace(&self) -> f64 {
55 self.tensor[(0, 0)] + self.tensor[(1, 1)] + self.tensor[(2, 2)]
56 }
57}
58
59pub struct AnisotropicDiffusion2D {
60 pub n_x: usize,
61 pub n_y: usize,
62 pub dx: f64,
63 pub dy: f64,
64}
65
66impl AnisotropicDiffusion2D {
67 pub fn new(n_x: usize, n_y: usize, dx: f64, dy: f64) -> Self { Self { n_x, n_y, dx, dy } }
68
69 pub fn step_perona_malik(&self, u: &[f64], dt: f64, kappa: f64) -> Vec<f64> {
70 let nx = self.n_x; let ny = self.n_y;
71 let mut u_new = u.to_vec();
72 for i in 1..nx - 1 {
73 for j in 1..ny - 1 {
74 let idx = i * ny + j;
75 let grad_x = (u[idx + ny] - u[idx - ny]) / (2.0 * self.dx);
76 let grad_y = (u[idx + 1] - u[idx - 1]) / (2.0 * self.dy);
77 let grad_sq = grad_x * grad_x + grad_y * grad_y;
78 let g = 1.0 / (1.0 + grad_sq / (kappa * kappa));
79 let laplacian_x = (u[idx + ny] - 2.0 * u[idx] + u[idx - ny]) / (self.dx * self.dx);
80 let laplacian_y = (u[idx + 1] - 2.0 * u[idx] + u[idx - 1]) / (self.dy * self.dy);
81 u_new[idx] = u[idx] + dt * g * (laplacian_x + laplacian_y);
82 }
83 }
84 u_new
85 }
86
87 pub fn evolve_perona_malik(&self, u0: &[f64], dt: f64, kappa: f64, n_steps: usize) -> Vec<f64> {
88 let mut u = u0.to_vec();
89 for _ in 0..n_steps { u = self.step_perona_malik(&u, dt, kappa); }
90 u
91 }
92
93 pub fn step_tensor(&self, u: &[f64], dxx: &[f64], dxy: &[f64], dyy: &[f64], dt: f64) -> Vec<f64> {
94 let nx = self.n_x; let ny = self.n_y;
95 let mut u_new = u.to_vec();
96 for i in 1..nx - 1 {
97 for j in 1..ny - 1 {
98 let idx = i * ny + j;
99 let ux = (u[idx + ny] - u[idx - ny]) / (2.0 * self.dx);
100 let uy = (u[idx + 1] - u[idx - 1]) / (2.0 * self.dy);
101 let div_x = (dxx[idx + ny] * ux + dxy[idx + ny] * uy - dxx[idx - ny] * ux - dxy[idx - ny] * uy) / (2.0 * self.dx);
102 let div_y = (dxy[idx + 1] * ux + dyy[idx + 1] * uy - dxy[idx - 1] * ux - dyy[idx - 1] * uy) / (2.0 * self.dy);
103 u_new[idx] = u[idx] + dt * (div_x + div_y);
104 }
105 }
106 u_new
107 }
108}
109
110pub fn structure_tensor(u: &[f64], nx: usize, ny: usize, dx: f64, dy: f64) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
111 let mut dxx = vec![0.0; nx * ny];
112 let mut dxy = vec![0.0; nx * ny];
113 let mut dyy = vec![0.0; nx * ny];
114 for i in 1..nx - 1 {
115 for j in 1..ny - 1 {
116 let idx = i * ny + j;
117 let ux = (u[idx + ny] - u[idx - ny]) / (2.0 * dx);
118 let uy = (u[idx + 1] - u[idx - 1]) / (2.0 * dy);
119 dxx[idx] = ux * ux;
120 dxy[idx] = ux * uy;
121 dyy[idx] = uy * uy;
122 }
123 }
124 (dxx, dxy, dyy)
125}
126
127#[cfg(test)]
128mod tests {
129 use super::*;
130 use approx::assert_relative_eq;
131
132 #[test]
133 fn test_isotropic_tensor() {
134 let t = DiffusionTensor::isotropic(2.0);
135 assert_relative_eq!(t.tensor[(0, 0)], 2.0);
136 assert_relative_eq!(t.trace(), 6.0);
137 }
138
139 #[test]
140 fn test_axis_aligned_tensor() {
141 let t = DiffusionTensor::axis_aligned(1.0, 2.0, 3.0);
142 assert_relative_eq!(t.mean_diffusivity(), 2.0);
143 }
144
145 #[test]
146 fn test_fractional_anisotropy_isotropic() {
147 let t = DiffusionTensor::isotropic(1.0);
148 assert_relative_eq!(t.fractional_anisotropy(), 0.0, epsilon = 1e-10);
149 }
150
151 #[test]
152 fn test_fractional_anisotropy_anisotropic() {
153 let t = DiffusionTensor::axis_aligned(3.0, 1.0, 1.0);
154 let fa = t.fractional_anisotropy();
155 assert!(fa > 0.0 && fa <= 1.0, "FA={}", fa);
156 }
157
158 #[test]
159 fn test_diffusion_tensor_apply() {
160 let t = DiffusionTensor::isotropic(2.0);
161 let g = Vector3::new(1.0, 0.0, 0.0);
162 let result = t.apply(&g);
163 assert_relative_eq!(result[0], 2.0);
164 }
165
166 #[test]
167 fn test_radial_diffusivity() {
168 let t = DiffusionTensor::axis_aligned(3.0, 1.0, 1.0);
169 assert_relative_eq!(t.radial_diffusivity(), 1.0);
170 }
171
172 #[test]
173 fn test_perona_malik() {
174 let diff = AnisotropicDiffusion2D::new(10, 10, 0.1, 0.1);
175 let u0: Vec<f64> = (0..100).map(|i| {
176 let x = (i / 10) as f64 / 10.0;
177 (-((x - 0.5) * 10.0).powi(2)).exp()
178 }).collect();
179 let result = diff.evolve_perona_malik(&u0, 0.001, 0.1, 10);
180 assert!(result.iter().all(|&v| v.is_finite()));
181 }
182
183 #[test]
184 fn test_structure_tensor() {
185 let u: Vec<f64> = (0..100).map(|i| { let x = (i / 10) as f64; let y = (i % 10) as f64; x + y }).collect();
186 let (dxx, _, _) = structure_tensor(&u, 10, 10, 1.0, 1.0);
187 assert_eq!(dxx.len(), 100);
188 }
189
190 #[test]
191 fn test_tensor_step() {
192 let diff = AnisotropicDiffusion2D::new(10, 10, 0.1, 0.1);
193 let u0 = vec![0.5; 100];
194 let dxx = vec![1.0; 100]; let dxy = vec![0.0; 100]; let dyy = vec![1.0; 100];
195 let result = diff.step_tensor(&u0, &dxx, &dxy, &dyy, 0.001);
196 for &v in &result { assert_relative_eq!(v, 0.5, epsilon = 0.01); }
197 }
198
199 #[test]
200 fn test_mean_diffusivity() {
201 let t = DiffusionTensor::axis_aligned(1.0, 2.0, 3.0);
202 assert_relative_eq!(t.mean_diffusivity(), 2.0);
203 }
204}