1pub 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 assert_eq!(mask[5 + 5 * 10 + 5 * 100], 1);
53
54 assert_eq!(mask[0], 0);
56
57 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 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 let count: usize = mask.iter().map(|&m| m as usize).sum();
76 assert_eq!(count, 1);
77 }
78}