qsm_core/inversion/

tfi.rs

//! TFI (Preconditioned Total Field Inversion)
//!
//! Single-step QSM: inverts the TOTAL field (before background removal)
//! directly to susceptibility over the whole field-of-view, jointly handling
//! background removal and dipole inversion.
//!
//! It reuses MEDI's nonlinear Gauss-Newton + CG + IRLS(L1) machinery, with
//! three differences from MEDI (local-field inversion):
//!
//! 1. Input is the **total field** (not background-removed), solved over the
//!    **whole FOV** (not just the brain mask).
//! 2. A **preconditioner** change-of-variables χ = P⊙y is used, where P[i] = 1
//!    inside the brain mask and P[i] = `precond` (default 30) outside. Solving
//!    for y is better conditioned because background susceptibility (e.g. air
//!    ≈ 9 ppm) is far larger than tissue. With χ = P⊙y the Gauss-Newton normal
//!    equation operator on δy is `A_tfi(δy) = P ⊙ A_medi(P ⊙ δy)` and the RHS
//!    is `b_tfi = P ⊙ b_medi(χ = P⊙y)`.
//! 3. Regularization (L1 morphology) is applied over the whole FOV: the
//!    magnitude edge mask is used inside the brain and set to 1 (regularize)
//!    OUTSIDE the brain. The data-fidelity weight is SNR/brain-based (≈0 outside
//!    the brain), so outside-brain χ is constrained by regularization + the
//!    preconditioner.
//!
//! Reference:
//! Liu, Z., Kee, Y., Zhou, D., Wang, Y., Spincemaille, P. (2017).
//! "Preconditioned total field inversion (TFI) algorithm for quantitative
//! susceptibility mapping." Magnetic Resonance in Medicine, 78(1):303-315.
//! https://doi.org/10.1002/mrm.26946

use num_complex::Complex32;
use crate::kernels::dipole::dipole_kernel_f32;
use crate::utils::simd_ops::{
    dot_product_f32, norm_squared_f32, axpy_f32, xpby_f32,
    compute_p_weights_f32, negate_f32,
};
use crate::Grid;

use super::medi::{
    MediWorkspace, MediOpBuffers,
    apply_medi_operator_core, apply_dipole_conv, compute_rhs_inplace,
    dataterm_mask_f32, gradient_mask_f32, fgrad_periodic_inplace_f32,
};

/// TFI algorithm parameters
#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
#[derive(Clone, Debug)]
pub struct TfiParams {
    /// Regularization weight
    pub lambda: f64,
    /// Preconditioner value outside the brain mask (1.0 inside)
    pub precond: f64,
    /// Enable MERIT (outlier adjustment) — currently a no-op placeholder
    pub merit: bool,
    /// Data weighting mode (1 = SNR)
    pub data_weighting: i32,
    /// Fraction of voxels considered edges (0.0-1.0)
    pub percentage: f64,
    /// CG convergence tolerance
    pub cg_tol: f64,
    /// Maximum CG iterations
    pub cg_max_iter: usize,
    /// Maximum outer (Gauss-Newton) iterations
    pub max_iter: usize,
    /// Outer convergence tolerance
    pub tol: f64,
}

impl Default for TfiParams {
    fn default() -> Self {
        Self {
            // TFI shares MEDI's L1 machinery and λ convention, so we use MEDI's default λ
            // (7.5e-5) rather than a benchmark-fitted value. precond=30 is the standard TFI
            // preconditioner (Liu 2017). Neither is tuned to any specific dataset.
            lambda: 7.5e-5,
            precond: 30.0,
            merit: false,
            data_weighting: 1,
            percentage: 0.9,
            cg_tol: 0.01,
            cg_max_iter: 100,
            max_iter: 10,
            tol: 0.1,
        }
    }
}

