qsm_core/inversion/

heidi.rs

//! HEIDI — Homogeneity Enabled Incremental Dipole Inversion
//!
//! HEIDI splits k-space at the dipole cone and treats the two halves differently:
//!
//! * **Well-conditioned coefficients** (`|D(k)| > threshold`) are taken straight
//!   from an already-computed χ map — normally the minimally regularised LSQR
//!   solution of [`super::lsqr_qsm`] — and held as *exact* constraints.
//! * **Ill-conditioned coefficients** inside the cone, where direct division
//!   amplifies noise without bound, are recovered by minimising a weighted total
//!   variation subject to those constraints.
//!
//! ```text
//! minimise  Σ_v ‖ (M_x ∂_x χ, M_y ∂_y χ, M_z ∂_z χ) ‖   subject to   (F χ)|cone = (F χ_init)|cone
//! ```
//!
//! The "homogeneity" of the name is the weighting `M`. Each direction gets its own
//! 0/1 mask, built by thresholding the gradient of the *field* map: where the field
//! is locally flat the susceptibility is expected to be flat too, and TV is applied;
//! across an anatomical edge the mask opens and the gradient goes unpenalised. A
//! Laplacian test additionally opens the mask wherever the field curves sharply.
//! Voxels that are an edge in every direction get the weight
//! [`HeidiParams::gradient_mask_floor`], 0 in the shipped code (see that field).
//!
//! The constrained minimisation is solved with NESTA (Nesterov's smoothing with
//! continuation): the TV term is smoothed by a parameter μ that is annealed down
//! over several outer steps, each of which runs a fixed number of accelerated
//! gradient iterations. Because the data constraint is an orthogonal projection
//! (`A·Aᵀ = I` for a restricted FFT), the projection step is exact and costs one
//! FFT round trip.
//!
//! Ported from the authors' code, qsm_heidi's `inversion_heidi.m` and
//! `computesusceptibility_heidi.m` (<https://gitlab.com/R01NS114227/qsm_heidi>), run after
//! its `inversion_lsqr.m` ([`super::lsqr_qsm`]):
//!
//! * the seed is the LSQR solution on LSQR's grid (the mask's bounding box with a 40-voxel
//!   anti-aliasing frame, odd extents), unmasked, and HEIDI runs on that same grid;
//! * the cone is `|D| > 0.1` (the driver's `PostProcCone.threshold`) with `D` the
//!   spatial-domain unit dipole response ([`super::lsqr_qsm::spatial_dipole_kernel`]);
//! * NESTA runs with the driver's tolerance, machine epsilon, so every continuation step uses
//!   all of its inner iterations (8 × 500);
//! * the homogeneity masks are built from the phase denoised by ITK 3.20's gradient
//!   anisotropic diffusion (the executable `diffusionimagefilter.m` runs, 5 iterations, time
//!   step 0.05, conductance 1), in the phase's units. That executable works in `float` and sums
//!   the mean squared gradient that sets its conductance scale in `float`, which underestimates
//!   it (by about 1 % on a whole brain); QSM.rs computes it in double by default, and
//!   [`AnisotropicDiffusionParams::reference_itk`] reproduces the executable bit for bit. The
//!   two give the same masks but for the voxels at the gradient thresholds (about 0.07 % on an
//!   in-vivo brain, QSM-CI parity case `s005001-cbic-prisma-bridge-run1`).
//!
//! Units: the field is in ppm and χ comes back in ppm. The gradient and Laplacian
//! thresholds below are therefore in ppm/mm and ppm/voxel², independent of echo
//! time and field strength — upstream divides the phase by TE[ms]·B₀[T] before
//! thresholding, which is the same quantity. The cone fill is not scale-free: NESTA's
//! smoothing ends at an absolute μ (machine epsilon) and upstream's χ is in units of the
//! phase (`ppm · γ·TE·B₀·1e-6`, `γ` in rad/s/T), so the seed is converted to those units
//! before NESTA and back after, with [`HeidiParams::b0`] and [`HeidiParams::echo_times`] (scan
//! parameters the pipeline fills from metadata) at the echo time [`HeidiParams::te`] picks.
//!
//! Reference:
//! Schweser, F., Sommer, K., Deistung, A., Reichenbach, J.R. (2012).
//! "Quantitative susceptibility mapping for investigating subtle susceptibility
//! variations in the human brain." NeuroImage, 62(3):2083-2100.
//! https://doi.org/10.1016/j.neuroimage.2012.05.067
//!
//! NESTA:
//! Becker, S., Bobin, J., Candès, E.J. (2011). "NESTA: A fast and accurate
//! first-order method for sparse recovery." SIAM Journal on Imaging Sciences,
//! 4(1):1-39. https://doi.org/10.1137/090756855
//!
//! # Any seed: HEIDI as streaking reduction
//!
//! The cone fill does not depend on where the seed came from. The authors'
//! `computesusceptibility_heidi(suscInterim, …)` (qsm_heidi 4592631) takes the initial χ map
//! from outside and only refills its cone (the changelog calls the step "PostProcCone"), and
//! SEPIA's HEIDI add-on (Salman, Schweser and Chan, 2025) offers it after any dipole inversion
//! as "Streaking reduction by HEIDI": `Wrapper_QSM_HEIDI4all(localField, chiRaw, …)`, called from
//! `wrapper/QSMMacro.m` on the method's result. [`heidi_cone_fill`] is that step:
//!
//! * **Grid:** as `Wrapper_QSM_HEIDI4all`, the field, mask and seed are cropped to the mask's
//!   bounding box with the 40-voxel anti-aliasing frame (the grid [`heidi`] runs on), and χ is
//!   masked at the end (`chi .* mask`). The seed is the method's (masked) result as it is.
//! * **Units:** the qsm_heidi pipeline's, phase at the chosen echo time, for every seed, so
//!   [`heidi`] and [`heidi_cone_fill`] fill the cone the same way. `Wrapper_QSM_HEIDI4all` instead
//!   passes the field and the seed in Hz/2π, a factor `(2π)²·TE` (≈ 0.8 at TE = 20 ms) from
//!   those; the factor moves the gradient thresholds and NESTA's absolute final μ only.
//! * **Seed suitability.** The paper takes the well-conditioned region from a minimally
//!   regularised inversion and only the cone from the TV prior. A seed whose region outside the
//!   cone is close to data-consistent fits that: LSQR (the reference, which [`heidi`] runs), and
//!   TKD, TSVD, small-λ Tikhonov, NDI and RTS (RTS already splits k-space the same way). A
//!   regularised seed (MEDI, TV and the FANSI family, NLTV, iLSQR, STAR, AMP-PE, the
//!   deep-learning methods) keeps its regularisation outside the cone, as only the cone is
//!   replaced. The pipeline publishes this as the `heidi.seed-*` rules
//!   ([`crate::pipeline::rules::RULES`]).
//! * **Truncated seeds.** TKD and TSVD zero or clip `|D| < threshold`; with that threshold above
//!   the cone threshold, truncated coefficients would be kept as if they were data. Upstream
//!   stops with an error there ("DipoleFilter parameter must not be larger that Cone threshold",
//!   `computesusceptibility_heidi.m` L922-923), and so does the pipeline
//!   ([`crate::pipeline::heidi_seed_problem`]).
//! * **Not offered:** methods without a local field (the total-field and phase methods: the
//!   homogeneity weights come from the local field), χ-separation and multi-orientation
//!   reconstructions (upstream refuses more than one field: "You cannot use HEIDI with COSMOS").
//!
//! # Reference pipeline
//! Same inputs as [`super::lsqr_qsm`](mod@super::lsqr_qsm) (qsm_heidi commit 4592631): a SHARP-corrected local field
//! (`makefile` L14-20, `background_correction_sharp.mat`), the magnitude, the LSQR result, and
//! the mask `mask_tissue .* mask_phase .* phsInput.mask` (`inversion_heidi.m`). The echo time for
//! the phase units is the last one (`phase2susceptibilityfactor.m` L70-71), as the pipeline sets.
//!
//! # Pipeline rules
//!
//! The facts above that a configuration can depart from are published as rules in
//! [`crate::pipeline::rules::RULES`] (each with its strength, basis, reason and source): `lsqr.bg-sharp`, `heidi.te-last`.

