qsm_core/bgremove/

pdf.rs

//! Projection onto Dipole Fields (PDF) background field removal
//!
//! Projects the field onto dipole fields generated by sources outside
//! the brain mask, separating background and local fields.
//!
//! Reference:
//! Liu, T., Khalidov, I., de Rochefort, L., Spincemaille, P., Liu, J., Tsiouris, A.J.,
//! Wang, Y. (2011). "A novel background field removal method for MRI using projection
//! onto dipole fields." NMR in Biomedicine, 24(9):1129-1136. https://doi.org/10.1002/nbm.1670
//!
//! Reference implementation: https://github.com/kamesy/QSM.jl

use num_complex::Complex64;
use crate::Grid;
use crate::fft::{fft3d, ifft3d};
use crate::kernels::dipole::dipole_kernel;
use crate::utils::vec_norm;

/// PDF algorithm parameters
#[cfg_attr(feature = "introspection", derive(serde::Serialize))]
#[derive(Clone, Debug)]
pub struct PdfParams {
    /// Convergence tolerance
    pub tol: f64,
    /// Maximum LSMR iterations. `None` uses an automatic default of
    /// `sqrt(n_voxels)`.
    pub max_iter: Option<usize>,
}

impl Default for PdfParams {
    fn default() -> Self {
        Self { tol: 1e-5, max_iter: None }
    }
}

/// PDF background field removal
///
/// # Arguments
/// * `field` - Total field (nx * ny * nz)
/// * `mask` - Binary mask (nx * ny * nz), 1 = brain, 0 = background
/// * `grid` - Volume dimensions and voxel sizes
/// * `bdir` - B0 field direction
/// * `params` - PDF parameters (tolerance, optional max iterations)
/// * `progress` - Progress callback (iteration, max_iter)
///
/// # Returns
/// Local field with background removed
pub fn pdf(
    field: &[f64],
    mask: &[u8],
    grid: &Grid,
    bdir: (f64, f64, f64),
    params: &PdfParams,
    progress: impl FnMut(usize, usize),
) -> Vec<f64> {
    // Default max iterations to sqrt(n_voxels) when unspecified.
    let max_iter = params.max_iter.unwrap_or_else(|| (grid.n_total() as f64).sqrt() as usize);
    pdf_core(field, mask, grid, bdir, params.tol, max_iter, progress)
}

