qsm_core/utils/

epg.rs

//! EPG-based T2/R2 mapping from multi-echo spin-echo (MESE) magnitude data.
//!
//! A multi-echo spin-echo (CPMG) train uses repeated refocusing pulses. Real
//! refocusing pulses are not perfect 180° flips (B1 inhomogeneity, slice
//! profile), so they generate stimulated and indirect echoes that add signal at
//! later echoes. A naive mono-exponential fit (e.g. [`arlo`](crate::utils::r2star::arlo)) reads that
//! extra signal as slower decay and therefore *overestimates* T2 (underestimates
//! R2). The Extended Phase Graph (EPG) formalism models the full echo train,
//! including the stimulated-echo pathways, as a function of `(T2, T1, B1)` and so
//! removes that bias.
//!
//! This module provides:
//! - [`epg_cpmg_echoes`]: the forward model — simulated CPMG echo amplitudes.
//! - [`r2_epg`]: per-voxel R2 mapping by EPG dictionary matching.
//! - [`r2prime`]: R2' = R2* − R2 (the input chi-separation needs).
//!
//! When the refocusing flip angle is exactly 180° (`b1 = 1.0`), the EPG echo
//! train reduces exactly to the mono-exponential `S(TE_n) = exp(-TE_n / T2)`,
//! so on perfectly-refocused data EPG and ARLO agree (see the unit tests).
//!
//! **What the χ-separation references used for R2.** χ-sepnet and R2PRIMEnet (Kim et al. 2025,
//! arXiv:2409.14030, §2.1.1–2.1.2) matched multi-echo SE (or dual-echo TSE) magnitude against a
//! StimFit dictionary with stimulated-echo correction, the approach [`r2_epg`] takes, assuming
//! B1+ = 1 for the patient data. R2' = R2* − R2 was then clamped at 0, as [`r2prime`] does.
//! SUSEP-Net's in-vivo R2 was a plain mono-exponential STI Suite fit (arXiv:2506.13293, §3.2).
//!
//! References:
//! - Weigel, M. (2015). "Extended phase graphs: dephasing, RF pulses, and echoes
//!   — pure and simple." JMRI 41(2):266-295.
//! - Ben-Eliezer, N., et al. (2015). "Rapid and accurate T2 mapping from
//!   multi-spin-echo data using Bloch-simulation-based reconstruction." MRM
//!   73(2):809-817.

use num_complex::Complex64;
use crate::error::{check_len, InputError};
#[cfg(feature = "parallel")]
use rayon::prelude::*;

const PI: f64 = std::f64::consts::PI;

/// EPG configuration state: three complex vectors indexed by dephasing order
/// `k = 0..=k_max`. `fp[k]` is the F+ (positive dephasing) coherence at order
/// `k`, `fm[k]` the F- (negative dephasing) coherence, and `z[k]` the
/// longitudinal coherence.
struct EpgState {
    fp: Vec<Complex64>,
    fm: Vec<Complex64>,
    z: Vec<Complex64>,
}

impl EpgState {
    fn new(k_max: usize) -> Self {
        let zero = Complex64::new(0.0, 0.0);
        EpgState {
            fp: vec![zero; k_max + 1],
            fm: vec![zero; k_max + 1],
            z: vec![zero; k_max + 1],
        }
    }

    /// Relaxation + longitudinal regrowth (toward M0 = 1) over one interval.
    /// `e1 = exp(-dt/T1)`, `e2 = exp(-dt/T2)`.
    #[inline]
    fn relax(&mut self, e1: f64, e2: f64) {
        for v in self.fp.iter_mut() {
            *v *= e2;
        }
        for v in self.fm.iter_mut() {
            *v *= e2;
        }
        for v in self.z.iter_mut() {
            *v *= e1;
        }
        self.z[0] += 1.0 - e1;
    }

    /// Unit positive dephasing gradient: shift coherence orders. F+ moves to
    /// higher orders, F- to lower orders, and the observable F+(0) is refilled
    /// from conj(F-(0)) (Weigel 2015; Hargreaves EPG `epg_grad`).
    #[inline]
    fn grad(&mut self) {
        let k_max = self.fp.len() - 1;
        // F+ up: fp[k] = fp[k-1]
        for k in (1..=k_max).rev() {
            self.fp[k] = self.fp[k - 1];
        }
        // F- down: fm[k] = fm[k+1]
        for k in 0..k_max {
            self.fm[k] = self.fm[k + 1];
        }
        self.fm[k_max] = Complex64::new(0.0, 0.0);
        // Refill the observable from the conjugate-symmetric partner.
        self.fp[0] = self.fm[0].conj();
    }

