1use crate::context::shared_wgpu_context;
6use crate::core::{
7 BoundingBox, DrawCall, GpuVertexBuffer, Material, PipelineType, RenderData, Vertex,
8};
9use crate::gpu::axis::OwnedAxisData;
10use crate::gpu::{util::readback_scalar_buffer_f64, ScalarType};
11use glam::{Vec3, Vec4};
12use std::sync::Arc;
13
14#[derive(Debug, Clone)]
16pub struct SurfacePlot {
17 pub x_data: Vec<f64>,
19 pub y_data: Vec<f64>,
20 pub z_data: Option<Vec<Vec<f64>>>, x_len: usize,
23 y_len: usize,
24
25 pub colormap: ColorMap,
27 pub shading_mode: ShadingMode,
28 pub wireframe: bool,
29 pub alpha: f32,
30 pub flatten_z: bool,
32
33 pub image_mode: bool,
35
36 pub color_limits: Option<(f64, f64)>,
38
39 pub color_grid: Option<Vec<Vec<Vec4>>>, pub lighting_enabled: bool,
44 pub ambient_strength: f32,
45 pub diffuse_strength: f32,
46 pub specular_strength: f32,
47 pub shininess: f32,
48
49 pub label: Option<String>,
51 pub visible: bool,
52
53 vertices: Option<Vec<Vertex>>,
55 indices: Option<Vec<u32>>,
56 bounds: Option<BoundingBox>,
57 dirty: bool,
58 gpu_vertices: Option<GpuVertexBuffer>,
59 gpu_vertex_count: Option<usize>,
60 gpu_bounds: Option<BoundingBox>,
61 gpu_source: Option<SurfaceGpuSource>,
62 gpu_color_grid_source: Option<SurfaceGpuColorGridSource>,
63}
64
65#[derive(Clone, Debug)]
66pub struct SurfaceGpuSource {
67 pub x_axis: OwnedAxisData,
68 pub y_axis: OwnedAxisData,
69 pub z_buffer: Arc<wgpu::Buffer>,
70 pub x_len: usize,
71 pub y_len: usize,
72 pub scalar: ScalarType,
73}
74
75#[derive(Clone, Debug)]
76pub struct SurfaceGpuColorGridSource {
77 pub image_buffer: Arc<wgpu::Buffer>,
78 pub rows: usize,
79 pub cols: usize,
80 pub channels: usize,
81 pub scalar: ScalarType,
82}
83
84#[derive(Debug, Clone, Copy, PartialEq)]
86pub enum ColorMap {
87 Jet,
89 Hot,
90 Cool,
91 Spring,
92 Summer,
93 Autumn,
94 Winter,
95 Gray,
96 Bone,
97 Copper,
98 Pink,
99 Lines,
100
101 Viridis,
103 Plasma,
104 Inferno,
105 Magma,
106 Turbo,
107
108 Parula,
110
111 Custom(Vec4, Vec4), }
114
115impl ColorMap {
116 pub const CANONICAL_NAMES: &[&str] = &[
117 "parula", "viridis", "plasma", "inferno", "magma", "turbo", "jet", "hot", "cool", "spring",
118 "summer", "autumn", "winter", "gray", "bone", "copper", "pink", "lines",
119 ];
120
121 pub const ALIASES: &[&str] = &["grey"];
122
123 pub fn from_name(name: &str) -> Option<Self> {
124 match name.trim().to_ascii_lowercase().as_str() {
125 "parula" => Some(Self::Parula),
126 "viridis" => Some(Self::Viridis),
127 "plasma" => Some(Self::Plasma),
128 "inferno" => Some(Self::Inferno),
129 "magma" => Some(Self::Magma),
130 "turbo" => Some(Self::Turbo),
131 "jet" => Some(Self::Jet),
132 "hot" => Some(Self::Hot),
133 "cool" => Some(Self::Cool),
134 "spring" => Some(Self::Spring),
135 "summer" => Some(Self::Summer),
136 "autumn" => Some(Self::Autumn),
137 "winter" => Some(Self::Winter),
138 "gray" | "grey" => Some(Self::Gray),
139 "bone" => Some(Self::Bone),
140 "copper" => Some(Self::Copper),
141 "pink" => Some(Self::Pink),
142 "lines" => Some(Self::Lines),
143 _ => None,
144 }
145 }
146}
147
148#[derive(Debug, Clone, Copy, PartialEq)]
150pub enum ShadingMode {
151 Flat,
153 Smooth,
155 Faceted,
157 None,
159}
160
161impl Default for ColorMap {
162 fn default() -> Self {
163 Self::Viridis
164 }
165}
166
167impl Default for ShadingMode {
168 fn default() -> Self {
169 Self::Smooth
170 }
171}
172
173impl SurfacePlot {
174 pub async fn export_scene_grid_data(
175 &self,
176 ) -> Result<(Vec<f64>, Vec<f64>, Vec<Vec<f64>>), String> {
177 if let Some(z) = &self.z_data {
178 return Ok((self.x_data.clone(), self.y_data.clone(), z.clone()));
179 }
180
181 if let Some(source) = &self.gpu_source {
182 let context = shared_wgpu_context().ok_or_else(|| {
183 "surface plot has GPU source data but no shared WGPU context is installed"
184 .to_string()
185 })?;
186 let x = source
187 .x_axis
188 .export_f64(&context.device, &context.queue, source.x_len, source.scalar)
189 .await?;
190 let y = source
191 .y_axis
192 .export_f64(&context.device, &context.queue, source.y_len, source.scalar)
193 .await?;
194 let z_flat = readback_scalar_buffer_f64(
195 &context.device,
196 &context.queue,
197 &source.z_buffer,
198 source.x_len * source.y_len,
199 source.scalar,
200 )
201 .await?;
202 let mut z = Vec::with_capacity(source.x_len);
203 for row in 0..source.x_len {
204 let start = row * source.y_len;
205 z.push(
206 z_flat
207 .get(start..start + source.y_len)
208 .ok_or_else(|| "surface GPU source grid is out of range".to_string())?
209 .to_vec(),
210 );
211 }
212 return Ok((x, y, z));
213 }
214
215 if self.gpu_vertices.is_some() {
216 return Err(
217 "surface plot has GPU render vertices but no exportable source data".to_string(),
218 );
219 }
220
221 Ok((Vec::new(), Vec::new(), Vec::new()))
222 }
223
224 pub async fn export_scene_color_grid(&self) -> Result<Option<Vec<Vec<Vec4>>>, String> {
225 if let Some(grid) = &self.color_grid {
226 return Ok(Some(grid.clone()));
227 }
228
229 let Some(source) = &self.gpu_color_grid_source else {
230 return Ok(None);
231 };
232 let context = shared_wgpu_context().ok_or_else(|| {
233 "surface image has GPU color data but no shared WGPU context is installed".to_string()
234 })?;
235 let values = readback_scalar_buffer_f64(
236 &context.device,
237 &context.queue,
238 &source.image_buffer,
239 source.rows * source.cols * source.channels,
240 source.scalar,
241 )
242 .await?;
243 let mut grid = vec![vec![Vec4::ZERO; source.rows]; source.cols];
244 let plane = source.rows * source.cols;
245 for (col, grid_row) in grid.iter_mut().enumerate() {
246 for (row, color) in grid_row.iter_mut().enumerate() {
247 let base = row + source.rows * col;
248 let r = values.get(base).copied().unwrap_or(0.0) as f32;
249 let g = values.get(base + plane).copied().unwrap_or(0.0) as f32;
250 let b = values.get(base + (2 * plane)).copied().unwrap_or(0.0) as f32;
251 let a = if source.channels == 4 {
252 values.get(base + (3 * plane)).copied().unwrap_or(1.0) as f32
253 } else {
254 1.0
255 };
256 *color = Vec4::new(r, g, b, a);
257 }
258 }
259 Ok(Some(grid))
260 }
261
262 pub fn new(x_data: Vec<f64>, y_data: Vec<f64>, z_data: Vec<Vec<f64>>) -> Result<Self, String> {
264 if z_data.len() != x_data.len() {
266 return Err(format!(
267 "Z data rows ({}) must match X data length ({})",
268 z_data.len(),
269 x_data.len()
270 ));
271 }
272
273 for (i, row) in z_data.iter().enumerate() {
274 if row.len() != y_data.len() {
275 return Err(format!(
276 "Z data row {} length ({}) must match Y data length ({})",
277 i,
278 row.len(),
279 y_data.len()
280 ));
281 }
282 }
283
284 Ok(Self {
285 x_len: x_data.len(),
286 y_len: y_data.len(),
287 x_data,
288 y_data,
289 z_data: Some(z_data),
290 colormap: ColorMap::default(),
291 shading_mode: ShadingMode::default(),
292 wireframe: false,
293 alpha: 1.0,
294 flatten_z: false,
295 image_mode: false,
296 color_limits: None,
297 color_grid: None,
298 lighting_enabled: true,
299 ambient_strength: 0.2,
300 diffuse_strength: 0.8,
301 specular_strength: 0.5,
302 shininess: 32.0,
303 label: None,
304 visible: true,
305 vertices: None,
306 indices: None,
307 bounds: None,
308 dirty: true,
309 gpu_vertices: None,
310 gpu_vertex_count: None,
311 gpu_bounds: None,
312 gpu_source: None,
313 gpu_color_grid_source: None,
314 })
315 }
316
317 pub fn from_gpu_buffer(
319 x_len: usize,
320 y_len: usize,
321 buffer: GpuVertexBuffer,
322 vertex_count: usize,
323 bounds: BoundingBox,
324 ) -> Self {
325 Self {
326 x_data: Vec::new(),
327 y_data: Vec::new(),
328 z_data: None,
329 x_len,
330 y_len,
331 colormap: ColorMap::default(),
332 shading_mode: ShadingMode::default(),
333 wireframe: false,
334 alpha: 1.0,
335 flatten_z: false,
336 image_mode: false,
337 color_limits: None,
338 color_grid: None,
339 lighting_enabled: true,
340 ambient_strength: 0.2,
341 diffuse_strength: 0.8,
342 specular_strength: 0.5,
343 shininess: 32.0,
344 label: None,
345 visible: true,
346 vertices: None,
347 indices: None,
348 bounds: Some(bounds),
349 dirty: false,
350 gpu_vertices: Some(buffer),
351 gpu_vertex_count: Some(vertex_count),
352 gpu_bounds: Some(bounds),
353 gpu_source: None,
354 gpu_color_grid_source: None,
355 }
356 }
357
358 pub fn with_gpu_source(mut self, source: SurfaceGpuSource) -> Self {
359 self.gpu_source = Some(source);
360 self
361 }
362
363 pub fn with_gpu_color_grid_source(mut self, source: SurfaceGpuColorGridSource) -> Self {
364 self.gpu_color_grid_source = Some(source);
365 self
366 }
367
368 fn drop_gpu_if_possible(&mut self) {
369 if self.gpu_vertices.is_some() && self.z_data.is_some() {
370 self.invalidate_gpu_data();
371 }
372 }
373
374 pub fn from_function<F>(
376 x_range: (f64, f64),
377 y_range: (f64, f64),
378 resolution: (usize, usize),
379 func: F,
380 ) -> Result<Self, String>
381 where
382 F: Fn(f64, f64) -> f64,
383 {
384 let (x_res, y_res) = resolution;
385 if x_res < 2 || y_res < 2 {
386 return Err("Resolution must be at least 2x2".to_string());
387 }
388
389 let x_data: Vec<f64> = (0..x_res)
390 .map(|i| x_range.0 + (x_range.1 - x_range.0) * i as f64 / (x_res - 1) as f64)
391 .collect();
392
393 let y_data: Vec<f64> = (0..y_res)
394 .map(|j| y_range.0 + (y_range.1 - y_range.0) * j as f64 / (y_res - 1) as f64)
395 .collect();
396
397 let z_data: Vec<Vec<f64>> = x_data
398 .iter()
399 .map(|&x| y_data.iter().map(|&y| func(x, y)).collect())
400 .collect();
401
402 Self::new(x_data, y_data, z_data)
403 }
404
405 fn invalidate_gpu_data(&mut self) {
406 self.gpu_vertices = None;
407 self.gpu_vertex_count = None;
408 self.gpu_bounds = None;
409 self.gpu_source = None;
410 }
411
412 pub fn with_colormap(mut self, colormap: ColorMap) -> Self {
414 self.colormap = colormap;
415 self.dirty = true;
416 self.drop_gpu_if_possible();
417 self
418 }
419
420 pub fn with_shading(mut self, shading: ShadingMode) -> Self {
422 self.shading_mode = shading;
423 self.dirty = true;
424 self.drop_gpu_if_possible();
425 self
426 }
427
428 pub fn with_wireframe(mut self, enabled: bool) -> Self {
430 self.wireframe = enabled;
431 self.dirty = true;
432 self.drop_gpu_if_possible();
433 self
434 }
435
436 pub fn with_alpha(mut self, alpha: f32) -> Self {
438 self.alpha = alpha.clamp(0.0, 1.0);
439 self.dirty = true;
440 self.drop_gpu_if_possible();
441 self
442 }
443
444 pub fn with_flatten_z(mut self, enabled: bool) -> Self {
446 self.flatten_z = enabled;
447 self.dirty = true;
448 self.drop_gpu_if_possible();
449 self
450 }
451
452 pub fn with_image_mode(mut self, enabled: bool) -> Self {
453 self.image_mode = enabled;
454 self.dirty = true;
455 self.drop_gpu_if_possible();
456 self
457 }
458
459 pub fn with_color_limits(mut self, limits: Option<(f64, f64)>) -> Self {
461 self.color_limits = limits;
462 self.dirty = true;
463 self.drop_gpu_if_possible();
464 self
465 }
466
467 pub fn set_color_limits(&mut self, limits: Option<(f64, f64)>) {
469 self.color_limits = limits;
470 self.dirty = true;
471 self.drop_gpu_if_possible();
472 }
473
474 pub fn with_color_grid(mut self, grid: Vec<Vec<Vec4>>) -> Self {
476 self.color_grid = Some(grid);
477 self.gpu_color_grid_source = None;
478 self.dirty = true;
479 self.drop_gpu_if_possible();
480 self
481 }
482
483 pub fn with_label<S: Into<String>>(mut self, label: S) -> Self {
485 self.label = Some(label.into());
486 self
487 }
488
489 pub fn len(&self) -> usize {
491 self.x_len * self.y_len
492 }
493
494 pub fn is_empty(&self) -> bool {
496 self.x_len == 0 || self.y_len == 0
497 }
498
499 pub fn bounds(&mut self) -> BoundingBox {
501 if self.dirty || self.bounds.is_none() {
502 self.compute_bounds();
503 }
504 self.bounds.unwrap()
505 }
506
507 fn compute_bounds(&mut self) {
509 if let Some(bounds) = self.gpu_bounds {
510 self.bounds = Some(bounds);
511 return;
512 }
513
514 let mut min_x = f32::INFINITY;
515 let mut max_x = f32::NEG_INFINITY;
516 let mut min_y = f32::INFINITY;
517 let mut max_y = f32::NEG_INFINITY;
518 let mut min_z = f32::INFINITY;
519 let mut max_z = f32::NEG_INFINITY;
520
521 for &x in &self.x_data {
522 min_x = min_x.min(x as f32);
523 max_x = max_x.max(x as f32);
524 }
525
526 for &y in &self.y_data {
527 min_y = min_y.min(y as f32);
528 max_y = max_y.max(y as f32);
529 }
530
531 if let Some(rows) = &self.z_data {
532 for row in rows {
533 for &z in row {
534 if z.is_finite() {
535 min_z = min_z.min(z as f32);
536 max_z = max_z.max(z as f32);
537 }
538 }
539 }
540 }
541
542 self.bounds = Some(BoundingBox::new(
543 Vec3::new(min_x, min_y, min_z),
544 Vec3::new(max_x, max_y, max_z),
545 ));
546 }
547
548 pub fn statistics(&self) -> SurfaceStatistics {
550 let grid_size = self.x_len * self.y_len;
551 let triangle_count = if self.x_len > 1 && self.y_len > 1 {
552 (self.x_len - 1) * (self.y_len - 1) * 2
553 } else {
554 0
555 };
556
557 SurfaceStatistics {
558 grid_points: grid_size,
559 triangle_count,
560 x_resolution: self.x_len,
561 y_resolution: self.y_len,
562 memory_usage: self.estimated_memory_usage(),
563 }
564 }
565
566 pub fn estimated_memory_usage(&self) -> usize {
568 let data_size = std::mem::size_of::<f64>()
569 * (self.x_data.len()
570 + self.y_data.len()
571 + self
572 .z_data
573 .as_ref()
574 .map_or(0, |z| z.len() * self.y_data.len()));
575
576 let vertices_size = self
577 .vertices
578 .as_ref()
579 .map_or(0, |v| v.len() * std::mem::size_of::<Vertex>());
580
581 let indices_size = self
582 .indices
583 .as_ref()
584 .map_or(0, |i| i.len() * std::mem::size_of::<u32>());
585
586 let gpu_size = self.gpu_vertex_count.unwrap_or(0) * std::mem::size_of::<Vertex>();
587
588 data_size + vertices_size + indices_size + gpu_size
589 }
590
591 fn generate_vertices(&mut self) -> &Vec<Vertex> {
593 if self.gpu_vertices.is_some() {
594 if self.vertices.is_none() {
595 self.vertices = Some(Vec::new());
596 }
597 return self.vertices.as_ref().unwrap();
598 }
599
600 if self.dirty || self.vertices.is_none() {
601 log::trace!(
602 target: "runmat_plot",
603 "surface gen vertices {} x {}",
604 self.x_data.len(),
605 self.y_data.len()
606 );
607
608 let mut vertices = Vec::new();
609
610 let z_rows = self
612 .z_data
613 .as_ref()
614 .expect("CPU surface data missing during vertex generation");
615 let (min_z, max_z) = if let Some((lo, hi)) = self.color_limits {
616 (lo, hi)
617 } else {
618 let mut min_z = f64::INFINITY;
619 let mut max_z = f64::NEG_INFINITY;
620 for row in z_rows {
621 for &z in row {
622 if z.is_finite() {
623 min_z = min_z.min(z);
624 max_z = max_z.max(z);
625 }
626 }
627 }
628 (min_z, max_z)
629 };
630 let z_range = (max_z - min_z).max(f64::MIN_POSITIVE);
631
632 for (i, &x) in self.x_data.iter().enumerate() {
634 for (j, &y) in self.y_data.iter().enumerate() {
635 let z = z_rows[i][j];
636 let z_pos = if self.flatten_z { 0.0 } else { z as f32 };
637 let position = Vec3::new(x as f32, y as f32, z_pos);
638
639 let normal = Vec3::new(0.0, 0.0, 1.0); let color = if let Some(grid) = &self.color_grid {
644 let c = grid[i][j];
645 Vec4::new(c.x, c.y, c.z, c.w)
646 } else {
647 let t = ((z - min_z) / z_range) as f32;
648 let color_rgb = self.colormap.map_value(t.clamp(0.0, 1.0));
649 Vec4::new(color_rgb.x, color_rgb.y, color_rgb.z, self.alpha)
650 };
651
652 vertices.push(Vertex {
653 position: position.to_array(),
654 normal: normal.to_array(),
655 color: color.to_array(),
656 tex_coords: [
657 i as f32 / (self.x_data.len() - 1).max(1) as f32,
658 j as f32 / (self.y_data.len() - 1).max(1) as f32,
659 ],
660 });
661 }
662 }
663
664 log::trace!(target: "runmat_plot", "surface vertices={}", vertices.len());
665 self.vertices = Some(vertices);
666 }
667 self.vertices.as_ref().unwrap()
668 }
669
670 fn generate_indices(&mut self) -> &Vec<u32> {
672 if self.dirty || self.indices.is_none() {
673 log::trace!(target: "runmat_plot", "surface generating indices");
674
675 let mut indices = Vec::new();
676 let x_res = self.x_len;
677 let y_res = self.y_len;
678
679 for i in 0..x_res - 1 {
681 for j in 0..y_res - 1 {
682 let base = (i * y_res + j) as u32;
683 let next_row = base + y_res as u32;
684
685 indices.push(base);
688 indices.push(next_row);
689 indices.push(base + 1);
690
691 indices.push(next_row);
693 indices.push(next_row + 1);
694 indices.push(base + 1);
695 }
696 }
697
698 log::trace!(target: "runmat_plot", "surface indices={}", indices.len());
699 self.indices = Some(indices);
700 self.dirty = false;
701 }
702 self.indices.as_ref().unwrap()
703 }
704
705 fn generate_wireframe_indices(&self) -> Vec<u32> {
706 let mut indices = Vec::new();
707 if self.x_len < 2 || self.y_len < 2 {
708 return indices;
709 }
710
711 for i in 0..self.x_len {
713 for j in 0..(self.y_len - 1) {
714 let a = (i * self.y_len + j) as u32;
715 let b = (i * self.y_len + j + 1) as u32;
716 indices.push(a);
717 indices.push(b);
718 }
719 }
720
721 for i in 0..(self.x_len - 1) {
723 for j in 0..self.y_len {
724 let a = (i * self.y_len + j) as u32;
725 let b = ((i + 1) * self.y_len + j) as u32;
726 indices.push(a);
727 indices.push(b);
728 }
729 }
730
731 indices
732 }
733
734 pub fn render_data(&mut self) -> RenderData {
736 log::debug!(
737 target: "runmat_plot",
738 "surface render_data start: {} x {}",
739 self.x_len,
740 self.y_len
741 );
742
743 let using_gpu = self.gpu_vertices.is_some();
744 let bounds = self.bounds();
745 let vertices = if using_gpu {
746 Vec::new()
747 } else {
748 self.generate_vertices().clone()
749 };
750 let indices = if self.wireframe {
751 self.generate_wireframe_indices()
752 } else {
753 self.generate_indices().clone()
754 };
755
756 let material = Material {
757 albedo: Vec4::new(1.0, 1.0, 1.0, self.alpha),
758 ..Default::default()
759 };
760
761 let vertex_count = if using_gpu {
762 self.gpu_vertex_count.unwrap_or(0)
763 } else {
764 vertices.len()
765 };
766
767 log::debug!(
768 target: "runmat_plot",
769 "surface render_data generated: vertex_count={} (gpu={}), indices={}",
770 vertex_count,
771 using_gpu,
772 indices.len()
773 );
774
775 let draw_call = DrawCall {
776 vertex_offset: 0,
777 vertex_count,
778 index_offset: Some(0),
779 index_count: Some(indices.len()),
780 instance_count: 1,
781 };
782
783 log::trace!(target: "runmat_plot", "surface render_data done");
784
785 RenderData {
786 pipeline_type: if self.wireframe {
787 PipelineType::Lines
788 } else {
789 PipelineType::Triangles
790 },
791 vertices,
792 indices: Some(indices),
793
794 gpu_vertices: self.gpu_vertices.clone(),
795 bounds: Some(bounds),
796 material,
797 draw_calls: vec![draw_call],
798 image: None,
799 }
800 }
801}
802
803#[derive(Debug, Clone)]
805pub struct SurfaceStatistics {
806 pub grid_points: usize,
807 pub triangle_count: usize,
808 pub x_resolution: usize,
809 pub y_resolution: usize,
810 pub memory_usage: usize,
811}
812
813impl ColorMap {
814 pub fn map_value(&self, t: f32) -> Vec3 {
816 let t = t.clamp(0.0, 1.0);
817
818 match self {
819 ColorMap::Jet => self.jet_colormap(t),
820 ColorMap::Hot => self.hot_colormap(t),
821 ColorMap::Cool => self.cool_colormap(t),
822 ColorMap::Spring => self.spring_colormap(t),
823 ColorMap::Summer => self.summer_colormap(t),
824 ColorMap::Autumn => self.autumn_colormap(t),
825 ColorMap::Winter => self.winter_colormap(t),
826 ColorMap::Gray => Vec3::splat(t),
827 ColorMap::Bone => self.bone_colormap(t),
828 ColorMap::Copper => self.copper_colormap(t),
829 ColorMap::Pink => self.pink_colormap(t),
830 ColorMap::Lines => self.lines_colormap(t),
831 ColorMap::Viridis => self.viridis_colormap(t),
832 ColorMap::Plasma => self.plasma_colormap(t),
833 ColorMap::Inferno => self.inferno_colormap(t),
834 ColorMap::Magma => self.magma_colormap(t),
835 ColorMap::Turbo => self.turbo_colormap(t),
836 ColorMap::Parula => self.parula_colormap(t),
837 ColorMap::Custom(min_color, max_color) => {
838 min_color.truncate().lerp(max_color.truncate(), t)
839 }
840 }
841 }
842
843 fn jet_colormap(&self, t: f32) -> Vec3 {
845 let r = (1.5 - 4.0 * (t - 0.75).abs()).clamp(0.0, 1.0);
846 let g = (1.5 - 4.0 * (t - 0.5).abs()).clamp(0.0, 1.0);
847 let b = (1.5 - 4.0 * (t - 0.25).abs()).clamp(0.0, 1.0);
848 Vec3::new(r, g, b)
849 }
850
851 fn hot_colormap(&self, t: f32) -> Vec3 {
853 if t < 1.0 / 3.0 {
854 Vec3::new(3.0 * t, 0.0, 0.0)
855 } else if t < 2.0 / 3.0 {
856 Vec3::new(1.0, 3.0 * t - 1.0, 0.0)
857 } else {
858 Vec3::new(1.0, 1.0, 3.0 * t - 2.0)
859 }
860 }
861
862 fn cool_colormap(&self, t: f32) -> Vec3 {
864 Vec3::new(t, 1.0 - t, 1.0)
865 }
866
867 fn viridis_colormap(&self, t: f32) -> Vec3 {
869 let r = (0.267004 + t * (0.993248 - 0.267004)).clamp(0.0, 1.0);
871 let g = (0.004874 + t * (0.906157 - 0.004874)).clamp(0.0, 1.0);
872 let b = (0.329415 + t * (0.143936 - 0.329415) + t * t * 0.5).clamp(0.0, 1.0);
873 Vec3::new(r, g, b)
874 }
875
876 fn plasma_colormap(&self, t: f32) -> Vec3 {
878 let r = (0.050383 + t * (0.940015 - 0.050383)).clamp(0.0, 1.0);
880 let g = (0.029803 + t * (0.975158 - 0.029803) * (1.0 - t)).clamp(0.0, 1.0);
881 let b = (0.527975 + t * (0.131326 - 0.527975)).clamp(0.0, 1.0);
882 Vec3::new(r, g, b)
883 }
884
885 fn spring_colormap(&self, t: f32) -> Vec3 {
887 Vec3::new(1.0, t, 1.0 - t)
888 }
889
890 fn summer_colormap(&self, t: f32) -> Vec3 {
892 Vec3::new(t, 0.5 + 0.5 * t, 0.4)
893 }
894
895 fn autumn_colormap(&self, t: f32) -> Vec3 {
897 Vec3::new(1.0, t, 0.0)
898 }
899
900 fn winter_colormap(&self, t: f32) -> Vec3 {
902 Vec3::new(0.0, t, 1.0 - 0.5 * t)
903 }
904
905 fn bone_colormap(&self, t: f32) -> Vec3 {
907 if t < 3.0 / 8.0 {
908 Vec3::new(7.0 / 8.0 * t, 7.0 / 8.0 * t, 29.0 / 24.0 * t)
909 } else {
910 Vec3::new(
911 (29.0 + 7.0 * t) / 24.0,
912 (29.0 + 7.0 * t) / 24.0,
913 (29.0 + 7.0 * t) / 24.0,
914 )
915 }
916 }
917
918 fn copper_colormap(&self, t: f32) -> Vec3 {
920 Vec3::new((1.25 * t).min(1.0), 0.7812 * t, 0.4975 * t)
921 }
922
923 fn pink_colormap(&self, t: f32) -> Vec3 {
925 let sqrt_t = t.sqrt();
926 if t < 3.0 / 8.0 {
927 Vec3::new(14.0 / 9.0 * sqrt_t, 2.0 / 3.0 * sqrt_t, 2.0 / 3.0 * sqrt_t)
928 } else {
929 Vec3::new(
930 2.0 * sqrt_t - 1.0 / 3.0,
931 8.0 / 9.0 * sqrt_t + 1.0 / 3.0,
932 8.0 / 9.0 * sqrt_t + 1.0 / 3.0,
933 )
934 }
935 }
936
937 fn lines_colormap(&self, t: f32) -> Vec3 {
939 let _phase = (t * 7.0) % 1.0; let index = (t * 7.0) as usize % 7;
941 match index {
942 0 => Vec3::new(0.0, 0.0, 1.0), 1 => Vec3::new(0.0, 0.5, 0.0), 2 => Vec3::new(1.0, 0.0, 0.0), 3 => Vec3::new(0.0, 0.75, 0.75), 4 => Vec3::new(0.75, 0.0, 0.75), 5 => Vec3::new(0.75, 0.75, 0.0), _ => Vec3::new(0.25, 0.25, 0.25), }
950 }
951
952 fn inferno_colormap(&self, t: f32) -> Vec3 {
954 let r = (0.001462 + t * (0.988362 - 0.001462)).clamp(0.0, 1.0);
956 let g = (0.000466 + t * t * (0.982895 - 0.000466)).clamp(0.0, 1.0);
957 let b = (0.013866 + t * (1.0 - t) * (0.416065 - 0.013866)).clamp(0.0, 1.0);
958 Vec3::new(r, g, b)
959 }
960
961 fn magma_colormap(&self, t: f32) -> Vec3 {
963 let r = (0.001462 + t * (0.987053 - 0.001462)).clamp(0.0, 1.0);
965 let g = (0.000466 + t * t * (0.991438 - 0.000466)).clamp(0.0, 1.0);
966 let b = (0.013866 + t * (0.644237 - 0.013866) * (1.0 - t)).clamp(0.0, 1.0);
967 Vec3::new(r, g, b)
968 }
969
970 fn turbo_colormap(&self, t: f32) -> Vec3 {
972 let r = if t < 0.5 {
974 (0.13 + 0.87 * (2.0 * t).powf(0.25)).clamp(0.0, 1.0)
975 } else {
976 (0.8685 + 0.1315 * (2.0 * (1.0 - t)).powf(0.25)).clamp(0.0, 1.0)
977 };
978
979 let g = if t < 0.25 {
980 4.0 * t
981 } else if t < 0.75 {
982 1.0
983 } else {
984 1.0 - 4.0 * (t - 0.75)
985 }
986 .clamp(0.0, 1.0);
987
988 let b = if t < 0.5 {
989 (0.8 * (1.0 - 2.0 * t).powf(0.25)).clamp(0.0, 1.0)
990 } else {
991 (0.1 + 0.9 * (2.0 * t - 1.0).powf(0.25)).clamp(0.0, 1.0)
992 };
993
994 Vec3::new(r, g, b)
995 }
996
997 fn parula_colormap(&self, t: f32) -> Vec3 {
999 let r = if t < 0.25 {
1001 0.2081 * (1.0 - t)
1002 } else if t < 0.5 {
1003 t - 0.25
1004 } else if t < 0.75 {
1005 1.0
1006 } else {
1007 1.0 - 0.5 * (t - 0.75)
1008 }
1009 .clamp(0.0, 1.0);
1010
1011 let g = if t < 0.125 {
1012 0.1663 * t / 0.125
1013 } else if t < 0.375 {
1014 0.1663 + (0.7079 - 0.1663) * (t - 0.125) / 0.25
1015 } else if t < 0.625 {
1016 0.7079 + (0.9839 - 0.7079) * (t - 0.375) / 0.25
1017 } else {
1018 0.9839 * (1.0 - (t - 0.625) / 0.375)
1019 }
1020 .clamp(0.0, 1.0);
1021
1022 let b = if t < 0.25 {
1023 0.5 + 0.5 * t / 0.25
1024 } else if t < 0.5 {
1025 1.0
1026 } else {
1027 1.0 - 2.0 * (t - 0.5)
1028 }
1029 .clamp(0.0, 1.0);
1030
1031 Vec3::new(r, g, b)
1032 }
1033
1034 #[allow(dead_code)] fn default_colormap(&self, t: f32) -> Vec3 {
1037 if t < 0.5 {
1039 Vec3::new(0.0, 2.0 * t, 1.0 - 2.0 * t)
1040 } else {
1041 Vec3::new(2.0 * (t - 0.5), 1.0 - 2.0 * (t - 0.5), 0.0)
1042 }
1043 }
1044}
1045
1046pub mod matlab_compat {
1048 use super::*;
1049
1050 pub fn surf(x: Vec<f64>, y: Vec<f64>, z: Vec<Vec<f64>>) -> Result<SurfacePlot, String> {
1052 SurfacePlot::new(x, y, z)
1053 }
1054
1055 pub fn mesh(x: Vec<f64>, y: Vec<f64>, z: Vec<Vec<f64>>) -> Result<SurfacePlot, String> {
1057 Ok(SurfacePlot::new(x, y, z)?
1058 .with_wireframe(true)
1059 .with_shading(ShadingMode::None))
1060 }
1061
1062 pub fn meshgrid_surf(
1064 x_range: (f64, f64),
1065 y_range: (f64, f64),
1066 resolution: (usize, usize),
1067 func: impl Fn(f64, f64) -> f64,
1068 ) -> Result<SurfacePlot, String> {
1069 SurfacePlot::from_function(x_range, y_range, resolution, func)
1070 }
1071
1072 pub fn surf_with_colormap(
1074 x: Vec<f64>,
1075 y: Vec<f64>,
1076 z: Vec<Vec<f64>>,
1077 colormap: &str,
1078 ) -> Result<SurfacePlot, String> {
1079 let cmap =
1080 ColorMap::from_name(colormap).ok_or_else(|| format!("Unknown colormap: {colormap}"))?;
1081
1082 Ok(SurfacePlot::new(x, y, z)?.with_colormap(cmap))
1083 }
1084}
1085
1086#[cfg(test)]
1087mod tests {
1088 use super::*;
1089
1090 #[test]
1091 fn test_surface_plot_creation() {
1092 let x = vec![0.0, 1.0, 2.0];
1093 let y = vec![0.0, 1.0];
1094 let z = vec![vec![0.0, 1.0], vec![1.0, 2.0], vec![2.0, 3.0]];
1095
1096 let surface = SurfacePlot::new(x, y, z).unwrap();
1097
1098 assert_eq!(surface.x_data.len(), 3);
1099 assert_eq!(surface.y_data.len(), 2);
1100 let rows = surface.z_data.as_ref().unwrap();
1101 assert_eq!(rows.len(), 3);
1102 assert_eq!(rows[0].len(), 2);
1103 assert!(surface.visible);
1104 }
1105
1106 #[test]
1107 fn test_surface_from_function() {
1108 let surface =
1109 SurfacePlot::from_function((-2.0, 2.0), (-2.0, 2.0), (10, 10), |x, y| x * x + y * y)
1110 .unwrap();
1111
1112 assert_eq!(surface.x_data.len(), 10);
1113 assert_eq!(surface.y_data.len(), 10);
1114 let rows = surface.z_data.as_ref().unwrap();
1115 assert_eq!(rows.len(), 10);
1116
1117 assert_eq!(rows[0][0], 8.0); }
1120
1121 #[test]
1122 fn test_surface_validation() {
1123 let x = vec![0.0, 1.0];
1124 let y = vec![0.0, 1.0, 2.0];
1125 let z = vec![
1126 vec![0.0, 1.0], vec![1.0, 2.0],
1128 ];
1129
1130 assert!(SurfacePlot::new(x, y, z).is_err());
1131 }
1132
1133 #[test]
1134 fn test_surface_styling() {
1135 let x = vec![0.0, 1.0];
1136 let y = vec![0.0, 1.0];
1137 let z = vec![vec![0.0, 1.0], vec![1.0, 2.0]];
1138
1139 let surface = SurfacePlot::new(x, y, z)
1140 .unwrap()
1141 .with_colormap(ColorMap::Hot)
1142 .with_wireframe(true)
1143 .with_alpha(0.8)
1144 .with_label("Test Surface");
1145
1146 assert_eq!(surface.colormap, ColorMap::Hot);
1147 assert!(surface.wireframe);
1148 assert_eq!(surface.alpha, 0.8);
1149 assert_eq!(surface.label, Some("Test Surface".to_string()));
1150 }
1151
1152 #[test]
1153 fn test_colormap_mapping() {
1154 let jet = ColorMap::Jet;
1155
1156 let color_0 = jet.map_value(0.0);
1158 let color_1 = jet.map_value(1.0);
1159
1160 assert!(color_0.x >= 0.0 && color_0.x <= 1.0);
1161 assert!(color_1.x >= 0.0 && color_1.x <= 1.0);
1162
1163 let color_mid = jet.map_value(0.5);
1165 assert_ne!(color_0, color_mid);
1166 assert_ne!(color_mid, color_1);
1167 }
1168
1169 #[test]
1170 fn test_surface_statistics() {
1171 let x = vec![0.0, 1.0, 2.0, 3.0];
1172 let y = vec![0.0, 1.0, 2.0];
1173 let z = vec![
1174 vec![0.0, 1.0, 2.0],
1175 vec![1.0, 2.0, 3.0],
1176 vec![2.0, 3.0, 4.0],
1177 vec![3.0, 4.0, 5.0],
1178 ];
1179
1180 let surface = SurfacePlot::new(x, y, z).unwrap();
1181 let stats = surface.statistics();
1182
1183 assert_eq!(stats.grid_points, 12); assert_eq!(stats.triangle_count, 12); assert_eq!(stats.x_resolution, 4);
1186 assert_eq!(stats.y_resolution, 3);
1187 assert!(stats.memory_usage > 0);
1188 }
1189
1190 #[test]
1191 fn test_matlab_compat() {
1192 use super::matlab_compat::*;
1193
1194 let x = vec![0.0, 1.0];
1195 let y = vec![0.0, 1.0];
1196 let z = vec![vec![0.0, 1.0], vec![1.0, 2.0]];
1197
1198 let surface = surf(x.clone(), y.clone(), z.clone()).unwrap();
1199 assert!(!surface.wireframe);
1200
1201 let mesh_plot = mesh(x.clone(), y.clone(), z.clone()).unwrap();
1202 assert!(mesh_plot.wireframe);
1203
1204 let colormap_surface = surf_with_colormap(x, y, z, "viridis").unwrap();
1205 assert_eq!(colormap_surface.colormap, ColorMap::Viridis);
1206 }
1207
1208 #[test]
1209 fn colormap_from_name_accepts_canonical_names_and_aliases() {
1210 let cases = [
1211 ("parula", ColorMap::Parula),
1212 ("viridis", ColorMap::Viridis),
1213 ("plasma", ColorMap::Plasma),
1214 ("inferno", ColorMap::Inferno),
1215 ("magma", ColorMap::Magma),
1216 ("turbo", ColorMap::Turbo),
1217 ("jet", ColorMap::Jet),
1218 ("hot", ColorMap::Hot),
1219 ("cool", ColorMap::Cool),
1220 ("spring", ColorMap::Spring),
1221 ("summer", ColorMap::Summer),
1222 ("autumn", ColorMap::Autumn),
1223 ("winter", ColorMap::Winter),
1224 ("gray", ColorMap::Gray),
1225 ("grey", ColorMap::Gray),
1226 ("bone", ColorMap::Bone),
1227 ("copper", ColorMap::Copper),
1228 ("pink", ColorMap::Pink),
1229 ("lines", ColorMap::Lines),
1230 ];
1231
1232 for (name, expected) in cases {
1233 assert_eq!(ColorMap::from_name(name), Some(expected), "{name}");
1234 }
1235 for name in ColorMap::CANONICAL_NAMES
1236 .iter()
1237 .chain(ColorMap::ALIASES.iter())
1238 .copied()
1239 {
1240 assert!(
1241 ColorMap::from_name(name).is_some(),
1242 "colormap table entry should parse: {name}"
1243 );
1244 }
1245 }
1246
1247 #[test]
1248 fn colormap_from_name_normalizes_and_rejects_unknown_names() {
1249 assert_eq!(ColorMap::from_name(" Turbo "), Some(ColorMap::Turbo));
1250 assert_eq!(ColorMap::from_name("GREY"), Some(ColorMap::Gray));
1251 assert_eq!(ColorMap::from_name("hsv"), None);
1252 assert!(!ColorMap::CANONICAL_NAMES.contains(&"hsv"));
1253 assert!(!ColorMap::ALIASES.contains(&"hsv"));
1254 assert_eq!(ColorMap::from_name("not-a-colormap"), None);
1255 }
1256}