qsm_core/bgremove/

vsharp.rs

//! V-SHARP background field removal
//!
//! Variable kernel SHARP uses multiple SMV kernel radii
//! to preserve more brain tissue at edges while still
//! removing background fields.
//!
//! Reference:
//! Wu, B., Li, W., Guidon, A., Liu, C. (2012).
//! "Whole brain susceptibility mapping using compressed sensing."
//! Magnetic Resonance in Medicine, 67(1):137-147. https://doi.org/10.1002/mrm.23000
//!
//! Reference implementation: https://github.com/kamesy/QSM.jl

use num_complex::Complex64;
use crate::fft::{fft3d, ifft3d};
use crate::kernels::smv::smv_kernel;

/// V-SHARP background field removal
///
/// Uses multiple SMV kernel radii, starting from largest and decreasing.
/// At each voxel, uses the smallest radius that doesn't touch the boundary.
///
/// # Arguments
/// * `field` - Unwrapped total field (nx * ny * nz)
/// * `mask` - Binary mask (nx * ny * nz), 1 = inside ROI
/// * `nx`, `ny`, `nz` - Array dimensions
/// * `vsx`, `vsy`, `vsz` - Voxel sizes in mm
/// * `radii` - SMV kernel radii in mm (should be sorted large to small)
/// * `threshold` - High-pass filter threshold (typically 0.05)
///
/// V-SHARP algorithm parameters
#[derive(Clone, Debug)]
pub struct VsharpParams {
    /// Deconvolution threshold
    pub threshold: f64,
    /// Maximum kernel radius in mm (as multiple of min voxel size; default: 18.0)
    pub max_radius_factor: f64,
    /// Minimum kernel radius in mm (as multiple of max voxel size; default: 2.0)
    pub min_radius_factor: f64,
}

impl Default for VsharpParams {
    fn default() -> Self {
        Self {
            threshold: 0.001,
            max_radius_factor: 18.0,
            min_radius_factor: 2.0,
        }
    }
}

/// # Returns
/// (local_field, eroded_mask)
pub fn vsharp(
    field: &[f64],
    mask: &[u8],
    nx: usize, ny: usize, nz: usize,
    vsx: f64, vsy: f64, vsz: f64,
    radii: &[f64],
    threshold: f64,
) -> (Vec<f64>, Vec<u8>) {
    if radii.is_empty() {
        return (vec![0.0; nx * ny * nz], mask.to_vec());
    }

    // If only one radius, use regular SHARP
    if radii.len() == 1 {
        return crate::bgremove::sharp::sharp(
            field, mask, nx, ny, nz, vsx, vsy, vsz, radii[0], threshold
        );
    }

    let n_total = nx * ny * nz;

    // Sort radii from largest to smallest
    let mut sorted_radii = radii.to_vec();
    sorted_radii.sort_by(|a, b| b.partial_cmp(a).unwrap());

    // FFT of field
    let mut field_complex: Vec<Complex64> = field.iter()
        .map(|&x| Complex64::new(x, 0.0))
        .collect();
    fft3d(&mut field_complex, nx, ny, nz);
    let field_fft = field_complex.clone();

    // Track which voxels have been processed and final mask
    let mut processed = vec![false; n_total];
    let mut local_field = vec![0.0; n_total];
    let mut final_mask = vec![0u8; n_total];

    // Process each radius (large to small)
    let delta = 1.0 - 1e-7_f64.sqrt();

    // Store the first (largest) kernel's inverse for deconvolution
    let mut inverse_kernel: Option<Vec<f64>> = None;

    for &radius in &sorted_radii {
        // Generate SMV kernel
        let s_kernel = smv_kernel(nx, ny, nz, vsx, vsy, vsz, radius);

        // FFT of SMV kernel
        let mut s_complex: Vec<Complex64> = s_kernel.iter()
            .map(|&x| Complex64::new(x, 0.0))
            .collect();
        fft3d(&mut s_complex, nx, ny, nz);
        let s_fft: Vec<f64> = s_complex.iter().map(|c| c.re).collect();

        // Store inverse of first (largest) kernel
        if inverse_kernel.is_none() {
            inverse_kernel = Some(s_fft.iter().map(|&s| {
                let one_minus_s = 1.0 - s;
                if one_minus_s.abs() < threshold {
                    0.0
                } else {
                    1.0 / one_minus_s
                }
            }).collect());
        }

        // Erode mask for this radius
        let mask_f64: Vec<f64> = mask.iter().map(|&m| m as f64).collect();
        let mut mask_complex: Vec<Complex64> = mask_f64.iter()
            .map(|&x| Complex64::new(x, 0.0))
            .collect();

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

        // Convolve mask with SMV kernel
        for i in 0..n_total {
            mask_complex[i] *= s_fft[i];
        }

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

        // Current eroded mask
        let current_mask: Vec<bool> = mask_complex.iter()
            .map(|c| c.re > delta)
            .collect();

        // Apply high-pass filter: multiply by (1-S)
        let mut filtered = field_fft.clone();
        for i in 0..n_total {
            filtered[i] *= 1.0 - s_fft[i];
        }

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

        // For voxels that are in current mask but not yet processed,
        // store the filtered value
        for i in 0..n_total {
            if current_mask[i] && !processed[i] {
                local_field[i] = filtered[i].re;
                processed[i] = true;
                final_mask[i] = 1;
            }
        }
    }

    // Deconvolution with largest kernel's inverse
    if let Some(inv_kernel) = inverse_kernel {
        let mut local_complex: Vec<Complex64> = local_field.iter()
            .map(|&x| Complex64::new(x, 0.0))
            .collect();

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

        for i in 0..n_total {
            local_complex[i] *= inv_kernel[i];
        }

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

        // Apply final mask
        for i in 0..n_total {
            local_field[i] = if final_mask[i] == 1 { local_complex[i].re } else { 0.0 };
        }
    }

    (local_field, final_mask)
}

