qsm_core/inversion/

tgv.rs

//! TGV-QSM: Total Generalized Variation for Quantitative Susceptibility Mapping
//!
//! Single-step QSM: background removal and dipole inversion (and, for a wrapped echo, phase
//! unwrapping) in one primal-dual iteration with TGV regularisation.
//!
//! References:
//! Langkammer, C., Bredies, K., Poser, B.A., et al. (2015).
//! "Fast quantitative susceptibility mapping using 3D EPI and total generalized variation."
//! NeuroImage, 111:622-630. https://doi.org/10.1016/j.neuroimage.2015.02.041
//!
//! Chatnuntawech, I., McDaniel, P., et al. (2017).
//! "Single-step quantitative susceptibility mapping with variational penalties."
//! NMR in Biomedicine, 30(4):e3570. https://doi.org/10.1002/nbm.3570
//!
//! Reference implementation: https://github.com/korbinian90/QuantitativeSusceptibilityMappingTGV.jl
//! (`qsm_tgv`). The iteration, step sizes, dipole stencil, mask erosions, Laplacian mean
//! correction and the defaults (α from `get_default_alpha(2)`, iterations from
//! `get_default_iterations`, `erosions = 3`, `step_size = 3`) are that package's.
//!
//! # Units
//!
//! The reference solves
//!
//! ```text
//!   min_{χ,w,ψ}  α₁‖∇χ − w‖₁ + α₀‖ε(w)‖₁ + ½‖ψ‖²   s.t.  ∇²φ = (⅓∇² − ∂²_z)χ + ∇²ψ
//! ```
//!
//! on the phase Laplacian `∇²φ` in rad/mm², with α in those units, and divides χ by
//! `c = 2π·γ·B0·TE` ([`rad_per_ppm`]) at the end. QSM.rs does the same: the input (a field in
//! ppm, or one wrapped echo in rad) becomes phase in rad at the given TE and B0 as the first
//! step, and χ is divided by `c` as the last. α are therefore in the reference's rad units and
//! mean what they mean there, at the scan's actual TE and B0. What the caller has decides how
//! `∇²φ` is formed (see [`TgvInput`]):
//!
//! - [`TgvInput::FieldPpm`]: an unwrapped field in ppm (for example a multi-echo field map),
//!   scaled to phase at the echo time [`TgvParams::te`] picks. `∇²φ` is the ordinary discrete
//!   Laplacian: the reference's `get_laplace_phase_romeo` with the unwrapping already done.
//! - [`TgvInput::WrappedPhase`]: one wrapped echo in rad, the original single-step TGV-QSM.
//!   `∇²φ` is the wrap-insensitive sin/cos (DEL) Laplacian, the reference's default
//!   `get_laplace_phase_del`.
//!
//! The output is χ in ppm.
//!
//! # Agreement with the reference
//!
//! On the QSM-CI in-vivo travelling-head scan (Siemens Prisma 2.89 T, 1 mm, 184×256×128, ROMEO
//! total field, brain mask touching the first z plane) against the package in the QSM-CI
//! `julia-qsm` image (`sha256:b3c048ac…`, 6c0886a, 4 threads), both at the package's defaults
//! (1655 iterations) and with both [`TgvParams::reference_divergence`] and
//! [`TgvParams::reference_edge_erosion`], over the voxels both return: NRMSE 4.0e-4 % for a wrapped echo (the field as phase
//! at TE1, the package's `get_laplace_phase_del`), and 3.7e-4 % for [`TgvInput::FieldPpm`]
//! against the package with its own discrete Laplacian kernel on the unwrapped phase. Both are
//! float32 rounding over 1655 iterations. `tests/tgv_tgvjl_parity.rs` checks both inputs,
//! anisotropic voxels and oblique B0 against the package's output on small phantoms (relative
//! L2 ≤ 3e-6), with all three reference options on (the comparison covers the package's whole
//! output box). The defaults differ; see below.
//!
//! Against the package's `get_laplace_phase_romeo`, which ROMEO-unwraps the given phase again
//! before its discrete Laplacian, the in-vivo field differs by NRMSE 12 %: on that scan the
//! re-unwrap adds ±2π at 2 voxels inside the data-term mask, which spreads through the
//! inversion. [`TgvInput::FieldPpm`] takes the Laplacian of the field as given.
//!
//! # Where QSM.rs differs from the reference
//!
//! Two details of the reference code are fixed by default, and one of its outputs is trimmed.
//! Each is an independent option that restores the reference's behaviour (QSM.rs's reference
//! rule: unequivocal fixes are the default, and the original's behaviour stays available for
//! parity):
//! - **Divergence in the `w` update** ([`TgvParams::reference_divergence`]). The reference scales
//!   each row of `q` by that row's own resolution for all three of its differences
//!   (`update_w_kernel!`; the same in the original Python `tgv_deconv_update_w`). Chambolle–Pock
//!   converges to the problem the paper states only when the operator in the primal update is
//!   the exact adjoint of the one in the dual update. On anisotropic voxels the reference's is not
//!   the adjoint of its symmetric gradient; a unit test checks the identity for the default and
//!   its failure for the reference's. By default the `w` update uses the exact adjoint. This is a
//!   pure correctness fix: it affects only anisotropic voxels, and on isotropic voxels the two
//!   are the same operator.
//! - **Erosion at the array's edge** ([`TgvParams::reference_edge_erosion`]). The reference skips
//!   neighbours outside the array (`erode_mask`, `src/tgv.jl:218-222`; the original Python's
//!   `erode_mask` treats them as inside), so a mask that reaches an array face is not eroded
//!   from it. There the Laplacian and gradient stencils have no valid neighbour (the reference
//!   replicates the edge value or drops the term), so the data term at those voxels is not the
//!   operator it is elsewhere. By default a neighbour outside the array counts as outside the
//!   mask, which removes those voxels. This affects only masks that reach an array face, and
//!   costs coverage at those faces.
//! - **Output support** ([`TgvParams::reference_output_box`]). χ is returned inside the eroded
//!   mask (`erosions` erosions of `mask`) and is zero elsewhere. The reference returns its whole
//!   working box (the eroded mask's bounding box plus one voxel, `reduce_to_mask_box`), whose χ
//!   outside that mask is the stencils' spill-over where the data term does not act: a rim of
//!   up to two voxels around the eroded mask (no data there; e.g. 44,530 voxels of the in-vivo
//!   brain mask). Inside the eroded mask the two are the same. By default QSM.rs trims that rim,
//!   since no data constrain it; the option returns it.
//!   - **What it changes in a comparison with the package.** Inside an uneroded brain mask the
//!     rim is part of the comparison. QSM-CI's `tgv` pair (QSMxT/QSM-CI `tgv-tgvjl` vs
//!     `tgv-qsmrs`, compared inside the dataset's mask, QSM.rs main 7f8271e, run 41019310531)
//!     measured NRMSE 3.84 % (sim, 41,907 rim voxels) and 6.74 % (KKI travelling head, 46,101):
//!     all of it the rim. Over the voxels both return, the same outputs differ by 3.8e-4 %
//!     (sim) and 4.9e-3 % (KKI): float32 rounding over 1655 iterations (#369).
//!
//! Measured against the package and, on the QSM-CI phantom, against the true χ. The total field
//! is the input; the defaults are QSM.rs's, and "both" turns both options on. NRMSE against the
//! package is over the voxels both return. NRMSE and corr against the truth are over the voxels
//! every setting returns.
//!
//! | case | mask reaches | default | divergence only | edge erosion only | both |
//! |---|---|---|---|---|---|
//! | in vivo, 1 mm (travelling head, brain mask) | z = 0 | 11.0 % | 11.0 % | 3.7e-4 % | 3.7e-4 % |
//! | in vivo, 1 × 1 × 2 mm (same, z pairs averaged) | z = 0 | 14.6 % | 14.4 % | 2.3 % | 2.0e-4 % |
//! | phantom, 1 mm (CI brain mask) | no face | 3.8e-4 % | 3.8e-4 % | 3.8e-4 % | 3.8e-4 % |
//! | phantom slab, 48 slices at 1 mm | both z faces | 99.8 % | 99.8 % | 4.4e-4 % | 4.4e-4 % |
//! | phantom slab, 48 slices at 1 × 1 × 2 mm | both z faces | 97.9 % | 97.9 % | 1.7 % | 1.1e-4 % |
//!
//! Against the true χ:
//!
//! | case | default | divergence only | edge erosion only | both (= the package) |
//! |---|---|---|---|---|
//! | phantom, 1 mm | 59.4 % / 0.807 | same | same | same |
//! | phantom slab, 1 mm | **68.5 % / 0.730** | 68.5 % / 0.730 | 1334 % / 0.034 | 1334 % / 0.034 |
//! | phantom slab, 1 × 1 × 2 mm | **65.6 % / 0.757** | 65.5 % / 0.757 | 349 % / 0.082 | 347 % / 0.082 |
//!
//! - **Coverage.** The edge erosion's cost is the voxels it removes at the faces: about 1,010 of
//!   the in-vivo scan's 1.5 million (one face), and 70,418 of 617,917 on the 1 mm phantom slab
//!   (two faces, 11 %). On a head mask (χ ≠ 0) of the phantom reaching three faces, it removes
//!   308,014 voxels.
//! - **Why it is still the default.** When the mask is cut by a face, the reference's unmasked
//!   stencils there spread through the whole slab: χ against the truth goes from NRMSE 69 % to
//!   1334 %.
//! - **Divergence.** It changes χ only on anisotropic voxels: by 1.1 % (relative L2) on a
//!   0.9 × 1.0 × 1.3 mm phantom. On the 1 × 1 × 2 mm in-vivo scan with the reference's erosion,
//!   the package's divergence accounts for the 2.3 % that separates "edge erosion only" from the
//!   package. Against the truth it does not help either way: on the 1 × 1 × 2 mm slab the
//!   package's divergence scores 0.1 points lower NRMSE (65.5 % vs 65.6 %, corr equal). The
//!   default stays the exact adjoint because that is the operator the TGV problem needs, not
//!   for a measured gain.
//!
//! The phantom cases use QSM-CI's in-silico phantom (QSM Challenge 2.0 head model, 164 × 205 × 205
//! at 1 mm, 7 T) with its true total field. They were run in the QSM.rs Actions on the CI's data
//! (QSM.rs#324). The slabs are 48 slices centred on the brain, cut from the field, the mask and
//! the truth; the 1 × 1 × 2 mm slab averages pairs of 96 slices.
//!
//! Optimizations:
//! - Bounding box reduction: only process the region containing the mask
//! - Pre-allocated buffers: all temporary arrays allocated once outside the loop
//!
//! # Reference pipeline
//! QuantitativeSusceptibilityMappingTGV.jl, commit 6c0886a:
//! - Input: one echo's phase in rad, its `TE` and the field strength, plus a brain mask
//!   (README "Example to run TGV"; `qsm_tgv(phase, mask, res; TE, fieldstrength, ...)`,
//!   `src/tgv.jl` L10). No unwrapping or background removal beforehand: the Laplacian of the
//!   phase does both (default `get_laplace_phase_del`; `get_laplace_phase_romeo` is the option
//!   for an unwrapped phase). χ is scaled by `1/(2π·TE·B0·42.5781)` at the end (L224-228).
//! - Mask: `erosions = 3` erosions of the given mask with a diamond (6-neighbour) element, then
//!   one more for the data-term mask `mask0` (L14-20, L218-222). The README's "Masking for QSM"
//!   points to phase-quality masking (ROMEO voxel quality, `robustmask(romeovoxelquality(...))`)
//!   combined with a brain mask.
//! - Resolution and orientation: no resampling; anisotropic voxels and oblique `B0_dir` are
//!   handled directly (README intro; the 27-point oblique stencil, L160). The default iteration
//!   count scales with the voxel volume (`get_default_iterations`, L108-115), and the README says
//!   the defaults "should lead to good results independently of the acquired resolution".
//!   `dedimensionalize = false` by default.
//! - Defaults: `regularization = 2` (α₀ = 0.002, α₁ = 0.003), `step_size = 3`,
//!   `correct_laplacian = true`. The settings that reproduce the original Python version are
//!   `step_size = 1`, `iterations = 1000`, `alpha = [0.003, 0.001]`, `erosions = 5`,
//!   `laplacian = get_laplace_phase3`, `correct_laplacian = false` (README).
//! - Langkammer et al. 2015 (abstract): 3D EPI at 3T and 7T, 1 mm isotropic at 3T.

#[cfg(feature = "parallel")]
use rayon::prelude::*;