/// PDF with an explicit LSMR iteration count.
///
/// Internal entry point; the public [`pdf`] wrapper supplies the default
/// iteration count from [`PdfParams`].
pub(crate) fn pdf_core(
    field: &[f64],
    mask: &[u8],
    grid: &Grid,
    bdir: (f64, f64, f64),
    tol: f64,
    max_iter: usize,
    mut progress: impl FnMut(usize, usize),
) -> Vec<f64> {
    let (nx, ny, nz) = grid.dims;
    let n_total = nx * ny * nz;

    // Generate dipole kernel
    let d_kernel = dipole_kernel(grid, bdir);

    // Create background mask (complement of brain mask)
    let bg_mask: Vec<f64> = mask.iter()
        .map(|&m| if m == 0 { 1.0 } else { 0.0 })
        .collect();

    // Brain mask as f64
    let brain_mask: Vec<f64> = mask.iter()
        .map(|&m| if m != 0 { 1.0 } else { 0.0 })
        .collect();

    // RHS: b = W * f where W is brain mask weights
    let b: Vec<f64> = field.iter()
        .zip(brain_mask.iter())
        .map(|(&f, &w)| f * w)
        .collect();

    // Initialize solution
    let mut x = vec![0.0; n_total];

    let mut u = b.clone();
    let mut beta = vec_norm(&u);

    if beta < 1e-20 {
        return vec![0.0; n_total];
    }

    for i in 0..n_total {
        u[i] /= beta;
    }

    let mut v = apply_at(&u, &brain_mask, &d_kernel, &bg_mask, nx, ny, nz);
    let mut alpha = vec_norm(&v);

    if alpha < 1e-20 {
        return vec![0.0; n_total];
    }

    for i in 0..n_total {
        v[i] /= alpha;
    }

    // LSMR variables (following Fong & Saunders 2011 / QSM.jl implementation)
    let norm_b = beta;
    let mut w = v.clone();
    let mut phi_bar = beta;
    let mut rho_bar = alpha;

    // Variables for ||r|| estimation (matching Julia's lsmr.jl)
    let mut beta_dd = beta;
    let mut beta_d = 0.0;
    let mut rho_d_old = 1.0;
    let mut tau_tilde_old = 0.0;
    let mut theta_tilde = 0.0;
    let mut zeta = 0.0;
    let d_accum = 0.0; // beta_check^2 accumulator; always 0 when lambda=0

    // Variables for ||A|| estimation
    let mut norm_a2 = alpha * alpha;

    // Variables for the QR factorization (needed for ||r|| recurrence)
    let mut zeta_bar = alpha * beta;
    let mut alpha_bar = alpha;
    let mut c_bar = 1.0;
    let mut s_bar = 0.0;
    let mut rho_val = 1.0;
    let mut rho_bar_lsmr = 1.0;

    for iter in 0..max_iter {
        // Report progress
        progress(iter + 1, max_iter);

        // Bidiagonalization
        let av = apply_a(&v, &bg_mask, &d_kernel, &brain_mask, nx, ny, nz);
        for i in 0..n_total {
            u[i] = av[i] - alpha * u[i];
        }
        beta = vec_norm(&u);

        if beta < 1e-20 {
            progress(iter + 1, iter + 1);
            break;
        }

        for i in 0..n_total {
            u[i] /= beta;
        }

        let atu = apply_at(&u, &brain_mask, &d_kernel, &bg_mask, nx, ny, nz);
        for i in 0..n_total {
            v[i] = atu[i] - beta * v[i];
        }
        alpha = vec_norm(&v);

        if alpha < 1e-20 {
            progress(iter + 1, iter + 1);
            break;
        }

        for i in 0..n_total {
            v[i] /= alpha;
        }

        // Construct and apply rotation Q_i (for solution update)
        let rho = (rho_bar * rho_bar + beta * beta).sqrt();
        let c = rho_bar / rho;
        let s = beta / rho;
        let theta = s * alpha;
        rho_bar = -c * alpha;
        let phi = c * phi_bar;
        phi_bar *= s;

        // Update x and w
        let phi_rho = phi / rho;
        let theta_rho = theta / rho;

        for i in 0..n_total {
            x[i] += phi_rho * w[i];
            w[i] = v[i] - theta_rho * w[i];
        }

        // LSMR QR factorization (matching Julia lsmr.jl, with lambda=0)
        let _rho_old = rho_val;
        rho_val = (alpha_bar * alpha_bar + beta * beta).sqrt();
        let c_lsmr = alpha_bar / rho_val;
        let s_lsmr = beta / rho_val;
        let theta_new = s_lsmr * alpha;
        alpha_bar = c_lsmr * alpha;

        let _rho_bar_old = rho_bar_lsmr;
        let zeta_old = zeta;
        let theta_bar = s_bar * rho_val;
        let rho_tmp = c_bar * rho_val;
        rho_bar_lsmr = (rho_tmp * rho_tmp + theta_new * theta_new).sqrt();
        c_bar = rho_tmp / rho_bar_lsmr;
        s_bar = theta_new / rho_bar_lsmr;
        zeta = c_bar * zeta_bar;
        zeta_bar *= -s_bar;

        // Estimate ||r|| (matching Julia lsmr.jl lines 264-287, with lambda=0)
        let beta_hat = c_lsmr * beta_dd;
        beta_dd *= -s_lsmr;

        let theta_tilde_old = theta_tilde;
        let rho_tilde_old = (rho_d_old * rho_d_old + theta_bar * theta_bar).sqrt();
        let c_tilde_old = rho_d_old / rho_tilde_old;
        let s_tilde_old = theta_bar / rho_tilde_old;
        theta_tilde = s_tilde_old * rho_bar_lsmr;
        rho_d_old = c_tilde_old * rho_bar_lsmr;
        beta_d = -s_tilde_old * beta_d + c_tilde_old * beta_hat;

        tau_tilde_old = (zeta_old - theta_tilde_old * tau_tilde_old) / rho_tilde_old;
        let tau_d = (zeta - theta_tilde * tau_tilde_old) / rho_d_old;

        let norm_r = (d_accum + (beta_d - tau_d).powi(2) + beta_dd.powi(2)).sqrt();

        // Estimate ||A||
        norm_a2 += beta * beta;
        let norm_a = norm_a2.sqrt();
        norm_a2 += alpha * alpha;

        // ||A'r|| estimate
        let norm_ar = zeta_bar.abs();

        // Convergence tests (matching Julia lsmr.jl lines 308-318)
        let norm_x = vec_norm(&x);
        let test1 = norm_r / norm_b;
        let test2 = if norm_a * norm_r > 0.0 { norm_ar / (norm_a * norm_r) } else { 0.0 };
        let eps_r = tol + tol * norm_a * norm_x / norm_b;

        if test1 <= eps_r || test2 <= tol {
            progress(iter + 1, iter + 1);
            break;
        }
    }

    // Compute background field: b_field = D * (M_bg * x)
    let mut bg_source: Vec<Complex64> = x.iter()
        .zip(bg_mask.iter())
        .map(|(&xi, &m)| Complex64::new(xi * m, 0.0))
        .collect();

    fft3d(&mut bg_source, nx, ny, nz);

    for i in 0..n_total {
        bg_source[i] *= d_kernel[i];
    }

    ifft3d(&mut bg_source, nx, ny, nz);

    // Local field = total field - background field, masked
    let mut local_field = vec![0.0; n_total];
    for i in 0..n_total {
        if mask[i] != 0 {
            local_field[i] = field[i] - bg_source[i].re;
        }
    }

    local_field
}