    /// Instantaneous RF pulse of flip angle `alpha` (rad) and phase `phi` (rad),
    /// applied to every coherence order (Weigel 2015, Hargreaves `epg_rf`).
    #[inline]
    fn rf(&mut self, alpha: f64, phi: f64) {
        let c = (alpha / 2.0).cos();
        let s = (alpha / 2.0).sin();
        let c2 = c * c;
        let s2 = s * s;
        let sa = alpha.sin();
        let ca = alpha.cos();
        let i = Complex64::i();
        let eip = Complex64::from_polar(1.0, phi); // e^{i phi}
        let ei2p = Complex64::from_polar(1.0, 2.0 * phi); // e^{2i phi}

        let m00 = Complex64::new(c2, 0.0);
        let m01 = ei2p * s2;
        let m02 = -i * eip * sa;
        let m10 = ei2p.conj() * s2;
        let m11 = Complex64::new(c2, 0.0);
        let m12 = i * eip.conj() * sa;
        let m20 = -i * 0.5 * eip.conj() * sa;
        let m21 = i * 0.5 * eip * sa;
        let m22 = Complex64::new(ca, 0.0);

        for k in 0..self.fp.len() {
            let a = self.fp[k];
            let b = self.fm[k];
            let d = self.z[k];
            self.fp[k] = m00 * a + m01 * b + m02 * d;
            self.fm[k] = m10 * a + m11 * b + m12 * d;
            self.z[k] = m20 * a + m21 * b + m22 * d;
        }
    }
}

/// Simulate the CPMG multi-echo spin-echo magnitude train.
///
/// Models a 90° excitation followed by `n_echoes` refocusing pulses of nominal
/// 180° scaled by `b1` (so the actual refocusing flip is `b1 * 180°`), spaced by
/// echo spacing `esp` (seconds). Echoes form at `TE_n = n * esp`.
///
/// # Arguments
/// * `t2` - Transverse relaxation time (seconds)
/// * `t1` - Longitudinal relaxation time (seconds)
/// * `b1` - Refocusing efficiency: fraction of the nominal 180° flip (1.0 = perfect)
/// * `esp` - Echo spacing (seconds)
/// * `n_echoes` - Number of echoes to simulate
///
/// # Returns
/// `n_echoes` echo magnitudes, normalized so equilibrium magnetization M0 = 1.
///
/// With `b1 = 1.0` this returns exactly `exp(-n * esp / t2)` (perfect refocusing).
pub fn epg_cpmg_echoes(t2: f64, t1: f64, b1: f64, esp: f64, n_echoes: usize) -> Vec<f64> {
    // Coherence can reach up to one order per gradient (two per echo interval).
    let k_max = 2 * n_echoes + 2;
    let mut st = EpgState::new(k_max);

    // Equilibrium then 90° excitation about y (phi = pi/2). CPMG requires the
    // excitation and refocusing pulses to be 90° out of phase.
    st.z[0] = Complex64::new(1.0, 0.0);
    st.rf(PI / 2.0, PI / 2.0);

    let e1 = (-esp / 2.0 / t1).exp();
    let e2 = (-esp / 2.0 / t2).exp();
    let alpha = b1 * PI; // refocusing flip angle (rad)

    let mut echoes = Vec::with_capacity(n_echoes);
    for _ in 0..n_echoes {
        // First half-interval: relax + dephase.
        st.relax(e1, e2);
        st.grad();
        // Refocusing pulse about x (phi = 0).
        st.rf(alpha, 0.0);
        // Second half-interval: relax + rephase.
        st.relax(e1, e2);
        st.grad();
        // Echo forms at the observable coherence order 0.
        echoes.push(st.fp[0].norm());
    }
    echoes
}

/// Parameters for [`r2_epg`] dictionary matching.
pub struct R2EpgParams {
    /// Assumed T1 in seconds (weak influence on the T2 estimate).
    pub t1: f64,
    /// Candidate T2 values (seconds) — the dictionary's T2 axis.
    pub t2_grid: Vec<f64>,
    /// Candidate B1 (refocusing efficiency) values — the dictionary's B1 axis.
    /// When a B1 map is supplied to [`r2_epg`], map values are snapped to the
    /// nearest entry of this grid (densify it for finer quantization).
    pub b1_grid: Vec<f64>,
}

impl Default for R2EpgParams {
    fn default() -> Self {
        // T2 log-spaced 5 ms .. 500 ms (R2 ~ 2 .. 200 Hz).
        let n_t2 = 150;
        let (t2_lo, t2_hi) = (0.005_f64, 0.5_f64);
        let ln_lo = t2_lo.ln();
        let ln_hi = t2_hi.ln();
        let t2_grid: Vec<f64> = (0..n_t2)
            .map(|i| (ln_lo + (ln_hi - ln_lo) * i as f64 / (n_t2 - 1) as f64).exp())
            .collect();

        // B1 (refocusing efficiency) 0.6 .. 1.0 in steps of 0.02.
        let n_b1 = 21;
        let (b1_lo, b1_hi) = (0.6_f64, 1.0_f64);
        let b1_grid: Vec<f64> = (0..n_b1)
            .map(|i| b1_lo + (b1_hi - b1_lo) * i as f64 / (n_b1 - 1) as f64)
            .collect();

        R2EpgParams {
            t1: 1.0,
            t2_grid,
            b1_grid,
        }
    }
}