/// One erosion with the 6-neighbour (`strel_diamond`) element. A neighbour outside the array
/// counts as outside the mask, so a mask that reaches the array's edge is eroded from it (the
/// default: the Laplacian and gradient stencils have no neighbour there). With `skip_outside`,
/// neighbours outside the array are skipped instead, as the reference's `erode_mask` does
/// (`src/tgv.jl:218-222`, 6c0886a, `checkbounds`), so such a mask is not eroded from the edge
/// ([`TgvParams::reference_edge_erosion`]).
fn erode_mask(mask: &[u8], nx: usize, ny: usize, nz: usize, skip_outside: bool) -> Vec<u8> {
    let edge = skip_outside;
    let nxy = nx * ny;
    let mut out = vec![0u8; mask.len()];
    for k in 0..nz {
        for j in 0..ny {
            for i in 0..nx {
                let l = i + j * nx + k * nxy;
                out[l] = u8::from(
                    mask[l] != 0
                        && (if i == 0 { edge } else { mask[l - 1] != 0 })
                        && (if i + 1 == nx { edge } else { mask[l + 1] != 0 })
                        && (if j == 0 { edge } else { mask[l - nx] != 0 })
                        && (if j + 1 == ny { edge } else { mask[l + nx] != 0 })
                        && (if k == 0 { edge } else { mask[l - nxy] != 0 })
                        && (if k + 1 == nz { edge } else { mask[l + nxy] != 0 }),
                );
            }
        }
    }
    out
}

/// The divergence of the symmetric tensor `q` in the `w` update, as the reference computes it
/// (`update_w_kernel!`, `src/tgv_helper.jl:88-105`, 6c0886a; the same in the original Python
/// `tgv_deconv_update_w`, `qsm_tgv_cython_helper.pyx`): row `d` of the tensor (`(qxx, qxy, qxz)`
/// for `w_x`, ...) takes its backward differences along x, y and z all scaled by `1/res_d`, the
/// row's own resolution, not by the resolution of the axis each difference is along. On
/// isotropic voxels that is the ordinary divergence; on anisotropic voxels it is not the adjoint
/// of the symmetric gradient used in the `q` update. Used only with
/// [`TgvParams::reference_divergence`]; the default is the adjoint
/// ([`crate::utils::gradient::symdiv_inplace_f32`]). See the module docs.
#[allow(clippy::too_many_arguments)] // the six tensor components, as in symdiv_inplace_f32
fn symdiv_rows_f32(
    divx: &mut [f32], divy: &mut [f32], divz: &mut [f32],
    sxx: &[f32], sxy: &[f32], sxz: &[f32],
    syy: &[f32], syz: &[f32], szz: &[f32],
    grid: &crate::Grid,
) {
    let (nx, ny, _) = grid.dims;
    let (r1, r2, r3) = (1.0 / grid.vsx() as f32, 1.0 / grid.vsy() as f32, 1.0 / grid.vsz() as f32);
    let nxy = nx * ny;
    // Backward differences with zero flux before the first plane (`div_local`).
    let d = |a: &[f32], l: usize, has_prev: bool, step: usize| a[l] - if has_prev { a[l - step] } else { 0.0 };
    maybe_par_chunks_mut!(divx, nxy)
        .zip(maybe_par_chunks_mut!(divy, nxy))
        .zip(maybe_par_chunks_mut!(divz, nxy))
        .enumerate()
        .for_each(|(k, ((dxs, dys), dzs))| {
            for j in 0..ny {
                for i in 0..nx {
                    let local = i + j * nx;
                    let l = local + k * nxy;
                    let (bi, bj, bk) = (i > 0, j > 0, k > 0);
                    dxs[local] = (d(sxx, l, bi, 1) + d(sxy, l, bj, nx) + d(sxz, l, bk, nxy)) * r1;
                    dys[local] = (d(sxy, l, bi, 1) + d(syy, l, bj, nx) + d(syz, l, bk, nxy)) * r2;
                    dzs[local] = (d(sxz, l, bi, 1) + d(syz, l, bj, nx) + d(szz, l, bk, nxy)) * r3;
                }
            }
        });
}

/// γ/2π in MHz/T, as used by TGV-QSM (Julia and Python reference implementations): the
/// reference implementation's value, not `crate::constants`. Keep it.
pub const TGV_GAMMA_MHZ_PER_T: f64 = 42.5781;

/// Phase (rad) per ppm of field at echo time `te_s` (s) and field strength `b0_t` (T):
/// `2π·γ·B0·TE` with γ/2π = [`TGV_GAMMA_MHZ_PER_T`].
pub fn rad_per_ppm(te_s: f64, b0_t: f64) -> f64 {
    2.0 * std::f64::consts::PI * TGV_GAMMA_MHZ_PER_T * b0_t * te_s
}

/// What TGV-QSM reconstructs from. Both carry the echo time and field strength that turn
/// the input into phase in rad, the units the reference (and α) work in.
#[derive(Clone, Copy, Debug)]
pub enum TgvInput<'a> {
    /// Unwrapped total (or local) field in ppm, e.g. a multi-echo field map. It is scaled to
    /// phase at the echo time [`TgvParams::te`] picks from `echo_times`, and its Laplacian is
    /// the ordinary discrete Laplacian.
    FieldPpm {
        /// Field in ppm.
        field: &'a [f64],
        /// The scan's echo times in seconds, in acquisition order.
        echo_times: &'a [f64],
        /// Field strength in Tesla.
        b0_t: f64,
    },
    /// One wrapped echo in rad (the original single-step TGV-QSM). Its Laplacian is the
    /// sin/cos (DEL) identity.
    WrappedPhase {
        /// Wrapped phase in rad.
        phase: &'a [f64],
        /// Echo time of `phase` in seconds.
        te_s: f64,
        /// Field strength in Tesla.
        b0_t: f64,
    },
}

impl TgvInput<'_> {
    fn data(&self) -> &[f64] {
        match self {
            TgvInput::FieldPpm { field, .. } => field,
            TgvInput::WrappedPhase { phase, .. } => phase,
        }
    }

    /// The echo time the input is phase at (for a field, the one `te` picks) and B0.
    fn te_b0(&self, te: crate::EchoTime) -> Result<(f64, f64), String> {
        match *self {
            TgvInput::FieldPpm { echo_times, b0_t, .. } => Ok((te.resolve(echo_times)?, b0_t)),
            TgvInput::WrappedPhase { te_s, b0_t, .. } => Ok((te_s, b0_t)),
        }
    }
}

/// TGV parameters. `alpha0`/`alpha1` are in the reference's rad units (see the module docs).
#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
#[derive(Clone, Debug)]
pub struct TgvParams {
    /// First-order TGV weight (gradient term), rad units
    pub alpha1: f32,
    /// Second-order TGV weight (symmetric gradient term), rad units
    pub alpha0: f32,
    /// Number of primal-dual iterations; `None` uses [`get_default_iterations`] for the
    /// voxel size and `step_size`, as the reference does.
    pub iterations: Option<usize>,
    /// Number of mask erosions
    pub erosions: usize,
    /// Primal step size multiplier (larger = faster but less stable; 1 is the publication)
    pub step_size: f32,
    /// Which echo time a field ([`TgvInput::FieldPpm`]) is converted to phase at; a wrapped
    /// echo is at its own. Default [`EchoTime::First`](crate::EchoTime::First):
    /// QuantitativeSusceptibilityMappingTGV.jl takes the phase of one echo and its `TE`
    /// (`qsm_tgv(phase, mask, res; TE, ...)`, README, af7650b) and has no rule for a field
    /// combined over echoes, so the first echo time, the default where the reference does not
    /// say.
    pub te: crate::EchoTime,
    /// Use the reference's divergence in the `w` update (default `false`: the exact adjoint of
    /// the symmetric gradient). The reference scales each row of `q` by that row's own
    /// resolution; on anisotropic voxels that is not the adjoint, so Chambolle–Pock does not
    /// solve the stated TGV problem. No effect on isotropic voxels. See the module docs,
    /// *Where QSM.rs differs from the reference*.
    pub reference_divergence: bool,
    /// Erode the mask as the reference does (default `false`: erode it from the array's edge
    /// too). The reference skips neighbours outside the array, so a mask that reaches an array
    /// face is not eroded from it, and keeps voxels whose Laplacian and gradient stencils have
    /// no neighbour there. No effect unless the mask reaches an array face.
    ///
    /// With this, [`reference_divergence`](Self::reference_divergence) and
    /// [`reference_output_box`](Self::reference_output_box), QSM.rs matches the package to
    /// float32 rounding (`tests/tgv_tgvjl_parity.rs`).
    pub reference_edge_erosion: bool,
    /// Return χ over the reference's whole working box (default `false`: inside the eroded
    /// mask only, zero elsewhere). The reference returns every voxel of the eroded mask's
    /// bounding box plus one voxel (`reduce_to_mask_box`, `src/tgv.jl` L21, L86-88), so a rim
    /// of up to two voxels just outside the eroded mask carries the stencils' spill-over, where
    /// the data term does not act. A comparison with the package inside an uneroded mask
    /// counts that rim. See the module docs, *Where QSM.rs differs from the reference*.
    pub reference_output_box: bool,
}

impl Default for TgvParams {
    fn default() -> Self {
        // QuantitativeSusceptibilityMappingTGV.jl defaults: regularization = 2,
        // erosions = 3, step_size = 3, iterations from the voxel size.
        let (alpha0, alpha1) = get_default_alpha(2.0);
        Self {
            alpha1,
            alpha0,
            iterations: None,
            erosions: 3,
            step_size: 3.0,
            te: crate::EchoTime::First,
            reference_divergence: false,
            reference_edge_erosion: false,
            reference_output_box: false,
        }
    }
}

/// Default `(alpha0, alpha1)` for a regularization strength (the reference's
/// `get_default_alpha`): `max([0.001, 0.001] + [0.001, 0.002]·(regularization − 1), 0)`.
/// 2 is the default, giving α₀ = 0.002, α₁ = 0.003.
pub fn get_default_alpha(regularization: f32) -> (f32, f32) {
    let r = regularization - 1.0;
    ((0.001 + 0.001 * r).max(0.0), (0.001 + 0.002 * r).max(0.0))
}

/// Get default number of iterations based on voxel size and step size.
///
/// Matches the Julia reference: `max(1000, 3200 / prod(res)^0.42) / step_size^0.6`
pub fn get_default_iterations(res: (f32, f32, f32), step_size: f32) -> usize {
    let prod_res = res.0 * res.1 * res.2;
    let it = (1000.0_f32).max(3200.0 / prod_res.powf(0.42)) / step_size.powf(0.6);
    it.round() as usize
}

/// Bounding box for mask region
#[derive(Clone, Debug)]
struct BoundingBox {
    i_min: usize,
    i_max: usize,
    j_min: usize,
    j_max: usize,
    k_min: usize,
    k_max: usize,
}

impl BoundingBox {
    /// Find minimal bounding box containing the mask with padding
    fn from_mask(mask: &[u8], nx: usize, ny: usize, nz: usize, padding: usize) -> Self {
        let mut i_min = nx;
        let mut i_max = 0;
        let mut j_min = ny;
        let mut j_max = 0;
        let mut k_min = nz;
        let mut k_max = 0;

        for k in 0..nz {
            for j in 0..ny {
                for i in 0..nx {
                    if mask[i + j * nx + k * nx * ny] != 0 {
                        i_min = i_min.min(i);
                        i_max = i_max.max(i);
                        j_min = j_min.min(j);
                        j_max = j_max.max(j);
                        k_min = k_min.min(k);
                        k_max = k_max.max(k);
                    }
                }
            }
        }

        // Add padding and clamp to bounds
        let i_min = i_min.saturating_sub(padding);
        let j_min = j_min.saturating_sub(padding);
        let k_min = k_min.saturating_sub(padding);
        let i_max = (i_max + padding + 1).min(nx);
        let j_max = (j_max + padding + 1).min(ny);
        let k_max = (k_max + padding + 1).min(nz);

        Self { i_min, i_max, j_min, j_max, k_min, k_max }
    }

    fn dims(&self) -> (usize, usize, usize) {
        (self.i_max - self.i_min, self.j_max - self.j_min, self.k_max - self.k_min)
    }

    fn total(&self) -> usize {
        let (bx, by, bz) = self.dims();
        bx * by * bz
    }
}

