Skip to main content

spatialrust_voxelize/
occupancy.rs

1//! Dense voxel occupancy / count grids.
2
3use spatialrust_core::{HasPositions3, PointCloud, SpatialError, SpatialResult};
4
5/// How a voxel's value is filled from the points that fall in it.
6#[derive(Clone, Copy, Debug, PartialEq, Eq)]
7pub enum VoxelFill {
8    /// `1.0` if any point falls in the voxel, else `0.0`.
9    Occupancy,
10    /// The number of points that fall in the voxel.
11    Count,
12}
13
14/// Configuration for [`voxelize`].
15#[derive(Clone, Copy, Debug, PartialEq)]
16pub struct VoxelGridConfig {
17    /// Side length of each cubic voxel.
18    pub voxel_size: f32,
19    /// Grid origin (lower corner). `None` uses the cloud's minimum corner.
20    pub origin: Option<[f32; 3]>,
21    /// Grid dimensions `(nx, ny, nz)`. `None` derives them from the cloud bounds.
22    pub dims: Option<[usize; 3]>,
23    /// How to fill each voxel.
24    pub fill: VoxelFill,
25}
26
27impl Default for VoxelGridConfig {
28    fn default() -> Self {
29        Self { voxel_size: 0.1, origin: None, dims: None, fill: VoxelFill::Occupancy }
30    }
31}
32
33impl VoxelGridConfig {
34    /// Creates a config from the voxel size (auto origin/dims, occupancy fill).
35    #[must_use]
36    pub fn new(voxel_size: f32) -> Self {
37        Self { voxel_size, ..Self::default() }
38    }
39}
40
41/// A dense 3D grid in row-major `(z, y, x)` order.
42#[derive(Clone, Debug, PartialEq)]
43pub struct OccupancyGrid {
44    /// Grid dimensions `(nx, ny, nz)`.
45    pub dims: [usize; 3],
46    /// Lower corner of voxel `(0, 0, 0)`.
47    pub origin: [f32; 3],
48    /// Voxel side length.
49    pub voxel_size: f32,
50    /// Values, indexed `z * (ny * nx) + y * nx + x`.
51    pub data: Vec<f32>,
52}
53
54impl OccupancyGrid {
55    /// Total number of voxels (`nx * ny * nz`).
56    #[must_use]
57    pub fn len(&self) -> usize {
58        self.data.len()
59    }
60
61    /// Whether the grid has no voxels.
62    #[must_use]
63    pub fn is_empty(&self) -> bool {
64        self.data.is_empty()
65    }
66
67    /// Number of voxels with a non-zero value.
68    #[must_use]
69    pub fn occupied_count(&self) -> usize {
70        self.data.iter().filter(|&&v| v != 0.0).count()
71    }
72
73    /// Value at voxel `(x, y, z)`, or `None` if out of range.
74    #[must_use]
75    pub fn get(&self, x: usize, y: usize, z: usize) -> Option<f32> {
76        let [nx, ny, nz] = self.dims;
77        if x >= nx || y >= ny || z >= nz {
78            return None;
79        }
80        Some(self.data[z * ny * nx + y * nx + x])
81    }
82}
83
84/// Voxelizes a cloud into a dense occupancy / count grid.
85pub fn voxelize(cloud: &PointCloud, config: VoxelGridConfig) -> SpatialResult<OccupancyGrid> {
86    if config.voxel_size <= 0.0 || config.voxel_size.is_nan() {
87        return Err(SpatialError::InvalidArgument("voxel_size must be positive".to_owned()));
88    }
89    let (x, y, z) = cloud.positions3()?;
90    let len = cloud.len();
91
92    let origin = match config.origin {
93        Some(o) => o,
94        None => {
95            if len == 0 {
96                [0.0, 0.0, 0.0]
97            } else {
98                [
99                    x.iter().copied().fold(f32::INFINITY, f32::min),
100                    y.iter().copied().fold(f32::INFINITY, f32::min),
101                    z.iter().copied().fold(f32::INFINITY, f32::min),
102                ]
103            }
104        }
105    };
106
107    let inv = 1.0 / config.voxel_size;
108    let dims = match config.dims {
109        Some(d) => d,
110        None => {
111            if len == 0 {
112                [1, 1, 1]
113            } else {
114                let span = |vals: &[f32], o: f32| {
115                    let max = vals.iter().copied().fold(f32::NEG_INFINITY, f32::max);
116                    (((max - o) * inv).floor() as usize) + 1
117                };
118                [span(x, origin[0]), span(y, origin[1]), span(z, origin[2])]
119            }
120        }
121    };
122    let [nx, ny, nz] = dims;
123    let total = nx.checked_mul(ny).and_then(|v| v.checked_mul(nz)).ok_or_else(|| {
124        SpatialError::InvalidArgument("voxel grid dimensions overflow".to_owned())
125    })?;
126
127    let mut data = vec![0.0_f32; total];
128    for i in 0..len {
129        let vx = ((x[i] - origin[0]) * inv).floor();
130        let vy = ((y[i] - origin[1]) * inv).floor();
131        let vz = ((z[i] - origin[2]) * inv).floor();
132        if vx < 0.0 || vy < 0.0 || vz < 0.0 {
133            continue;
134        }
135        let (vx, vy, vz) = (vx as usize, vy as usize, vz as usize);
136        if vx >= nx || vy >= ny || vz >= nz {
137            continue;
138        }
139        let idx = vz * ny * nx + vy * nx + vx;
140        match config.fill {
141            VoxelFill::Occupancy => data[idx] = 1.0,
142            VoxelFill::Count => data[idx] += 1.0,
143        }
144    }
145
146    Ok(OccupancyGrid { dims, origin, voxel_size: config.voxel_size, data })
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152    use spatialrust_core::{PointCloudBuilder, StandardSchemas};
153
154    fn cloud(points: &[[f32; 3]]) -> PointCloud {
155        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
156        for p in points {
157            builder.push_point(*p).unwrap();
158        }
159        builder.build().unwrap()
160    }
161
162    #[test]
163    fn occupancy_marks_filled_voxels() {
164        // Two points one voxel apart along x.
165        let c = cloud(&[[0.05, 0.05, 0.05], [0.15, 0.05, 0.05]]);
166        let grid = voxelize(&c, VoxelGridConfig::new(0.1)).unwrap();
167        assert_eq!(grid.dims, [2, 1, 1]);
168        assert_eq!(grid.occupied_count(), 2);
169        assert_eq!(grid.get(0, 0, 0), Some(1.0));
170        assert_eq!(grid.get(1, 0, 0), Some(1.0));
171    }
172
173    #[test]
174    fn count_accumulates_points_per_voxel() {
175        // Three points in the same voxel, one in another.
176        let c = cloud(&[
177            [0.01, 0.01, 0.01],
178            [0.02, 0.02, 0.02],
179            [0.03, 0.03, 0.03],
180            [0.55, 0.01, 0.01],
181        ]);
182        let config = VoxelGridConfig { fill: VoxelFill::Count, ..VoxelGridConfig::new(0.1) };
183        let grid = voxelize(&c, config).unwrap();
184        assert_eq!(grid.get(0, 0, 0), Some(3.0));
185        assert_eq!(grid.occupied_count(), 2);
186        // The total count equals the number of in-bounds points.
187        let total: f32 = grid.data.iter().sum();
188        assert_eq!(total, 4.0);
189    }
190
191    #[test]
192    fn fixed_dims_drop_out_of_bounds_points() {
193        let c = cloud(&[[0.05, 0.05, 0.05], [100.0, 0.0, 0.0]]);
194        let config = VoxelGridConfig {
195            origin: Some([0.0, 0.0, 0.0]),
196            dims: Some([2, 2, 2]),
197            ..VoxelGridConfig::new(0.1)
198        };
199        let grid = voxelize(&c, config).unwrap();
200        assert_eq!(grid.dims, [2, 2, 2]);
201        assert_eq!(grid.len(), 8);
202        // The far point is outside the fixed grid and is dropped.
203        assert_eq!(grid.occupied_count(), 1);
204    }
205
206    #[test]
207    fn rejects_bad_voxel_size() {
208        let c = cloud(&[[0.0, 0.0, 0.0]]);
209        assert!(voxelize(&c, VoxelGridConfig::new(0.0)).is_err());
210    }
211}