qsm_core/inversion/

nextqsm.rs

//! NeXtQSM single-step deep-learning reconstruction (`onnx` feature).
//!
//! NeXtQSM (Cognolato 2023) is a *hybrid* method: a background-removal U-Net
//! followed by a **6-step variational** dipole inversion. Each variational step
//! is a gradient-descent update
//! `x ← x − (λ_k·∇E_D + ∇E_R)·mask`, where `E_D` is an RMSE data-consistency
//! term through the FFT dipole forward and `E_R = mean(|VarNet(x)|)` is a learned
//! regularizer whose gradient is a backprop through the VarNet U-Net.
//!
//! The two U-Net pieces run as ONNX (`nextqsm-bf.onnx` = BFR forward;
//! `nextqsm-vjp.onnx` = the regularizer gradient `∇ₓ mean(|VarNet(x)|)`, hand-coded
//! as a forward graph). The FFT data-consistency gradient and the unroll live here
//! in Rust (`rustfft`), since tract can't do in-graph FFT.
//!
//! Weights are not bundled; the caller passes both ONNX byte buffers.
//!
//! # Reference pipeline
//!
//! From Cognolato et al., *Medical Image Analysis* 84:102700 (2023; arXiv:2107.07752) and the
//! authors' package (`QSMxT/NeXtQSM`):
//!
//! - **Training data:** hybrid 1 mm isotropic brains from 3T MP2RAGE segmentations with simulated
//!   susceptibilities (paper §3.1); tested on QSM challenge 2.0 data and 0.75 mm 7T data, with an
//!   FSL BET mask (paper §3.2); the paper reports robustness at 0.64 and 0.95 mm (§4.2).
//! - **Input:** "an unwrapped frequency map (unitless and scaled to ppm) and brain mask"
//!   (README) — a total field in ppm; the package masks it and zero-pads to a multiple of 64
//!   (`processing/data_loader.py`). The unwrapping method is not specified.
//! - **Resampling:** none; the data-consistency dipole kernel uses the NIfTI voxel size
//!   (`predict_all.py:60-61`).
//! - **Padding:** an odd axis gets one zero voxel in front, then each axis is padded equally on
//!   both sides to a multiple of 64 (`tf_utils/misc.py` `add_padding`/`remove_padding`).
//!
//! # Parity
//!
//! Measured against the `nextqsm` 1.0.6 package (QSM-CI's `nextqsm` entry, `predict_all.py`, run
//! on the same CPU) on in-vivo boxes of `prisma-bridge-run1`, all at 1 mm (QSM.rs#330):
//!
//! - Per step, given the same inputs, the BFR output, the data-consistency gradient and the
//!   regulariser gradient each match TensorFlow to ~1e-6 relative.
//! - End to end: 0.13 % NRMSE (r 0.9999992, max |Δ| 0.045 ppm) on 36×32×28, and 0.038 % (r
//!   0.9999999, max |Δ| 0.022 ppm) on 35×31×27 (odd axes). Before the padding fix below, the odd
//!   box was 70.6 % (r 0.73).
//! - That residual is the method's own sensitivity to float32 rounding, not a port difference:
//!   the regulariser's gradient goes through `sign(VarNet(x))`, so voxels whose output is near
//!   zero flip sign between implementations. The package itself moves by 0.10 % NRMSE (max |Δ|
//!   0.058 ppm) on the 36×32×28 box between an AVX-512 Intel and an AVX2 AMD CPU, and the same
//!   glue on ONNX Runtime instead of tract lands 0.0057 % from it.
//!
//! # Agreement with the original
//!
//! `tests/models_onnx/field.rs` runs it on a 36×32×28 box of in-vivo data (1 mm, 2.89 T, across the
//! brain's edge) against its original's verified QSM-CI entry, nextqsm 1.0.6 (`nextqsm`)
//! (QSMxT/QSM-CI#316; the fixture is regenerated in Actions by
//! `tests/data/dl_field_qsmci_small.py`): **NRMSE 0.12 %**, r 0.9999993, max |Δ| 0.058 ppm
//! (measured 2026-10-30; 0.97 % on QSM-CI's 128×160×96 case, 3.1 % on a 64³ box). Not rounding, but
//! the size of the package's own CPU-to-CPU spread (0.10 %, max |Δ| 0.058 ppm on this box, under
//! Parity above): the method's sensitivity to float32 rounding, not a port difference
//! (astewartau/QSM.rs#330). The test holds it just above that floor.
//!
//! # 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): `nextqsm.mask-bet`.

use num_complex::Complex64;

use crate::fft::{fft3d_real, ifft3d_real};
use crate::grid::Grid;
use crate::models::onnx::{OnnxModel, OnnxError, Tensor};

/// Trained data-consistency weights `λ_k` (one per variational step) shipped with
/// the NeXtQSM checkpoint. Scalars, like a normalization constant.
pub const NEXTQSM_LAMBDAS: [f64; 6] =
    [15.141516, 15.005543, 16.733437, 14.574243, 7.072418, 0.7588475];

