Module nextqsm 

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.

Constants

NEXTQSM_LAMBDAS

Trained data-consistency weights λ_k (one per variational step) shipped with the NeXtQSM checkpoint. Scalars, like a normalization constant.

Functions

dipole_forward 🔒

Dipole forward D(y) = real(ifft3(fft3(y) · kernel_shifted)).

dipole_kernel_shifted 🔒

fftshift(get_dipole_kernel_fourier(...)) in column-major layout.

l2 🔒

nextqsm

Run NeXtQSM on a total field map.

nextqsm_padded

NeXtQSM on a grid already sized to a multiple of 64 (no internal padding). Exposed for validation against the reference trajectory.

nextqsm_tiled

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.

pad64 🔒

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).

run_unet 🔒

Run a single-input/single-output 3D U-Net ONNX on a column-major volume.