/// Extract sub-volume from full volume
fn extract_subvolume<T: Copy + Default>(
    full: &[T],
    bbox: &BoundingBox,
    nx: usize, ny: usize, _nz: usize,
) -> Vec<T> {
    let (bx, by, bz) = bbox.dims();
    let mut sub = vec![T::default(); bx * by * bz];

    for k in 0..bz {
        for j in 0..by {
            for i in 0..bx {
                let full_idx = (bbox.i_min + i) + (bbox.j_min + j) * nx + (bbox.k_min + k) * nx * ny;
                let sub_idx = i + j * bx + k * bx * by;
                sub[sub_idx] = full[full_idx];
            }
        }
    }
    sub
}

/// Insert sub-volume back into full volume
fn insert_subvolume<T: Copy>(
    full: &mut [T],
    sub: &[T],
    bbox: &BoundingBox,
    nx: usize, ny: usize, _nz: usize,
) {
    let (bx, by, bz) = bbox.dims();

    for k in 0..bz {
        for j in 0..by {
            for i in 0..bx {
                let full_idx = (bbox.i_min + i) + (bbox.j_min + j) * nx + (bbox.k_min + k) * nx * ny;
                let sub_idx = i + j * bx + k * bx * by;
                full[full_idx] = sub[sub_idx];
            }
        }
    }
}

/// Compute dipole stencil (27-point spatial kernel)
pub fn compute_dipole_stencil(
    res: (f32, f32, f32),
    b0_dir: (f32, f32, f32),
) -> [[[f32; 3]; 3]; 3] {
    let (dx, dy, dz) = res;
    let (bx, by, bz) = b0_dir;

    // Normalize B0 direction
    let b_norm = (bx * bx + by * by + bz * bz).sqrt();
    let bx = bx / b_norm;
    let by = by / b_norm;
    let bz = bz / b_norm;

    let mut stencil = [[[0.0f32; 3]; 3]; 3];

    let hx2 = 1.0 / (dx * dx);
    let hy2 = 1.0 / (dy * dy);
    let hz2 = 1.0 / (dz * dz);
    let factor = 1.0 / 3.0;

    // Center (i=1, j=1, k=1)
    stencil[1][1][1] = -2.0 * (hx2 + hy2 + hz2) * factor
                     + 2.0 * (bx * bx * hx2 + by * by * hy2 + bz * bz * hz2);

    // X neighbors
    stencil[0][1][1] = hx2 * factor - bx * bx * hx2;
    stencil[2][1][1] = hx2 * factor - bx * bx * hx2;

    // Y neighbors
    stencil[1][0][1] = hy2 * factor - by * by * hy2;
    stencil[1][2][1] = hy2 * factor - by * by * hy2;

    // Z neighbors
    stencil[1][1][0] = hz2 * factor - bz * bz * hz2;
    stencil[1][1][2] = hz2 * factor - bz * bz * hz2;

    // Cross terms for oblique B0
    let hxy = 1.0 / (dx * dy);
    let hxz = 1.0 / (dx * dz);
    let hyz = 1.0 / (dy * dz);

    let xy_factor = -bx * by * hxy;
    stencil[0][0][1] = xy_factor;
    stencil[2][2][1] = xy_factor;
    stencil[0][2][1] = -xy_factor;
    stencil[2][0][1] = -xy_factor;

    let xz_factor = -bx * bz * hxz;
    stencil[0][1][0] = xz_factor;
    stencil[2][1][2] = xz_factor;
    stencil[0][1][2] = -xz_factor;
    stencil[2][1][0] = -xz_factor;

    let yz_factor = -by * bz * hyz;
    stencil[1][0][0] = yz_factor;
    stencil[1][2][2] = yz_factor;
    stencil[1][0][2] = -yz_factor;
    stencil[1][2][0] = -yz_factor;

    stencil
}

/// Compute SVD-fitted oblique 27-point dipole stencil via DST-based Poisson solves.
///
/// Matches the Julia reference implementation from QuantitativeSusceptibilityMappingTGV.jl.
/// For each unique centrosymmetric pair of the 26 off-center stencil positions, a Poisson
/// equation is solved using DST-I to obtain Green's function values. These are then fitted
/// to the analytical dipole field via least squares, and resolution-weighted to produce
/// the final stencil.
///
/// Note: The delta function for offsets (di,dj,dk) and (-di,-dj,-dk) are identical
/// (both place +1 at mid±I and -2 at mid), so we only compute 13 unique Poisson
/// solutions and solve a full-rank 13×13 system. This avoids the rank deficiency
/// that would occur with all 26 positions.
pub fn compute_oblique_stencil(
    res: (f32, f32, f32),
    b0_dir: (f32, f32, f32),
) -> [[[f32; 3]; 3]; 3] {
    let n: usize = 64;
    let singularity_cutout = 4.0_f64;
    let mid = n / 2; // 32 (0-indexed center)
    let n2 = n * n;
    let n3 = n * n * n;

    let (dx, dy, dz) = (res.0 as f64, res.1 as f64, res.2 as f64);

    // Normalize B0 direction
    let b_norm = ((b0_dir.0 * b0_dir.0 + b0_dir.1 * b0_dir.1 + b0_dir.2 * b0_dir.2) as f64).sqrt();
    let bdir = (b0_dir.0 as f64 / b_norm, b0_dir.1 as f64 / b_norm, b0_dir.2 as f64 / b_norm);

    // Compute dipole field on 64³ grid
    let mut d = vec![f64::NAN; n3];
    let mut d_mask = vec![false; n3];

    for k in 0..n {
        for j in 0..n {
            for i in 0..n {
                let x = i as f64 - mid as f64;
                let y = j as f64 - mid as f64;
                let z = k as f64 - mid as f64;
                let r = (x * x + y * y + z * z).sqrt();

                let idx = i + j * n + k * n2;
                if r < singularity_cutout {
                    // d[idx] remains NAN, d_mask[idx] remains false
                } else {
                    let xz = (bdir.0 * x + bdir.1 * y + bdir.2 * z) / r;
                    let kappa = (3.0 * xz * xz - 1.0) / (4.0 * std::f64::consts::PI * r * r * r);
                    d[idx] = kappa;
                    d_mask[idx] = true;
                }
            }
        }
    }

    // DST-I eigenvalue components: coord2[k] = 2 * sin(π*(k+1) / (2*(N+1)))
    let coord2_sq: Vec<f64> = (0..n)
        .map(|k| {
            let v = 2.0 * (std::f64::consts::PI * (k as f64 + 1.0) / (2.0 * (n as f64 + 1.0))).sin();
            v * v
        })
        .collect();

    // 3D eigenvalue grid: coord2_grid[i,j,k] = coord2[i]² + coord2[j]² + coord2[k]²
    let mut coord2_grid = vec![0.0_f64; n3];
    for k in 0..n {
        for j in 0..n {
            for i in 0..n {
                coord2_grid[i + j * n + k * n2] = coord2_sq[i] + coord2_sq[j] + coord2_sq[k];
            }
        }
    }

    let dst = rustdct::DctPlanner::<f64>::new().plan_dst1(n);
    // Inverse DST-I scale: our DST computes Σ x[n]*sin(...) (no factor of 2),
    // so the inverse is (2/(N+1))^3 for 3D (unlike FFTW which uses 1/(2*(N+1))^3).
    let idst_scale = (2.0 / (n as f64 + 1.0)).powi(3);

    // Enumerate all 26 stencil positions (excluding center), column-major order
    let mut stencil_positions: Vec<(i32, i32, i32)> = Vec::with_capacity(26);
    for dk in -1..=1_i32 {
        for dj in -1..=1_i32 {
            for di in -1..=1_i32 {
                if di == 0 && dj == 0 && dk == 0 {
                    continue;
                }
                stencil_positions.push((di, dj, dk));
            }
        }
    }

    // Collect valid dipole point indices
    let valid_indices: Vec<usize> = d_mask
        .iter()
        .enumerate()
        .filter(|(_, &m)| m)
        .map(|(idx, _)| idx)
        .collect();

    // Exploit centrosymmetry: delta for (di,dj,dk) is identical to (-di,-dj,-dk),
    // so only compute 13 unique Poisson solutions (one per centrosymmetric pair).
    // Pair p corresponds to stencil_positions[p] and stencil_positions[25-p].
    let num_pairs = 13;
    let sin_table = dst_sin_table(n);
    let pairs: Vec<usize> = (0..num_pairs).collect();
    let a_rows: Vec<Vec<f64>> = maybe_par_iter!(pairs).map(|&p| {
        let (di, dj, dk) = stencil_positions[p];

        // Delta function: 1 at mid+I and at mid-I, -2 at mid. Its forward 3D DST-I is the
        // weighted sum of the separable transforms of its three points, written down directly.
        let at = |o: i32| (mid as i32 + o) as usize;
        let points = [
            ((at(di), at(dj), at(dk)), 1.0),
            ((at(-di), at(-dj), at(-dk)), 1.0),
            ((mid, mid, mid), -2.0),
        ];
        let mut fdelta = vec![0.0_f64; n3];
        for &((a, b, c), w) in &points {
            for k in 0..n {
                let sk = w * sin_table[c][k];
                for j in 0..n {
                    let sjk = sk * sin_table[b][j];
                    let base = j * n + k * n2;
                    for (i, v) in fdelta[base..base + n].iter_mut().enumerate() {
                        *v += sjk * sin_table[a][i];
                    }
                }
            }
        }

        // Poisson solve: divide by -coord2_grid
        for idx in 0..n3 {
            fdelta[idx] /= -coord2_grid[idx];
        }

        // Inverse 3D DST-I (= forward DST-I * scale)
        let vdelta = dst3d(&fdelta, n, dst.as_ref());

        // Extract values at valid dipole points, applying inverse scale
        valid_indices.iter().map(|&idx| vdelta[idx] * idst_scale).collect()
    }).collect();

    // Extract dipole values at valid points
    let d_valid: Vec<f64> = valid_indices.iter().map(|&idx| d[idx]).collect();

    // Solve least squares: B^T * y ≈ d_valid via normal equations (13 unknowns, full rank)
    // G = B * B^T (13×13), h = B * d_valid (13×1)
    let mut g = vec![vec![0.0_f64; num_pairs]; num_pairs];
    for i in 0..num_pairs {
        for j in 0..=i {
            let dot: f64 = a_rows[i]
                .iter()
                .zip(a_rows[j].iter())
                .map(|(&a, &b)| a * b)
                .sum();
            g[i][j] = dot;
            g[j][i] = dot;
        }
    }

    let mut h = vec![0.0_f64; num_pairs];
    for i in 0..num_pairs {
        h[i] = a_rows[i]
            .iter()
            .zip(d_valid.iter())
            .map(|(&a, &b)| a * b)
            .sum();
    }

    // Solve G * y = h via eigendecomposition with thresholding (handles rank deficiency
    // from centrosymmetry and other spatial symmetries like x-y equivalence)
    let y = solve_symmetric_pseudoinverse(&g, &h);

    // Assemble stencil: assign each pair's coefficient to both positions
    // y[p] = x[p] + x[25-p] = 2*x[p] (by symmetry), so stencil[I] = 2*x[I] = y[p]
    let mut result = [[[0.0f32; 3]; 3]; 3];
    for p in 0..num_pairs {
        let coeff = y[p] as f32;

        // First member of pair
        let (di, dj, dk) = stencil_positions[p];
        result[(di + 1) as usize][(dj + 1) as usize][(dk + 1) as usize] = coeff;

        // Second member (centrosymmetric partner)
        let (di2, dj2, dk2) = stencil_positions[25 - p];
        result[(di2 + 1) as usize][(dj2 + 1) as usize][(dk2 + 1) as usize] = coeff;
    }

    // Apply resolution weights: w = (i²/dx² + j²/dy² + k²/dz²) / (i² + j² + k²)
    for dk in -1..=1_i32 {
        for dj in -1..=1_i32 {
            for di in -1..=1_i32 {
                if di == 0 && dj == 0 && dk == 0 {
                    continue;
                }
                let si = (di + 1) as usize;
                let sj = (dj + 1) as usize;
                let sk = (dk + 1) as usize;

                let i2 = (di * di) as f64;
                let j2 = (dj * dj) as f64;
                let k2 = (dk * dk) as f64;
                let weight = (i2 / (dx * dx) + j2 / (dy * dy) + k2 / (dz * dz)) / (i2 + j2 + k2);
                result[si][sj][sk] *= weight as f32;
            }
        }
    }

    // Center = -sum(all others)
    let mut total = 0.0f32;
    for dk in 0..3 {
        for dj in 0..3 {
            for di in 0..3 {
                if !(di == 1 && dj == 1 && dk == 1) {
                    total += result[di][dj][dk];
                }
            }
        }
    }
    result[1][1][1] = -total;

    result
}

