qsm_core/inversion/

hdqsm.rs

//! HD-QSM: Hybrid data-fidelity two-stage linear dipole inversion.
//!
//! A linear dipole-inversion method that runs in two stages. Stage 1 solves an
//! L1 data-fidelity problem (robust to phase/model errors) and derives a
//! spatially-varying discrepancy weighting. Stage 2 solves an L2 data-fidelity
//! problem reweighted by that discrepancy map. Both stages use ADMM with an
//! L1 total-variation regularizer.
//!
//! # Data weight
//! `w` is the magnitude normalised as in HDQSM.m's usage example,
//! `mask.*mag/max(mag)`, or the mask when there is no magnitude (HDQSM.m: "Not used if not
//! specified"). The magnitude is the paper's Eq. 2 echo combination,
//! `sum(TE_i*mag_i^2) / sum(TE_i*mag_i)` ([`hdqsm_magnitude`]), which the QSM-CI original
//! (HD-QSM c456a94) uses too; the pipeline makes it from every echo's magnitude and requires
//! them (QSM.rs#403: QSMxT gave the port the root sum of squares, part of a 9.1 % / 7.1 % gap
//! on sim / KKI).
//!
//! # Input
//! The field is masked first, as in HDQSM.m's usage example (`params.input =
//! mask.*phase_use/phase_scale`) and the QSM-CI original. It matters where a local field is
//! non-zero outside the mask (stage 1's L1 term and the stage-2 discrepancy see it).
//!
//! # Dipole kernel
//! HDQSM.m takes the kernel as an input and ships none; it is "based on" FANSI, and the
//! original is run with FANSI's `dipole_kernel_angulated` (mode 1, discrete). That is the
//! default here ([`HdQsmParams::kernel`], [`FansiKernel`] has the decision and the
//! measurements); the continuous kernel is an option.
//!
//! # Finite differences
//! As in HDQSM.m, the gradients are unit-spacing forward differences whatever
//! the voxel size; only the dipole kernel uses the voxel size.
//!
//! # Units
//! Like the other inversions, HD-QSM takes the local field in ppm and returns
//! ppm. HDQSM.m's usage example passes `phase/phase_scale`, and the authors'
//! scripts on the 2019 challenge data divide the frequency map by
//! `gamma*B0`, i.e. they work in ppm, so no conversion is made. Stage 2 is not
//! scale invariant, so `alpha_l2` is a value for fields in ppm.
//!
//! Reference:
//! Lambert, M., Tejos, C., Langkammer, C., Milovic, C. (2022).
//! "Hybrid data fidelity term approach for quantitative susceptibility mapping."
//! Magnetic Resonance in Medicine, 88(2):962-972.
//! https://doi.org/10.1002/mrm.29218
//!
//! Reference implementation: HDQSM.m, https://github.com/mglambert/HD-QSM (MIT).
//!
//! # Reference pipeline
//! Lambert et al. 2022, Section 3.3 (in vivo): Siemens 3T, 6 echoes, TE1 = spacing = 4.92 ms,
//! 0.8x0.8x2 mm; SEGUE unwrapping, PDF background removal, a magnitude-weighted least-squares
//! echo fit for the local field, and the residual background removed with WH-QSM's harmonic
//! estimate. For the 2019-challenge simulations the field was the magnitude-weighted
//! least-squares fit, zero-padded to 256^3, with 300 iterations (Section 3.2); this port does not
//! pad. HDQSM.m's header says `tol_update` defaults to 1.0, its code to 1e-2 (L68-72); the code
//! value is used here.
//!
//! # Parity
//! Against the verified original (QSM-CI pair `hdqsm`; FANSI-toolbox 93044a3 / HD-QSM c456a94, 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.0015 %, KKI 0.00086 %, from
//! 9.1 % / 7.1 % (sim / KKI) before. Three causes: the kernel, the Eq. 2 magnitude (QSMxT gave the root sum of squares) and the masked input. What remains is
//! single (MATLAB) vs double (here) precision (HD-QSM works in ppm, so no γ). 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): `hdqsm.bg-pdf`.

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

