Module amp_pe 

AMP-PE: Approximate Message Passing with built-in Parameter Estimation for QSM.

Nonlinear dipole inversion solved with Generalized Approximate Message Passing (GAMP) over a linearized wrapped-phase (complex-exponential) forward model, using a Laplace sparse-wavelet prior and a Gaussian-mixture noise model that absorbs phase outliers.

Ported from the reference MATLAB implementation (https://github.com/EmoryCN2L/QSM_AMP_PE, Huang et al., Magn. Reson. Med. 2023) as packaged for the QSM-CI dipole stage (recon.m + amp_pe_mri_qsm_awgn.m / amp_pe_mri_qsm_awgn_mix.m). This is the dipole-inversion stage only: the local (tissue) field is provided, turned into a simulated single-echo phase, and fed to the two-step AMP-PE solve (single-Gaussian warm-up → Gaussian-mixture final).

Design notes / fidelity

• The forward dipole operator, wavelet transform (periodic db1/db2), and erfcx all match the MATLAB reference (see utils::wavelet, utils::special).
• The reference estimates operator Frobenius norms with 2 random probes. Here they are computed exactly (the wavelet transform is orthonormal, so its norm is sqrt(coef_len); the weighted-dipole norm has a closed form). This is deterministic and a strict improvement; the verification harness injects the same exact values into MATLAB so the two agree to numerical precision.
• Input local field is ppm-scale (crate convention); mut_cst converts ppm to radians at the simulated echo time.
• The L2 (chiL2) seed is computed on the field masked to the ROI (as recon.m does) — background outside the mask must not enter the seed.
• Verified against the MATLAB reference to ~2e-10 on the real 164x205x205 phantom when both use a double-precision seed. The reference’s dipole_kernel_angulated casts the seed kernel to single; in phase-wrapping regions (|phase|>pi) the nonlinear solve is multistable, so that ~1e-7 seed perturbation can select a different local solution there. This crate uses full double precision (a strict improvement).
• Against the unmodified original, the verified QSM-CI amp-pe container (QSMxT/QSM-CI#309, bit-identical to licensed MATLAB), tests/amppe_qsmci.rs measures NRMSE 1.09 % (30×28×22, db1, nlevel 3, with magnitude) and 1.41 % (32×32×24, db2, nlevel 2, oblique B0), 6 linearisations per stage (2026-10-30; the fixture is regenerated in Actions by tests/data/amp_pe_qsmci_small.py). The exact norms above are the likely cause: scaling ‖A_qsm‖²_F alone by 1 % moves the result by 0.3–0.45 % NRMSE, and two random probes miss it by more. Decision (Ashley Stewart, 2026-10-30, astewartau/QSM.rs#363): the exact norms are an unequivocal improvement on a two-probe random estimate, so they stay, with no option for the original’s probes (MATLAB’s randn after rng('default')). The documented parity with the original is these ~1.1–1.4 %, held at that size by tests/amppe_qsmci.rs and by the amp-pe parity pair in QSM-CI.

Reference pipeline

The authors’ scripts (https://github.com/EmoryCN2L/QSM_AMP_PE, commit 15bf8dc; README “Overview of the pipeline”, and qsm_multi_echo_combined.m, which the README says to try first):

• Mask: BET (MEDI-toolbox wrapper, default threshold 0.5) on the root-sum-of-squares magnitude over echoes, eroded by 2 mm (MaskErode, an SMV erosion in mm), and by another 2 mm after background removal (L71-81, L145-147, erosion_width_1/2 = 2).
• Unwrapping: MEDI’s Fit_ppm_complex_TE for the initial phase, which is removed, then 3D best-path unwrapping of each echo (L72, L83-96).
• Echo combination: the plain mean over echoes of phase_e / TE_e (L100-104), multiplied by simulated_TE = 8 ms to give one simulated phase (L33, L107); the input here mirrors that.
• Background removal: PDF (their zero-padding implementation by default, pdf_method = 1, 100 iterations, tolerance 1e-3; L37-39, L112-139).
• Data weight: the first echo’s magnitude in the combined script (L161-169), as the pipeline note in run_dipole_inversion says.
• Scan assumptions in the script’s comments: “Ideally, the voxel size should be isotropic” (L31) and a straight acquisition (B0 along z) is recommended (L34); the README advises trying both db1 and db2 wavelets, db1 being at risk of pixelation when B0 is tilted.

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): amp-pe.mask-bet, amp-pe.echo-combination, amp-pe.echo-weights, amp-pe.bg-pdf.

Structs

AmpPeParams

Parameters for the AMP-PE inversion.

DipoleOp 🔒

Dipole forward/adjoint operator over a padded grid, with mask gather/scatter.

GampCfg 🔒

Static GAMP configuration.

GampSnapshot 🔒

A restore point for backing out of a divergent GAMP sweep.

GampState 🔒

Persistent GAMP state carried across linearization iterations.

MixState 🔒

Gaussian-mixture noise parameters (single component + outlier component).

Constants

DAMP_DIVERGE_FACTOR 🔒

Cost blow-up factor (relative to the best fit so far) treated as divergence.

DAMP_GROW 🔒

Factor by which damping relaxes back towards the configured value after an accepted sweep.

DAMP_MIN 🔒

Damping floor; below this the iteration is making no progress and we stop with the last good state.

DAMP_SHRINK 🔒

Factor applied to the damping when a sweep is rejected.

EPS 🔒

MAX_BACKTRACKS 🔒

Cap on rejected sweeps per GAMP call, so a pathological case cannot spin.

MAX_NLEVEL

Largest nlevel whose padding multiple 2^nlevel fits in a usize: 63 on 64-bit targets, 31 on wasm32.

PARAM_CHUNK 🔒

Chunk length for the deterministic parallel reductions in the GAMP parameter estimation. Fixed size ⇒ the result is independent of the thread count.

Functions

a_qsm_mult 🔒

A_qsm forward: image -> masked complex measurement.

a_qsm_multtr_real 🔒

real(A_qsm^H * s): masked complex -> full-volume real.

amp_pe

AMP-PE nonlinear dipole inversion.

build_wave_mask 🔒

Build the morphology mask: coefficients whose magnitude exceeds a cumulative-energy threshold of the magnitude image’s wavelet coefficients.

chi_l2_seed 🔒

Direct L2 (Tikhonov, gradient) QSM solver used only to seed distributions. Matches chiL2.m: real(ifftn(conj(K) fftn(phase) / (|K|^2 + beta*E2 + eps))) * mask.

crop_volume 🔒

data_fit 🔒

Plain data fit ||A x - y||^2 / m, the quantity the divergence guard watches.

frob_qsm_sq 🔒

Frobenius norm squared of the weighted-dipole operator A_qsm.

gamp_awgn 🔒

One GAMP call with single-Gaussian (AWGN) noise. Returns updated tau_w_1.

gamp_awgn_mix 🔒

One GAMP call with Gaussian-mixture noise (final stage).

input_function 🔒

Laplace input function: soft-threshold with the morphology mask passing large coefficients through unshrunk. Returns (x_hat_psi, tau_x).

input_parameter_est 🔒

EM update of the Laplace scale parameter lambda (single cluster).

linearize 🔒

Build the linearized measurement for the current estimate.

mix_output_function 🔒

Gaussian-mixture output function (single Gaussian + zero-mean outlier Gaussian). Returns (x_hat, tau_x) where x_hat is the posterior-mean noise estimate.

mix_output_parameter_est 🔒

EM update of the mixture parameters omega, phi, psi (theta, gamma fixed).

note_guard_firing 🔒

Record that the divergence guard rejected a sweep. Compiled away outside tests.

pad_volume 🔒

var 🔒

Sample variance (MATLAB var, normalized by n-1).

var_complex 🔒

Sample variance of a complex vector (MATLAB var on complex data: mean(|x - mean|^2) * n/(n-1)).

wavelet_block 🔒

Shared wavelet-domain GAMP block (identical between AWGN and mixture stages).