qsm_core/inversion/
nltv.rs

1//! Nonlinear Total Variation (NLTV) regularized dipole inversion
2//!
3//! NLTV extends standard TV by using iteratively reweighted minimization,
4//! which produces sharper edges and better preserves fine details.
5//!
6//! The method solves:
7//! min_x ||Dx - f||_2^2 + lambda * sum w_i |grad(x)|_i
8//!
9//! where weights w_i are iteratively updated based on the current solution.
10//!
11//! Reference:
12//! Kames, C., Wiggermann, V., Rauscher, A. (2018).
13//! "Rapid two-step dipole inversion for susceptibility mapping with sparsity priors."
14//! NeuroImage, 167:276-283. https://doi.org/10.1016/j.neuroimage.2017.11.018
15//!
16//! Reference implementation: https://github.com/kamesy/QSM.jl
17
18use crate::utils::gradient::fgrad_inplace;
19use crate::utils::{weighted_shrink, apply_mask_zero};
20use crate::Grid;
21use super::admm::{AdmmBuffers, admm_step, prepare_admm_spectral};
22
23/// NLTV algorithm parameters
24#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
25#[derive(Clone, Debug)]
26pub struct NltvParams {
27    /// Regularization parameter
28    pub lambda: f64,
29    /// Penalty parameter
30    pub mu: f64,
31    /// Convergence tolerance
32    pub tol: f64,
33    /// Maximum ADMM iterations
34    pub max_iter: usize,
35    /// Newton iterations for weight update
36    pub newton_iter: usize,
37}
38
39impl Default for NltvParams {
40    fn default() -> Self {
41        Self {
42            lambda: 1e-3,
43            mu: 1.0,
44            tol: 1e-3,
45            max_iter: 250,
46            newton_iter: 10,
47        }
48    }
49}
50
51/// NLTV dipole inversion using iteratively reweighted ADMM
52///
53/// # Arguments
54/// * `local_field` - Local field values (nx * ny * nz)
55/// * `mask` - Binary mask (nx * ny * nz), 1 = inside ROI
56/// * `grid` - Volume grid (dimensions and voxel sizes)
57/// * `bdir` - B0 field direction
58/// * `params` - NLTV parameters
59/// * `progress` - Progress callback `(iteration, total_iterations)`
60///
61/// # Returns
62/// Susceptibility map
63pub fn nltv(
64    local_field: &[f64],
65    mask: &[u8],
66    grid: &Grid,
67    bdir: (f64, f64, f64),
68    params: &NltvParams,
69    mut progress: impl FnMut(usize, usize),
70) -> Vec<f64> {
71    let n_total = grid.n_total();
72    let eps = 1e-6; // Small constant to avoid division by zero
73
74    // Compute rho adaptively (for ADMM)
75    let rho = 100.0 * params.lambda;
76
77    // Pre-compute spectral operators
78    let (mut fft_ws, inv_a, f_hat) = prepare_admm_spectral(local_field, grid, bdir, rho);
79
80    // Pre-allocate working buffers and run ADMM iterations
81    let mut buf = AdmmBuffers::new(n_total);
82    let mut weights = vec![1.0; n_total];
83
84    let total_iter = params.max_iter * params.newton_iter;
85    let mut current_iter = 0;
86
87    // Outer loop: Newton-like reweighting
88    for _newton in 0..params.newton_iter {
89        let lambda_over_rho = params.lambda / rho;
90
91        // Inner loop: ADMM with current weights
92        for _iter in 0..params.max_iter {
93            current_iter += 1;
94            progress(current_iter, total_iter);
95
96            let converged = admm_step(
97                &mut buf, &mut fft_ws, &f_hat, &inv_a, rho, grid, params.tol,
98                |vx, vy, vz, i| (
99                    weighted_shrink(vx, lambda_over_rho, weights[i]),
100                    weighted_shrink(vy, lambda_over_rho, weights[i]),
101                    weighted_shrink(vz, lambda_over_rho, weights[i]),
102                ),
103            );
104
105            if converged {
106                break;
107            }
108        }
109
110        // Update weights based on current gradient magnitude
111        fgrad_inplace(&mut buf.gx, &mut buf.gy, &mut buf.gz, &buf.x, grid);
112
113        for i in 0..n_total {
114            let grad_mag = (buf.gx[i] * buf.gx[i] + buf.gy[i] * buf.gy[i] + buf.gz[i] * buf.gz[i]).sqrt();
115            weights[i] = 1.0 / (grad_mag + params.mu * eps);
116        }
117
118        // Normalize weights to prevent explosion
119        let max_weight: f64 = weights.iter().cloned().fold(0.0, f64::max);
120        if max_weight > 1.0 {
121            for w in weights.iter_mut() {
122                *w /= max_weight;
123            }
124        }
125    }
126
127    apply_mask_zero(&mut buf.x, mask);
128
129    buf.x
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135    use crate::utils::weighted_shrink;
136
137    #[test]
138    fn test_nltv_zero_field() {
139        let n = 8;
140        let field = vec![0.0; n * n * n];
141        let mask = vec![1u8; n * n * n];
142        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
143        let params = NltvParams { lambda: 1e-3, mu: 1.0, tol: 1e-2, max_iter: 10, newton_iter: 2 };
144
145        let chi = nltv(&field, &mask, &grid, (0.0, 0.0, 1.0), &params, |_, _| {});
146
147        for &val in chi.iter() {
148            assert!(val.abs() < 1e-8, "Zero field should give zero chi");
149        }
150    }
151
152    #[test]
153    fn test_nltv_finite() {
154        let n = 8;
155        let field: Vec<f64> = (0..n*n*n).map(|i| (i as f64) * 0.001).collect();
156        let mask = vec![1u8; n * n * n];
157        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
158        let params = NltvParams { lambda: 1e-3, mu: 1.0, tol: 1e-2, max_iter: 10, newton_iter: 2 };
159
160        let chi = nltv(&field, &mask, &grid, (0.0, 0.0, 1.0), &params, |_, _| {});
161
162        for (i, &val) in chi.iter().enumerate() {
163            assert!(val.is_finite(), "Chi should be finite at index {}", i);
164        }
165    }
166
167    #[test]
168    fn test_weighted_shrink() {
169        // w=1 should behave like regular shrink
170        assert!((weighted_shrink(1.0, 0.5, 1.0) - 0.5).abs() < 1e-10);
171        assert!((weighted_shrink(-1.0, 0.5, 1.0) - (-0.5)).abs() < 1e-10);
172        assert!((weighted_shrink(0.3, 0.5, 1.0) - 0.0).abs() < 1e-10);
173
174        // w=0.5 should have half the threshold
175        assert!((weighted_shrink(1.0, 0.5, 0.5) - 0.75).abs() < 1e-10);
176        assert!((weighted_shrink(0.3, 0.5, 0.5) - 0.05).abs() < 1e-10);
177    }
178}
qsm_core/inversion/nltv.rs

qsm_core/inversion/
nltv.rs