/// One dictionary atom: a normalized echo train and its `(T2, B1)` labels.
struct DictAtom {
    signal: Vec<f64>, // L2-normalized echo amplitudes
    t2: f64,
    b1: f64,
}

/// Build the EPG dictionary for a given echo spacing and echo count, grouped by
/// B1: `dict[i]` holds the atoms for `params.b1_grid[i]` (over all T2 values),
/// so a fixed-B1 search can restrict itself to a single group.
fn build_dictionary(params: &R2EpgParams, esp: f64, n_echoes: usize) -> Vec<Vec<DictAtom>> {
    params
        .b1_grid
        .iter()
        .map(|&b1| {
            let mut group = Vec::with_capacity(params.t2_grid.len());
            for &t2 in &params.t2_grid {
                let mut sig = epg_cpmg_echoes(t2, params.t1, b1, esp, n_echoes);
                let norm: f64 = sig.iter().map(|&v| v * v).sum::<f64>().sqrt();
                if norm > 1e-30 {
                    for v in sig.iter_mut() {
                        *v /= norm;
                    }
                    group.push(DictAtom { signal: sig, t2, b1 });
                }
            }
            group
        })
        .collect()
}

/// How far, in seconds, [`r2_epg`]'s echo times may be from `TEₙ = n·ESP` (each spacing from
/// the first, and the first echo from the spacing): 0.1 ms, which covers the rounding of echo
/// times in DICOM and BIDS sidecars.
pub const EPG_TE_TOLERANCE_S: f64 = 1e-4;

