polyscope_structures/surface_mesh/
intrinsic_vector_quantity.rs

1//! Intrinsic (tangent-space) vector quantities for surface meshes.
2
3use glam::{Vec2, Vec3, Vec4};
4use polyscope_core::quantity::{FaceQuantity, Quantity, QuantityKind, VertexQuantity};
5use polyscope_render::{VectorRenderData, VectorUniforms};
6
7/// A vertex intrinsic vector quantity on a surface mesh.
8///
9/// Stores 2D tangent-space vectors along with a per-vertex tangent basis.
10/// These are projected to 3D world space for rendering.
11pub struct MeshVertexIntrinsicVectorQuantity {
12    name: String,
13    structure_name: String,
14    vectors: Vec<Vec2>, // 2D tangent-space vectors
15    basis_x: Vec<Vec3>, // Per-element X axis of tangent frame
16    basis_y: Vec<Vec3>, // Per-element Y axis of tangent frame
17    n_sym: u32,         // Symmetry: 1=vector, 2=line, 4=cross
18    enabled: bool,
19    length_scale: f32,
20    radius: f32,
21    color: Vec4,
22    render_data: Option<VectorRenderData>,
23}
24
25impl MeshVertexIntrinsicVectorQuantity {
26    /// Creates a new vertex intrinsic vector quantity.
27    pub fn new(
28        name: impl Into<String>,
29        structure_name: impl Into<String>,
30        vectors: Vec<Vec2>,
31        basis_x: Vec<Vec3>,
32        basis_y: Vec<Vec3>,
33    ) -> Self {
34        Self {
35            name: name.into(),
36            structure_name: structure_name.into(),
37            vectors,
38            basis_x,
39            basis_y,
40            n_sym: 1,
41            enabled: false,
42            length_scale: 1.0,
43            radius: 0.005,
44            color: Vec4::new(0.8, 0.2, 0.8, 1.0),
45            render_data: None,
46        }
47    }
48
49    /// Returns the 2D tangent-space vectors.
50    #[must_use]
51    pub fn vectors(&self) -> &[Vec2] {
52        &self.vectors
53    }
54
55    /// Returns the tangent basis X axes.
56    #[must_use]
57    pub fn basis_x(&self) -> &[Vec3] {
58        &self.basis_x
59    }
60
61    /// Returns the tangent basis Y axes.
62    #[must_use]
63    pub fn basis_y(&self) -> &[Vec3] {
64        &self.basis_y
65    }
66
67    /// Gets the symmetry order.
68    #[must_use]
69    pub fn n_sym(&self) -> u32 {
70        self.n_sym
71    }
72
73    /// Sets the symmetry order (1=vector, 2=line, 4=cross).
74    pub fn set_n_sym(&mut self, n: u32) -> &mut Self {
75        self.n_sym = n;
76        self
77    }
78
79    /// Gets the length scale.
80    #[must_use]
81    pub fn length_scale(&self) -> f32 {
82        self.length_scale
83    }
84
85    /// Sets the length scale.
86    pub fn set_length_scale(&mut self, scale: f32) -> &mut Self {
87        self.length_scale = scale;
88        self
89    }
90
91    /// Gets the radius.
92    #[must_use]
93    pub fn radius(&self) -> f32 {
94        self.radius
95    }
96
97    /// Sets the radius.
98    pub fn set_radius(&mut self, r: f32) -> &mut Self {
99        self.radius = r;
100        self
101    }
102
103    /// Gets the color.
104    #[must_use]
105    pub fn color(&self) -> Vec4 {
106        self.color
107    }
108
109    /// Sets the color.
110    pub fn set_color(&mut self, c: Vec3) -> &mut Self {
111        self.color = c.extend(1.0);
112        self
113    }
114
115    /// Project 2D tangent-space vectors to 3D world space.
116    #[must_use]
117    pub fn compute_world_vectors(&self) -> Vec<Vec3> {
118        self.vectors
119            .iter()
120            .enumerate()
121            .map(|(i, v2d)| self.basis_x[i] * v2d.x + self.basis_y[i] * v2d.y)
122            .collect()
123    }
124
125    /// Generate symmetry-rotated copies of the world vectors.
126    ///
127    /// For `n_sym = 1`, returns the original vectors.
128    /// For `n_sym = 2`, returns each vector and its negation (line field).
129    /// For `n_sym = 4`, returns each vector rotated by 0, 90, 180, 270 degrees in the tangent plane.
130    #[must_use]
131    pub fn compute_symmetric_world_vectors(&self) -> Vec<(usize, Vec3)> {
132        let mut result = Vec::new();
133        for (i, v2d) in self.vectors.iter().enumerate() {
134            for k in 0..self.n_sym {
135                let angle = k as f32 * std::f32::consts::TAU / self.n_sym as f32;
136                let cos_a = angle.cos();
137                let sin_a = angle.sin();
138                let rotated =
139                    Vec2::new(v2d.x * cos_a - v2d.y * sin_a, v2d.x * sin_a + v2d.y * cos_a);
140                let world_vec = self.basis_x[i] * rotated.x + self.basis_y[i] * rotated.y;
141                result.push((i, world_vec));
142            }
143        }
144        result
145    }
146
147    /// Auto-scales length and radius based on the structure's bounding box diagonal.
148    pub fn auto_scale(&mut self, structure_length_scale: f32) {
149        let world_vecs = self.compute_world_vectors();
150        let avg_length: f32 = if world_vecs.is_empty() {
151            1.0
152        } else {
153            let sum: f32 = world_vecs.iter().map(|v| v.length()).sum();
154            sum / world_vecs.len() as f32
155        };
156        if avg_length > 1e-8 {
157            self.length_scale = 0.02 * structure_length_scale / avg_length;
158        }
159        self.radius = 0.002 * structure_length_scale;
160    }
161
162    /// Initializes GPU resources for this vector quantity.
163    ///
164    /// Projects 2D tangent-space vectors to 3D world space (with symmetry)
165    /// and creates GPU buffers for rendering as arrows.
166    pub fn init_gpu_resources(
167        &mut self,
168        device: &wgpu::Device,
169        bind_group_layout: &wgpu::BindGroupLayout,
170        camera_buffer: &wgpu::Buffer,
171        base_positions: &[Vec3],
172    ) {
173        let sym_vectors = self.compute_symmetric_world_vectors();
174        let bases: Vec<Vec3> = sym_vectors
175            .iter()
176            .map(|(idx, _)| base_positions[*idx])
177            .collect();
178        let vecs: Vec<Vec3> = sym_vectors.iter().map(|(_, v)| *v).collect();
179        self.render_data = Some(VectorRenderData::new(
180            device,
181            bind_group_layout,
182            camera_buffer,
183            &bases,
184            &vecs,
185        ));
186    }
187
188    /// Returns the render data if initialized.
189    #[must_use]
190    pub fn render_data(&self) -> Option<&VectorRenderData> {
191        self.render_data.as_ref()
192    }
193
194    /// Updates GPU uniforms with the given model transform.
195    pub fn update_uniforms(&self, queue: &wgpu::Queue, model: &glam::Mat4) {
196        if let Some(render_data) = &self.render_data {
197            let uniforms = VectorUniforms {
198                model: model.to_cols_array(),
199                length_scale: self.length_scale,
200                radius: self.radius,
201                _padding: [0.0; 2],
202                color: self.color.to_array(),
203            };
204            render_data.update_uniforms(queue, &uniforms);
205        }
206    }
207
208    /// Builds the egui UI for this quantity.
209    pub fn build_egui_ui(&mut self, ui: &mut egui::Ui) -> bool {
210        let mut color = [self.color.x, self.color.y, self.color.z];
211        let changed = polyscope_ui::build_intrinsic_vector_quantity_ui(
212            ui,
213            &self.name,
214            &mut self.enabled,
215            &mut self.length_scale,
216            &mut self.radius,
217            &mut color,
218            &mut self.n_sym,
219        );
220        if changed {
221            self.color = Vec4::new(color[0], color[1], color[2], self.color.w);
222        }
223        changed
224    }
225}
226
227impl Quantity for MeshVertexIntrinsicVectorQuantity {
228    fn as_any(&self) -> &dyn std::any::Any {
229        self
230    }
231    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
232        self
233    }
234    fn name(&self) -> &str {
235        &self.name
236    }
237    fn structure_name(&self) -> &str {
238        &self.structure_name
239    }
240    fn kind(&self) -> QuantityKind {
241        QuantityKind::Vector
242    }
243    fn is_enabled(&self) -> bool {
244        self.enabled
245    }
246    fn set_enabled(&mut self, enabled: bool) {
247        self.enabled = enabled;
248    }
249    fn build_ui(&mut self, _ui: &dyn std::any::Any) {}
250    fn refresh(&mut self) {}
251    fn clear_gpu_resources(&mut self) {
252        self.render_data = None;
253    }
254    fn data_size(&self) -> usize {
255        self.vectors.len()
256    }
257}
258
259impl VertexQuantity for MeshVertexIntrinsicVectorQuantity {}
260
261/// A face intrinsic vector quantity on a surface mesh.
262///
263/// Stores 2D tangent-space vectors along with a per-face tangent basis.
264pub struct MeshFaceIntrinsicVectorQuantity {
265    name: String,
266    structure_name: String,
267    vectors: Vec<Vec2>,
268    basis_x: Vec<Vec3>,
269    basis_y: Vec<Vec3>,
270    n_sym: u32,
271    enabled: bool,
272    length_scale: f32,
273    radius: f32,
274    color: Vec4,
275    render_data: Option<VectorRenderData>,
276}
277
278impl MeshFaceIntrinsicVectorQuantity {
279    /// Creates a new face intrinsic vector quantity.
280    pub fn new(
281        name: impl Into<String>,
282        structure_name: impl Into<String>,
283        vectors: Vec<Vec2>,
284        basis_x: Vec<Vec3>,
285        basis_y: Vec<Vec3>,
286    ) -> Self {
287        Self {
288            name: name.into(),
289            structure_name: structure_name.into(),
290            vectors,
291            basis_x,
292            basis_y,
293            n_sym: 1,
294            enabled: false,
295            length_scale: 1.0,
296            radius: 0.005,
297            color: Vec4::new(0.2, 0.8, 0.8, 1.0),
298            render_data: None,
299        }
300    }
301
302    /// Returns the 2D tangent-space vectors.
303    #[must_use]
304    pub fn vectors(&self) -> &[Vec2] {
305        &self.vectors
306    }
307
308    /// Returns the tangent basis X axes.
309    #[must_use]
310    pub fn basis_x(&self) -> &[Vec3] {
311        &self.basis_x
312    }
313
314    /// Returns the tangent basis Y axes.
315    #[must_use]
316    pub fn basis_y(&self) -> &[Vec3] {
317        &self.basis_y
318    }
319
320    /// Gets the symmetry order.
321    #[must_use]
322    pub fn n_sym(&self) -> u32 {
323        self.n_sym
324    }
325
326    /// Sets the symmetry order.
327    pub fn set_n_sym(&mut self, n: u32) -> &mut Self {
328        self.n_sym = n;
329        self
330    }
331
332    /// Gets the length scale.
333    #[must_use]
334    pub fn length_scale(&self) -> f32 {
335        self.length_scale
336    }
337
338    /// Sets the length scale.
339    pub fn set_length_scale(&mut self, scale: f32) -> &mut Self {
340        self.length_scale = scale;
341        self
342    }
343
344    /// Gets the radius.
345    #[must_use]
346    pub fn radius(&self) -> f32 {
347        self.radius
348    }
349
350    /// Sets the radius.
351    pub fn set_radius(&mut self, r: f32) -> &mut Self {
352        self.radius = r;
353        self
354    }
355
356    /// Gets the color.
357    #[must_use]
358    pub fn color(&self) -> Vec4 {
359        self.color
360    }
361
362    /// Sets the color.
363    pub fn set_color(&mut self, c: Vec3) -> &mut Self {
364        self.color = c.extend(1.0);
365        self
366    }
367
368    /// Project 2D tangent-space vectors to 3D world space.
369    #[must_use]
370    pub fn compute_world_vectors(&self) -> Vec<Vec3> {
371        self.vectors
372            .iter()
373            .enumerate()
374            .map(|(i, v2d)| self.basis_x[i] * v2d.x + self.basis_y[i] * v2d.y)
375            .collect()
376    }
377
378    /// Generate symmetry-rotated copies of the world vectors.
379    #[must_use]
380    pub fn compute_symmetric_world_vectors(&self) -> Vec<(usize, Vec3)> {
381        let mut result = Vec::new();
382        for (i, v2d) in self.vectors.iter().enumerate() {
383            for k in 0..self.n_sym {
384                let angle = k as f32 * std::f32::consts::TAU / self.n_sym as f32;
385                let cos_a = angle.cos();
386                let sin_a = angle.sin();
387                let rotated =
388                    Vec2::new(v2d.x * cos_a - v2d.y * sin_a, v2d.x * sin_a + v2d.y * cos_a);
389                let world_vec = self.basis_x[i] * rotated.x + self.basis_y[i] * rotated.y;
390                result.push((i, world_vec));
391            }
392        }
393        result
394    }
395
396    /// Auto-scales length and radius based on the structure's bounding box diagonal.
397    pub fn auto_scale(&mut self, structure_length_scale: f32) {
398        let world_vecs = self.compute_world_vectors();
399        let avg_length: f32 = if world_vecs.is_empty() {
400            1.0
401        } else {
402            let sum: f32 = world_vecs.iter().map(|v| v.length()).sum();
403            sum / world_vecs.len() as f32
404        };
405        if avg_length > 1e-8 {
406            self.length_scale = 0.02 * structure_length_scale / avg_length;
407        }
408        self.radius = 0.002 * structure_length_scale;
409    }
410
411    /// Initializes GPU resources for this vector quantity.
412    pub fn init_gpu_resources(
413        &mut self,
414        device: &wgpu::Device,
415        bind_group_layout: &wgpu::BindGroupLayout,
416        camera_buffer: &wgpu::Buffer,
417        base_positions: &[Vec3],
418    ) {
419        let sym_vectors = self.compute_symmetric_world_vectors();
420        let bases: Vec<Vec3> = sym_vectors
421            .iter()
422            .map(|(idx, _)| base_positions[*idx])
423            .collect();
424        let vecs: Vec<Vec3> = sym_vectors.iter().map(|(_, v)| *v).collect();
425        self.render_data = Some(VectorRenderData::new(
426            device,
427            bind_group_layout,
428            camera_buffer,
429            &bases,
430            &vecs,
431        ));
432    }
433
434    /// Returns the render data if initialized.
435    #[must_use]
436    pub fn render_data(&self) -> Option<&VectorRenderData> {
437        self.render_data.as_ref()
438    }
439
440    /// Updates GPU uniforms with the given model transform.
441    pub fn update_uniforms(&self, queue: &wgpu::Queue, model: &glam::Mat4) {
442        if let Some(render_data) = &self.render_data {
443            let uniforms = VectorUniforms {
444                model: model.to_cols_array(),
445                length_scale: self.length_scale,
446                radius: self.radius,
447                _padding: [0.0; 2],
448                color: self.color.to_array(),
449            };
450            render_data.update_uniforms(queue, &uniforms);
451        }
452    }
453
454    /// Builds the egui UI for this quantity.
455    pub fn build_egui_ui(&mut self, ui: &mut egui::Ui) -> bool {
456        let mut color = [self.color.x, self.color.y, self.color.z];
457        let changed = polyscope_ui::build_intrinsic_vector_quantity_ui(
458            ui,
459            &self.name,
460            &mut self.enabled,
461            &mut self.length_scale,
462            &mut self.radius,
463            &mut color,
464            &mut self.n_sym,
465        );
466        if changed {
467            self.color = Vec4::new(color[0], color[1], color[2], self.color.w);
468        }
469        changed
470    }
471}
472
473impl Quantity for MeshFaceIntrinsicVectorQuantity {
474    fn as_any(&self) -> &dyn std::any::Any {
475        self
476    }
477    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
478        self
479    }
480    fn name(&self) -> &str {
481        &self.name
482    }
483    fn structure_name(&self) -> &str {
484        &self.structure_name
485    }
486    fn kind(&self) -> QuantityKind {
487        QuantityKind::Vector
488    }
489    fn is_enabled(&self) -> bool {
490        self.enabled
491    }
492    fn set_enabled(&mut self, enabled: bool) {
493        self.enabled = enabled;
494    }
495    fn build_ui(&mut self, _ui: &dyn std::any::Any) {}
496    fn refresh(&mut self) {}
497    fn clear_gpu_resources(&mut self) {
498        self.render_data = None;
499    }
500    fn data_size(&self) -> usize {
501        self.vectors.len()
502    }
503}
504
505impl FaceQuantity for MeshFaceIntrinsicVectorQuantity {}
506
507#[cfg(test)]
508mod tests {
509    use super::*;
510
511    #[test]
512    fn test_vertex_intrinsic_creation() {
513        let vectors = vec![Vec2::new(1.0, 0.0), Vec2::new(0.0, 1.0)];
514        let basis_x = vec![Vec3::X, Vec3::X];
515        let basis_y = vec![Vec3::Y, Vec3::Y];
516
517        let q = MeshVertexIntrinsicVectorQuantity::new(
518            "tangent_field",
519            "mesh",
520            vectors,
521            basis_x,
522            basis_y,
523        );
524
525        assert_eq!(q.name(), "tangent_field");
526        assert_eq!(q.data_size(), 2);
527        assert_eq!(q.kind(), QuantityKind::Vector);
528        assert_eq!(q.n_sym(), 1);
529        assert!(!q.is_enabled());
530    }
531
532    #[test]
533    fn test_world_vector_projection() {
534        // Vector (1, 0) in tangent space with basis X=X, Y=Y -> world vector = X
535        let vectors = vec![Vec2::new(1.0, 0.0), Vec2::new(0.0, 1.0)];
536        let basis_x = vec![Vec3::X, Vec3::X];
537        let basis_y = vec![Vec3::Y, Vec3::Y];
538
539        let q = MeshVertexIntrinsicVectorQuantity::new("test", "mesh", vectors, basis_x, basis_y);
540
541        let world = q.compute_world_vectors();
542        assert_eq!(world.len(), 2);
543        assert!((world[0] - Vec3::X).length() < 1e-5);
544        assert!((world[1] - Vec3::Y).length() < 1e-5);
545    }
546
547    #[test]
548    fn test_world_vector_projection_rotated_basis() {
549        // Vector (1, 0) in tangent space with basis X=Y, Y=Z -> world vector = Y
550        let vectors = vec![Vec2::new(1.0, 0.0)];
551        let basis_x = vec![Vec3::Y];
552        let basis_y = vec![Vec3::Z];
553
554        let q = MeshVertexIntrinsicVectorQuantity::new("test", "mesh", vectors, basis_x, basis_y);
555
556        let world = q.compute_world_vectors();
557        assert!((world[0] - Vec3::Y).length() < 1e-5);
558    }
559
560    #[test]
561    fn test_symmetry_n1() {
562        let vectors = vec![Vec2::new(1.0, 0.0)];
563        let basis_x = vec![Vec3::X];
564        let basis_y = vec![Vec3::Y];
565
566        let q = MeshVertexIntrinsicVectorQuantity::new("test", "mesh", vectors, basis_x, basis_y);
567
568        let sym = q.compute_symmetric_world_vectors();
569        assert_eq!(sym.len(), 1); // 1 vector * 1 symmetry = 1
570        assert_eq!(sym[0].0, 0);
571        assert!((sym[0].1 - Vec3::X).length() < 1e-5);
572    }
573
574    #[test]
575    fn test_symmetry_n2() {
576        let vectors = vec![Vec2::new(1.0, 0.0)];
577        let basis_x = vec![Vec3::X];
578        let basis_y = vec![Vec3::Y];
579
580        let mut q =
581            MeshVertexIntrinsicVectorQuantity::new("test", "mesh", vectors, basis_x, basis_y);
582        q.set_n_sym(2);
583
584        let sym = q.compute_symmetric_world_vectors();
585        assert_eq!(sym.len(), 2); // 1 vector * 2 symmetry = 2
586        // First copy: original direction (+X)
587        assert!((sym[0].1 - Vec3::X).length() < 1e-5);
588        // Second copy: 180 degree rotation (-X)
589        assert!((sym[1].1 + Vec3::X).length() < 1e-5);
590    }
591
592    #[test]
593    fn test_symmetry_n4() {
594        let vectors = vec![Vec2::new(1.0, 0.0)];
595        let basis_x = vec![Vec3::X];
596        let basis_y = vec![Vec3::Y];
597
598        let mut q =
599            MeshVertexIntrinsicVectorQuantity::new("test", "mesh", vectors, basis_x, basis_y);
600        q.set_n_sym(4);
601
602        let sym = q.compute_symmetric_world_vectors();
603        assert_eq!(sym.len(), 4); // 1 vector * 4 symmetry = 4
604        // 0 deg: +X
605        assert!((sym[0].1 - Vec3::X).length() < 1e-5);
606        // 90 deg: +Y
607        assert!((sym[1].1 - Vec3::Y).length() < 1e-5);
608        // 180 deg: -X
609        assert!((sym[2].1 + Vec3::X).length() < 1e-5);
610        // 270 deg: -Y
611        assert!((sym[3].1 + Vec3::Y).length() < 1e-5);
612    }
613
614    #[test]
615    fn test_face_intrinsic_creation() {
616        let vectors = vec![Vec2::new(1.0, 0.0)];
617        let basis_x = vec![Vec3::X];
618        let basis_y = vec![Vec3::Y];
619
620        let q =
621            MeshFaceIntrinsicVectorQuantity::new("face_tangent", "mesh", vectors, basis_x, basis_y);
622
623        assert_eq!(q.name(), "face_tangent");
624        assert_eq!(q.data_size(), 1);
625        assert_eq!(q.kind(), QuantityKind::Vector);
626    }
627
628    #[test]
629    fn test_face_intrinsic_world_vectors() {
630        let vectors = vec![Vec2::new(0.5, 0.5)];
631        let basis_x = vec![Vec3::X];
632        let basis_y = vec![Vec3::Y];
633
634        let q = MeshFaceIntrinsicVectorQuantity::new("test", "mesh", vectors, basis_x, basis_y);
635
636        let world = q.compute_world_vectors();
637        let expected = Vec3::new(0.5, 0.5, 0.0);
638        assert!((world[0] - expected).length() < 1e-5);
639    }
640}
polyscope_structures/surface_mesh/intrinsic_vector_quantity.rs

polyscope_structures/surface_mesh/
intrinsic_vector_quantity.rs