/// V-SHARP with progress callback
///
/// Same as `vsharp` but calls `progress_callback(radius_index, total_radii)` for each radius.
pub fn vsharp_with_progress<F>(
    field: &[f64],
    mask: &[u8],
    nx: usize, ny: usize, nz: usize,
    vsx: f64, vsy: f64, vsz: f64,
    radii: &[f64],
    threshold: f64,
    mut progress_callback: F,
) -> (Vec<f64>, Vec<u8>)
where
    F: FnMut(usize, usize),
{
    if radii.is_empty() {
        return (vec![0.0; nx * ny * nz], mask.to_vec());
    }

    // If only one radius, use regular SHARP
    if radii.len() == 1 {
        progress_callback(1, 1);
        return crate::bgremove::sharp::sharp(
            field, mask, nx, ny, nz, vsx, vsy, vsz, radii[0], threshold
        );
    }

    let n_total = nx * ny * nz;
    let n_radii = radii.len();

    // Sort radii from largest to smallest
    let mut sorted_radii = radii.to_vec();
    sorted_radii.sort_by(|a, b| b.partial_cmp(a).unwrap());

    // FFT of field
    let mut field_complex: Vec<Complex64> = field.iter()
        .map(|&x| Complex64::new(x, 0.0))
        .collect();
    fft3d(&mut field_complex, nx, ny, nz);
    let field_fft = field_complex.clone();

    // Track which voxels have been processed and final mask
    let mut processed = vec![false; n_total];
    let mut local_field = vec![0.0; n_total];
    let mut final_mask = vec![0u8; n_total];

    let delta = 1.0 - 1e-7_f64.sqrt();
    let mut inverse_kernel: Option<Vec<f64>> = None;

    for (idx, &radius) in sorted_radii.iter().enumerate() {
        // Report progress
        progress_callback(idx + 1, n_radii);

        // Generate SMV kernel
        let s_kernel = smv_kernel(nx, ny, nz, vsx, vsy, vsz, radius);

        // FFT of SMV kernel
        let mut s_complex: Vec<Complex64> = s_kernel.iter()
            .map(|&x| Complex64::new(x, 0.0))
            .collect();
        fft3d(&mut s_complex, nx, ny, nz);
        let s_fft: Vec<f64> = s_complex.iter().map(|c| c.re).collect();

        // Store inverse of first (largest) kernel
        if inverse_kernel.is_none() {
            inverse_kernel = Some(s_fft.iter().map(|&s| {
                let one_minus_s = 1.0 - s;
                if one_minus_s.abs() < threshold {
                    0.0
                } else {
                    1.0 / one_minus_s
                }
            }).collect());
        }

        // Erode mask for this radius
        let mask_f64: Vec<f64> = mask.iter().map(|&m| m as f64).collect();
        let mut mask_complex: Vec<Complex64> = mask_f64.iter()
            .map(|&x| Complex64::new(x, 0.0))
            .collect();

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

        for i in 0..n_total {
            mask_complex[i] *= s_fft[i];
        }

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

        let current_mask: Vec<bool> = mask_complex.iter()
            .map(|c| c.re > delta)
            .collect();

        // Apply high-pass filter
        let mut filtered = field_fft.clone();
        for i in 0..n_total {
            filtered[i] *= 1.0 - s_fft[i];
        }

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

        for i in 0..n_total {
            if current_mask[i] && !processed[i] {
                local_field[i] = filtered[i].re;
                processed[i] = true;
                final_mask[i] = 1;
            }
        }
    }

    // Deconvolution
    if let Some(inv_kernel) = inverse_kernel {
        let mut local_complex: Vec<Complex64> = local_field.iter()
            .map(|&x| Complex64::new(x, 0.0))
            .collect();

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

        for i in 0..n_total {
            local_complex[i] *= inv_kernel[i];
        }

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

        for i in 0..n_total {
            local_field[i] = if final_mask[i] == 1 { local_complex[i].re } else { 0.0 };
        }
    }

    (local_field, final_mask)
}