/// Run NeXtQSM on a total field map.
///
/// * `total_field` — total field (ppm), column-major `(nx,ny,nz)`.
/// * `mask` — binary brain mask (same layout).
/// * `bdir` — B0 direction; `bf_onnx`/`vjp_onnx` — the two exported graphs.
///
/// The volume is zero-padded to a multiple of 64 (six pooling levels) as the package pads it,
/// reconstructed, and cropped back. Returns susceptibility (ppm), masked.
pub fn nextqsm(
    total_field: &[f64],
    mask: &[u8],
    grid: &Grid,
    bdir: (f64, f64, f64),
    bf_onnx: &[u8],
    vjp_onnx: &[u8],
) -> Result<Vec<f64>, OnnxError> {
    let (nx, ny, nz) = grid.dims;
    let n = nx * ny * nz;
    assert_eq!(total_field.len(), n);
    assert_eq!(mask.len(), n);

    let (bx, px) = pad64(nx);
    let (by, py) = pad64(ny);
    let (bz, pz) = pad64(nz);
    let pgrid = Grid { dims: (px, py, pz), voxel_size: grid.voxel_size };
    let np = px * py * pz;

    let mut field_p = vec![0.0f64; np];
    let mut mask_p = vec![0u8; np];
    for z in 0..nz {
        for y in 0..ny {
            for x in 0..nx {
                let s = x + nx * (y + ny * z);
                let d = (x + bx) + px * ((y + by) + py * (z + bz));
                field_p[d] = total_field[s];
                mask_p[d] = mask[s];
            }
        }
    }

    let chi_p = nextqsm_padded(&field_p, &mask_p, &pgrid, bdir, bf_onnx, vjp_onnx, &NEXTQSM_LAMBDAS)?;

    // Crop back.
    let mut chi = vec![0.0f64; n];
    for z in 0..nz {
        for y in 0..ny {
            for x in 0..nx {
                chi[x + nx * (y + ny * z)] = chi_p[(x + bx) + px * ((y + by) + py * (z + bz))];
            }
        }
    }
    Ok(chi)
}

/// Where an axis of `s` voxels starts in its padded copy, and the padded length (the next
/// multiple of 64), as the package pads it (`tf_utils/misc.py` `add_padding`): an odd axis first
/// gets one zero voxel in front, then each axis is padded equally on both sides. So the volume
/// starts at ⌈total/2⌉, not ⌊total/2⌋, on odd axes. The networks pool six times, so a one-voxel
/// shift moves every pooling window: on a 35×31×27 in-vivo box, ⌊total/2⌋ gave 70.6 % NRMSE
/// against the package and ⌈total/2⌉ gives 0.038 % (QSM.rs#330).
fn pad64(s: usize) -> (usize, usize) {
    let total = (64 - s % 64) % 64;
    (total.div_ceil(2), s + total)
}

/// NeXtQSM on a grid already sized to a multiple of 64 (no internal padding).
/// Exposed for validation against the reference trajectory.
#[allow(clippy::too_many_arguments)]
pub fn nextqsm_padded(
    field: &[f64],
    mask: &[u8],
    grid: &Grid,
    bdir: (f64, f64, f64),
    bf_onnx: &[u8],
    vjp_onnx: &[u8],
    lambdas: &[f64],
) -> Result<Vec<f64>, OnnxError> {
    let (nx, ny, nz) = grid.dims;
    let n = nx * ny * nz;
    let maskf: Vec<f64> = mask.iter().map(|&m| m as f64).collect();

    // fftshift(dipole kernel) in the crate's column-major layout.
    let kernel = dipole_kernel_shifted(grid, bdir);

    // bf_logits = BFRnet(field·mask)·mask
    let bf = OnnxModel::load(bf_onnx)?;
    let masked_field: Vec<f64> = field.iter().zip(&maskf).map(|(&f, &m)| f * m).collect();
    let bf_out = run_unet(&bf, &masked_field, grid)?;
    let bf_logits: Vec<f64> = bf_out.iter().zip(&maskf).map(|(&v, &m)| v * m).collect();
    let norm_bf = l2(&bf_logits);

    let vjp = OnnxModel::load(vjp_onnx)?;
    let mut x = bf_logits.clone();
    for &lam in lambdas.iter() {
        let dx = dipole_forward(&x, &kernel, grid);
        // u = bf_logits − D(x);  ∇E_D = 100/(‖bf‖·‖u‖) · D(D(x) − bf)
        let u: Vec<f64> = bf_logits.iter().zip(&dx).map(|(&b, &d)| b - d).collect();
        let norm_u = l2(&u);
        let resid: Vec<f64> = dx.iter().zip(&bf_logits).map(|(&d, &b)| d - b).collect();
        let d_resid = dipole_forward(&resid, &kernel, grid);
        let scale = if norm_bf > 0.0 && norm_u > 0.0 { 100.0 / (norm_bf * norm_u) } else { 0.0 };

        // The VJP graph returns the unnormalized gradient ∇ₓ Σ|VarNet(x)|; divide
        // by N to get the mean-gradient ∇ₓ mean(|VarNet(x)|) the regularizer uses.
        let grad_r = run_unet(&vjp, &x, grid)?;
        let inv_n = 1.0 / (n as f64);

        for i in 0..n {
            let de = lam * (scale * d_resid[i]) + grad_r[i] * inv_n;
            x[i] = (x[i] - de) * maskf[i];
        }
    }
    Ok(x)
}

