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

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.

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