/// R2 mapping from multi-echo spin-echo magnitude via EPG dictionary matching.
///
/// For each masked voxel the (normalized) measured echo train is matched against
/// a dictionary of EPG-simulated trains over a `(T2, B1)` grid; the R2 = 1/T2 of
/// the best match is returned. Normalization removes the proton-density / receive
/// gain, so only the decay *shape* is fit. This corrects the stimulated-echo bias
/// that a mono-exponential fit ([`arlo`](crate::utils::r2star::arlo)) suffers on imperfectly-refocused
/// (B1 < 1) data.
///
/// When a per-voxel `b1_map` is supplied (e.g. from a separate B1 acquisition),
/// the B1 dimension of the search is fixed: each voxel's B1 is snapped to the
/// nearest `params.b1_grid` entry and only that column of the dictionary is
/// searched (a T2-only fit). This keeps EPG's stimulated-echo bias correction
/// while removing the free B1 parameter that overfits noise at low SNR.
///
/// # Arguments
/// * `magnitude` - MESE magnitude, flattened `[v0_e0, v0_e1, ..., v1_e0, ...]`
///   (row-major `(n_voxels, n_echoes)`)
/// * `mask` - Binary brain mask `[nx*ny*nz]` (1 = process, 0 = skip)
/// * `echo_times` - Spin-echo times in seconds `[n_echoes]` (equi-spaced, ≥3)
/// * `grid` - Volume grid (dimensions and voxel sizes)
/// * `params` - Dictionary grids and assumed T1
/// * `b1_map` - Optional known refocusing-efficiency map `[nx*ny*nz]`; `None`
///   fits B1 per voxel over the full `b1_grid`
///
/// # Returns
/// `(r2_map, b1_map)` - R2 in Hz and the B1 (refocusing efficiency) used per
/// voxel — fitted, or the snapped input when `b1_map` was given — both
/// `[nx*ny*nz]`.
///
/// # Errors
/// - fewer than 3 echoes;
/// - echo times that are not equally spaced, or whose first echo is not at the echo spacing
///   (`TE₁ = ESP`), each to within [`EPG_TE_TOLERANCE_S`]. The dictionary simulates a CPMG
///   train with echoes at `n·ESP`; a train that starts elsewhere has a different first
///   interval, and fitting it against this dictionary gives a wrong R2 rather than a noisy one;
/// - an empty dictionary (`params` grids with no usable `(T2, B1)` pair);
/// - `magnitude`/`mask`/`b1_map` lengths inconsistent with `grid`.
pub fn r2_epg(
    magnitude: &[f64],
    mask: &[u8],
    echo_times: &[f64],
    grid: &crate::Grid,
    params: &R2EpgParams,
    b1_map: Option<&[f64]>,
) -> Result<(Vec<f64>, Vec<f64>), InputError> {
    let n_echoes = echo_times.len();
    if n_echoes < 3 {
        return Err(InputError(format!("EPG R2 fitting needs at least 3 echoes, got {n_echoes}")));
    }
    if echo_times.iter().any(|t| !t.is_finite()) {
        return Err(InputError(format!("EPG R2 fitting needs finite echo times, got {echo_times:?}")));
    }
    let n_voxels = grid.n_total();
    check_len("magnitude (n_voxels * n_echoes)", magnitude.len(), n_voxels * n_echoes)?;
    check_len("mask", mask.len(), n_voxels)?;
    if let Some(b1) = b1_map {
        check_len("b1_map", b1.len(), n_voxels)?;
    }

    // Sort echoes by time and require (approximately) uniform spacing.
    let sort_indices: Vec<usize> = {
        let mut idx: Vec<usize> = (0..n_echoes).collect();
        idx.sort_by(|&a, &b| echo_times[a].total_cmp(&echo_times[b]));
        idx
    };
    let mut te_sorted: Vec<f64> = echo_times.to_vec();
    te_sorted.sort_by(f64::total_cmp);
    let esp = te_sorted[1] - te_sorted[0];
    let diffs: Vec<f64> = te_sorted.windows(2).map(|w| w[1] - w[0]).collect();
    let max_dev = diffs.iter().map(|&d| (d - esp).abs()).fold(0.0_f64, f64::max);
    if esp <= 0.0 || max_dev > EPG_TE_TOLERANCE_S {
        return Err(InputError(format!(
            "EPG R2 fitting needs distinct, equally spaced echo times, got {echo_times:?} s"
        )));
    }
    // EPG forms echoes at TE_n = n*esp; require the acquisition to start at esp.
    // (A constant offset would need a different first-interval model.)
    if (te_sorted[0] - esp).abs() > EPG_TE_TOLERANCE_S {
        return Err(InputError(format!(
            "EPG R2 fitting needs the first echo at the echo spacing (TE1 = ESP), got \
             TE1 = {:.3} ms, ESP = {:.3} ms",
            te_sorted[0] * 1e3, esp * 1e3
        )));
    }

    let dict = build_dictionary(params, esp, n_echoes);
    if dict.iter().all(|g| g.is_empty()) {
        return Err(InputError(
            "EPG dictionary is empty: no (T2, B1) pair in the parameter grids is usable".into(),
        ));
    }

    let mut r2_out = vec![0.0_f64; n_voxels];
    let mut b1_out = vec![0.0_f64; n_voxels];

    // Match each voxel against the dictionary (parallel over voxels).
    let mut out: Vec<(f64, f64)> = vec![(0.0, 0.0); n_voxels];
    maybe_par_chunks_mut!(out.as_mut_slice(), 1)
        .enumerate()
        .for_each(|(v, slot)| {
            if mask[v] == 0 {
                return;
            }
            // Extract and normalize the measured echo train.
            let mut sig: Vec<f64> = sort_indices
                .iter()
                .map(|&ei| magnitude[v * n_echoes + ei])
                .collect();
            let norm: f64 = sig.iter().map(|&x| x * x).sum::<f64>().sqrt();
            if norm < 1e-20 {
                return;
            }
            for x in sig.iter_mut() {
                *x /= norm;
            }

            // With a known B1, restrict the search to the nearest B1 column.
            let groups: &[Vec<DictAtom>] = match b1_map {
                Some(b1) => {
                    let gi = params
                        .b1_grid
                        .iter()
                        .enumerate()
                        .min_by(|(_, a), (_, b)| {
                            (*a - b1[v]).abs().partial_cmp(&(*b - b1[v]).abs()).unwrap()
                        })
                        .map(|(i, _)| i)
                        .unwrap();
                    std::slice::from_ref(&dict[gi])
                }
                None => &dict,
            };

            // Best match = maximum normalized dot product (cosine similarity).
            let mut best_dot = f64::NEG_INFINITY;
            let mut best = (0.0_f64, 0.0_f64);
            for atom in groups.iter().flatten() {
                let dot: f64 = sig
                    .iter()
                    .zip(atom.signal.iter())
                    .map(|(&a, &b)| a * b)
                    .sum();
                if dot > best_dot {
                    best_dot = dot;
                    best = (1.0 / atom.t2, atom.b1);
                }
            }
            slot[0] = best;
        });

    for v in 0..n_voxels {
        r2_out[v] = out[v].0;
        b1_out[v] = out[v].1;
    }

    Ok((r2_out, b1_out))
}

/// R2' and the region over which it is a measurement.
///
/// The two are returned together because they cannot be recovered from one another: a zero in
/// [`r2prime`](Self::r2prime) means either "R2* equalled R2 here" or "nothing was measured here",
/// and only [`coverage`](Self::coverage) tells them apart.
#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
#[derive(Clone, Debug)]
pub struct R2PrimeMap {
    /// R2' = R2* − R2 in Hz, clamped at zero, and zero wherever `coverage` is zero.
    pub r2prime: Vec<f64>,
    /// 1 where R2 was measured — and R2' is therefore a measurement — 0 elsewhere.
    pub coverage: Vec<u8>,
}

