qsm_core/inversion/
ndi.rs

1//! Nonlinear Dipole Inversion (NDI) for QSM
2//!
3//! Solves the dipole inversion problem with a nonlinear data-consistency term,
4//! modelling the wrapped phase directly via `sin(Dx - phase)` and minimizing it
5//! with a simple L2-regularized gradient descent:
6//!
7//! min_x || W .* (exp(i*Dx) - exp(i*phase)) ||^2 + alpha ||x||^2
8//!
9//! which yields the update (weight W = mask here):
10//!
11//! x <- x - tau * D^H(W .* sin(Dx - phase)) - tau*alpha*x
12//!
13//! Because the dipole kernel `D` is real, `conj(D) = D` and the adjoint is just
14//! another forward dipole application.
15//!
16//! Reference:
17//! Polak, D., Chatnuntawech, I., Yoon, J., Iyer, S.S., Milovic, C., Lee, J.,
18//! Bachert, P., Adalsteinsson, E., Setsompop, K., Bilgic, B. (2020).
19//! "Nonlinear dipole inversion (NDI) enables robust quantitative susceptibility
20//! mapping (QSM)." NMR in Biomedicine, 33(12):e4271.
21//! https://doi.org/10.1002/nbm.4271
22//!
23//! Ported from FANSI's `ndi.m`
24//! (https://gitlab.com/cmilovic/FANSI-toolbox).
25//!
26//! # Note on `phase_scale`
27//! FANSI's NDI operates on the phase in radians. In this crate the input local
28//! field is provided at ppm-scale. With the default `phase_scale = 1.0` the
29//! solver runs at the native ppm scale; for small ppm values `sin(x) ~ x`, so
30//! the behaviour is numerically well-conditioned and scale-consistent with the
31//! rest of the crate. A user may set `phase_scale` to convert the input to
32//! radians (e.g. via the Hz->rad / ppm->rad factor for their acquisition) to
33//! recover the true nonlinear behaviour; the output is divided back by the same
34//! factor so the returned susceptibility stays at ppm-scale.
35
36use crate::fft::Fft3dWorkspace;
37use crate::utils::apply_mask_zero;
38use crate::Grid;
39use super::admm::prepare_fansi_spectral;
40use num_complex::Complex64;
41
42/// NDI algorithm parameters
43#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
44#[derive(Clone, Debug)]
45pub struct NdiParams {
46    /// Gradient-descent step size
47    pub tau: f64,
48    /// L2 regularization weight
49    pub alpha: f64,
50    /// Number of iterations
51    pub max_iter: usize,
52    /// ppm -> working-scale multiplier (see module docs)
53    pub phase_scale: f64,
54}
55
56impl Default for NdiParams {
57    fn default() -> Self {
58        Self {
59            tau: 2.0,
60            alpha: 1e-5,
61            max_iter: 200,
62            phase_scale: 1.0,
63        }
64    }
65}
66
67/// Apply the forward dipole operator `D * x` in-place buffers.
68///
69/// Equivalent to MATLAB `susc2field(kernel, x) = real(ifftn(D .* fftn(x)))`.
70/// Since the dipole kernel `k` is real, `conj(k) = k`, so the same routine
71/// computes the adjoint `D^H`.
72fn apply_dipole(
73    fft_ws: &mut Fft3dWorkspace,
74    k: &[f64],
75    x: &[f64],
76    buf: &mut [Complex64],
77    out: &mut [f64],
78) {
79    for i in 0..x.len() {
80        buf[i] = Complex64::new(x[i], 0.0);
81    }
82    fft_ws.fft3d(buf);
83    for i in 0..x.len() {
84        buf[i] = buf[i] * k[i];
85    }
86    fft_ws.ifft3d(buf);
87    for i in 0..x.len() {
88        out[i] = buf[i].re;
89    }
90}
91
92/// Nonlinear Dipole Inversion (NDI)
93///
94/// # Arguments
95/// * `local_field` - Local field values, ppm-scale (nx * ny * nz)
96/// * `mask` - Binary mask (nx * ny * nz), non-zero = inside ROI
97/// * `grid` - Volume grid (dimensions and voxel sizes)
98/// * `bdir` - B0 field direction
99/// * `params` - NDI parameters
100/// * `progress` - Progress callback `(iteration, max_iter)`
101///
102/// # Returns
103/// Estimated susceptibility map (ppm-scale), masked to the ROI.
104pub fn ndi(
105    local_field: &[f64],
106    mask: &[u8],
107    grid: &Grid,
108    bdir: (f64, f64, f64),
109    params: &NdiParams,
110    mut progress: impl FnMut(usize, usize),
111) -> Vec<f64> {
112    let n = grid.n_total();
113
114    let (mut fft_ws, k_kernel, _ee2) = prepare_fansi_spectral(grid, bdir);
115
116    // Scaled phase and weight (W = mask.*mask = mask for binary mask).
117    let phase: Vec<f64> = local_field.iter().map(|&f| f * params.phase_scale).collect();
118    let w: Vec<f64> = mask
119        .iter()
120        .map(|&m| if m != 0 { 1.0 } else { 0.0 })
121        .collect();
122
123    let mut x = vec![0.0f64; n];
124
125    // Reusable buffers.
126    let mut buf = vec![Complex64::new(0.0, 0.0); n];
127    let mut phix = vec![0.0f64; n];
128    let mut resid = vec![0.0f64; n];
129    let mut grad = vec![0.0f64; n];
130
131    for t in 0..params.max_iter {
132        progress(t + 1, params.max_iter);
133
134        // phix = D * x
135        apply_dipole(&mut fft_ws, &k_kernel, &x, &mut buf, &mut phix);
136
137        // resid = W .* sin(phix - phase)
138        for i in 0..n {
139            resid[i] = w[i] * (phix[i] - phase[i]).sin();
140        }
141
142        // grad = D^H * resid  (conj(kernel) = kernel, real)
143        apply_dipole(&mut fft_ws, &k_kernel, &resid, &mut buf, &mut grad);
144
145        // x <- x - tau*grad - tau*alpha*x
146        for i in 0..n {
147            x[i] -= params.tau * grad[i] + params.tau * params.alpha * x[i];
148        }
149    }
150
151    // Undo the working-scale so the output stays at ppm-scale.
152    if params.phase_scale != 1.0 {
153        for v in &mut x {
154            *v /= params.phase_scale;
155        }
156    }
157
158    apply_mask_zero(&mut x, mask);
159    x
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165
166    #[test]
167    fn test_ndi_zero_field() {
168        // Zero field should give zero susceptibility.
169        let n = 8;
170        let field = vec![0.0; n * n * n];
171        let mask = vec![1u8; n * n * n];
172        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
173        let params = NdiParams { max_iter: 20, ..Default::default() };
174
175        let chi = ndi(&field, &mask, &grid, (0.0, 0.0, 1.0), &params, |_, _| {});
176
177        for &val in chi.iter() {
178            assert!(val.abs() < 1e-6, "Zero field should give zero chi, got {}", val);
179        }
180    }
181
182    #[test]
183    fn test_ndi_finite() {
184        // Result should be finite for a small ramp field.
185        let n = 8;
186        let field: Vec<f64> = (0..n * n * n).map(|i| (i as f64) * 0.001).collect();
187        let mask = vec![1u8; n * n * n];
188        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
189        let params = NdiParams { max_iter: 20, ..Default::default() };
190
191        let chi = ndi(&field, &mask, &grid, (0.0, 0.0, 1.0), &params, |_, _| {});
192
193        for (i, &val) in chi.iter().enumerate() {
194            assert!(val.is_finite(), "Chi should be finite at index {}", i);
195        }
196    }
197}
qsm_core/inversion/ndi.rs

qsm_core/inversion/
ndi.rs