use num_complex::Complex64;

use crate::fft::Fft3dWorkspace;
use crate::kernels::laplacian::laplacian_kernel;
use crate::utils::anisotropic_diffusion::{
    gradient_anisotropic_diffusion, AnisotropicDiffusionParams,
};
use crate::utils::apply_mask_zero;
use crate::{EchoTime, Grid};

use super::lsqr_qsm::{crop_grid, crop_pad, lsqr_qsm_padded, spatial_dipole_kernel, uncrop, LsqrQsmParams};
#[cfg(feature = "parallel")]
use rayon::prelude::*;

/// Gyromagnetic ratio of the proton, in rad/s/T: the reference implementation's value
/// (`getconstants.m`), not `crate::constants`. Keep it.
const GAMMA: f64 = 2.67522128e8;

/// Upstream expresses its mask thresholds per unit of `phase / (TE[ms] · B₀[T])`.
/// One ppm of field is `γ · 1e-3 · 1e-6` of that, so dividing an upstream threshold
/// by this constant converts it to ppm.
const UPSTREAM_PPM_SCALE: f64 = GAMMA * 1e-9;

/// Upstream `DEFAULT_GRADIENTMASK_THRESHOLD`, converted to ppm/mm.
pub const DEFAULT_GRADIENT_THRESHOLD_PPM: f64 = 0.00105 / UPSTREAM_PPM_SCALE;

/// Upstream `DEFAULT_GRADIENTMASK_LAPLACIANTRHESHOLD`, converted to ppm/voxel².
pub const DEFAULT_LAPLACIAN_THRESHOLD_PPM: f64 = 0.0159 / UPSTREAM_PPM_SCALE;

/// Parameters for [`heidi`].
#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
#[derive(Clone, Debug)]
pub struct HeidiParams {
    /// `|D(k)|` above which a coefficient counts as well-conditioned and is taken
    /// from `chi_init` unchanged.
    pub cone_threshold: f64,
    /// Field-gradient threshold in ppm/mm. Below it, a direction is considered
    /// homogeneous and its TV weight is switched on.
    pub gradient_threshold: f64,
    /// Open the homogeneity masks where the field's Laplacian is large.
    pub apply_laplacian_correction: bool,
    /// Field-Laplacian threshold in ppm per voxel² (the stencil is in voxel units,
    /// matching upstream).
    pub laplacian_threshold: f64,
    /// Anisotropic-diffusion smoothing applied to the field (in upstream's phase units) before
    /// the gradients are taken. `None` skips it. Its `reference_itk` reproduces the ITK
    /// executable upstream runs (see the module docs).
    pub denoise: Option<AnisotropicDiffusionParams>,
    /// TV weight kept at voxels that read as an edge in *every* direction.
    ///
    /// The default, 0, is what the shipped code does: `computesusceptibility_heidi` sets
    /// these voxels to `gradientMaskReplaceValue` (0.1) and then passes the masks through
    /// `fourierbasedresize(...) > 0.5`, which rounds the 0.1 away. (Its changelog, v10.6,
    /// says the masks "must contain values of 0.1"; set 0.1 to follow that instead.)
    pub gradient_mask_floor: f64,
    /// NESTA continuation steps on μ.
    ///
    /// This is the dominant quality knob, and more is not better. μ always lands
    /// on `mu_min`, so the count controls how many levels of TV minimisation are
    /// actually performed — past the first couple of levels μ is already far
    /// below the field's gradient scale and the term is effectively pure TV, so
    /// extra levels just keep flattening. On the CI phantom, at otherwise-default
    /// settings: 4 steps → r=0.8984, 6 → r=0.9017, 8 (upstream's
    /// `DEFAULT_TV_MAXINTITER`, kept here) → r=0.8757. Worth sweeping on your own
    /// data before trusting the default.
    pub continuation_steps: usize,
    /// Accelerated-gradient iterations per continuation step.
    pub inner_iterations: usize,
    /// Final (smallest) smoothing parameter μ.
    pub mu_min: f64,
    /// Relative objective change that ends an inner loop early.
    ///
    /// The driver `inversion_heidi.m` sets `PostProcCone.tol = eps`, which never
    /// triggers: every continuation step runs all `inner_iterations`.
    /// (`computesusceptibility_heidi`'s own default, used when the driver sets none, is
    /// `DEFAULT_TV_TOLVAR` = 1e-3.)
    pub tol: f64,
    /// Field strength in tesla (scan parameter, for the conversion to upstream's units)
    pub b0: f64,
    /// Echo times in seconds, in acquisition order (scan parameter, for the conversion to
    /// upstream's units).
    pub echo_times: Vec<f64>,
    /// Which echo time the field stands for the phase of, in the conversion to upstream's units. Default [`EchoTime::Last`]: qsm_heidi's
    /// `phase2susceptibilityfactor.m` uses the last echo time unless an echo is named
    /// (`heidi/phase2susceptibilityfactor.m` lines 70-71, qsm_heidi 4592631), and
    /// `inversion_heidi.m` names none (line 58).
    pub te: EchoTime,
}

impl Default for HeidiParams {
    /// `inversion_heidi.m`'s settings.
    fn default() -> Self {
        Self {
            // Upstream's `computesusceptibility_heidi` defaults to 0.14; the
            // shipped driver `inversion_heidi.m` overrides it to 0.1.
            cone_threshold: 0.1,
            gradient_threshold: DEFAULT_GRADIENT_THRESHOLD_PPM,
            apply_laplacian_correction: true,
            laplacian_threshold: DEFAULT_LAPLACIAN_THRESHOLD_PPM,
            denoise: Some(AnisotropicDiffusionParams {
                iterations: 5,
                time_step: 0.05,
                conductance: 1.0,
                reference_itk: false,
            }),
            gradient_mask_floor: 0.0,
            continuation_steps: 8,
            inner_iterations: 500,
            mu_min: f64::EPSILON,
            tol: f64::EPSILON,
            b0: 3.0,
            echo_times: vec![0.02],
            te: EchoTime::Last,
        }
    }
}