/// HD-QSM algorithm parameters.
#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
#[derive(Clone, Debug)]
pub struct HdQsmParams {
    /// L2-stage TV weight, for a field in ppm and the magnitude weight (default 3e-5).
    ///
    /// HDQSM.m requires it and gives no default; its example uses 10^-4.785 (1.6e-5) and the
    /// paper's 2019-challenge reconstructions 6.3e-6, both on simulated data. 3e-5 was chosen
    /// from a sweep (1e-6 to 1e-4) on the 2016 challenge (COSMOS reference: best correlation,
    /// 0.76, NRMSE 73%), the 2019-style simulation (NRMSE 39%) and ten travelling-head scans
    /// (stable on all of them, with less noise than the mask-weighted 1e-4 the port used).
    pub alpha_l2: f64,
    /// L2-stage gradient-consistency ADMM weight; `None` is `100 * alpha_l2`.
    ///
    /// The paper ties it to alpha_l2 with 10 <= mu1_l2/alpha_l2 <= 100; HDQSM.m's example and
    /// the paper's challenge reconstructions use 10. On in vivo travelling-head scans the ADMM
    /// diverges (in HDQSM.m as well) once mu1_l2 falls to about 3e-4 or below, which a ratio of
    /// 10 reaches at the alpha_l2 that suits in vivo data, so the default uses the paper's 100.
    pub mu1_l2: Option<f64>,
    /// Fidelity consistency weight (HDQSM.m `mu2L1`, default 1.0).
    pub mu2: f64,
    /// Stage-1 (L1) iterations (HDQSM.m `maxOuterIterL1`, default 20).
    pub max_iter_l1: usize,
    /// Stage-2 (L2) iterations (HDQSM.m `maxOuterIterL2`, default 80).
    pub max_iter_l2: usize,
    /// Stage-2 percent-update stopping tolerance (HDQSM.m `tol_update`, default 1e-2).
    pub tol_update: f64,
    /// The dipole kernel. Default [`FansiKernel::Discrete`], FANSI's
    /// `dipole_kernel_angulated` as the QSM-CI original is run with (HDQSM.m ships no kernel); [`FansiKernel::Continuous`] is
    /// the crate's continuous kernel (see [`FansiKernel`] for the decision and the measurements).
    pub kernel: FansiKernel,
    // Stage-1 alpha_l1 / mu1_l1 are sqrt(alpha_l2) / sqrt(mu1_l2), as in HDQSM.m.
}

impl Default for HdQsmParams {
    fn default() -> Self {
        Self {
            alpha_l2: 3e-5,
            mu1_l2: None,
            mu2: 1.0,
            max_iter_l1: 20,
            max_iter_l2: 80,
            tol_update: 1e-2,
            kernel: FansiKernel::Discrete,
        }
    }
}

/// HD-QSM's magnitude, the paper's Eq. 2 echo combination (Lambert et al. 2022), as the
/// QSM-CI original computes it: `mag_use = sum(TE.*mag.^2, 4) ./ sum(TE.*mag, 4)`, with
/// non-finite values (no signal in any echo) set to 0.
///
/// `echoes` holds every echo's magnitude stacked with the echo slowest (`echo_times.len()`
/// volumes); the result is one volume.
///
/// # Panics
/// If `echoes.len()` is not a multiple of `echo_times.len()`, or there are no echo times.
pub fn hdqsm_magnitude(echoes: &[f64], echo_times: &[f64]) -> Vec<f64> {
    let ne = echo_times.len();
    assert!(ne > 0 && echoes.len().is_multiple_of(ne), "HD-QSM: {} magnitude values for {ne} echoes", echoes.len());
    let n = echoes.len() / ne;
    (0..n)
        .map(|i| {
            let (mut num, mut den) = (0.0, 0.0);
            for (e, &te) in echo_times.iter().enumerate() {
                let m = echoes[e * n + i];
                num += te * m * m;
                den += te * m;
            }
            let v = num / den;
            if v.is_finite() { v } else { 0.0 }
        })
        .collect()
}

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