/// Apply A = W * D * M_bg
fn apply_a(
    x: &[f64],
    bg_mask: &[f64],
    d_kernel: &[f64],
    brain_mask: &[f64],
    nx: usize, ny: usize, nz: usize,
) -> Vec<f64> {
    let n_total = nx * ny * nz;

    // Apply background mask
    let mut temp: Vec<Complex64> = x.iter()
        .zip(bg_mask.iter())
        .map(|(&xi, &m)| Complex64::new(xi * m, 0.0))
        .collect();

    // FFT
    fft3d(&mut temp, nx, ny, nz);

    // Apply dipole kernel
    for i in 0..n_total {
        temp[i] *= d_kernel[i];
    }

    // IFFT
    ifft3d(&mut temp, nx, ny, nz);

    // Apply brain mask weights
    temp.iter()
        .zip(brain_mask.iter())
        .map(|(t, &w)| t.re * w)
        .collect()
}

/// Apply A^T = M_bg * D * W
fn apply_at(
    u: &[f64],
    brain_mask: &[f64],
    d_kernel: &[f64],
    bg_mask: &[f64],
    nx: usize, ny: usize, nz: usize,
) -> Vec<f64> {
    let n_total = nx * ny * nz;

    // Apply brain mask weights
    let mut temp: Vec<Complex64> = u.iter()
        .zip(brain_mask.iter())
        .map(|(&ui, &w)| Complex64::new(ui * w, 0.0))
        .collect();

    // FFT
    fft3d(&mut temp, nx, ny, nz);

    // Apply dipole kernel (D is real and symmetric)
    for i in 0..n_total {
        temp[i] *= d_kernel[i];
    }

    // IFFT
    ifft3d(&mut temp, nx, ny, nz);

    // Apply background mask
    temp.iter()
        .zip(bg_mask.iter())
        .map(|(t, &m)| t.re * m)
        .collect()
}

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

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

        let local = pdf_core(
            &field, &mask, &grid,
            (0.0, 0.0, 1.0), 1e-5, 10, |_, _| {}
        );

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

    #[test]
    fn test_pdf_finite() {
        let n = 8;
        let field: Vec<f64> = (0..n*n*n).map(|i| (i as f64) * 0.001).collect();

        // Create a spherical mask in the center
        let mut mask = vec![0u8; n * n * n];
        let center = n / 2;
        let radius = n / 4;

        for i in 0..n {
            for j in 0..n {
                for k in 0..n {
                    let di = (i as i32) - (center as i32);
                    let dj = (j as i32) - (center as i32);
                    let dk = (k as i32) - (center as i32);
                    if di*di + dj*dj + dk*dk <= (radius * radius) as i32 {
                        mask[i * n * n + j * n + k] = 1;
                    }
                }
            }
        }

        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
        let local = pdf_core(
            &field, &mask, &grid,
            (0.0, 0.0, 1.0), 1e-5, 20, |_, _| {}
        );

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

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

        let local = pdf_core(
            &field, &mask, &grid,
            (0.0, 0.0, 1.0), 1e-5, 10, |_, _| {}
        );

        assert_eq!(local[0], 0.0, "Masked voxel should be zero");
        assert_eq!(local[10], 0.0, "Masked voxel should be zero");
    }

    #[test]
    fn test_pdf_nonuniform_voxels() {
        let n = 8;
        let field: Vec<f64> = (0..n*n*n).map(|i| (i as f64) * 0.001).collect();

        // Create a spherical mask in the center
        let mut mask = vec![0u8; n * n * n];
        let center = n / 2;
        let radius = n / 4;

        for z in 0..n {
            for y in 0..n {
                for x in 0..n {
                    let dx = (x as i32) - (center as i32);
                    let dy = (y as i32) - (center as i32);
                    let dz = (z as i32) - (center as i32);
                    if dx*dx + dy*dy + dz*dz <= (radius * radius) as i32 {
                        mask[x + y * n + z * n * n] = 1;
                    }
                }
            }
        }

        // Anisotropic voxel sizes
        let grid = Grid::new(n, n, n, 0.5, 1.0, 2.0);
        let local = pdf_core(
            &field, &mask, &grid,
            (0.0, 0.0, 1.0), 1e-5, 20, |_, _| {}
        );

        for (i, &val) in local.iter().enumerate() {
            assert!(val.is_finite(), "PDF with nonuniform voxels should be finite at index {}", i);
        }

        // Masked voxels should be zero
        for i in 0..n*n*n {
            if mask[i] == 0 {
                assert_eq!(local[i], 0.0, "Outside mask should be zero");
            }
        }
    }

    #[test]
    fn test_pdf_varying_field() {
        let n = 8;

        // Create a spatially varying field (quadratic in z)
        let mut field = vec![0.0; n * n * n];
        for z in 0..n {
            for y in 0..n {
                for x in 0..n {
                    let idx = x + y * n + z * n * n;
                    let zf = (z as f64) / (n as f64);
                    field[idx] = zf * zf * 0.5;
                }
            }
        }

        // Spherical mask
        let mut mask = vec![0u8; n * n * n];
        let center = n / 2;
        let radius = n / 4;
        for z in 0..n {
            for y in 0..n {
                for x in 0..n {
                    let dx = (x as i32) - (center as i32);
                    let dy = (y as i32) - (center as i32);
                    let dz = (z as i32) - (center as i32);
                    if dx*dx + dy*dy + dz*dz <= (radius * radius) as i32 {
                        mask[x + y * n + z * n * n] = 1;
                    }
                }
            }
        }

        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
        let local = pdf_core(
            &field, &mask, &grid,
            (0.0, 0.0, 1.0), 1e-5, 30, |_, _| {}
        );

        // All values should be finite
        for (i, &val) in local.iter().enumerate() {
            assert!(val.is_finite(), "Varying field should produce finite results at index {}", i);
        }

        // The local field inside the mask should differ from the total field
        // (some background was removed)
        let mut any_changed = false;
        for i in 0..n*n*n {
            if mask[i] != 0 && (local[i] - field[i]).abs() > 1e-10 {
                any_changed = true;
                break;
            }
        }
        assert!(any_changed, "PDF should modify the field inside the mask for a varying field");
    }

    #[test]
    fn test_pdf_larger_volume() {
        // Use 16x16x16 to exercise more of the LSMR loop
        let n = 16;

        // Create a dipole-like field pattern
        let mut field = vec![0.0; n * n * n];
        let center = n / 2;
        for z in 0..n {
            for y in 0..n {
                for x in 0..n {
                    let dx = (x as f64) - (center as f64);
                    let dy = (y as f64) - (center as f64);
                    let dz = (z as f64) - (center as f64);
                    let r2 = dx * dx + dy * dy + dz * dz;
                    if r2 > 0.5 {
                        // Dipole-like field: (3*cos^2(theta) - 1) / r^3
                        let r = r2.sqrt();
                        let cos_theta = dz / r;
                        field[x + y * n + z * n * n] = (3.0 * cos_theta * cos_theta - 1.0) / (r * r * r) * 0.01;
                    }
                }
            }
        }

        // Spherical mask
        let mut mask = vec![0u8; n * n * n];
        let radius = n / 3;
        for z in 0..n {
            for y in 0..n {
                for x in 0..n {
                    let dx = (x as i32) - (center as i32);
                    let dy = (y as i32) - (center as i32);
                    let dz = (z as i32) - (center as i32);
                    if dx * dx + dy * dy + dz * dz <= (radius * radius) as i32 {
                        mask[x + y * n + z * n * n] = 1;
                    }
                }
            }
        }

        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
        let local = pdf_core(
            &field, &mask, &grid,
            (0.0, 0.0, 1.0), 1e-4, 50, |_, _| {}
        );

        assert_eq!(local.len(), n * n * n);
        for (i, &val) in local.iter().enumerate() {
            assert!(val.is_finite(), "PDF larger volume: finite at index {}", i);
        }

        // Masked-out voxels must be zero
        for i in 0..n * n * n {
            if mask[i] == 0 {
                assert_eq!(local[i], 0.0);
            }
        }
    }

    #[test]
    fn test_pdf_more_iterations() {
        // Test with more LSMR iterations to exercise convergence check
        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];
        let center = n / 2;
        let radius = n / 4;
        for z in 0..n {
            for y in 0..n {
                for x in 0..n {
                    let dx = (x as i32) - (center as i32);
                    let dy = (y as i32) - (center as i32);
                    let dz = (z as i32) - (center as i32);
                    if dx * dx + dy * dy + dz * dz <= (radius * radius) as i32 {
                        mask[x + y * n + z * n * n] = 1;
                    }
                }
            }
        }

        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);

        // With many iterations the result should converge
        let local_many = pdf_core(
            &field, &mask, &grid,
            (0.0, 0.0, 1.0), 1e-8, 100, |_, _| {}
        );

        let local_few = pdf_core(
            &field, &mask, &grid,
            (0.0, 0.0, 1.0), 1e-8, 5, |_, _| {}
        );

        // Both should be finite
        for &val in local_many.iter().chain(local_few.iter()) {
            assert!(val.is_finite());
        }
    }

    #[test]
    fn test_pdf_different_bdir() {
        // Test with non-standard B0 direction
        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];
        let center = n / 2;
        let radius = n / 4;
        for z in 0..n {
            for y in 0..n {
                for x in 0..n {
                    let dx = (x as i32) - (center as i32);
                    let dy = (y as i32) - (center as i32);
                    let dz = (z as i32) - (center as i32);
                    if dx * dx + dy * dy + dz * dz <= (radius * radius) as i32 {
                        mask[x + y * n + z * n * n] = 1;
                    }
                }
            }
        }

        // Tilted B0 direction
        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
        let local = pdf_core(
            &field, &mask, &grid,
            (0.1, 0.2, 0.97), 1e-5, 20, |_, _| {}
        );

        for &val in &local {
            assert!(val.is_finite());
        }
    }

    #[test]
    fn test_pdf_with_progress() {
        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];
        let center = n / 2;
        let radius = n / 4;
        for z in 0..n {
            for y in 0..n {
                for x in 0..n {
                    let dx = (x as i32) - (center as i32);
                    let dy = (y as i32) - (center as i32);
                    let dz = (z as i32) - (center as i32);
                    if dx * dx + dy * dy + dz * dz <= (radius * radius) as i32 {
                        mask[x + y * n + z * n * n] = 1;
                    }
                }
            }
        }

        let grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);
        let mut progress_calls = Vec::new();
        let local = pdf_core(
            &field, &mask, &grid,
            (0.0, 0.0, 1.0), 1e-5, 20,
            |iter, max| { progress_calls.push((iter, max)); }
        );

        assert_eq!(local.len(), n * n * n);
        assert!(!progress_calls.is_empty(), "Progress should be called at least once");
        for &val in &local {
            assert!(val.is_finite());
        }
    }

    #[test]
    fn test_pdf_all_mask() {
        // All voxels masked (no background) - should still work
        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 grid = Grid::new(n, n, n, 1.0, 1.0, 1.0);

        let local = pdf_core(
            &field, &mask, &grid,
            (0.0, 0.0, 1.0), 1e-5, 10, |_, _| {}
        );

        // With all voxels as "brain", the background mask is empty
        // so there's nothing to project onto => local should approximate field
        for &val in &local {
            assert!(val.is_finite());
        }
    }
}