qsm_core/inversion/

l1qsm.rs

//! L1-QSM: nonlinear L1 data-fidelity dipole inversion with TV regularization
//!
//! Solves the nonlinear QSM inverse problem with an L1 data-fidelity term
//! (robust to phase outliers / noise) and a total-variation prior:
//!
//! min_x  || W .* ( exp(i * D x) - exp(i * phi) ) ||_1  +  alpha1 * ||grad(x)||_1
//!
//! where `D` is the dipole operator and `phi` the local field phase in radians.
//! The nonlinear fidelity is linearized around the current auxiliary phase and
//! solved with an inner complex-argument Newton iteration, all embedded in an
//! ADMM splitting.
//!
//! Reference:
//! Milovic C, Lambert M, Langkammer C, Bredies K, Irarrazaval P, Tejos C.
//! "Streaking artifact suppression of quantitative susceptibility mapping
//! reconstructions via L1-norm data fidelity optimization (L1-QSM)."
//! Magnetic Resonance in Medicine, 2022;87:457-473.
//! https://doi.org/10.1002/mrm.28957
//!
//! Ported from the FANSI toolbox `nlL1TV.m`
//! (https://gitlab.com/cmilovic/FANSI-toolbox).
//!
//! # Units and weights
//! Like the other inversions, L1-QSM takes the local field in ppm and returns ppm.
//! nlL1TV.m works on the phase in radians, so the field is converted with
//! `2*pi*gamma*b0*te` first and the result converted back at the end. The
//! scale matters here: the L1 proximal step thresholds `|exp(i*Dx) - exp(i*phi)|`
//! at `W/mu3`, so in ppm every residual falls below it and the data are fitted
//! exactly. `b0` and `echo_times` are scan parameters
//! that the pipeline fills from metadata, and [`L1QsmParams::te`] picks the echo time the
//! conversion uses (by default the reference's choice; see there).
//!
//! `W = lambda * min(2*mask*mag/max(mag), 1)`: the magnitude in [0, 1] as in
//! FANSI's in vivo example, scaled by lambda (nlL1TV.m: "lambda*magnitude_data
//! or lambda*mask"). The magnitude is required: with a flat (mask) weight the
//! nonlinear data term can diverge where the phase nears +-pi.
//!
//! The finite differences are FANSI's, on unit spacing whatever the voxel size.
//!
//! # Reference pipeline
//! FANSI `script_DICOM_Philips.m` (commit 93044a3) runs `nlL1TV` on the same local field and
//! TE scaling as nlTV (see [`super::fansi`](mod@super::fansi)), with weight `lambda*m0.*mask./cnoise / max`,
//! `lambda = 0.8` (L443-450), and says the regularisation weight must be re-tuned for L1-QSM,
//! not reused from FANSI or WH-QSM (L442-443). `nlL1TV.m` defaults: mu1 = 100*alpha1, mu2 = mu3
//! = 1, 50 iterations, tolUpdate 1, weight ones (L48-112).
//!
//! # Parity
//! Against the verified original (QSM-CI pair `l1qsm`; FANSI-toolbox 93044a3, identical to
//! licensed MATLAB R2023b), final-map χ NRMSE at both sides' defaults on QSM-CI's simulated head
//! and the KKI travelling head (QSM-CI parity run 41403629314, QSM.rs#399–#403): sim 0.0022 %, KKI 0.0026 %, from
//! 14.4 % / 15.5 % (sim / KKI) before. The kernel was the whole gap (a numpy `nlL1TV.m` with FANSI's kernel is 0.0004 % from the original on sim; the `isPrecond` start and the 0.05 floor already matched). What remains is
//! single (MATLAB) vs double (here) precision and FANSI's γ, 42.58 against the crate's 42.576
//! (the wrappers keep theirs). Fixture: `tests/fansi_qsmci_parity.rs` (relative L2 below 1e-5
//! on an oblique, anisotropic phantom).
//!
//! # Pipeline rules
//!
//! The facts above that a configuration can depart from are published as rules in
//! [`crate::pipeline::rules::RULES`] (each with its strength, basis, reason and source): `fansi.no-laplacian-multi-echo`, `l1qsm.te-first`.

use crate::inversion::admm::{fansi_magnitude_weight, fansi_ppm_to_rad, fansi_unit_ee2, prepare_fansi_kernel, FansiKernel};
use crate::utils::gradient::{bdiv_inplace, fgrad_inplace};
use crate::utils::{apply_mask_zero, shrink};
use crate::{EchoTime, Grid};
use num_complex::Complex64;