/// The three per-direction homogeneity weights, as used by the TV term.
pub struct HomogeneityMasks {
    pub x: Vec<f64>,
    pub y: Vec<f64>,
    pub z: Vec<f64>,
}

/// HEIDI dipole inversion: LSQR ([`super::lsqr_qsm`]) followed by HEIDI's cone fill, as
/// qsm_heidi's `inversion_lsqr.m` then `inversion_heidi.m`.
///
/// # Arguments
/// * `local_field` - Local field in ppm (`nx * ny * nz`, Fortran order)
/// * `mask` - Binary brain mask
/// * `magnitude` - Optional magnitude image, the LSQR seed's data weight
/// * `grid` - Volume grid
/// * `bdir` - B0 direction
/// * `lsqr` - Parameters of the LSQR seed (its `mask_output` is ignored: the seed stays
///   unmasked, as upstream's `susceptibilityRaw`)
/// * `params` - HEIDI parameters
/// * `progress` - Progress callback `(iteration, total_iterations)`, over both stages
///
/// # Returns
/// Susceptibility map in ppm, masked.
#[allow(clippy::too_many_arguments)] // the dipole-inversion inputs plus the seed's parameters
pub fn heidi(
    local_field: &[f64],
    mask: &[u8],
    magnitude: Option<&[f64]>,
    grid: &Grid,
    bdir: (f64, f64, f64),
    lsqr: &LsqrQsmParams,
    params: &HeidiParams,
    mut progress: impl FnMut(usize, usize),
) -> Vec<f64> {
    let lsqr_iters = lsqr.max_iter;
    let total = lsqr_iters + cone_fill_iterations(params);
    let seed = lsqr_qsm_padded(local_field, mask, magnitude, grid, bdir, lsqr, |i, _| progress(i, total));
    let Some((pgrid, crop, chi_seed)) = seed else {
        return vec![0.0; grid.n_total()];
    };
    let field = crop_pad(local_field, grid.dims, &crop);
    let pmask = crop_pad(mask, grid.dims, &crop);
    let chi = heidi_on_grid(&field, &pmask, &chi_seed, &pgrid, bdir, params, |i, _| progress(lsqr_iters + i, total));
    let mut chi = uncrop(&chi, grid.dims, &crop);
    apply_mask_zero(&mut chi, mask);
    chi
}

/// HEIDI's cone fill on any seed χ map: streaking reduction after a dipole inversion, as the
/// authors' `computesusceptibility_heidi(suscInterim, …)` and SEPIA's `Wrapper_QSM_HEIDI4all`.
///
/// The well-conditioned coefficients (`|D(k)| > cone_threshold`) of `chi_seed` are kept as exact
/// constraints and the dipole cone is refilled under HEIDI's homogeneity-weighted TV, exactly as
/// [`heidi`] does to its LSQR seed. See the module docs, *Any seed*, for which seeds suit it.
///
/// The grid is upstream's: the field, the mask and the seed are cropped to the mask's bounding
/// box and given LSQR's 40-voxel anti-aliasing frame (odd extents), as `Wrapper_QSM_HEIDI4all`
/// does with `zealouscrop(mask)` and `prepareforconvolution(…, 40)` on every input, the seed
/// included. A seed is whatever its method returned, normally masked; the frame around it is 0.
///
/// # Arguments
/// * `local_field` - Local field in ppm (`nx * ny * nz`, Fortran order), the one the seed was
///   inverted from: HEIDI's homogeneity weights come from it
/// * `mask` - Binary brain mask
/// * `chi_seed` - The seed χ map in ppm
/// * `grid` - Volume grid
/// * `bdir` - B0 direction
/// * `params` - HEIDI parameters ([`HeidiParams::b0`] and [`HeidiParams::echo_times`] are the
///   scan's)
/// * `progress` - Progress callback `(iteration, total_iterations)`
///
/// # Returns
/// Susceptibility map in ppm, masked (as `Wrapper_QSM_HEIDI4all`'s `chi .* mask`).
pub fn heidi_cone_fill(
    local_field: &[f64],
    mask: &[u8],
    chi_seed: &[f64],
    grid: &Grid,
    bdir: (f64, f64, f64),
    params: &HeidiParams,
    progress: impl FnMut(usize, usize),
) -> Vec<f64> {
    let n_total = grid.n_total();
    assert_eq!(local_field.len(), n_total, "field length must match grid");
    assert_eq!(mask.len(), n_total, "mask length must match grid");
    assert_eq!(chi_seed.len(), n_total, "seed length must match grid");
    let Some(crop) = crop_grid(mask, grid.dims) else {
        return vec![0.0; n_total];
    };
    let (vsx, vsy, vsz) = grid.voxel_size;
    let pgrid = Grid::new(crop[0].2, crop[1].2, crop[2].2, vsx, vsy, vsz);
    let field = crop_pad(local_field, grid.dims, &crop);
    let pmask = crop_pad(mask, grid.dims, &crop);
    let seed = crop_pad(chi_seed, grid.dims, &crop);
    let chi = heidi_on_grid(&field, &pmask, &seed, &pgrid, bdir, params, progress);
    let mut chi = uncrop(&chi, grid.dims, &crop);
    apply_mask_zero(&mut chi, mask);
    chi
}

/// Progress steps of the cone fill: [`heidi_cone_fill`], and [`heidi`]'s second stage.
pub fn cone_fill_iterations(params: &HeidiParams) -> usize {
    params.continuation_steps.max(1) * params.inner_iterations
}

/// HEIDI's cone fill on a grid with odd extents (upstream's), from the seed `chi_init`
/// (unmasked). Returns χ unmasked.
fn heidi_on_grid(
    local_field: &[f64],
    mask: &[u8],
    chi_init: &[f64],
    grid: &Grid,
    bdir: (f64, f64, f64),
    params: &HeidiParams,
    progress: impl FnMut(usize, usize),
) -> Vec<f64> {
    let n_total = grid.n_total();
    assert_eq!(local_field.len(), n_total, "field length must match grid");
    assert_eq!(mask.len(), n_total, "mask length must match grid");
    assert_eq!(chi_init.len(), n_total, "chi_init length must match grid");

    let masks = homogeneity_masks(local_field, mask, grid, params);
    let cone = cone_mask(grid, bdir, params.cone_threshold);
    // ppm -> upstream's phase units for χ (see the module docs, *Units*), and back.
    let to_phase = GAMMA * params.te.seconds(&params.echo_times, "HEIDI") * params.b0 * 1e-6;
    let seed: Vec<f64> = chi_init.iter().map(|v| v * to_phase).collect();
    let mut chi = nesta_cone_fill(&seed, &cone, &masks, grid, params, progress);
    for v in &mut chi {
        *v /= to_phase;
    }
    chi
}

