Skip to main content

qsm_core/utils/
mask.rs

1//! Mask generation utilities
2//!
3//! Provides functions for creating geometric masks (e.g. spheres) on 3D volumes.
4
5/// Create a binary sphere mask on a 3D volume
6///
7/// Generates a mask where voxels within the specified radius of the center
8/// are set to 1, and all others are 0. Uses Fortran (column-major) ordering
9/// to match NIfTI convention: index = x + y*nx + z*nx*ny.
10///
11/// # Arguments
12/// * `nx`, `ny`, `nz` - Volume dimensions
13/// * `center_x`, `center_y`, `center_z` - Sphere center in voxel coordinates
14/// * `radius` - Sphere radius in voxels
15///
16/// # Returns
17/// Flattened binary mask of length nx*ny*nz
18pub fn create_sphere_mask(
19    nx: usize, ny: usize, nz: usize,
20    center_x: f64, center_y: f64, center_z: f64,
21    radius: f64,
22) -> Vec<u8> {
23    let mut mask = vec![0u8; nx * ny * nz];
24    let r2 = radius * radius;
25
26    for k in 0..nz {
27        for j in 0..ny {
28            for i in 0..nx {
29                let dx = i as f64 - center_x;
30                let dy = j as f64 - center_y;
31                let dz = k as f64 - center_z;
32                if dx * dx + dy * dy + dz * dz <= r2 {
33                    mask[i + j * nx + k * nx * ny] = 1;
34                }
35            }
36        }
37    }
38
39    mask
40}
41
42#[cfg(test)]
43mod tests {
44    use super::*;
45
46    #[test]
47    fn test_sphere_mask_basic() {
48        let mask = create_sphere_mask(10, 10, 10, 5.0, 5.0, 5.0, 3.0);
49        assert_eq!(mask.len(), 1000);
50
51        // Center voxel should be inside
52        assert_eq!(mask[5 + 5 * 10 + 5 * 100], 1);
53
54        // Corner should be outside
55        assert_eq!(mask[0], 0);
56
57        // Count should be reasonable for a sphere of radius 3
58        let count: usize = mask.iter().map(|&m| m as usize).sum();
59        assert!(count > 50 && count < 200, "Sphere voxel count {} seems wrong", count);
60    }
61
62    #[test]
63    fn test_sphere_mask_non_cubic() {
64        let mask = create_sphere_mask(20, 10, 5, 10.0, 5.0, 2.5, 2.0);
65        assert_eq!(mask.len(), 1000);
66
67        // Center should be inside
68        assert_eq!(mask[10 + 5 * 20 + 2 * 20 * 10], 1);
69    }
70
71    #[test]
72    fn test_sphere_mask_zero_radius() {
73        let mask = create_sphere_mask(5, 5, 5, 2.0, 2.0, 2.0, 0.0);
74        // Only the exact center voxel (distance 0 <= 0)
75        let count: usize = mask.iter().map(|&m| m as usize).sum();
76        assert_eq!(count, 1);
77    }
78}