/// 3D separable DST-I, unnormalised: X[k] = Σ x[j]·sin(π·(j+1)·(k+1)/(N+1)) along each axis.
fn dst3d(input: &[f64], n: usize, dst: &dyn rustdct::Dst1<f64>) -> Vec<f64> {
    let n2 = n * n;
    let mut data = input.to_vec();
    let mut scratch = vec![0.0_f64; dst.get_scratch_len()];
    let mut line = vec![0.0_f64; n];
    // rustdct's FFT-based DST-I takes the zero padding of its FFT input from the scratch
    // buffer, so the scratch is cleared before every line.
    let mut dst1 = |buf: &mut [f64]| {
        scratch.fill(0.0);
        dst.process_dst1_with_scratch(buf, &mut scratch);
    };
    for row in data.chunks_mut(n) {
        dst1(row);
    }
    // y lines start at i + k·n² with stride n; z lines start at i + j·n with stride n².
    for (stride, outer_step) in [(n, n2), (n2, n)] {
        for a in 0..n {
            for i in 0..n {
                let base = i + a * outer_step;
                for t in 0..n {
                    line[t] = data[base + t * stride];
                }
                dst1(&mut line);
                for t in 0..n {
                    data[base + t * stride] = line[t];
                }
            }
        }
    }
    data
}

/// Build sin table for DST-I of given length.
/// sin_table[n][k] = sin(π*(n+1)*(k+1)/(N+1))
fn dst_sin_table(n: usize) -> Vec<Vec<f64>> {
    let scale = std::f64::consts::PI / (n as f64 + 1.0);
    let mut table = vec![vec![0.0_f64; n]; n];
    for j in 0..n {
        for k in 0..n {
            table[j][k] = ((j as f64 + 1.0) * (k as f64 + 1.0) * scale).sin();
        }
    }
    table
}

/// 3D separable DST-I by the direct O(N⁴) sum, the reference [`dst3d`] is checked against.
#[cfg(test)]
fn dst3d_direct(input: &[f64], n: usize, sin_table: &[Vec<f64>]) -> Vec<f64> {
    let n2 = n * n;
    let mut data = input.to_vec();
    let mut buf_in = vec![0.0_f64; n];
    let mut buf_out = vec![0.0_f64; n];

    // Transform along x (contiguous dimension)
    for k in 0..n {
        for j in 0..n {
            let base = j * n + k * n2;
            buf_in.copy_from_slice(&data[base..base + n]);
            dst1(&buf_in, sin_table, &mut buf_out);
            data[base..base + n].copy_from_slice(&buf_out);
        }
    }

    // Transform along y (stride = n)
    for k in 0..n {
        for i in 0..n {
            for j in 0..n {
                buf_in[j] = data[i + j * n + k * n2];
            }
            dst1(&buf_in, sin_table, &mut buf_out);
            for j in 0..n {
                data[i + j * n + k * n2] = buf_out[j];
            }
        }
    }

    // Transform along z (stride = n*n)
    for j in 0..n {
        for i in 0..n {
            for k in 0..n {
                buf_in[k] = data[i + j * n + k * n2];
            }
            dst1(&buf_in, sin_table, &mut buf_out);
            for k in 0..n {
                data[i + j * n + k * n2] = buf_out[k];
            }
        }
    }

    data
}

#[cfg(test)]
/// 1D DST-I: X[k] = Σ x[n] * sin(π*(n+1)*(k+1)/(N+1))
fn dst1(input: &[f64], sin_table: &[Vec<f64>], output: &mut [f64]) {
    let n = input.len();
    for k in 0..n {
        let mut sum = 0.0_f64;
        for j in 0..n {
            sum += input[j] * sin_table[j][k];
        }
        output[k] = sum;
    }
}

/// Solve a symmetric positive semi-definite system via eigendecomposition with thresholding.
///
/// Computes the pseudo-inverse solution: x = V * diag(1/λ, thresholded) * V^T * h
/// where G = V Λ V^T is the eigendecomposition. Small eigenvalues (< threshold * max_eigenvalue)
/// are zeroed, matching Julia's SVD-based approach for handling rank-deficient systems.
fn solve_symmetric_pseudoinverse(g: &[Vec<f64>], h: &[f64]) -> Vec<f64> {
    let n = h.len();

    // Copy g for eigendecomposition (Jacobi method modifies in place)
    let mut a: Vec<Vec<f64>> = g.to_vec();
    let mut v = vec![vec![0.0_f64; n]; n];
    for i in 0..n {
        v[i][i] = 1.0;
    }

    // Jacobi eigendecomposition for symmetric matrices
    let max_sweeps = 100;
    let tol = 1e-15;

    for _ in 0..max_sweeps {
        // Find largest off-diagonal element
        let mut max_off = 0.0_f64;
        for i in 0..n {
            for j in (i + 1)..n {
                max_off = max_off.max(a[i][j].abs());
            }
        }
        if max_off < tol {
            break;
        }

        // Sweep all off-diagonal pairs
        for p in 0..n {
            for q in (p + 1)..n {
                if a[p][q].abs() < tol {
                    continue;
                }

                // Compute Givens rotation angle to zero a[p][q]
                let app = a[p][p];
                let aqq = a[q][q];
                let apq = a[p][q];
                let tau = (aqq - app) / (2.0 * apq);
                let t = if tau >= 0.0 {
                    1.0 / (tau + (1.0 + tau * tau).sqrt())
                } else {
                    -1.0 / (-tau + (1.0 + tau * tau).sqrt())
                };
                let c = 1.0 / (1.0 + t * t).sqrt();
                let s = t * c;

                // Update matrix A' = G^T A G (only rows/cols p, q change)
                // First update off-diagonal rows
                for i in 0..n {
                    if i == p || i == q {
                        continue;
                    }
                    let aip = a[i][p];
                    let aiq = a[i][q];
                    a[i][p] = c * aip - s * aiq;
                    a[p][i] = a[i][p];
                    a[i][q] = s * aip + c * aiq;
                    a[q][i] = a[i][q];
                }

                // Update 2×2 block
                a[p][p] = c * c * app - 2.0 * c * s * apq + s * s * aqq;
                a[q][q] = s * s * app + 2.0 * c * s * apq + c * c * aqq;
                a[p][q] = 0.0;
                a[q][p] = 0.0;

                // Accumulate eigenvectors: V' = V * G
                for i in 0..n {
                    let vip = v[i][p];
                    let viq = v[i][q];
                    v[i][p] = c * vip - s * viq;
                    v[i][q] = s * vip + c * viq;
                }
            }
        }
    }

    // Eigenvalues are on diagonal
    let eigenvalues: Vec<f64> = (0..n).map(|i| a[i][i]).collect();
    let max_eigen = eigenvalues.iter().cloned().fold(0.0_f64, |a, b| a.max(b.abs()));
    let threshold = 1e-10 * max_eigen;

    // Solve: x = V * diag(1/λ_thresholded) * V^T * h
    let mut vt_h = vec![0.0_f64; n];
    for i in 0..n {
        for j in 0..n {
            vt_h[i] += v[j][i] * h[j];
        }
    }

    for i in 0..n {
        if eigenvalues[i].abs() > threshold {
            vt_h[i] /= eigenvalues[i];
        } else {
            vt_h[i] = 0.0;
        }
    }

    let mut x = vec![0.0_f64; n];
    for i in 0..n {
        for j in 0..n {
            x[i] += v[i][j] * vt_h[j];
        }
    }

    x
}

/// Apply dipole stencil to a 3D volume
/// Uses Neumann BC at boundaries (matching Julia's wave_local)
///
/// With `mask`, the output is zero outside it; without, it is computed at every voxel (the
/// reference's `wave_local`, zero only on the box boundary).
fn apply_stencil(
    output: &mut [f32],
    input: &[f32],
    stencil: &[[[f32; 3]; 3]; 3],
    mask: Option<&[u8]>,
    nx: usize, ny: usize, nz: usize,
) {
    let nxy = nx * ny;
    // One rayon task per z-slab (disjoint output chunk); reads input by global index.
    maybe_par_chunks_mut!(output, nxy).enumerate().for_each(|(k, out_slab)| {
        for j in 0..ny {
            for i in 0..nx {
                let local = i + j * nx;
                let idx = local + k * nxy;

                if mask.is_some_and(|m| m[idx] == 0) {
                    out_slab[local] = 0.0;
                    continue;
                }

                // Julia's wave_local only computes if not at boundary
                // If at boundary, result is 0
                if i == 0 || j == 0 || k == 0 || i + 1 >= nx || j + 1 >= ny || k + 1 >= nz {
                    out_slab[local] = 0.0;
                    continue;
                }

                let mut sum = 0.0f32;

                for dk in 0..3i32 {
                    for dj in 0..3i32 {
                        for di in 0..3i32 {
                            let ni = (i as i32 + di - 1) as usize;
                            let nj = (j as i32 + dj - 1) as usize;
                            let nk = (k as i32 + dk - 1) as usize;

                            let nidx = ni + nj * nx + nk * nxy;
                            sum += stencil[di as usize][dj as usize][dk as usize] * input[nidx];
                        }
                    }
                }

                out_slab[local] = sum;
            }
        }
    });
}

/// Laplacian of an unwrapped field (any units), zero outside `mask`.
///
/// Ordinary 7-point discrete Laplacian with replicated boundaries. This is what
/// [`TgvInput::FieldPpm`] uses: unlike [`compute_phase_laplacian`] it is linear, so it is
/// exact for fields whose neighbour differences are not small compared with 1.
pub fn compute_field_laplacian(
    field: &[f32],
    mask: &[u8],
    nx: usize, ny: usize, nz: usize,
    vsx: f32, vsy: f32, vsz: f32,
) -> Vec<f32> {
    let mut lap = compute_laplacian(field, nx, ny, nz, vsx, vsy, vsz);
    for (v, &m) in lap.iter_mut().zip(mask.iter()) {
        if m == 0 {
            *v = 0.0;
        }
    }
    lap
}

/// Compute Laplacian of wrapped phase (rad) using the DEL method.
///
/// `cos φ·∇²sin φ − sin φ·∇²cos φ`, which on the discrete grid is
/// `Σ_neighbours sin(φ_n − φ)/h²`: insensitive to 2π wraps, and equal to the ordinary
/// Laplacian only while neighbour differences are small. Only meaningful for phase in rad.
pub fn compute_phase_laplacian(
    phase: &[f32],
    mask: &[u8],
    nx: usize, ny: usize, nz: usize,
    vsx: f32, vsy: f32, vsz: f32,
) -> Vec<f32> {
    let n_total = nx * ny * nz;

    let sin_phase: Vec<f32> = phase.iter().map(|&p| p.sin()).collect();
    let cos_phase: Vec<f32> = phase.iter().map(|&p| p.cos()).collect();

    let lap_sin = compute_laplacian(&sin_phase, nx, ny, nz, vsx, vsy, vsz);
    let lap_cos = compute_laplacian(&cos_phase, nx, ny, nz, vsx, vsy, vsz);

    let mut laplacian = vec![0.0f32; n_total];
    for i in 0..n_total {
        if mask[i] != 0 {
            laplacian[i] = lap_sin[i] * cos_phase[i] - lap_cos[i] * sin_phase[i];
        }
    }

    laplacian
}

/// Compute discrete Laplacian of a 3D array
fn compute_laplacian(
    input: &[f32],
    nx: usize, ny: usize, nz: usize,
    vsx: f32, vsy: f32, vsz: f32,
) -> Vec<f32> {
    let n_total = nx * ny * nz;
    let mut output = vec![0.0f32; n_total];

    let hx2 = 1.0 / (vsx * vsx);
    let hy2 = 1.0 / (vsy * vsy);
    let hz2 = 1.0 / (vsz * vsz);
    let center = -2.0 * (hx2 + hy2 + hz2);

    for k in 0..nz {
        let km1 = if k == 0 { 0 } else { k - 1 };
        let kp1 = if k + 1 >= nz { nz - 1 } else { k + 1 };

        for j in 0..ny {
            let jm1 = if j == 0 { 0 } else { j - 1 };
            let jp1 = if j + 1 >= ny { ny - 1 } else { j + 1 };

            for i in 0..nx {
                let im1 = if i == 0 { 0 } else { i - 1 };
                let ip1 = if i + 1 >= nx { nx - 1 } else { i + 1 };

                let idx = i + j * nx + k * nx * ny;

                output[idx] = center * input[idx]
                    + hx2 * (input[im1 + j * nx + k * nx * ny] + input[ip1 + j * nx + k * nx * ny])
                    + hy2 * (input[i + jm1 * nx + k * nx * ny] + input[i + jp1 * nx + k * nx * ny])
                    + hz2 * (input[i + j * nx + km1 * nx * ny] + input[i + j * nx + kp1 * nx * ny]);
            }
        }
    }

    output
}