/// L1-QSM algorithm parameters (FANSI `nlL1TV.m` defaults where it has them).
#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
#[derive(Clone, Debug)]
pub struct L1QsmParams {
    /// Gradient (TV) L1 penalty weight. nlL1TV.m gives no default, and FANSI's example says
    /// L1-QSM needs its own tuning. The default 5e-3 was chosen from an alpha1 sweep on the
    /// EMTP travelling-head data: 2e-4 streaks, 1e-2 starts to oversmooth.
    pub alpha1: f64,
    /// Gradient-consistency ADMM weight; `None` is nlL1TV.m's `100 * alpha1`.
    pub mu1: Option<f64>,
    /// Fidelity-consistency ADMM weight.
    pub mu2: f64,
    /// L1 proximal ADMM weight.
    pub mu3: f64,
    /// Phase-rejection strength: the data weight is `lambda` times the normalised
    /// magnitude. Smaller values reject more inconsistent voxels. Default 0.8, the value of
    /// FANSI's in-vivo demo (`script_DICOM_Philips.m` l. 445-447: `lambda = 0.8;
    /// params.weight = lambda*params.weight/max(params.weight(:))`); `nlL1TV.m` itself has no
    /// default ("recommended = lambda*magnitude_data", l. 18-19). QSM.rs previously used 1.0.
    /// Measured against 1.0 (ROMEO field, default V-SHARP): on the qsm-forward phantom χ XSIM
    /// 0.338 → 0.339, NRMSE 78.2 → 78.1 % (on its true local field 0.808 → 0.809, NRMSE
    /// 38.8 → 39.2 %); on an in-vivo 3T Philips scan, agreement with the UK Biobank pipeline's χ
    /// XSIM 0.446 → 0.461, NRMSE 103.0 → 97.7 %.
    pub lambda: f64,
    /// Number of outer ADMM iterations.
    pub max_iter: usize,
    /// Percent-update convergence stopping tolerance.
    pub tol_update: f64,
    /// Inner Newton convergence tolerance.
    pub tol_delta: f64,
    /// Field strength in Tesla (scan parameter, for the radian conversion)
    pub b0: f64,
    /// Echo times in seconds, in acquisition order (scan parameter, for the radian
    /// conversion).
    pub echo_times: Vec<f64>,
    /// Which echo time the radian conversion uses. Default [`EchoTime::First`]. FANSI's
    /// multi-echo example, `script_DICOM_Philips.m`, scales the field by `deltaTE = TE(echo2) -
    /// TE(echo1)` (FANSI-toolbox 93044a3, line 364) before `nlL1TV` (line 450), and its comment
    /// on that line says "You may use also any individual echo time". 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 magnitude weight at the first
    /// echo does best (0.776 / 38.9% / 0.923), against 0.764 / 46.9% / 0.889 at the echo
    /// spacing. (With the continuous kernel, #276: 0.782 / 37.4% / 0.928 against 0.776 / 40.2% /
    /// 0.916.)
    ///
    /// The magnitude weight is required, although `nlL1TV.m` weights by ones when given none
    /// (lines 109-113): with a flat weight L1-QSM is worse on that phantom (0.767 / 42.0% / 0.912
    /// at the first echo, 0.763 / 43.1% / 0.906 at the spacing). With the continuous kernel the
    /// flat weight's NRMSE exceeded 100% at either echo time (0.738 / 108% / 0.636 at the first
    /// echo, #276).
    pub te: EchoTime,
    /// The dipole kernel. Default [`FansiKernel::Discrete`], FANSI's
    /// `dipole_kernel_angulated` as the original is run with; [`FansiKernel::Continuous`] is
    /// the crate's continuous kernel (see [`FansiKernel`] for the decision and the measurements).
    pub kernel: FansiKernel,
}

impl Default for L1QsmParams {
    fn default() -> Self {
        Self {
            alpha1: 5e-3,
            mu1: None,
            mu2: 1.0,
            mu3: 1.0,
            lambda: 0.8,
            max_iter: 50,
            tol_update: 1.0,
            tol_delta: 1e-6,
            b0: 3.0,
            echo_times: vec![0.02],
            te: EchoTime::First,
            kernel: FansiKernel::Discrete,
        }
    }
}

/// L2 norm of a real slice.
fn norm2(v: &[f64]) -> f64 {
    v.iter().map(|&a| a * a).sum::<f64>().sqrt()
}