/// Dipole forward `D(y) = real(ifft3(fft3(y) · kernel_shifted))`.
pub(crate) fn dipole_forward(y: &[f64], kernel_shifted: &[f64], grid: &Grid) -> Vec<f64> {
    let (nx, ny, nz) = grid.dims;
    let mut spec = fft3d_real(y, nx, ny, nz);
    for (s, &k) in spec.iter_mut().zip(kernel_shifted) {
        *s *= Complex64::new(k, 0.0);
    }
    ifft3d_real(&spec, nx, ny, nz)
}

/// `fftshift(get_dipole_kernel_fourier(...))` in column-major layout.
pub(crate) fn dipole_kernel_shifted(grid: &Grid, bdir: (f64, f64, f64)) -> Vec<f64> {
    let (nx, ny, nz) = grid.dims;
    let (vx, vy, vz) = grid.voxel_size;
    let eps = f32::EPSILON as f64;
    let (bx, by, bz) = bdir;
    let mut k = vec![0.0f64; nx * ny * nz];
    // Centered grid, then fftshift => index shift by half in each dim.
    let sh = |i: usize, nsz: usize| (i + nsz / 2) % nsz; // fftshift index
    for z in 0..nz {
        for y in 0..ny {
            for x in 0..nx {
                // centered coordinate = (idx − n/2); normalized by (n/2) and (2·vox)
                let rx = ((x as f64) - (nx as f64) / 2.0) / ((nx as f64) / 2.0) / (2.0 * vx);
                let ry = ((y as f64) - (ny as f64) / 2.0) / ((ny as f64) / 2.0) / (2.0 * vy);
                let rz = ((z as f64) - (nz as f64) / 2.0) / ((nz as f64) / 2.0) / (2.0 * vz);
                let r2 = rx * rx + ry * ry + rz * rz;
                let dot = rx * bx + ry * by + rz * bz;
                let val = 1.0 / 3.0 - (dot * dot) / (r2 + eps);
                // write to fftshifted position
                let (sx, sy, sz) = (sh(x, nx), sh(y, ny), sh(z, nz));
                k[sx + nx * (sy + ny * sz)] = val;
            }
        }
    }
    k
}

/// Run a single-input/single-output 3D U-Net ONNX on a column-major volume.
fn run_unet(model: &OnnxModel, vol: &[f64], grid: &Grid) -> Result<Vec<f64>, OnnxError> {
    let (nx, ny, nz) = grid.dims;
    let mut input = vec![0.0f32; nx * ny * nz];
    for z in 0..nz {
        for y in 0..ny {
            for x in 0..nx {
                input[(x * ny + y) * nz + z] = vol[x + nx * (y + ny * z)] as f32;
            }
        }
    }
    let out = model.run_single(&Tensor::new(vec![1, 1, nx, ny, nz], input))?;
    let mut v = vec![0.0f64; nx * ny * nz];
    for z in 0..nz {
        for y in 0..ny {
            for x in 0..nx {
                v[x + nx * (y + ny * z)] = out.data[(x * ny + y) * nz + z] as f64;
            }
        }
    }
    Ok(v)
}

fn l2(v: &[f64]) -> f64 {
    v.iter().map(|&a| a * a).sum::<f64>().sqrt()
}

/// Memory-bounded NeXtQSM via whole-algorithm overlap-tiling — the full VarNet gradient-descent
/// loop runs on each patch sub-volume (each internally padded to /64). NeXtQSM's forward model is
/// global, so tiling is **strongly off-design** and approximate; prefer a whole-volume run (e.g.
/// QSMxT). See [`crate::inversion::tiled::tiled_volume_algorithm`].
#[allow(clippy::too_many_arguments)]
pub fn nextqsm_tiled(
    total_field: &[f64],
    mask: &[u8],
    grid: &Grid,
    bdir: (f64, f64, f64),
    bf_onnx: &[u8],
    vjp_onnx: &[u8],
    cfg: &super::tiled::TileConfig,
    progress: impl FnMut(usize, usize),
) -> Result<Vec<f64>, OnnxError> {
    super::tiled::tiled_volume_algorithm(
        total_field, mask, grid, &super::tiled::tile_cost_of("nextqsm"), cfg,
        |f, m, g| nextqsm(f, m, g, bdir, bf_onnx, vjp_onnx),
        progress,
    )
}

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

    #[test]
    fn pads_like_the_package() {
        // Even axes: equal padding on both sides.
        assert_eq!(pad64(64), (0, 64));
        assert_eq!(pad64(36), (14, 64));
        assert_eq!(pad64(160), (16, 192));
        // Odd axes: one voxel in front first, then equal padding (misc.add_padding).
        assert_eq!(pad64(35), (15, 64));
        assert_eq!(pad64(27), (19, 64));
        assert_eq!(pad64(63), (1, 64));
        assert_eq!(pad64(65), (32, 128));
    }
}