polyscope_structures/volume_grid/
mod.rs

1//! Volume grid structure for visualizing regular 3D grids.
2
3mod scalar_quantity;
4
5pub use scalar_quantity::*;
6
7use glam::{Mat4, UVec3, Vec3, Vec4};
8use polyscope_core::pick::PickResult;
9use polyscope_core::quantity::Quantity;
10use polyscope_core::structure::{HasQuantities, RenderContext, Structure};
11use polyscope_render::CurveNetworkRenderData;
12
13/// A regular 3D grid structure.
14///
15/// `VolumeGrid` represents a regular axis-aligned 3D grid defined by:
16/// - Grid dimensions (number of nodes in X, Y, Z)
17/// - Bounding box (min and max corners in world space)
18///
19/// Node values are at grid vertices, cell values are at grid cell centers.
20pub struct VolumeGrid {
21    name: String,
22
23    // Grid parameters
24    node_dim: UVec3, // Number of nodes in each dimension
25    bound_min: Vec3, // Minimum corner of the grid
26    bound_max: Vec3, // Maximum corner of the grid
27
28    // Common structure fields
29    enabled: bool,
30    transform: Mat4,
31    quantities: Vec<Box<dyn Quantity>>,
32
33    // Visualization parameters
34    color: Vec4,
35    edge_color: Vec4,
36    edge_width: f32,
37    cube_size_factor: f32,
38
39    // GPU resources (bounding box wireframe)
40    render_data: Option<CurveNetworkRenderData>,
41}
42
43impl VolumeGrid {
44    /// Creates a new volume grid.
45    ///
46    /// # Arguments
47    /// * `name` - The name of the grid
48    /// * `node_dim` - Number of nodes in each dimension (X, Y, Z)
49    /// * `bound_min` - Minimum corner of the grid bounding box
50    /// * `bound_max` - Maximum corner of the grid bounding box
51    pub fn new(name: impl Into<String>, node_dim: UVec3, bound_min: Vec3, bound_max: Vec3) -> Self {
52        Self {
53            name: name.into(),
54            node_dim,
55            bound_min,
56            bound_max,
57            enabled: true,
58            transform: Mat4::IDENTITY,
59            quantities: Vec::new(),
60            color: Vec4::new(0.5, 0.5, 0.5, 1.0),
61            edge_color: Vec4::new(0.0, 0.0, 0.0, 1.0),
62            edge_width: 1.0,
63            cube_size_factor: 0.0,
64            render_data: None,
65        }
66    }
67
68    /// Creates a volume grid with uniform dimensions.
69    pub fn new_uniform(
70        name: impl Into<String>,
71        dim: u32,
72        bound_min: Vec3,
73        bound_max: Vec3,
74    ) -> Self {
75        Self::new(name, UVec3::splat(dim), bound_min, bound_max)
76    }
77
78    /// Returns the number of nodes in each dimension.
79    #[must_use]
80    pub fn node_dim(&self) -> UVec3 {
81        self.node_dim
82    }
83
84    /// Returns the number of cells in each dimension.
85    #[must_use]
86    pub fn cell_dim(&self) -> UVec3 {
87        self.node_dim.saturating_sub(UVec3::ONE)
88    }
89
90    /// Returns the total number of nodes.
91    #[must_use]
92    pub fn num_nodes(&self) -> u64 {
93        u64::from(self.node_dim.x) * u64::from(self.node_dim.y) * u64::from(self.node_dim.z)
94    }
95
96    /// Returns the total number of cells.
97    #[must_use]
98    pub fn num_cells(&self) -> u64 {
99        let cell_dim = self.cell_dim();
100        u64::from(cell_dim.x) * u64::from(cell_dim.y) * u64::from(cell_dim.z)
101    }
102
103    /// Returns the minimum bound.
104    #[must_use]
105    pub fn bound_min(&self) -> Vec3 {
106        self.bound_min
107    }
108
109    /// Returns the maximum bound.
110    #[must_use]
111    pub fn bound_max(&self) -> Vec3 {
112        self.bound_max
113    }
114
115    /// Returns the grid spacing (distance between adjacent nodes).
116    #[must_use]
117    pub fn grid_spacing(&self) -> Vec3 {
118        let cell_dim = self.cell_dim().as_vec3();
119        (self.bound_max - self.bound_min) / cell_dim.max(Vec3::ONE)
120    }
121
122    /// Flattens a 3D node index to a linear index.
123    #[must_use]
124    pub fn flatten_node_index(&self, i: u32, j: u32, k: u32) -> u64 {
125        u64::from(i)
126            + (u64::from(j) * u64::from(self.node_dim.x))
127            + (u64::from(k) * u64::from(self.node_dim.x) * u64::from(self.node_dim.y))
128    }
129
130    /// Unflattens a linear node index to a 3D index.
131    #[must_use]
132    pub fn unflatten_node_index(&self, idx: u64) -> UVec3 {
133        let x = idx % u64::from(self.node_dim.x);
134        let y = (idx / u64::from(self.node_dim.x)) % u64::from(self.node_dim.y);
135        let z = idx / (u64::from(self.node_dim.x) * u64::from(self.node_dim.y));
136        UVec3::new(x as u32, y as u32, z as u32)
137    }
138
139    /// Returns the world position of a node at the given 3D index.
140    #[must_use]
141    pub fn position_of_node(&self, i: u32, j: u32, k: u32) -> Vec3 {
142        let cell_dim = self.cell_dim().as_vec3().max(Vec3::ONE);
143        let t = Vec3::new(i as f32, j as f32, k as f32) / cell_dim;
144        self.bound_min + t * (self.bound_max - self.bound_min)
145    }
146
147    /// Gets the grid color.
148    #[must_use]
149    pub fn color(&self) -> Vec4 {
150        self.color
151    }
152
153    /// Sets the grid color.
154    pub fn set_color(&mut self, color: Vec3) -> &mut Self {
155        self.color = color.extend(1.0);
156        self
157    }
158
159    /// Gets the edge color.
160    #[must_use]
161    pub fn edge_color(&self) -> Vec4 {
162        self.edge_color
163    }
164
165    /// Sets the edge color.
166    pub fn set_edge_color(&mut self, color: Vec3) -> &mut Self {
167        self.edge_color = color.extend(1.0);
168        self
169    }
170
171    /// Gets the edge width.
172    #[must_use]
173    pub fn edge_width(&self) -> f32 {
174        self.edge_width
175    }
176
177    /// Sets the edge width.
178    pub fn set_edge_width(&mut self, width: f32) -> &mut Self {
179        self.edge_width = width;
180        self
181    }
182
183    /// Gets the cube size factor.
184    #[must_use]
185    pub fn cube_size_factor(&self) -> f32 {
186        self.cube_size_factor
187    }
188
189    /// Sets the cube size factor (0 = no cubes, 1 = full size).
190    pub fn set_cube_size_factor(&mut self, factor: f32) -> &mut Self {
191        self.cube_size_factor = factor.clamp(0.0, 1.0);
192        self
193    }
194
195    /// Returns a mutable reference to the quantities list.
196    pub fn quantities_mut(&mut self) -> &mut [Box<dyn Quantity>] {
197        &mut self.quantities
198    }
199
200    /// Generates the bounding box wireframe geometry.
201    fn generate_bbox_wireframe(&self) -> (Vec<Vec3>, Vec<[u32; 2]>) {
202        let min = self.bound_min;
203        let max = self.bound_max;
204
205        // 8 corners of the bounding box
206        let nodes = vec![
207            Vec3::new(min.x, min.y, min.z), // 0
208            Vec3::new(max.x, min.y, min.z), // 1
209            Vec3::new(max.x, max.y, min.z), // 2
210            Vec3::new(min.x, max.y, min.z), // 3
211            Vec3::new(min.x, min.y, max.z), // 4
212            Vec3::new(max.x, min.y, max.z), // 5
213            Vec3::new(max.x, max.y, max.z), // 6
214            Vec3::new(min.x, max.y, max.z), // 7
215        ];
216
217        // 12 edges of the bounding box
218        let edges = vec![
219            // Bottom face
220            [0, 1],
221            [1, 2],
222            [2, 3],
223            [3, 0],
224            // Top face
225            [4, 5],
226            [5, 6],
227            [6, 7],
228            [7, 4],
229            // Vertical edges
230            [0, 4],
231            [1, 5],
232            [2, 6],
233            [3, 7],
234        ];
235
236        (nodes, edges)
237    }
238
239    /// Initializes GPU render data.
240    pub fn init_render_data(
241        &mut self,
242        device: &wgpu::Device,
243        bind_group_layout: &wgpu::BindGroupLayout,
244        camera_buffer: &wgpu::Buffer,
245        queue: &wgpu::Queue,
246    ) {
247        let (nodes, edges) = self.generate_bbox_wireframe();
248
249        let edge_tail_inds: Vec<u32> = edges.iter().map(|e| e[0]).collect();
250        let edge_tip_inds: Vec<u32> = edges.iter().map(|e| e[1]).collect();
251
252        let render_data = CurveNetworkRenderData::new(
253            device,
254            bind_group_layout,
255            camera_buffer,
256            &nodes,
257            &edge_tail_inds,
258            &edge_tip_inds,
259        );
260
261        // Update uniforms with edge color and width
262        let uniforms = polyscope_render::CurveNetworkUniforms {
263            color: self.edge_color.to_array(),
264            radius: self.edge_width * 0.002,
265            radius_is_relative: 1,
266            render_mode: 0,
267            _padding: 0.0,
268        };
269        render_data.update_uniforms(queue, &uniforms);
270
271        self.render_data = Some(render_data);
272    }
273
274    /// Returns the render data if available.
275    #[must_use]
276    pub fn render_data(&self) -> Option<&CurveNetworkRenderData> {
277        self.render_data.as_ref()
278    }
279
280    /// Updates GPU buffers per-frame (wireframe uniforms: edge color, edge width).
281    pub fn update_gpu_buffers(&self, queue: &wgpu::Queue) {
282        if let Some(render_data) = &self.render_data {
283            let uniforms = polyscope_render::CurveNetworkUniforms {
284                color: self.edge_color.to_array(),
285                radius: self.edge_width * 0.002,
286                radius_is_relative: 1,
287                render_mode: 0,
288                _padding: 0.0,
289            };
290            render_data.update_uniforms(queue, &uniforms);
291        }
292    }
293
294    /// Adds a node scalar quantity to the grid.
295    pub fn add_node_scalar_quantity(
296        &mut self,
297        name: impl Into<String>,
298        values: Vec<f32>,
299    ) -> &mut Self {
300        let quantity = VolumeGridNodeScalarQuantity::new(
301            name,
302            self.name.clone(),
303            values,
304            self.node_dim,
305            self.bound_min,
306            self.bound_max,
307        );
308        self.add_quantity(Box::new(quantity));
309        self
310    }
311
312    /// Adds a cell scalar quantity to the grid.
313    pub fn add_cell_scalar_quantity(
314        &mut self,
315        name: impl Into<String>,
316        values: Vec<f32>,
317    ) -> &mut Self {
318        let quantity = VolumeGridCellScalarQuantity::new(
319            name,
320            self.name.clone(),
321            values,
322            self.cell_dim(),
323            self.bound_min,
324            self.bound_max,
325        );
326        self.add_quantity(Box::new(quantity));
327        self
328    }
329
330    /// Builds the egui UI for this volume grid.
331    ///
332    /// # Arguments
333    /// * `colormap_names` - Available colormap names for quantity visualization.
334    ///   If empty, defaults to built-in names.
335    pub fn build_egui_ui(&mut self, ui: &mut egui::Ui, colormap_names: &[&str]) {
336        let default_names = ["viridis", "blues", "reds", "coolwarm", "rainbow"];
337        let names = if colormap_names.is_empty() {
338            &default_names[..]
339        } else {
340            colormap_names
341        };
342        // Grid info
343        ui.label(format!(
344            "{}x{}x{} grid ({} nodes)",
345            self.node_dim.x,
346            self.node_dim.y,
347            self.node_dim.z,
348            self.num_nodes()
349        ));
350
351        // Color picker
352        ui.horizontal(|ui| {
353            ui.label("Edge Color:");
354            let mut color = [self.edge_color.x, self.edge_color.y, self.edge_color.z];
355            if ui.color_edit_button_rgb(&mut color).changed() {
356                self.edge_color = Vec4::new(color[0], color[1], color[2], self.edge_color.w);
357            }
358        });
359
360        // Show quantities
361        if !self.quantities.is_empty() {
362            ui.separator();
363            ui.label("Quantities:");
364            for quantity in &mut self.quantities {
365                if let Some(sq) = quantity
366                    .as_any_mut()
367                    .downcast_mut::<VolumeGridNodeScalarQuantity>()
368                {
369                    sq.build_egui_ui(ui, names);
370                } else if let Some(sq) = quantity
371                    .as_any_mut()
372                    .downcast_mut::<VolumeGridCellScalarQuantity>()
373                {
374                    sq.build_egui_ui(ui, names);
375                }
376            }
377        }
378    }
379}
380
381impl Structure for VolumeGrid {
382    fn as_any(&self) -> &dyn std::any::Any {
383        self
384    }
385
386    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
387        self
388    }
389
390    fn name(&self) -> &str {
391        &self.name
392    }
393
394    fn type_name(&self) -> &'static str {
395        "VolumeGrid"
396    }
397
398    fn bounding_box(&self) -> Option<(Vec3, Vec3)> {
399        Some((self.bound_min, self.bound_max))
400    }
401
402    fn length_scale(&self) -> f32 {
403        (self.bound_max - self.bound_min).length()
404    }
405
406    fn transform(&self) -> Mat4 {
407        self.transform
408    }
409
410    fn set_transform(&mut self, transform: Mat4) {
411        self.transform = transform;
412    }
413
414    fn is_enabled(&self) -> bool {
415        self.enabled
416    }
417
418    fn set_enabled(&mut self, enabled: bool) {
419        self.enabled = enabled;
420    }
421
422    fn draw(&self, _ctx: &mut dyn RenderContext) {
423        // Drawing is handled externally
424    }
425
426    fn draw_pick(&self, _ctx: &mut dyn RenderContext) {
427        // Picking not implemented
428    }
429
430    fn build_ui(&mut self, _ui: &dyn std::any::Any) {
431        // UI is built via build_egui_ui
432    }
433
434    fn build_pick_ui(&self, _ui: &dyn std::any::Any, _pick: &PickResult) {
435        // Pick UI not implemented
436    }
437
438    fn clear_gpu_resources(&mut self) {
439        self.render_data = None;
440        for quantity in &mut self.quantities {
441            quantity.clear_gpu_resources();
442        }
443    }
444
445    fn refresh(&mut self) {
446        self.render_data = None;
447        for quantity in &mut self.quantities {
448            quantity.refresh();
449        }
450    }
451}
452
453impl HasQuantities for VolumeGrid {
454    fn add_quantity(&mut self, quantity: Box<dyn Quantity>) {
455        self.quantities.push(quantity);
456    }
457
458    fn get_quantity(&self, name: &str) -> Option<&dyn Quantity> {
459        self.quantities
460            .iter()
461            .find(|q| q.name() == name)
462            .map(std::convert::AsRef::as_ref)
463    }
464
465    fn get_quantity_mut(&mut self, name: &str) -> Option<&mut Box<dyn Quantity>> {
466        self.quantities.iter_mut().find(|q| q.name() == name)
467    }
468
469    fn remove_quantity(&mut self, name: &str) -> Option<Box<dyn Quantity>> {
470        let idx = self.quantities.iter().position(|q| q.name() == name)?;
471        Some(self.quantities.remove(idx))
472    }
473
474    fn quantities(&self) -> &[Box<dyn Quantity>] {
475        &self.quantities
476    }
477}
478
479#[cfg(test)]
480mod tests {
481    use super::*;
482
483    #[test]
484    fn test_volume_grid_creation() {
485        let grid = VolumeGrid::new("test", UVec3::new(10, 20, 30), Vec3::ZERO, Vec3::ONE);
486        assert_eq!(grid.node_dim(), UVec3::new(10, 20, 30));
487        assert_eq!(grid.cell_dim(), UVec3::new(9, 19, 29));
488        assert_eq!(grid.num_nodes(), 10 * 20 * 30);
489        assert_eq!(grid.num_cells(), 9 * 19 * 29);
490    }
491
492    #[test]
493    fn test_index_conversion() {
494        let grid = VolumeGrid::new("test", UVec3::new(5, 6, 7), Vec3::ZERO, Vec3::ONE);
495
496        let idx = grid.flatten_node_index(2, 3, 4);
497        let uvec = grid.unflatten_node_index(idx);
498        assert_eq!(uvec, UVec3::new(2, 3, 4));
499    }
500
501    #[test]
502    fn test_node_position() {
503        let grid = VolumeGrid::new(
504            "test",
505            UVec3::new(3, 3, 3),
506            Vec3::ZERO,
507            Vec3::new(2.0, 2.0, 2.0),
508        );
509
510        let p = grid.position_of_node(0, 0, 0);
511        assert!((p - Vec3::ZERO).length() < 1e-6);
512
513        let p = grid.position_of_node(2, 2, 2);
514        assert!((p - Vec3::new(2.0, 2.0, 2.0)).length() < 1e-6);
515
516        let p = grid.position_of_node(1, 1, 1);
517        assert!((p - Vec3::ONE).length() < 1e-6);
518    }
519}
polyscope_structures/volume_grid/mod.rs

polyscope_structures/volume_grid/
mod.rs