/// `|D(k)| > threshold`: the k-space region the dipole model conditions well, with `D` the
/// spatial-domain unit dipole response, as upstream (`generateunitdipoleresponse`).
fn cone_mask(grid: &Grid, bdir: (f64, f64, f64), threshold: f64) -> Vec<bool> {
    spatial_dipole_kernel(grid, bdir)
        .into_iter()
        .map(|d| d.abs() > threshold)
        .collect()
}

/// Build the three per-direction homogeneity weights from the field map.
///
/// Public so callers can inspect or reuse the weighting; [`heidi`] calls it itself.
pub fn homogeneity_masks(
    local_field: &[f64],
    mask: &[u8],
    grid: &Grid,
    params: &HeidiParams,
) -> HomogeneityMasks {
    let (nx, ny, nz) = grid.dims;
    let n_total = grid.n_total();

    // Upstream masks the field before denoising, so the diffusion sees a clean
    // zero outside the brain rather than background junk.
    let mut field = local_field.to_vec();
    apply_mask_zero(&mut field, mask);

    // Upstream denoises the phase (`diffusionimagefilter` on `phase .* mask`), so the
    // diffusion runs in its units; it is scale-free but for rounding, which ITK's float
    // arithmetic (`reference_itk`) makes visible.
    let smoothed = match &params.denoise {
        Some(d) => {
            let to_phase = -GAMMA * params.te.seconds(&params.echo_times, "HEIDI") * params.b0 * 1e-6;
            let phase: Vec<f64> = field.iter().map(|v| v * to_phase).collect();
            gradient_anisotropic_diffusion(&phase, grid, d).into_iter().map(|v| v / to_phase).collect()
        }
        None => field,
    };

    // A direction is "homogeneous" where the field gradient is small.
    let mut mx = threshold_gradient(&smoothed, grid, 0, params.gradient_threshold);
    let mut my = threshold_gradient(&smoothed, grid, 1, params.gradient_threshold);
    let mut mz = threshold_gradient(&smoothed, grid, 2, params.gradient_threshold);

    if params.apply_laplacian_correction {
        // Sharp field curvature means an unresolved source: drop the TV weight
        // there in every direction. The stencil is in voxel units, as upstream's.
        let unit_grid = Grid::new(nx, ny, nz, 1.0, 1.0, 1.0);
        let lap_kernel = laplacian_kernel(&unit_grid, false);

        let mut ws = Fft3dWorkspace::new(nx, ny, nz);
        let mut kernel_ft: Vec<Complex64> =
            lap_kernel.iter().map(|&v| Complex64::new(v, 0.0)).collect();
        ws.fft3d(&mut kernel_ft);

        let mut buf: Vec<Complex64> =
            smoothed.iter().map(|&v| Complex64::new(v, 0.0)).collect();
        ws.fft3d(&mut buf);
        for (b, k) in buf.iter_mut().zip(&kernel_ft) {
            *b *= *k;
        }
        ws.ifft3d(&mut buf);

        for i in 0..n_total {
            if buf[i].re.abs() > params.laplacian_threshold {
                mx[i] = 0.0;
                my[i] = 0.0;
                mz[i] = 0.0;
            }
        }
    }

    // Restrict to the brain, then give all-edge voxels a residual weight so they
    // are not left completely unconstrained.
    for i in 0..n_total {
        if mask[i] == 0 {
            mx[i] = 0.0;
            my[i] = 0.0;
            mz[i] = 0.0;
        } else if mx[i] == 0.0 && my[i] == 0.0 && mz[i] == 0.0 {
            mx[i] = params.gradient_mask_floor;
            my[i] = params.gradient_mask_floor;
            mz[i] = params.gradient_mask_floor;
        }
    }

    HomogeneityMasks { x: mx, y: my, z: mz }
}

/// 1.0 where `|∂_axis field| < threshold`, else 0.0.
fn threshold_gradient(field: &[f64], grid: &Grid, axis: usize, threshold: f64) -> Vec<f64> {
    let g = discrete_gradient(field, grid, axis, false);
    g.into_iter()
        .map(|v| if v.abs() < threshold { 1.0 } else { 0.0 })
        .collect()
}

/// Backward difference along `axis`, scaled by voxel size, zero on the leading face.
///
/// `transpose` gives the operator upstream pairs with it: a forward difference,
/// zero on the *trailing* face. The two agree with a true adjoint pair except on
/// the two boundary faces, which carry no TV weight (the brain mask is zero there),
/// so the discrepancy never reaches the solution.
fn discrete_gradient_into(x: &[f64], out: &mut [f64], grid: &Grid, axis: usize, transpose: bool) {
    let (nx, ny, nz) = grid.dims;
    let (vsx, vsy, vsz) = grid.voxel_size;
    let (n_axis, voxel_size) = match axis {
        0 => (nx, vsx),
        1 => (ny, vsy),
        _ => (nz, vsz),
    };

    out.fill(0.0);
    if n_axis < 2 {
        return;
    }
    let inv = 1.0 / voxel_size;
    let slab = nx * ny;

    // Split the volume into independent pieces: one row per chunk along x, one
    // z-slab along y and z. Each chunk writes only its own output but may read
    // its neighbour's input, which is fine — `x` is shared immutably.
    let chunk = if axis == 0 { nx } else { slab };

    crate::maybe_par_chunks_mut!(out, chunk)
        .enumerate()
        .for_each(|(c, o)| {
            let base = c * chunk;
            match axis {
                0 => {
                    if transpose {
                        for p in 0..nx - 1 {
                            o[p] = (x[base + p] - x[base + p + 1]) * inv;
                        }
                    } else {
                        for p in 1..nx {
                            o[p] = (x[base + p] - x[base + p - 1]) * inv;
                        }
                    }
                }
                1 => {
                    if transpose {
                        for j in 0..ny - 1 {
                            for i in 0..nx {
                                let q = i + j * nx;
                                o[q] = (x[base + q] - x[base + q + nx]) * inv;
                            }
                        }
                    } else {
                        for j in 1..ny {
                            for i in 0..nx {
                                let q = i + j * nx;
                                o[q] = (x[base + q] - x[base + q - nx]) * inv;
                            }
                        }
                    }
                }
                _ => {
                    // `c` is the slice index along z.
                    if transpose {
                        if c + 1 < nz {
                            for p in 0..slab {
                                o[p] = (x[base + p] - x[base + p + slab]) * inv;
                            }
                        }
                    } else if c >= 1 {
                        for p in 0..slab {
                            o[p] = (x[base + p] - x[base + p - slab]) * inv;
                        }
                    }
                }
            }
        });
}

/// Allocating wrapper around [`discrete_gradient_into`].
fn discrete_gradient(x: &[f64], grid: &Grid, axis: usize, transpose: bool) -> Vec<f64> {
    let mut out = vec![0.0; x.len()];
    discrete_gradient_into(x, &mut out, grid, axis, transpose);
    out
}

