1use glam::{Vec3, Vec4};
4use polyscope_core::quantity::{FaceQuantity, Quantity, QuantityKind, VertexQuantity};
5use polyscope_render::{ColorMap, VectorRenderData, VectorUniforms};
6
7pub struct MeshVertexScalarQuantity {
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 MeshVertexScalarQuantity {
19 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 #[must_use]
41 pub fn values(&self) -> &[f32] {
42 &self.values
43 }
44
45 #[must_use]
47 pub fn colormap_name(&self) -> &str {
48 &self.colormap_name
49 }
50
51 pub fn set_colormap(&mut self, name: impl Into<String>) {
53 self.colormap_name = name.into();
54 }
55
56 #[must_use]
58 pub fn range_min(&self) -> f32 {
59 self.range_min
60 }
61
62 #[must_use]
64 pub fn range_max(&self) -> f32 {
65 self.range_max
66 }
67
68 pub fn set_range(&mut self, min: f32, max: f32) {
70 self.range_min = min;
71 self.range_max = max;
72 }
73
74 #[must_use]
76 pub fn compute_colors(&self, colormap: &ColorMap) -> Vec<Vec4> {
77 let range = self.range_max - self.range_min;
78 let range = if range.abs() < 1e-10 { 1.0 } else { range };
79
80 self.values
81 .iter()
82 .map(|&v| {
83 let t = (v - self.range_min) / range;
84 colormap.sample(t).extend(1.0)
85 })
86 .collect()
87 }
88
89 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 MeshVertexScalarQuantity {
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
135 fn refresh(&mut self) {}
136
137 fn data_size(&self) -> usize {
138 self.values.len()
139 }
140}
141
142impl VertexQuantity for MeshVertexScalarQuantity {}
143
144pub struct MeshFaceScalarQuantity {
146 name: String,
147 structure_name: String,
148 values: Vec<f32>,
149 enabled: bool,
150 colormap_name: String,
151 range_min: f32,
152 range_max: f32,
153}
154
155impl MeshFaceScalarQuantity {
156 pub fn new(
158 name: impl Into<String>,
159 structure_name: impl Into<String>,
160 values: Vec<f32>,
161 ) -> Self {
162 let min = values.iter().copied().fold(f32::INFINITY, f32::min);
163 let max = values.iter().copied().fold(f32::NEG_INFINITY, f32::max);
164
165 Self {
166 name: name.into(),
167 structure_name: structure_name.into(),
168 values,
169 enabled: false,
170 colormap_name: "viridis".to_string(),
171 range_min: min,
172 range_max: max,
173 }
174 }
175
176 #[must_use]
178 pub fn values(&self) -> &[f32] {
179 &self.values
180 }
181
182 #[must_use]
184 pub fn colormap_name(&self) -> &str {
185 &self.colormap_name
186 }
187
188 pub fn set_colormap(&mut self, name: impl Into<String>) {
190 self.colormap_name = name.into();
191 }
192
193 #[must_use]
195 pub fn range_min(&self) -> f32 {
196 self.range_min
197 }
198
199 #[must_use]
201 pub fn range_max(&self) -> f32 {
202 self.range_max
203 }
204
205 pub fn set_range(&mut self, min: f32, max: f32) {
207 self.range_min = min;
208 self.range_max = max;
209 }
210
211 #[must_use]
214 pub fn compute_vertex_colors(
215 &self,
216 faces: &[Vec<u32>],
217 num_vertices: usize,
218 colormap: &ColorMap,
219 ) -> Vec<Vec4> {
220 let range = self.range_max - self.range_min;
221 let range = if range.abs() < 1e-10 { 1.0 } else { range };
222
223 let mut colors = vec![Vec4::splat(0.5); num_vertices];
224
225 for (face_idx, face) in faces.iter().enumerate() {
226 let t = (self.values[face_idx] - self.range_min) / range;
227 let color = colormap.sample(t).extend(1.0);
228 for &vi in face {
229 colors[vi as usize] = color;
230 }
231 }
232
233 colors
234 }
235
236 pub fn build_egui_ui(&mut self, ui: &mut egui::Ui) -> bool {
238 let colormaps = ["viridis", "blues", "reds", "coolwarm", "rainbow"];
239 polyscope_ui::build_scalar_quantity_ui(
240 ui,
241 &self.name,
242 &mut self.enabled,
243 &mut self.colormap_name,
244 &mut self.range_min,
245 &mut self.range_max,
246 &colormaps,
247 )
248 }
249}
250
251impl Quantity for MeshFaceScalarQuantity {
252 fn as_any(&self) -> &dyn std::any::Any {
253 self
254 }
255
256 fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
257 self
258 }
259
260 fn name(&self) -> &str {
261 &self.name
262 }
263
264 fn structure_name(&self) -> &str {
265 &self.structure_name
266 }
267
268 fn kind(&self) -> QuantityKind {
269 QuantityKind::Scalar
270 }
271
272 fn is_enabled(&self) -> bool {
273 self.enabled
274 }
275
276 fn set_enabled(&mut self, enabled: bool) {
277 self.enabled = enabled;
278 }
279
280 fn build_ui(&mut self, _ui: &dyn std::any::Any) {}
281
282 fn refresh(&mut self) {}
283
284 fn data_size(&self) -> usize {
285 self.values.len()
286 }
287}
288
289impl FaceQuantity for MeshFaceScalarQuantity {}
290
291pub struct MeshVertexColorQuantity {
293 name: String,
294 structure_name: String,
295 colors: Vec<Vec4>,
296 enabled: bool,
297 has_transparency: bool,
298}
299
300impl MeshVertexColorQuantity {
301 pub fn new(
303 name: impl Into<String>,
304 structure_name: impl Into<String>,
305 colors: Vec<Vec3>,
306 ) -> Self {
307 Self {
308 name: name.into(),
309 structure_name: structure_name.into(),
310 colors: colors.into_iter().map(|c| c.extend(1.0)).collect(),
311 enabled: false,
312 has_transparency: false,
313 }
314 }
315
316 pub fn new_with_alpha(
318 name: impl Into<String>,
319 structure_name: impl Into<String>,
320 colors: Vec<Vec4>,
321 ) -> Self {
322 let has_transparency = colors.iter().any(|c| c.w < 0.999);
323 Self {
324 name: name.into(),
325 structure_name: structure_name.into(),
326 colors,
327 enabled: false,
328 has_transparency,
329 }
330 }
331
332 #[must_use]
334 pub fn colors(&self) -> &[Vec4] {
335 &self.colors
336 }
337
338 #[must_use]
340 pub fn has_transparency(&self) -> bool {
341 self.has_transparency
342 }
343
344 pub fn build_egui_ui(&mut self, ui: &mut egui::Ui) -> bool {
346 polyscope_ui::build_color_quantity_ui(ui, &self.name, &mut self.enabled, self.colors.len())
347 }
348}
349
350impl Quantity for MeshVertexColorQuantity {
351 fn as_any(&self) -> &dyn std::any::Any {
352 self
353 }
354
355 fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
356 self
357 }
358
359 fn name(&self) -> &str {
360 &self.name
361 }
362
363 fn structure_name(&self) -> &str {
364 &self.structure_name
365 }
366
367 fn kind(&self) -> QuantityKind {
368 QuantityKind::Color
369 }
370
371 fn is_enabled(&self) -> bool {
372 self.enabled
373 }
374
375 fn set_enabled(&mut self, enabled: bool) {
376 self.enabled = enabled;
377 }
378
379 fn build_ui(&mut self, _ui: &dyn std::any::Any) {}
380
381 fn refresh(&mut self) {}
382
383 fn data_size(&self) -> usize {
384 self.colors.len()
385 }
386}
387
388impl VertexQuantity for MeshVertexColorQuantity {}
389
390pub struct MeshFaceColorQuantity {
392 name: String,
393 structure_name: String,
394 colors: Vec<Vec4>,
395 enabled: bool,
396 has_transparency: bool,
397}
398
399impl MeshFaceColorQuantity {
400 pub fn new(
402 name: impl Into<String>,
403 structure_name: impl Into<String>,
404 colors: Vec<Vec3>,
405 ) -> Self {
406 Self {
407 name: name.into(),
408 structure_name: structure_name.into(),
409 colors: colors.into_iter().map(|c| c.extend(1.0)).collect(),
410 enabled: false,
411 has_transparency: false,
412 }
413 }
414
415 pub fn new_with_alpha(
417 name: impl Into<String>,
418 structure_name: impl Into<String>,
419 colors: Vec<Vec4>,
420 ) -> Self {
421 let has_transparency = colors.iter().any(|c| c.w < 0.999);
422 Self {
423 name: name.into(),
424 structure_name: structure_name.into(),
425 colors,
426 enabled: false,
427 has_transparency,
428 }
429 }
430
431 #[must_use]
433 pub fn colors(&self) -> &[Vec4] {
434 &self.colors
435 }
436
437 #[must_use]
439 pub fn has_transparency(&self) -> bool {
440 self.has_transparency
441 }
442
443 #[must_use]
446 pub fn compute_vertex_colors(&self, faces: &[Vec<u32>], num_vertices: usize) -> Vec<Vec4> {
447 let mut colors = vec![Vec4::splat(0.5); num_vertices];
448
449 for (face_idx, face) in faces.iter().enumerate() {
450 let color = self.colors[face_idx];
451 for &vi in face {
452 colors[vi as usize] = color;
453 }
454 }
455
456 colors
457 }
458
459 pub fn build_egui_ui(&mut self, ui: &mut egui::Ui) -> bool {
461 polyscope_ui::build_color_quantity_ui(ui, &self.name, &mut self.enabled, self.colors.len())
462 }
463}
464
465impl Quantity for MeshFaceColorQuantity {
466 fn as_any(&self) -> &dyn std::any::Any {
467 self
468 }
469
470 fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
471 self
472 }
473
474 fn name(&self) -> &str {
475 &self.name
476 }
477
478 fn structure_name(&self) -> &str {
479 &self.structure_name
480 }
481
482 fn kind(&self) -> QuantityKind {
483 QuantityKind::Color
484 }
485
486 fn is_enabled(&self) -> bool {
487 self.enabled
488 }
489
490 fn set_enabled(&mut self, enabled: bool) {
491 self.enabled = enabled;
492 }
493
494 fn build_ui(&mut self, _ui: &dyn std::any::Any) {}
495
496 fn refresh(&mut self) {}
497
498 fn data_size(&self) -> usize {
499 self.colors.len()
500 }
501}
502
503impl FaceQuantity for MeshFaceColorQuantity {}
504
505pub struct MeshVertexVectorQuantity {
507 name: String,
508 structure_name: String,
509 vectors: Vec<Vec3>,
510 enabled: bool,
511 length_scale: f32,
512 radius: f32,
513 color: Vec4,
514 render_data: Option<VectorRenderData>,
515}
516
517impl MeshVertexVectorQuantity {
518 pub fn new(
520 name: impl Into<String>,
521 structure_name: impl Into<String>,
522 vectors: Vec<Vec3>,
523 ) -> Self {
524 Self {
525 name: name.into(),
526 structure_name: structure_name.into(),
527 vectors,
528 enabled: false,
529 length_scale: 1.0,
530 radius: 0.005,
531 color: Vec4::new(0.8, 0.2, 0.2, 1.0),
532 render_data: None,
533 }
534 }
535
536 #[must_use]
538 pub fn vectors(&self) -> &[Vec3] {
539 &self.vectors
540 }
541
542 #[must_use]
544 pub fn length_scale(&self) -> f32 {
545 self.length_scale
546 }
547
548 pub fn set_length_scale(&mut self, scale: f32) {
550 self.length_scale = scale;
551 }
552
553 #[must_use]
555 pub fn radius(&self) -> f32 {
556 self.radius
557 }
558
559 pub fn set_radius(&mut self, r: f32) {
561 self.radius = r;
562 }
563
564 #[must_use]
566 pub fn color(&self) -> Vec4 {
567 self.color
568 }
569
570 pub fn set_color(&mut self, c: Vec3) {
572 self.color = c.extend(1.0);
573 }
574
575 pub fn auto_scale(&mut self, structure_length_scale: f32) {
580 let avg_length: f32 = if self.vectors.is_empty() {
581 1.0
582 } else {
583 let sum: f32 = self.vectors.iter().map(|v| v.length()).sum();
584 sum / self.vectors.len() as f32
585 };
586 if avg_length > 1e-8 {
587 self.length_scale = 0.02 * structure_length_scale / avg_length;
589 }
590 self.radius = 0.002 * structure_length_scale;
592 }
593
594 pub fn init_gpu_resources(
596 &mut self,
597 device: &wgpu::Device,
598 bind_group_layout: &wgpu::BindGroupLayout,
599 camera_buffer: &wgpu::Buffer,
600 base_positions: &[Vec3],
601 ) {
602 self.render_data = Some(VectorRenderData::new(
603 device,
604 bind_group_layout,
605 camera_buffer,
606 base_positions,
607 &self.vectors,
608 ));
609 }
610
611 #[must_use]
613 pub fn render_data(&self) -> Option<&VectorRenderData> {
614 self.render_data.as_ref()
615 }
616
617 pub fn update_uniforms(&self, queue: &wgpu::Queue, model: &glam::Mat4) {
619 if let Some(render_data) = &self.render_data {
620 let uniforms = VectorUniforms {
621 model: model.to_cols_array(),
622 length_scale: self.length_scale,
623 radius: self.radius,
624 _padding: [0.0; 2],
625 color: self.color.to_array(),
626 };
627 render_data.update_uniforms(queue, &uniforms);
628 }
629 }
630
631 pub fn build_egui_ui(&mut self, ui: &mut egui::Ui) -> bool {
633 let mut color = [self.color.x, self.color.y, self.color.z];
634 let changed = polyscope_ui::build_vector_quantity_ui(
635 ui,
636 &self.name,
637 &mut self.enabled,
638 &mut self.length_scale,
639 &mut self.radius,
640 &mut color,
641 );
642 if changed {
643 self.color = Vec4::new(color[0], color[1], color[2], self.color.w);
644 }
645 changed
646 }
647}
648
649impl Quantity for MeshVertexVectorQuantity {
650 fn as_any(&self) -> &dyn std::any::Any {
651 self
652 }
653
654 fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
655 self
656 }
657
658 fn name(&self) -> &str {
659 &self.name
660 }
661
662 fn structure_name(&self) -> &str {
663 &self.structure_name
664 }
665
666 fn kind(&self) -> QuantityKind {
667 QuantityKind::Vector
668 }
669
670 fn is_enabled(&self) -> bool {
671 self.enabled
672 }
673
674 fn set_enabled(&mut self, enabled: bool) {
675 self.enabled = enabled;
676 }
677
678 fn build_ui(&mut self, _ui: &dyn std::any::Any) {}
679
680 fn refresh(&mut self) {}
681
682 fn clear_gpu_resources(&mut self) {
683 self.render_data = None;
684 }
685
686 fn data_size(&self) -> usize {
687 self.vectors.len()
688 }
689}
690
691impl VertexQuantity for MeshVertexVectorQuantity {}
692
693pub struct MeshFaceVectorQuantity {
695 name: String,
696 structure_name: String,
697 vectors: Vec<Vec3>,
698 enabled: bool,
699 length_scale: f32,
700 radius: f32,
701 color: Vec4,
702 render_data: Option<VectorRenderData>,
703}
704
705impl MeshFaceVectorQuantity {
706 pub fn new(
708 name: impl Into<String>,
709 structure_name: impl Into<String>,
710 vectors: Vec<Vec3>,
711 ) -> Self {
712 Self {
713 name: name.into(),
714 structure_name: structure_name.into(),
715 vectors,
716 enabled: false,
717 length_scale: 1.0,
718 radius: 0.005,
719 color: Vec4::new(0.2, 0.2, 0.8, 1.0),
720 render_data: None,
721 }
722 }
723
724 #[must_use]
726 pub fn vectors(&self) -> &[Vec3] {
727 &self.vectors
728 }
729
730 #[must_use]
732 pub fn length_scale(&self) -> f32 {
733 self.length_scale
734 }
735
736 pub fn set_length_scale(&mut self, scale: f32) {
738 self.length_scale = scale;
739 }
740
741 #[must_use]
743 pub fn radius(&self) -> f32 {
744 self.radius
745 }
746
747 pub fn set_radius(&mut self, r: f32) {
749 self.radius = r;
750 }
751
752 #[must_use]
754 pub fn color(&self) -> Vec4 {
755 self.color
756 }
757
758 pub fn set_color(&mut self, c: Vec3) {
760 self.color = c.extend(1.0);
761 }
762
763 pub fn auto_scale(&mut self, structure_length_scale: f32) {
768 let avg_length: f32 = if self.vectors.is_empty() {
769 1.0
770 } else {
771 let sum: f32 = self.vectors.iter().map(|v| v.length()).sum();
772 sum / self.vectors.len() as f32
773 };
774 if avg_length > 1e-8 {
775 self.length_scale = 0.02 * structure_length_scale / avg_length;
777 }
778 self.radius = 0.002 * structure_length_scale;
780 }
781
782 pub fn init_gpu_resources(
784 &mut self,
785 device: &wgpu::Device,
786 bind_group_layout: &wgpu::BindGroupLayout,
787 camera_buffer: &wgpu::Buffer,
788 base_positions: &[Vec3],
789 ) {
790 self.render_data = Some(VectorRenderData::new(
791 device,
792 bind_group_layout,
793 camera_buffer,
794 base_positions,
795 &self.vectors,
796 ));
797 }
798
799 #[must_use]
801 pub fn render_data(&self) -> Option<&VectorRenderData> {
802 self.render_data.as_ref()
803 }
804
805 pub fn update_uniforms(&self, queue: &wgpu::Queue, model: &glam::Mat4) {
807 if let Some(render_data) = &self.render_data {
808 let uniforms = VectorUniforms {
809 model: model.to_cols_array(),
810 length_scale: self.length_scale,
811 radius: self.radius,
812 _padding: [0.0; 2],
813 color: self.color.to_array(),
814 };
815 render_data.update_uniforms(queue, &uniforms);
816 }
817 }
818
819 pub fn build_egui_ui(&mut self, ui: &mut egui::Ui) -> bool {
821 let mut color = [self.color.x, self.color.y, self.color.z];
822 let changed = polyscope_ui::build_vector_quantity_ui(
823 ui,
824 &self.name,
825 &mut self.enabled,
826 &mut self.length_scale,
827 &mut self.radius,
828 &mut color,
829 );
830 if changed {
831 self.color = Vec4::new(color[0], color[1], color[2], self.color.w);
832 }
833 changed
834 }
835}
836
837impl Quantity for MeshFaceVectorQuantity {
838 fn as_any(&self) -> &dyn std::any::Any {
839 self
840 }
841
842 fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
843 self
844 }
845
846 fn name(&self) -> &str {
847 &self.name
848 }
849
850 fn structure_name(&self) -> &str {
851 &self.structure_name
852 }
853
854 fn kind(&self) -> QuantityKind {
855 QuantityKind::Vector
856 }
857
858 fn is_enabled(&self) -> bool {
859 self.enabled
860 }
861
862 fn set_enabled(&mut self, enabled: bool) {
863 self.enabled = enabled;
864 }
865
866 fn build_ui(&mut self, _ui: &dyn std::any::Any) {}
867
868 fn refresh(&mut self) {}
869
870 fn clear_gpu_resources(&mut self) {
871 self.render_data = None;
872 }
873
874 fn data_size(&self) -> usize {
875 self.vectors.len()
876 }
877}
878
879impl FaceQuantity for MeshFaceVectorQuantity {}