/// Apply the Laplacian (Neumann boundary). With `mask`, the output is zero outside it;
/// without, it is computed at every voxel.
fn apply_laplacian_inplace(
    output: &mut [f32],
    input: &[f32],
    mask: Option<&[u8]>,
    nx: usize, ny: usize, nz: usize,
    vsx: f32, vsy: f32, vsz: f32,
) {
    let hx2 = 1.0 / (vsx * vsx);
    let hy2 = 1.0 / (vsy * vsy);
    let hz2 = 1.0 / (vsz * vsz);
    let nxy = nx * ny;

    maybe_par_chunks_mut!(output, nxy).enumerate().for_each(|(k, out_slab)| {
        let k_offset = k * nxy;

        for j in 0..ny {
            let j_offset = j * nx;

            for i in 0..nx {
                let local = i + j_offset;
                let idx = local + k_offset;

                if mask.is_some_and(|m| m[idx] == 0) {
                    out_slab[local] = 0.0;
                    continue;
                }

                let a0 = input[idx];

                // Neumann BC: use center value at boundary (matching Julia)
                let a_xm = if i > 0 { input[(i - 1) + j_offset + k_offset] } else { a0 };
                let a_xp = if i + 1 < nx { input[(i + 1) + j_offset + k_offset] } else { a0 };
                let a_ym = if j > 0 { input[i + (j - 1) * nx + k_offset] } else { a0 };
                let a_yp = if j + 1 < ny { input[i + (j + 1) * nx + k_offset] } else { a0 };
                let a_zm = if k > 0 { input[i + j_offset + (k - 1) * nxy] } else { a0 };
                let a_zp = if k + 1 < nz { input[i + j_offset + (k + 1) * nxy] } else { a0 };

                // Laplacian: sum of second derivatives
                out_slab[local] = hx2 * (a_xm - 2.0 * a0 + a_xp)
                            + hy2 * (a_ym - 2.0 * a0 + a_yp)
                            + hz2 * (a_zm - 2.0 * a0 + a_zp);
            }
        }
    });
}


/// Compute gradient norm squared
#[inline]
fn grad_norm_sq(res: (f32, f32, f32)) -> f32 {
    let (dx, dy, dz) = res;
    4.0 * (1.0 / (dx * dx) + 1.0 / (dy * dy) + 1.0 / (dz * dz))
}

/// Compute the squared spectral norm of the TGV operator matrix via power iteration
///
/// The operator matrix M is:
/// [0,   g,  1 ]
/// [0,   0,  g ]
/// [g², w,  0 ]
///
/// We compute the largest eigenvalue of M^T * M using power iteration.
fn compute_operator_norm_sqr(g: f32, g2: f32, w: f32) -> f32 {
    // M^T * M = [g⁴,   g²w,    0    ]
    //           [g²w,  g²+w²,  g    ]
    //           [0,    g,      g²+1 ]
    let g4 = g2 * g2;
    let g2w = g2 * w;
    let g2_w2 = g2 + w * w;
    let g2_1 = g2 + 1.0;

    // Power iteration to find largest eigenvalue
    let mut v = [1.0f32, 1.0, 1.0];

    for _ in 0..20 {
        // Matrix-vector multiply: y = (M^T * M) * v
        let y0 = g4 * v[0] + g2w * v[1];
        let y1 = g2w * v[0] + g2_w2 * v[1] + g * v[2];
        let y2 = g * v[1] + g2_1 * v[2];

        // Compute norm
        let norm = (y0 * y0 + y1 * y1 + y2 * y2).sqrt();
        if norm < 1e-10 {
            break;
        }

        // Normalize
        v[0] = y0 / norm;
        v[1] = y1 / norm;
        v[2] = y2 / norm;
    }

    // Rayleigh quotient: eigenvalue = v^T * (M^T * M) * v
    let y0 = g4 * v[0] + g2w * v[1];
    let y1 = g2w * v[0] + g2_w2 * v[1] + g * v[2];
    let y2 = g * v[1] + g2_1 * v[2];

    v[0] * y0 + v[1] * y1 + v[2] * y2
}

/// L2 norm of 3-component vector
#[inline]
fn norm3(x: f32, y: f32, z: f32) -> f32 {
    (x * x + y * y + z * z).sqrt()
}

/// Frobenius norm of symmetric 3x3 tensor (6 components)
#[inline]
fn frobenius_norm(sxx: f32, sxy: f32, sxz: f32, syy: f32, syz: f32, szz: f32) -> f32 {
    (sxx * sxx + syy * syy + szz * szz + 2.0 * (sxy * sxy + sxz * sxz + syz * syz)).sqrt()
}

/// L-infinity projection for 3-component vector
#[inline]
fn project_linf3(px: &mut f32, py: &mut f32, pz: &mut f32, threshold: f32) {
    let norm = norm3(*px, *py, *pz);
    if norm > threshold {
        let scale = threshold / norm;
        *px *= scale;
        *py *= scale;
        *pz *= scale;
    }
}

/// L-infinity projection for 6-component symmetric tensor
#[inline]
fn project_linf6(
    qxx: &mut f32, qxy: &mut f32, qxz: &mut f32,
    qyy: &mut f32, qyz: &mut f32, qzz: &mut f32,
    threshold: f32,
) {
    let norm = frobenius_norm(*qxx, *qxy, *qxz, *qyy, *qyz, *qzz);
    if norm > threshold {
        let scale = threshold / norm;
        *qxx *= scale;
        *qxy *= scale;
        *qxz *= scale;
        *qyy *= scale;
        *qyz *= scale;
        *qzz *= scale;
    }
}

/// Pre-allocated workspace for TGV iteration
struct TgvWorkspace {
    // Primal variables
    chi: Vec<f32>,
    chi_: Vec<f32>,
    phi: Vec<f32>,
    phi_: Vec<f32>,
    wx: Vec<f32>,
    wy: Vec<f32>,
    wz: Vec<f32>,
    wx_: Vec<f32>,
    wy_: Vec<f32>,
    wz_: Vec<f32>,

    // Dual variables
    eta: Vec<f32>,
    px: Vec<f32>,
    py: Vec<f32>,
    pz: Vec<f32>,
    qxx: Vec<f32>,
    qxy: Vec<f32>,
    qxz: Vec<f32>,
    qyy: Vec<f32>,
    qyz: Vec<f32>,
    qzz: Vec<f32>,

    // Temporary buffers
    temp1: Vec<f32>,
    temp2: Vec<f32>,
    gx: Vec<f32>,
    gy: Vec<f32>,
    gz: Vec<f32>,

    // Symmetric gradient buffers (reused)
    sxx: Vec<f32>,
    sxy: Vec<f32>,
    sxz: Vec<f32>,
    syy: Vec<f32>,
    syz: Vec<f32>,
    szz: Vec<f32>,

    // Divergence buffers (reused)
    divqx: Vec<f32>,
    divqy: Vec<f32>,
    divqz: Vec<f32>,
}

impl TgvWorkspace {
    fn new(n: usize) -> Self {
        Self {
            chi: vec![0.0; n],
            chi_: vec![0.0; n],
            phi: vec![0.0; n],
            phi_: vec![0.0; n],
            wx: vec![0.0; n],
            wy: vec![0.0; n],
            wz: vec![0.0; n],
            wx_: vec![0.0; n],
            wy_: vec![0.0; n],
            wz_: vec![0.0; n],
            eta: vec![0.0; n],
            px: vec![0.0; n],
            py: vec![0.0; n],
            pz: vec![0.0; n],
            qxx: vec![0.0; n],
            qxy: vec![0.0; n],
            qxz: vec![0.0; n],
            qyy: vec![0.0; n],
            qyz: vec![0.0; n],
            qzz: vec![0.0; n],
            temp1: vec![0.0; n],
            temp2: vec![0.0; n],
            gx: vec![0.0; n],
            gy: vec![0.0; n],
            gz: vec![0.0; n],
            sxx: vec![0.0; n],
            sxy: vec![0.0; n],
            sxz: vec![0.0; n],
            syy: vec![0.0; n],
            syz: vec![0.0; n],
            szz: vec![0.0; n],
            divqx: vec![0.0; n],
            divqy: vec![0.0; n],
            divqz: vec![0.0; n],
        }
    }
}

