runmat_plot/plots/
scatter.rs

1//! Scatter plot implementation
2//!
3//! High-performance scatter plotting with GPU acceleration.
4
5use crate::context::shared_wgpu_context;
6use crate::core::{
7    vertex_utils, BoundingBox, DrawCall, GpuVertexBuffer, Material, PipelineType, RenderData,
8    Vertex,
9};
10use crate::gpu::scatter2::Scatter2GpuInputs;
11use crate::gpu::util::readback_scalar_buffer_f64;
12use crate::plots::surface::ColorMap;
13use glam::{Vec3, Vec4};
14
15/// High-performance GPU-accelerated scatter plot
16#[derive(Debug, Clone)]
17pub struct ScatterPlot {
18    /// Raw data points (x, y coordinates)
19    pub x_data: Vec<f64>,
20    pub y_data: Vec<f64>,
21    pub theta_data: Option<Vec<f64>>,
22    pub r_data: Option<Vec<f64>>,
23
24    /// Visual styling
25    pub color: Vec4,
26    pub edge_color: Vec4,
27    pub edge_thickness: f32,
28    pub marker_size: f32,
29    pub marker_style: MarkerStyle,
30    pub per_point_sizes: Option<Vec<f32>>, // pixel diameters per point
31    pub per_point_colors: Option<Vec<Vec4>>, // per-point RGBA
32    pub color_values: Option<Vec<f64>>,    // scalar values mapped by colormap
33    pub color_limits: Option<(f64, f64)>,
34    pub colormap: ColorMap,
35    pub filled: bool,
36    pub edge_color_from_vertex_colors: bool,
37
38    /// Metadata
39    pub label: Option<String>,
40    pub visible: bool,
41
42    /// Generated rendering data (cached)
43    vertices: Option<Vec<Vertex>>,
44    bounds: Option<BoundingBox>,
45    dirty: bool,
46    gpu_vertices: Option<GpuVertexBuffer>,
47    gpu_point_count: Option<usize>,
48    gpu_inputs: Option<Scatter2GpuInputs>,
49    gpu_has_per_point_sizes: bool,
50    gpu_has_per_point_colors: bool,
51}
52
53/// Marker styles for scatter plots
54#[derive(Debug, Clone, Copy, PartialEq, Eq)]
55pub enum MarkerStyle {
56    Circle,
57    Square,
58    Triangle,
59    Diamond,
60    Plus,
61    Cross,
62    Star,
63    Hexagon,
64}
65
66impl Default for MarkerStyle {
67    fn default() -> Self {
68        Self::Circle
69    }
70}
71
72#[derive(Clone, Copy, Debug)]
73pub struct ScatterGpuStyle {
74    pub color: Vec4,
75    pub edge_color: Vec4,
76    pub edge_thickness: f32,
77    pub marker_size: f32,
78    pub marker_style: MarkerStyle,
79    pub filled: bool,
80    pub has_per_point_sizes: bool,
81    pub has_per_point_colors: bool,
82    pub edge_from_vertex_colors: bool,
83}
84
85impl ScatterPlot {
86    pub async fn export_scene_xy_data(&self) -> Result<(Vec<f64>, Vec<f64>), String> {
87        if !self.x_data.is_empty() && self.x_data.len() == self.y_data.len() {
88            return Ok((self.x_data.clone(), self.y_data.clone()));
89        }
90        if !self.x_data.is_empty() || !self.y_data.is_empty() {
91            return Err(format!(
92                "scatter plot has partial CPU source data: x has {} values, y has {} values",
93                self.x_data.len(),
94                self.y_data.len()
95            ));
96        }
97
98        if let Some(inputs) = &self.gpu_inputs {
99            let context = shared_wgpu_context().ok_or_else(|| {
100                "scatter plot has GPU source data but no shared WGPU context is installed"
101                    .to_string()
102            })?;
103            let len = inputs.len as usize;
104            let x = readback_scalar_buffer_f64(
105                &context.device,
106                &context.queue,
107                &inputs.x_buffer,
108                len,
109                inputs.scalar,
110            )
111            .await?;
112            let y = readback_scalar_buffer_f64(
113                &context.device,
114                &context.queue,
115                &inputs.y_buffer,
116                len,
117                inputs.scalar,
118            )
119            .await?;
120            return Ok((x, y));
121        }
122
123        if self.gpu_vertices.is_some() {
124            return Err(
125                "scatter plot has GPU render vertices but no exportable source data".to_string(),
126            );
127        }
128
129        Ok((Vec::new(), Vec::new()))
130    }
131
132    /// Create a new scatter plot with data
133    pub fn new(x_data: Vec<f64>, y_data: Vec<f64>) -> Result<Self, String> {
134        if x_data.len() != y_data.len() {
135            return Err(format!(
136                "Data length mismatch: x_data has {} points, y_data has {} points",
137                x_data.len(),
138                y_data.len()
139            ));
140        }
141
142        if x_data.is_empty() {
143            return Err("Cannot create scatter plot with empty data".to_string());
144        }
145
146        Ok(Self {
147            x_data,
148            y_data,
149            theta_data: None,
150            r_data: None,
151            color: Vec4::new(1.0, 0.2, 0.2, 1.0), // Brighter red
152            edge_color: Vec4::new(0.0, 0.0, 0.0, 1.0),
153            edge_thickness: 1.0,
154            marker_size: 12.0,
155            marker_style: MarkerStyle::default(),
156            per_point_sizes: None,
157            per_point_colors: None,
158            color_values: None,
159            color_limits: None,
160            colormap: ColorMap::Parula,
161            filled: false,
162            edge_color_from_vertex_colors: false,
163            label: None,
164            visible: true,
165            vertices: None,
166            bounds: None,
167            dirty: true,
168            gpu_vertices: None,
169            gpu_point_count: None,
170            gpu_inputs: None,
171            gpu_has_per_point_sizes: false,
172            gpu_has_per_point_colors: false,
173        })
174    }
175
176    /// Build a scatter plot directly from a GPU vertex buffer to avoid CPU copies.
177    pub fn from_gpu_buffer(
178        buffer: GpuVertexBuffer,
179        point_count: usize,
180        bounds: BoundingBox,
181        style: ScatterGpuStyle,
182    ) -> Self {
183        Self {
184            x_data: Vec::new(),
185            y_data: Vec::new(),
186            theta_data: None,
187            r_data: None,
188            color: style.color,
189            edge_color: style.edge_color,
190            edge_thickness: style.edge_thickness,
191            marker_size: style.marker_size,
192            marker_style: style.marker_style,
193            per_point_sizes: None,
194            per_point_colors: None,
195            color_values: None,
196            color_limits: None,
197            colormap: ColorMap::Parula,
198            filled: style.filled,
199            edge_color_from_vertex_colors: style.edge_from_vertex_colors,
200            label: None,
201            visible: true,
202            vertices: None,
203            bounds: Some(bounds),
204            dirty: false,
205            gpu_vertices: Some(buffer),
206            gpu_point_count: Some(point_count),
207            gpu_inputs: None,
208            gpu_has_per_point_sizes: style.has_per_point_sizes,
209            gpu_has_per_point_colors: style.has_per_point_colors,
210        }
211    }
212
213    pub fn with_gpu_source_inputs(mut self, inputs: Scatter2GpuInputs) -> Self {
214        self.gpu_inputs = Some(inputs);
215        self
216    }
217
218    pub fn with_polar_data(mut self, theta: Vec<f64>, r: Vec<f64>) -> Self {
219        self.theta_data = Some(theta);
220        self.r_data = Some(r);
221        self
222    }
223
224    fn invalidate_gpu_vertices(&mut self) {
225        self.gpu_vertices = None;
226        self.gpu_point_count = None;
227    }
228
229    fn clear_gpu_source_inputs(&mut self) {
230        self.gpu_inputs = None;
231        self.gpu_has_per_point_sizes = false;
232        self.gpu_has_per_point_colors = false;
233    }
234
235    /// Create a scatter plot with custom styling
236    pub fn with_style(mut self, color: Vec4, marker_size: f32, marker_style: MarkerStyle) -> Self {
237        self.color = color;
238        self.marker_size = marker_size;
239        self.marker_style = marker_style;
240        self.dirty = true;
241        self.invalidate_gpu_vertices();
242        self.gpu_has_per_point_sizes = false;
243        self.gpu_has_per_point_colors = false;
244        self
245    }
246
247    /// Set the plot label for legends
248    pub fn with_label<S: Into<String>>(mut self, label: S) -> Self {
249        self.label = Some(label.into());
250        self
251    }
252
253    /// Set marker face color
254    pub fn set_face_color(&mut self, color: Vec4) {
255        self.color = color;
256        self.dirty = true;
257        self.invalidate_gpu_vertices();
258    }
259    /// Set marker edge color
260    pub fn set_edge_color(&mut self, color: Vec4) {
261        self.edge_color = color;
262        self.dirty = true;
263        self.invalidate_gpu_vertices();
264    }
265    pub fn set_edge_color_from_vertex(&mut self, enabled: bool) {
266        self.edge_color_from_vertex_colors = enabled;
267    }
268    /// Set marker edge thickness (pixels)
269    pub fn set_edge_thickness(&mut self, px: f32) {
270        self.edge_thickness = px.max(0.0);
271        self.dirty = true;
272        self.invalidate_gpu_vertices();
273    }
274    pub fn set_sizes(&mut self, sizes: Vec<f32>) {
275        self.per_point_sizes = Some(sizes);
276        self.dirty = true;
277        self.invalidate_gpu_vertices();
278        self.gpu_has_per_point_sizes = false;
279    }
280    pub fn set_colors(&mut self, colors: Vec<Vec4>) {
281        self.per_point_colors = Some(colors);
282        self.dirty = true;
283        self.invalidate_gpu_vertices();
284        self.gpu_has_per_point_colors = false;
285    }
286    pub fn set_color_values(&mut self, values: Vec<f64>, limits: Option<(f64, f64)>) {
287        self.color_values = Some(values);
288        self.color_limits = limits;
289        self.dirty = true;
290        self.invalidate_gpu_vertices();
291        self.gpu_has_per_point_colors = false;
292    }
293    pub fn with_colormap(mut self, cmap: ColorMap) -> Self {
294        self.colormap = cmap;
295        self.dirty = true;
296        self.invalidate_gpu_vertices();
297        self
298    }
299    pub fn set_filled(&mut self, filled: bool) {
300        self.filled = filled;
301        self.dirty = true;
302        self.invalidate_gpu_vertices();
303    }
304
305    /// Update the data points
306    pub fn update_data(&mut self, x_data: Vec<f64>, y_data: Vec<f64>) -> Result<(), String> {
307        if x_data.len() != y_data.len() {
308            return Err(format!(
309                "Data length mismatch: x_data has {} points, y_data has {} points",
310                x_data.len(),
311                y_data.len()
312            ));
313        }
314
315        if x_data.is_empty() {
316            return Err("Cannot update with empty data".to_string());
317        }
318
319        self.x_data = x_data;
320        self.y_data = y_data;
321        self.dirty = true;
322        self.invalidate_gpu_vertices();
323        self.clear_gpu_source_inputs();
324        Ok(())
325    }
326
327    /// Set the color of the markers
328    pub fn set_color(&mut self, color: Vec4) {
329        self.color = color;
330        self.dirty = true;
331        self.invalidate_gpu_vertices();
332    }
333
334    /// Set the marker size
335    pub fn set_marker_size(&mut self, size: f32) {
336        self.marker_size = size.max(0.1); // Minimum marker size
337        self.dirty = true;
338        self.invalidate_gpu_vertices();
339    }
340
341    /// Set the marker style
342    pub fn set_marker_style(&mut self, style: MarkerStyle) {
343        self.marker_style = style;
344        self.dirty = true;
345        self.invalidate_gpu_vertices();
346    }
347
348    /// Show or hide the plot
349    pub fn set_visible(&mut self, visible: bool) {
350        self.visible = visible;
351    }
352
353    /// Get the number of data points
354    pub fn len(&self) -> usize {
355        if !self.x_data.is_empty() {
356            self.x_data.len()
357        } else {
358            self.gpu_point_count.unwrap_or(0)
359        }
360    }
361
362    /// Check if the plot has no data
363    pub fn is_empty(&self) -> bool {
364        self.len() == 0
365    }
366
367    /// Generate vertices for GPU rendering
368    pub fn generate_vertices(&mut self) -> &Vec<Vertex> {
369        if self.gpu_vertices.is_some() {
370            if self.vertices.is_none() {
371                self.vertices = Some(Vec::new());
372            }
373            return self.vertices.as_ref().unwrap();
374        }
375        if self.dirty || self.vertices.is_none() {
376            let base_color = self.color;
377            if self.per_point_colors.is_some() || self.color_values.is_some() { /* vertex color takes precedence; shader blends by face alpha */
378            }
379            let mut verts =
380                vertex_utils::create_scatter_plot(&self.x_data, &self.y_data, base_color);
381            // per-point colors
382            if let Some(ref colors) = self.per_point_colors {
383                let m = colors.len().min(verts.len());
384                for i in 0..m {
385                    verts[i].color = colors[i].to_array();
386                }
387            } else if let Some(ref vals) = self.color_values {
388                let n = verts.len();
389                let (mut cmin, mut cmax) = if let Some(lims) = self.color_limits {
390                    lims
391                } else {
392                    let mut lo = f64::INFINITY;
393                    let mut hi = f64::NEG_INFINITY;
394                    for &v in vals {
395                        if v.is_finite() {
396                            if v < lo {
397                                lo = v;
398                            }
399                            if v > hi {
400                                hi = v;
401                            }
402                        }
403                    }
404                    if !lo.is_finite() || !hi.is_finite() || hi <= lo {
405                        (0.0, 1.0)
406                    } else {
407                        (lo, hi)
408                    }
409                };
410                if !(cmin.is_finite() && cmax.is_finite()) || cmax <= cmin {
411                    cmin = 0.0;
412                    cmax = 1.0;
413                }
414                let denom = (cmax - cmin).max(f64::EPSILON);
415                for (i, vert) in verts.iter_mut().enumerate().take(n) {
416                    let t = ((vals[i] - cmin) / denom) as f32;
417                    let rgb = self.colormap.map_value(t);
418                    vert.color = [rgb.x, rgb.y, rgb.z, 1.0];
419                }
420            }
421            // Store marker size in normal.z for direct point expansion
422            if let Some(ref sizes) = self.per_point_sizes {
423                for (i, vert) in verts.iter_mut().enumerate() {
424                    let s = sizes.get(i).copied().unwrap_or(self.marker_size);
425                    vert.normal[2] = s.max(1.0);
426                }
427            } else {
428                for v in &mut verts {
429                    v.normal[2] = self.marker_size.max(1.0);
430                }
431            }
432            self.vertices = Some(verts);
433            self.dirty = false;
434        }
435        self.vertices.as_ref().unwrap()
436    }
437
438    /// Get the bounding box of the data
439    pub fn bounds(&mut self) -> BoundingBox {
440        if self.gpu_vertices.is_some() {
441            return self.bounds.unwrap_or_default();
442        }
443        if self.dirty || self.bounds.is_none() {
444            let points: Vec<Vec3> = self
445                .x_data
446                .iter()
447                .zip(self.y_data.iter())
448                .map(|(&x, &y)| Vec3::new(x as f32, y as f32, 0.0))
449                .collect();
450            self.bounds = Some(BoundingBox::from_points(&points));
451        }
452        self.bounds.unwrap()
453    }
454
455    /// Generate complete render data for the graphics pipeline
456    pub fn render_data(&mut self) -> RenderData {
457        let using_gpu = self.gpu_vertices.is_some();
458        let gpu_vertices = self.gpu_vertices.clone();
459        let bounds = self.bounds();
460        let (vertices, vertex_count) = if using_gpu {
461            let count = self
462                .gpu_point_count
463                .or_else(|| gpu_vertices.as_ref().map(|buf| buf.vertex_count))
464                .unwrap_or(0);
465            (Vec::new(), count)
466        } else {
467            let verts = self.generate_vertices().clone();
468            let count = verts.len();
469            (verts, count)
470        };
471
472        let mut material = Material {
473            albedo: self.color,
474            ..Default::default()
475        };
476        // If vertex colors vary across points, prefer per-vertex colors (alpha=0)
477        let is_multi_color = if using_gpu {
478            self.gpu_has_per_point_colors
479                || self.per_point_colors.is_some()
480                || self.color_values.is_some()
481        } else if vertices.is_empty() {
482            false
483        } else {
484            let first = vertices[0].color;
485            vertices.iter().any(|v| v.color != first)
486        };
487        if is_multi_color {
488            material.albedo.w = 0.0;
489        } else if self.filled {
490            material.albedo.w = 1.0;
491        }
492        material.emissive = self.edge_color; // stash edge color
493        material.roughness = self.edge_thickness; // stash thickness in roughness
494        material.metallic = match self.marker_style {
495            MarkerStyle::Circle => 0.0,
496            MarkerStyle::Square => 1.0,
497            MarkerStyle::Triangle => 2.0,
498            MarkerStyle::Diamond => 3.0,
499            MarkerStyle::Plus => 4.0,
500            MarkerStyle::Cross => 5.0,
501            MarkerStyle::Star => 6.0,
502            MarkerStyle::Hexagon => 7.0,
503        };
504        let has_vertex_colors = if using_gpu {
505            self.gpu_has_per_point_colors
506        } else {
507            self.per_point_colors.is_some() || self.color_values.is_some()
508        };
509        let use_vertex_edge_color = self.edge_color_from_vertex_colors && has_vertex_colors;
510        material.emissive.w = if use_vertex_edge_color { 0.0 } else { 1.0 };
511
512        let draw_call = DrawCall {
513            vertex_offset: 0,
514            vertex_count,
515            index_offset: None,
516            index_count: None,
517            instance_count: 1,
518        };
519
520        RenderData {
521            pipeline_type: PipelineType::Points,
522            vertices,
523            indices: None,
524            gpu_vertices,
525            bounds: Some(bounds),
526            material,
527            draw_calls: vec![draw_call],
528            image: None,
529        }
530    }
531
532    /// Get plot statistics for debugging
533    pub fn statistics(&self) -> PlotStatistics {
534        let (min_x, max_x, min_y, max_y) = if !self.x_data.is_empty() {
535            let (min_x, max_x) = self
536                .x_data
537                .iter()
538                .fold((f64::INFINITY, f64::NEG_INFINITY), |(min, max), &x| {
539                    (min.min(x), max.max(x))
540                });
541            let (min_y, max_y) = self
542                .y_data
543                .iter()
544                .fold((f64::INFINITY, f64::NEG_INFINITY), |(min, max), &y| {
545                    (min.min(y), max.max(y))
546                });
547            (min_x, max_x, min_y, max_y)
548        } else if let Some(bounds) = &self.bounds {
549            (
550                bounds.min.x as f64,
551                bounds.max.x as f64,
552                bounds.min.y as f64,
553                bounds.max.y as f64,
554            )
555        } else {
556            (0.0, 0.0, 0.0, 0.0)
557        };
558
559        PlotStatistics {
560            point_count: self.len(),
561            x_range: (min_x, max_x),
562            y_range: (min_y, max_y),
563            memory_usage: self.estimated_memory_usage(),
564        }
565    }
566
567    /// Estimate memory usage in bytes
568    pub fn estimated_memory_usage(&self) -> usize {
569        std::mem::size_of::<f64>() * (self.x_data.len() + self.y_data.len())
570            + self
571                .vertices
572                .as_ref()
573                .map_or(0, |v| v.len() * std::mem::size_of::<Vertex>())
574            + self.gpu_point_count.unwrap_or(0) * std::mem::size_of::<Vertex>()
575    }
576}
577
578/// Plot performance and data statistics
579#[derive(Debug, Clone)]
580pub struct PlotStatistics {
581    pub point_count: usize,
582    pub x_range: (f64, f64),
583    pub y_range: (f64, f64),
584    pub memory_usage: usize,
585}
586
587/// MATLAB-compatible scatter plot creation utilities
588pub mod matlab_compat {
589    use super::*;
590
591    /// Create a simple scatter plot (equivalent to MATLAB's `scatter(x, y)`)
592    pub fn scatter(x: Vec<f64>, y: Vec<f64>) -> Result<ScatterPlot, String> {
593        ScatterPlot::new(x, y)
594    }
595
596    /// Create a scatter plot with specified color and size (`scatter(x, y, size, color)`)
597    pub fn scatter_with_style(
598        x: Vec<f64>,
599        y: Vec<f64>,
600        size: f32,
601        color: &str,
602    ) -> Result<ScatterPlot, String> {
603        let color_vec = parse_matlab_color(color)?;
604        Ok(ScatterPlot::new(x, y)?.with_style(color_vec, size, MarkerStyle::Circle))
605    }
606
607    /// Parse MATLAB color specifications
608    fn parse_matlab_color(color: &str) -> Result<Vec4, String> {
609        match color {
610            "r" | "red" => Ok(Vec4::new(1.0, 0.0, 0.0, 1.0)),
611            "g" | "green" => Ok(Vec4::new(0.0, 1.0, 0.0, 1.0)),
612            "b" | "blue" => Ok(Vec4::new(0.0, 0.0, 1.0, 1.0)),
613            "c" | "cyan" => Ok(Vec4::new(0.0, 1.0, 1.0, 1.0)),
614            "m" | "magenta" => Ok(Vec4::new(1.0, 0.0, 1.0, 1.0)),
615            "y" | "yellow" => Ok(Vec4::new(1.0, 1.0, 0.0, 1.0)),
616            "k" | "black" => Ok(Vec4::new(0.0, 0.0, 0.0, 1.0)),
617            "w" | "white" => Ok(Vec4::new(1.0, 1.0, 1.0, 1.0)),
618            _ => Err(format!("Unknown color: {color}")),
619        }
620    }
621}
622
623#[cfg(test)]
624mod tests {
625    use super::*;
626
627    #[test]
628    fn test_scatter_plot_creation() {
629        let x = vec![0.0, 1.0, 2.0, 3.0];
630        let y = vec![0.0, 1.0, 4.0, 9.0];
631
632        let plot = ScatterPlot::new(x.clone(), y.clone()).unwrap();
633
634        assert_eq!(plot.x_data, x);
635        assert_eq!(plot.y_data, y);
636        assert_eq!(plot.len(), 4);
637        assert!(!plot.is_empty());
638        assert!(plot.visible);
639    }
640
641    #[test]
642    fn test_scatter_plot_styling() {
643        let x = vec![0.0, 1.0, 2.0];
644        let y = vec![1.0, 2.0, 1.5];
645        let color = Vec4::new(0.0, 1.0, 0.0, 1.0);
646
647        let plot = ScatterPlot::new(x, y)
648            .unwrap()
649            .with_style(color, 5.0, MarkerStyle::Square)
650            .with_label("Test Scatter");
651
652        assert_eq!(plot.color, color);
653        assert_eq!(plot.marker_size, 5.0);
654        assert_eq!(plot.marker_style, MarkerStyle::Square);
655        assert_eq!(plot.label, Some("Test Scatter".to_string()));
656    }
657
658    #[test]
659    fn test_scatter_plot_render_data() {
660        let x = vec![0.0, 1.0, 2.0];
661        let y = vec![1.0, 2.0, 1.0];
662
663        let mut plot = ScatterPlot::new(x, y).unwrap();
664        let render_data = plot.render_data();
665
666        assert_eq!(render_data.pipeline_type, PipelineType::Points);
667        assert_eq!(render_data.vertices.len(), 3); // One vertex per point
668        assert!(render_data.indices.is_none());
669        assert_eq!(render_data.draw_calls.len(), 1);
670    }
671
672    #[test]
673    fn test_matlab_compat_scatter() {
674        use super::matlab_compat::*;
675
676        let x = vec![0.0, 1.0];
677        let y = vec![0.0, 1.0];
678
679        let basic_scatter = scatter(x.clone(), y.clone()).unwrap();
680        assert_eq!(basic_scatter.len(), 2);
681
682        let styled_scatter = scatter_with_style(x.clone(), y.clone(), 5.0, "g").unwrap();
683        assert_eq!(styled_scatter.color, Vec4::new(0.0, 1.0, 0.0, 1.0));
684        assert_eq!(styled_scatter.marker_size, 5.0);
685    }
686}
runmat_plot/plots/scatter.rs

runmat_plot/plots/
scatter.rs