/// Scratch buffers for the NESTA loop, so the inner iterations allocate nothing.
struct TvWorkspace {
    /// Per-direction weighted gradients, reused to carry `M·u` into the adjoint pass.
    d1: Vec<f64>,
    d2: Vec<f64>,
    d3: Vec<f64>,
    /// One `(semi_norm, ‖u‖²)` partial per z-slab, summed after the parallel pass.
    partials: Vec<(f64, f64)>,
}

impl TvWorkspace {
    fn new(grid: &Grid) -> Self {
        let n = grid.n_total();
        Self {
            d1: vec![0.0; n],
            d2: vec![0.0; n],
            d3: vec![0.0; n],
            partials: vec![(0.0, 0.0); grid.nz()],
        }
    }
}

/// Smoothed weighted-TV objective and its gradient.
///
/// Returns `f_μ(x)` and writes `∇f_μ(x)` into `df`. Two passes over the volume,
/// each parallel over z-slabs: the first builds the masked gradients and the
/// smoothing weight `w`, the second applies the adjoint differences.
fn tv_objective(
    x: &[f64],
    mu: f64,
    masks: &HomogeneityMasks,
    grid: &Grid,
    ws: &mut TvWorkspace,
    df: &mut [f64],
) -> f64 {
    let (nx, ny, nz) = grid.dims;
    let (vsx, vsy, vsz) = grid.voxel_size;
    let slab = nx * ny;
    let (ix, iy, iz) = (1.0 / vsx, 1.0 / vsy, 1.0 / vsz);

    // ── Pass 1: masked gradients, Nesterov smoothing, objective partials ──
    //
    // `w = max(μ, ‖∇x‖)` replaces the TV norm with a Huber-like function that is
    // quadratic below μ and linear above it. The buffers come out holding
    // `M·u = M²·Dx / w`, which is exactly what the adjoint pass needs.
    crate::maybe_par_chunks_mut!(ws.d1, slab)
        .zip(crate::maybe_par_chunks_mut!(ws.d2, slab))
        .zip(crate::maybe_par_chunks_mut!(ws.d3, slab))
        .zip(crate::maybe_par_chunks_mut!(ws.partials, 1))
        .enumerate()
        .for_each(|(k, (((d1s, d2s), d3s), acc))| {
            let base = k * slab;
            let (mut semi, mut nsq) = (0.0, 0.0);

            for j in 0..ny {
                for i in 0..nx {
                    let q = i + j * nx;
                    let idx = base + q;

                    // Backward differences, zero on the leading face of each axis.
                    let gx = if i >= 1 { (x[idx] - x[idx - 1]) * ix } else { 0.0 };
                    let gy = if j >= 1 { (x[idx] - x[idx - nx]) * iy } else { 0.0 };
                    let gz = if k >= 1 { (x[idx] - x[idx - slab]) * iz } else { 0.0 };

                    let (a, b, c) = (
                        masks.x[idx] * gx,
                        masks.y[idx] * gy,
                        masks.z[idx] * gz,
                    );
                    let mag_sq = a * a + b * b + c * c;
                    let w = mag_sq.sqrt().max(mu);
                    semi += mag_sq / w;
                    nsq += mag_sq / (w * w);

                    let inv_w = 1.0 / w;
                    d1s[q] = masks.x[idx] * a * inv_w;
                    d2s[q] = masks.y[idx] * b * inv_w;
                    d3s[q] = masks.z[idx] * c * inv_w;
                }
            }

            acc[0] = (semi, nsq);
        });

    let semi_norm: f64 = ws.partials.iter().map(|p| p.0).sum();
    let norm_u_sq: f64 = ws.partials.iter().map(|p| p.1).sum();

    // ── Pass 2: ∇f = Dᵀ(M u), summed over the three directions ──
    let (d1, d2, d3) = (&ws.d1, &ws.d2, &ws.d3);
    crate::maybe_par_chunks_mut!(df, slab)
        .enumerate()
        .for_each(|(k, out)| {
            let base = k * slab;
            for j in 0..ny {
                for i in 0..nx {
                    let q = i + j * nx;
                    let idx = base + q;

                    // Forward differences, zero on the trailing face of each axis.
                    let mut v = 0.0;
                    if i + 1 < nx {
                        v += (d1[idx] - d1[idx + 1]) * ix;
                    }
                    if j + 1 < ny {
                        v += (d2[idx] - d2[idx + nx]) * iy;
                    }
                    if k + 1 < nz {
                        v += (d3[idx] - d3[idx + slab]) * iz;
                    }
                    out[q] = v;
                }
            }
        });

    semi_norm - 0.5 * mu * norm_u_sq
}