/// Main TGV-QSM reconstruction
///
/// # Arguments
/// * `input` - Field in ppm or one wrapped echo in rad (see [`TgvInput`])
/// * `mask` - Binary mask
/// * `grid` - Volume grid (dimensions and voxel sizes)
/// * `params` - TGV parameters (α in the reference's rad units)
/// * `b0_dir` - B0 field direction
/// * `progress` - Progress callback `(iteration, total_iterations)`
///
/// Returns χ in ppm.
///
/// # Errors
/// If the input or mask length does not match `grid`, if the echo time (`te_s`, or the one
/// [`TgvParams::te`] picks from `echo_times`) or `b0_t` is not positive and finite, or if the
/// mask is empty after `params.erosions` erosions.
///
/// The reconstruction runs in single precision internally; inputs and
/// outputs use `f64` to match the rest of the crate's API.
pub fn tgv_qsm(
    input: TgvInput,
    mask: &[u8],
    grid: &crate::Grid,
    params: &TgvParams,
    b0_dir: (f64, f64, f64),
    mut progress: impl FnMut(usize, usize),
) -> Result<Vec<f64>, crate::error::InputError> {
    use crate::error::{check_len, InputError};
    let n_total = grid.n_total();
    check_len("TGV input", input.data().len(), n_total)?;
    check_len("TGV mask", mask.len(), n_total)?;
    let (te_s, b0_t) = input.te_b0(params.te).map_err(|e| InputError(format!("TGV: {e}")))?;
    if !(te_s.is_finite() && te_s > 0.0 && b0_t.is_finite() && b0_t > 0.0) {
        return Err(InputError(format!(
            "TGV needs te_s > 0 and b0_t > 0, got TE = {te_s} s, B0 = {b0_t} T"
        )));
    }
    // rad per ppm at this TE and B0: the input becomes phase in rad (the reference's units)
    // first, and χ is divided by it last (the reference's `scale`).
    let c = rad_per_ppm(te_s, b0_t);

    // The algorithm runs in single precision; convert the f64 boundary inputs.
    let data_f32: Vec<f32> = match input {
        TgvInput::FieldPpm { field, .. } => field.iter().map(|&v| (v * c) as f32).collect(),
        TgvInput::WrappedPhase { phase, .. } => phase.iter().map(|&v| v as f32).collect(),
    };
    let b0_dir = (b0_dir.0 as f32, b0_dir.1 as f32, b0_dir.2 as f32);

    let (nx, ny, nz) = grid.dims;
    let vsx = grid.vsx() as f32;
    let vsy = grid.vsy() as f32;
    let vsz = grid.vsz() as f32;
    let res = (vsx, vsy, vsz);
    let iterations = params.iterations.unwrap_or_else(|| get_default_iterations(res, params.step_size));

    // Erode the mask (from the array's edge too, unless `reference_edge_erosion`)
    let mut mask_eroded = mask.to_vec();
    for _ in 0..params.erosions {
        mask_eroded = erode_mask(&mask_eroded, nx, ny, nz, params.reference_edge_erosion);
    }
    // The bounding box below has no extent for an empty mask (its size underflowed).
    if !mask_eroded.iter().any(|&m| m != 0) {
        return Err(if mask.iter().any(|&m| m != 0) {
            InputError::param(
                "TGV erosions",
                params.erosions,
                "small enough to leave part of the mask (this mask erodes away)",
            )
        } else {
            InputError::new("TGV mask is empty")
        });
    }

    // Create mask0 (one more erosion for internal computations)
    let mask0 = erode_mask(&mask_eroded, nx, ny, nz, params.reference_edge_erosion);

    // Bounding box of the eroded mask with one voxel of padding (the reference's
    // `mask_box_indices`)
    let bbox = BoundingBox::from_mask(&mask_eroded, nx, ny, nz, 1);
    let (bx, by, bz) = bbox.dims();
    let b_total = bbox.total();
    let sub_grid = crate::Grid::new(bx, by, bz, vsx as f64, vsy as f64, vsz as f64);

    // Extract sub-volumes for the bounding box region
    let data_sub = extract_subvolume(&data_f32, &bbox, nx, ny, nz);
    let mask0_sub = extract_subvolume(&mask0, &bbox, nx, ny, nz);
    let mask_eroded_sub = extract_subvolume(&mask_eroded, &bbox, nx, ny, nz);

    // Laplacian of the phase in rad/mm² on the sub-volume
    let mut laplace_phi0 = match input {
        TgvInput::FieldPpm { .. } => {
            compute_field_laplacian(&data_sub, &mask0_sub, bx, by, bz, vsx, vsy, vsz)
        }
        TgvInput::WrappedPhase { .. } => {
            compute_phase_laplacian(&data_sub, &mask0_sub, bx, by, bz, vsx, vsy, vsz)
        }
    };

    // Subtract mean within mask
    let (sum, count): (f32, usize) = laplace_phi0.iter().zip(mask0_sub.iter())
        .filter(|(_, &m)| m != 0)
        .fold((0.0, 0), |(s, c), (&v, _)| (s + v, c + 1));
    if count > 0 {
        let mean = sum / count as f32;
        for (v, &m) in laplace_phi0.iter_mut().zip(mask0_sub.iter()) {
            if m != 0 {
                *v -= mean;
            }
        }
    }

    // Compute SVD-fitted oblique dipole stencil (matching Julia reference)
    let stencil = compute_oblique_stencil(res, b0_dir);

    // Compute step sizes for convergence (matching Julia implementation)
    let grad_norm_squared = grad_norm_sq(res);
    let grad_norm = grad_norm_squared.sqrt();
    let wave_norm: f32 = stencil.iter().flatten().flatten().map(|x| x.abs()).sum();
    let norm_sqr = compute_operator_norm_sqr(grad_norm, grad_norm_squared, wave_norm);

    let tau = 1.0 / norm_sqr.sqrt();
    let sigma = tau;

    // step_size is applied selectively in the updates:
    // - eta, phi: use base sigma/tau
    // - p, q, chi, w: use sigma * step_size / tau * step_size
    let sigma_step = sigma * params.step_size;
    let tau_step = tau * params.step_size;

    // Projection thresholds are alpha values (NOT 1/alpha!)
    // Julia: projects p to ||p|| <= alpha1, q to ||q|| <= alpha0
    let alpha = (params.alpha0, params.alpha1);

    // Pre-allocate all workspace buffers
    let mut ws = TgvWorkspace::new(b_total);

    // Main iteration loop (a fixed number of iterations, as in the reference)
    for iter in 0..iterations {
        progress(iter, iterations);

        // === DUAL UPDATE ===

        // 1. Update eta (data term dual)
        apply_laplacian_inplace(&mut ws.temp1, &ws.phi_, Some(&mask0_sub), bx, by, bz, vsx, vsy, vsz);
        apply_stencil(&mut ws.temp2, &ws.chi_, &stencil, Some(&mask0_sub), bx, by, bz);

        for i in 0..b_total {
            if mask0_sub[i] != 0 {
                ws.eta[i] += sigma * (-ws.temp1[i] + ws.temp2[i] - laplace_phi0[i]);
            }
        }

        // 2. Update p (gradient dual)
        // Julia: p += mask0 * sigma * grad(chi) - mask * sigma * w
        // Compute unmasked gradient first
        crate::utils::gradient::fgrad_inplace_f32(
            &mut ws.gx, &mut ws.gy, &mut ws.gz, &ws.chi_, &sub_grid
        );

        for i in 0..b_total {
            let in_mask0 = mask0_sub[i] != 0;
            let in_mask = mask_eroded_sub[i] != 0;

            if in_mask0 || in_mask {
                // gradient term scaled by mask0, w term scaled by mask
                let sigmaw0 = if in_mask0 { sigma_step } else { 0.0 };
                let sigmaw = if in_mask { sigma_step } else { 0.0 };

                ws.px[i] += sigmaw0 * ws.gx[i] - sigmaw * ws.wx_[i];
                ws.py[i] += sigmaw0 * ws.gy[i] - sigmaw * ws.wy_[i];
                ws.pz[i] += sigmaw0 * ws.gz[i] - sigmaw * ws.wz_[i];

                project_linf3(&mut ws.px[i], &mut ws.py[i], &mut ws.pz[i], alpha.1);
            }
        }

        // 3. Update q (symmetric gradient dual)
        crate::utils::gradient::symgrad_inplace_f32(
            &mut ws.sxx, &mut ws.sxy, &mut ws.sxz, &mut ws.syy, &mut ws.syz, &mut ws.szz,
            &ws.wx_, &ws.wy_, &ws.wz_, &sub_grid
        );

        for i in 0..b_total {
            if mask0_sub[i] != 0 {
                ws.qxx[i] += sigma_step * ws.sxx[i];
                ws.qxy[i] += sigma_step * ws.sxy[i];
                ws.qxz[i] += sigma_step * ws.sxz[i];
                ws.qyy[i] += sigma_step * ws.syy[i];
                ws.qyz[i] += sigma_step * ws.syz[i];
                ws.qzz[i] += sigma_step * ws.szz[i];

                project_linf6(
                    &mut ws.qxx[i], &mut ws.qxy[i], &mut ws.qxz[i],
                    &mut ws.qyy[i], &mut ws.qyz[i], &mut ws.qzz[i],
                    alpha.0
                );
            }
        }

        // === VARIABLE SWAP ===
        std::mem::swap(&mut ws.phi, &mut ws.phi_);
        std::mem::swap(&mut ws.chi, &mut ws.chi_);
        std::mem::swap(&mut ws.wx, &mut ws.wx_);
        std::mem::swap(&mut ws.wy, &mut ws.wy_);
        std::mem::swap(&mut ws.wz, &mut ws.wz_);

        // === PRIMAL UPDATE ===

        // 1. Update phi
        for i in 0..b_total {
            ws.temp1[i] = if mask0_sub[i] != 0 { ws.eta[i] } else { 0.0 };
        }
        // Julia: phi = (phi + tau*laplace(mask0*eta)) / (1 + mask*tau), at every voxel
        apply_laplacian_inplace(&mut ws.temp2, &ws.temp1, None, bx, by, bz, vsx, vsy, vsz);

        for i in 0..b_total {
            let denom = 1.0 + if mask_eroded_sub[i] != 0 { tau } else { 0.0 };
            ws.phi[i] = (ws.phi_[i] + tau * ws.temp2[i]) / denom;
        }

        // 2. Update chi
        crate::utils::gradient::bdiv_masked_inplace_f32(
            &mut ws.temp1, &ws.px, &ws.py, &ws.pz, &mask0_sub, &sub_grid
        );

        for i in 0..b_total {
            ws.gx[i] = if mask0_sub[i] != 0 { ws.eta[i] } else { 0.0 };
        }
        // Julia: chi = chi + tau*(div(mask0*p) - wave(mask0*eta)), at every voxel
        apply_stencil(&mut ws.temp2, &ws.gx, &stencil, None, bx, by, bz);

        for i in 0..b_total {
            ws.chi[i] = ws.chi_[i] + tau_step * (ws.temp1[i] - ws.temp2[i]);
        }

        // 3. Update w
        for i in 0..b_total {
            let m = if mask0_sub[i] != 0 { 1.0 } else { 0.0 };
            ws.sxx[i] = ws.qxx[i] * m;
            ws.sxy[i] = ws.qxy[i] * m;
            ws.sxz[i] = ws.qxz[i] * m;
            ws.syy[i] = ws.qyy[i] * m;
            ws.syz[i] = ws.qyz[i] * m;
            ws.szz[i] = ws.qzz[i] * m;
        }

        // The exact adjoint of the symmetric gradient in the q update (Chambolle–Pock needs
        // it); the reference's row-scaled version with `reference_divergence`.
        if params.reference_divergence {
            symdiv_rows_f32(
                &mut ws.divqx, &mut ws.divqy, &mut ws.divqz,
                &ws.sxx, &ws.sxy, &ws.sxz, &ws.syy, &ws.syz, &ws.szz,
                &sub_grid
            );
        } else {
            crate::utils::gradient::symdiv_inplace_f32(
                &mut ws.divqx, &mut ws.divqy, &mut ws.divqz,
                &ws.sxx, &ws.sxy, &ws.sxz, &ws.syy, &ws.syz, &ws.szz,
                &sub_grid
            );
        }

        // Julia: w_dest = w; if mask: w_dest += tau*(p + div(mask0*q))
        for i in 0..b_total {
            ws.wx[i] = ws.wx_[i];
            ws.wy[i] = ws.wy_[i];
            ws.wz[i] = ws.wz_[i];
            if mask_eroded_sub[i] != 0 {
                ws.wx[i] += tau_step * (ws.px[i] + ws.divqx[i]);
                ws.wy[i] += tau_step * (ws.py[i] + ws.divqy[i]);
                ws.wz[i] += tau_step * (ws.pz[i] + ws.divqz[i]);
            }
        }

        // === EXTRAGRADIENT UPDATE ===
        for i in 0..b_total {
            ws.phi_[i] = 2.0 * ws.phi[i] - ws.phi_[i];
            ws.chi_[i] = 2.0 * ws.chi[i] - ws.chi_[i];
            ws.wx_[i] = 2.0 * ws.wx[i] - ws.wx_[i];
            ws.wy_[i] = 2.0 * ws.wy[i] - ws.wy_[i];
            ws.wz_[i] = 2.0 * ws.wz[i] - ws.wz_[i];
        }
    }

    progress(iterations, iterations);

    // Full-size result inside the eroded mask, or over the whole box with
    // `reference_output_box`; chi is in rad units, divide by c for ppm
    let mut result = vec![0.0f32; n_total];
    if !params.reference_output_box {
        for (v, &m) in ws.chi.iter_mut().zip(mask_eroded_sub.iter()) {
            if m == 0 {
                *v = 0.0;
            }
        }
    }

    // Insert back into full volume
    insert_subvolume(&mut result, &ws.chi, &bbox, nx, ny, nz);

    Ok(result.iter().map(|&v| v as f64 / c).collect())
}

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

    #[test]
    fn test_dipole_stencil() {
        let stencil = compute_dipole_stencil((1.0, 1.0, 1.0), (0.0, 0.0, 1.0));

        let mut sum = 0.0f32;
        for k in 0..3 {
            for j in 0..3 {
                for i in 0..3 {
                    sum += stencil[i][j][k];
                }
            }
        }
        assert!(sum.abs() < 1e-6, "Stencil sum should be ~0, got {}", sum);
    }

    #[test]
    fn test_phase_laplacian() {
        let nx = 4;
        let ny = 4;
        let nz = 4;
        let n = nx * ny * nz;

        let phase = vec![1.0f32; n];
        let mask = vec![1u8; n];

        let lap = compute_phase_laplacian(&phase, &mask, nx, ny, nz, 1.0, 1.0, 1.0);

        let max_val = lap.iter().map(|x| x.abs()).fold(0.0f32, f32::max);
        assert!(max_val < 1e-5, "Laplacian of constant should be ~0, got max {}", max_val);
    }

    fn adjoint_gap(rows: bool) -> f64 {
        // <E w, q> (off-diagonals counted twice, as the q projection's norm counts them) against
        // -<w, div q>, on anisotropic voxels, with q zero on the box faces (as mask0 keeps it).
        let (nx, ny, nz) = (9, 8, 7);
        let n = nx * ny * nz;
        let grid = crate::Grid::new(nx, ny, nz, 0.9, 1.0, 1.4);
        let val = |s: usize, l: usize| ((l * 2654435761 + s * 40503) % 1000) as f32 / 500.0 - 1.0;
        let w: Vec<Vec<f32>> = (0..3).map(|s| (0..n).map(|l| val(s, l)).collect()).collect();
        let face = |l: usize| {
            let (i, j, k) = (l % nx, (l / nx) % ny, l / (nx * ny));
            i == 0 || j == 0 || k == 0 || i + 1 == nx || j + 1 == ny || k + 1 == nz
        };
        let q: Vec<Vec<f32>> = (0..6).map(|s| (0..n).map(|l| if face(l) { 0.0 } else { val(s + 3, l) }).collect()).collect();
        let mut e: Vec<Vec<f32>> = vec![vec![0.0; n]; 6];
        let [exx, exy, exz, eyy, eyz, ezz] = &mut e[..] else { unreachable!() };
        crate::utils::gradient::symgrad_inplace_f32(exx, exy, exz, eyy, eyz, ezz, &w[0], &w[1], &w[2], &grid);
        let mut d: Vec<Vec<f32>> = vec![vec![0.0; n]; 3];
        let [dx, dy, dz] = &mut d[..] else { unreachable!() };
        if rows {
            symdiv_rows_f32(dx, dy, dz, &q[0], &q[1], &q[2], &q[3], &q[4], &q[5], &grid);
        } else {
            crate::utils::gradient::symdiv_inplace_f32(dx, dy, dz, &q[0], &q[1], &q[2], &q[3], &q[4], &q[5], &grid);
        }
        let weight = [1.0, 2.0, 2.0, 1.0, 2.0, 1.0];
        let lhs: f64 = (0..6).map(|s| (0..n).map(|l| weight[s] * e[s][l] as f64 * q[s][l] as f64).sum::<f64>()).sum();
        let rhs: f64 = -(0..3).map(|s| (0..n).map(|l| w[s][l] as f64 * d[s][l] as f64).sum::<f64>()).sum::<f64>();
        (lhs - rhs).abs() / lhs.abs().max(rhs.abs())
    }

    #[test]
    fn default_w_update_divergence_is_the_adjoint_and_the_references_is_not() {
        assert!(adjoint_gap(false) < 1e-5, "default: relative gap {}", adjoint_gap(false));
        assert!(adjoint_gap(true) > 1e-2, "reference: relative gap {}", adjoint_gap(true));
    }

    #[test]
    fn test_erode_mask() {
        let nx = 5;
        let ny = 5;
        let nz = 5;

        // Default: outside the array counts as outside the mask, so a full mask erodes from
        // the edge; the reference (skip_outside) does not erode it.
        let mask = vec![1u8; nx * ny * nz];
        let eroded = erode_mask(&mask, nx, ny, nz, false);
        assert_eq!(eroded.iter().filter(|&&m| m == 1).count(), 27);
        assert!(erode_mask(&mask, nx, ny, nz, true).iter().all(|&m| m == 1));

        // A hole erodes its 6 neighbours, not the diagonals.
        let mut mask = mask;
        let center = 2 + 2 * nx + 2 * nx * ny;
        mask[center] = 0;
        let eroded = erode_mask(&mask, nx, ny, nz, true);
        assert_eq!(eroded.iter().filter(|&&m| m == 0).count(), 7);
        assert_eq!(eroded[center + 1], 0);
        assert_eq!(eroded[center + 1 + nx], 1);
    }

    #[test]
    fn test_default_alpha() {
        // The reference's get_default_alpha: [0.001, 0.001] + [0.001, 0.002]·(r − 1), ≥ 0.
        let close = |a: (f32, f32), b: (f32, f32)| (a.0 - b.0).abs() < 1e-9 && (a.1 - b.1).abs() < 1e-9;
        assert!(close(get_default_alpha(1.0), (0.001, 0.001)));
        assert!(close(get_default_alpha(2.0), (0.002, 0.003)));
        assert!(close(get_default_alpha(3.0), (0.003, 0.005)));
        assert!(close(get_default_alpha(4.0), (0.004, 0.007)));
        assert!(close(get_default_alpha(0.0), (0.0, 0.0)));
        let d = TgvParams::default();
        assert!(close((d.alpha0, d.alpha1), (0.002, 0.003)));
        assert_eq!((d.iterations, d.erosions, d.step_size), (None, 3, 3.0));
        assert!((rad_per_ppm(0.004, 7.0) - 7.4907).abs() < 1e-3);
    }

    #[test]
    fn test_bounding_box() {
        let nx = 10;
        let ny = 10;
        let nz = 10;
        let mut mask = vec![0u8; nx * ny * nz];

        // Set a small region in the center
        for k in 3..7 {
            for j in 3..7 {
                for i in 3..7 {
                    mask[i + j * nx + k * nx * ny] = 1;
                }
            }
        }

        let bbox = BoundingBox::from_mask(&mask, nx, ny, nz, 1);

        // Should be 3-1=2 to 6+1+1=8 (with padding 1)
        assert_eq!(bbox.i_min, 2);
        assert_eq!(bbox.i_max, 8);
        assert_eq!(bbox.j_min, 2);
        assert_eq!(bbox.j_max, 8);
    }

    #[test]
    fn test_tgv_qsm_small() {
        let n = 12;
        let n_total = n * n * n;
        let center = 6.0_f32;
        let radius = 4.0_f32;

        // Build a sphere mask
        let mut mask = vec![0u8; n_total];
        for k in 0..n {
            for j in 0..n {
                for i in 0..n {
                    let dx = i as f32 - center;
                    let dy = j as f32 - center;
                    let dz = k as f32 - center;
                    if (dx * dx + dy * dy + dz * dz).sqrt() < radius {
                        mask[i + j * n + k * n * n] = 1;
                    }
                }
            }
        }

        // Fill phase with a linear ramp (z direction) masked
        let mut phase = vec![0.0f64; n_total];
        for k in 0..n {
            for j in 0..n {
                for i in 0..n {
                    let idx = i + j * n + k * n * n;
                    if mask[idx] != 0 {
                        phase[idx] = 0.1 * k as f64;
                    }
                }
            }
        }

        let params = TgvParams {
            iterations: Some(10),
            erosions: 1,
            ..TgvParams::default()
        };

        let grid = crate::Grid::new(n, n, n, 1.0, 1.0, 1.0);
        let input = TgvInput::WrappedPhase { phase: &phase, te_s: 0.02, b0_t: 3.0 };
        let result = tgv_qsm(input, &mask, &grid, &params, (0.0, 0.0, 1.0), |_, _| {}).unwrap();

        assert_eq!(result.len(), n_total);

        // All values should be finite
        for &v in &result {
            assert!(v.is_finite(), "Result contains non-finite value: {}", v);
        }

        // There should be at least some non-zero values within the (eroded) mask
        let has_nonzero = result.iter().any(|&v| v.abs() > 1e-20);
        assert!(has_nonzero, "Result is entirely zero; expected non-zero values within mask");
    }

    #[test]
    fn test_oblique_stencil() {
        let stencil = compute_oblique_stencil((1.0, 1.0, 1.0), (0.0, 0.0, 1.0));

        // Sum of all elements should be ~0 (Laplacian-like operator)
        let mut sum = 0.0f32;
        for k in 0..3 {
            for j in 0..3 {
                for i in 0..3 {
                    sum += stencil[i][j][k];
                }
            }
        }
        assert!(
            sum.abs() < 1e-4,
            "Oblique stencil sum should be ~0, got {}",
            sum
        );

        // Center element should be approximately -sum(others), so verify it matches
        let mut off_sum = 0.0f32;
        for k in 0..3 {
            for j in 0..3 {
                for i in 0..3 {
                    if !(i == 1 && j == 1 && k == 1) {
                        off_sum += stencil[i][j][k];
                    }
                }
            }
        }
        assert!(
            (stencil[1][1][1] + off_sum).abs() < 1e-6,
            "Center should be -sum(others): center={}, off_sum={}",
            stencil[1][1][1], off_sum
        );
    }

    #[test]
    fn test_oblique_stencil_aniso() {
        let stencil = compute_oblique_stencil((1.0, 1.0, 2.0), (0.2, 0.3, 0.9));

        // Sum of all elements should be ~0
        let mut sum = 0.0f32;
        for k in 0..3 {
            for j in 0..3 {
                for i in 0..3 {
                    sum += stencil[i][j][k];
                }
            }
        }
        assert!(
            sum.abs() < 1e-4,
            "Anisotropic oblique stencil sum should be ~0, got {}",
            sum
        );
    }

    #[test]
    fn test_get_default_iterations() {
        // With isotropic 1mm voxels and step_size=1.0
        let it = get_default_iterations((1.0, 1.0, 1.0), 1.0);
        assert!(it >= 1000, "Iterations should be >= 1000 for 1mm iso, got {}", it);

        // With larger voxels the count should decrease (prod_res is larger)
        let it_large = get_default_iterations((2.0, 2.0, 2.0), 1.0);
        assert!(it_large >= 1000, "Iterations should still be >= 1000 for 2mm iso");

        // With very small voxels the count should be larger than 1mm iso
        let it_small = get_default_iterations((0.5, 0.5, 0.5), 1.0);
        assert!(it_small > it, "Smaller voxels should need more iterations: {} vs {}", it_small, it);

        // Higher step_size should reduce iterations
        let it_fast = get_default_iterations((1.0, 1.0, 1.0), 3.0);
        assert!(it_fast < it, "Higher step_size should give fewer iterations: {} vs {}", it_fast, it);
    }

    #[test]
    fn test_tgv_zero_input() {
        // Zero field => chi should be ~0, and progress reports every iteration.
        let n = 12;
        let n_total = n * n * n;
        let center = 6.0_f32;
        let radius = 4.0_f32;

        let mut mask = vec![0u8; n_total];
        for k in 0..n {
            for j in 0..n {
                for i in 0..n {
                    let dx = i as f32 - center;
                    let dy = j as f32 - center;
                    let dz = k as f32 - center;
                    if (dx * dx + dy * dy + dz * dz).sqrt() < radius {
                        mask[i + j * n + k * n * n] = 1;
                    }
                }
            }
        }

        let phase = vec![0.0f64; n_total];

        let params = TgvParams {
            iterations: Some(100),
            erosions: 1,
            ..TgvParams::default()
        };

        let grid = crate::Grid::new(n, n, n, 1.0, 1.0, 1.0);
        let progress_iters = std::cell::RefCell::new(Vec::new());
        let result = tgv_qsm(
            TgvInput::FieldPpm { field: &phase, echo_times: &[0.02], b0_t: 3.0 }, &mask, &grid, &params, (0.0, 0.0, 1.0),
            |iter, _total| { progress_iters.borrow_mut().push(iter); }
        ).unwrap();

        assert_eq!(result.len(), n_total);

        // With zero phase, all output values should be very close to zero
        let max_abs = result.iter().map(|v| v.abs()).fold(0.0f64, f64::max);
        assert!(
            max_abs < 1e-3,
            "Zero-phase TGV result should be ~0, got max abs {}",
            max_abs
        );

        // A fixed number of iterations, as in the reference: no early stop
        let iters = progress_iters.borrow();
        assert_eq!(iters.len(), 101);
        assert_eq!(*iters.last().unwrap(), 100);
    }

    /// The default returns χ inside the eroded mask only; `reference_output_box` returns the
    /// same χ there plus the reference's rim (nonzero spill-over just outside it).
    #[test]
    fn output_box_adds_the_rim_and_keeps_the_eroded_mask() {
        let n = 20;
        let mask = sphere_mask(n, 8.0);
        let field = test_field(n, 0.01);
        let grid = crate::Grid::new(n, n, n, 1.0, 1.0, 1.0);
        let run = |box_out: bool| {
            let p = TgvParams { iterations: Some(50), erosions: 1, reference_output_box: box_out, ..Default::default() };
            tgv_qsm(TgvInput::FieldPpm { field: &field, echo_times: &[0.02], b0_t: 3.0 }, &mask, &grid, &p, (0.0, 0.0, 1.0), |_, _| {})
                .unwrap()
        };
        let (trimmed, boxed) = (run(false), run(true));
        let eroded = erode_mask(&mask, n, n, n, false);
        let mut rim = 0;
        for l in 0..n * n * n {
            if eroded[l] != 0 {
                assert_eq!(trimmed[l], boxed[l]);
            } else {
                assert_eq!(trimmed[l], 0.0);
                rim += (boxed[l] != 0.0) as usize;
            }
        }
        assert!(rim > 0, "no χ outside the eroded mask with reference_output_box");
    }

    fn sphere_mask(n: usize, radius: f64) -> Vec<u8> {
        let c = (n as f64 - 1.0) / 2.0;
        let mut mask = vec![0u8; n * n * n];
        for k in 0..n {
            for j in 0..n {
                for i in 0..n {
                    let (dx, dy, dz) = (i as f64 - c, j as f64 - c, k as f64 - c);
                    if (dx * dx + dy * dy + dz * dz).sqrt() < radius {
                        mask[i + j * n + k * n * n] = 1;
                    }
                }
            }
        }
        mask
    }

    /// A smooth, non-harmonic field in ppm: a Gaussian bump plus a quadratic.
    fn test_field(n: usize, scale: f64) -> Vec<f64> {
        let c = (n as f64 - 1.0) / 2.0;
        let mut f = vec![0.0; n * n * n];
        for k in 0..n {
            for j in 0..n {
                for i in 0..n {
                    let (x, y, z) = (i as f64 - c, j as f64 - c, k as f64 - c);
                    let r2 = x * x + y * y + z * z;
                    f[i + j * n + k * n * n] = scale * ((-r2 / 8.0).exp() + 0.002 * (x * x - z * z));
                }
            }
        }
        f
    }

    /// FieldPpm uses the plain Laplacian (the reference's `get_laplace_phase_romeo` on an
    /// already unwrapped phase). On a field whose rad-equivalent exceeds π and whose neighbour
    /// differences are not small, DEL would be wrong, while the plain Laplacian is exact.
    #[test]
    fn field_ppm_uses_plain_laplacian_where_del_would_be_wrong() {
        let n = 16;
        let mask = vec![1u8; n * n * n];
        // Quadratic field f = a·x² ppm, ∇²f = 2a exactly with the 3-point stencil.
        let a = 0.05f32;
        let c = (n as f32 - 1.0) / 2.0;
        let field: Vec<f32> = (0..n * n * n)
            .map(|idx| { let x = (idx % n) as f32 - c; a * x * x })
            .collect();
        let rad_per = rad_per_ppm(0.020, 3.0) as f32; // ≈ 16.05
        let phase: Vec<f32> = field.iter().map(|&v| v * rad_per).collect();
        let max_rad = phase.iter().cloned().fold(0.0f32, f32::max);
        assert!(max_rad > std::f32::consts::PI, "test needs rad-equivalent > π, max {max_rad}");

        let plain = compute_field_laplacian(&field, &mask, n, n, n, 1.0, 1.0, 1.0);
        let del = compute_phase_laplacian(&phase, &mask, n, n, n, 1.0, 1.0, 1.0);

        // Interior voxels, away from the replicated boundary
        let mut max_err_plain = 0.0f32;
        let mut max_err_del = 0.0f32;
        for k in 1..n - 1 {
            for j in 1..n - 1 {
                for i in 1..n - 1 {
                    let idx = i + j * n + k * n * n;
                    max_err_plain = max_err_plain.max((plain[idx] - 2.0 * a).abs());
                    max_err_del = max_err_del.max((del[idx] / rad_per - 2.0 * a).abs());
                }
            }
        }
        assert!(max_err_plain < 1e-4, "plain Laplacian error {max_err_plain}");
        // DEL sees sin of the neighbour differences, which reach several rad here.
        assert!(max_err_del > 0.5 * 2.0 * a, "DEL error unexpectedly small: {max_err_del}");
    }

    /// α are in the reference's rad units at the given TE: the same α at twice the TE
    /// regularises like half the α at the original TE, as in the reference, where the phase
    /// (and so the data term) doubles with TE.
    #[test]
    fn field_ppm_alpha_is_in_rad_units_at_te() {
        let n = 20;
        let mask = sphere_mask(n, 8.5);
        let grid = crate::Grid::new(n, n, n, 1.0, 1.0, 1.0);
        let field = test_field(n, 0.1);
        let params = TgvParams { iterations: Some(300), erosions: 1, ..TgvParams::default() };
        let at = |te: f64, p: &TgvParams| {
            tgv_qsm(TgvInput::FieldPpm { field: &field, echo_times: &[te], b0_t: 3.0 }, &mask, &grid, p, (0.0, 0.0, 1.0), |_, _| {}).unwrap()
        };
        let chi_2te = at(0.02, &params);
        let half = TgvParams { alpha0: params.alpha0 / 2.0, alpha1: params.alpha1 / 2.0, ..params.clone() };
        let chi_half = at(0.01, &half);
        let chi_te = at(0.01, &params);

        let norm = chi_2te.iter().map(|v| v * v).sum::<f64>().sqrt();
        let rel = |a: &[f64], b: &[f64]| a.iter().zip(b).map(|(x, y)| (x - y).powi(2)).sum::<f64>().sqrt() / norm;
        assert!(norm > 0.0);
        assert!(rel(&chi_2te, &chi_half) < 1e-5, "relative difference {}", rel(&chi_2te, &chi_half));
        // and the TE does matter for a fixed α
        assert!(rel(&chi_2te, &chi_te) > 1e-3, "TE had no effect: {}", rel(&chi_2te, &chi_te));
    }

    /// WrappedPhase of a gently varying phase equals FieldPpm of the same field at the same
    /// TE (DEL ≈ plain Laplacian when neighbour differences are small).
    #[test]
    fn wrapped_phase_matches_field_ppm() {
        let n = 20;
        let mask = sphere_mask(n, 8.5);
        let grid = crate::Grid::new(n, n, n, 1.0, 1.0, 1.0);
        let field = test_field(n, 0.02);
        let params = TgvParams { iterations: Some(300), erosions: 1, ..TgvParams::default() };
        for te in [0.005, 0.020] {
            let chi_f = tgv_qsm(TgvInput::FieldPpm { field: &field, echo_times: &[te], b0_t: 3.0 }, &mask, &grid, &params, (0.0, 0.0, 1.0), |_, _| {}).unwrap();
            let norm = chi_f.iter().map(|v| v * v).sum::<f64>().sqrt();
            let c = rad_per_ppm(te, 3.0);
            let wrap = |x: f64| (x + std::f64::consts::PI).rem_euclid(2.0 * std::f64::consts::PI) - std::f64::consts::PI;
            let phase: Vec<f64> = field.iter().map(|v| wrap(v * c)).collect();
            let input = TgvInput::WrappedPhase { phase: &phase, te_s: te, b0_t: 3.0 };
            let chi_p = tgv_qsm(input, &mask, &grid, &params, (0.0, 0.0, 1.0), |_, _| {}).unwrap();
            let diff = chi_f.iter().zip(&chi_p).map(|(a, b)| (a - b).powi(2)).sum::<f64>().sqrt();
            assert!(diff / norm < 0.02, "TE {te}: relative difference {}", diff / norm);
        }
    }

    /// The fast DST-I is the direct sum, to rounding.
    #[test]
    fn dst3d_matches_the_direct_sum() {
        let n = 9;
        let x: Vec<f64> = (0..n * n * n).map(|i| ((i * 37 % 101) as f64 / 50.0) - 1.0).collect();
        let fast = dst3d(&x, n, rustdct::DctPlanner::<f64>::new().plan_dst1(n).as_ref());
        let direct = dst3d_direct(&x, n, &dst_sin_table(n));
        let err = fast.iter().zip(&direct).map(|(a, b)| (a - b).abs()).fold(0.0, f64::max);
        assert!(err < 1e-9, "max difference {err}");
    }

    /// Least-squares slope of `got` against `truth` over the voxels where `keep` holds:
    /// 1 when χ comes out in the units of the phantom, offset-invariant.
    fn slope(got: &[f64], truth: &[f64], keep: impl Fn(usize) -> bool) -> f64 {
        let idx: Vec<usize> = (0..truth.len()).filter(|&i| keep(i)).collect();
        let k = idx.len() as f64;
        let mt = idx.iter().map(|&i| truth[i]).sum::<f64>() / k;
        let mg = idx.iter().map(|&i| got[i]).sum::<f64>() / k;
        let cov: f64 = idx.iter().map(|&i| (truth[i] - mt) * (got[i] - mg)).sum();
        let var: f64 = idx.iter().map(|&i| (truth[i] - mt).powi(2)).sum();
        cov / var
    }

    /// TGV returns χ in ppm on the scale of the field it was given: three spheres of known χ,
    /// forward-simulated through the dipole kernel, come back with a slope near 1 against the
    /// truth, whether TGV gets the field in ppm or one wrapped echo with its TE and B0 (the
    /// default α shrink these small spheres to about 0.66). A field in ppm read as rad at TE1
    /// (QSM.rs#170) gave a slope of about 0.2 here.
    #[test]
    fn chi_comes_back_on_the_phantom_scale() {
        let n = 24;
        let c = n as f64 / 2.0;
        let at = |i: usize| ((i % n) as f64 + 0.5, ((i / n) % n) as f64 + 0.5, (i / (n * n)) as f64 + 0.5);
        let blobs = [((c - 4.0, c, c), 3.0, 0.15), ((c + 4.0, c + 2.0, c), 2.5, -0.10), ((c, c - 4.0, c + 3.0), 2.0, 0.20)];
        let chi: Vec<f64> = (0..n * n * n).map(|i| {
            let p = at(i);
            blobs.iter().rev()
                .find(|&&((x, y, z), r, _)| ((p.0 - x).powi(2) + (p.1 - y).powi(2) + (p.2 - z).powi(2)).sqrt() < r)
                .map_or(0.0, |b| b.2)
        }).collect();
        let r = |i: usize| { let p = at(i); ((p.0 - c).powi(2) + (p.1 - c).powi(2) + (p.2 - c).powi(2)).sqrt() };
        let mask: Vec<u8> = (0..n * n * n).map(|i| (r(i) < c - 1.0) as u8).collect();
        let grid = crate::Grid::new(n, n, n, 1.0, 1.0, 1.0);
        let kernel = crate::kernels::dipole::dipole_kernel(&grid, (0.0, 0.0, 1.0));
        let field = crate::fft::apply_real_kernel(&chi, &kernel, n, n, n);
        // Scored well inside the mask, clear of the eroded rim TGV leaves at zero.
        let scored = |i: usize| r(i) < c - 5.0;

        let params = TgvParams::default();
        let (te, b0) = (0.010, 3.0);
        let ppm = tgv_qsm(TgvInput::FieldPpm { field: &field, echo_times: &[te], b0_t: b0 }, &mask, &grid, &params, (0.0, 0.0, 1.0), |_, _| {}).unwrap();
        let s_ppm = slope(&ppm, &chi, scored);
        assert!((0.5..1.5).contains(&s_ppm), "FieldPpm slope {s_ppm:.3} against the phantom");

        let wrap = |x: f64| (x + std::f64::consts::PI).rem_euclid(2.0 * std::f64::consts::PI) - std::f64::consts::PI;
        let phase: Vec<f64> = field.iter().map(|v| wrap(v * rad_per_ppm(te, b0))).collect();
        let input = TgvInput::WrappedPhase { phase: &phase, te_s: te, b0_t: b0 };
        let rad = tgv_qsm(input, &mask, &grid, &params, (0.0, 0.0, 1.0), |_, _| {}).unwrap();
        let s_rad = slope(&rad, &chi, scored);
        eprintln!("phantom χ slope vs truth: FieldPpm {s_ppm:.3}, WrappedPhase {s_rad:.3}");
        assert!((0.5..1.5).contains(&s_rad), "WrappedPhase slope {s_rad:.3} against the phantom");
    }

    #[test]
    fn rejects_zero_te() {
        let n = 8;
        let mask = vec![1u8; n * n * n];
        let phase = vec![0.0; n * n * n];
        let grid = crate::Grid::new(n, n, n, 1.0, 1.0, 1.0);
        let input = TgvInput::WrappedPhase { phase: &phase, te_s: 0.0, b0_t: 3.0 };
        let err = tgv_qsm(input, &mask, &grid, &TgvParams::default(), (0.0, 0.0, 1.0), |_, _| {})
            .unwrap_err();
        assert!(err.to_string().contains("te_s > 0"), "{err}");
        let input = TgvInput::FieldPpm { field: &phase, echo_times: &[0.0], b0_t: 3.0 };
        assert!(tgv_qsm(input, &mask, &grid, &TgvParams::default(), (0.0, 0.0, 1.0), |_, _| {}).is_err());
        // and a mask of the wrong size is an error, not a panic
        let input = TgvInput::FieldPpm { field: &phase, echo_times: &[0.02], b0_t: 3.0 };
        assert!(tgv_qsm(input, &mask[1..], &grid, &TgvParams::default(), (0.0, 0.0, 1.0), |_, _| {}).is_err());
    }

    /// #199: erosions that empty the mask underflowed the bounding-box size; an empty mask, or
    /// one eroded away, is an error now.
    #[test]
    fn mask_eroded_away_is_an_error() {
        let n = 8;
        let grid = crate::Grid::new(n, n, n, 1.0, 1.0, 1.0);
        let field = vec![0.0; n * n * n];
        let mut mask = vec![0u8; n * n * n];
        mask[3 + 3 * n + 3 * n * n] = 1; // one voxel: gone after a single erosion
        let p = TgvParams { erosions: 1, iterations: Some(1), ..Default::default() };
        let input = TgvInput::FieldPpm { field: &field, echo_times: &[0.02], b0_t: 3.0 };
        let err = tgv_qsm(input, &mask, &grid, &p, (0.0, 0.0, 1.0), |_, _| {}).unwrap_err();
        assert!(err.0.contains("TGV erosions = 1"), "{err}");
        let input = TgvInput::FieldPpm { field: &field, echo_times: &[0.02], b0_t: 3.0 };
        let err = tgv_qsm(input, &[0u8; 512], &grid, &p, (0.0, 0.0, 1.0), |_, _| {}).unwrap_err();
        assert!(err.0.contains("TGV mask is empty"), "{err}");
    }
}