Skip to main content

polyscope_structures/point_cloud/
quantities.rs

1//! Point cloud quantity implementations.
2
3use glam::{Vec3, Vec4};
4use polyscope_core::quantity::{Quantity, QuantityKind, VertexQuantity};
5use polyscope_render::{ColorMap, PointCloudRenderData, VectorRenderData, VectorUniforms};
6
7/// A scalar quantity on a point cloud.
8pub struct PointCloudScalarQuantity {
9    name: String,
10    structure_name: String,
11    values: Vec<f32>,
12    enabled: bool,
13    colormap_name: String,
14    range_min: f32,
15    range_max: f32,
16}
17
18impl PointCloudScalarQuantity {
19    /// Creates a new scalar quantity.
20    pub fn new(
21        name: impl Into<String>,
22        structure_name: impl Into<String>,
23        values: Vec<f32>,
24    ) -> Self {
25        let min = values.iter().copied().fold(f32::INFINITY, f32::min);
26        let max = values.iter().copied().fold(f32::NEG_INFINITY, f32::max);
27
28        Self {
29            name: name.into(),
30            structure_name: structure_name.into(),
31            values,
32            enabled: false,
33            colormap_name: "viridis".to_string(),
34            range_min: min,
35            range_max: max,
36        }
37    }
38
39    /// Returns the scalar values.
40    #[must_use]
41    pub fn values(&self) -> &[f32] {
42        &self.values
43    }
44
45    /// Maps scalar values to colors using the colormap.
46    #[must_use]
47    pub fn compute_colors(&self, colormap: &ColorMap) -> Vec<Vec4> {
48        let range = self.range_max - self.range_min;
49        let range = if range.abs() < 1e-10 { 1.0 } else { range };
50
51        self.values
52            .iter()
53            .map(|&v| {
54                let t = (v - self.range_min) / range;
55                colormap.sample(t).extend(1.0)
56            })
57            .collect()
58    }
59
60    /// Gets the colormap name.
61    #[must_use]
62    pub fn colormap_name(&self) -> &str {
63        &self.colormap_name
64    }
65
66    /// Sets the colormap name.
67    pub fn set_colormap(&mut self, name: impl Into<String>) {
68        self.colormap_name = name.into();
69    }
70
71    /// Gets the range minimum.
72    #[must_use]
73    pub fn range_min(&self) -> f32 {
74        self.range_min
75    }
76
77    /// Gets the range maximum.
78    #[must_use]
79    pub fn range_max(&self) -> f32 {
80        self.range_max
81    }
82
83    /// Sets the range.
84    pub fn set_range(&mut self, min: f32, max: f32) {
85        self.range_min = min;
86        self.range_max = max;
87    }
88
89    /// Builds the egui UI for this scalar quantity.
90    pub fn build_egui_ui(&mut self, ui: &mut egui::Ui) -> bool {
91        let colormaps = ["viridis", "blues", "reds", "coolwarm", "rainbow"];
92        polyscope_ui::build_scalar_quantity_ui(
93            ui,
94            &self.name,
95            &mut self.enabled,
96            &mut self.colormap_name,
97            &mut self.range_min,
98            &mut self.range_max,
99            &colormaps,
100        )
101    }
102}
103
104impl Quantity for PointCloudScalarQuantity {
105    fn as_any(&self) -> &dyn std::any::Any {
106        self
107    }
108
109    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
110        self
111    }
112
113    fn name(&self) -> &str {
114        &self.name
115    }
116
117    fn structure_name(&self) -> &str {
118        &self.structure_name
119    }
120
121    fn kind(&self) -> QuantityKind {
122        QuantityKind::Scalar
123    }
124
125    fn is_enabled(&self) -> bool {
126        self.enabled
127    }
128
129    fn set_enabled(&mut self, enabled: bool) {
130        self.enabled = enabled;
131    }
132
133    fn build_ui(&mut self, _ui: &dyn std::any::Any) {
134        // UI is handled by polyscope-ui/src/structure_ui.rs
135    }
136
137    fn refresh(&mut self) {
138        // GPU refresh is handled by polyscope/src/app/render.rs
139    }
140
141    fn data_size(&self) -> usize {
142        self.values.len()
143    }
144}
145
146impl VertexQuantity for PointCloudScalarQuantity {}
147
148/// A vector quantity on a point cloud.
149pub struct PointCloudVectorQuantity {
150    name: String,
151    structure_name: String,
152    vectors: Vec<Vec3>,
153    enabled: bool,
154    length_scale: f32,
155    radius: f32,
156    color: Vec4,
157    render_data: Option<VectorRenderData>,
158}
159
160impl PointCloudVectorQuantity {
161    /// Creates a new vector quantity.
162    pub fn new(
163        name: impl Into<String>,
164        structure_name: impl Into<String>,
165        vectors: Vec<Vec3>,
166    ) -> Self {
167        Self {
168            name: name.into(),
169            structure_name: structure_name.into(),
170            vectors,
171            enabled: false,
172            length_scale: 1.0,
173            radius: 0.005,
174            color: Vec4::new(0.8, 0.2, 0.2, 1.0), // Red
175            render_data: None,
176        }
177    }
178
179    /// Returns the vectors.
180    #[must_use]
181    pub fn vectors(&self) -> &[Vec3] {
182        &self.vectors
183    }
184
185    /// Initializes GPU resources for this vector quantity.
186    pub fn init_gpu_resources(
187        &mut self,
188        device: &wgpu::Device,
189        bind_group_layout: &wgpu::BindGroupLayout,
190        camera_buffer: &wgpu::Buffer,
191        base_positions: &[Vec3],
192    ) {
193        self.render_data = Some(VectorRenderData::new(
194            device,
195            bind_group_layout,
196            camera_buffer,
197            base_positions,
198            &self.vectors,
199        ));
200    }
201
202    /// Returns the render data if initialized.
203    #[must_use]
204    pub fn render_data(&self) -> Option<&VectorRenderData> {
205        self.render_data.as_ref()
206    }
207
208    /// Updates GPU uniforms with the given model transform.
209    pub fn update_uniforms(&self, queue: &wgpu::Queue, model: &glam::Mat4) {
210        if let Some(render_data) = &self.render_data {
211            let uniforms = VectorUniforms {
212                model: model.to_cols_array(),
213                length_scale: self.length_scale,
214                radius: self.radius,
215                _padding: [0.0; 2],
216                color: self.color.to_array(),
217            };
218            render_data.update_uniforms(queue, &uniforms);
219        }
220    }
221
222    /// Sets the length scale.
223    pub fn set_length_scale(&mut self, scale: f32) {
224        self.length_scale = scale;
225    }
226
227    /// Sets the radius.
228    pub fn set_radius(&mut self, radius: f32) {
229        self.radius = radius;
230    }
231
232    /// Sets the color.
233    pub fn set_color(&mut self, color: Vec3) {
234        self.color = color.extend(1.0);
235    }
236
237    /// Gets the length scale.
238    #[must_use]
239    pub fn length_scale(&self) -> f32 {
240        self.length_scale
241    }
242
243    /// Gets the radius.
244    #[must_use]
245    pub fn radius(&self) -> f32 {
246        self.radius
247    }
248
249    /// Gets the color.
250    #[must_use]
251    pub fn color(&self) -> Vec4 {
252        self.color
253    }
254
255    /// Builds the egui UI for this vector quantity.
256    pub fn build_egui_ui(&mut self, ui: &mut egui::Ui) -> bool {
257        let mut color = [self.color.x, self.color.y, self.color.z];
258        let changed = polyscope_ui::build_vector_quantity_ui(
259            ui,
260            &self.name,
261            &mut self.enabled,
262            &mut self.length_scale,
263            &mut self.radius,
264            &mut color,
265        );
266        if changed {
267            self.color = Vec4::new(color[0], color[1], color[2], self.color.w);
268        }
269        changed
270    }
271}
272
273impl Quantity for PointCloudVectorQuantity {
274    fn as_any(&self) -> &dyn std::any::Any {
275        self
276    }
277
278    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
279        self
280    }
281
282    fn name(&self) -> &str {
283        &self.name
284    }
285
286    fn structure_name(&self) -> &str {
287        &self.structure_name
288    }
289
290    fn kind(&self) -> QuantityKind {
291        QuantityKind::Vector
292    }
293
294    fn is_enabled(&self) -> bool {
295        self.enabled
296    }
297
298    fn set_enabled(&mut self, enabled: bool) {
299        self.enabled = enabled;
300    }
301
302    fn build_ui(&mut self, _ui: &dyn std::any::Any) {
303        // UI is handled by polyscope-ui/src/structure_ui.rs
304    }
305
306    fn refresh(&mut self) {
307        // GPU refresh is handled by polyscope/src/app/render.rs
308    }
309
310    fn clear_gpu_resources(&mut self) {
311        self.render_data = None;
312    }
313
314    fn data_size(&self) -> usize {
315        self.vectors.len()
316    }
317}
318
319impl VertexQuantity for PointCloudVectorQuantity {}
320
321/// A color quantity on a point cloud.
322pub struct PointCloudColorQuantity {
323    name: String,
324    structure_name: String,
325    colors: Vec<Vec4>,
326    enabled: bool,
327}
328
329impl PointCloudColorQuantity {
330    /// Creates a new color quantity.
331    pub fn new(
332        name: impl Into<String>,
333        structure_name: impl Into<String>,
334        colors: Vec<Vec3>,
335    ) -> Self {
336        Self {
337            name: name.into(),
338            structure_name: structure_name.into(),
339            colors: colors.into_iter().map(|c| c.extend(1.0)).collect(),
340            enabled: false,
341        }
342    }
343
344    /// Returns the colors.
345    #[must_use]
346    pub fn colors(&self) -> &[Vec4] {
347        &self.colors
348    }
349
350    /// Applies this color quantity to the point cloud render data.
351    pub fn apply_to_render_data(&self, queue: &wgpu::Queue, render_data: &PointCloudRenderData) {
352        render_data.update_colors(queue, &self.colors);
353    }
354
355    /// Builds the egui UI for this color quantity.
356    pub fn build_egui_ui(&mut self, ui: &mut egui::Ui) -> bool {
357        polyscope_ui::build_color_quantity_ui(ui, &self.name, &mut self.enabled, self.colors.len())
358    }
359}
360
361impl Quantity for PointCloudColorQuantity {
362    fn as_any(&self) -> &dyn std::any::Any {
363        self
364    }
365
366    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
367        self
368    }
369
370    fn name(&self) -> &str {
371        &self.name
372    }
373
374    fn structure_name(&self) -> &str {
375        &self.structure_name
376    }
377
378    fn kind(&self) -> QuantityKind {
379        QuantityKind::Color
380    }
381
382    fn is_enabled(&self) -> bool {
383        self.enabled
384    }
385
386    fn set_enabled(&mut self, enabled: bool) {
387        self.enabled = enabled;
388    }
389
390    fn build_ui(&mut self, _ui: &dyn std::any::Any) {
391        // UI is handled by polyscope-ui/src/structure_ui.rs
392    }
393
394    fn refresh(&mut self) {
395        // GPU refresh is handled by polyscope/src/app/render.rs
396    }
397
398    fn data_size(&self) -> usize {
399        self.colors.len()
400    }
401}
402
403impl VertexQuantity for PointCloudColorQuantity {}