Skip to main content

polyscope_structures/surface_mesh/
parameterization_quantity.rs

1//! Parameterization (UV) quantities for surface meshes.
2
3use glam::{Vec2, Vec3, Vec4};
4use polyscope_core::quantity::{
5    FaceQuantity, ParamCoordsType, ParamVizStyle, Quantity, QuantityKind, VertexQuantity,
6};
7
8/// Simple HSV to RGB conversion helper.
9#[allow(clippy::many_single_char_names)]
10fn hsv_to_rgb(h: f32, s: f32, v: f32) -> Vec3 {
11    let c = v * s;
12    let x = c * (1.0 - ((h * 6.0) % 2.0 - 1.0).abs());
13    let m = v - c;
14    let (r, g, b) = match (h * 6.0) as u32 {
15        0 => (c, x, 0.0),
16        1 => (x, c, 0.0),
17        2 => (0.0, c, x),
18        3 => (0.0, x, c),
19        4 => (x, 0.0, c),
20        _ => (c, 0.0, x),
21    };
22    Vec3::new(r + m, g + m, b + m)
23}
24
25/// A vertex parameterization (UV) quantity on a surface mesh.
26pub struct MeshVertexParameterizationQuantity {
27    name: String,
28    structure_name: String,
29    coords: Vec<Vec2>,
30    enabled: bool,
31    // Visualization parameters
32    style: ParamVizStyle,
33    coords_type: ParamCoordsType,
34    checker_size: f32,
35    checker_colors: [Vec3; 2],
36    grid_line_width: f32,
37}
38
39impl MeshVertexParameterizationQuantity {
40    /// Creates a new vertex parameterization quantity.
41    pub fn new(
42        name: impl Into<String>,
43        structure_name: impl Into<String>,
44        coords: Vec<Vec2>,
45    ) -> Self {
46        Self {
47            name: name.into(),
48            structure_name: structure_name.into(),
49            coords,
50            enabled: false,
51            style: ParamVizStyle::default(),
52            coords_type: ParamCoordsType::default(),
53            checker_size: 0.1,
54            checker_colors: [Vec3::new(1.0, 0.4, 0.4), Vec3::new(0.4, 0.4, 1.0)],
55            grid_line_width: 0.02,
56        }
57    }
58
59    /// Returns the UV coordinates.
60    #[must_use]
61    pub fn coords(&self) -> &[Vec2] {
62        &self.coords
63    }
64
65    /// Gets the visualization style.
66    #[must_use]
67    pub fn style(&self) -> ParamVizStyle {
68        self.style
69    }
70
71    /// Sets the visualization style.
72    pub fn set_style(&mut self, style: ParamVizStyle) -> &mut Self {
73        self.style = style;
74        self
75    }
76
77    /// Gets the coordinate type.
78    #[must_use]
79    pub fn coords_type(&self) -> ParamCoordsType {
80        self.coords_type
81    }
82
83    /// Sets the coordinate type.
84    pub fn set_coords_type(&mut self, ct: ParamCoordsType) -> &mut Self {
85        self.coords_type = ct;
86        self
87    }
88
89    /// Gets the checker size.
90    #[must_use]
91    pub fn checker_size(&self) -> f32 {
92        self.checker_size
93    }
94
95    /// Sets the checker size.
96    pub fn set_checker_size(&mut self, size: f32) -> &mut Self {
97        self.checker_size = size;
98        self
99    }
100
101    /// Gets the checker colors.
102    #[must_use]
103    pub fn checker_colors(&self) -> [Vec3; 2] {
104        self.checker_colors
105    }
106
107    /// Sets the checker colors.
108    pub fn set_checker_colors(&mut self, colors: [Vec3; 2]) -> &mut Self {
109        self.checker_colors = colors;
110        self
111    }
112
113    /// Compute per-vertex colors based on the current visualization style.
114    #[must_use]
115    pub fn compute_colors(&self) -> Vec<Vec4> {
116        match self.style {
117            ParamVizStyle::Checker => self.compute_checker_colors(),
118            ParamVizStyle::Grid => self.compute_grid_colors(),
119            ParamVizStyle::LocalCheck => self.compute_local_check_colors(),
120            ParamVizStyle::LocalRad => self.compute_local_rad_colors(),
121        }
122    }
123
124    fn compute_checker_colors(&self) -> Vec<Vec4> {
125        self.coords
126            .iter()
127            .map(|uv| {
128                let u_cell = (uv.x / self.checker_size).floor() as i32;
129                let v_cell = (uv.y / self.checker_size).floor() as i32;
130                if (u_cell + v_cell) % 2 == 0 {
131                    self.checker_colors[0].extend(1.0)
132                } else {
133                    self.checker_colors[1].extend(1.0)
134                }
135            })
136            .collect()
137    }
138
139    fn compute_grid_colors(&self) -> Vec<Vec4> {
140        self.coords
141            .iter()
142            .map(|uv| {
143                let u_frac = (uv.x / self.checker_size).fract().abs();
144                let v_frac = (uv.y / self.checker_size).fract().abs();
145                let on_line = u_frac < self.grid_line_width
146                    || u_frac > (1.0 - self.grid_line_width)
147                    || v_frac < self.grid_line_width
148                    || v_frac > (1.0 - self.grid_line_width);
149                if on_line {
150                    self.checker_colors[1].extend(1.0)
151                } else {
152                    self.checker_colors[0].extend(1.0)
153                }
154            })
155            .collect()
156    }
157
158    fn compute_local_check_colors(&self) -> Vec<Vec4> {
159        self.coords
160            .iter()
161            .map(|uv| {
162                let r = uv.length();
163                let angle = uv.y.atan2(uv.x);
164                let hue = (angle / std::f32::consts::TAU + 1.0) % 1.0;
165                let base = hsv_to_rgb(hue, 0.7, 0.9);
166                let u_cell = (uv.x / self.checker_size).floor() as i32;
167                let v_cell = (uv.y / self.checker_size).floor() as i32;
168                let dim = if (u_cell + v_cell) % 2 == 0 { 1.0 } else { 0.6 };
169                (base * dim * (1.0 - (-r * 2.0).exp() * 0.5)).extend(1.0)
170            })
171            .collect()
172    }
173
174    fn compute_local_rad_colors(&self) -> Vec<Vec4> {
175        self.coords
176            .iter()
177            .map(|uv| {
178                let r = uv.length();
179                let angle = uv.y.atan2(uv.x);
180                let hue = (angle / std::f32::consts::TAU + 1.0) % 1.0;
181                let base = hsv_to_rgb(hue, 0.7, 0.9);
182                let stripe = f32::from(u8::from((r / self.checker_size).floor() as i32 % 2 == 0));
183                let dim = 0.6 + 0.4 * stripe;
184                (base * dim).extend(1.0)
185            })
186            .collect()
187    }
188
189    /// Builds the egui UI for this quantity.
190    pub fn build_egui_ui(&mut self, ui: &mut egui::Ui) -> bool {
191        polyscope_ui::build_parameterization_quantity_ui(
192            ui,
193            &self.name,
194            &mut self.enabled,
195            &mut self.style,
196            &mut self.checker_size,
197            &mut self.checker_colors,
198        )
199    }
200}
201
202impl Quantity for MeshVertexParameterizationQuantity {
203    fn as_any(&self) -> &dyn std::any::Any {
204        self
205    }
206
207    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
208        self
209    }
210
211    fn name(&self) -> &str {
212        &self.name
213    }
214
215    fn structure_name(&self) -> &str {
216        &self.structure_name
217    }
218
219    fn kind(&self) -> QuantityKind {
220        QuantityKind::Parameterization
221    }
222
223    fn is_enabled(&self) -> bool {
224        self.enabled
225    }
226
227    fn set_enabled(&mut self, enabled: bool) {
228        self.enabled = enabled;
229    }
230
231    fn build_ui(&mut self, _ui: &dyn std::any::Any) {}
232
233    fn refresh(&mut self) {}
234
235    fn data_size(&self) -> usize {
236        self.coords.len()
237    }
238}
239
240impl VertexQuantity for MeshVertexParameterizationQuantity {}
241
242/// A corner (per-face-vertex) parameterization quantity.
243/// Used when UV islands are disconnected (different UV at shared vertices).
244pub struct MeshCornerParameterizationQuantity {
245    name: String,
246    structure_name: String,
247    coords: Vec<Vec2>, // One per corner (3 * num_triangles for triangle meshes)
248    enabled: bool,
249    style: ParamVizStyle,
250    coords_type: ParamCoordsType,
251    checker_size: f32,
252    checker_colors: [Vec3; 2],
253    grid_line_width: f32,
254}
255
256impl MeshCornerParameterizationQuantity {
257    /// Creates a new corner parameterization quantity.
258    pub fn new(
259        name: impl Into<String>,
260        structure_name: impl Into<String>,
261        coords: Vec<Vec2>,
262    ) -> Self {
263        Self {
264            name: name.into(),
265            structure_name: structure_name.into(),
266            coords,
267            enabled: false,
268            style: ParamVizStyle::default(),
269            coords_type: ParamCoordsType::default(),
270            checker_size: 0.1,
271            checker_colors: [Vec3::new(1.0, 0.4, 0.4), Vec3::new(0.4, 0.4, 1.0)],
272            grid_line_width: 0.02,
273        }
274    }
275
276    /// Returns the UV coordinates.
277    #[must_use]
278    pub fn coords(&self) -> &[Vec2] {
279        &self.coords
280    }
281
282    /// Gets the visualization style.
283    #[must_use]
284    pub fn style(&self) -> ParamVizStyle {
285        self.style
286    }
287
288    /// Sets the visualization style.
289    pub fn set_style(&mut self, style: ParamVizStyle) -> &mut Self {
290        self.style = style;
291        self
292    }
293
294    /// Gets the coordinate type.
295    #[must_use]
296    pub fn coords_type(&self) -> ParamCoordsType {
297        self.coords_type
298    }
299
300    /// Sets the coordinate type.
301    pub fn set_coords_type(&mut self, ct: ParamCoordsType) -> &mut Self {
302        self.coords_type = ct;
303        self
304    }
305
306    /// Gets the checker size.
307    #[must_use]
308    pub fn checker_size(&self) -> f32 {
309        self.checker_size
310    }
311
312    /// Sets the checker size.
313    pub fn set_checker_size(&mut self, size: f32) -> &mut Self {
314        self.checker_size = size;
315        self
316    }
317
318    /// Gets the checker colors.
319    #[must_use]
320    pub fn checker_colors(&self) -> [Vec3; 2] {
321        self.checker_colors
322    }
323
324    /// Sets the checker colors.
325    pub fn set_checker_colors(&mut self, colors: [Vec3; 2]) -> &mut Self {
326        self.checker_colors = colors;
327        self
328    }
329
330    /// Compute per-corner colors based on the current visualization style.
331    /// Returns one color per corner (same length as self.coords).
332    #[must_use]
333    pub fn compute_colors(&self) -> Vec<Vec4> {
334        self.coords
335            .iter()
336            .map(|uv| match self.style {
337                ParamVizStyle::Checker => {
338                    let u_cell = (uv.x / self.checker_size).floor() as i32;
339                    let v_cell = (uv.y / self.checker_size).floor() as i32;
340                    if (u_cell + v_cell) % 2 == 0 {
341                        self.checker_colors[0].extend(1.0)
342                    } else {
343                        self.checker_colors[1].extend(1.0)
344                    }
345                }
346                ParamVizStyle::Grid => {
347                    let u_frac = (uv.x / self.checker_size).fract().abs();
348                    let v_frac = (uv.y / self.checker_size).fract().abs();
349                    let on_line = u_frac < self.grid_line_width
350                        || u_frac > (1.0 - self.grid_line_width)
351                        || v_frac < self.grid_line_width
352                        || v_frac > (1.0 - self.grid_line_width);
353                    if on_line {
354                        self.checker_colors[1].extend(1.0)
355                    } else {
356                        self.checker_colors[0].extend(1.0)
357                    }
358                }
359                ParamVizStyle::LocalCheck => {
360                    let angle = uv.y.atan2(uv.x);
361                    let hue = (angle / std::f32::consts::TAU + 1.0) % 1.0;
362                    let base = hsv_to_rgb(hue, 0.7, 0.9);
363                    let u_cell = (uv.x / self.checker_size).floor() as i32;
364                    let v_cell = (uv.y / self.checker_size).floor() as i32;
365                    let dim = if (u_cell + v_cell) % 2 == 0 { 1.0 } else { 0.6 };
366                    let r = uv.length();
367                    (base * dim * (1.0 - (-r * 2.0).exp() * 0.5)).extend(1.0)
368                }
369                ParamVizStyle::LocalRad => {
370                    let angle = uv.y.atan2(uv.x);
371                    let hue = (angle / std::f32::consts::TAU + 1.0) % 1.0;
372                    let base = hsv_to_rgb(hue, 0.7, 0.9);
373                    let r = uv.length();
374                    let stripe =
375                        f32::from(u8::from((r / self.checker_size).floor() as i32 % 2 == 0));
376                    let dim = 0.6 + 0.4 * stripe;
377                    (base * dim).extend(1.0)
378                }
379            })
380            .collect()
381    }
382
383    /// Builds the egui UI for this quantity.
384    pub fn build_egui_ui(&mut self, ui: &mut egui::Ui) -> bool {
385        polyscope_ui::build_parameterization_quantity_ui(
386            ui,
387            &self.name,
388            &mut self.enabled,
389            &mut self.style,
390            &mut self.checker_size,
391            &mut self.checker_colors,
392        )
393    }
394}
395
396impl Quantity for MeshCornerParameterizationQuantity {
397    fn as_any(&self) -> &dyn std::any::Any {
398        self
399    }
400
401    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
402        self
403    }
404
405    fn name(&self) -> &str {
406        &self.name
407    }
408
409    fn structure_name(&self) -> &str {
410        &self.structure_name
411    }
412
413    fn kind(&self) -> QuantityKind {
414        QuantityKind::Parameterization
415    }
416
417    fn is_enabled(&self) -> bool {
418        self.enabled
419    }
420
421    fn set_enabled(&mut self, enabled: bool) {
422        self.enabled = enabled;
423    }
424
425    fn build_ui(&mut self, _ui: &dyn std::any::Any) {}
426
427    fn refresh(&mut self) {}
428
429    fn data_size(&self) -> usize {
430        self.coords.len()
431    }
432}
433
434impl FaceQuantity for MeshCornerParameterizationQuantity {}
435
436#[cfg(test)]
437mod tests {
438    use super::*;
439
440    #[test]
441    fn test_vertex_parameterization_creation() {
442        let coords = vec![
443            Vec2::new(0.0, 0.0),
444            Vec2::new(1.0, 0.0),
445            Vec2::new(0.5, 1.0),
446        ];
447        let q = MeshVertexParameterizationQuantity::new("uv", "mesh", coords);
448
449        assert_eq!(q.name(), "uv");
450        assert_eq!(q.structure_name(), "mesh");
451        assert_eq!(q.data_size(), 3);
452        assert_eq!(q.kind(), QuantityKind::Parameterization);
453        assert!(!q.is_enabled());
454        assert_eq!(q.style(), ParamVizStyle::Checker);
455        assert_eq!(q.coords_type(), ParamCoordsType::Unit);
456    }
457
458    #[test]
459    fn test_checker_colors_computation() {
460        let coords = vec![
461            Vec2::new(0.0, 0.0),   // cell (0,0) -> even -> color[0]
462            Vec2::new(0.15, 0.0),  // cell (1,0) -> odd -> color[1]
463            Vec2::new(0.0, 0.15),  // cell (0,1) -> odd -> color[1]
464            Vec2::new(0.15, 0.15), // cell (1,1) -> even -> color[0]
465        ];
466        let q = MeshVertexParameterizationQuantity::new("uv", "mesh", coords);
467
468        let colors = q.compute_colors();
469        assert_eq!(colors.len(), 4);
470        assert_eq!(colors[0], q.checker_colors()[0].extend(1.0));
471        assert_eq!(colors[1], q.checker_colors()[1].extend(1.0));
472        assert_eq!(colors[2], q.checker_colors()[1].extend(1.0));
473        assert_eq!(colors[3], q.checker_colors()[0].extend(1.0));
474    }
475
476    #[test]
477    fn test_grid_colors_computation() {
478        let coords = vec![
479            Vec2::new(0.001, 0.001), // near grid intersection -> on line
480            Vec2::new(0.05, 0.05),   // center of cell -> not on line
481        ];
482        let mut q = MeshVertexParameterizationQuantity::new("uv", "mesh", coords);
483        q.set_style(ParamVizStyle::Grid);
484
485        let colors = q.compute_colors();
486        assert_eq!(colors.len(), 2);
487        assert_eq!(colors[0], q.checker_colors()[1].extend(1.0)); // on line
488        assert_eq!(colors[1], q.checker_colors()[0].extend(1.0)); // off line
489    }
490
491    #[test]
492    fn test_local_check_colors_computation() {
493        let coords = vec![Vec2::new(0.5, 0.0), Vec2::new(0.0, 0.5)];
494        let mut q = MeshVertexParameterizationQuantity::new("uv", "mesh", coords);
495        q.set_style(ParamVizStyle::LocalCheck);
496
497        let colors = q.compute_colors();
498        assert_eq!(colors.len(), 2);
499        // Colors should be non-zero (derived from HSV)
500        assert!(colors[0].length() > 0.0);
501        assert!(colors[1].length() > 0.0);
502    }
503
504    #[test]
505    fn test_local_rad_colors_computation() {
506        let coords = vec![Vec2::new(0.5, 0.0), Vec2::new(0.0, 0.5)];
507        let mut q = MeshVertexParameterizationQuantity::new("uv", "mesh", coords);
508        q.set_style(ParamVizStyle::LocalRad);
509
510        let colors = q.compute_colors();
511        assert_eq!(colors.len(), 2);
512        assert!(colors[0].length() > 0.0);
513        assert!(colors[1].length() > 0.0);
514    }
515
516    #[test]
517    fn test_corner_parameterization_creation() {
518        // 1 triangle = 3 corners
519        let coords = vec![
520            Vec2::new(0.0, 0.0),
521            Vec2::new(1.0, 0.0),
522            Vec2::new(0.5, 1.0),
523        ];
524        let q = MeshCornerParameterizationQuantity::new("uv_corners", "mesh", coords);
525
526        assert_eq!(q.name(), "uv_corners");
527        assert_eq!(q.data_size(), 3);
528        assert_eq!(q.kind(), QuantityKind::Parameterization);
529    }
530
531    #[test]
532    fn test_corner_parameterization_compute_colors() {
533        let coords = vec![
534            Vec2::new(0.0, 0.0),
535            Vec2::new(0.15, 0.0),
536            Vec2::new(0.0, 0.15),
537        ];
538        let q = MeshCornerParameterizationQuantity::new("uv_corners", "mesh", coords);
539
540        let colors = q.compute_colors();
541        assert_eq!(colors.len(), 3);
542        assert_eq!(colors[0], q.checker_colors()[0].extend(1.0)); // cell (0,0) even
543        assert_eq!(colors[1], q.checker_colors()[1].extend(1.0)); // cell (1,0) odd
544        assert_eq!(colors[2], q.checker_colors()[1].extend(1.0)); // cell (0,1) odd
545    }
546
547    #[test]
548    fn test_hsv_to_rgb() {
549        // Red
550        let red = hsv_to_rgb(0.0, 1.0, 1.0);
551        assert!((red.x - 1.0).abs() < 1e-5);
552        assert!(red.y.abs() < 1e-5);
553        assert!(red.z.abs() < 1e-5);
554
555        // Green
556        let green = hsv_to_rgb(1.0 / 3.0, 1.0, 1.0);
557        assert!(green.x.abs() < 1e-5);
558        assert!((green.y - 1.0).abs() < 1e-5);
559        assert!(green.z.abs() < 1e-5);
560    }
561}