/// Compute R2' = R2* − R2, clamped at zero, over the region where R2 was measured.
///
/// R2* (from a gradient-echo acquisition) captures reversible + irreversible dephasing; R2 (from a
/// spin-echo acquisition) captures only the irreversible part. Their difference R2' is the
/// reversible dephasing that chi-separation uses to constrain the paramagnetic/diamagnetic split.
/// Noise can make R2* < R2 in some voxels, so negative values are clamped to zero, as
/// [`r2primenet`](super::r2primenet::r2primenet) does: every R2′ QSM.rs makes is non-negative,
/// and the χ-separation methods are given that map.
///
/// **R2' is only defined where R2 exists.** R2 comes from a multi-echo spin-echo acquisition, and
/// a MESE is routinely a *slab* — a band of slices rather than the whole head — because
/// whole-brain MESE is expensive. Subtracting over the whole brain mask would then compute
/// `R2' = R2* − 0 = R2*` outside the slab, overstating R2' by roughly the tissue R2 (~10-20 Hz in
/// brain at 3 T) with nothing in the output to show it: the map looks plausible, and every
/// R2'-consuming method inherits the error. So this returns the coverage alongside the map, and
/// leaves R2' at zero outside it; what to do with the remainder — restrict the downstream outputs
/// to it, or fill it with [`r2primenet`](super::r2primenet::r2primenet) — is a policy decision for
/// the caller, not this function.
///
/// Coverage is the intersection of three things:
/// - `mask` — where an R2' is wanted at all (the brain mask);
/// - `r2_coverage`, when given — where the R2 acquisition reached. Only the caller knows this; it
///   is a geometric fact about the acquisition, not something the R2 map reveals;
/// - `r2 > 0` — a finite, positive R2. R2 = 0 means T2 = ∞, which no tissue has, so a zero is an
///   absent measurement rather than a measured zero. This is exact for an R2 map from
///   [`r2_epg`], which returns exactly zero outside its own fit mask; for an R2 map resampled
///   from another grid, interpolation can leave small non-zero values just outside the acquired
///   FOV, and only an explicit `r2_coverage` excludes those.
///
/// Passing `r2_coverage = None` therefore falls back to the `r2 > 0` heuristic — conservative, and
/// right for an EPG-derived R2, but weaker than telling the function where the acquisition was.
///
/// # Arguments
/// * `r2star` - R2* map in Hz `[nx*ny*nz]`
/// * `r2` - R2 map in Hz `[nx*ny*nz]` (co-registered to `r2star`)
/// * `mask` - Binary brain mask `[nx*ny*nz]`
/// * `r2_coverage` - Optional binary mask of where the R2 acquisition reached `[nx*ny*nz]`
///
/// # Returns
/// [`R2PrimeMap`]: the R2' map in Hz and the coverage it is valid over.
///
/// # Errors
/// If `r2`, `mask` or `r2_coverage` is not the same length as `r2star`: R2 from a separate
/// spin-echo acquisition has to be resampled onto the R2\* grid first.
pub fn r2prime(
    r2star: &[f64],
    r2: &[f64],
    mask: &[u8],
    r2_coverage: Option<&[u8]>,
) -> Result<R2PrimeMap, InputError> {
    check_len("r2 (must be on the r2star grid)", r2.len(), r2star.len())?;
    check_len("mask", mask.len(), r2star.len())?;
    if let Some(c) = r2_coverage {
        check_len("r2_coverage", c.len(), r2star.len())?;
    }

    let n = r2star.len();
    let mut out = vec![0.0_f64; n];
    let mut coverage = vec![0_u8; n];
    for i in 0..n {
        let measured = mask[i] != 0
            && r2_coverage.is_none_or(|c| c[i] != 0)
            && r2[i].is_finite()
            && r2[i] > 0.0;
        if measured {
            coverage[i] = 1;
            out[i] = (r2star[i] - r2[i]).max(0.0);
        }
    }
    Ok(R2PrimeMap { r2prime: out, coverage })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Grid;

    /// With perfect 180° refocusing (b1 = 1), the EPG train must reduce exactly
    /// to the mono-exponential S(TE_n) = exp(-TE_n / T2).
    #[test]
    fn test_epg_perfect_refocus_is_monoexponential() {
        let t2 = 0.060; // 60 ms
        let t1 = 1.2;
        let esp = 0.010; // 10 ms
        let n = 8;
        let echoes = epg_cpmg_echoes(t2, t1, 1.0, esp, n);
        for k in 0..n {
            let te = (k + 1) as f64 * esp;
            let expected = (-te / t2).exp();
            let err = (echoes[k] - expected).abs() / expected;
            assert!(
                err < 1e-6,
                "echo {} EPG {} vs mono-exp {} (rel err {:.2e})",
                k,
                echoes[k],
                expected,
                err
            );
        }
    }

    /// With imperfect refocusing (b1 < 1), stimulated-echo pathways make later
    /// echoes DECAY MORE SLOWLY than a mono-exponential — i.e. the apparent T2 is
    /// biased upward. Verify the later echoes sit above the mono-exp curve.
    #[test]
    fn test_epg_imperfect_refocus_slower_decay() {
        let t2 = 0.060;
        let t1 = 1.2;
        let esp = 0.010;
        let n = 8;
        let echoes = epg_cpmg_echoes(t2, t1, 0.7, esp, n);
        // Normalize both to first echo for a fair shape comparison.
        let first = echoes[0];
        for k in 2..n {
            let te = (k + 1) as f64 * esp;
            let te1 = esp;
            let mono = (-(te - te1) / t2).exp(); // relative to echo 1
            let epg_rel = echoes[k] / first;
            assert!(
                epg_rel > mono,
                "echo {}: EPG rel {} should exceed mono-exp rel {}",
                k,
                epg_rel,
                mono
            );
        }
    }

    /// The headline result: on imperfectly-refocused data, EPG recovers the true
    /// T2 while a mono-exponential (log-linear) fit is biased. We simulate a
    /// b1 = 0.75 train, fit R2 with both, and check EPG is far closer to truth.
    #[test]
    fn test_epg_beats_monoexp_on_imperfect_refocus() {
        let t2_true = 0.070; // 70 ms  -> R2 = 14.29 Hz
        let r2_true = 1.0 / t2_true;
        let t1 = 1.2;
        let b1_true = 0.75;
        let esp = 0.010;
        let n = 10;
        let te: Vec<f64> = (1..=n).map(|k| k as f64 * esp).collect();

        // Simulate one voxel's MESE train (scaled by an arbitrary M0).
        let m0 = 850.0;
        let train = epg_cpmg_echoes(t2_true, t1, b1_true, esp, n);
        let mag: Vec<f64> = train.iter().map(|&v| v * m0).collect();

        let grid = Grid::new(1, 1, 1, 1.0, 1.0, 1.0);
        let mask = vec![1u8];

        // EPG fit (dictionary includes b1 = 0.75 and t2 near 70 ms).
        let params = R2EpgParams::default();
        let (r2_epg_map, b1_map) = r2_epg(&mag, &mask, &te, &grid, &params, None).unwrap();
        let r2_epg_val = r2_epg_map[0];
        let epg_err = (r2_epg_val - r2_true).abs() / r2_true;

        // Mono-exponential (ARLO) fit for comparison.
        // one voxel: voxel-major and echo-slowest are the same
        let r2_arlo_map = crate::utils::r2star::arlo(&mag, &te, 1).unwrap();
        let arlo_err = (r2_arlo_map[0] - r2_true).abs() / r2_true;

        // EPG should be within ~5% (grid-limited) and clearly better than ARLO.
        assert!(
            epg_err < 0.06,
            "EPG R2 {:.3} Hz vs true {:.3} Hz (err {:.1}%)",
            r2_epg_val,
            r2_true,
            epg_err * 100.0
        );
        assert!(
            epg_err < arlo_err,
            "EPG err {:.1}% should beat ARLO err {:.1}% (EPG {:.2} Hz, ARLO {:.2} Hz, true {:.2} Hz)",
            epg_err * 100.0,
            arlo_err * 100.0,
            r2_epg_val,
            r2_arlo_map[0],
            r2_true
        );
        // The fitted B1 should land near the truth.
        assert!(
            (b1_map[0] - b1_true).abs() <= 0.06,
            "fitted B1 {} vs true {}",
            b1_map[0],
            b1_true
        );
    }

    /// On perfectly-refocused data, EPG and ARLO should agree (both correct).
    #[test]
    fn test_epg_matches_arlo_on_perfect_refocus() {
        let t2_true = 0.050;
        let r2_true = 1.0 / t2_true;
        let esp = 0.012;
        let n = 6;
        let te: Vec<f64> = (1..=n).map(|k| k as f64 * esp).collect();
        let train = epg_cpmg_echoes(t2_true, 1.0, 1.0, esp, n);
        let mag: Vec<f64> = train.iter().map(|&v| v * 500.0).collect();

        let grid = Grid::new(1, 1, 1, 1.0, 1.0, 1.0);
        let mask = vec![1u8];
        let (r2_epg_map, _) = r2_epg(&mag, &mask, &te, &grid, &R2EpgParams::default(), None).unwrap();
        let err = (r2_epg_map[0] - r2_true).abs() / r2_true;
        assert!(err < 0.05, "EPG R2 {} vs true {}", r2_epg_map[0], r2_true);
    }

    /// With a known B1 map the fit is T2-only: R2 accuracy should match the full
    /// (T2, B1) search and the returned B1 must be the snapped input value.
    #[test]
    fn test_epg_fixed_b1_map_recovers_r2() {
        let t2_true = 0.070;
        let r2_true = 1.0 / t2_true;
        let t1 = 1.2;
        let b1_true = 0.7; // on the default b1_grid (0.6 + 5*0.02)
        let esp = 0.010;
        let n = 10;
        let te: Vec<f64> = (1..=n).map(|k| k as f64 * esp).collect();
        let train = epg_cpmg_echoes(t2_true, t1, b1_true, esp, n);
        let mag: Vec<f64> = train.iter().map(|&v| v * 850.0).collect();

        let grid = Grid::new(1, 1, 1, 1.0, 1.0, 1.0);
        let mask = vec![1u8];
        let b1_known = vec![b1_true];
        let (r2_map, b1_out) = r2_epg(
            &mag,
            &mask,
            &te,
            &grid,
            &R2EpgParams::default(),
            Some(&b1_known),
        ).unwrap();
        let err = (r2_map[0] - r2_true).abs() / r2_true;
        assert!(
            err < 0.06,
            "fixed-B1 EPG R2 {:.3} Hz vs true {:.3} Hz (err {:.1}%)",
            r2_map[0],
            r2_true,
            err * 100.0
        );
        assert!(
            (b1_out[0] - b1_true).abs() < 1e-12,
            "returned B1 {} should be the snapped input {}",
            b1_out[0],
            b1_true
        );
    }

    /// Off-grid B1 values must snap to the nearest b1_grid entry, and a wrong
    /// fixed B1 must actually constrain the fit (biasing R2 versus the truth) —
    /// proving the search really is restricted to the supplied column.
    #[test]
    fn test_epg_fixed_b1_snaps_and_constrains() {
        let t2_true = 0.070;
        let r2_true = 1.0 / t2_true;
        let t1 = 1.2;
        let b1_true = 0.7;
        let esp = 0.010;
        let n = 10;
        let te: Vec<f64> = (1..=n).map(|k| k as f64 * esp).collect();
        let train = epg_cpmg_echoes(t2_true, t1, b1_true, esp, n);
        let mag: Vec<f64> = train.iter().map(|&v| v * 850.0).collect();

        let grid = Grid::new(1, 1, 1, 1.0, 1.0, 1.0);
        let mask = vec![1u8];
        let params = R2EpgParams::default();

        // Off-grid input snaps to the nearest grid entry (0.695 -> 0.70).
        let (_, b1_out) = r2_epg(&mag, &mask, &te, &grid, &params, Some(&[0.695])).unwrap();
        assert!(
            (b1_out[0] - 0.70).abs() < 1e-12,
            "B1 0.695 should snap to 0.70, got {}",
            b1_out[0]
        );

        // Forcing perfect refocusing on a b1 = 0.7 train mis-reads the
        // stimulated-echo signal as slower decay: R2 biased low vs truth.
        let (r2_right, _) = r2_epg(&mag, &mask, &te, &grid, &params, Some(&[b1_true])).unwrap();
        let (r2_wrong, _) = r2_epg(&mag, &mask, &te, &grid, &params, Some(&[1.0])).unwrap();
        let err_right = (r2_right[0] - r2_true).abs() / r2_true;
        let err_wrong = (r2_wrong[0] - r2_true).abs() / r2_true;
        assert!(
            r2_wrong[0] < r2_true && err_wrong > err_right + 0.05,
            "wrong fixed B1 should bias R2 low: right {:.2} Hz (err {:.1}%), wrong {:.2} Hz (err {:.1}%), true {:.2} Hz",
            r2_right[0],
            err_right * 100.0,
            r2_wrong[0],
            err_wrong * 100.0,
            r2_true
        );
    }

    #[test]
    fn test_r2prime_subtract_and_clamp() {
        let r2star = vec![50.0, 30.0, 10.0, 0.0];
        let r2 = vec![20.0, 35.0, 10.0, 5.0];
        let mask = vec![1u8, 1, 0, 1];
        let out = r2prime(&r2star, &r2, &mask, None).unwrap();
        assert!((out.r2prime[0] - 30.0).abs() < 1e-12); // 50 - 20
        assert!(out.r2prime[1].abs() < 1e-12); // clamp negative (30 - 35)
        assert!(out.r2prime[2].abs() < 1e-12); // masked out
        assert!(out.r2prime[3].abs() < 1e-12); // clamp negative (0 - 5)
        // Clamping to zero is still a measurement; being outside the mask is not.
        assert_eq!(out.coverage, vec![1u8, 1, 0, 1]);
    }

    /// An explicit coverage mask must bound where R2' is reported, and the uncovered
    /// voxels must come back as zero *and* flagged — never as R2* − 0 = R2*.
    ///
    /// The uncovered voxels carry a *non-zero* R2 here on purpose: that is what a slab R2 map
    /// resampled onto the GRE grid leaves just outside its FOV, and it is the case the `r2 > 0`
    /// heuristic cannot catch. So only the supplied mask can exclude them, and dropping the
    /// mask from the implementation fails this test rather than slipping past it.
    #[test]
    fn test_r2prime_explicit_coverage_bounds_the_map() {
        // Four in-brain voxels; the R2 acquisition reached only the first two.
        let r2star = vec![40.0, 35.0, 30.0, 25.0];
        let r2 = vec![15.0, 12.0, 0.4, 0.1]; // interpolation bleed outside the slab
        let mask = vec![1u8; 4];
        let coverage = vec![1u8, 1, 0, 0];

        let out = r2prime(&r2star, &r2, &mask, Some(&coverage)).unwrap();
        assert_eq!(out.coverage, coverage, "coverage should be the supplied one");
        // Inside: R2* − R2, computed independently of the function under test.
        assert!((out.r2prime[0] - 25.0).abs() < 1e-12);
        assert!((out.r2prime[1] - 23.0).abs() < 1e-12);
        // Outside: zero. Subtracting the bleed-through would have given 29.6 and 24.9 —
        // essentially the whole of R2* reported as reversible.
        assert_eq!(out.r2prime[2], 0.0);
        assert_eq!(out.r2prime[3], 0.0);
    }

    /// Explicit coverage must *narrow*, not widen: a voxel the acquisition reached whose R2
    /// came back as zero is still an absent measurement, so it stays uncovered.
    #[test]
    fn test_r2prime_coverage_never_admits_a_zero_r2() {
        let r2star = vec![40.0, 40.0, 40.0];
        let r2 = vec![15.0, 0.0, f64::NAN];
        let mask = vec![1u8; 3];
        // The caller claims the acquisition covered everything.
        let out = r2prime(&r2star, &r2, &mask, Some(&[1u8, 1, 1])).unwrap();
        assert_eq!(out.coverage, vec![1u8, 0, 0], "a zero or non-finite R2 is not a measurement");
        assert!((out.r2prime[0] - 25.0).abs() < 1e-12);
        assert_eq!(out.r2prime[1], 0.0);
        assert_eq!(out.r2prime[2], 0.0);
    }

    /// With no coverage mask the `r2 > 0` heuristic has to stand in for it — which is exactly
    /// what an EPG-derived R2 supports, since [`r2_epg`] leaves un-fitted voxels at zero.
    #[test]
    fn test_r2prime_heuristic_coverage_without_a_mask() {
        let r2star = vec![40.0, 35.0, 30.0];
        let r2 = vec![15.0, 12.0, 0.0]; // third voxel never fitted
        let mask = vec![1u8; 3];
        let out = r2prime(&r2star, &r2, &mask, None).unwrap();
        assert_eq!(out.coverage, vec![1u8, 1, 0]);
        assert_eq!(out.r2prime[2], 0.0, "an unfitted voxel must not report R2* as R2'");
    }

    /// Echo times the dictionary cannot model are refused, not asserted on or fitted
    /// wrongly (#273).
    #[test]
    fn r2_epg_refuses_echo_times_it_cannot_model() {
        let grid = crate::Grid::new(1, 1, 1, 1.0, 1.0, 1.0);
        let params = R2EpgParams::default();
        let run = |te: &[f64]| {
            let mag: Vec<f64> = te.iter().map(|t| (-t / 0.08).exp()).collect();
            r2_epg(&mag, &[1], te, &grid, &params, None)
        };
        assert!(run(&[0.01, 0.02, 0.03, 0.04]).is_ok());
        for (te, why) in [
            (&[0.01, 0.02][..], "at least 3"),
            (&[0.01, 0.02, 0.04][..], "equally spaced"),
            (&[0.01, 0.01, 0.01][..], "equally spaced"),
            (&[0.015, 0.025, 0.035][..], "TE1 = ESP"),
            (&[0.01, f64::NAN, 0.03][..], "finite"),
        ] {
            let err = run(te).unwrap_err().to_string();
            assert!(err.contains(why), "{te:?}: {err}");
        }
        let te = [0.01, 0.02, 0.03];
        assert!(r2_epg(&[1.0; 2], &[1], &te, &grid, &params, None).is_err());
        assert!(r2_epg(&[1.0; 3], &[1], &te, &grid, &params, Some(&[1.0, 1.0])).is_err());
        let empty = R2EpgParams { t2_grid: vec![], ..R2EpgParams::default() };
        assert!(r2_epg(&[1.0; 3], &[1], &te, &grid, &empty, None).unwrap_err().to_string().contains("empty"));
    }

    #[test]
    fn r2prime_refuses_maps_on_different_grids() {
        assert!(r2prime(&[1.0; 4], &[1.0; 3], &[1; 4], None).is_err());
        assert!(r2prime(&[1.0; 4], &[1.0; 4], &[1; 3], None).is_err());
        assert!(r2prime(&[1.0; 4], &[1.0; 4], &[1; 4], Some(&[1; 5])).is_err());
    }
}