qsm_core/inversion/
l1qsm.rs

1//! L1-QSM: nonlinear L1 data-fidelity dipole inversion with TV regularization
2//!
3//! Solves the nonlinear QSM inverse problem with an L1 data-fidelity term
4//! (robust to phase outliers / noise) and a total-variation prior:
5//!
6//! min_x  || W .* ( exp(i * D x) - exp(i * phi) ) ||_1  +  alpha1 * ||grad(x)||_1
7//!
8//! where `D` is the dipole operator, `phi` is the (scaled) local field phase,
9//! and `W = lambda * mask`. The nonlinear fidelity is linearized around the
10//! current auxiliary phase and solved with an inner complex-argument Newton
11//! iteration, all embedded in an ADMM splitting.
12//!
13//! Reference:
14//! Milovic, C., Tejos, C., Acosta-Cabronero, J., et al. (2022).
15//! "The 2016 QSM Challenge: L1-QSM and PI-QSM — comparison of nonlinear
16//! L1 and phase-integral data-fidelity QSM reconstructions."
17//! Magnetic Resonance in Medicine (MRM), 2022.
18//!
19//! Ported from the FANSI toolbox `nlL1TV.m`
20//! (https://gitlab.com/cmilovic/FANSI-toolbox).
21
22use crate::inversion::admm::prepare_fansi_spectral;
23use crate::utils::gradient::{bdiv_inplace, fgrad_inplace};
24use crate::utils::{apply_mask_zero, shrink};
25use crate::Grid;
26use num_complex::Complex64;
27
28/// L1-QSM algorithm parameters.
29#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
30#[derive(Clone, Debug)]
31pub struct L1QsmParams {
32    /// Gradient (TV) L1 penalty weight.
33    pub alpha1: f64,
34    /// Gradient-consistency ADMM weight (typically 100 * alpha1).
35    pub mu1: f64,
36    /// Fidelity-consistency ADMM weight.
37    pub mu2: f64,
38    /// L1 proximal ADMM weight.
39    pub mu3: f64,
40    /// L1 fidelity strength; effective weight is `lambda * mask`.
41    pub lambda: f64,
42    /// Number of outer ADMM iterations.
43    pub max_iter: usize,
44    /// Percent-update convergence stopping tolerance.
45    pub tol_update: f64,
46    /// Inner Newton convergence tolerance.
47    pub tol_delta: f64,
48    /// ppm -> working (phase) scale applied to the input local field.
49    pub phase_scale: f64,
50}
51
52impl Default for L1QsmParams {
53    fn default() -> Self {
54        Self {
55            alpha1: 2e-4,
56            mu1: 2e-2,
57            mu2: 1.0,
58            mu3: 1.0,
59            lambda: 1.0,
60            max_iter: 50,
61            tol_update: 1.0,
62            tol_delta: 1e-6,
63            phase_scale: 1.0,
64        }
65    }
66}
67
68/// L2 norm of a real slice.
69fn norm2(v: &[f64]) -> f64 {
70    v.iter().map(|&a| a * a).sum::<f64>().sqrt()
71}
72
73/// L1-QSM nonlinear dipole inversion.
74///
75/// # Arguments
76/// * `local_field` - Local field values, ppm-scale (nx * ny * nz).
77/// * `mask` - Binary mask (nx * ny * nz), non-zero = inside ROI.
78/// * `grid` - Volume grid (dimensions and voxel sizes).
79/// * `bdir` - B0 field direction.
80/// * `params` - L1-QSM parameters.
81/// * `progress` - Progress callback `(iteration, max_iter)`.
82///
83/// # Returns
84/// Estimated susceptibility map (ppm-scale), masked to the ROI.
85pub fn l1qsm(
86    local_field: &[f64],
87    mask: &[u8],
88    grid: &Grid,
89    bdir: (f64, f64, f64),
90    params: &L1QsmParams,
91    mut progress: impl FnMut(usize, usize),
92) -> Vec<f64> {
93    let n = grid.n_total();
94
95    let (mut fft_ws, k, ee2) = prepare_fansi_spectral(grid, bdir);
96
97    // Scaled phase and L1 fidelity weight W = lambda * mask.
98    let phase: Vec<f64> = local_field.iter().map(|&f| f * params.phase_scale).collect();
99    let w: Vec<f64> = mask
100        .iter()
101        .map(|&m| if m != 0 { params.lambda } else { 0.0 })
102        .collect();
103
104    // IS = exp(1i * phase)
105    let is: Vec<Complex64> = phase
106        .iter()
107        .map(|&p| Complex64::new(p.cos(), p.sin()))
108        .collect();
109
110    let mu1 = params.mu1;
111    let mu2 = params.mu2;
112    let mu3 = params.mu3;
113    let alpha_over_mu = params.alpha1 / mu1;
114
115    // ADMM variables.
116    let mut x = vec![0.0f64; n];
117    let mut x_prev = vec![0.0f64; n];
118
119    // Gradient-consistency split (real).
120    let mut z_dx = vec![0.0f64; n];
121    let mut z_dy = vec![0.0f64; n];
122    let mut z_dz = vec![0.0f64; n];
123    let mut s_dx = vec![0.0f64; n];
124    let mut s_dy = vec![0.0f64; n];
125    let mut s_dz = vec![0.0f64; n];
126
127    // Fidelity-consistency split (real phase auxiliary) and its multiplier.
128    // isPrecond=true: z2 = W .* phase / max(W)
129    let w_max = w.iter().cloned().fold(0.0f64, f64::max);
130    let mut z2 = vec![0.0f64; n];
131    if w_max > 0.0 {
132        for i in 0..n {
133            z2[i] = w[i] * phase[i] / w_max;
134        }
135    }
136    let mut s2 = vec![0.0f64; n];
137
138    // L1 proximal split (complex) and its multiplier.
139    let mut z3 = vec![Complex64::new(0.0, 0.0); n];
140    let mut s3 = vec![Complex64::new(0.0, 0.0); n];
141
142    // Reusable buffers.
143    let mut fdiv = vec![Complex64::new(0.0, 0.0); n];
144    let mut fd2 = vec![Complex64::new(0.0, 0.0); n];
145    let mut xhat = vec![Complex64::new(0.0, 0.0); n];
146    let mut fx = vec![Complex64::new(0.0, 0.0); n];
147
148    let mut gxc = vec![0.0f64; n];
149    let mut gyc = vec![0.0f64; n];
150    let mut gzc = vec![0.0f64; n];
151    let mut x_dx = vec![0.0f64; n];
152    let mut x_dy = vec![0.0f64; n];
153    let mut x_dz = vec![0.0f64; n];
154    let mut div = vec![0.0f64; n];
155    let mut dx = vec![0.0f64; n]; // Dx = real(ifft(k .* fft(x)))
156    let mut rhs_z2 = vec![0.0f64; n];
157    let mut diff = vec![0.0f64; n];
158
159    for t in 0..params.max_iter {
160        progress(t + 1, params.max_iter);
161
162        // ---- x-subproblem (gradient consistency) --------------------------
163        // gradient side: mu1 * sum(E_t .* fft(z_d - s_d)) == mu1 * fft(bdiv(z_d - s_d))
164        for i in 0..n {
165            gxc[i] = z_dx[i] - s_dx[i];
166            gyc[i] = z_dy[i] - s_dy[i];
167            gzc[i] = z_dz[i] - s_dz[i];
168        }
169        bdiv_inplace(&mut div, &gxc, &gyc, &gzc, grid);
170        for i in 0..n {
171            fdiv[i] = Complex64::new(div[i], 0.0);
172        }
173        fft_ws.fft3d(&mut fdiv);
174
175        // fidelity side: mu2 * conj(K) .* fft(z2 - s2); K real so conj(K) = K,
176        // applied after the FFT below.
177        for i in 0..n {
178            fd2[i] = Complex64::new(z2[i] - s2[i], 0.0);
179        }
180        fft_ws.fft3d(&mut fd2);
181
182        // Combine in k-space.
183        for i in 0..n {
184            // Minus on the gradient term: adjoint of crate `fgrad` is `-bdiv` (matches
185            // QSM.rs TV-ADMM). `+bdiv` doubles the effective regularization. See fansi.rs.
186            let num = -mu1 * fdiv[i] + mu2 * k[i] * fd2[i];
187            // Guard the dipole null-space (DC/singular bins): both dipole kernel
188            // and Laplacian vanish there. Zero it instead of dividing FFT
189            // round-off by ~0, which would otherwise create a huge DC pedestal.
190            let den = mu2 * k[i] * k[i] + mu1 * ee2[i];
191            xhat[i] = if den > 1e-20 { num / den } else { Complex64::new(0.0, 0.0) };
192        }
193        fft_ws.ifft3d(&mut xhat);
194        x_prev.copy_from_slice(&x);
195        for i in 0..n {
196            x[i] = xhat[i].re;
197        }
198
199        // ---- convergence check (percent update) ---------------------------
200        let xnorm = norm2(&x);
201        if xnorm > 0.0 {
202            for i in 0..n {
203                diff[i] = x[i] - x_prev[i];
204            }
205            let x_update = 100.0 * norm2(&diff) / xnorm;
206            if x_update < params.tol_update || x_update.is_nan() {
207                progress(t + 1, t + 1);
208                break;
209            }
210        }
211
212        if t + 1 >= params.max_iter {
213            break;
214        }
215
216        // ---- gradient split update (TV shrink) ----------------------------
217        // Fx = fft(x)  (reused for the dipole rhs below)
218        for i in 0..n {
219            fx[i] = Complex64::new(x[i], 0.0);
220        }
221        fft_ws.fft3d(&mut fx);
222
223        // x_d{x,y,z} = fgrad(x)
224        fgrad_inplace(&mut x_dx, &mut x_dy, &mut x_dz, &x, grid);
225        for i in 0..n {
226            z_dx[i] = shrink(x_dx[i] + s_dx[i], alpha_over_mu);
227            z_dy[i] = shrink(x_dy[i] + s_dy[i], alpha_over_mu);
228            z_dz[i] = shrink(x_dz[i] + s_dz[i], alpha_over_mu);
229            s_dx[i] += x_dx[i] - z_dx[i];
230            s_dy[i] += x_dy[i] - z_dy[i];
231            s_dz[i] += x_dz[i] - z_dz[i];
232        }
233
234        // ---- L1 fidelity proximal (z3) ------------------------------------
235        // Y3 = exp(1i*z2) - IS + s3 ;  z3 = soft_threshold_complex(Y3, W/mu3)
236        for i in 0..n {
237            let ez2 = Complex64::new(z2[i].cos(), z2[i].sin());
238            let y3 = ez2 - is[i] + s3[i];
239            let mag = y3.norm();
240            let thr = w[i] / (mu3 + f64::EPSILON);
241            let shr = (mag - thr).max(0.0);
242            z3[i] = if mag > 0.0 {
243                y3 * (shr / mag)
244            } else {
245                Complex64::new(0.0, 0.0)
246            };
247        }
248
249        // ---- fidelity auxiliary (z2) via complex-arg Newton ---------------
250        // rhs_z2 = mu2 * (Dx + s2), with Dx = real(ifft(K .* Fx))
251        for i in 0..n {
252            xhat[i] = fx[i] * k[i];
253        }
254        fft_ws.ifft3d(&mut xhat);
255        for i in 0..n {
256            dx[i] = xhat[i].re;
257            rhs_z2[i] = mu2 * (dx[i] + s2[i]);
258            // init z2 = rhs_z2 / mu2 (== Dx + s2)
259            z2[i] = rhs_z2[i] / mu2;
260        }
261
262        // Newton per-voxel: minimize mu3*|exp(iz2)-(IS+z3-s3)| + mu2/2*(z2-...)^2
263        // Uses complex-argument trigonometric linearization (see module notes).
264        // yphase = angle(IS+z3-s3), ym = abs(IS+z3-s3), b = ln(ym).
265        let mut yphase = vec![0.0f64; n];
266        let mut cosh_b = vec![0.0f64; n];
267        let mut sinh_b = vec![0.0f64; n];
268        for i in 0..n {
269            let yc = is[i] + z3[i] - s3[i];
270            yphase[i] = yc.arg();
271            let m = yc.norm();
272            if m > 0.0 {
273                cosh_b[i] = 0.5 * (m + 1.0 / m); // cosh(ln m)
274                sinh_b[i] = 0.5 * (m - 1.0 / m); // sinh(ln m)
275            } else {
276                cosh_b[i] = 1.0;
277                sinh_b[i] = 0.0;
278            }
279        }
280
281        let mut delta = f64::INFINITY;
282        let mut inn = 0usize;
283        let mut update = vec![0.0f64; n];
284        while delta > params.tol_delta && inn < 4 {
285            inn += 1;
286            let norm_old = norm2(&z2);
287
288            for i in 0..n {
289                let a = z2[i] - yphase[i];
290                let (sa, ca) = a.sin_cos();
291                // sin(a - i b) = sin(a)cosh(b) - i cos(a)sinh(b)
292                let sin_arg = Complex64::new(sa * cosh_b[i], -ca * sinh_b[i]);
293                // cos(a - i b) = cos(a)cosh(b) + i sin(a)sinh(b)
294                let cos_arg = Complex64::new(ca * cosh_b[i], sa * sinh_b[i]);
295
296                let temp = mu3 * cos_arg + Complex64::new(mu2 + f64::EPSILON, 0.0);
297                let numer = mu3 * sin_arg + Complex64::new(mu2 * z2[i] - rhs_z2[i], 0.0);
298
299                // max(abs(temp), 0.05) .* sign(temp): floor magnitude at 0.05.
300                let tm = temp.norm();
301                let denom = if tm < 0.05 {
302                    if tm > 0.0 {
303                        temp * (0.05 / tm)
304                    } else {
305                        Complex64::new(0.05, 0.0)
306                    }
307                } else {
308                    temp
309                };
310
311                update[i] = (numer / denom).re;
312                z2[i] -= update[i];
313            }
314
315            let delta_new = if norm_old > 0.0 {
316                norm2(&update) / norm_old
317            } else {
318                0.0
319            };
320            if delta_new > delta {
321                break;
322            }
323            delta = delta_new;
324        }
325
326        // ---- multiplier updates -------------------------------------------
327        // s2 = s2 + Dx - z2
328        for i in 0..n {
329            s2[i] += dx[i] - z2[i];
330        }
331        // s3 = exp(1i*z2) - IS + s3 - z3
332        for i in 0..n {
333            let ez2 = Complex64::new(z2[i].cos(), z2[i].sin());
334            s3[i] = ez2 - is[i] + s3[i] - z3[i];
335        }
336    }
337
338    // Undo working-scale so the output is ppm-scale.
339    if params.phase_scale != 1.0 {
340        for v in &mut x {
341            *v /= params.phase_scale;
342        }
343    }
344
345    apply_mask_zero(&mut x, mask);
346    x
347}
348
349#[cfg(test)]
350mod tests {
351    use super::*;
352
353    #[test]
354    fn test_l1qsm_zero_field() {
355        // Zero field should give (approximately) zero susceptibility.
356        let n = 8;
357        let field = vec![0.0; n * n * n];
358        let mask = vec![1u8; n * n * n];
359        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
360        let params = L1QsmParams {
361            max_iter: 10,
362            ..L1QsmParams::default()
363        };
364
365        let chi = l1qsm(&field, &mask, &grid, (0.0, 0.0, 1.0), &params, |_, _| {});
366
367        for &val in chi.iter() {
368            assert!(val.abs() < 1e-6, "Zero field should give ~zero chi, got {}", val);
369        }
370    }
371
372    #[test]
373    fn test_l1qsm_finite() {
374        // A small ramp field should produce all-finite output.
375        let n = 8;
376        let field: Vec<f64> = (0..n * n * n).map(|i| (i as f64) * 0.001).collect();
377        let mask = vec![1u8; n * n * n];
378        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
379        let params = L1QsmParams {
380            max_iter: 10,
381            ..L1QsmParams::default()
382        };
383
384        let chi = l1qsm(&field, &mask, &grid, (0.0, 0.0, 1.0), &params, |_, _| {});
385
386        for (i, &val) in chi.iter().enumerate() {
387            assert!(val.is_finite(), "Chi should be finite at index {}", i);
388        }
389    }
390}
qsm_core/inversion/l1qsm.rs

qsm_core/inversion/
l1qsm.rs
