Expand description
MoDL-QSM dipole inversion (onnx feature).
MoDL-QSM (Feng 2021) solves the dipole deconvolution by unrolling a
model-based gradient descent (num_iter = 3): each iteration takes a
k-space data-consistency step with the dipole kernel D (the A/A^H
operators), then applies a learned 3D-CNN prior. Only the CNN prior is a
network (exported ONNX — a conv/BN/ReLU residual stack, 2-channel in/out); the
FFT data-consistency, the unroll, the learned step size alpha, and the
per-channel mean/std normalization live here in Rust.
§Field & susceptibility layout
Input is the background-removed local field in ppm (single orientation).
MoDL-QSM’s phi input is the tissue field normalized to ppm (the repo’s example
test_data.mat fields span ~±0.1–0.2 ppm), so the field is consumed directly.
The network is STI-flavored: it works on a 2-channel susceptibility estimate
— channel 0 is the STI tensor component χ33 (comparable to scalar QSM) and
channel 1 is the field induced by the χ13/χ23 terms. The A^H operator maps a
1-channel field to 2 channels [ifft(D·fft(φ)), φ]; the A operator maps the
2-channel susceptibility back to a 1-channel field ifft(D·fft(χ33) + fft(χ13/23)).
We return channel 0 (χ33) as the QSM susceptibility, masked (matching the
authors’ recon.py, which keeps Y[...,0]).
§FFT convention
MoDL-QSM’s A/A^H use an ortho FFT (fft/√N, ifft·√N); the √N factors
cancel through the linear (i)FFT for both operators, so we use the crate’s
standard normalized fft3d_real/ifft3d_real pair directly. The dipole kernel
D = 1/3 − (k·B̂)²/|k|² uses the fftfreq convention (DC at the array corner,
D[0,0,0]=0) — bit-identical to test_tools.dipole_kernel (which builds a
centered D then fftshifts it). This is the same kernel LPCNN uses.
§Grid requirements
model_test requires an isotropic 1 mm grid with even dimensions (odd dims are
cropped, non-unit voxels are k-space interpolated to 1 mm before inference). The
Rust glue therefore assumes 1 mm even-dimension input (the QSM-CI/dev grid);
callers needing resampling should do so upstream. Weights are not bundled; the
caller passes modl-qsm.onnx.
§Parity
Measured against the authors’ Keras graph (define_generator + model_test, QSM-CI’s
modl-qsm image, MoDL-QSM 111a493, logs/last.h5) on a 36×32×28 in-vivo box of
prisma-bridge-run1 (QSM.rs#330):
- Every intermediate of the unroll (
A^H φ, the step, the normalisation, each CNN-prior pass,-α A^H A x) matches the Keras layer outputs to 1e-7–3e-6 relative. - End to end: 0.0002 % NRMSE (r 1.0000000, max |Δ| 3.2e-7 ppm).
- QSM-CI’s
recon.pywrote χ through the mask’s uint8 NIfTI header, so nibabel stored it in 8 bits; that quantisation was the whole 1.49 % gap first measured against that entry. The port does not reproduce it.
§Reference pipeline
The authors’ test/test_tools.py:210-249 (model_test) is what the grid requirements above
paraphrase: inter_data zero-pads/truncates k-space to 1 mm before the net and back to the
input voxel size afterwards, then masks. The README describes the input as the “tissue phase”;
test/test_demo.py feeds the shipped phi with voxel size [1,1,1] and B0 [0,0,1].
§Agreement with the original
tests/models_onnx/dipole.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, MoDL-QSM 111a493 (modl-qsm)
(QSMxT/QSM-CI#316; the fixture is regenerated in Actions by
tests/data/dl_dipole_qsmci_small.py, from the entry that writes χ as float32,
QSMxT/QSM-CI#331): NRMSE 0.0002 %, r 1.0000000, max |Δ| 3.2e-7 ppm (measured 2026-11-02),
float32 rounding. The 1.5 % first measured was the entry’s uint8 output (above;
astewartau/QSM.rs#330).
Constants§
- MODL_
ALPHA - Learned data-consistency step size (
Alpha, checkpointlogs/last.h5MyLayer; initializer was 4.0, this is the trained value). - MODL_
MEAN - Train-set per-channel mean (
NormFactor.matCosTrnMean): [χ33, χ13/23 field]. - MODL_
STD - Train-set per-channel std (
NormFactor.matCosTrnStd): [χ33, χ13/23 field]. - NUM_
ITER 🔒
Functions§
- a_op 🔒
Aoperator: 2-channel susceptibility → 1-channel fieldreal(ifft(D·fft(χ33) + fft(χ13/23))). Input is interleaved[ch0, ch1].- ah_op 🔒
A^Hoperator: 1-channel field φ → 2-channel[real(ifft(D·fft(φ))), φ]. Interleaved output[ch0, ch1]per voxel.- dipole_
apply 🔒 D(v) = real(ifft3(fft3(v)·dk))withdkin fftfreq (unshifted) layout.- modl_
dipole_ 🔒kernel - MoDL-QSM dipole kernel
D = 1/3 − (k·B̂)²/|k|²withkfromfftfreq(DC at the array corner,D[0,0,0]=0). Bit-identical totest_tools.dipole_kernel. - modl_
qsm - Run MoDL-QSM on a background-removed local field (ppm), column-major
(nx,ny,nz).bdiris the B0 direction (for the dipole kernel). Returns χ33 (ppm), masked. - modl_
qsm_ tiled - Memory-bounded MoDL-QSM via whole-algorithm overlap-tiling — the full unrolled net runs on
each patch sub-volume. MoDL’s data-consistency step is a global k-space operation, so tiling
is strongly off-design and approximate; prefer a whole-volume run (e.g. QSMxT). See
crate::inversion::tiled::tiled_volume_algorithm. - run_
prior 🔒 - Run the 2-in / 2-out CNN prior on a column-major volume.
volis interleaved[ch0, ch1]per voxel (length2·n); output is the same interleaved layout.