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::fft::{fft3d, ifft3d};
use crate::kernels::dipole::dipole_kernel;

#[cfg(feature = "parallel")]
use crate::par::*;

/// PDF background field removal
///
/// # Arguments
/// * `field` - Total field (nx * ny * nz)
/// * `mask` - Binary mask (nx * ny * nz), 1 = brain, 0 = background
/// * `nx`, `ny`, `nz` - Array dimensions
/// * `vsx`, `vsy`, `vsz` - Voxel sizes in mm
/// * `bdir` - B0 field direction
/// * `tol` - Convergence tolerance for LSMR
/// * `max_iter` - Maximum iterations for LSMR
///
/// PDF algorithm parameters
#[derive(Clone, Debug)]
pub struct PdfParams {
    /// Convergence tolerance
    pub tol: f64,
}

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

/// # Returns
/// Local field with background removed
pub fn pdf(
    field: &[f64],
    mask: &[u8],
    nx: usize, ny: usize, nz: usize,
    vsx: f64, vsy: f64, vsz: f64,
    bdir: (f64, f64, f64),
    tol: f64,
    max_iter: usize,
) -> Vec<f64> {
    pdf_with_progress(field, mask, nx, ny, nz, vsx, vsy, vsz, bdir, tol, max_iter, |_, _| {})
}



/// 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()
}

/// Vector 2-norm
fn vec_norm(v: &[f64]) -> f64 {
    v.iter().map(|&x| x * x).sum::<f64>().sqrt()
}

/// PDF with progress callback
///
/// Same as `pdf` but calls `progress_callback(iteration, max_iter)` each iteration.
pub fn pdf_with_progress<F>(
    field: &[f64],
    mask: &[u8],
    nx: usize, ny: usize, nz: usize,
    vsx: f64, vsy: f64, vsz: f64,
    bdir: (f64, f64, f64),
    tol: f64,
    max_iter: usize,
    mut progress_callback: F,
) -> Vec<f64>
where
    F: FnMut(usize, usize),
{
    let n_total = nx * ny * nz;

    // Generate dipole kernel
    let d_kernel = dipole_kernel(nx, ny, nz, vsx, vsy, vsz, 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_callback(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_callback(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_callback(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 * phi_bar;

        // 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 * zeta_bar;

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

        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_callback(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
}

/// PDF with default parameters (adaptive max_iter matching QSM.jl: ceil(sqrt(nx*ny*nz)))
pub fn pdf_default(
    field: &[f64],
    mask: &[u8],
    nx: usize, ny: usize, nz: usize,
    vsx: f64, vsy: f64, vsz: f64,
) -> Vec<f64> {
    let max_iter = ((nx * ny * nz) as f64).sqrt().ceil() as usize;
    pdf(
        field, mask, nx, ny, nz, vsx, vsy, vsz,
        (0.0, 0.0, 1.0),  // bdir
        1e-5,              // tol
        max_iter           // adaptive max_iter
    )
}

#[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 local = pdf(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0,
            (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 local = pdf(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0,
            (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 local = pdf(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0,
            (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 local = pdf(
            &field, &mask, n, n, n, 0.5, 1.0, 2.0,
            (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 local = pdf(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0,
            (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 local = pdf(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0,
            (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;
                    }
                }
            }
        }

        // With many iterations the result should converge
        let local_many = pdf(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0,
            (0.0, 0.0, 1.0), 1e-8, 100
        );

        let local_few = pdf(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0,
            (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 local = pdf(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0,
            (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 mut progress_calls = Vec::new();
        let local = pdf_with_progress(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0,
            (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_default_wrapper() {
        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 local = pdf_default(&field, &mask, n, n, n, 1.0, 1.0, 1.0);
        assert_eq!(local.len(), n * n * n);
        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 local = pdf(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0,
            (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());
        }
    }
}