1mod quantities;
4
5use glam::{Mat4, Vec3, Vec4};
6use polyscope_core::pick::PickResult;
7use polyscope_core::quantity::Quantity;
8use polyscope_core::structure::{HasQuantities, RenderContext, Structure};
9use polyscope_render::{
10 ColorMapRegistry, CurveNetworkRenderData, CurveNetworkUniforms, PickUniforms, PointUniforms,
11 TubePickUniforms,
12};
13use wgpu::util::DeviceExt;
14
15pub use quantities::*;
16
17pub struct CurveNetwork {
19 name: String,
20
21 node_positions: Vec<Vec3>,
23 edge_tail_inds: Vec<u32>,
24 edge_tip_inds: Vec<u32>,
25
26 edge_centers: Vec<Vec3>,
28 node_degrees: Vec<usize>,
29
30 enabled: bool,
32 transform: Mat4,
33 quantities: Vec<Box<dyn Quantity>>,
34
35 color: Vec4,
37 radius: f32,
38 radius_is_relative: bool,
39 material: String,
40 render_mode: u32,
42
43 #[allow(dead_code)]
45 node_radius_quantity_name: Option<String>,
46 #[allow(dead_code)]
47 edge_radius_quantity_name: Option<String>,
48 #[allow(dead_code)]
49 node_radius_autoscale: bool,
50 #[allow(dead_code)]
51 edge_radius_autoscale: bool,
52
53 render_data: Option<CurveNetworkRenderData>,
55
56 pick_uniform_buffer: Option<wgpu::Buffer>,
58 pick_bind_group: Option<wgpu::BindGroup>,
59 global_start: u32,
60
61 tube_pick_uniform_buffer: Option<wgpu::Buffer>,
63 tube_pick_bind_group: Option<wgpu::BindGroup>,
64}
65
66impl CurveNetwork {
67 pub fn new(name: impl Into<String>, nodes: Vec<Vec3>, edges: Vec<[u32; 2]>) -> Self {
69 let edge_tail_inds: Vec<u32> = edges.iter().map(|e| e[0]).collect();
70 let edge_tip_inds: Vec<u32> = edges.iter().map(|e| e[1]).collect();
71
72 let mut cn = Self {
73 name: name.into(),
74 node_positions: nodes,
75 edge_tail_inds,
76 edge_tip_inds,
77 edge_centers: Vec::new(),
78 node_degrees: Vec::new(),
79 enabled: true,
80 transform: Mat4::IDENTITY,
81 quantities: Vec::new(),
82 color: Vec4::new(0.2, 0.5, 0.8, 1.0),
83 radius: 0.005,
84 radius_is_relative: true,
85 material: "default".to_string(),
86 render_mode: 0, node_radius_quantity_name: None,
88 edge_radius_quantity_name: None,
89 node_radius_autoscale: true,
90 edge_radius_autoscale: true,
91 render_data: None,
92 pick_uniform_buffer: None,
93 pick_bind_group: None,
94 global_start: 0,
95 tube_pick_uniform_buffer: None,
96 tube_pick_bind_group: None,
97 };
98 cn.recompute_geometry();
99 cn
100 }
101
102 pub fn new_line(name: impl Into<String>, nodes: Vec<Vec3>) -> Self {
104 let n = nodes.len();
105 let edges: Vec<[u32; 2]> = (0..n.saturating_sub(1))
106 .map(|i| [i as u32, (i + 1) as u32])
107 .collect();
108 Self::new(name, nodes, edges)
109 }
110
111 pub fn new_loop(name: impl Into<String>, nodes: Vec<Vec3>) -> Self {
113 let n = nodes.len();
114 let edges: Vec<[u32; 2]> = (0..n).map(|i| [i as u32, ((i + 1) % n) as u32]).collect();
115 Self::new(name, nodes, edges)
116 }
117
118 pub fn new_segments(name: impl Into<String>, nodes: Vec<Vec3>) -> Self {
120 let n = nodes.len();
121 let edges: Vec<[u32; 2]> = (0..n / 2)
122 .map(|i| [(i * 2) as u32, (i * 2 + 1) as u32])
123 .collect();
124 Self::new(name, nodes, edges)
125 }
126
127 #[must_use]
129 pub fn name(&self) -> &str {
130 &self.name
131 }
132
133 #[must_use]
135 pub fn num_nodes(&self) -> usize {
136 self.node_positions.len()
137 }
138
139 #[must_use]
141 pub fn num_edges(&self) -> usize {
142 self.edge_tail_inds.len()
143 }
144
145 #[must_use]
147 pub fn nodes(&self) -> &[Vec3] {
148 &self.node_positions
149 }
150
151 #[must_use]
153 pub fn edge_tail_inds(&self) -> &[u32] {
154 &self.edge_tail_inds
155 }
156
157 #[must_use]
159 pub fn edge_tip_inds(&self) -> &[u32] {
160 &self.edge_tip_inds
161 }
162
163 #[must_use]
165 pub fn edge_centers(&self) -> &[Vec3] {
166 &self.edge_centers
167 }
168
169 #[must_use]
171 pub fn node_degrees(&self) -> &[usize] {
172 &self.node_degrees
173 }
174
175 #[must_use]
177 pub fn color(&self) -> Vec4 {
178 self.color
179 }
180
181 pub fn set_color(&mut self, color: Vec3) -> &mut Self {
183 self.color = color.extend(1.0);
184 self
185 }
186
187 #[must_use]
189 pub fn radius(&self) -> f32 {
190 self.radius
191 }
192
193 #[must_use]
195 pub fn radius_is_relative(&self) -> bool {
196 self.radius_is_relative
197 }
198
199 pub fn set_radius(&mut self, radius: f32, is_relative: bool) -> &mut Self {
201 self.radius = radius;
202 self.radius_is_relative = is_relative;
203 self
204 }
205
206 #[must_use]
208 pub fn material(&self) -> &str {
209 &self.material
210 }
211
212 pub fn set_material(&mut self, material: impl Into<String>) -> &mut Self {
214 self.material = material.into();
215 self
216 }
217
218 #[must_use]
220 pub fn render_mode(&self) -> u32 {
221 self.render_mode
222 }
223
224 pub fn set_render_mode(&mut self, mode: u32) -> &mut Self {
226 self.render_mode = mode.min(1); self
228 }
229
230 pub fn update_node_positions(&mut self, nodes: Vec<Vec3>) {
232 self.node_positions = nodes;
233 self.recompute_geometry();
234 self.refresh();
235 }
236
237 pub fn add_node_scalar_quantity(
239 &mut self,
240 name: impl Into<String>,
241 values: Vec<f32>,
242 ) -> &mut Self {
243 let quantity = CurveNodeScalarQuantity::new(name, self.name.clone(), values);
244 self.add_quantity(Box::new(quantity));
245 self
246 }
247
248 pub fn add_edge_scalar_quantity(
250 &mut self,
251 name: impl Into<String>,
252 values: Vec<f32>,
253 ) -> &mut Self {
254 let quantity = CurveEdgeScalarQuantity::new(name, self.name.clone(), values);
255 self.add_quantity(Box::new(quantity));
256 self
257 }
258
259 pub fn add_node_color_quantity(
261 &mut self,
262 name: impl Into<String>,
263 colors: Vec<Vec3>,
264 ) -> &mut Self {
265 let quantity = CurveNodeColorQuantity::new(name, self.name.clone(), colors);
266 self.add_quantity(Box::new(quantity));
267 self
268 }
269
270 pub fn add_edge_color_quantity(
272 &mut self,
273 name: impl Into<String>,
274 colors: Vec<Vec3>,
275 ) -> &mut Self {
276 let quantity = CurveEdgeColorQuantity::new(name, self.name.clone(), colors);
277 self.add_quantity(Box::new(quantity));
278 self
279 }
280
281 pub fn add_node_vector_quantity(
283 &mut self,
284 name: impl Into<String>,
285 vectors: Vec<Vec3>,
286 ) -> &mut Self {
287 let quantity = CurveNodeVectorQuantity::new(name, self.name.clone(), vectors);
288 self.add_quantity(Box::new(quantity));
289 self
290 }
291
292 pub fn add_edge_vector_quantity(
294 &mut self,
295 name: impl Into<String>,
296 vectors: Vec<Vec3>,
297 ) -> &mut Self {
298 let quantity = CurveEdgeVectorQuantity::new(name, self.name.clone(), vectors);
299 self.add_quantity(Box::new(quantity));
300 self
301 }
302
303 #[must_use]
305 pub fn active_node_scalar_quantity(&self) -> Option<&CurveNodeScalarQuantity> {
306 use polyscope_core::quantity::QuantityKind;
307
308 for q in &self.quantities {
309 if q.is_enabled() && q.kind() == QuantityKind::Scalar {
310 if let Some(sq) = q.as_any().downcast_ref::<CurveNodeScalarQuantity>() {
311 return Some(sq);
312 }
313 }
314 }
315 None
316 }
317
318 #[must_use]
320 pub fn active_edge_scalar_quantity(&self) -> Option<&CurveEdgeScalarQuantity> {
321 use polyscope_core::quantity::QuantityKind;
322
323 for q in &self.quantities {
324 if q.is_enabled() && q.kind() == QuantityKind::Scalar {
325 if let Some(sq) = q.as_any().downcast_ref::<CurveEdgeScalarQuantity>() {
326 return Some(sq);
327 }
328 }
329 }
330 None
331 }
332
333 #[must_use]
335 pub fn active_node_color_quantity(&self) -> Option<&CurveNodeColorQuantity> {
336 use polyscope_core::quantity::QuantityKind;
337
338 for q in &self.quantities {
339 if q.is_enabled() && q.kind() == QuantityKind::Color {
340 if let Some(cq) = q.as_any().downcast_ref::<CurveNodeColorQuantity>() {
341 return Some(cq);
342 }
343 }
344 }
345 None
346 }
347
348 #[must_use]
350 pub fn active_edge_color_quantity(&self) -> Option<&CurveEdgeColorQuantity> {
351 use polyscope_core::quantity::QuantityKind;
352
353 for q in &self.quantities {
354 if q.is_enabled() && q.kind() == QuantityKind::Color {
355 if let Some(cq) = q.as_any().downcast_ref::<CurveEdgeColorQuantity>() {
356 return Some(cq);
357 }
358 }
359 }
360 None
361 }
362
363 pub fn build_egui_ui(&mut self, ui: &mut egui::Ui, available_materials: &[&str]) {
365 let mut color = [self.color.x, self.color.y, self.color.z];
366 let mut radius = self.radius;
367 let mut radius_is_relative = self.radius_is_relative;
368 let mut render_mode = self.render_mode;
369
370 if polyscope_ui::build_curve_network_ui(
371 ui,
372 self.node_positions.len(),
373 self.edge_tail_inds.len(),
374 &mut radius,
375 &mut radius_is_relative,
376 &mut color,
377 &mut render_mode,
378 &mut self.material,
379 available_materials,
380 ) {
381 self.color = Vec4::new(color[0], color[1], color[2], self.color.w);
382 self.radius = radius;
383 self.radius_is_relative = radius_is_relative;
384 self.render_mode = render_mode;
385 }
386
387 if !self.quantities.is_empty() {
389 ui.separator();
390 ui.label("Quantities:");
391 for quantity in &mut self.quantities {
392 if let Some(sq) = quantity
394 .as_any_mut()
395 .downcast_mut::<CurveNodeScalarQuantity>()
396 {
397 sq.build_egui_ui(ui);
398 } else if let Some(sq) = quantity
399 .as_any_mut()
400 .downcast_mut::<CurveEdgeScalarQuantity>()
401 {
402 sq.build_egui_ui(ui);
403 } else if let Some(cq) = quantity
404 .as_any_mut()
405 .downcast_mut::<CurveNodeColorQuantity>()
406 {
407 cq.build_egui_ui(ui);
408 } else if let Some(cq) = quantity
409 .as_any_mut()
410 .downcast_mut::<CurveEdgeColorQuantity>()
411 {
412 cq.build_egui_ui(ui);
413 } else if let Some(vq) = quantity
414 .as_any_mut()
415 .downcast_mut::<CurveNodeVectorQuantity>()
416 {
417 vq.build_egui_ui(ui);
418 } else if let Some(vq) = quantity
419 .as_any_mut()
420 .downcast_mut::<CurveEdgeVectorQuantity>()
421 {
422 vq.build_egui_ui(ui);
423 }
424 }
425 }
426 }
427
428 pub fn init_gpu_resources(
430 &mut self,
431 device: &wgpu::Device,
432 bind_group_layout: &wgpu::BindGroupLayout,
433 camera_buffer: &wgpu::Buffer,
434 ) {
435 self.render_data = Some(CurveNetworkRenderData::new(
436 device,
437 bind_group_layout,
438 camera_buffer,
439 &self.node_positions,
440 &self.edge_tail_inds,
441 &self.edge_tip_inds,
442 ));
443 }
444
445 #[must_use]
447 pub fn render_data(&self) -> Option<&CurveNetworkRenderData> {
448 self.render_data.as_ref()
449 }
450
451 pub fn init_tube_resources(
453 &mut self,
454 device: &wgpu::Device,
455 compute_bind_group_layout: &wgpu::BindGroupLayout,
456 render_bind_group_layout: &wgpu::BindGroupLayout,
457 camera_buffer: &wgpu::Buffer,
458 ) {
459 if let Some(render_data) = &mut self.render_data {
460 render_data.init_tube_resources(
461 device,
462 compute_bind_group_layout,
463 render_bind_group_layout,
464 camera_buffer,
465 );
466 }
467 }
468
469 pub fn init_node_render_resources(
471 &mut self,
472 device: &wgpu::Device,
473 point_bind_group_layout: &wgpu::BindGroupLayout,
474 camera_buffer: &wgpu::Buffer,
475 ) {
476 if let Some(render_data) = &mut self.render_data {
477 render_data.init_node_render_resources(device, point_bind_group_layout, camera_buffer);
478 }
479 }
480
481 pub fn init_pick_resources(
483 &mut self,
484 device: &wgpu::Device,
485 pick_bind_group_layout: &wgpu::BindGroupLayout,
486 camera_buffer: &wgpu::Buffer,
487 global_start: u32,
488 ) {
489 self.global_start = global_start;
490
491 let pick_uniforms = PickUniforms {
493 global_start,
494 point_radius: self.radius, _padding: [0.0; 2],
496 };
497 let pick_uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
498 label: Some("curve network pick uniforms"),
499 contents: bytemuck::cast_slice(&[pick_uniforms]),
500 usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
501 });
502
503 if let Some(render_data) = &self.render_data {
505 let pick_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
506 label: Some("curve network pick bind group"),
507 layout: pick_bind_group_layout,
508 entries: &[
509 wgpu::BindGroupEntry {
510 binding: 0,
511 resource: camera_buffer.as_entire_binding(),
512 },
513 wgpu::BindGroupEntry {
514 binding: 1,
515 resource: pick_uniform_buffer.as_entire_binding(),
516 },
517 wgpu::BindGroupEntry {
518 binding: 2,
519 resource: render_data.edge_vertex_buffer.as_entire_binding(),
520 },
521 ],
522 });
523 self.pick_bind_group = Some(pick_bind_group);
524 }
525
526 self.pick_uniform_buffer = Some(pick_uniform_buffer);
527 }
528
529 #[must_use]
531 pub fn pick_bind_group(&self) -> Option<&wgpu::BindGroup> {
532 self.pick_bind_group.as_ref()
533 }
534
535 pub fn update_pick_uniforms(&self, queue: &wgpu::Queue) {
537 if let Some(buffer) = &self.pick_uniform_buffer {
538 let pick_uniforms = PickUniforms {
539 global_start: self.global_start,
540 point_radius: self.radius,
541 _padding: [0.0; 2],
542 };
543 queue.write_buffer(buffer, 0, bytemuck::cast_slice(&[pick_uniforms]));
544 }
545
546 if let Some(buffer) = &self.tube_pick_uniform_buffer {
548 let tube_pick_uniforms = TubePickUniforms {
549 global_start: self.global_start,
550 radius: self.radius,
551 min_pick_radius: 0.02, _padding: 0.0,
553 };
554 queue.write_buffer(buffer, 0, bytemuck::cast_slice(&[tube_pick_uniforms]));
555 }
556 }
557
558 pub fn init_tube_pick_resources(
561 &mut self,
562 device: &wgpu::Device,
563 tube_pick_bind_group_layout: &wgpu::BindGroupLayout,
564 camera_buffer: &wgpu::Buffer,
565 ) {
566 let tube_pick_uniforms = TubePickUniforms {
568 global_start: self.global_start,
569 radius: self.radius,
570 min_pick_radius: 0.02, _padding: 0.0,
572 };
573 let tube_pick_uniform_buffer =
574 device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
575 label: Some("curve network tube pick uniforms"),
576 contents: bytemuck::cast_slice(&[tube_pick_uniforms]),
577 usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
578 });
579
580 if let Some(render_data) = &self.render_data {
582 let tube_pick_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
583 label: Some("curve network tube pick bind group"),
584 layout: tube_pick_bind_group_layout,
585 entries: &[
586 wgpu::BindGroupEntry {
587 binding: 0,
588 resource: camera_buffer.as_entire_binding(),
589 },
590 wgpu::BindGroupEntry {
591 binding: 1,
592 resource: tube_pick_uniform_buffer.as_entire_binding(),
593 },
594 wgpu::BindGroupEntry {
595 binding: 2,
596 resource: render_data.edge_vertex_buffer.as_entire_binding(),
597 },
598 ],
599 });
600 self.tube_pick_bind_group = Some(tube_pick_bind_group);
601 }
602
603 self.tube_pick_uniform_buffer = Some(tube_pick_uniform_buffer);
604 }
605
606 #[must_use]
608 pub fn tube_pick_bind_group(&self) -> Option<&wgpu::BindGroup> {
609 self.tube_pick_bind_group.as_ref()
610 }
611
612 #[must_use]
614 pub fn has_tube_pick_resources(&self) -> bool {
615 self.tube_pick_bind_group.is_some()
616 }
617
618 pub fn update_gpu_buffers(&self, queue: &wgpu::Queue, color_maps: &ColorMapRegistry) {
620 let Some(render_data) = &self.render_data else {
621 return;
622 };
623
624 let transformed_nodes: Vec<Vec3> = self
626 .node_positions
627 .iter()
628 .map(|p| (self.transform * p.extend(1.0)).truncate())
629 .collect();
630 render_data.update_node_positions(queue, &transformed_nodes);
631 render_data.update_edge_vertices(
632 queue,
633 &transformed_nodes,
634 &self.edge_tail_inds,
635 &self.edge_tip_inds,
636 );
637
638 let uniforms = CurveNetworkUniforms {
639 color: self.color.to_array(),
640 radius: self.radius,
641 radius_is_relative: u32::from(self.radius_is_relative),
642 render_mode: self.render_mode,
643 _padding: 0.0,
644 };
645
646 render_data.update_uniforms(queue, &uniforms);
647
648 if self.render_mode == 1 && render_data.has_node_render_resources() {
650 let model_matrix = self.transform.to_cols_array_2d();
651 let node_uniforms = PointUniforms {
652 model_matrix,
653 point_radius: self.radius * 1.02,
655 use_per_point_color: 0, _padding: [0.0; 2],
657 base_color: self.color.to_array(),
658 };
659 render_data.update_node_uniforms(queue, &node_uniforms);
660 }
661
662 if let Some(color_q) = self.active_node_color_quantity() {
664 color_q.apply_to_render_data(queue, render_data);
665 } else if let Some(scalar_q) = self.active_node_scalar_quantity() {
666 if let Some(colormap) = color_maps.get(scalar_q.colormap_name()) {
667 let colors = scalar_q.compute_colors(colormap);
668 render_data.update_node_colors(queue, &colors);
669 }
670 }
671
672 if let Some(color_q) = self.active_edge_color_quantity() {
674 color_q.apply_to_render_data(queue, render_data);
675 } else if let Some(scalar_q) = self.active_edge_scalar_quantity() {
676 if let Some(colormap) = color_maps.get(scalar_q.colormap_name()) {
677 let colors = scalar_q.compute_colors(colormap);
678 render_data.update_edge_colors(queue, &colors);
679 }
680 }
681 }
682
683 fn recompute_geometry(&mut self) {
685 self.edge_centers.clear();
687 for i in 0..self.edge_tail_inds.len() {
688 let tail = self.node_positions[self.edge_tail_inds[i] as usize];
689 let tip = self.node_positions[self.edge_tip_inds[i] as usize];
690 self.edge_centers.push((tail + tip) * 0.5);
691 }
692
693 self.node_degrees = vec![0; self.node_positions.len()];
695 for &tail in &self.edge_tail_inds {
696 self.node_degrees[tail as usize] += 1;
697 }
698 for &tip in &self.edge_tip_inds {
699 self.node_degrees[tip as usize] += 1;
700 }
701 }
702}
703
704impl Structure for CurveNetwork {
705 fn as_any(&self) -> &dyn std::any::Any {
706 self
707 }
708
709 fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
710 self
711 }
712
713 fn name(&self) -> &str {
714 &self.name
715 }
716
717 fn type_name(&self) -> &'static str {
718 "CurveNetwork"
719 }
720
721 fn bounding_box(&self) -> Option<(Vec3, Vec3)> {
722 if self.node_positions.is_empty() {
723 return None;
724 }
725
726 let mut min = Vec3::splat(f32::MAX);
727 let mut max = Vec3::splat(f32::MIN);
728
729 for &p in &self.node_positions {
730 min = min.min(p);
731 max = max.max(p);
732 }
733
734 let transform = self.transform;
736 let corners = [
737 transform.transform_point3(Vec3::new(min.x, min.y, min.z)),
738 transform.transform_point3(Vec3::new(max.x, min.y, min.z)),
739 transform.transform_point3(Vec3::new(min.x, max.y, min.z)),
740 transform.transform_point3(Vec3::new(max.x, max.y, min.z)),
741 transform.transform_point3(Vec3::new(min.x, min.y, max.z)),
742 transform.transform_point3(Vec3::new(max.x, min.y, max.z)),
743 transform.transform_point3(Vec3::new(min.x, max.y, max.z)),
744 transform.transform_point3(Vec3::new(max.x, max.y, max.z)),
745 ];
746
747 let mut world_min = Vec3::splat(f32::MAX);
748 let mut world_max = Vec3::splat(f32::MIN);
749 for corner in corners {
750 world_min = world_min.min(corner);
751 world_max = world_max.max(corner);
752 }
753
754 Some((world_min, world_max))
755 }
756
757 fn length_scale(&self) -> f32 {
758 self.bounding_box()
759 .map_or(1.0, |(min, max)| (max - min).length())
760 }
761
762 fn transform(&self) -> Mat4 {
763 self.transform
764 }
765
766 fn set_transform(&mut self, transform: Mat4) {
767 self.transform = transform;
768 }
769
770 fn is_enabled(&self) -> bool {
771 self.enabled
772 }
773
774 fn set_enabled(&mut self, enabled: bool) {
775 self.enabled = enabled;
776 }
777
778 fn material(&self) -> &str {
779 &self.material
780 }
781
782 fn set_material(&mut self, material_name: &str) {
783 self.material = material_name.to_string();
784 }
785
786 fn draw(&self, _ctx: &mut dyn RenderContext) {
787 }
789
790 fn draw_pick(&self, _ctx: &mut dyn RenderContext) {
791 }
793
794 fn build_ui(&mut self, _ui: &dyn std::any::Any) {
795 }
797
798 fn build_pick_ui(&self, _ui: &dyn std::any::Any, _pick: &PickResult) {
799 }
801
802 fn clear_gpu_resources(&mut self) {
803 self.render_data = None;
804 self.pick_uniform_buffer = None;
805 self.pick_bind_group = None;
806 self.tube_pick_uniform_buffer = None;
807 self.tube_pick_bind_group = None;
808 for quantity in &mut self.quantities {
809 quantity.clear_gpu_resources();
810 }
811 }
812
813 fn refresh(&mut self) {
814 self.recompute_geometry();
815 for quantity in &mut self.quantities {
816 quantity.refresh();
817 }
818 }
819}
820
821impl HasQuantities for CurveNetwork {
822 fn add_quantity(&mut self, quantity: Box<dyn Quantity>) {
823 self.quantities.push(quantity);
824 }
825
826 fn get_quantity(&self, name: &str) -> Option<&dyn Quantity> {
827 self.quantities
828 .iter()
829 .find(|q| q.name() == name)
830 .map(std::convert::AsRef::as_ref)
831 }
832
833 fn get_quantity_mut(&mut self, name: &str) -> Option<&mut Box<dyn Quantity>> {
834 self.quantities.iter_mut().find(|q| q.name() == name)
835 }
836
837 fn remove_quantity(&mut self, name: &str) -> Option<Box<dyn Quantity>> {
838 let idx = self.quantities.iter().position(|q| q.name() == name)?;
839 Some(self.quantities.remove(idx))
840 }
841
842 fn quantities(&self) -> &[Box<dyn Quantity>] {
843 &self.quantities
844 }
845}
846
847#[cfg(test)]
848mod tests {
849 use super::*;
850
851 #[test]
852 fn test_curve_network_creation() {
853 let nodes = vec![
854 Vec3::new(0.0, 0.0, 0.0),
855 Vec3::new(1.0, 0.0, 0.0),
856 Vec3::new(1.0, 1.0, 0.0),
857 ];
858 let edges = vec![[0, 1], [1, 2]];
859
860 let cn = CurveNetwork::new("test", nodes.clone(), edges);
861
862 assert_eq!(cn.name(), "test");
863 assert_eq!(cn.num_nodes(), 3);
864 assert_eq!(cn.num_edges(), 2);
865 assert_eq!(cn.nodes(), &nodes);
866 }
867
868 #[test]
869 fn test_curve_network_line_connectivity() {
870 let nodes = vec![
871 Vec3::new(0.0, 0.0, 0.0),
872 Vec3::new(1.0, 0.0, 0.0),
873 Vec3::new(2.0, 0.0, 0.0),
874 Vec3::new(3.0, 0.0, 0.0),
875 ];
876
877 let cn = CurveNetwork::new_line("line", nodes);
878
879 assert_eq!(cn.num_edges(), 3); assert_eq!(cn.edge_tail_inds(), &[0, 1, 2]);
881 assert_eq!(cn.edge_tip_inds(), &[1, 2, 3]);
882 }
883
884 #[test]
885 fn test_curve_network_loop_connectivity() {
886 let nodes = vec![
887 Vec3::new(0.0, 0.0, 0.0),
888 Vec3::new(1.0, 0.0, 0.0),
889 Vec3::new(1.0, 1.0, 0.0),
890 ];
891
892 let cn = CurveNetwork::new_loop("loop", nodes);
893
894 assert_eq!(cn.num_edges(), 3); assert_eq!(cn.edge_tail_inds(), &[0, 1, 2]);
896 assert_eq!(cn.edge_tip_inds(), &[1, 2, 0]);
897 }
898
899 #[test]
900 fn test_curve_network_segments_connectivity() {
901 let nodes = vec![
902 Vec3::new(0.0, 0.0, 0.0),
903 Vec3::new(1.0, 0.0, 0.0),
904 Vec3::new(2.0, 0.0, 0.0),
905 Vec3::new(3.0, 0.0, 0.0),
906 ];
907
908 let cn = CurveNetwork::new_segments("segments", nodes);
909
910 assert_eq!(cn.num_edges(), 2); assert_eq!(cn.edge_tail_inds(), &[0, 2]);
912 assert_eq!(cn.edge_tip_inds(), &[1, 3]);
913 }
914
915 #[test]
916 fn test_curve_network_visualization_params() {
917 let nodes = vec![Vec3::ZERO, Vec3::X];
918 let edges = vec![[0, 1]];
919 let mut cn = CurveNetwork::new("test", nodes, edges);
920
921 assert_eq!(cn.color(), Vec4::new(0.2, 0.5, 0.8, 1.0));
923 assert_eq!(cn.radius(), 0.005);
924 assert!(cn.radius_is_relative());
925
926 cn.set_color(Vec3::new(1.0, 0.0, 0.0));
928 assert_eq!(cn.color(), Vec4::new(1.0, 0.0, 0.0, 1.0));
929
930 cn.set_radius(0.1, false);
931 assert_eq!(cn.radius(), 0.1);
932 assert!(!cn.radius_is_relative());
933
934 assert_eq!(cn.material(), "default");
936 cn.set_material("clay");
937 assert_eq!(cn.material(), "clay");
938 }
939
940 #[test]
941 fn test_curve_network_vector_quantities() {
942 use polyscope_core::quantity::QuantityKind;
943
944 let nodes = vec![Vec3::ZERO, Vec3::X, Vec3::Y];
945 let edges = vec![[0, 1], [1, 2]];
946 let mut cn = CurveNetwork::new("test", nodes, edges);
947
948 cn.add_node_vector_quantity("node_vecs", vec![Vec3::X, Vec3::Y, Vec3::Z]);
949 cn.add_edge_vector_quantity("edge_vecs", vec![Vec3::X, Vec3::Y]);
950
951 let nq = cn.get_quantity("node_vecs").expect("node vector not found");
952 assert_eq!(nq.data_size(), 3);
953 assert_eq!(nq.kind(), QuantityKind::Vector);
954
955 let eq = cn.get_quantity("edge_vecs").expect("edge vector not found");
956 assert_eq!(eq.data_size(), 2);
957 assert_eq!(eq.kind(), QuantityKind::Vector);
958 }
959}