qsm_core/inversion/
tikhonov.rs

1//! Tikhonov regularization for QSM
2//!
3//! Tikhonov regularization adds an L2 penalty term to stabilize the inversion:
4//!
5//! chi = argmin_x ||Dx - f||_2^2 + lambda||Gamma*x||_2^2
6//!
7//! This has a closed-form solution in k-space:
8//! chi_hat = D* * f_hat / (|D|^2 + lambda|Gamma|^2)
9//!
10//! Reference:
11//! Bilgic, B., et al. (2014). "Fast image reconstruction with L2-regularization."
12//! Journal of Magnetic Resonance Imaging, 40(1):181-191. https://doi.org/10.1002/jmri.24365
13//!
14//! Reference implementation: https://github.com/kamesy/QSM.jl
15
16use num_complex::Complex64;
17use crate::fft::{fft3d, ifft3d};
18use crate::kernels::dipole::dipole_kernel;
19use crate::kernels::laplacian::laplacian_kernel;
20use crate::utils::apply_mask_zero;
21use crate::Grid;
22
23/// Regularization type for Tikhonov
24#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
25#[derive(Clone, Copy, Debug)]
26pub enum Regularization {
27    /// Identity: lambda||x||_2^2
28    Identity,
29    /// Gradient: lambda||grad(x)||_2^2 (uses negative Laplacian)
30    Gradient,
31    /// Laplacian: lambda||lap(x)||_2^2
32    Laplacian,
33}
34
35/// Tikhonov algorithm parameters
36#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
37#[derive(Clone, Debug)]
38pub struct TikhonovParams {
39    /// Regularization weight
40    pub lambda: f64,
41    /// Regularization type
42    pub reg: Regularization,
43}
44
45impl Default for TikhonovParams {
46    fn default() -> Self {
47        Self { lambda: 0.01, reg: Regularization::Identity }
48    }
49}
50
51/// Tikhonov regularization for dipole inversion
52///
53/// # Arguments
54/// * `local_field` - Local field values
55/// * `mask` - Binary mask (1 = inside ROI, 0 = outside)
56/// * `grid` - Volume grid (dimensions and voxel sizes)
57/// * `bdir` - B0 field direction
58/// * `params` - Tikhonov parameters (lambda, regularization type)
59///
60/// # Returns
61/// Susceptibility map
62pub fn tikhonov(
63    local_field: &[f64],
64    mask: &[u8],
65    grid: &Grid,
66    bdir: (f64, f64, f64),
67    params: &TikhonovParams,
68) -> Vec<f64> {
69    let (nx, ny, nz) = grid.dims;
70    let n_total = grid.n_total();
71    let lambda = params.lambda;
72
73    // Generate dipole kernel
74    let d = dipole_kernel(grid, bdir);
75
76    // Generate regularization kernel and FFT it
77    let gamma: Vec<f64> = match params.reg {
78        Regularization::Identity => {
79            vec![1.0; n_total]
80        }
81        Regularization::Gradient => {
82            // Negative Laplacian for gradient regularization
83            let l = laplacian_kernel(grid, true);
84            // FFT to get frequency response
85            let mut l_complex: Vec<Complex64> = l.iter()
86                .map(|&x| Complex64::new(x, 0.0))
87                .collect();
88            fft3d(&mut l_complex, nx, ny, nz);
89            // Take real part (Laplacian FFT is real)
90            l_complex.iter().map(|c| c.re).collect()
91        }
92        Regularization::Laplacian => {
93            // Laplacian squared
94            let l = laplacian_kernel(grid, false);
95            let mut l_complex: Vec<Complex64> = l.iter()
96                .map(|&x| Complex64::new(x, 0.0))
97                .collect();
98            fft3d(&mut l_complex, nx, ny, nz);
99            // |Gamma|^2
100            l_complex.iter().map(|c| c.re * c.re).collect()
101        }
102    };
103
104    // Compute Tikhonov inverse: D / (D^2 + lambda*Gamma)
105    let inv_d: Vec<f64> = d.iter().zip(gamma.iter()).map(|(&dval, &gval)| {
106        let denom = dval * dval + lambda * gval;
107        if denom.abs() > 1e-20 {
108            dval / denom
109        } else {
110            0.0
111        }
112    }).collect();
113
114    // Convert local field to complex
115    let mut field_complex: Vec<Complex64> = local_field.iter()
116        .map(|&x| Complex64::new(x, 0.0))
117        .collect();
118
119    // FFT of local field
120    fft3d(&mut field_complex, nx, ny, nz);
121
122    // Multiply by Tikhonov inverse
123    for i in 0..n_total {
124        field_complex[i] *= inv_d[i];
125    }
126
127    // IFFT to get susceptibility
128    ifft3d(&mut field_complex, nx, ny, nz);
129
130    // Extract real part and apply mask
131    let mut chi: Vec<f64> = field_complex.iter()
132        .map(|c| c.re)
133        .collect();
134
135    apply_mask_zero(&mut chi, mask);
136
137    chi
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143
144    #[test]
145    fn test_tikhonov_zero_field() {
146        let n = 8;
147        let field = vec![0.0; n * n * n];
148        let mask = vec![1u8; n * n * n];
149        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
150
151        let chi = tikhonov(&field, &mask, &grid, (0.0, 0.0, 1.0), &TikhonovParams::default());
152
153        for val in chi.iter() {
154            assert!(val.abs() < 1e-10, "Zero field should give zero chi");
155        }
156    }
157
158    #[test]
159    fn test_tikhonov_finite() {
160        let n = 8;
161        let field: Vec<f64> = (0..n*n*n).map(|i| (i as f64) * 0.01).collect();
162        let mask = vec![1u8; n * n * n];
163        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
164
165        let chi = tikhonov(&field, &mask, &grid, (0.0, 0.0, 1.0), &TikhonovParams::default());
166
167        for (i, val) in chi.iter().enumerate() {
168            assert!(val.is_finite(), "Chi should be finite at index {}", i);
169        }
170    }
171
172    #[test]
173    fn test_tikhonov_regularization_types() {
174        let n = 8;
175        let field: Vec<f64> = (0..n*n*n).map(|i| ((i as f64) * 0.1).sin()).collect();
176        let mask = vec![1u8; n * n * n];
177        let bdir = (0.0, 0.0, 1.0);
178        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
179
180        let chi_id = tikhonov(&field, &mask, &grid, bdir,
181                             &TikhonovParams { lambda: 0.01, reg: Regularization::Identity });
182        let chi_grad = tikhonov(&field, &mask, &grid, bdir,
183                               &TikhonovParams { lambda: 0.01, reg: Regularization::Gradient });
184        let chi_lap = tikhonov(&field, &mask, &grid, bdir,
185                              &TikhonovParams { lambda: 0.01, reg: Regularization::Laplacian });
186
187        // All should be different
188        let diff_ig: f64 = chi_id.iter().zip(chi_grad.iter())
189            .map(|(a, b)| (a - b).abs()).sum();
190        let diff_gl: f64 = chi_grad.iter().zip(chi_lap.iter())
191            .map(|(a, b)| (a - b).abs()).sum();
192
193        assert!(diff_ig > 1e-10, "Identity and Gradient should differ");
194        assert!(diff_gl > 1e-10, "Gradient and Laplacian should differ");
195    }
196}
qsm_core/inversion/tikhonov.rs

qsm_core/inversion/
tikhonov.rs
