Skip to main content

Module modl_qsm

Module modl_qsm 

Source
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.py wrote χ 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, checkpoint logs/last.h5 MyLayer; initializer was 4.0, this is the trained value).
MODL_MEAN
Train-set per-channel mean (NormFactor.mat CosTrnMean): [χ33, χ13/23 field].
MODL_STD
Train-set per-channel std (NormFactor.mat CosTrnStd): [χ33, χ13/23 field].
NUM_ITER 🔒

Functions§

a_op 🔒
A operator: 2-channel susceptibility → 1-channel field real(ifft(D·fft(χ33) + fft(χ13/23))). Input is interleaved [ch0, ch1].
ah_op 🔒
A^H operator: 1-channel field φ → 2-channel [real(ifft(D·fft(φ))), φ]. Interleaved output [ch0, ch1] per voxel.
dipole_apply 🔒
D(v) = real(ifft3(fft3(v)·dk)) with dk in fftfreq (unshifted) layout.
modl_dipole_kernel 🔒
MoDL-QSM dipole kernel D = 1/3 − (k·B̂)²/|k|² with k from fftfreq (DC at the array corner, D[0,0,0]=0). Bit-identical to test_tools.dipole_kernel.
modl_qsm
Run MoDL-QSM on a background-removed local field (ppm), column-major (nx,ny,nz). bdir is 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. vol is interleaved [ch0, ch1] per voxel (length 2·n); output is the same interleaved layout.