/// Conjugate gradient solver for the preconditioned TFI operator.
///
/// Solves `A_tfi(y) = b` where `A_tfi(dy) = P ⊙ A_medi(P ⊙ dy)`.
/// The MEDI operator itself is reused unchanged.
#[allow(clippy::too_many_arguments)]
fn cg_solve_tfi<F>(
    ws: &mut MediWorkspace,
    precond: &[f32],
    w: &[Complex32],
    d_kernel: &[f32],
    mx: &[f32],
    my: &[f32],
    mz: &[f32],
    vr: &[f32],
    lambda: f32,
    b: &[f32],
    x: &mut [f32],
    tol: f32,
    max_iter: usize,
    mut progress_callback: F,
) where
    F: FnMut(usize, usize),
{
    let n = ws.n_total;
    let (nx, ny, nz) = (ws.nx, ws.ny, ws.nz);
    let (vsx, vsy, vsz) = (ws.vsx, ws.vsy, ws.vsz);

    x.fill(0.0);
    ws.cg_r.copy_from_slice(b);
    ws.cg_p.copy_from_slice(&ws.cg_r);

    let mut rsold: f32 = norm_squared_f32(&ws.cg_r);
    let b_norm: f32 = norm_squared_f32(b).sqrt();
    if b_norm < 1e-10 {
        return;
    }

    // Scratch buffers: p_scaled = P ⊙ p; ap holds A_medi(p_scaled) then P ⊙ ap
    let mut p_scaled = vec![0.0f32; n];

    for cg_iter in 0..max_iter {
        progress_callback(cg_iter + 1, max_iter);

        // p_scaled = P ⊙ p
        for i in 0..n {
            p_scaled[i] = precond[i] * ws.cg_p[i];
        }

        // ap = A_medi(p_scaled)
        {
            let mut bufs = MediOpBuffers {
                gx: &mut ws.gx,
                gy: &mut ws.gy,
                gz: &mut ws.gz,
                reg_x: &mut ws.reg_x,
                reg_y: &mut ws.reg_y,
                reg_z: &mut ws.reg_z,
                div_buf: &mut ws.div_buf,
                dipole_buf: &mut ws.dipole_buf,
                complex_buf: &mut ws.complex_buf,
                complex_buf2: &mut ws.complex_buf2,
            };
            apply_medi_operator_core(
                &mut ws.fft_ws, &mut bufs, n, nx, ny, nz, vsx, vsy, vsz,
                &p_scaled, w, d_kernel, mx, my, mz, vr, lambda, &mut ws.cg_ap,
            );
        }

        // ap = P ⊙ ap
        for i in 0..n {
            ws.cg_ap[i] *= precond[i];
        }

        let pap: f32 = dot_product_f32(&ws.cg_p, &ws.cg_ap);
        if pap.abs() < 1e-15 {
            break;
        }

        let alpha = rsold / pap;
        axpy_f32(x, alpha, &ws.cg_p);
        axpy_f32(&mut ws.cg_r, -alpha, &ws.cg_ap);

        let rsnew: f32 = norm_squared_f32(&ws.cg_r);
        let residual = rsnew.sqrt();
        if residual < tol * b_norm {
            break;
        }

        let beta = rsnew / rsold;
        xpby_f32(&mut ws.cg_p, &ws.cg_r, beta);
        rsold = rsnew;
    }
}