/// L1-QSM nonlinear dipole inversion.
///
/// # Arguments
/// * `local_field` - Local field in ppm (nx * ny * nz).
/// * `mask` - Binary mask (nx * ny * nz), non-zero = inside ROI.
/// * `magnitude` - Magnitude image, the data weight (see module docs).
/// * `grid` - Volume grid (dimensions and voxel sizes).
/// * `bdir` - B0 field direction.
/// * `params` - L1-QSM parameters.
/// * `progress` - Progress callback `(iteration, max_iter)`.
///
/// # Returns
/// Susceptibility map in ppm, masked to the ROI.
pub fn l1qsm(
    local_field: &[f64],
    mask: &[u8],
    magnitude: &[f64],
    grid: &Grid,
    bdir: (f64, f64, f64),
    params: &L1QsmParams,
    mut progress: impl FnMut(usize, usize),
) -> Vec<f64> {
    let n = grid.n_total();

    let (mut fft_ws, k) = prepare_fansi_kernel(grid, bdir, params.kernel);
    // FANSI's E1..E3 are unit-spacing differences.
    let ee2 = fansi_unit_ee2(grid.dims);
    let ugrid = Grid::new(grid.nx(), grid.ny(), grid.nz(), 1.0, 1.0, 1.0);
    let grid = &ugrid;

    // ppm -> radians, as nlL1TV.m expects; W = lambda * normalised magnitude (or mask).
    let ppm_to_rad = fansi_ppm_to_rad(params.b0, params.te.seconds(&params.echo_times, "L1-QSM"));
    let phase: Vec<f64> = local_field.iter().map(|&f| f * ppm_to_rad).collect();
    let w: Vec<f64> = fansi_magnitude_weight(mask, magnitude)
        .into_iter()
        .map(|v| params.lambda * v)
        .collect();

    // IS = exp(1i * phase)
    let is: Vec<Complex64> = phase
        .iter()
        .map(|&p| Complex64::new(p.cos(), p.sin()))
        .collect();

    let mu1 = params.mu1.unwrap_or(100.0 * params.alpha1);
    let mu2 = params.mu2;
    let mu3 = params.mu3;
    let alpha_over_mu = params.alpha1 / mu1;

    // ADMM variables.
    let mut x = vec![0.0f64; n];
    let mut x_prev = vec![0.0f64; n];

    // Gradient-consistency split (real).
    let mut z_dx = vec![0.0f64; n];
    let mut z_dy = vec![0.0f64; n];
    let mut z_dz = vec![0.0f64; n];
    let mut s_dx = vec![0.0f64; n];
    let mut s_dy = vec![0.0f64; n];
    let mut s_dz = vec![0.0f64; n];

    // Fidelity-consistency split (real phase auxiliary) and its multiplier.
    // isPrecond=true: z2 = W .* phase / max(W)
    let w_max = w.iter().cloned().fold(0.0f64, f64::max);
    let mut z2 = vec![0.0f64; n];
    if w_max > 0.0 {
        for i in 0..n {
            z2[i] = w[i] * phase[i] / w_max;
        }
    }
    let mut s2 = vec![0.0f64; n];

    // L1 proximal split (complex) and its multiplier.
    let mut z3 = vec![Complex64::new(0.0, 0.0); n];
    let mut s3 = vec![Complex64::new(0.0, 0.0); n];

    // Reusable buffers.
    let mut fdiv = vec![Complex64::new(0.0, 0.0); n];
    let mut fd2 = vec![Complex64::new(0.0, 0.0); n];
    let mut xhat = vec![Complex64::new(0.0, 0.0); n];
    let mut fx = vec![Complex64::new(0.0, 0.0); n];

    let mut gxc = vec![0.0f64; n];
    let mut gyc = vec![0.0f64; n];
    let mut gzc = vec![0.0f64; n];
    let mut x_dx = vec![0.0f64; n];
    let mut x_dy = vec![0.0f64; n];
    let mut x_dz = vec![0.0f64; n];
    let mut div = vec![0.0f64; n];
    let mut dx = vec![0.0f64; n]; // Dx = real(ifft(k .* fft(x)))
    let mut rhs_z2 = vec![0.0f64; n];
    let mut diff = vec![0.0f64; n];

    for t in 0..params.max_iter {
        progress(t + 1, params.max_iter);

        // ---- x-subproblem (gradient consistency) --------------------------
        // gradient side: mu1 * sum(E_t .* fft(z_d - s_d)) == mu1 * fft(bdiv(z_d - s_d))
        for i in 0..n {
            gxc[i] = z_dx[i] - s_dx[i];
            gyc[i] = z_dy[i] - s_dy[i];
            gzc[i] = z_dz[i] - s_dz[i];
        }
        bdiv_inplace(&mut div, &gxc, &gyc, &gzc, grid);
        for i in 0..n {
            fdiv[i] = Complex64::new(div[i], 0.0);
        }
        fft_ws.fft3d(&mut fdiv);

        // fidelity side: mu2 * conj(K) .* fft(z2 - s2); K real so conj(K) = K,
        // applied after the FFT below.
        for i in 0..n {
            fd2[i] = Complex64::new(z2[i] - s2[i], 0.0);
        }
        fft_ws.fft3d(&mut fd2);

        // Combine in k-space.
        for i in 0..n {
            // Minus on the gradient term: adjoint of crate `fgrad` is `-bdiv` (matches
            // QSM.rs TV-ADMM). `+bdiv` doubles the effective regularization. See fansi.rs.
            let num = -mu1 * fdiv[i] + mu2 * k[i] * fd2[i];
            // Guard the dipole null-space (DC/singular bins): both dipole kernel
            // and Laplacian vanish there. Zero it instead of dividing FFT
            // round-off by ~0, which would otherwise create a huge DC pedestal.
            let den = mu2 * k[i] * k[i] + mu1 * ee2[i];
            xhat[i] = if den > 1e-20 { num / den } else { Complex64::new(0.0, 0.0) };
        }
        fft_ws.ifft3d(&mut xhat);
        x_prev.copy_from_slice(&x);
        for i in 0..n {
            x[i] = xhat[i].re;
        }

        // ---- convergence check (percent update) ---------------------------
        // As nlL1TV.m, a zero solution gives NaN and stops.
        for i in 0..n {
            diff[i] = x[i] - x_prev[i];
        }
        let x_update = 100.0 * norm2(&diff) / norm2(&x);
        if x_update < params.tol_update || x_update.is_nan() {
            progress(t + 1, t + 1);
            break;
        }

        if t + 1 >= params.max_iter {
            break;
        }

        // ---- gradient split update (TV shrink) ----------------------------
        // Fx = fft(x)  (reused for the dipole rhs below)
        for i in 0..n {
            fx[i] = Complex64::new(x[i], 0.0);
        }
        fft_ws.fft3d(&mut fx);

        // x_d{x,y,z} = fgrad(x)
        fgrad_inplace(&mut x_dx, &mut x_dy, &mut x_dz, &x, grid);
        for i in 0..n {
            z_dx[i] = shrink(x_dx[i] + s_dx[i], alpha_over_mu);
            z_dy[i] = shrink(x_dy[i] + s_dy[i], alpha_over_mu);
            z_dz[i] = shrink(x_dz[i] + s_dz[i], alpha_over_mu);
            s_dx[i] += x_dx[i] - z_dx[i];
            s_dy[i] += x_dy[i] - z_dy[i];
            s_dz[i] += x_dz[i] - z_dz[i];
        }

        // ---- L1 fidelity proximal (z3) ------------------------------------
        // Y3 = exp(1i*z2) - IS + s3 ;  z3 = soft_threshold_complex(Y3, W/mu3)
        for i in 0..n {
            let ez2 = Complex64::new(z2[i].cos(), z2[i].sin());
            let y3 = ez2 - is[i] + s3[i];
            let mag = y3.norm();
            let thr = w[i] / (mu3 + f64::EPSILON);
            let shr = (mag - thr).max(0.0);
            z3[i] = if mag > 0.0 {
                y3 * (shr / mag)
            } else {
                Complex64::new(0.0, 0.0)
            };
        }

        // ---- fidelity auxiliary (z2) via complex-arg Newton ---------------
        // rhs_z2 = mu2 * (Dx + s2), with Dx = real(ifft(K .* Fx))
        for i in 0..n {
            xhat[i] = fx[i] * k[i];
        }
        fft_ws.ifft3d(&mut xhat);
        for i in 0..n {
            dx[i] = xhat[i].re;
            rhs_z2[i] = mu2 * (dx[i] + s2[i]);
            // init z2 = rhs_z2 / mu2 (== Dx + s2)
            z2[i] = rhs_z2[i] / mu2;
        }

        // Newton per-voxel: minimize mu3*|exp(iz2)-(IS+z3-s3)| + mu2/2*(z2-...)^2
        // Uses complex-argument trigonometric linearization (see module notes).
        // yphase = angle(IS+z3-s3), ym = abs(IS+z3-s3), b = ln(ym).
        let mut yphase = vec![0.0f64; n];
        let mut cosh_b = vec![0.0f64; n];
        let mut sinh_b = vec![0.0f64; n];
        for i in 0..n {
            let yc = is[i] + z3[i] - s3[i];
            yphase[i] = yc.arg();
            let m = yc.norm();
            if m > 0.0 {
                cosh_b[i] = 0.5 * (m + 1.0 / m); // cosh(ln m)
                sinh_b[i] = 0.5 * (m - 1.0 / m); // sinh(ln m)
            } else {
                cosh_b[i] = 1.0;
                sinh_b[i] = 0.0;
            }
        }

        let mut delta = f64::INFINITY;
        let mut inn = 0usize;
        let mut update = vec![0.0f64; n];
        while delta > params.tol_delta && inn < 4 {
            inn += 1;
            let norm_old = norm2(&z2);

            for i in 0..n {
                let a = z2[i] - yphase[i];
                let (sa, ca) = a.sin_cos();
                // sin(a - i b) = sin(a)cosh(b) - i cos(a)sinh(b)
                let sin_arg = Complex64::new(sa * cosh_b[i], -ca * sinh_b[i]);
                // cos(a - i b) = cos(a)cosh(b) + i sin(a)sinh(b)
                let cos_arg = Complex64::new(ca * cosh_b[i], sa * sinh_b[i]);

                let temp = mu3 * cos_arg + Complex64::new(mu2 + f64::EPSILON, 0.0);
                let numer = mu3 * sin_arg + Complex64::new(mu2 * z2[i] - rhs_z2[i], 0.0);

                // max(abs(temp), 0.05) .* sign(temp): floor magnitude at 0.05.
                let tm = temp.norm();
                let denom = if tm < 0.05 {
                    if tm > 0.0 {
                        temp * (0.05 / tm)
                    } else {
                        Complex64::new(0.05, 0.0)
                    }
                } else {
                    temp
                };

                update[i] = (numer / denom).re;
                z2[i] -= update[i];
            }

            let delta_new = if norm_old > 0.0 {
                norm2(&update) / norm_old
            } else {
                0.0
            };
            if delta_new > delta {
                break;
            }
            delta = delta_new;
        }

        // ---- multiplier updates -------------------------------------------
        // s2 = s2 + Dx - z2
        for i in 0..n {
            s2[i] += dx[i] - z2[i];
        }
        // s3 = exp(1i*z2) - IS + s3 - z3
        for i in 0..n {
            let ez2 = Complex64::new(z2[i].cos(), z2[i].sin());
            s3[i] = ez2 - is[i] + s3[i] - z3[i];
        }
    }

    // Back to ppm.
    for v in &mut x {
        *v /= ppm_to_rad;
    }

    apply_mask_zero(&mut x, mask);
    x
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_l1qsm_zero_field() {
        // Zero field should give (approximately) zero susceptibility.
        let n = 8;
        let field = vec![0.0; n * n * n];
        let mask = vec![1u8; n * n * n];
        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
        let params = L1QsmParams {
            max_iter: 10,
            ..L1QsmParams::default()
        };

        let chi = l1qsm(&field, &mask, &mask.iter().map(|_| 1.0).collect::<Vec<f64>>(), &grid, (0.0, 0.0, 1.0), &params, |_, _| {});

        for &val in chi.iter() {
            assert!(val.abs() < 1e-6, "Zero field should give ~zero chi, got {}", val);
        }
    }

    #[test]
    fn test_l1qsm_finite() {
        // A small ramp field should produce all-finite output.
        let n = 8;
        let field: Vec<f64> = (0..n * n * n).map(|i| (i as f64) * 0.001).collect();
        let mask = vec![1u8; n * n * n];
        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
        let params = L1QsmParams {
            max_iter: 10,
            ..L1QsmParams::default()
        };

        let chi = l1qsm(&field, &mask, &mask.iter().map(|_| 1.0).collect::<Vec<f64>>(), &grid, (0.0, 0.0, 1.0), &params, |_, _| {});

        for (i, &val) in chi.iter().enumerate() {
            assert!(val.is_finite(), "Chi should be finite at index {}", i);
        }
    }
}