runmat_plot/plots/

scatter3.rs

//! 3D scatter plot implementation for MATLAB's `scatter3`.

use crate::core::{
    vertex_utils, BoundingBox, DrawCall, GpuVertexBuffer, Material, PipelineType, RenderData,
    Vertex,
};
use crate::plots::scatter::MarkerStyle;
use glam::{Vec3, Vec4};

#[derive(Clone, Copy, Debug)]
pub struct Scatter3GpuStyle {
    pub color: Vec4,
    pub edge_color: Vec4,
    pub edge_thickness: f32,
    pub marker_style: MarkerStyle,
    pub filled: bool,
    pub has_per_point_colors: bool,
    pub edge_from_vertex_colors: bool,
}

/// GPU-accelerated scatter3 plot for MATLAB semantics.
#[derive(Debug, Clone)]
pub struct Scatter3Plot {
    /// Point positions in 3D space.
    pub points: Vec<Vec3>,
    /// Per-point RGBA colors.
    pub colors: Vec<Vec4>,
    /// Marker size in pixels.
    pub point_size: f32,
    /// Optional per-point marker sizes.
    pub point_sizes: Option<Vec<f32>>,
    /// Marker edge color.
    pub edge_color: Vec4,
    /// Marker edge thickness in pixels.
    pub edge_thickness: f32,
    /// Marker shape.
    pub marker_style: MarkerStyle,
    /// Whether marker faces are filled.
    pub filled: bool,
    /// Whether edge color should come from per-vertex colors.
    pub edge_color_from_vertex_colors: bool,
    /// Legend label.
    pub label: Option<String>,
    /// Visibility flag.
    pub visible: bool,
    vertices: Option<Vec<Vertex>>,
    bounds: Option<BoundingBox>,
    gpu_vertices: Option<GpuVertexBuffer>,
    gpu_point_count: Option<usize>,
    gpu_has_per_point_colors: bool,
}

impl Scatter3Plot {
    /// Create a new scatter3 plot. Colors default to a blue colormap.
    pub fn new(points: Vec<Vec3>) -> Result<Self, String> {
        let default_color = Vec4::new(0.1, 0.7, 0.3, 1.0);
        let colors = vec![default_color; points.len()];
        Ok(Self {
            points,
            colors,
            point_size: 8.0,
            point_sizes: None,
            edge_color: default_color,
            edge_thickness: 1.0,
            marker_style: MarkerStyle::Circle,
            filled: true,
            edge_color_from_vertex_colors: false,
            label: None,
            visible: true,
            vertices: None,
            bounds: None,
            gpu_vertices: None,
            gpu_point_count: None,
            gpu_has_per_point_colors: false,
        })
    }

    /// Build a scatter plot directly from a GPU vertex buffer, bypassing CPU copies.
    pub fn from_gpu_buffer(
        buffer: GpuVertexBuffer,
        point_count: usize,
        style: Scatter3GpuStyle,
        point_size: f32,
        bounds: BoundingBox,
    ) -> Self {
        Self {
            points: Vec::new(),
            colors: vec![style.color],
            point_size,
            point_sizes: None,
            edge_color: style.edge_color,
            edge_thickness: style.edge_thickness,
            marker_style: style.marker_style,
            filled: style.filled,
            edge_color_from_vertex_colors: style.edge_from_vertex_colors,
            label: None,
            visible: true,
            vertices: None,
            bounds: Some(bounds),
            gpu_vertices: Some(buffer),
            gpu_point_count: Some(point_count),
            gpu_has_per_point_colors: style.has_per_point_colors,
        }
    }

    /// Override all point colors with a single RGBA value.
    pub fn with_color(mut self, color: Vec4) -> Self {
        self.colors = vec![color; self.points.len()];
        self.vertices = None;
        self.gpu_vertices = None;
        self.gpu_point_count = None;
        self.gpu_has_per_point_colors = false;
        self
    }

    /// Supply per-point colors. Length must match the number of points.
    pub fn with_colors(mut self, colors: Vec<Vec4>) -> Result<Self, String> {
        if colors.len() != self.points.len() {
            return Err(format!(
                "Point cloud color count ({}) must match point count ({})",
                colors.len(),
                self.points.len()
            ));
        }
        self.colors = colors;
        self.vertices = None;
        self.gpu_vertices = None;
        self.gpu_point_count = None;
        self.gpu_has_per_point_colors = false;
        Ok(self)
    }

