qsm_core/inversion/
fansi.rs

1//! FANSI nonlinear TV / TGV dipole inversion.
2//!
3//! Nonlinear total-variation (nlTV) and nonlinear total-generalized-variation
4//! (nlTGV) QSM dipole inversion with a nonlinear (wrapped-phase) data-fidelity
5//! term, solved with ADMM plus an inner complex-argument Newton iteration.
6//!
7//! The nonlinear fidelity models the field data as `exp(i * D x)` and matches
8//! it to the wrapped local-field phase, which makes the reconstruction robust to
9//! phase-wrap / high-field regimes.
10//!
11//! References:
12//! - nlTV: Milovic, C., Bilgic, B., Zhao, B., et al. (2018).
13//!   "Fast nonlinear susceptibility inversion with variational regularization."
14//!   Magnetic Resonance in Medicine, 80(2):814-821.
15//!   <https://doi.org/10.1002/mrm.27073>
16//! - nlTGV: the total-generalized-variation variant of the above.
17//!
18//! Ported faithfully from the FANSI toolbox `nlTV.m` and `nlTGV.m`
19//! (<https://gitlab.com/cmilovic/FANSI-toolbox>).
20//!
21//! # Units and weights
22//! Like the other inversions, these take the local field in ppm and return ppm. nlTV.m and nlTGV.m
23//! work on the phase in radians, so the field is converted with `2*pi*gamma*b0*te` first and the
24//! result converted back at the end; left in ppm, `sin(x) ~ x` and the data term is not nonlinear.
25//! `b0` and `echo_times` are scan parameters that the pipeline fills from metadata, and
26//! [`FansiParams::te`] picks the echo time the conversion uses (by default the reference's choice;
27//! see there).
28//!
29//! The data weight is the magnitude, normalised as in FANSI's in vivo example
30//! (`min(2*mask*mag/max(mag), 1)`), and squared as nlTV.m/nlTGV.m square
31//! `params.weight`. The magnitude is required: with a flat (mask) weight, the Newton
32//! step's denominator `W*cos(z2 - phase) + mu2` reaches 0 where the phase nears +-pi
33//! and the solver diverges (as nlTV.m does).
34//!
35//! The finite differences are FANSI's, on unit spacing whatever the voxel size.
36//!
37//! # Reference pipeline
38//! FANSI-toolbox at commit 93044a3 (<https://gitlab.com/cmilovic/FANSI-toolbox>). Its only
39//! full pipeline is `script_DICOM_Philips.m` (3T Philips, 5 echoes, TE1 4.0 ms, spacing 3.7 ms,
40//! 0.73x0.73x1 mm, B0 along z; header L1-6), which says it is "NOT an optimized pipeline":
41//! - Mask: magnitude (M0) threshold plus an R2* range, then morphology, ending with a 2-voxel
42//!   erosion "to avoid data corruption in the cortical areas" (L125-170); L126 "We encourage to
43//!   use dedicated tools as BET2".
44//! - Unwrapping: SEGUE per echo (L190-194); L183-187 also recommend ROMEO, and L196-197 say
45//!   "Avoid Laplacian-based unwrapping for multiecho processing, as they remove harmonic
46//!   components of the phase".
47//! - Echo combination: a nonlinear multi-echo fit (`nlme_L2gradient`, weight `mag*TE`,
48//!   L259-276), which also gives the noise map `cnoise`.
49//! - Background removal: FANSI's own multiscale SMV, `msmv_noextrapolation(..., 5)` (L335);
50//!   L290 allows LBV, PDF or V-SHARP instead. (This MSMV is not the mSMV of
51//!   `bg_removal.msmv_refine`.)
52//! - Input: `local*deltaTE`, `deltaTE = TE(4) - TE(2)` with the comment "You may use also any
53//!   individual echo time" (L364, L370). Weight `m0.*mask./cnoise`, divided by its maximum
54//!   (L371-372). nlTGV is run with the same parameters as wTV (L421).
55//! - alpha1: no default in `nlTV.m`/`nlTGV.m` (required field); the toolbox's demos use 2e-4
56//!   (`script_fansi_phantom.m` L196, `script_qsmchallenge.m` L73) and 3e-4 with mu1 = 1e-2 for
57//!   nlTV on the 2016 challenge data (`script_qsmchallenge.m` L111-112), which is the
58//!   challenge's own LBV local field and BET mask eroded by 5 voxels (L22-24). The default
59//!   here is that 2e-4; the in vivo pipeline picks alpha1 by L-curve / frequency analysis
60//!   (L374-410).
61//! - Default weight in `nlTV.m`/`nlTGV.m` is ones (`defaultW = ones(N)`, with a warning).
62//!   Every demo call of `nlTV`/`nlTGV` passes a magnitude-derived weight
63//!   (`script_fansi_phantom.m` L199 "Recommended for nonlinear algorithms",
64//!   `script_cosmos_phantom.m` L88, `script_qsmchallenge.m` L106-108); the mask-only weights in
65//!   the demos go to the linear `wTV`/`wTGV`.
66//!
67//! # Parity
68//! Against the verified original (QSM-CI pair `fansi-nltv`, `fansi-nltgv`; FANSI-toolbox 93044a3, identical to
69//! licensed MATLAB R2023b), final-map χ NRMSE at both sides' defaults on QSM-CI's simulated head
70//! and the KKI travelling head (QSM-CI parity run 41403629314, QSM.rs#399–#403): nlTV sim 0.0069 %, KKI 0.0030 %; nlTGV sim 0.0063 %, KKI 0.0039 %, from
71//! nlTV 11.0 % / 9.5 %, nlTGV 11.0 % / 9.6 % (sim / KKI) before. The kernel was the whole gap. What remains is
72//! single (MATLAB) vs double (here) precision and FANSI's γ, 42.58 against the crate's 42.576
73//! (the wrappers keep theirs). Fixture: `tests/fansi_qsmci_parity.rs` (relative L2 below 1e-5
74//! on an oblique, anisotropic phantom).
75//!
76//! # Pipeline rules
77//!
78//! The facts above that a configuration can depart from are published as rules in
79//! [`crate::pipeline::rules::RULES`] (each with its strength, basis, reason and source): `fansi.no-laplacian-multi-echo`, `fansi.te-spacing`.
80
81use crate::inversion::admm::{fansi_magnitude_weight, fansi_ppm_to_rad, fansi_unit_ee2, prepare_fansi_kernel, FansiKernel};
82use crate::utils::gradient::{bdiv_inplace, fgrad_inplace};
83use crate::utils::{apply_mask_zero, shrink};
84use crate::{EchoTime, Grid};
85use num_complex::Complex64;
86use std::f64::consts::PI;
87
88/// FANSI nlTV/nlTGV parameters (nlTV.m / nlTGV.m defaults where they have them).
89#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
90#[derive(Clone, Debug)]
91pub struct FansiParams {
92    /// First-order (TV / TGV gradient) L1 penalty weight.
93    pub alpha1: f64,
94    /// Gradient-consistency ADMM weight; `None` is FANSI's `100 * alpha1`.
95    pub mu1: Option<f64>,
96    /// Fidelity-consistency ADMM weight.
97    pub mu2: f64,
98    /// Second-order (symmetric-gradient) L1 penalty weight (nlTGV only); `None` is
99    /// nlTGV.m's `2 * alpha1`.
100    pub alpha0: Option<f64>,
101    /// Second-order consistency ADMM weight (nlTGV only); `None` is nlTGV.m's `2 * mu1`.
102    pub mu0: Option<f64>,
103    /// Number of outer ADMM iterations.
104    pub max_iter: usize,
105    /// Percent-update convergence stopping tolerance.
106    pub tol_update: f64,
107    /// Inner Newton convergence tolerance.
108    pub tol_delta: f64,
109    /// Field strength in Tesla (scan parameter, for the radian conversion)
110    pub b0: f64,
111    /// Echo times in seconds, in acquisition order (scan parameter, for the radian
112    /// conversion).
113    pub echo_times: Vec<f64>,
114    /// Which echo time the radian conversion uses. Default [`EchoTime::Spacing`]: FANSI's
115    /// multi-echo example, `script_DICOM_Philips.m`, scales the field to radians by `deltaTE =
116    /// TE(echo2) - TE(echo1)` (FANSI-toolbox 93044a3, line 364) before `nlTV`/`nlTGV` (lines 370,
117    /// 421). That script's `echo1 = 2`, `echo2 = 4` (lines 224-225) make it twice the spacing of
118    /// its evenly spaced echoes; its comment adds "You may use also any individual echo time".
119    /// On the qsm-forward phantom (echoes 4/12/20/28 ms; FANSI's discrete kernel, measured for QSM.rs#418 with the integration tests' settings; XSIM / NRMSE / r), the
120    /// magnitude weight does better at the echo spacing than at the first echo on every metric:
121    /// nlTV 0.801 / 40.7% / 0.914 vs 0.799 / 43.6% / 0.901, nlTGV (50 iterations) 0.768 / 41.0% /
122    /// 0.912 vs 0.767 / 43.0% / 0.903. (With the continuous kernel, #276: nlTV 0.802 / 39.8% /
123    /// 0.918 vs 0.797 / 43.0% / 0.903, nlTGV 0.767 / 40.1% / 0.916 vs 0.765 / 42.8% / 0.904.)
124    ///
125    /// The magnitude weight is required, although `nlTV.m` and `nlTGV.m` weight by ones when
126    /// given none (`nlTV.m` lines 213, 274; `nlTGV.m` lines 350, 365): with a flat weight both
127    /// diverge on that phantom at either echo time, with either kernel.
128    pub te: EchoTime,
129    /// Select nlTGV (`true`) or nlTV (`false`).
130    pub is_tgv: bool,
131    /// The dipole kernel. Default [`FansiKernel::Discrete`], FANSI's
132    /// `dipole_kernel_angulated` as the original is run with; [`FansiKernel::Continuous`] is
133    /// the crate's continuous kernel (see [`FansiKernel`] for the decision and the measurements).
134    pub kernel: FansiKernel,
135}
136
137impl Default for FansiParams {
138    fn default() -> Self {
139        Self {
140            alpha1: 2e-4,
141            mu1: None,
142            mu2: 1.0,
143            alpha0: None,
144            mu0: None,
145            max_iter: 150,
146            tol_update: 0.1,
147            tol_delta: 1e-6,
148            b0: 3.0,
149            echo_times: vec![0.02],
150            te: EchoTime::Spacing,
151            is_tgv: false,
152            kernel: FansiKernel::Discrete,
153        }
154    }
155}
156
157impl FansiParams {
158    fn mu1(&self) -> f64 {
159        self.mu1.unwrap_or(100.0 * self.alpha1)
160    }
161}
162
163/// ppm -> radians, and the squared data weight (see module docs).
164fn phase_and_weight(
165    local_field: &[f64],
166    mask: &[u8],
167    magnitude: &[f64],
168    params: &FansiParams,
169) -> (f64, Vec<f64>, Vec<f64>) {
170    let ppm_to_rad = fansi_ppm_to_rad(params.b0, params.te.seconds(&params.echo_times, "FANSI"));
171    let phase = local_field.iter().map(|&f| f * ppm_to_rad).collect();
172    let w = fansi_magnitude_weight(mask, magnitude).into_iter().map(|v| v * v).collect();
173    (ppm_to_rad, phase, w)
174}
175
176/// L2 norm of a real slice.
177fn norm2(v: &[f64]) -> f64 {
178    v.iter().map(|&a| a * a).sum::<f64>().sqrt()
179}
180
181/// FANSI nlTV/nlTGV dipole inversion.
182///
183/// # Arguments
184/// * `local_field` - Local field in ppm (nx * ny * nz).
185/// * `mask` - Binary mask (nx * ny * nz), non-zero = inside ROI.
186/// * `magnitude` - Magnitude image, the data weight (see module docs).
187/// * `grid` - Volume grid (dimensions and voxel sizes).
188/// * `bdir` - B0 field direction.
189/// * `params` - FANSI parameters (`is_tgv` selects nlTGV vs nlTV).
190/// * `progress` - Progress callback `(iteration, max_iter)`.
191///
192/// # Returns
193/// Susceptibility map in ppm, masked to the ROI.
194pub fn fansi(
195    local_field: &[f64],
196    mask: &[u8],
197    magnitude: &[f64],
198    grid: &Grid,
199    bdir: (f64, f64, f64),
200    params: &FansiParams,
201    progress: impl FnMut(usize, usize),
202) -> Vec<f64> {
203    if params.is_tgv {
204        nltgv(local_field, mask, magnitude, grid, bdir, params, progress)
205    } else {
206        nltv(local_field, mask, magnitude, grid, bdir, params, progress)
207    }
208}
209
210/// Nonlinear total-variation dipole inversion (FANSI `nlTV.m`).
211fn nltv(
212    local_field: &[f64],
213    mask: &[u8],
214    magnitude: &[f64],
215    grid: &Grid,
216    bdir: (f64, f64, f64),
217    params: &FansiParams,
218    mut progress: impl FnMut(usize, usize),
219) -> Vec<f64> {
220    let n = grid.n_total();
221
222    let (mut fft_ws, k) = prepare_fansi_kernel(grid, bdir, params.kernel);
223    // FANSI's E1..E3 are unit-spacing differences.
224    let ee2 = fansi_unit_ee2(grid.dims);
225    let ugrid = Grid::new(grid.nx(), grid.ny(), grid.nz(), 1.0, 1.0, 1.0);
226    let grid = &ugrid;
227
228    let (ppm_to_rad, phase, w) = phase_and_weight(local_field, mask, magnitude, params);
229
230    let mu1 = params.mu1();
231    let mu2 = params.mu2;
232    let alpha_over_mu = params.alpha1 / mu1;
233
234    // ADMM variables.
235    let mut x = vec![0.0f64; n];
236    let mut x_prev = vec![0.0f64; n];
237
238    // Gradient-consistency split (real) and its multipliers.
239    let mut z_dx = vec![0.0f64; n];
240    let mut z_dy = vec![0.0f64; n];
241    let mut z_dz = vec![0.0f64; n];
242    let mut s_dx = vec![0.0f64; n];
243    let mut s_dy = vec![0.0f64; n];
244    let mut s_dz = vec![0.0f64; n];
245
246    // Fidelity-consistency auxiliary z2 = W .* phase ./ (W + mu2), and multiplier.
247    let mut z2 = vec![0.0f64; n];
248    for i in 0..n {
249        let den = w[i] + mu2;
250        z2[i] = if den != 0.0 { w[i] * phase[i] / den } else { 0.0 };
251    }
252    let mut s2 = vec![0.0f64; n];
253
254    // Reusable buffers.
255    let mut fdiv = vec![Complex64::new(0.0, 0.0); n];
256    let mut fd2 = vec![Complex64::new(0.0, 0.0); n];
257    let mut xhat = vec![Complex64::new(0.0, 0.0); n];
258    let mut fx = vec![Complex64::new(0.0, 0.0); n];
259
260    let mut gxc = vec![0.0f64; n];
261    let mut gyc = vec![0.0f64; n];
262    let mut gzc = vec![0.0f64; n];
263    let mut x_dx = vec![0.0f64; n];
264    let mut x_dy = vec![0.0f64; n];
265    let mut x_dz = vec![0.0f64; n];
266    let mut div = vec![0.0f64; n];
267    let mut dx = vec![0.0f64; n]; // Dx = real(ifft(k .* fft(x)))
268    let mut rhs_z2 = vec![0.0f64; n];
269    let mut diff = vec![0.0f64; n];
270    let mut update = vec![0.0f64; n];
271
272    for t in 0..params.max_iter {
273        progress(t + 1, params.max_iter);
274
275        // ---- x-subproblem -------------------------------------------------
276        // Gradient side: mu1 * sum(E_t .* fft(z_d - s_d)) == mu1 * fft(bdiv(z_d - s_d)).
277        for i in 0..n {
278            gxc[i] = z_dx[i] - s_dx[i];
279            gyc[i] = z_dy[i] - s_dy[i];
280            gzc[i] = z_dz[i] - s_dz[i];
281        }
282        bdiv_inplace(&mut div, &gxc, &gyc, &gzc, grid);
283        for i in 0..n {
284            fdiv[i] = Complex64::new(div[i], 0.0);
285        }
286        fft_ws.fft3d(&mut fdiv);
287
288        // Fidelity side: mu2 * conj(K) .* fft(z2 - s2) (K real -> conj = K).
289        for i in 0..n {
290            fd2[i] = Complex64::new(z2[i] - s2[i], 0.0);
291        }
292        fft_ws.fft3d(&mut fd2);
293
294        for i in 0..n {
295            // NOTE the minus on the gradient-consistency term: the adjoint of the
296            // crate's forward-difference `fgrad` is `-bdiv` (not `+bdiv`), so the
297            // spectral term mu1*sum(E_t.*F(z_d-s_d)) = -mu1*F(bdiv(z_d-s_d)).
298            // Matches QSM.rs's own TV-ADMM (`f_hat - rho*fft(bdiv...)`). Using +bdiv
299            // fails to cancel mu1*∇*∇ at the fixed point, doubling the effective
300            // regularization and damping the susceptibility amplitude.
301            let num = -mu1 * fdiv[i] + mu2 * k[i] * fd2[i];
302            let den = mu2 * k[i] * k[i] + mu1 * ee2[i];
303            // Guard the dipole null-space (DC and any singular bin): both the
304            // dipole kernel and the Laplacian vanish there, so susceptibility is
305            // undetermined up to a constant. Zero it (matches TV's inv_a guard);
306            // dividing FFT round-off by ~0 would otherwise create a huge DC pedestal.
307            xhat[i] = if den > 1e-20 { num / den } else { Complex64::new(0.0, 0.0) };
308        }
309        fft_ws.ifft3d(&mut xhat);
310        x_prev.copy_from_slice(&x);
311        for i in 0..n {
312            x[i] = xhat[i].re;
313        }
314
315        // ---- convergence check --------------------------------------------
316        // As nlTV.m, a zero solution gives NaN and stops.
317        for i in 0..n {
318            diff[i] = x[i] - x_prev[i];
319        }
320        let x_update = 100.0 * norm2(&diff) / norm2(&x);
321        if x_update < params.tol_update || x_update.is_nan() {
322            progress(t + 1, t + 1);
323            break;
324        }
325
326        if t + 1 >= params.max_iter {
327            break;
328        }
329
330        // ---- gradient split update (TV shrink) ----------------------------
331        // Fx = fft(x) (reused for the dipole rhs below).
332        for i in 0..n {
333            fx[i] = Complex64::new(x[i], 0.0);
334        }
335        fft_ws.fft3d(&mut fx);
336
337        fgrad_inplace(&mut x_dx, &mut x_dy, &mut x_dz, &x, grid);
338        for i in 0..n {
339            z_dx[i] = shrink(x_dx[i] + s_dx[i], alpha_over_mu);
340            z_dy[i] = shrink(x_dy[i] + s_dy[i], alpha_over_mu);
341            z_dz[i] = shrink(x_dz[i] + s_dz[i], alpha_over_mu);
342            s_dx[i] += x_dx[i] - z_dx[i];
343            s_dy[i] += x_dy[i] - z_dy[i];
344            s_dz[i] += x_dz[i] - z_dz[i];
345        }
346
347        // ---- fidelity auxiliary (z2) via Newton ---------------------------
348        // Dx = real(ifft(K .* Fx)) ; rhs_z2 = mu2 * (Dx + s2) ; z2 init = Dx + s2.
349        for i in 0..n {
350            xhat[i] = fx[i] * k[i];
351        }
352        fft_ws.ifft3d(&mut xhat);
353        for i in 0..n {
354            dx[i] = xhat[i].re;
355            rhs_z2[i] = mu2 * (dx[i] + s2[i]);
356            z2[i] = rhs_z2[i] / mu2;
357        }
358
359        let mut delta = f64::INFINITY;
360        let mut inn = 0usize;
361        while delta > params.tol_delta && inn < 10 {
362            inn += 1;
363            let norm_old = norm2(&z2);
364            for i in 0..n {
365                let a = z2[i] - phase[i];
366                let numer = w[i] * a.sin() + mu2 * z2[i] - rhs_z2[i];
367                let denom = w[i] * a.cos() + mu2;
368                update[i] = numer / denom;
369                z2[i] -= update[i];
370            }
371            delta = if norm_old > 0.0 {
372                norm2(&update) / norm_old
373            } else {
374                0.0
375            };
376        }
377
378        // ---- multiplier update --------------------------------------------
379        // s2 = s2 + Dx - z2.
380        for i in 0..n {
381            s2[i] += dx[i] - z2[i];
382        }
383    }
384
385    // Back to ppm.
386    for v in &mut x {
387        *v /= ppm_to_rad;
388    }
389
390    apply_mask_zero(&mut x, mask);
391    x
392}
393
394/// Multiply a complex spectral array in place by a complex multiplier array.
395#[inline]
396fn spectral_mul_assign(dst: &mut [Complex64], m: &[Complex64]) {
397    for (d, &mm) in dst.iter_mut().zip(m.iter()) {
398        *d *= mm;
399    }
400}
401
402/// Nonlinear total-generalized-variation dipole inversion (FANSI `nlTGV.m`).
403///
404/// Everything (operators + normal-equation cofactors) is built from the local,
405/// *unscaled* spectral gradient multipliers `E1,E2,E3` to stay self-consistent
406/// with the MATLAB Cramer-rule algebra.
407#[allow(clippy::too_many_lines)]
408fn nltgv(
409    local_field: &[f64],
410    mask: &[u8],
411    magnitude: &[f64],
412    grid: &Grid,
413    bdir: (f64, f64, f64),
414    params: &FansiParams,
415    mut progress: impl FnMut(usize, usize),
416) -> Vec<f64> {
417    let n = grid.n_total();
418    let (nx, ny, nz) = (grid.nx(), grid.ny(), grid.nz());
419
420    let (mut fft_ws, k) = prepare_fansi_kernel(grid, bdir, params.kernel);
421
422    // Local unscaled spectral gradient multipliers, MATLAB order i+j*nx+k*nx*ny.
423    let mut e1 = vec![Complex64::new(0.0, 0.0); n];
424    let mut e2 = vec![Complex64::new(0.0, 0.0); n];
425    let mut e3 = vec![Complex64::new(0.0, 0.0); n];
426    let two_pi = 2.0 * PI;
427    for kk in 0..nz {
428        let ez = Complex64::new(0.0, 1.0) * (two_pi * (kk as f64) / (nz as f64));
429        let e3v = Complex64::new(1.0, 0.0) - ez.exp();
430        for jj in 0..ny {
431            let ey = Complex64::new(0.0, 1.0) * (two_pi * (jj as f64) / (ny as f64));
432            let e2v = Complex64::new(1.0, 0.0) - ey.exp();
433            for ii in 0..nx {
434                let ex = Complex64::new(0.0, 1.0) * (two_pi * (ii as f64) / (nx as f64));
435                let e1v = Complex64::new(1.0, 0.0) - ex.exp();
436                let idx = ii + jj * nx + kk * nx * ny;
437                e1[idx] = e1v;
438                e2[idx] = e2v;
439                e3[idx] = e3v;
440            }
441        }
442    }
443
444    let (ppm_to_rad, phase, w) = phase_and_weight(local_field, mask, magnitude, params);
445
446    let mu1 = params.mu1();
447    let mu0 = params.mu0.unwrap_or(2.0 * mu1);
448    let alpha0 = params.alpha0.unwrap_or(2.0 * params.alpha1);
449    let mu2 = params.mu2;
450
451    // Precompute per-voxel normal-equation cofactors (all complex).
452    let mut d11 = vec![Complex64::new(0.0, 0.0); n];
453    let mut d21 = d11.clone();
454    let mut d31 = d11.clone();
455    let mut d41 = d11.clone();
456    let mut d12 = d11.clone();
457    let mut d22 = d11.clone();
458    let mut d32 = d11.clone();
459    let mut d42 = d11.clone();
460    let mut d13 = d11.clone();
461    let mut d23 = d11.clone();
462    let mut d33 = d11.clone();
463    let mut d43 = d11.clone();
464    let mut d14 = d11.clone();
465    let mut d24 = d11.clone();
466    let mut d34 = d11.clone();
467    let mut d44 = d11.clone();
468    let mut det_ainv = d11.clone();
469
470    let half = 0.5;
471    for i in 0..n {
472        let e1i = e1[i];
473        let e2i = e2[i];
474        let e3i = e3[i];
475        let et1 = e1i.conj();
476        let et2 = e2i.conj();
477        let et3 = e3i.conj();
478
479        let e1te1 = et1 * e1i;
480        let e2te2 = et2 * e2i;
481        let e3te3 = et3 * e3i;
482        let mu0h_e1te2 = (mu0 * half) * et1 * e2i;
483        let mu0h_e1te3 = (mu0 * half) * et1 * e3i;
484        let mu0h_e2te3 = (mu0 * half) * et2 * e3i;
485
486        // a0 = mu2 * conj(K) .* K = mu2 * k * k (real).
487        let a0 = Complex64::new(mu2 * k[i] * k[i], 0.0);
488        let a1 = a0 + mu1 * (e1te1 + e2te2 + e3te3);
489        let a2 = Complex64::new(mu1, 0.0) + mu0 * (e1te1 + (e2te2 + e3te3) * half);
490        let a3 = Complex64::new(mu1, 0.0) + mu0 * (e1te1 * half + e2te2 + e3te3 * half);
491        let a4 = Complex64::new(mu1, 0.0) + mu0 * ((e1te1 + e2te2) * half + e3te3);
492        let a5 = -mu1 * e1i;
493        let a6 = -mu1 * e2i;
494        let a7 = mu0h_e1te2;
495        let a8 = -mu1 * e3i;
496        let a9 = mu0h_e1te3;
497        let a10 = mu0h_e2te3;
498        let a5t = a5.conj();
499        let a6t = a6.conj();
500        let a7t = a7.conj();
501        let a8t = a8.conj();
502        let a9t = a9.conj();
503        let a10t = a10.conj();
504
505        let c11 = a2 * a3 * a4 + a7t * a9 * a10t + a7 * a9t * a10
506            - a3 * a9 * a9t
507            - a2 * a10 * a10t
508            - a4 * a7 * a7t;
509        let c21 = a3 * a4 * a5t + a6t * a9 * a10t + a7 * a8t * a10
510            - a3 * a8t * a9
511            - a5t * a10 * a10t
512            - a4 * a6t * a7;
513        let c31 = a4 * a5t * a7t + a6t * a9 * a9t + a2 * a8t * a10
514            - a7t * a8t * a9
515            - a5t * a9t * a10
516            - a2 * a4 * a6t;
517        let c41 = a5t * a7t * a10t + a6t * a7 * a9t + a2 * a3 * a8t
518            - a7 * a7t * a8t
519            - a3 * a5t * a9t
520            - a2 * a6t * a10t;
521        let c12 = a3 * a4 * a5 + a7t * a8 * a10t + a6 * a9t * a10
522            - a3 * a8 * a9t
523            - a5 * a10 * a10t
524            - a4 * a6 * a7t;
525        let c22 = a1 * a3 * a4 + a6t * a8 * a10t + a6 * a8t * a10
526            - a3 * a8 * a8t
527            - a1 * a10 * a10t
528            - a4 * a6 * a6t;
529        let c32 = a1 * a4 * a7t + a6t * a8 * a9t + a5 * a8t * a10
530            - a7t * a8 * a8t
531            - a1 * a9t * a10
532            - a4 * a5 * a6t;
533        let c42 = a1 * a7t * a10t + a6 * a6t * a9t + a3 * a5 * a8t
534            - a6 * a7t * a8t
535            - a1 * a3 * a9t
536            - a5 * a6t * a10t;
537        let c13 = a4 * a5 * a7 + a2 * a8 * a10t + a6 * a9 * a9t
538            - a7 * a8 * a9t
539            - a5 * a9 * a10t
540            - a2 * a4 * a6;
541        let c23 = a1 * a4 * a7 + a5t * a8 * a10t + a6 * a8t * a9
542            - a7 * a8 * a8t
543            - a1 * a9 * a10t
544            - a4 * a5t * a6;
545        let c33 = a1 * a2 * a4 + a5t * a8 * a9t + a5 * a8t * a9
546            - a2 * a8 * a8t
547            - a1 * a9 * a9t
548            - a4 * a5 * a5t;
549        let c43 = a1 * a2 * a10t + a5t * a6 * a9t + a5 * a7 * a8t
550            - a2 * a6 * a8t
551            - a1 * a7 * a9t
552            - a5 * a5t * a10t;
553        let c14 = a5 * a7 * a10 + a2 * a3 * a8 + a6 * a7t * a9
554            - a7 * a7t * a8
555            - a3 * a5 * a9
556            - a2 * a6 * a10;
557        let c24 = a1 * a7 * a10 + a3 * a5t * a8 + a6 * a6t * a9
558            - a6t * a7 * a8
559            - a1 * a3 * a9
560            - a5t * a6 * a10;
561        let c34 = a1 * a2 * a10 + a5t * a7t * a8 + a5 * a6t * a9
562            - a2 * a6t * a8
563            - a1 * a7t * a9
564            - a5 * a5t * a10;
565        let c44 = a1 * a2 * a3 + a5t * a6 * a7t + a5 * a6t * a7
566            - a2 * a6 * a6t
567            - a1 * a7 * a7t
568            - a3 * a5 * a5t;
569
570        let det_a = a1 * c11 - a5 * c21 + a6 * c31 - a8 * c41;
571
572        d11[i] = c11;
573        d21[i] = c21;
574        d31[i] = c31;
575        d41[i] = c41;
576        d12[i] = c12;
577        d22[i] = c22;
578        d32[i] = c32;
579        d42[i] = c42;
580        d13[i] = c13;
581        d23[i] = c23;
582        d33[i] = c33;
583        d43[i] = c43;
584        d14[i] = c14;
585        d24[i] = c24;
586        d34[i] = c34;
587        d44[i] = c44;
588        // Guard the singular (DC null-space) bins: at k=0 every spectral
589        // operator vanishes so det_a -> 0. Zero the solve there instead of
590        // dividing FFT round-off by ~eps (which would create a huge DC pedestal).
591        det_ainv[i] = if det_a.norm() > 1e-20 {
592            Complex64::new(1.0, 0.0) / det_a
593        } else {
594            Complex64::new(0.0, 0.0)
595        };
596    }
597
598    // Conjugates of E used in RHS assembly.
599    let et1: Vec<Complex64> = e1.iter().map(|c| c.conj()).collect();
600    let et2: Vec<Complex64> = e2.iter().map(|c| c.conj()).collect();
601    let et3: Vec<Complex64> = e3.iter().map(|c| c.conj()).collect();
602
603    // ADMM state.
604    let mut x = vec![0.0f64; n];
605    let mut x_prev = vec![0.0f64; n];
606    let mut v1 = vec![0.0f64; n];
607    let mut v2 = vec![0.0f64; n];
608    let mut v3 = vec![0.0f64; n];
609
610    // First-order splits.
611    let mut z1_1 = vec![0.0f64; n];
612    let mut z1_2 = vec![0.0f64; n];
613    let mut z1_3 = vec![0.0f64; n];
614    let mut s1_1 = vec![0.0f64; n];
615    let mut s1_2 = vec![0.0f64; n];
616    let mut s1_3 = vec![0.0f64; n];
617
618    // Symmetric (second-order) splits.
619    let mut z0_1 = vec![0.0f64; n];
620    let mut z0_2 = vec![0.0f64; n];
621    let mut z0_3 = vec![0.0f64; n];
622    let mut z0_4 = vec![0.0f64; n];
623    let mut z0_5 = vec![0.0f64; n];
624    let mut z0_6 = vec![0.0f64; n];
625    let mut s0_1 = vec![0.0f64; n];
626    let mut s0_2 = vec![0.0f64; n];
627    let mut s0_3 = vec![0.0f64; n];
628    let mut s0_4 = vec![0.0f64; n];
629    let mut s0_5 = vec![0.0f64; n];
630    let mut s0_6 = vec![0.0f64; n];
631
632    // Fidelity auxiliary z2 = W .* phase ./ (W + mu2).
633    let mut z2 = vec![0.0f64; n];
634    for i in 0..n {
635        let den = w[i] + mu2;
636        z2[i] = if den != 0.0 { w[i] * phase[i] / den } else { 0.0 };
637    }
638    let mut s2 = vec![0.0f64; n];
639
640    let alpha1_over_mu1 = params.alpha1 / mu1;
641    let alpha0_over_mu0 = alpha0 / mu0;
642
643    // Reusable complex/real buffers.
644    let mut rhs1 = vec![Complex64::new(0.0, 0.0); n];
645    let mut rhs2 = rhs1.clone();
646    let mut rhs3 = rhs1.clone();
647    let mut rhs4 = rhs1.clone();
648    let mut t1 = rhs1.clone();
649    let mut t2 = rhs1.clone();
650    let mut t3 = rhs1.clone();
651    let mut fx = rhs1.clone();
652    let mut fv1 = rhs1.clone();
653    let mut fv2 = rhs1.clone();
654    let mut fv3 = rhs1.clone();
655    let mut cbuf = rhs1.clone();
656
657    let mut dx1 = vec![0.0f64; n];
658    let mut dx2 = vec![0.0f64; n];
659    let mut dx3 = vec![0.0f64; n];
660    let mut ev1 = vec![0.0f64; n];
661    let mut ev2 = vec![0.0f64; n];
662    let mut ev3 = vec![0.0f64; n];
663    let mut ev4 = vec![0.0f64; n];
664    let mut ev5 = vec![0.0f64; n];
665    let mut ev6 = vec![0.0f64; n];
666    let mut dx = vec![0.0f64; n];
667    let mut rhs_z2 = vec![0.0f64; n];
668    let mut diff = vec![0.0f64; n];
669    let mut update = vec![0.0f64; n];
670
671    // Scratch real buffer for building fft inputs.
672    let mut rbuf = vec![0.0f64; n];
673
674    // real(fft) of a real field `src` into complex buffer `dst`.
675    macro_rules! fft_real {
676        ($dst:expr, $src:expr) => {{
677            for i in 0..n {
678                $dst[i] = Complex64::new($src[i], 0.0);
679            }
680            fft_ws.fft3d(&mut $dst);
681        }};
682    }
683    // real(fft) of (a - b) elementwise into complex buffer `dst`.
684    macro_rules! fft_real_diff {
685        ($dst:expr, $a:expr, $b:expr) => {{
686            for i in 0..n {
687                rbuf[i] = $a[i] - $b[i];
688            }
689            fft_real!($dst, rbuf);
690        }};
691    }
692
693    for t in 0..params.max_iter {
694        progress(t + 1, params.max_iter);
695
696        // ---- assemble RHS (spectral) --------------------------------------
697        // rhs1 = mu2*conj(K)*fft(z2 - s2) + mu1*(Et*fft(z1 - s1))
698        for i in 0..n {
699            cbuf[i] = Complex64::new(z2[i] - s2[i], 0.0);
700        }
701        fft_ws.fft3d(&mut cbuf);
702        for i in 0..n {
703            rhs1[i] = mu2 * k[i] * cbuf[i];
704        }
705        fft_real_diff!(t1, z1_1, s1_1);
706        fft_real_diff!(t2, z1_2, s1_2);
707        fft_real_diff!(t3, z1_3, s1_3);
708        for i in 0..n {
709            rhs1[i] += mu1 * (et1[i] * t1[i] + et2[i] * t2[i] + et3[i] * t3[i]);
710        }
711
712        // rhs2 = -mu1*fft(z1_1 - s1_1) + mu0*(Et1*fft(z0_1 - s0_1) + Et2*fft(z0_4 - s0_4) + Et3*fft(z0_5 - s0_5))
713        // t1 already holds fft(z1_1 - s1_1).
714        for i in 0..n {
715            rhs2[i] = -mu1 * t1[i];
716            rhs3[i] = -mu1 * t2[i];
717            rhs4[i] = -mu1 * t3[i];
718        }
719        // rhs2 second-order part.
720        fft_real_diff!(t1, z0_1, s0_1);
721        fft_real_diff!(t2, z0_4, s0_4);
722        fft_real_diff!(t3, z0_5, s0_5);
723        for i in 0..n {
724            rhs2[i] += mu0 * (et1[i] * t1[i] + et2[i] * t2[i] + et3[i] * t3[i]);
725        }
726        // rhs3 second-order part: mu0*(Et2*fft(z0_2 - s0_2) + Et1*fft(z0_4 - s0_4) + Et3*fft(z0_6 - s0_6))
727        fft_real_diff!(t1, z0_2, s0_2);
728        fft_real_diff!(t2, z0_4, s0_4);
729        fft_real_diff!(t3, z0_6, s0_6);
730        for i in 0..n {
731            rhs3[i] += mu0 * (et2[i] * t1[i] + et1[i] * t2[i] + et3[i] * t3[i]);
732        }
733        // rhs4 second-order part: mu0*(Et3*fft(z0_3 - s0_3) + Et1*fft(z0_5 - s0_5) + Et2*fft(z0_6 - s0_6))
734        fft_real_diff!(t1, z0_3, s0_3);
735        fft_real_diff!(t2, z0_5, s0_5);
736        fft_real_diff!(t3, z0_6, s0_6);
737        for i in 0..n {
738            rhs4[i] += mu0 * (et3[i] * t1[i] + et1[i] * t2[i] + et2[i] * t3[i]);
739        }
740
741        // ---- Cramer solve -------------------------------------------------
742        for i in 0..n {
743            let r1 = rhs1[i];
744            let r2 = rhs2[i];
745            let r3 = rhs3[i];
746            let r4 = rhs4[i];
747            let da = det_ainv[i];
748            fx[i] = (r1 * d11[i] - r2 * d21[i] + r3 * d31[i] - r4 * d41[i]) * da;
749            fv1[i] = (-r1 * d12[i] + r2 * d22[i] - r3 * d32[i] + r4 * d42[i]) * da;
750            fv2[i] = (r1 * d13[i] - r2 * d23[i] + r3 * d33[i] - r4 * d43[i]) * da;
751            fv3[i] = (-r1 * d14[i] + r2 * d24[i] - r3 * d34[i] + r4 * d44[i]) * da;
752        }
753        // x, v via ifft (take real part). Use t1 as scratch to preserve the
754        // spectral fx/fv (re-FFT'd after convergence anyway).
755        t1.copy_from_slice(&fx);
756        fft_ws.ifft3d(&mut t1);
757        x_prev.copy_from_slice(&x);
758        for i in 0..n {
759            x[i] = t1[i].re;
760        }
761        t1.copy_from_slice(&fv1);
762        fft_ws.ifft3d(&mut t1);
763        for i in 0..n {
764            v1[i] = t1[i].re;
765        }
766        t1.copy_from_slice(&fv2);
767        fft_ws.ifft3d(&mut t1);
768        for i in 0..n {
769            v2[i] = t1[i].re;
770        }
771        t1.copy_from_slice(&fv3);
772        fft_ws.ifft3d(&mut t1);
773        for i in 0..n {
774            v3[i] = t1[i].re;
775        }
776
777        // ---- convergence --------------------------------------------------
778        let xnorm = norm2(&x);
779        if xnorm > 0.0 {
780            for i in 0..n {
781                diff[i] = x[i] - x_prev[i];
782            }
783            let x_update = 100.0 * norm2(&diff) / xnorm;
784            if x_update < params.tol_update || x_update.is_nan() {
785                progress(t + 1, t + 1);
786                break;
787            }
788        }
789
790        if t + 1 >= params.max_iter {
791            break;
792        }
793
794        // ---- re-FFT for stability -----------------------------------------
795        fft_real!(fx, x);
796        fft_real!(fv1, v1);
797        fft_real!(fv2, v2);
798        fft_real!(fv3, v3);
799
800        // Dx1 = real(ifft(E1.*Fx)), etc.
801        cbuf.copy_from_slice(&fx);
802        spectral_mul_assign(&mut cbuf, &e1);
803        fft_ws.ifft3d(&mut cbuf);
804        for i in 0..n {
805            dx1[i] = cbuf[i].re;
806        }
807        cbuf.copy_from_slice(&fx);
808        spectral_mul_assign(&mut cbuf, &e2);
809        fft_ws.ifft3d(&mut cbuf);
810        for i in 0..n {
811            dx2[i] = cbuf[i].re;
812        }
813        cbuf.copy_from_slice(&fx);
814        spectral_mul_assign(&mut cbuf, &e3);
815        fft_ws.ifft3d(&mut cbuf);
816        for i in 0..n {
817            dx3[i] = cbuf[i].re;
818        }
819
820        // E_v1 = real(ifft(E1.*Fv1)), E_v2 = ...E2.*Fv2, E_v3 = ...E3.*Fv3.
821        cbuf.copy_from_slice(&fv1);
822        spectral_mul_assign(&mut cbuf, &e1);
823        fft_ws.ifft3d(&mut cbuf);
824        for i in 0..n {
825            ev1[i] = cbuf[i].re;
826        }
827        cbuf.copy_from_slice(&fv2);
828        spectral_mul_assign(&mut cbuf, &e2);
829        fft_ws.ifft3d(&mut cbuf);
830        for i in 0..n {
831            ev2[i] = cbuf[i].re;
832        }
833        cbuf.copy_from_slice(&fv3);
834        spectral_mul_assign(&mut cbuf, &e3);
835        fft_ws.ifft3d(&mut cbuf);
836        for i in 0..n {
837            ev3[i] = cbuf[i].re;
838        }
839
840        // E_v4 = real(ifft(E1.*Fv2 + E2.*Fv1))/2, etc.
841        for i in 0..n {
842            cbuf[i] = e1[i] * fv2[i] + e2[i] * fv1[i];
843        }
844        fft_ws.ifft3d(&mut cbuf);
845        for i in 0..n {
846            ev4[i] = cbuf[i].re * 0.5;
847        }
848        for i in 0..n {
849            cbuf[i] = e1[i] * fv3[i] + e3[i] * fv1[i];
850        }
851        fft_ws.ifft3d(&mut cbuf);
852        for i in 0..n {
853            ev5[i] = cbuf[i].re * 0.5;
854        }
855        for i in 0..n {
856            cbuf[i] = e2[i] * fv3[i] + e3[i] * fv2[i];
857        }
858        fft_ws.ifft3d(&mut cbuf);
859        for i in 0..n {
860            ev6[i] = cbuf[i].re * 0.5;
861        }
862
863        // ---- symmetric-gradient split (shrink) ----------------------------
864        for i in 0..n {
865            z0_1[i] = shrink(ev1[i] + s0_1[i], alpha0_over_mu0);
866            z0_2[i] = shrink(ev2[i] + s0_2[i], alpha0_over_mu0);
867            z0_3[i] = shrink(ev3[i] + s0_3[i], alpha0_over_mu0);
868            z0_4[i] = shrink(ev4[i] + s0_4[i], alpha0_over_mu0);
869            z0_5[i] = shrink(ev5[i] + s0_5[i], alpha0_over_mu0);
870            z0_6[i] = shrink(ev6[i] + s0_6[i], alpha0_over_mu0);
871        }
872
873        // ---- first-order split (shrink) -----------------------------------
874        for i in 0..n {
875            z1_1[i] = shrink(dx1[i] - v1[i] + s1_1[i], alpha1_over_mu1);
876            z1_2[i] = shrink(dx2[i] - v2[i] + s1_2[i], alpha1_over_mu1);
877            z1_3[i] = shrink(dx3[i] - v3[i] + s1_3[i], alpha1_over_mu1);
878        }
879
880        // ---- fidelity auxiliary (z2) via Newton ---------------------------
881        // Dx = real(ifft(K.*Fx)) ; rhs_z2 = mu2*(Dx + s2) ; z2 init = Dx + s2.
882        for i in 0..n {
883            cbuf[i] = fx[i] * k[i];
884        }
885        fft_ws.ifft3d(&mut cbuf);
886        for i in 0..n {
887            dx[i] = cbuf[i].re;
888            rhs_z2[i] = mu2 * (dx[i] + s2[i]);
889            z2[i] = rhs_z2[i] / mu2;
890        }
891        let mut delta = f64::INFINITY;
892        let mut inn = 0usize;
893        while delta > params.tol_delta && inn < 50 {
894            inn += 1;
895            let norm_old = norm2(&z2);
896            for i in 0..n {
897                let a = z2[i] - phase[i];
898                let numer = w[i] * a.sin() + mu2 * z2[i] - rhs_z2[i];
899                let denom = w[i] * a.cos() + mu2;
900                update[i] = numer / denom;
901                z2[i] -= update[i];
902            }
903            delta = if norm_old > 0.0 {
904                norm2(&update) / norm_old
905            } else {
906                0.0
907            };
908        }
909
910        // ---- multiplier updates -------------------------------------------
911        for i in 0..n {
912            s0_1[i] += ev1[i] - z0_1[i];
913            s0_2[i] += ev2[i] - z0_2[i];
914            s0_3[i] += ev3[i] - z0_3[i];
915            s0_4[i] += ev4[i] - z0_4[i];
916            s0_5[i] += ev5[i] - z0_5[i];
917            s0_6[i] += ev6[i] - z0_6[i];
918            s1_1[i] += dx1[i] - v1[i] - z1_1[i];
919            s1_2[i] += dx2[i] - v2[i] - z1_2[i];
920            s1_3[i] += dx3[i] - v3[i] - z1_3[i];
921            s2[i] += dx[i] - z2[i];
922        }
923    }
924
925    // Back to ppm.
926    for v in &mut x {
927        *v /= ppm_to_rad;
928    }
929
930    apply_mask_zero(&mut x, mask);
931    x
932}
933
934#[cfg(test)]
935mod tests {
936    use super::*;
937
938    #[test]
939    fn test_fansi_nltv_zero_field() {
940        // Zero field should give (approximately) zero susceptibility.
941        let n = 8;
942        let field = vec![0.0; n * n * n];
943        let mask = vec![1u8; n * n * n];
944        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
945        let params = FansiParams {
946            max_iter: 10,
947            is_tgv: false,
948            ..FansiParams::default()
949        };
950
951        let chi = fansi(&field, &mask, &mask.iter().map(|_| 1.0).collect::<Vec<f64>>(), &grid, (0.0, 0.0, 1.0), &params, |_, _| {});
952
953        for &val in chi.iter() {
954            assert!(val.abs() < 1e-6, "Zero field should give ~zero chi, got {}", val);
955        }
956    }
957
958    #[test]
959    fn test_fansi_nltv_finite() {
960        // A small ramp field should produce all-finite output.
961        let n = 8;
962        let field: Vec<f64> = (0..n * n * n).map(|i| (i as f64) * 0.001).collect();
963        let mask = vec![1u8; n * n * n];
964        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
965        let params = FansiParams {
966            max_iter: 10,
967            is_tgv: false,
968            ..FansiParams::default()
969        };
970
971        let chi = fansi(&field, &mask, &mask.iter().map(|_| 1.0).collect::<Vec<f64>>(), &grid, (0.0, 0.0, 1.0), &params, |_, _| {});
972
973        for (i, &val) in chi.iter().enumerate() {
974            assert!(val.is_finite(), "Chi should be finite at index {}", i);
975        }
976    }
977
978    #[test]
979    fn test_fansi_nltgv_zero_field() {
980        // Zero field should give (approximately) zero susceptibility.
981        let n = 8;
982        let field = vec![0.0; n * n * n];
983        let mask = vec![1u8; n * n * n];
984        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
985        let params = FansiParams {
986            max_iter: 10,
987            is_tgv: true,
988            ..FansiParams::default()
989        };
990
991        let chi = fansi(&field, &mask, &mask.iter().map(|_| 1.0).collect::<Vec<f64>>(), &grid, (0.0, 0.0, 1.0), &params, |_, _| {});
992
993        for &val in chi.iter() {
994            assert!(val.abs() < 1e-6, "Zero field should give ~zero chi, got {}", val);
995        }
996    }
997
998    #[test]
999    fn test_fansi_nltgv_finite() {
1000        // A small ramp field should produce all-finite output.
1001        let n = 8;
1002        let field: Vec<f64> = (0..n * n * n).map(|i| (i as f64) * 0.001).collect();
1003        let mask = vec![1u8; n * n * n];
1004        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
1005        let params = FansiParams {
1006            max_iter: 10,
1007            is_tgv: true,
1008            ..FansiParams::default()
1009        };
1010
1011        let chi = fansi(&field, &mask, &mask.iter().map(|_| 1.0).collect::<Vec<f64>>(), &grid, (0.0, 0.0, 1.0), &params, |_, _| {});
1012
1013        for (i, &val) in chi.iter().enumerate() {
1014            assert!(val.is_finite(), "Chi should be finite at index {}", i);
1015        }
1016    }
1017}
qsm_core/inversion/fansi.rs

qsm_core/inversion/
fansi.rs