/// Percent update `100 * ||x - x_prev|| / ||x||`.
#[inline]
fn percent_update(x: &[f64], x_prev: &[f64]) -> f64 {
    let diff: f64 = x.iter().zip(x_prev).map(|(&a, &b)| (a - b) * (a - b)).sum::<f64>().sqrt();
    let nx = norm2(x);
    if nx > 0.0 { 100.0 * diff / nx } else { 0.0 }
}

/// Forward dipole convolution: `Dx = real(ifft(kernel .* fft(x)))`.
///
/// `cbuf` is scratch of length n_total; result written to `out`.
fn apply_dipole(
    fft_ws: &mut crate::fft::Fft3dWorkspace,
    k: &[f64],
    x: &[f64],
    cbuf: &mut [Complex64],
    out: &mut [f64],
) {
    for (c, &xv) in cbuf.iter_mut().zip(x.iter()) {
        *c = Complex64::new(xv, 0.0);
    }
    fft_ws.fft3d(cbuf);
    for (c, &kv) in cbuf.iter_mut().zip(k.iter()) {
        *c *= kv;
    }
    fft_ws.ifft3d(cbuf);
    for (o, c) in out.iter_mut().zip(cbuf.iter()) {
        *o = c.re;
    }
}

/// HD-QSM dipole inversion.
///
/// # Arguments
/// * `local_field` - Local field values in ppm (nx * ny * nz).
/// * `mask` - Binary ROI mask (nx * ny * nz), 1 = inside.
/// * `magnitude` - Magnitude image, the data weight (see module docs; the pipeline passes
///   [`hdqsm_magnitude`]); `None` weights every voxel in the mask by 1
/// * `grid` - Volume grid (dimensions and voxel sizes).
/// * `bdir` - B0 field direction.
/// * `params` - HD-QSM parameters.
/// * `progress` - Progress callback `(global_iter, max_iter_l1 + max_iter_l2)`.
///
/// # Returns
/// Susceptibility map in ppm.
pub fn hdqsm(
    local_field: &[f64],
    mask: &[u8],
    magnitude: Option<&[f64]>,
    grid: &Grid,
    bdir: (f64, f64, f64),
    params: &HdQsmParams,
    mut progress: impl FnMut(usize, usize),
) -> Vec<f64> {
    let n = grid.n_total();
    let mu1_l2 = params.mu1_l2.unwrap_or(100.0 * params.alpha_l2);
    let (mut fft_ws, k) = prepare_fansi_kernel(grid, bdir, params.kernel);
    // HDQSM.m's finite differences are on unit spacing whatever the voxel size
    // (E_j = 1 - exp(2*pi*i*k_j/N_j)); the dipole kernel keeps the voxel size.
    let ee2 = fansi_unit_ee2(grid.dims);
    let ugrid = Grid::new(grid.nx(), grid.ny(), grid.nz(), 1.0, 1.0, 1.0);
    // K2 = abs(kernel).^2 = k^2 (kernel is real). mu*EE2 is added per stage.
    let denom_base: Vec<f64> = k.iter()
        .map(|&kk| 1e-30 + params.mu2 * kk * kk)
        .collect();

    // weight = mask.*(mag/max(mag(:))) (HDQSM.m usage example); the mask when there is no
    // magnitude.
    let mag_max = magnitude.map_or(0.0, |m| m.iter().cloned().fold(0.0f64, f64::max));
    let weight_mask: Vec<f64> = (0..n)
        .map(|i| {
            if mask[i] == 0 {
                return 0.0;
            }
            match magnitude {
                Some(m) if mag_max > 0.0 => m[i] / mag_max,
                _ => 1.0,
            }
        })
        .collect();

    // input = mask.*field (HDQSM.m usage example).
    let masked_field: Vec<f64> = (0..n).map(|i| if mask[i] != 0 { local_field[i] } else { 0.0 }).collect();
    let local_field = &masked_field[..];

    let total_iter = params.max_iter_l1 + params.max_iter_l2;
    let mut global_iter = 0usize;

    // Working buffers.
    let mut x = vec![0.0f64; n];
    let mut x_prev = vec![0.0f64; n];

    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];
    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 z2 = vec![0.0f64; n];
    let mut s2 = vec![0.0f64; n];
    let mut dx = vec![0.0f64; n]; // Dx = real(ifft(K.*Fx))

    let mut div = vec![0.0f64; n];
    let mut gx = vec![0.0f64; n];
    let mut gy = vec![0.0f64; n];
    let mut gz = vec![0.0f64; n];

    let mut cbuf = vec![Complex64::new(0.0, 0.0); n];
    let mut fdiv = vec![Complex64::new(0.0, 0.0); n];
    let mut fdt = vec![Complex64::new(0.0, 0.0); n];

    // ---- Helper closures cannot capture &mut fft_ws + buffers simultaneously,
    //      so the x-subproblem is inlined below. ----

    // =========================================================================
    // STAGE 1: L1 data fidelity.
    // =========================================================================
    {
        let mu = mu1_l2.sqrt();
        let alpha = params.alpha_l2.sqrt();
        let mu2 = params.mu2;
        let ll = alpha / mu;

        // Denominator for stage 1: eps + mu2*K2 + mu*EE2.
        let denom: Vec<f64> = (0..n).map(|i| denom_base[i] + mu * ee2[i]).collect();

        // Wy = input.
        let wy = local_field;

        for t in 0..params.max_iter_l1 {
            x_prev.copy_from_slice(&x);

            // Dt_kspace_arg = z2 - s2 + Wy.
            for i in 0..n {
                gx[i] = z2[i] - s2[i] + wy[i]; // reuse gx as dt_arg scratch
            }
            // fft(dt_arg) -> fdt
            for (c, &v) in fdt.iter_mut().zip(gx.iter()) {
                *c = Complex64::new(v, 0.0);
            }
            fft_ws.fft3d(&mut fdt);

            // fft(bdiv(z_d - s_d)) -> fdiv
            for i in 0..n {
                gx[i] = z_dx[i] - s_dx[i];
                gy[i] = z_dy[i] - s_dy[i];
                gz[i] = z_dz[i] - s_dz[i];
            }
            bdiv_inplace(&mut div, &gx, &gy, &gz, &ugrid);
            for (c, &v) in fdiv.iter_mut().zip(div.iter()) {
                *c = Complex64::new(v, 0.0);
            }
            fft_ws.fft3d(&mut fdiv);

            // num = mu*fdiv + mu2*conj(K)*fdt  (K real -> conj(K)=K); divide by denom.
            // Guard the dipole null-space (DC/singular bins): denom -> ~0 there
            // (dipole kernel and Laplacian both vanish). Zero it instead of
            // dividing FFT round-off by ~0, which would create a huge DC pedestal.
            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 = -fdiv[i] * mu + fdt[i] * (mu2 * k[i]);
                cbuf[i] = if denom[i] > 1e-20 { num / denom[i] } else { Complex64::new(0.0, 0.0) };
            }
            fft_ws.ifft3d(&mut cbuf);
            for i in 0..n {
                x[i] = cbuf[i].re;
            }

            global_iter += 1;
            progress(global_iter, total_iter);

            if t < params.max_iter_l1 - 1 {
                // fgrad(x) -> x_d
                fgrad_inplace(&mut x_dx, &mut x_dy, &mut x_dz, &x, &ugrid);
                // z_d = shrink(x_d + s_d, ll); s_d += x_d - z_d.
                for i in 0..n {
                    let ax = x_dx[i] + s_dx[i];
                    let ay = x_dy[i] + s_dy[i];
                    let az = x_dz[i] + s_dz[i];
                    z_dx[i] = shrink(ax, ll);
                    z_dy[i] = shrink(ay, ll);
                    z_dz[i] = shrink(az, ll);
                    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];
                }

                // Dx = real(ifft(K.*Fx))
                apply_dipole(&mut fft_ws, &k, &x, &mut cbuf, &mut dx);
                // z2_inner = Dx + s2 - Wy; z2 = shrink(z2_inner, weight/mu2); s2 = z2_inner - z2.
                for i in 0..n {
                    let z2_inner = dx[i] + s2[i] - wy[i];
                    z2[i] = shrink(z2_inner, weight_mask[i] / mu2);
                    s2[i] = z2_inner - z2[i];
                }
            }
        }
    }

    // Discrepancy factor: dphi = ifft(fft(input) - fft(x).*K); dphi = abs(dphi).*mask; dphi /= max.
    // = |input - Dx| .* mask, normalized.  (fft(input) - K.*fft(x) -> ifft = input - Dx)
    let mut dphi = vec![0.0f64; n];
    {
        // Dx = real(ifft(K.*Fx)) with current x.
        apply_dipole(&mut fft_ws, &k, &x, &mut cbuf, &mut dx);
        let mut maxv = 0.0f64;
        for i in 0..n {
            let v = (local_field[i] - dx[i]).abs() * (if mask[i] != 0 { 1.0 } else { 0.0 });
            dphi[i] = v;
            if v > maxv {
                maxv = v;
            }
        }
        if maxv > 0.0 {
            for v in dphi.iter_mut() {
                *v /= maxv;
            }
        }
    }

    // =========================================================================
    // STAGE 2: L2 data fidelity, reweighted by discrepancy.
    // =========================================================================
    {
        let mu = mu1_l2;
        let alpha = params.alpha_l2;
        let mu2 = params.mu2;
        let ll = alpha / mu;

        // Denominator for stage 2: eps + mu2*K2 + mu*EE2.
        let denom: Vec<f64> = (0..n).map(|i| denom_base[i] + mu * ee2[i]).collect();

        // weight = weight.*weight.*(1-dphi).*(1-dphi).
        let weight: Vec<f64> = (0..n).map(|i| {
            let w = weight_mask[i];
            let d = 1.0 - dphi[i];
            w * w * d * d
        }).collect();
        // Wy = weight.*input./(weight+mu2).
        let wy: Vec<f64> = (0..n).map(|i| weight[i] * local_field[i] / (weight[i] + mu2)).collect();

        // Reset dual/aux for stage 2 (x carried over).
        for i in 0..n {
            z_dx[i] = 0.0; z_dy[i] = 0.0; z_dz[i] = 0.0;
            s_dx[i] = 0.0; s_dy[i] = 0.0; s_dz[i] = 0.0;
            s2[i] = 0.0;
        }

        for _t in 0..params.max_iter_l2 {
            x_prev.copy_from_slice(&x);

            // z/s updates BEFORE the x-update (unlike stage 1).
            fgrad_inplace(&mut x_dx, &mut x_dy, &mut x_dz, &x, &ugrid);
            for i in 0..n {
                let ax = x_dx[i] + s_dx[i];
                let ay = x_dy[i] + s_dy[i];
                let az = x_dz[i] + s_dz[i];
                z_dx[i] = shrink(ax, ll);
                z_dy[i] = shrink(ay, ll);
                z_dz[i] = shrink(az, ll);
                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];
            }

            // Dx = real(ifft(K.*Fx))
            apply_dipole(&mut fft_ws, &k, &x, &mut cbuf, &mut dx);
            // z2 = Wy + mu2*(Dx + s2)./(weight + mu2); s2 = s2 + Dx - z2.
            for i in 0..n {
                z2[i] = wy[i] + mu2 * (dx[i] + s2[i]) / (weight[i] + mu2);
                s2[i] = s2[i] + dx[i] - z2[i];
            }

            // x update. Dt_kspace_arg = z2 - s2.
            for i in 0..n {
                gx[i] = z2[i] - s2[i];
            }
            for (c, &v) in fdt.iter_mut().zip(gx.iter()) {
                *c = Complex64::new(v, 0.0);
            }
            fft_ws.fft3d(&mut fdt);

            for i in 0..n {
                gx[i] = z_dx[i] - s_dx[i];
                gy[i] = z_dy[i] - s_dy[i];
                gz[i] = z_dz[i] - s_dz[i];
            }
            bdiv_inplace(&mut div, &gx, &gy, &gz, &ugrid);
            for (c, &v) in fdiv.iter_mut().zip(div.iter()) {
                *c = Complex64::new(v, 0.0);
            }
            fft_ws.fft3d(&mut fdiv);

            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 = -fdiv[i] * mu + fdt[i] * (mu2 * k[i]);
                // Guard the dipole null-space (DC/singular bins) — see stage 1.
                cbuf[i] = if denom[i] > 1e-20 { num / denom[i] } else { Complex64::new(0.0, 0.0) };
            }
            fft_ws.ifft3d(&mut cbuf);
            for i in 0..n {
                x[i] = cbuf[i].re;
            }

            global_iter += 1;
            progress(global_iter, total_iter);

            let upd = percent_update(&x, &x_prev);
            if upd < params.tol_update {
                break;
            }
        }
    }

    // Ensure progress reaches the total even if stage 2 stopped early.
    if global_iter < total_iter {
        progress(total_iter, total_iter);
    }

    apply_mask_zero(&mut x, mask);
    x
}

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

    /// Eq. 2: sum(TE*mag^2)/sum(TE*mag) per voxel, 0 where there is no signal.
    #[test]
    fn test_hdqsm_magnitude_eq2() {
        // two voxels, two echoes (echo slowest)
        let m = hdqsm_magnitude(&[2.0, 0.0, 1.0, 0.0], &[0.01, 0.02]);
        let want = (0.01 * 4.0 + 0.02 * 1.0) / (0.01 * 2.0 + 0.02 * 1.0);
        assert!((m[0] - want).abs() < 1e-15);
        assert_eq!(m[1], 0.0);
    }

    #[test]
    fn test_hdqsm_zero_field() {
        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 = HdQsmParams {
            max_iter_l1: 5,
            max_iter_l2: 10,
            ..HdQsmParams::default()
        };

        let chi = hdqsm(&field, &mask, None, &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_hdqsm_finite() {
        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 = HdQsmParams {
            max_iter_l1: 5,
            max_iter_l2: 10,
            ..HdQsmParams::default()
        };

        let chi = hdqsm(&field, &mask, None, &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);
        }
    }

    #[test]
    fn test_hdqsm_magnitude_weight() {
        let n = 8;
        let field: Vec<f64> = (0..n * n * n).map(|i| ((i % 7) as f64 - 3.0) * 0.01).collect();
        let mask = vec![1u8; n * n * n];
        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
        let params = HdQsmParams { max_iter_l1: 5, max_iter_l2: 10, ..HdQsmParams::default() };
        let run = |mag: Option<&[f64]>| hdqsm(&field, &mask, mag, &grid, (0.0, 0.0, 1.0), &params, |_, _| {});

        // A flat magnitude normalises to 1 everywhere: the same as no magnitude (mask weight).
        let flat = vec![5.0; n * n * n];
        let a = run(None);
        let b = run(Some(&flat));
        assert!(a.iter().zip(&b).all(|(x, y)| (x - y).abs() < 1e-12));

        // A varying magnitude changes the weight, and so the result.
        let ramp: Vec<f64> = (0..n * n * n).map(|i| 1.0 + (i % n) as f64).collect();
        let c = run(Some(&ramp));
        assert!(c.iter().all(|v| v.is_finite()));
        assert!(a.iter().zip(&c).any(|(x, y)| (x - y).abs() > 1e-9));
    }
}