/// V-SHARP with default parameters
pub fn vsharp_default(
    field: &[f64],
    mask: &[u8],
    nx: usize, ny: usize, nz: usize,
    vsx: f64, vsy: f64, vsz: f64,
) -> (Vec<f64>, Vec<u8>) {
    let min_vox = vsx.min(vsy).min(vsz);
    let max_vox = vsx.max(vsy).max(vsz);

    // Default radii: 18*min_vox down to 2*max_vox in steps of 2*max_vox
    let mut radii = Vec::new();
    let mut r = 18.0 * min_vox;
    while r >= 2.0 * max_vox {
        radii.push(r);
        r -= 2.0 * max_vox;
    }

    if radii.is_empty() {
        radii.push(18.0 * min_vox);
    }

    vsharp(field, mask, nx, ny, nz, vsx, vsy, vsz, &radii, 0.05)
}

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

    #[test]
    fn test_vsharp_zero_field() {
        let n = 8;
        let field = vec![0.0; n * n * n];
        let mask = vec![1u8; n * n * n];

        let radii = vec![4.0, 3.0, 2.0];
        let (local, _) = vsharp(&field, &mask, n, n, n, 1.0, 1.0, 1.0, &radii, 0.05);

        for &val in local.iter() {
            assert!(val.abs() < 1e-10);
        }
    }

    #[test]
    fn test_vsharp_preserves_more_than_sharp() {
        let n = 16;
        let field = vec![0.0; n * n * n];
        let mask = vec![1u8; n * n * n];

        // V-SHARP with multiple radii
        let radii = vec![5.0, 4.0, 3.0, 2.0];
        let (_, vsharp_mask) = vsharp(&field, &mask, n, n, n, 1.0, 1.0, 1.0, &radii, 0.05);

        // SHARP with single large radius
        let (_, sharp_mask) = crate::bgremove::sharp::sharp(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0, 5.0, 0.05
        );

        let vsharp_count: usize = vsharp_mask.iter().map(|&m| m as usize).sum();
        let sharp_count: usize = sharp_mask.iter().map(|&m| m as usize).sum();

        // V-SHARP should preserve at least as many voxels as SHARP
        assert!(vsharp_count >= sharp_count,
            "V-SHARP {} should preserve at least as many as SHARP {}",
            vsharp_count, sharp_count);
    }

    #[test]
    fn test_vsharp_nonuniform_voxels() {
        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];

        // Anisotropic voxel sizes
        let radii = vec![4.0, 3.0, 2.0];
        let (local, final_mask) = vsharp(
            &field, &mask, n, n, n, 0.5, 1.0, 2.0, &radii, 0.05
        );

        // All values should be finite
        for (i, &val) in local.iter().enumerate() {
            assert!(val.is_finite(), "V-SHARP nonuniform voxels: finite at index {}", i);
        }

        // Final mask should have some voxels
        let mask_count: usize = final_mask.iter().map(|&m| m as usize).sum();
        assert!(mask_count > 0, "V-SHARP nonuniform: final mask should have some voxels");
    }

    #[test]
    fn test_vsharp_single_radius() {
        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];

        // Single radius should delegate to SHARP
        let radii = vec![3.0];
        let (local, final_mask) = vsharp(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0, &radii, 0.05
        );

        // All values should be finite
        for (i, &val) in local.iter().enumerate() {
            assert!(val.is_finite(), "V-SHARP single radius: finite at index {}", i);
        }

        // Result should match SHARP with same radius
        let (sharp_local, sharp_mask) = crate::bgremove::sharp::sharp(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0, 3.0, 0.05
        );

        for i in 0..n*n*n {
            assert!(
                (local[i] - sharp_local[i]).abs() < 1e-10,
                "Single-radius V-SHARP should match SHARP at index {}", i
            );
        }

        assert_eq!(final_mask, sharp_mask, "Single-radius V-SHARP mask should match SHARP mask");
    }

    #[test]
    fn test_vsharp_empty_radii() {
        // Empty radii should return zeros and the original mask
        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, returned_mask) = vsharp(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0, &[], 0.05
        );

        for &val in &local {
            assert_eq!(val, 0.0, "Empty radii should return zero local field");
        }
        assert_eq!(returned_mask, mask, "Empty radii should return original mask");
    }

    #[test]
    fn test_vsharp_larger_volume() {
        // 16x16x16 volume with spherical mask
        let n = 16;
        let mut field = vec![0.0; n * n * n];
        // Linear background field in z
        for z in 0..n {
            for y in 0..n {
                for x in 0..n {
                    field[x + y * n + z * n * n] = (z as f64) * 0.1;
                }
            }
        }

        // Spherical mask
        let mut mask = vec![0u8; n * n * n];
        let center = n / 2;
        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 radii = vec![6.0, 4.0, 3.0, 2.0];
        let (local, final_mask) = vsharp(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0, &radii, 0.05
        );

        assert_eq!(local.len(), n * n * n);
        for &val in &local {
            assert!(val.is_finite(), "V-SHARP larger volume values should be finite");
        }

        let mask_count: usize = final_mask.iter().map(|&m| m as usize).sum();
        assert!(mask_count > 0, "V-SHARP larger volume should have voxels in final mask");

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

    #[test]
    fn test_vsharp_with_progress() {
        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 radii = vec![4.0, 3.0, 2.0];
        let mut progress_calls = Vec::new();
        let (local, _) = vsharp_with_progress(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0, &radii, 0.05,
            |idx, total| { progress_calls.push((idx, total)); }
        );

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

    #[test]
    fn test_vsharp_with_progress_single_radius() {
        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 radii = vec![3.0];
        let mut progress_calls = Vec::new();
        let (local, _) = vsharp_with_progress(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0, &radii, 0.05,
            |idx, total| { progress_calls.push((idx, total)); }
        );

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

    #[test]
    fn test_vsharp_with_progress_empty_radii() {
        let n = 8;
        let field = vec![0.0; n * n * n];
        let mask = vec![1u8; n * n * n];

        let mut progress_calls = Vec::new();
        let (local, returned_mask) = vsharp_with_progress(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0, &[], 0.05,
            |idx, total| { progress_calls.push((idx, total)); }
        );

        for &val in &local {
            assert_eq!(val, 0.0);
        }
        assert_eq!(returned_mask, mask);
    }

    #[test]
    fn test_vsharp_default_wrapper() {
        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, final_mask) = vsharp_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());
        }
        // Final mask should have some voxels
        let count: usize = final_mask.iter().map(|&m| m as usize).sum();
        // May be 0 if volume is too small for default radii, but should not crash
        assert!(count <= n * n * n);
    }

    #[test]
    fn test_vsharp_unsorted_radii() {
        // Radii given in arbitrary order - should be sorted internally
        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 radii_sorted = vec![4.0, 3.0, 2.0];
        let radii_unsorted = vec![2.0, 4.0, 3.0];

        let (local_sorted, mask_sorted) = vsharp(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0, &radii_sorted, 0.05
        );
        let (local_unsorted, mask_unsorted) = vsharp(
            &field, &mask, n, n, n, 1.0, 1.0, 1.0, &radii_unsorted, 0.05
        );

        // Results should be the same regardless of input order
        assert_eq!(mask_sorted, mask_unsorted, "Sorted and unsorted radii should give same mask");
        for i in 0..n * n * n {
            assert!(
                (local_sorted[i] - local_unsorted[i]).abs() < 1e-10,
                "Results should match at index {}",
                i
            );
        }
    }
}