/// NESTA with continuation: fill the dipole cone under a weighted-TV prior while
/// holding the well-conditioned coefficients of `chi_init` fixed.
fn nesta_cone_fill(
    chi_init: &[f64],
    cone: &[bool],
    masks: &HomogeneityMasks,
    grid: &Grid,
    params: &HeidiParams,
    mut progress: impl FnMut(usize, usize),
) -> Vec<f64> {
    let (nx, ny, nz) = grid.dims;
    let n = grid.n_total();

    let mut fft = Fft3dWorkspace::new(nx, ny, nz);
    let mut cbuf: Vec<Complex64> = vec![Complex64::new(0.0, 0.0); n];

    // P(x) = F⁻¹( cone · F x ): projection onto the well-conditioned subspace.
    // A = restrict∘F and Aᵀ = F⁻¹∘zero-fill satisfy A·Aᵀ = I, so P is idempotent
    // and the NESTA data step below is an exact projection.
    let project = |x: &[f64], out: &mut [f64], fft: &mut Fft3dWorkspace, cbuf: &mut Vec<Complex64>| {
        for (c, &v) in cbuf.iter_mut().zip(x) {
            *c = Complex64::new(v, 0.0);
        }
        fft.fft3d(cbuf);
        for (c, &keep) in cbuf.iter_mut().zip(cone) {
            if !keep {
                *c = Complex64::new(0.0, 0.0);
            }
        }
        fft.ifft3d(cbuf);
        for (o, c) in out.iter_mut().zip(cbuf.iter()) {
            *o = c.re;
        }
    };

    // The data term: the cone content of chi_init, which every iterate must match.
    let mut p_b = vec![0.0; n];
    project(chi_init, &mut p_b, &mut fft, &mut cbuf);

    // TV gradients use relative voxel sizes, as upstream normalises them.
    let (vsx, vsy, vsz) = grid.voxel_size;
    let vmax = vsx.max(vsy).max(vsz);
    let tv_grid = Grid::new(nx, ny, nz, vsx / vmax, vsy / vmax, vsz / vmax);

    // Lipschitz constant of the smoothed TV gradient: ‖D‖² / μ, with ‖D‖² summed
    // over the three difference operators (2/v per direction).
    let lip_geometry = 4.0 / (vsx / vmax).powi(2)
        + 4.0 / (vsy / vmax).powi(2)
        + 4.0 / (vsz / vmax).powi(2);

    let mu0 = 0.9 * p_b.iter().fold(0.0f64, |acc, v| acc.max(v.abs()));
    if mu0 <= 0.0 {
        // Nothing in the cone — no data to honour, and TV alone gives zero.
        return vec![0.0; n];
    }

    let steps = params.continuation_steps.max(1);
    let gamma = (params.mu_min / mu0).powf(1.0 / steps as f64);

    let mut ws = TvWorkspace::new(&tv_grid);

    let mut xplug = p_b.clone();
    let mut df = vec![0.0; n];
    let mut wk = vec![0.0; n];
    let mut cp = vec![0.0; n];
    let mut proj = vec![0.0; n];
    let mut yk = vec![0.0; n];
    let mut zk = vec![0.0; n];

    let total_iters = steps * params.inner_iterations;
    let mut done = 0usize;
    let mut mu = mu0;

    for _ in 0..steps {
        mu *= gamma;
        let lmu = lip_geometry / mu;
        let inv_lmu = 1.0 / lmu;

        let mut xk = xplug.clone();
        wk.fill(0.0);
        // `fmean` is a short history of objective values; the stopping test
        // compares the current objective against their mean.
        let mut fmean: Vec<f64> = vec![f64::MIN_POSITIVE / 10.0];

        for k in 0..params.inner_iterations {
            let fx = tv_objective(&xk, mu, masks, &tv_grid, &mut ws, &mut df);

            // Upstream aborts the whole run on divergence. Rewinding to this
            // step's starting point instead keeps the previous continuation
            // level's answer, which is a usable map rather than an error.
            if !fx.is_finite() {
                xk.copy_from_slice(&xplug);
                break;
            }

            // yk: gradient step from xk, projected back onto the constraint set.
            for i in 0..n {
                cp[i] = xk[i] - inv_lmu * df[i];
            }
            project(&cp, &mut proj, &mut fft, &mut cbuf);
            for i in 0..n {
                yk[i] = p_b[i] + cp[i] - proj[i];
            }

            let fbar = fmean.iter().sum::<f64>() / fmean.len() as f64;
            let qp = (fx - fbar).abs() / fbar;

            done += 1;
            progress(done, total_iters);

            // Upstream breaks before the momentum update, keeping the current
            // iterate rather than the freshly projected `yk`.
            if qp <= params.tol {
                break;
            }

            fmean.insert(0, fx);
            fmean.truncate(10);

            // zk: step from the anchor point along the accumulated weighted
            // gradient — the Nesterov momentum term.
            let apk = 0.5 * (k as f64 + 1.0);
            let tauk = 2.0 / (k as f64 + 3.0);
            for i in 0..n {
                wk[i] += apk * df[i];
                cp[i] = xplug[i] - inv_lmu * wk[i];
            }
            project(&cp, &mut proj, &mut fft, &mut cbuf);
            for i in 0..n {
                zk[i] = p_b[i] + cp[i] - proj[i];
                xk[i] = tauk * zk[i] + (1.0 - tauk) * yk[i];
            }
        }

        xplug = xk;
    }

    xplug
}

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

    fn test_grid(n: usize) -> Grid {
        Grid::new(n, n, n, 1.0, 1.0, 1.0)
    }

    #[test]
    fn discrete_gradient_matches_upstream_edges() {
        let g = Grid::new(4, 1, 1, 1.0, 1.0, 1.0);
        let x = vec![1.0, 3.0, 6.0, 10.0];

        let fwd = discrete_gradient(&x, &g, 0, false);
        assert_eq!(fwd, vec![0.0, 2.0, 3.0, 4.0]);

        let tr = discrete_gradient(&x, &g, 0, true);
        assert_eq!(tr, vec![-2.0, -3.0, -4.0, 0.0]);
    }

    #[test]
    fn discrete_gradient_scales_by_voxel_size() {
        let g = Grid::new(3, 1, 1, 2.0, 1.0, 1.0);
        let x = vec![0.0, 4.0, 10.0];
        let fwd = discrete_gradient(&x, &g, 0, false);
        assert_eq!(fwd, vec![0.0, 2.0, 3.0]);
    }

    #[test]
    fn thresholds_convert_from_upstream_units() {
        // A 0.00105 threshold in phase/(TE[ms]·B0[T]) is ~3.93e-3 ppm/mm.
        assert!((DEFAULT_GRADIENT_THRESHOLD_PPM - 3.9249e-3).abs() < 1e-6);
        assert!((DEFAULT_LAPLACIAN_THRESHOLD_PPM - 0.059434).abs() < 1e-5);
    }

    #[test]
    fn homogeneity_masks_open_at_edges() {
        let n = 12;
        let g = test_grid(n);
        let n_total = g.n_total();

        // Flat left half, steep ramp on the right.
        let mut field = vec![0.0; n_total];
        for k in 0..n {
            for j in 0..n {
                for i in 0..n {
                    field[i + j * n + k * n * n] =
                        if i < n / 2 { 0.0 } else { 0.05 * (i - n / 2) as f64 };
                }
            }
        }
        let mask = vec![1u8; n_total];
        let params = HeidiParams {
            denoise: None,
            apply_laplacian_correction: false,
            ..Default::default()
        };
        let m = homogeneity_masks(&field, &mask, &g, &params);

        // Deep in the flat region the x mask is on.
        let flat = 2 + (n / 2) * n + (n / 2) * n * n;
        assert_eq!(m.x[flat], 1.0);
        // On the ramp the gradient is 0.05 ppm/mm, far above the 3.9e-3 threshold.
        let ramp = (n / 2 + 2) + (n / 2) * n + (n / 2) * n * n;
        assert_eq!(m.x[ramp], 0.0);
        // y and z are flat everywhere, so the all-edge floor never triggers.
        assert_eq!(m.y[ramp], 1.0);
    }

    #[test]
    fn all_edge_voxels_get_the_floor() {
        let n = 8;
        let g = test_grid(n);
        let n_total = g.n_total();

        // A field that ramps steeply in every direction.
        let mut field = vec![0.0; n_total];
        for k in 0..n {
            for j in 0..n {
                for i in 0..n {
                    field[i + j * n + k * n * n] = 0.1 * (i + j + k) as f64;
                }
            }
        }
        let mask = vec![1u8; n_total];
        let params = HeidiParams {
            denoise: None,
            apply_laplacian_correction: false,
            gradient_mask_floor: 0.1,
            ..Default::default()
        };
        let m = homogeneity_masks(&field, &mask, &g, &params);

        let inner = 4 + 4 * n + 4 * n * n;
        assert_eq!(m.x[inner], 0.1);
        assert_eq!(m.y[inner], 0.1);
        assert_eq!(m.z[inner], 0.1);
    }

    #[test]
    fn masks_are_zero_outside_the_brain() {
        let n = 8;
        let g = test_grid(n);
        let n_total = g.n_total();
        let field = vec![0.0; n_total];
        let mut mask = vec![0u8; n_total];
        mask[4 + 4 * n + 4 * n * n] = 1;

        let params = HeidiParams { denoise: None, ..Default::default() };
        let m = homogeneity_masks(&field, &mask, &g, &params);
        for (i, &inside) in mask.iter().enumerate() {
            if inside == 0 {
                assert_eq!(m.x[i], 0.0);
                assert_eq!(m.y[i], 0.0);
                assert_eq!(m.z[i], 0.0);
            }
        }
    }

    /// The defining property: coefficients outside the cone are free, but the
    /// well-conditioned ones must come back unchanged.
    #[test]
    fn cone_coefficients_are_preserved() {
        let n = 17;
        let g = test_grid(n);
        let n_total = g.n_total();
        let bdir = (0.0, 0.0, 1.0);

        let mut chi_init = vec![0.0; n_total];
        let c = n as f64 / 2.0;
        for k in 0..n {
            for j in 0..n {
                for i in 0..n {
                    let r = ((i as f64 - c).powi(2)
                        + (j as f64 - c).powi(2)
                        + (k as f64 - c).powi(2))
                    .sqrt();
                    chi_init[i + j * n + k * n * n] = if r < 4.0 { 0.1 } else { 0.0 };
                }
            }
        }

        let field = vec![0.0; n_total];
        let mask = vec![1u8; n_total];
        let params = HeidiParams {
            denoise: None,
            continuation_steps: 3,
            inner_iterations: 20,
            ..Default::default()
        };
        let chi = heidi_on_grid(&field, &mask, &chi_init, &g, bdir, &params, |_, _| {});

        // Compare the two maps' k-space inside the cone.
        let cone = cone_mask(&g, bdir, params.cone_threshold);
        let mut ws = Fft3dWorkspace::new(n, n, n);
        let mut a: Vec<Complex64> = chi_init.iter().map(|&v| Complex64::new(v, 0.0)).collect();
        let mut b: Vec<Complex64> = chi.iter().map(|&v| Complex64::new(v, 0.0)).collect();
        ws.fft3d(&mut a);
        ws.fft3d(&mut b);

        let mut worst: f64 = 0.0;
        let mut scale: f64 = 0.0;
        for i in 0..n_total {
            if cone[i] {
                worst = worst.max((a[i] - b[i]).norm());
                scale = scale.max(a[i].norm());
            }
        }
        assert!(
            worst < 1e-6 * scale.max(1.0),
            "cone coefficients drifted by {worst} (scale {scale})"
        );
    }

    /// The fused two-pass TV kernel must agree with the straightforward
    /// composition of `discrete_gradient` it replaced.
    #[test]
    fn fused_tv_matches_the_naive_composition() {
        let (nx, ny, nz) = (9usize, 7usize, 8usize);
        let g = Grid::new(nx, ny, nz, 1.0, 1.3, 0.8);
        let n = g.n_total();

        let x: Vec<f64> = (0..n).map(|i| ((i as f64) * 0.37).sin()).collect();
        let masks = HomogeneityMasks {
            x: (0..n).map(|i| if i % 5 == 0 { 0.1 } else { 1.0 }).collect(),
            y: (0..n).map(|i| if i % 3 == 0 { 0.0 } else { 1.0 }).collect(),
            z: (0..n).map(|i| if i % 7 == 0 { 0.1 } else { 1.0 }).collect(),
        };
        let mu = 0.05;

        let mut ws = TvWorkspace::new(&g);
        let mut df = vec![0.0; n];
        let fx = tv_objective(&x, mu, &masks, &g, &mut ws, &mut df);

        // Reference: build each direction separately, then apply the adjoints.
        let mut d1 = discrete_gradient(&x, &g, 0, false);
        let mut d2 = discrete_gradient(&x, &g, 1, false);
        let mut d3 = discrete_gradient(&x, &g, 2, false);
        for i in 0..n {
            d1[i] *= masks.x[i];
            d2[i] *= masks.y[i];
            d3[i] *= masks.z[i];
        }
        let mut semi = 0.0;
        let mut nsq = 0.0;
        for i in 0..n {
            let mag_sq = d1[i] * d1[i] + d2[i] * d2[i] + d3[i] * d3[i];
            let w = mag_sq.sqrt().max(mu);
            semi += mag_sq / w;
            nsq += mag_sq / (w * w);
            d1[i] = masks.x[i] * d1[i] / w;
            d2[i] = masks.y[i] * d2[i] / w;
            d3[i] = masks.z[i] * d3[i] / w;
        }
        let ref_fx = semi - 0.5 * mu * nsq;

        let gx = discrete_gradient(&d1, &g, 0, true);
        let gy = discrete_gradient(&d2, &g, 1, true);
        let gz = discrete_gradient(&d3, &g, 2, true);

        assert!((fx - ref_fx).abs() < 1e-12, "objective {fx} vs {ref_fx}");
        for i in 0..n {
            let expect = gx[i] + gy[i] + gz[i];
            assert!(
                (df[i] - expect).abs() < 1e-12,
                "gradient mismatch at {i}: {} vs {expect}",
                df[i]
            );
        }
    }

    /// End to end on an even, prime-factored grid: LSQR's cropped odd grid is used inside and
    /// the result comes back at the caller's dimensions, masked.
    #[test]
    fn heidi_runs_on_upstream_grid_and_returns_the_input_grid() {
        let (nx, ny, nz) = (12usize, 11usize, 13usize);
        let g = Grid::new(nx, ny, nz, 1.0, 1.0, 1.0);
        let n_total = g.n_total();
        let field: Vec<f64> = (0..n_total).map(|i| ((i as f64) * 0.31).sin() * 0.05).collect();
        let mut mask = vec![0u8; n_total];
        for k in 3..10 {
            for j in 2..9 {
                for i in 3..9 {
                    mask[i + nx * (j + ny * k)] = 1;
                }
            }
        }
        let lsqr = LsqrQsmParams { max_iter: 10, ..Default::default() };
        let params = HeidiParams { denoise: None, continuation_steps: 2, inner_iterations: 10, ..Default::default() };
        let mut last = 0;
        let chi = heidi(&field, &mask, None, &g, (0.0, 0.0, 1.0), &lsqr, &params, |i, total| {
            assert!(i <= total);
            last = i;
        });
        assert_eq!(chi.len(), n_total);
        assert!(chi.iter().all(|v| v.is_finite()));
        assert!(chi.iter().zip(&mask).all(|(v, &m)| m != 0 || *v == 0.0), "not masked");
        assert!(chi.iter().any(|v| v.abs() > 0.0), "empty map");
        assert!(last > 10, "progress did not cover the HEIDI stage");
        // empty mask
        let z = heidi(&field, &vec![0u8; n_total], None, &g, (0.0, 0.0, 1.0), &lsqr, &params, |_, _| {});
        assert!(z.iter().all(|v| *v == 0.0));
    }

    /// The cone fill runs on χ in upstream's phase units (ppm · γ·TE·B0·1e-6): NESTA's final
    /// μ is absolute, so the scale changes the answer.
    #[test]
    fn cone_fill_runs_in_phase_units() {
        let n = 11;
        let g = test_grid(n);
        let nt = g.n_total();
        let seed: Vec<f64> = (0..nt).map(|i| ((i as f64) * 0.13).sin() * 0.05).collect();
        let field: Vec<f64> = (0..nt).map(|i| ((i as f64) * 0.07).cos() * 0.01).collect();
        let mask = vec![1u8; nt];
        let p = HeidiParams { denoise: None, continuation_steps: 3, inner_iterations: 30, echo_times: vec![0.005, 0.025], b0: 7.0, ..Default::default() };
        let chi = heidi_on_grid(&field, &mask, &seed, &g, (0.0, 0.0, 1.0), &p, |_, _| {});
        let s = GAMMA * 0.025 * p.b0 * 1e-6; // the last echo time, HEIDI's default
        let masks = homogeneity_masks(&field, &mask, &g, &p);
        let cone = cone_mask(&g, (0.0, 0.0, 1.0), p.cone_threshold);
        let scaled: Vec<f64> = seed.iter().map(|v| v * s).collect();
        let manual = nesta_cone_fill(&scaled, &cone, &masks, &g, &p, |_, _| {});
        for (a, b) in chi.iter().zip(&manual) {
            assert!((a - b / s).abs() < 1e-12, "{a} vs {}", b / s);
        }
        let unscaled = nesta_cone_fill(&seed, &cone, &masks, &g, &p, |_, _| {});
        assert!(chi.iter().zip(&unscaled).any(|(a, b)| (a - b).abs() > 1e-9), "scale had no effect");
    }

    /// The cone fill on an outside seed runs on upstream's grid like `heidi`'s, returns the
    /// input grid masked, and agrees with `heidi_on_grid` on that grid (the seed cropped and
    /// framed as `Wrapper_QSM_HEIDI4all` frames it).
    #[test]
    fn cone_fill_takes_any_seed_on_upstreams_grid() {
        let (nx, ny, nz) = (12usize, 11usize, 13usize);
        let g = Grid::new(nx, ny, nz, 1.0, 1.0, 1.0);
        let n_total = g.n_total();
        let field: Vec<f64> = (0..n_total).map(|i| ((i as f64) * 0.31).sin() * 0.05).collect();
        let mut mask = vec![0u8; n_total];
        for k in 3..10 {
            for j in 2..9 {
                for i in 3..9 {
                    mask[i + nx * (j + ny * k)] = 1;
                }
            }
        }
        let seed: Vec<f64> = field.iter().zip(&mask).map(|(f, &m)| if m != 0 { 0.4 * f } else { 0.0 }).collect();
        let params = HeidiParams { denoise: None, continuation_steps: 2, inner_iterations: 10, ..Default::default() };
        let mut last = 0;
        let chi = heidi_cone_fill(&field, &mask, &seed, &g, (0.0, 0.0, 1.0), &params, |i, total| {
            assert!(i <= total);
            assert_eq!(total, cone_fill_iterations(&params));
            last = i;
        });
        assert_eq!(chi.len(), n_total);
        assert!(chi.iter().all(|v| v.is_finite()));
        assert!(chi.iter().zip(&mask).all(|(v, &m)| m != 0 || *v == 0.0), "not masked");
        assert!(chi.iter().any(|v| v.abs() > 0.0), "empty map");
        assert_eq!(last, cone_fill_iterations(&params));

        let crop = crop_grid(&mask, g.dims).unwrap();
        let pgrid = Grid::new(crop[0].2, crop[1].2, crop[2].2, 1.0, 1.0, 1.0);
        let on_grid = heidi_on_grid(
            &crop_pad(&field, g.dims, &crop),
            &crop_pad(&mask, g.dims, &crop),
            &crop_pad(&seed, g.dims, &crop),
            &pgrid,
            (0.0, 0.0, 1.0),
            &params,
            |_, _| {},
        );
        let mut expected = uncrop(&on_grid, g.dims, &crop);
        apply_mask_zero(&mut expected, &mask);
        assert_eq!(chi, expected);

        let z = heidi_cone_fill(&field, &vec![0u8; n_total], &seed, &g, (0.0, 0.0, 1.0), &params, |_, _| {});
        assert!(z.iter().all(|v| *v == 0.0));
    }

    #[test]
    fn zero_init_gives_zero_output() {
        let n = 9;
        let g = test_grid(n);
        let n_total = g.n_total();
        let params = HeidiParams {
            denoise: None,
            continuation_steps: 2,
            inner_iterations: 5,
            ..Default::default()
        };
        let chi = heidi_on_grid(
            &vec![0.0; n_total],
            &vec![1u8; n_total],
            &vec![0.0; n_total],
            &g,
            (0.0, 0.0, 1.0),
            &params,
            |_, _| {},
        );
        assert!(chi.iter().all(|v| v.abs() < 1e-12));
    }

    /// The cone is where the TV prior actually does work: HEIDI should reduce
    /// total variation relative to just low-passing the initial map.
    #[test]
    fn tv_is_reduced_against_the_bare_cone_projection() {
        let n = 17;
        let g = test_grid(n);
        let n_total = g.n_total();
        let bdir = (0.0, 0.0, 1.0);

        let mut chi_init = vec![0.0; n_total];
        let c = n as f64 / 2.0;
        for k in 0..n {
            for j in 0..n {
                for i in 0..n {
                    let r = ((i as f64 - c).powi(2)
                        + (j as f64 - c).powi(2)
                        + (k as f64 - c).powi(2))
                    .sqrt();
                    chi_init[i + j * n + k * n * n] = if r < 5.0 { 0.1 } else { 0.0 };
                }
            }
        }

        let field = vec![0.0; n_total];
        let mask = vec![1u8; n_total];
        let params = HeidiParams {
            denoise: None,
            continuation_steps: 4,
            inner_iterations: 60,
            ..Default::default()
        };
        let chi = heidi_on_grid(&field, &mask, &chi_init, &g, bdir, &params, |_, _| {});

        // Bare projection = what you get with no prior at all.
        let cone = cone_mask(&g, bdir, params.cone_threshold);
        let mut ws = Fft3dWorkspace::new(n, n, n);
        let mut buf: Vec<Complex64> =
            chi_init.iter().map(|&v| Complex64::new(v, 0.0)).collect();
        ws.fft3d(&mut buf);
        for (b, &keep) in buf.iter_mut().zip(&cone) {
            if !keep {
                *b = Complex64::new(0.0, 0.0);
            }
        }
        ws.ifft3d(&mut buf);
        let projected: Vec<f64> = buf.iter().map(|c| c.re).collect();

        let tv = |x: &[f64]| -> f64 {
            let dx = discrete_gradient(x, &g, 0, false);
            let dy = discrete_gradient(x, &g, 1, false);
            let dz = discrete_gradient(x, &g, 2, false);
            (0..n_total)
                .map(|i| (dx[i] * dx[i] + dy[i] * dy[i] + dz[i] * dz[i]).sqrt())
                .sum()
        };

        let tv_heidi = tv(&chi);
        let tv_proj = tv(&projected);
        assert!(
            tv_heidi < tv_proj,
            "HEIDI TV {tv_heidi} is not below the bare projection's {tv_proj}"
        );
    }
}