    /// Set the legend label.
    pub fn with_label<S: Into<String>>(mut self, label: S) -> Self {
        self.label = Some(label.into());
        self
    }

    /// Set marker size in pixels.
    pub fn with_point_size(mut self, size: f32) -> Self {
        self.point_size = size.max(1.0);
        self.point_sizes = None;
        self.gpu_vertices = None;
        self.gpu_point_count = None;
        self.gpu_has_per_point_colors = false;
        self
    }

    pub fn set_marker_style(&mut self, style: MarkerStyle) {
        self.marker_style = style;
        self.gpu_vertices = None;
        self.gpu_point_count = None;
        self.gpu_has_per_point_colors = false;
    }

    pub fn set_filled(&mut self, filled: bool) {
        self.filled = filled;
        self.gpu_vertices = None;
        self.gpu_point_count = None;
        self.gpu_has_per_point_colors = false;
    }

    pub fn set_edge_color(&mut self, color: Vec4) {
        self.edge_color = color;
        self.gpu_vertices = None;
        self.gpu_point_count = None;
        self.gpu_has_per_point_colors = false;
    }

    pub fn set_edge_thickness(&mut self, px: f32) {
        self.edge_thickness = px.max(0.0);
        self.gpu_vertices = None;
        self.gpu_point_count = None;
        self.gpu_has_per_point_colors = false;
    }

    pub fn set_edge_color_from_vertex(&mut self, enabled: bool) {
        self.edge_color_from_vertex_colors = enabled;
        self.gpu_vertices = None;
        self.gpu_point_count = None;
        self.gpu_has_per_point_colors = false;
    }

    /// Enable or disable visibility.
    pub fn set_visible(&mut self, visible: bool) {
        self.visible = visible;
    }

    /// Attach a GPU-resident vertex buffer that already encodes this point cloud in the renderer's vertex format.
    /// When provided, the renderer can skip per-frame uploads and reuse the supplied buffer directly.
    pub fn with_gpu_vertices(mut self, buffer: GpuVertexBuffer, point_count: usize) -> Self {
        self.gpu_vertices = Some(buffer);
        self.gpu_point_count = Some(point_count);
        self.vertices = None;
        self.gpu_has_per_point_colors = false;
        self
    }

    /// Supply per-point sizes in pixels.
    pub fn set_point_sizes(&mut self, sizes: Vec<f32>) {
        self.point_sizes = Some(sizes);
        self.vertices = None;
        self.gpu_vertices = None;
        self.gpu_point_count = None;
        self.gpu_has_per_point_colors = false;
    }

    fn ensure_vertices(&mut self) {
        if self.vertices.is_none() {
            let mut verts = vertex_utils::create_point_cloud(&self.points, &self.colors);
            if let Some(sizes) = self.point_sizes.as_ref() {
                for (idx, vertex) in verts.iter_mut().enumerate() {
                    let size = sizes.get(idx).copied().unwrap_or(self.point_size);
                    vertex.normal[2] = size;
                }
            } else {
                for vertex in &mut verts {
                    vertex.normal[2] = self.point_size;
                }
            }
            self.vertices = Some(verts);
        }
    }

    fn ensure_bounds(&mut self) {
        if self.bounds.is_none() {
            self.bounds = Some(BoundingBox::from_points(&self.points));
        }
    }

    /// Estimate memory required for this plot.
    pub fn estimated_memory_usage(&self) -> usize {
        let gpu_bytes = self
            .gpu_point_count
            .map(|count| count * std::mem::size_of::<Vertex>())
            .unwrap_or(0);
        self.points.len() * std::mem::size_of::<Vec3>()
            + self.colors.len() * std::mem::size_of::<Vec4>()
            + self
                .point_sizes
                .as_ref()
                .map(|sizes| sizes.len() * std::mem::size_of::<f32>())
                .unwrap_or(0)
            + gpu_bytes
    }

