mesh_shell/offset/
grid.rs

1//! Voxel grid for SDF (Signed Distance Field) computation.
2
3use nalgebra::Point3;
4use rayon::prelude::*;
5use tracing::{debug, info};
6
7use mesh_repair::Mesh;
8
9use crate::error::{ShellError, ShellResult};
10
11/// Parameters for SDF offset operation.
12#[derive(Debug, Clone)]
13pub struct SdfOffsetParams {
14    /// Voxel size in mm (smaller = more detail, more memory).
15    pub voxel_size_mm: f64,
16    /// Padding beyond mesh bounds in mm.
17    pub padding_mm: f64,
18    /// Maximum number of voxels before error (memory safety).
19    pub max_voxels: usize,
20    /// Number of nearest neighbors for offset interpolation.
21    pub offset_neighbors: usize,
22}
23
24impl Default for SdfOffsetParams {
25    fn default() -> Self {
26        Self {
27            voxel_size_mm: 0.75,
28            padding_mm: 12.0,
29            max_voxels: 50_000_000,
30            offset_neighbors: 8,
31        }
32    }
33}
34
35/// 3D voxel grid for SDF computation.
36pub struct SdfGrid {
37    /// Grid dimensions [x, y, z].
38    pub dims: [usize; 3],
39    /// Origin point (min corner) in world coordinates.
40    pub origin: Point3<f64>,
41    /// Voxel size in mm.
42    pub voxel_size: f64,
43    /// Signed distance values (negative = inside mesh).
44    pub values: Vec<f32>,
45    /// Offset values at each voxel (for variable offset).
46    pub offsets: Vec<f32>,
47}
48
49impl SdfGrid {
50    /// Create a grid sized to contain mesh with padding.
51    pub fn from_mesh_bounds(
52        mesh: &Mesh,
53        voxel_size: f64,
54        padding: f64,
55        max_voxels: usize,
56    ) -> ShellResult<Self> {
57        let (min, max) = mesh.bounds().ok_or(ShellError::EmptyMesh)?;
58
59        // Add padding to bounds
60        let origin = Point3::new(min.x - padding, min.y - padding, min.z - padding);
61        let max_padded = Point3::new(max.x + padding, max.y + padding, max.z + padding);
62
63        // Compute grid dimensions
64        let extent = max_padded - origin;
65        let dims = [
66            ((extent.x / voxel_size).ceil() as usize).max(1),
67            ((extent.y / voxel_size).ceil() as usize).max(1),
68            ((extent.z / voxel_size).ceil() as usize).max(1),
69        ];
70
71        let total_voxels = dims[0] * dims[1] * dims[2];
72        if total_voxels > max_voxels {
73            return Err(ShellError::GridTooLarge {
74                dims,
75                total: total_voxels,
76                max: max_voxels,
77            });
78        }
79
80        info!(
81            dims = ?dims,
82            total = total_voxels,
83            voxel_size_mm = voxel_size,
84            "Creating SDF grid"
85        );
86
87        Ok(Self {
88            dims,
89            origin,
90            voxel_size,
91            values: vec![0.0; total_voxels],
92            offsets: vec![0.0; total_voxels],
93        })
94    }
95
96    /// Total number of voxels in the grid.
97    #[inline]
98    pub fn total_voxels(&self) -> usize {
99        self.dims[0] * self.dims[1] * self.dims[2]
100    }
101
102    /// Convert 3D grid coordinates to linear index.
103    #[inline]
104    pub fn linearize(&self, x: usize, y: usize, z: usize) -> usize {
105        x + y * self.dims[0] + z * self.dims[0] * self.dims[1]
106    }
107
108    /// Convert linear index to 3D grid coordinates.
109    #[inline]
110    pub fn delinearize(&self, idx: usize) -> [usize; 3] {
111        let z = idx / (self.dims[0] * self.dims[1]);
112        let rem = idx % (self.dims[0] * self.dims[1]);
113        let y = rem / self.dims[0];
114        let x = rem % self.dims[0];
115        [x, y, z]
116    }
117
118    /// Get world position of voxel center.
119    #[inline]
120    pub fn voxel_center(&self, x: usize, y: usize, z: usize) -> Point3<f64> {
121        Point3::new(
122            self.origin.x + (x as f64 + 0.5) * self.voxel_size,
123            self.origin.y + (y as f64 + 0.5) * self.voxel_size,
124            self.origin.z + (z as f64 + 0.5) * self.voxel_size,
125        )
126    }
127
128    /// Compute SDF values using mesh_to_sdf crate.
129    pub fn compute_sdf(&mut self, mesh: &Mesh) {
130        use mesh_to_sdf::{generate_grid_sdf, Grid, SignMethod, Topology};
131
132        info!(vertices = mesh.vertices.len(), "Computing SDF");
133
134        // Convert Mesh to mesh_to_sdf format
135        let vertices: Vec<[f32; 3]> = mesh
136            .vertices
137            .iter()
138            .map(|v| [v.position.x as f32, v.position.y as f32, v.position.z as f32])
139            .collect();
140
141        let indices: Vec<u32> = mesh.faces.iter().flat_map(|f| f.iter().copied()).collect();
142
143        // Create grid specification for mesh_to_sdf
144        let grid = Grid::from_bounding_box(
145            &[
146                self.origin.x as f32,
147                self.origin.y as f32,
148                self.origin.z as f32,
149            ],
150            &[
151                (self.origin.x + self.dims[0] as f64 * self.voxel_size) as f32,
152                (self.origin.y + self.dims[1] as f64 * self.voxel_size) as f32,
153                (self.origin.z + self.dims[2] as f64 * self.voxel_size) as f32,
154            ],
155            [self.dims[0], self.dims[1], self.dims[2]],
156        );
157
158        // Generate SDF using Raycast for robust sign determination
159        let sdf_values = generate_grid_sdf(
160            &vertices,
161            Topology::TriangleList(Some(&indices)),
162            &grid,
163            SignMethod::Raycast,
164        );
165
166        self.values = sdf_values;
167
168        debug!(
169            min_sdf = self.values.iter().copied().fold(f32::INFINITY, f32::min),
170            max_sdf = self.values.iter().copied().fold(f32::NEG_INFINITY, f32::max),
171            "SDF computed"
172        );
173    }
174
175    /// Interpolate offset values from mesh vertices into the grid.
176    ///
177    /// Uses inverse distance weighted interpolation from K nearest vertices.
178    pub fn interpolate_offsets(&mut self, mesh: &Mesh, k_neighbors: usize) {
179        use kiddo::KdTree;
180
181        info!(
182            voxels = self.total_voxels(),
183            vertices = mesh.vertices.len(),
184            k = k_neighbors,
185            "Interpolating offsets (building KD-tree)"
186        );
187
188        // Build KD-tree for fast K-nearest-neighbor queries
189        let mut kdtree: KdTree<f64, 3> = KdTree::new();
190        for (i, v) in mesh.vertices.iter().enumerate() {
191            kdtree.add(&[v.position.x, v.position.y, v.position.z], i as u64);
192        }
193
194        // Pre-extract offsets for fast lookup
195        let vertex_offsets: Vec<f32> = mesh
196            .vertices
197            .iter()
198            .map(|v| v.offset.unwrap_or(0.0))
199            .collect();
200
201        info!("KD-tree built, interpolating {} voxels", self.total_voxels());
202
203        let [_dim_x, dim_y, dim_z] = self.dims;
204        let voxel_size = self.voxel_size;
205        let origin = self.origin;
206
207        // Parallel interpolation with KD-tree queries
208        // Use ZYX ordering to match mesh_to_sdf's SDF value array layout
209        let offsets: Vec<f32> = (0..self.total_voxels())
210            .into_par_iter()
211            .map(|idx| {
212                // Delinearize using ZYX order (Z varies fastest) to match mesh_to_sdf
213                let z = idx % dim_z;
214                let y = (idx / dim_z) % dim_y;
215                let x = idx / (dim_y * dim_z);
216
217                // Voxel center in world coordinates
218                let voxel_pos = nalgebra::Point3::new(
219                    origin.x + (x as f64 + 0.5) * voxel_size,
220                    origin.y + (y as f64 + 0.5) * voxel_size,
221                    origin.z + (z as f64 + 0.5) * voxel_size,
222                );
223
224                // Query K nearest neighbors from KD-tree
225                let query = [voxel_pos.x, voxel_pos.y, voxel_pos.z];
226                let nearest = kdtree.nearest_n::<kiddo::SquaredEuclidean>(&query, k_neighbors);
227
228                // Inverse distance weighted interpolation
229                let mut total_weight = 0.0;
230                let mut weighted_offset = 0.0;
231                for neighbor in nearest {
232                    let dist = neighbor.distance.sqrt();
233                    let w = 1.0 / (dist + 0.001);
234                    total_weight += w;
235                    weighted_offset += w * vertex_offsets[neighbor.item as usize] as f64;
236                }
237
238                if total_weight > 0.0 {
239                    (weighted_offset / total_weight) as f32
240                } else {
241                    0.0
242                }
243            })
244            .collect();
245
246        self.offsets = offsets;
247
248        debug!(
249            min_offset = self.offsets.iter().copied().fold(f32::INFINITY, f32::min),
250            max_offset = self.offsets.iter().copied().fold(f32::NEG_INFINITY, f32::max),
251            "Offsets interpolated"
252        );
253    }
254
255    /// Apply variable offset by adjusting SDF with offset values.
256    ///
257    /// After this, extracting isosurface at SDF=0 gives the variable offset surface.
258    pub fn apply_variable_offset(&mut self) {
259        info!("Applying variable offset to SDF");
260
261        // Subtract offset to expand surface (positive offset = outward)
262        self.values
263            .par_iter_mut()
264            .zip(self.offsets.par_iter())
265            .for_each(|(sdf, offset)| {
266                *sdf -= *offset;
267            });
268
269        debug!(
270            min_adjusted = self.values.iter().copied().fold(f32::INFINITY, f32::min),
271            max_adjusted = self.values.iter().copied().fold(f32::NEG_INFINITY, f32::max),
272            "Variable offset applied"
273        );
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use super::*;
280    use mesh_repair::Vertex;
281
282    fn create_unit_cube() -> Mesh {
283        let mut mesh = Mesh::new();
284
285        // Cube vertices (0-10mm)
286        mesh.vertices.push(Vertex::from_coords(0.0, 0.0, 0.0));
287        mesh.vertices.push(Vertex::from_coords(10.0, 0.0, 0.0));
288        mesh.vertices.push(Vertex::from_coords(10.0, 10.0, 0.0));
289        mesh.vertices.push(Vertex::from_coords(0.0, 10.0, 0.0));
290        mesh.vertices.push(Vertex::from_coords(0.0, 0.0, 10.0));
291        mesh.vertices.push(Vertex::from_coords(10.0, 0.0, 10.0));
292        mesh.vertices.push(Vertex::from_coords(10.0, 10.0, 10.0));
293        mesh.vertices.push(Vertex::from_coords(0.0, 10.0, 10.0));
294
295        // Faces
296        mesh.faces.push([0, 1, 2]);
297        mesh.faces.push([0, 2, 3]);
298        mesh.faces.push([4, 6, 5]);
299        mesh.faces.push([4, 7, 6]);
300        mesh.faces.push([0, 5, 1]);
301        mesh.faces.push([0, 4, 5]);
302        mesh.faces.push([2, 7, 3]);
303        mesh.faces.push([2, 6, 7]);
304        mesh.faces.push([0, 3, 7]);
305        mesh.faces.push([0, 7, 4]);
306        mesh.faces.push([1, 5, 6]);
307        mesh.faces.push([1, 6, 2]);
308
309        mesh
310    }
311
312    #[test]
313    fn test_grid_construction() {
314        let mesh = create_unit_cube();
315        let grid = SdfGrid::from_mesh_bounds(&mesh, 1.0, 5.0, 1_000_000).unwrap();
316
317        // Grid should be 20x20x20 (10mm cube + 5mm padding on each side = 20mm)
318        assert_eq!(grid.dims[0], 20);
319        assert_eq!(grid.dims[1], 20);
320        assert_eq!(grid.dims[2], 20);
321        assert_eq!(grid.total_voxels(), 8000);
322    }
323
324    #[test]
325    fn test_grid_too_large() {
326        let mesh = create_unit_cube();
327        let result = SdfGrid::from_mesh_bounds(&mesh, 0.1, 5.0, 1000);
328        assert!(result.is_err());
329    }
330}
mesh_shell/offset/grid.rs

mesh_shell/offset/
grid.rs
