qsm_core/inversion/
ir2qsm.rs

1//! IR2QSM dipole inversion (`onnx` feature).
2//!
3//! IR2QSM (Li et al., Med. Phys. 2025; arXiv:2406.12300) maps the **local (tissue)
4//! field in ppm** directly to **susceptibility in ppm** with a single "IR2U-net": a
5//! 3D U-net (`depth=4`) run for `iterations=4` unrolled passes with reverse
6//! concatenations and a recurrent SRU middle module, then a learned integration of
7//! the four per-iteration residual estimates (`latest_out`, the network's final
8//! product). Unlike LPCNN there is no separate physics/unroll here — the **entire**
9//! network is the exported ONNX graph, and this glue only does the surrounding I/O
10//! pipeline:
11//!
12//! 1. **No normalization.** IR2QSM consumes the ppm local field directly and emits
13//!    ppm susceptibility — no dataset mean/std, no `norm_factor` (there are no baked
14//!    scalar constants). Matches `IR2QSM/Evaluate/test_util.py`.
15//! 2. **Zero-pad to a multiple of 8.** The U-net has 3 pool/deconv levels, so each
16//!    spatial dim must be divisible by `2³ = 8`; we center-pad exactly as the
17//!    reference `zero_padding(image, 8)` (low offset `= ceil((target − shape)/2)`),
18//!    run the net, then crop back.
19//! 3. **Mask.** The result is multiplied by the supplied brain mask on the original
20//!    grid.
21//!
22//! **Determinism.** `IR2Unet.forward` has an ungated inference-time `AddNoise` in the
23//! decoder (`torch.rand(1) > 0.3` per iteration) that makes plain PyTorch inference
24//! mildly stochastic. The ONNX was exported with that call pinned to its noise-free
25//! branch (identity), so this net is deterministic. See `export_ir2qsm.py`.
26//!
27//! Input is the background-removed local field in ppm (single orientation; the net is
28//! orientation-agnostic — no dipole kernel, no `b_vec`). Weights are not bundled; the
29//! caller passes `ir2qsm.onnx`.
30
31use crate::grid::Grid;
32use crate::models::onnx::{OnnxError, OnnxModel, Tensor};
33
34/// U-net pool/deconv depth requirement: spatial dims are padded to a multiple of this.
35const SIZE_DIVISOR: usize = 8;
36
37/// Run IR2QSM on a background-removed local field (ppm), column-major `(nx,ny,nz)`.
38/// Returns χ (ppm), masked, on the original grid.
39pub fn ir2qsm(
40    local_field_ppm: &[f64],
41    mask: &[u8],
42    grid: &Grid,
43    onnx: &[u8],
44) -> Result<Vec<f64>, OnnxError> {
45    let (nx, ny, nz) = grid.dims;
46    let n = nx * ny * nz;
47    assert_eq!(local_field_ppm.len(), n, "field length must match grid");
48    assert_eq!(mask.len(), n, "mask length must match grid");
49
50    let model = OnnxModel::load(onnx)?;
51
52    // Center zero-pad each dim up to a multiple of 8 (low offset = ceil(pad/2)).
53    let pad = |d: usize| -> (usize, usize) {
54        let target = d.div_ceil(SIZE_DIVISOR) * SIZE_DIVISOR;
55        let total = target - d;
56        let lo = total.div_ceil(2); // ceil((target - d)/2), matching the reference
57        (lo, total - lo)
58    };
59    let (px0, _px1) = pad(nx);
60    let (py0, _py1) = pad(ny);
61    let (pz0, _pz1) = pad(nz);
62    let (pnx, pny, pnz) =
63        (nx.div_ceil(SIZE_DIVISOR) * SIZE_DIVISOR,
64         ny.div_ceil(SIZE_DIVISOR) * SIZE_DIVISOR,
65         nz.div_ceil(SIZE_DIVISOR) * SIZE_DIVISOR);
66
67    // Build the padded NCDHW (row-major) input from the column-major volume.
68    let mut input = vec![0.0f32; pnx * pny * pnz];
69    for z in 0..nz {
70        for y in 0..ny {
71            for x in 0..nx {
72                let src = x + nx * (y + ny * z);
73                let (px, py, pz) = (x + px0, y + py0, z + pz0);
74                input[((px * pny) + py) * pnz + pz] = local_field_ppm[src] as f32;
75            }
76        }
77    }
78
79    let out = model.run_single(&Tensor::new(vec![1, 1, pnx, pny, pnz], input))?;
80
81    // Crop back to the original grid, repack to column-major, and mask.
82    let mut chi = vec![0.0f64; n];
83    for z in 0..nz {
84        for y in 0..ny {
85            for x in 0..nx {
86                let dst = x + nx * (y + ny * z);
87                let (px, py, pz) = (x + px0, y + py0, z + pz0);
88                let v = out.data[((px * pny) + py) * pnz + pz] as f64;
89                chi[dst] = if mask[dst] != 0 { v } else { 0.0 };
90            }
91        }
92    }
93    Ok(chi)
94}
95
96/// Memory-bounded IR2QSM via overlap-tiling — the IR2U-net run patch-by-patch (for 32-bit WASM,
97/// where whole-volume [`ir2qsm`] overflows the heap on clinical data). No normalization; the
98/// net's 3 pool levels require a `/8` patch. Approximates whole-volume up to tile-boundary
99/// error; see [`crate::inversion::tiled`].
100pub fn ir2qsm_tiled(
101    local_field_ppm: &[f64],
102    mask: &[u8],
103    grid: &Grid,
104    onnx: &[u8],
105    cfg: &super::tiled::TileConfig,
106    progress: impl FnMut(usize, usize),
107) -> Result<Vec<f64>, OnnxError> {
108    let model = OnnxModel::load(onnx)?;
109    super::tiled::tiled_field_inversion(
110        local_field_ppm, mask, grid, &model, SIZE_DIVISOR, cfg, |v| v as f32, |o| o as f64, progress,
111    )
112}
qsm_core/inversion/ir2qsm.rs

qsm_core/inversion/
ir2qsm.rs