/// Preconditioned Total Field Inversion (TFI).
///
/// # Arguments
/// * `total_field` - Total field (before background removal) in **ppm** — same units convention as
///   NDI and the other inversions (NOT MEDI's radians). Do not pre-scale to radians: the total field
///   is large and would wrap in the `exp(i·field)` data term.
/// * `n_std` - Noise standard deviation map (same size as total_field)
/// * `magnitude` - Magnitude image for gradient weighting
/// * `mask` - Binary brain mask (1 = brain)
/// * `grid` - Volume grid (dimensions and voxel sizes)
/// * `bdir` - B0 field direction
/// * `params` - TFI parameters
/// * `progress` - Progress callback `(current_step, total_steps)`
///
/// # Returns
/// Susceptibility map χ (same units as input field), zeroed outside the brain mask.
/// The solve itself runs over the whole FOV (to absorb the background into out-of-brain
/// susceptibility), but that region is unconstrained/artefacty, so only the brain is returned.
#[allow(clippy::too_many_arguments)]
pub fn tfi(
    total_field: &[f64],
    n_std: &[f64],
    magnitude: &[f64],
    mask: &[u8],
    grid: &Grid,
    bdir: (f64, f64, f64),
    params: &TfiParams,
    mut progress: impl FnMut(usize, usize),
) -> Vec<f64> {
    let (nx, ny, nz) = grid.dims;
    let n_total = grid.n_total();

    let vsx_f32 = grid.vsx() as f32;
    let vsy_f32 = grid.vsy() as f32;
    let vsz_f32 = grid.vsz() as f32;
    let lambda_f32 = params.lambda as f32;
    let bdir_f32 = (bdir.0 as f32, bdir.1 as f32, bdir.2 as f32);
    let percentage_f32 = params.percentage as f32;
    let cg_tol_f32 = params.cg_tol as f32;
    let tol_f32 = params.tol as f32;
    let max_iter = params.max_iter;
    let cg_max_iter = params.cg_max_iter;
    let data_weighting = params.data_weighting;

    let field_f32: Vec<f32> = total_field.iter().map(|&v| v as f32).collect();
    let n_std_f32: Vec<f32> = n_std.iter().map(|&v| v as f32).collect();
    let magnitude_f32: Vec<f32> = magnitude.iter().map(|&v| v as f32).collect();

    let mut ws = MediWorkspace::new(grid);

    // Preconditioner: P[i] = 1 inside brain, precond outside.
    let precond_out = params.precond as f32;
    let precond: Vec<f32> = mask.iter()
        .map(|&m| if m != 0 { 1.0 } else { precond_out })
        .collect();

    // Dipole kernel (no SMV — TFI operates on the total field over whole FOV).
    let d_kernel = dipole_kernel_f32(grid, bdir_f32);

    // Data weighting: SNR inside brain, ~0 outside (dataterm_mask behavior).
    // n_std zeroed outside brain so m = 0 there.
    let mut tempn: Vec<f32> = n_std_f32.clone();
    for i in 0..n_total {
        if mask[i] == 0 {
            tempn[i] = 0.0;
        }
    }
    let m = dataterm_mask_f32(data_weighting, &tempn, mask);

    // b0 = m * exp(i * total_field)  (only nonzero inside the brain via m)
    let b0: Vec<Complex32> = field_f32.iter()
        .zip(m.iter())
        .map(|(&f, &mi)| {
            let phase = Complex32::new(0.0, f);
            mi * phase.exp()
        })
        .collect();

    // Gradient (morphology) mask: magnitude edges inside brain, then 1 OUTSIDE
    // the brain so the whole FOV is regularized (smoothed).
    let (mut w_gx, mut w_gy, mut w_gz) = gradient_mask_f32(
        &magnitude_f32, mask, nx, ny, nz, vsx_f32, vsy_f32, vsz_f32, percentage_f32,
    );
    // Fallback: if a mask is all zeros, use the magnitude image (matching MEDI).
    if !w_gx.iter().any(|&v| v != 0.0) { w_gx = magnitude_f32.clone(); }
    if !w_gy.iter().any(|&v| v != 0.0) { w_gy = magnitude_f32.clone(); }
    if !w_gz.iter().any(|&v| v != 0.0) { w_gz = magnitude_f32.clone(); }
    // Set to 1 OUTSIDE the brain: regularize/smooth the whole FOV.
    for i in 0..n_total {
        if mask[i] == 0 {
            w_gx[i] = 1.0;
            w_gy[i] = 1.0;
            w_gz[i] = 1.0;
        }
    }

    // State in the preconditioned variable y, where χ = P ⊙ y.
    let mut y = vec![0.0f32; n_total];
    let mut chi = vec![0.0f32; n_total];
    let mut dy = vec![0.0f32; n_total];
    let mut rhs = vec![0.0f32; n_total];
    let mut vr = vec![0.0f32; n_total];
    let mut w: Vec<Complex32> = vec![Complex32::new(0.0, 0.0); n_total];
    let mut y_prev = vec![0.0f32; n_total];

    let beta = 1.49e-8_f32;
    let total_steps = max_iter * cg_max_iter;

    for iter in 0..max_iter {
        y_prev.copy_from_slice(&y);

        // χ = P ⊙ y
        for i in 0..n_total {
            chi[i] = precond[i] * y[i];
        }

        // P weights (IRLS): P_irls = 1 / sqrt(|m_grad · grad(χ)|^2 + beta)
        fgrad_periodic_inplace_f32(
            &mut ws.gx, &mut ws.gy, &mut ws.gz,
            &chi, nx, ny, nz, vsx_f32, vsy_f32, vsz_f32,
        );
        compute_p_weights_f32(&mut vr, &w_gx, &w_gy, &w_gz, &ws.gx, &ws.gy, &ws.gz, beta);

        // w = m * exp(i * D*χ)
        apply_dipole_conv(&mut ws.fft_ws, &chi, &d_kernel, &mut ws.dipole_buf, &mut ws.complex_buf);
        for i in 0..n_total {
            let phase = Complex32::new(0.0, ws.dipole_buf[i]);
            w[i] = m[i] * phase.exp();
        }

        // b_medi(χ = P⊙y)
        compute_rhs_inplace(&chi, &w, &b0, &d_kernel, &w_gx, &w_gy, &w_gz, &vr, lambda_f32, &mut rhs, &mut ws);

        // b_tfi = P ⊙ b_medi
        for i in 0..n_total {
            rhs[i] *= precond[i];
        }

        // Solve A_tfi δy = -b_tfi
        negate_f32(&mut rhs);

        let gn_iter = iter;
        cg_solve_tfi(
            &mut ws, &precond, &w, &d_kernel, &w_gx, &w_gy, &w_gz, &vr, lambda_f32,
            &rhs, &mut dy, cg_tol_f32, cg_max_iter,
            |cg_iter, cg_total| {
                let current = gn_iter * cg_total + cg_iter;
                progress(current, total_steps);
            },
        );

        // y = y + dy
        axpy_f32(&mut y, 1.0, &dy);

        // Convergence check in y
        let norm_dy = norm_squared_f32(&dy).sqrt();
        let norm_y = norm_squared_f32(&y_prev).sqrt();
        let rel_change = norm_dy / (norm_y + 1e-6);
        if rel_change < tol_f32 {
            progress(total_steps, total_steps);
            break;
        }
    }

    // Final χ = P ⊙ y, zeroed outside the brain mask. The whole-FOV solve is only needed to
    // absorb the background field into out-of-brain susceptibility; that region is unconstrained
    // (artefacty), so the returned map keeps only the brain — matching MEDI/NDI/etc.
    let _ = params.merit;
    y.iter()
        .zip(precond.iter())
        .zip(mask.iter())
        .map(|((&yi, &pi), &m)| if m == 0 { 0.0 } else { (pi * yi) as f64 })
        .collect()
}

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

    fn test_tfi_params() -> TfiParams {
        TfiParams {
            lambda: 1e-3, percentage: 0.9, cg_tol: 0.1,
            cg_max_iter: 10, max_iter: 3, tol: 0.1,
            ..TfiParams::default()
        }
    }

    #[test]
    fn test_tfi_zero_field() {
        let n = 8;
        let field = vec![0.0; n * n * n];
        let mask = vec![1u8; n * n * n];
        let mag = vec![1.0; n * n * n];
        let n_std = vec![1.0; n * n * n];
        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);

        let chi = tfi(
            &field, &n_std, &mag, &mask, &grid,
            (0.0, 0.0, 1.0), &test_tfi_params(), |_, _| {},
        );

        for &val in chi.iter() {
            assert!(val.abs() < 1e-4, "Zero field should give near-zero chi, got {}", val);
        }
    }

    #[test]
    fn test_tfi_finite() {
        let n = 8;
        let field: Vec<f64> = (0..n*n*n).map(|i| (i as f64) * 0.001).collect();
        let mask = vec![1u8; n * n * n];
        let mag = vec![1.0; n * n * n];
        let n_std = vec![1.0; n * n * n];
        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);

        let chi = tfi(
            &field, &n_std, &mag, &mask, &grid,
            (0.0, 0.0, 1.0), &test_tfi_params(), |_, _| {},
        );

        for (i, &val) in chi.iter().enumerate() {
            assert!(val.is_finite(), "Chi should be finite at index {}", i);
        }
    }

    #[test]
    fn test_tfi_finite_with_background() {
        // Partial mask: exercise the preconditioner + whole-FOV regularization.
        let n = 8;
        let field: Vec<f64> = (0..n*n*n).map(|i| (i as f64) * 0.001).collect();
        let mut mask = vec![0u8; n * n * n];
        // Central 4x4x4 region is "brain".
        for z in 2..6 { for y in 2..6 { for x in 2..6 {
            mask[x + y*n + z*n*n] = 1;
        }}}
        let mag = vec![1.0; n * n * n];
        let n_std = vec![1.0; n * n * n];
        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);

        let chi = tfi(
            &field, &n_std, &mag, &mask, &grid,
            (0.0, 0.0, 1.0), &test_tfi_params(), |_, _| {},
        );

        assert_eq!(chi.len(), n * n * n);
        for (i, &val) in chi.iter().enumerate() {
            assert!(val.is_finite(), "Chi should be finite at index {} (whole FOV)", i);
        }
    }

    /// Numerical validation on the real dev phantom.
    /// Run with: cargo test --release test_tfi_phantom -- --ignored --nocapture
    #[test]
    #[ignore]
    fn test_tfi_phantom() {

        let base = "/home/ashley/repos/qsm/qsmci/qsmci/data/sim/dev";
        let field_path = format!("{}/groundtruth/totalfield.nii.gz", base);
        let mask_path = format!("{}/inputs/mask.nii.gz", base);
        let mag_path = format!("{}/inputs/magnitude.nii.gz", base);
        let chi_path = format!("{}/groundtruth/chimap.nii.gz", base);

        if !std::path::Path::new(&field_path).exists() {
            eprintln!("Skipping: {} not found", field_path);
            return;
        }

        // Load total field (ppm)
        let field_nii = crate::io::load_nifti(&std::fs::read(&field_path).unwrap()).unwrap();
        let (nx, ny, nz) = field_nii.dims;
        let (vsx, vsy, vsz) = field_nii.voxel_size;
        let n_total = nx * ny * nz;
        eprintln!("Dims: {}x{}x{}  voxel: {}x{}x{}", nx, ny, nz, vsx, vsy, vsz);

        let field_ppm = field_nii.data;

        // Mask
        let mask_nii = crate::io::load_nifti(&std::fs::read(&mask_path).unwrap()).unwrap();
        let mask: Vec<u8> = mask_nii.data.iter().map(|&v| if v > 0.5 { 1 } else { 0 }).collect();
        let mask_count = mask.iter().filter(|&&m| m != 0).count();
        eprintln!("Mask voxels: {} / {}", mask_count, n_total);

        // Magnitude (4D) — combine echoes as sqrt(sum of squares)
        let (mag4d, mag_dims, _, _) = crate::io::load_nifti_4d(&std::fs::read(&mag_path).unwrap()).unwrap();
        let (mnx, mny, mnz, nt) = mag_dims;
        assert_eq!(mnx * mny * mnz, n_total, "magnitude spatial dims mismatch");
        let mut magnitude = vec![0.0f64; n_total];
        for i in 0..n_total {
            let mut ss = 0.0;
            for t in 0..nt {
                let v = mag4d[i + t * n_total];
                ss += v * v;
            }
            magnitude[i] = ss.sqrt();
        }

        // Ground truth chi (ppm)
        let chi_gt_nii = crate::io::load_nifti(&std::fs::read(&chi_path).unwrap()).unwrap();
        let chi_gt_ppm = chi_gt_nii.data;

        // params.json
        let bdir = (0.0, 0.0, 1.0);

        // TFI takes the total field in ppm (no rad conversion — see fn docs).
        let n_std = vec![1.0f64; n_total];
        let grid = Grid::new(nx, ny, nz, vsx, vsy, vsz);

        // On this dev phantom the background outside the brain is tiny (±0.01 ppm), so a small
        // preconditioner (2-3) is best here; the default of 30 targets in-vivo air (~9 ppm).
        let lambda = std::env::var("TFI_LAMBDA").ok()
            .and_then(|s| s.parse::<f64>().ok()).unwrap_or(1e-4);
        let precond = std::env::var("TFI_PRECOND").ok()
            .and_then(|s| s.parse::<f64>().ok()).unwrap_or(3.0);
        let params = TfiParams {
            lambda,
            precond,
            percentage: 0.9,
            cg_tol: 0.01,
            cg_max_iter: 100,
            max_iter: 20,
            tol: 0.005,
            ..TfiParams::default()
        };
        eprintln!("params: lambda={} precond={}", lambda, precond);

        eprintln!("Running TFI...");
        let chi_ppm = tfi(
            &field_ppm, &n_std, &magnitude, &mask, &grid, bdir, &params,
            |c, t| { if c % 200 == 0 || c == t { eprintln!("  progress {}/{}", c, t); } },
        );

        // Metrics within the mask
        let idx: Vec<usize> = (0..n_total).filter(|&i| mask[i] != 0).collect();
        let a: Vec<f64> = idx.iter().map(|&i| chi_ppm[i]).collect();
        let b: Vec<f64> = idx.iter().map(|&i| chi_gt_ppm[i]).collect();
        let nm = a.len() as f64;

        let mean_a = a.iter().sum::<f64>() / nm;
        let mean_b = b.iter().sum::<f64>() / nm;

        // Pearson correlation
        let mut cov = 0.0; let mut va = 0.0; let mut vb = 0.0;
        for k in 0..a.len() {
            let da = a[k] - mean_a; let db = b[k] - mean_b;
            cov += da * db; va += da * da; vb += db * db;
        }
        let corr = cov / (va.sqrt() * vb.sqrt());

        // NRMSE = 100 * ||a - b|| / ||b||
        let mut num = 0.0; let mut den = 0.0;
        for k in 0..a.len() {
            num += (a[k] - b[k]).powi(2);
            den += b[k].powi(2);
        }
        let nrmse = 100.0 * (num / den).sqrt();

        // Detrended NRMSE (subtract per-map mean)
        let mut num_d = 0.0; let mut den_d = 0.0;
        for k in 0..a.len() {
            num_d += ((a[k] - mean_a) - (b[k] - mean_b)).powi(2);
            den_d += (b[k] - mean_b).powi(2);
        }
        let nrmse_d = 100.0 * (num_d / den_d).sqrt();

        eprintln!("=== TFI dev phantom results (within mask) ===");
        eprintln!("  Pearson correlation : {:.5}", corr);
        eprintln!("  NRMSE               : {:.3}", nrmse);
        eprintln!("  detrended NRMSE     : {:.3}", nrmse_d);
        eprintln!("  chi_out mean={:.4} min={:.4} max={:.4}",
            mean_a,
            a.iter().cloned().fold(f64::MAX, f64::min),
            a.iter().cloned().fold(f64::MIN, f64::max));
        eprintln!("  chi_gt  mean={:.4} min={:.4} max={:.4}",
            mean_b,
            b.iter().cloned().fold(f64::MAX, f64::min),
            b.iter().cloned().fold(f64::MIN, f64::max));

        assert!(corr > 0.95, "TFI correlation {:.5} below acceptance 0.95", corr);
    }

    /// Background-removal stress test on the REAL CI phantom (~/bids), which has a
    /// genuine in-brain background field (unlike the qsm-ci dev phantom). Sweeps
    /// `precond`/`lambda` on the TOTAL field and reports corr vs GT χ, plus a
    /// local-field run as the ceiling. Fields are fed raw (ppm), matching how the
    /// CI harness scores MEDI/NDI/etc.
    /// Run: cargo test --release --lib test_tfi_background -- --ignored --nocapture
    #[test]
    #[ignore]
    fn test_tfi_background() {
        let base = "/home/ashley/bids/derivatives/qsm-forward/sub-1/anat";
        let total_path = format!("{}/sub-1_fieldmap.nii", base);
        if !std::path::Path::new(&total_path).exists() {
            eprintln!("Skipping: {} not found", total_path);
            return;
        }
        let load = |p: &str| crate::io::load_nifti(&std::fs::read(p).unwrap()).unwrap();
        let total_nii = load(&total_path);
        let (nx, ny, nz) = total_nii.dims;
        let (vsx, vsy, vsz) = total_nii.voxel_size;
        let n = nx * ny * nz;
        let total_field = total_nii.data;                                   // ppm
        let local_field = load(&format!("{}/sub-1_fieldmap-local.nii", base)).data;
        let chi_gt = load(&format!("{}/sub-1_Chimap.nii", base)).data;
        let mask: Vec<u8> = load(&format!("{}/sub-1_mask.nii", base)).data
            .iter().map(|&v| if v > 0.5 { 1 } else { 0 }).collect();

        let grid = Grid::new(nx, ny, nz, vsx, vsy, vsz);
        let bdir = (0.0, 0.0, 1.0);
        let n_std = vec![1.0f64; n];
        let magnitude = vec![1.0f64; n];

        let std_in = |v: &[f64]| {
            let m: Vec<f64> = (0..n).filter(|&i| mask[i] != 0).map(|i| v[i]).collect();
            let mean = m.iter().sum::<f64>() / m.len() as f64;
            (m.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / m.len() as f64).sqrt()
        };
        eprintln!("dims {}x{}x{}  in-brain field std: total={:.5} local={:.5}",
            nx, ny, nz, std_in(&total_field), std_in(&local_field));

        let corr_in = |a: &[f64], b: &[f64]| {
            let idx: Vec<usize> = (0..n).filter(|&i| mask[i] != 0).collect();
            let (mut sa, mut sb) = (0.0, 0.0);
            for &i in &idx { sa += a[i]; sb += b[i]; }
            let (ma, mb) = (sa / idx.len() as f64, sb / idx.len() as f64);
            let (mut num, mut da, mut db) = (0.0, 0.0, 0.0);
            for &i in &idx {
                num += (a[i] - ma) * (b[i] - mb);
                da += (a[i] - ma).powi(2); db += (b[i] - mb).powi(2);
            }
            num / (da.sqrt() * db.sqrt())
        };

        // Ceiling: TFI on the LOCAL field (no background to remove, precond irrelevant).
        let p_local = TfiParams { lambda: 1e-4, precond: 1.0, ..TfiParams::default() };
        let chi_local = tfi(&local_field, &n_std, &magnitude, &mask, &grid, bdir, &p_local, |_, _| {});
        eprintln!("TFI on LOCAL field (ceiling): corr={:.4}", corr_in(&chi_local, &chi_gt));

        eprintln!("=== TFI on TOTAL field — precond × lambda sweep (corr vs GT χ, within mask) ===");
        for &lambda in &[5e-5f64, 7.5e-5, 1e-4] {
            for &precond in &[30.0f64] {
                let params = TfiParams { lambda, precond, ..TfiParams::default() };
                let chi = tfi(&total_field, &n_std, &magnitude, &mask, &grid, bdir, &params, |_, _| {});
                eprintln!("  lambda={:>6.0e} precond={:>5}  corr={:.4}", lambda, precond, corr_in(&chi, &chi_gt));
            }
        }
    }
}