    /// Generate render data for the renderer.
    pub fn render_data(&mut self) -> RenderData {
        let bounds = self.bounds();
        let vertex_count = self.gpu_point_count.unwrap_or_else(|| {
            self.ensure_vertices();
            self.vertices
                .as_ref()
                .map(|v| v.len())
                .unwrap_or(self.points.len())
        });

        let vertices = if self.gpu_vertices.is_some() {
            Vec::new()
        } else {
            self.ensure_vertices();
            self.vertices.clone().unwrap_or_default()
        };

        let is_multi_color = if self.gpu_vertices.is_some() {
            self.gpu_has_per_point_colors || self.colors.len() > 1
        } else if vertices.is_empty() {
            false
        } else {
            let first = vertices[0].color;
            vertices.iter().any(|v| v.color != first)
        };
        let has_vertex_colors = if self.gpu_vertices.is_some() {
            self.gpu_has_per_point_colors
        } else {
            self.colors.len() > 1
        };
        let use_vertex_edge_color = self.edge_color_from_vertex_colors && has_vertex_colors;
        let mut material = Material {
            albedo: self.colors.first().copied().unwrap_or(Vec4::ONE),
            roughness: self.edge_thickness,
            metallic: match self.marker_style {
                MarkerStyle::Circle => 0.0,
                MarkerStyle::Square => 1.0,
                MarkerStyle::Triangle => 2.0,
                MarkerStyle::Diamond => 3.0,
                MarkerStyle::Plus => 4.0,
                MarkerStyle::Cross => 5.0,
                MarkerStyle::Star => 6.0,
                MarkerStyle::Hexagon => 7.0,
            },
            emissive: self.edge_color,
            alpha_mode: crate::core::scene::AlphaMode::Blend,
            double_sided: true,
        };
        if is_multi_color {
            material.albedo.w = 0.0;
        } else if self.filled {
            material.albedo.w = 1.0;
        }
        material.emissive.w = if use_vertex_edge_color { 0.0 } else { 1.0 };

        RenderData {
            pipeline_type: PipelineType::Scatter3,
            vertices,
            indices: None,
            gpu_vertices: self.gpu_vertices.clone(),
            bounds: Some(bounds),
            material,
            draw_calls: vec![DrawCall {
                vertex_offset: 0,
                vertex_count,
                index_offset: None,
                index_count: None,
                instance_count: 1,
            }],
            image: None,
        }
    }

    /// Compute the axis-aligned bounding box.
    pub fn bounds(&mut self) -> BoundingBox {
        self.ensure_bounds();
        self.bounds.unwrap_or_default()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn scatter3_defaults() {
        let points = vec![Vec3::new(0.0, 0.0, 0.0), Vec3::new(1.0, 2.0, 3.0)];
        let cloud = Scatter3Plot::new(points.clone()).unwrap();
        assert_eq!(cloud.points.len(), points.len());
        assert_eq!(cloud.colors.len(), points.len());
        assert!(cloud.visible);
    }

    #[test]
    fn scatter3_custom_colors() {
        let points = vec![Vec3::new(0.0, 0.0, 0.0)];
        let colors = vec![Vec4::new(1.0, 0.0, 0.0, 1.0)];
        let cloud = Scatter3Plot::new(points)
            .unwrap()
            .with_colors(colors)
            .unwrap();
        assert_eq!(cloud.colors[0], Vec4::new(1.0, 0.0, 0.0, 1.0));
    }

    #[test]
    fn scatter3_render_data_contains_vertices() {
        let points = vec![Vec3::new(0.0, 0.0, 0.0), Vec3::new(1.0, 1.0, 1.0)];
        let mut cloud = Scatter3Plot::new(points).unwrap();
        let render_data = cloud.render_data();
        assert_eq!(render_data.vertices.len(), 2);
        assert_eq!(render_data.pipeline_type, PipelineType::Scatter3);
    }

    #[test]
    fn scatter3_marker_style_encodes_material_shape_channel() {
        let points = vec![Vec3::new(0.0, 0.0, 0.0)];
        let mut cloud = Scatter3Plot::new(points).unwrap();
        cloud.set_marker_style(MarkerStyle::Diamond);
        let render_data = cloud.render_data();
        assert_eq!(render_data.material.metallic, 3.0);
    }

    #[test]
    fn scatter3_default_material_uses_plot_color_not_white_override() {
        let points = vec![Vec3::new(0.0, 0.0, 0.0)];
        let mut cloud = Scatter3Plot::new(points).unwrap();
        let render_data = cloud.render_data();
        assert_ne!(render_data.material.albedo.truncate(), Vec4::ONE.truncate());
    }
}