polyscope_structures/volume_grid/
scalar_quantity.rs

1//! Scalar quantities for volume grids.
2
3use glam::{UVec3, Vec3};
4use polyscope_core::quantity::{Quantity, QuantityKind};
5use polyscope_core::{McmMesh, marching_cubes};
6use polyscope_render::{GridcubePickUniforms, GridcubeRenderData, IsosurfaceRenderData};
7use wgpu::util::DeviceExt;
8
9/// Visualization mode for volume grid scalar quantities.
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum VolumeGridVizMode {
12    /// Colored cubes at each grid node/cell.
13    Gridcube,
14    /// Isosurface extracted via marching cubes (node scalars only).
15    Isosurface,
16}
17
18/// A scalar quantity defined at grid nodes.
19pub struct VolumeGridNodeScalarQuantity {
20    name: String,
21    structure_name: String,
22    values: Vec<f32>,
23    node_dim: UVec3,
24    enabled: bool,
25
26    // Visualization parameters
27    color_map: String,
28    data_min: f32,
29    data_max: f32,
30
31    // Visualization mode
32    viz_mode: VolumeGridVizMode,
33
34    // Gridcube state
35    gridcube_render_data: Option<GridcubeRenderData>,
36    gridcube_dirty: bool,
37
38    // Isosurface state
39    isosurface_level: f32,
40    isosurface_color: Vec3,
41    isosurface_render_data: Option<IsosurfaceRenderData>,
42    isosurface_mesh_cache: Option<McmMesh>,
43    isosurface_dirty: bool,
44
45    // Grid geometry (needed for MC coordinate transform)
46    bound_min: Vec3,
47    bound_max: Vec3,
48
49    // Flag: user clicked "Register as Surface Mesh"
50    register_as_mesh_requested: bool,
51
52    // Pick state
53    pick_uniform_buffer: Option<wgpu::Buffer>,
54    pick_bind_group: Option<wgpu::BindGroup>,
55    global_start: u32,
56}
57
58impl VolumeGridNodeScalarQuantity {
59    /// Creates a new node scalar quantity.
60    pub fn new(
61        name: impl Into<String>,
62        structure_name: impl Into<String>,
63        values: Vec<f32>,
64        node_dim: UVec3,
65        bound_min: Vec3,
66        bound_max: Vec3,
67    ) -> Self {
68        let (data_min, data_max) = Self::compute_range(&values);
69        let isosurface_level = (data_min + data_max) * 0.5;
70        Self {
71            name: name.into(),
72            structure_name: structure_name.into(),
73            values,
74            node_dim,
75            enabled: false,
76            color_map: "viridis".to_string(),
77            data_min,
78            data_max,
79            viz_mode: VolumeGridVizMode::Gridcube,
80            gridcube_render_data: None,
81            gridcube_dirty: true,
82            isosurface_level,
83            isosurface_color: Vec3::new(0.047, 0.451, 0.690), // default blue
84            isosurface_render_data: None,
85            isosurface_mesh_cache: None,
86            isosurface_dirty: true,
87            bound_min,
88            bound_max,
89            register_as_mesh_requested: false,
90            pick_uniform_buffer: None,
91            pick_bind_group: None,
92            global_start: 0,
93        }
94    }
95
96    fn compute_range(values: &[f32]) -> (f32, f32) {
97        let mut min = f32::MAX;
98        let mut max = f32::MIN;
99        for &v in values {
100            if v.is_finite() {
101                min = min.min(v);
102                max = max.max(v);
103            }
104        }
105        if min > max { (0.0, 1.0) } else { (min, max) }
106    }
107
108    /// Returns the values.
109    #[must_use]
110    pub fn values(&self) -> &[f32] {
111        &self.values
112    }
113
114    /// Returns the grid node dimensions.
115    #[must_use]
116    pub fn node_dim(&self) -> UVec3 {
117        self.node_dim
118    }
119
120    /// Gets the value at a 3D index.
121    #[must_use]
122    pub fn get(&self, i: u32, j: u32, k: u32) -> f32 {
123        let idx = i as usize
124            + j as usize * self.node_dim.x as usize
125            + k as usize * self.node_dim.x as usize * self.node_dim.y as usize;
126        self.values.get(idx).copied().unwrap_or(0.0)
127    }
128
129    /// Gets the color map name.
130    #[must_use]
131    pub fn color_map(&self) -> &str {
132        &self.color_map
133    }
134
135    /// Sets the color map name.
136    pub fn set_color_map(&mut self, name: impl Into<String>) -> &mut Self {
137        self.color_map = name.into();
138        self.gridcube_dirty = true;
139        self
140    }
141
142    /// Gets the data range.
143    #[must_use]
144    pub fn data_range(&self) -> (f32, f32) {
145        (self.data_min, self.data_max)
146    }
147
148    /// Sets the data range.
149    pub fn set_data_range(&mut self, min: f32, max: f32) -> &mut Self {
150        self.data_min = min;
151        self.data_max = max;
152        self
153    }
154
155    // --- Visualization mode ---
156
157    /// Gets the current visualization mode.
158    #[must_use]
159    pub fn viz_mode(&self) -> VolumeGridVizMode {
160        self.viz_mode
161    }
162
163    /// Sets the visualization mode.
164    pub fn set_viz_mode(&mut self, mode: VolumeGridVizMode) -> &mut Self {
165        self.viz_mode = mode;
166        self
167    }
168
169    // --- Isosurface ---
170
171    /// Gets the isosurface level.
172    #[must_use]
173    pub fn isosurface_level(&self) -> f32 {
174        self.isosurface_level
175    }
176
177    /// Sets the isosurface level (invalidates cache).
178    pub fn set_isosurface_level(&mut self, level: f32) -> &mut Self {
179        self.isosurface_level = level;
180        self.isosurface_dirty = true;
181        // Keep old mesh cache and render data — replaced atomically in init phase
182        self
183    }
184
185    /// Gets the isosurface color.
186    #[must_use]
187    pub fn isosurface_color(&self) -> Vec3 {
188        self.isosurface_color
189    }
190
191    /// Sets the isosurface color.
192    pub fn set_isosurface_color(&mut self, color: Vec3) -> &mut Self {
193        self.isosurface_color = color;
194        self
195    }
196
197    /// Returns whether the isosurface needs re-extraction.
198    #[must_use]
199    pub fn isosurface_dirty(&self) -> bool {
200        self.isosurface_dirty
201    }
202
203    /// Returns whether the gridcube needs GPU re-init.
204    #[must_use]
205    pub fn gridcube_dirty(&self) -> bool {
206        self.gridcube_dirty
207    }
208
209    /// Extracts the isosurface mesh using marching cubes.
210    ///
211    /// MC output vertices are in grid index space: vertex (i,j,k) has coords
212    /// that need swizzle(z,y,x) * `grid_spacing` + `bound_min` to transform to world space.
213    pub fn extract_isosurface(&mut self) -> &McmMesh {
214        if self.isosurface_mesh_cache.is_none() || self.isosurface_dirty {
215            let nx = self.node_dim.x;
216            let ny = self.node_dim.y;
217            let nz = self.node_dim.z;
218
219            let mut mesh = marching_cubes(&self.values, self.isosurface_level, nx, ny, nz);
220
221            // Transform from MC index space to world space
222            // MC uses indexing (i * ny + j) * nz + k, output coords are in (i,j,k) space
223            // Need to map: x_world = x_mc * spacing_z + bound_min.z (swizzle z,y,x)
224            let cell_dim = Vec3::new(
225                (nx - 1).max(1) as f32,
226                (ny - 1).max(1) as f32,
227                (nz - 1).max(1) as f32,
228            );
229            let spacing = (self.bound_max - self.bound_min) / cell_dim;
230
231            for v in &mut mesh.vertices {
232                // MC output: v.x is in i-dimension, v.y in j-dimension, v.z in k-dimension
233                // Grid layout: i maps to x, j maps to y, k maps to z (no swizzle needed
234                // since our MC uses same indexing as the grid)
235                *v = Vec3::new(
236                    v.x * spacing.x + self.bound_min.x,
237                    v.y * spacing.y + self.bound_min.y,
238                    v.z * spacing.z + self.bound_min.z,
239                );
240            }
241
242            // Transform normals (only need to scale, then renormalize)
243            for n in &mut mesh.normals {
244                // Scale normals by inverse spacing to account for non-uniform grid
245                *n = Vec3::new(n.x / spacing.x, n.y / spacing.y, n.z / spacing.z);
246                let len = n.length();
247                if len > 0.0 {
248                    *n /= len;
249                }
250            }
251
252            self.isosurface_mesh_cache = Some(mesh);
253            self.isosurface_dirty = false;
254        }
255        self.isosurface_mesh_cache.as_ref().unwrap()
256    }
257
258    /// Returns the cached isosurface mesh, if available.
259    #[must_use]
260    pub fn isosurface_mesh(&self) -> Option<&McmMesh> {
261        self.isosurface_mesh_cache.as_ref()
262    }
263
264    // --- GPU resources ---
265
266    /// Returns the gridcube render data.
267    #[must_use]
268    pub fn gridcube_render_data(&self) -> Option<&GridcubeRenderData> {
269        self.gridcube_render_data.as_ref()
270    }
271
272    /// Returns a mutable reference to the gridcube render data.
273    pub fn gridcube_render_data_mut(&mut self) -> Option<&mut GridcubeRenderData> {
274        self.gridcube_render_data.as_mut()
275    }
276
277    /// Sets the gridcube render data.
278    pub fn set_gridcube_render_data(&mut self, data: GridcubeRenderData) {
279        self.gridcube_render_data = Some(data);
280        self.gridcube_dirty = false;
281    }
282
283    /// Returns the isosurface render data.
284    #[must_use]
285    pub fn isosurface_render_data(&self) -> Option<&IsosurfaceRenderData> {
286        self.isosurface_render_data.as_ref()
287    }
288
289    /// Returns a mutable reference to the isosurface render data.
290    pub fn isosurface_render_data_mut(&mut self) -> Option<&mut IsosurfaceRenderData> {
291        self.isosurface_render_data.as_mut()
292    }
293
294    /// Sets the isosurface render data.
295    pub fn set_isosurface_render_data(&mut self, data: IsosurfaceRenderData) {
296        self.isosurface_render_data = Some(data);
297        self.isosurface_dirty = false;
298    }
299
300    /// Clears the isosurface render data (e.g. when isovalue yields empty mesh).
301    pub fn clear_isosurface_render_data(&mut self) {
302        self.isosurface_render_data = None;
303        self.isosurface_dirty = false;
304    }
305
306    // --- Pick resources ---
307
308    /// Initializes pick resources for this quantity.
309    ///
310    /// Requires that `gridcube_render_data` is already initialized (needs the position buffer).
311    pub fn init_pick_resources(
312        &mut self,
313        device: &wgpu::Device,
314        pick_bind_group_layout: &wgpu::BindGroupLayout,
315        camera_buffer: &wgpu::Buffer,
316        global_start: u32,
317    ) {
318        self.global_start = global_start;
319
320        let Some(gridcube_rd) = &self.gridcube_render_data else {
321            return;
322        };
323
324        let uniforms = GridcubePickUniforms {
325            global_start,
326            cube_size_factor: 1.0,
327            ..Default::default()
328        };
329
330        let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
331            label: Some("gridcube node pick uniforms"),
332            contents: bytemuck::cast_slice(&[uniforms]),
333            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
334        });
335
336        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
337            label: Some("gridcube node pick bind group"),
338            layout: pick_bind_group_layout,
339            entries: &[
340                wgpu::BindGroupEntry {
341                    binding: 0,
342                    resource: camera_buffer.as_entire_binding(),
343                },
344                wgpu::BindGroupEntry {
345                    binding: 1,
346                    resource: uniform_buffer.as_entire_binding(),
347                },
348                wgpu::BindGroupEntry {
349                    binding: 2,
350                    resource: gridcube_rd.position_buffer.as_entire_binding(),
351                },
352            ],
353        });
354
355        self.pick_uniform_buffer = Some(uniform_buffer);
356        self.pick_bind_group = Some(bind_group);
357    }
358
359    /// Returns the pick bind group, if initialized.
360    #[must_use]
361    pub fn pick_bind_group(&self) -> Option<&wgpu::BindGroup> {
362        self.pick_bind_group.as_ref()
363    }
364
365    /// Updates the pick uniform buffer with current model transform and cube size factor.
366    pub fn update_pick_uniforms(
367        &self,
368        queue: &wgpu::Queue,
369        model: [[f32; 4]; 4],
370        cube_size_factor: f32,
371    ) {
372        if let Some(buffer) = &self.pick_uniform_buffer {
373            let uniforms = GridcubePickUniforms {
374                model,
375                global_start: self.global_start,
376                cube_size_factor,
377                ..Default::default()
378            };
379            queue.write_buffer(buffer, 0, bytemuck::cast_slice(&[uniforms]));
380        }
381    }
382
383    /// Returns the number of pick elements (= number of gridcube instances).
384    #[must_use]
385    pub fn num_pick_elements(&self) -> u32 {
386        self.gridcube_render_data
387            .as_ref()
388            .map_or(0, |rd| rd.num_instances)
389    }
390
391    /// Returns the total vertices for the pick draw call.
392    #[must_use]
393    pub fn pick_total_vertices(&self) -> u32 {
394        self.gridcube_render_data
395            .as_ref()
396            .map_or(0, GridcubeRenderData::total_vertices)
397    }
398
399    /// Returns whether the user has requested registering the isosurface as a mesh.
400    #[must_use]
401    pub fn register_as_mesh_requested(&self) -> bool {
402        self.register_as_mesh_requested
403    }
404
405    /// Clears the register-as-mesh request flag.
406    pub fn clear_register_as_mesh_request(&mut self) {
407        self.register_as_mesh_requested = false;
408    }
409
410    /// Builds egui UI for this quantity.
411    pub fn build_egui_ui(&mut self, ui: &mut egui::Ui, colormap_names: &[&str]) {
412        ui.horizontal(|ui| {
413            let mut enabled = self.enabled;
414            if ui.checkbox(&mut enabled, "").changed() {
415                self.enabled = enabled;
416            }
417            ui.label(&self.name);
418            ui.label(format!("[{:.3}, {:.3}]", self.data_min, self.data_max));
419        });
420
421        if self.enabled {
422            let indent_id = egui::Id::new(&self.name).with("node_scalar_indent");
423            ui.indent(indent_id, |ui| {
424                // Viz mode toggle
425                ui.horizontal(|ui| {
426                    ui.label("Mode:");
427                    if ui
428                        .selectable_label(self.viz_mode == VolumeGridVizMode::Gridcube, "Gridcube")
429                        .clicked()
430                    {
431                        self.viz_mode = VolumeGridVizMode::Gridcube;
432                    }
433                    if ui
434                        .selectable_label(
435                            self.viz_mode == VolumeGridVizMode::Isosurface,
436                            "Isosurface",
437                        )
438                        .clicked()
439                    {
440                        self.viz_mode = VolumeGridVizMode::Isosurface;
441                    }
442                });
443
444                match self.viz_mode {
445                    VolumeGridVizMode::Gridcube => {
446                        self.build_gridcube_ui(ui, colormap_names);
447                    }
448                    VolumeGridVizMode::Isosurface => {
449                        self.build_isosurface_ui(ui);
450                    }
451                }
452            });
453        }
454    }
455
456    fn build_gridcube_ui(&mut self, ui: &mut egui::Ui, colormap_names: &[&str]) {
457        // Colormap selector
458        if !colormap_names.is_empty() {
459            ui.horizontal(|ui| {
460                ui.label("Colormap:");
461                egui::ComboBox::from_id_salt(format!("{}_colormap", self.name))
462                    .selected_text(&self.color_map)
463                    .show_ui(ui, |ui| {
464                        for &name in colormap_names {
465                            if ui.selectable_label(self.color_map == name, name).clicked() {
466                                self.color_map = name.to_string();
467                                self.gridcube_dirty = true;
468                            }
469                        }
470                    });
471            });
472        }
473
474        // Data range
475        ui.horizontal(|ui| {
476            ui.label("Range:");
477            let mut min = self.data_min;
478            let mut max = self.data_max;
479            let speed = (max - min).abs() * 0.01;
480            let speed = if speed > 0.0 { speed } else { 0.01 };
481            ui.add(egui::DragValue::new(&mut min).speed(speed));
482            ui.label("–");
483            ui.add(egui::DragValue::new(&mut max).speed(speed));
484            if (min - self.data_min).abs() > f32::EPSILON
485                || (max - self.data_max).abs() > f32::EPSILON
486            {
487                self.data_min = min;
488                self.data_max = max;
489            }
490        });
491    }
492
493    fn build_isosurface_ui(&mut self, ui: &mut egui::Ui) {
494        egui::Grid::new(format!("{}_iso_grid", self.name))
495            .num_columns(2)
496            .show(ui, |ui| {
497                ui.label("Color:");
498                let mut color = [
499                    self.isosurface_color.x,
500                    self.isosurface_color.y,
501                    self.isosurface_color.z,
502                ];
503                if ui.color_edit_button_rgb(&mut color).changed() {
504                    self.isosurface_color = Vec3::new(color[0], color[1], color[2]);
505                }
506                ui.end_row();
507
508                ui.label("Level:");
509                let mut level = self.isosurface_level;
510                let (range_min, range_max) = (self.data_min, self.data_max);
511                if ui
512                    .add(egui::Slider::new(&mut level, range_min..=range_max))
513                    .changed()
514                {
515                    self.isosurface_level = level;
516                    self.isosurface_dirty = true;
517                }
518                ui.end_row();
519            });
520
521        // Triangle count
522        if let Some(mesh) = &self.isosurface_mesh_cache {
523            ui.label(format!("{} tris", mesh.indices.len() / 3));
524        }
525
526        // Buttons: equal-width columns, same row
527        let has_cache = self.isosurface_mesh_cache.is_some();
528        if has_cache {
529            ui.columns(2, |cols| {
530                let w = cols[0].available_width();
531                let h = cols[0].spacing().interact_size.y;
532                if cols[0]
533                    .add_sized([w, h], egui::Button::new("Refresh"))
534                    .clicked()
535                {
536                    self.isosurface_dirty = true;
537                    self.isosurface_mesh_cache = None;
538                    self.isosurface_render_data = None;
539                }
540                if cols[1]
541                    .add_sized([w, h], egui::Button::new("Register Mesh"))
542                    .clicked()
543                {
544                    self.register_as_mesh_requested = true;
545                }
546            });
547        } else if ui.button("Refresh").clicked() {
548            self.isosurface_dirty = true;
549            self.isosurface_mesh_cache = None;
550            self.isosurface_render_data = None;
551        }
552    }
553}
554
555impl Quantity for VolumeGridNodeScalarQuantity {
556    fn name(&self) -> &str {
557        &self.name
558    }
559
560    fn structure_name(&self) -> &str {
561        &self.structure_name
562    }
563
564    fn kind(&self) -> QuantityKind {
565        QuantityKind::Scalar
566    }
567
568    fn is_enabled(&self) -> bool {
569        self.enabled
570    }
571
572    fn set_enabled(&mut self, enabled: bool) {
573        self.enabled = enabled;
574    }
575
576    fn data_size(&self) -> usize {
577        self.values.len()
578    }
579
580    fn build_ui(&mut self, _ui: &dyn std::any::Any) {
581        // UI is built via build_egui_ui
582    }
583
584    fn refresh(&mut self) {
585        self.gridcube_render_data = None;
586        self.gridcube_dirty = true;
587        self.isosurface_render_data = None;
588        self.isosurface_mesh_cache = None;
589        self.isosurface_dirty = true;
590        self.pick_uniform_buffer = None;
591        self.pick_bind_group = None;
592    }
593
594    fn clear_gpu_resources(&mut self) {
595        self.gridcube_render_data = None;
596        self.isosurface_render_data = None;
597    }
598
599    fn as_any(&self) -> &dyn std::any::Any {
600        self
601    }
602
603    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
604        self
605    }
606}
607
608/// A scalar quantity defined at grid cells.
609pub struct VolumeGridCellScalarQuantity {
610    name: String,
611    structure_name: String,
612    values: Vec<f32>,
613    cell_dim: UVec3,
614    enabled: bool,
615
616    // Visualization parameters
617    color_map: String,
618    data_min: f32,
619    data_max: f32,
620
621    // Gridcube state (cell scalars only support gridcube, not isosurface)
622    gridcube_render_data: Option<GridcubeRenderData>,
623    gridcube_dirty: bool,
624
625    // Grid geometry (reserved for future gridcube world-space mapping)
626    #[allow(dead_code)]
627    bound_min: Vec3,
628    #[allow(dead_code)]
629    bound_max: Vec3,
630
631    // Pick state
632    pick_uniform_buffer: Option<wgpu::Buffer>,
633    pick_bind_group: Option<wgpu::BindGroup>,
634    global_start: u32,
635}
636
637impl VolumeGridCellScalarQuantity {
638    /// Creates a new cell scalar quantity.
639    pub fn new(
640        name: impl Into<String>,
641        structure_name: impl Into<String>,
642        values: Vec<f32>,
643        cell_dim: UVec3,
644        bound_min: Vec3,
645        bound_max: Vec3,
646    ) -> Self {
647        let (data_min, data_max) = Self::compute_range(&values);
648        Self {
649            name: name.into(),
650            structure_name: structure_name.into(),
651            values,
652            cell_dim,
653            enabled: false,
654            color_map: "viridis".to_string(),
655            data_min,
656            data_max,
657            gridcube_render_data: None,
658            gridcube_dirty: true,
659            bound_min,
660            bound_max,
661            pick_uniform_buffer: None,
662            pick_bind_group: None,
663            global_start: 0,
664        }
665    }
666
667    fn compute_range(values: &[f32]) -> (f32, f32) {
668        let mut min = f32::MAX;
669        let mut max = f32::MIN;
670        for &v in values {
671            if v.is_finite() {
672                min = min.min(v);
673                max = max.max(v);
674            }
675        }
676        if min > max { (0.0, 1.0) } else { (min, max) }
677    }
678
679    /// Returns the values.
680    #[must_use]
681    pub fn values(&self) -> &[f32] {
682        &self.values
683    }
684
685    /// Returns the grid cell dimensions.
686    #[must_use]
687    pub fn cell_dim(&self) -> UVec3 {
688        self.cell_dim
689    }
690
691    /// Gets the value at a 3D index.
692    #[must_use]
693    pub fn get(&self, i: u32, j: u32, k: u32) -> f32 {
694        let idx = i as usize
695            + j as usize * self.cell_dim.x as usize
696            + k as usize * self.cell_dim.x as usize * self.cell_dim.y as usize;
697        self.values.get(idx).copied().unwrap_or(0.0)
698    }
699
700    /// Gets the color map name.
701    #[must_use]
702    pub fn color_map(&self) -> &str {
703        &self.color_map
704    }
705
706    /// Sets the color map name.
707    pub fn set_color_map(&mut self, name: impl Into<String>) -> &mut Self {
708        self.color_map = name.into();
709        self.gridcube_dirty = true;
710        self
711    }
712
713    /// Gets the data range.
714    #[must_use]
715    pub fn data_range(&self) -> (f32, f32) {
716        (self.data_min, self.data_max)
717    }
718
719    /// Sets the data range.
720    pub fn set_data_range(&mut self, min: f32, max: f32) -> &mut Self {
721        self.data_min = min;
722        self.data_max = max;
723        self
724    }
725
726    /// Returns whether the gridcube needs GPU re-init.
727    #[must_use]
728    pub fn gridcube_dirty(&self) -> bool {
729        self.gridcube_dirty
730    }
731
732    /// Returns the gridcube render data.
733    #[must_use]
734    pub fn gridcube_render_data(&self) -> Option<&GridcubeRenderData> {
735        self.gridcube_render_data.as_ref()
736    }
737
738    /// Returns a mutable reference to the gridcube render data.
739    pub fn gridcube_render_data_mut(&mut self) -> Option<&mut GridcubeRenderData> {
740        self.gridcube_render_data.as_mut()
741    }
742
743    /// Sets the gridcube render data.
744    pub fn set_gridcube_render_data(&mut self, data: GridcubeRenderData) {
745        self.gridcube_render_data = Some(data);
746        self.gridcube_dirty = false;
747    }
748
749    // --- Pick resources ---
750
751    /// Initializes pick resources for this cell scalar quantity.
752    pub fn init_pick_resources(
753        &mut self,
754        device: &wgpu::Device,
755        pick_bind_group_layout: &wgpu::BindGroupLayout,
756        camera_buffer: &wgpu::Buffer,
757        global_start: u32,
758    ) {
759        self.global_start = global_start;
760
761        let Some(gridcube_rd) = &self.gridcube_render_data else {
762            return;
763        };
764
765        let uniforms = GridcubePickUniforms {
766            global_start,
767            cube_size_factor: 1.0,
768            ..Default::default()
769        };
770
771        let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
772            label: Some("gridcube cell pick uniforms"),
773            contents: bytemuck::cast_slice(&[uniforms]),
774            usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
775        });
776
777        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
778            label: Some("gridcube cell pick bind group"),
779            layout: pick_bind_group_layout,
780            entries: &[
781                wgpu::BindGroupEntry {
782                    binding: 0,
783                    resource: camera_buffer.as_entire_binding(),
784                },
785                wgpu::BindGroupEntry {
786                    binding: 1,
787                    resource: uniform_buffer.as_entire_binding(),
788                },
789                wgpu::BindGroupEntry {
790                    binding: 2,
791                    resource: gridcube_rd.position_buffer.as_entire_binding(),
792                },
793            ],
794        });
795
796        self.pick_uniform_buffer = Some(uniform_buffer);
797        self.pick_bind_group = Some(bind_group);
798    }
799
800    /// Returns the pick bind group, if initialized.
801    #[must_use]
802    pub fn pick_bind_group(&self) -> Option<&wgpu::BindGroup> {
803        self.pick_bind_group.as_ref()
804    }
805
806    /// Updates the pick uniform buffer with current model transform and cube size factor.
807    pub fn update_pick_uniforms(
808        &self,
809        queue: &wgpu::Queue,
810        model: [[f32; 4]; 4],
811        cube_size_factor: f32,
812    ) {
813        if let Some(buffer) = &self.pick_uniform_buffer {
814            let uniforms = GridcubePickUniforms {
815                model,
816                global_start: self.global_start,
817                cube_size_factor,
818                ..Default::default()
819            };
820            queue.write_buffer(buffer, 0, bytemuck::cast_slice(&[uniforms]));
821        }
822    }
823
824    /// Returns the number of pick elements (= number of gridcube instances).
825    #[must_use]
826    pub fn num_pick_elements(&self) -> u32 {
827        self.gridcube_render_data
828            .as_ref()
829            .map_or(0, |rd| rd.num_instances)
830    }
831
832    /// Returns the total vertices for the pick draw call.
833    #[must_use]
834    pub fn pick_total_vertices(&self) -> u32 {
835        self.gridcube_render_data
836            .as_ref()
837            .map_or(0, GridcubeRenderData::total_vertices)
838    }
839
840    /// Builds egui UI for this quantity.
841    pub fn build_egui_ui(&mut self, ui: &mut egui::Ui, colormap_names: &[&str]) {
842        ui.horizontal(|ui| {
843            let mut enabled = self.enabled;
844            if ui.checkbox(&mut enabled, "").changed() {
845                self.enabled = enabled;
846            }
847            ui.label(&self.name);
848            ui.label(format!("[{:.3}, {:.3}]", self.data_min, self.data_max));
849        });
850
851        if self.enabled {
852            let indent_id = egui::Id::new(&self.name).with("cell_scalar_indent");
853            ui.indent(indent_id, |ui| {
854                // Colormap selector
855                if !colormap_names.is_empty() {
856                    ui.horizontal(|ui| {
857                        ui.label("Colormap:");
858                        egui::ComboBox::from_id_salt(format!("{}_colormap", self.name))
859                            .selected_text(&self.color_map)
860                            .show_ui(ui, |ui| {
861                                for &name in colormap_names {
862                                    if ui.selectable_label(self.color_map == name, name).clicked() {
863                                        self.color_map = name.to_string();
864                                        self.gridcube_dirty = true;
865                                    }
866                                }
867                            });
868                    });
869                }
870
871                // Data range
872                ui.horizontal(|ui| {
873                    ui.label("Range:");
874                    let mut min = self.data_min;
875                    let mut max = self.data_max;
876                    let speed = (max - min).abs() * 0.01;
877                    let speed = if speed > 0.0 { speed } else { 0.01 };
878                    ui.add(egui::DragValue::new(&mut min).speed(speed));
879                    ui.label("–");
880                    ui.add(egui::DragValue::new(&mut max).speed(speed));
881                    if (min - self.data_min).abs() > f32::EPSILON
882                        || (max - self.data_max).abs() > f32::EPSILON
883                    {
884                        self.data_min = min;
885                        self.data_max = max;
886                    }
887                });
888            });
889        }
890    }
891}
892
893impl Quantity for VolumeGridCellScalarQuantity {
894    fn name(&self) -> &str {
895        &self.name
896    }
897
898    fn structure_name(&self) -> &str {
899        &self.structure_name
900    }
901
902    fn kind(&self) -> QuantityKind {
903        QuantityKind::Scalar
904    }
905
906    fn is_enabled(&self) -> bool {
907        self.enabled
908    }
909
910    fn set_enabled(&mut self, enabled: bool) {
911        self.enabled = enabled;
912    }
913
914    fn data_size(&self) -> usize {
915        self.values.len()
916    }
917
918    fn build_ui(&mut self, _ui: &dyn std::any::Any) {
919        // UI is built via build_egui_ui
920    }
921
922    fn refresh(&mut self) {
923        self.gridcube_render_data = None;
924        self.gridcube_dirty = true;
925        self.pick_uniform_buffer = None;
926        self.pick_bind_group = None;
927    }
928
929    fn clear_gpu_resources(&mut self) {
930        self.gridcube_render_data = None;
931    }
932
933    fn as_any(&self) -> &dyn std::any::Any {
934        self
935    }
936
937    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
938        self
939    }
940}
polyscope_structures/volume_grid/scalar_quantity.rs

polyscope_structures/volume_grid/
scalar_quantity.rs
