Skip to main content

runmat_plot/plots/
line3.rs

1use crate::context::shared_wgpu_context;
2use crate::core::{
3    BoundingBox, DrawCall, GpuPackContext, GpuVertexBuffer, Material, PipelineType, RenderData,
4    Vertex,
5};
6use crate::geometry::stroke3d::{
7    create_line_vertices_dashed, tessellate_polyline_tube, StrokeCap3D, StrokeStyle3D,
8};
9use crate::gpu::line3::{Line3GpuInputs, Line3GpuParams};
10use crate::gpu::util::readback_scalar_buffer_f64;
11use glam::{Vec3, Vec4};
12use log::warn;
13
14const POINTS_TO_PX: f32 = 96.0 / 72.0;
15const TUBE_RADIAL_SEGMENTS: usize = 8;
16
17#[derive(Debug, Clone)]
18pub struct Line3Plot {
19    pub x_data: Vec<f64>,
20    pub y_data: Vec<f64>,
21    pub z_data: Vec<f64>,
22    pub color: Vec4,
23    pub line_width: f32,
24    pub line_style: crate::plots::line::LineStyle,
25    pub label: Option<String>,
26    pub visible: bool,
27    vertices: Option<Vec<Vertex>>,
28    bounds: Option<BoundingBox>,
29    dirty: bool,
30    pub gpu_vertices: Option<GpuVertexBuffer>,
31    pub gpu_vertex_count: Option<usize>,
32    gpu_line_inputs: Option<Line3GpuInputs>,
33}
34
35impl Line3Plot {
36    #[inline]
37    fn line_width_px(&self) -> f32 {
38        (self.line_width.max(0.1)) * POINTS_TO_PX
39    }
40
41    pub fn new(x_data: Vec<f64>, y_data: Vec<f64>, z_data: Vec<f64>) -> Result<Self, String> {
42        if x_data.len() != y_data.len() || x_data.len() != z_data.len() {
43            return Err("Data length mismatch for plot3".to_string());
44        }
45        if x_data.is_empty() {
46            return Err("plot3 requires at least one point".to_string());
47        }
48        Ok(Self {
49            x_data,
50            y_data,
51            z_data,
52            color: Vec4::new(0.0, 0.5, 1.0, 1.0),
53            line_width: 1.0,
54            line_style: crate::plots::line::LineStyle::Solid,
55            label: None,
56            visible: true,
57            vertices: None,
58            bounds: None,
59            dirty: true,
60            gpu_vertices: None,
61            gpu_vertex_count: None,
62            gpu_line_inputs: None,
63        })
64    }
65
66    pub fn from_gpu_buffer(
67        buffer: GpuVertexBuffer,
68        vertex_count: usize,
69        color: Vec4,
70        line_width: f32,
71        line_style: crate::plots::line::LineStyle,
72        bounds: BoundingBox,
73    ) -> Self {
74        Self {
75            x_data: Vec::new(),
76            y_data: Vec::new(),
77            z_data: Vec::new(),
78            color,
79            line_width,
80            line_style,
81            label: None,
82            visible: true,
83            vertices: None,
84            bounds: Some(bounds),
85            dirty: false,
86            gpu_vertices: Some(buffer),
87            gpu_vertex_count: Some(vertex_count),
88            gpu_line_inputs: None,
89        }
90    }
91
92    pub fn from_gpu_xyz(
93        inputs: Line3GpuInputs,
94        color: Vec4,
95        line_width: f32,
96        line_style: crate::plots::line::LineStyle,
97        bounds: BoundingBox,
98    ) -> Self {
99        Self {
100            x_data: Vec::new(),
101            y_data: Vec::new(),
102            z_data: Vec::new(),
103            color,
104            line_width,
105            line_style,
106            label: None,
107            visible: true,
108            vertices: None,
109            bounds: Some(bounds),
110            dirty: false,
111            gpu_vertices: None,
112            gpu_vertex_count: None,
113            gpu_line_inputs: Some(inputs),
114        }
115    }
116
117    pub fn with_gpu_xyz_inputs(mut self, inputs: Line3GpuInputs, bounds: BoundingBox) -> Self {
118        self.gpu_line_inputs = Some(inputs);
119        self.bounds = Some(bounds);
120        self
121    }
122
123    pub async fn export_scene_xyz_data(&self) -> Result<(Vec<f64>, Vec<f64>, Vec<f64>), String> {
124        if !self.x_data.is_empty()
125            && self.x_data.len() == self.y_data.len()
126            && self.x_data.len() == self.z_data.len()
127        {
128            return Ok((
129                self.x_data.clone(),
130                self.y_data.clone(),
131                self.z_data.clone(),
132            ));
133        }
134        if !self.x_data.is_empty() || !self.y_data.is_empty() || !self.z_data.is_empty() {
135            return Err(format!(
136                "plot3 line has incomplete CPU data: x={}, y={}, z={}",
137                self.x_data.len(),
138                self.y_data.len(),
139                self.z_data.len()
140            ));
141        }
142
143        if let Some(inputs) = &self.gpu_line_inputs {
144            let context = shared_wgpu_context().ok_or_else(|| {
145                "plot3 line has GPU source data but no shared WGPU context is installed".to_string()
146            })?;
147            let len = inputs.len as usize;
148            let x = readback_scalar_buffer_f64(
149                &context.device,
150                &context.queue,
151                &inputs.x_buffer,
152                len,
153                inputs.scalar,
154            )
155            .await?;
156            let y = readback_scalar_buffer_f64(
157                &context.device,
158                &context.queue,
159                &inputs.y_buffer,
160                len,
161                inputs.scalar,
162            )
163            .await?;
164            let z = readback_scalar_buffer_f64(
165                &context.device,
166                &context.queue,
167                &inputs.z_buffer,
168                len,
169                inputs.scalar,
170            )
171            .await?;
172            return Ok((x, y, z));
173        }
174
175        if self.gpu_vertices.is_some() {
176            return Err(
177                "plot3 line has GPU render vertices but no exportable source data".to_string(),
178            );
179        }
180
181        Ok((Vec::new(), Vec::new(), Vec::new()))
182    }
183
184    pub fn with_label<S: Into<String>>(mut self, label: S) -> Self {
185        self.label = Some(label.into());
186        self
187    }
188
189    pub fn with_style(
190        mut self,
191        color: Vec4,
192        line_width: f32,
193        line_style: crate::plots::line::LineStyle,
194    ) -> Self {
195        self.color = color;
196        self.line_width = line_width;
197        self.line_style = line_style;
198        self.dirty = true;
199        self.gpu_vertices = None;
200        self.gpu_vertex_count = None;
201        self
202    }
203
204    pub fn update_data(
205        &mut self,
206        x_data: Vec<f64>,
207        y_data: Vec<f64>,
208        z_data: Vec<f64>,
209    ) -> Result<(), String> {
210        if x_data.len() != y_data.len() || x_data.len() != z_data.len() {
211            return Err("Data length mismatch for plot3".to_string());
212        }
213        if x_data.is_empty() {
214            return Err("plot3 requires at least one point".to_string());
215        }
216        self.x_data = x_data;
217        self.y_data = y_data;
218        self.z_data = z_data;
219        self.vertices = None;
220        self.bounds = None;
221        self.dirty = true;
222        self.gpu_vertices = None;
223        self.gpu_vertex_count = None;
224        self.gpu_line_inputs = None;
225        Ok(())
226    }
227
228    pub fn set_visible(&mut self, visible: bool) {
229        self.visible = visible;
230    }
231
232    fn generate_vertices(&mut self) -> &Vec<Vertex> {
233        if self.gpu_vertices.is_some() {
234            if self.vertices.is_none() {
235                self.vertices = Some(Vec::new());
236            }
237            return self.vertices.as_ref().unwrap();
238        }
239        if self.dirty || self.vertices.is_none() {
240            let points: Vec<Vec3> = self
241                .x_data
242                .iter()
243                .zip(self.y_data.iter())
244                .zip(self.z_data.iter())
245                .map(|((&x, &y), &z)| Vec3::new(x as f32, y as f32, z as f32))
246                .collect();
247            let vertices = if points.len() == 1 {
248                let mut vertex = Vertex::new(points[0], self.color);
249                vertex.normal[2] = (self.line_width_px().max(1.0) * 4.0).max(6.0);
250                vec![vertex]
251            } else if self.line_width_px() > 1.0 {
252                // No viewport hint: interpret width in data units for legacy/non-viewport paths.
253                let fallback_half_width_data = self.line_width_px() * 0.5;
254                tessellate_polyline_tube(
255                    &points,
256                    self.color,
257                    StrokeStyle3D::new(
258                        fallback_half_width_data,
259                        self.line_style,
260                        StrokeCap3D::Square,
261                    ),
262                    TUBE_RADIAL_SEGMENTS,
263                )
264            } else {
265                create_line_vertices_dashed(&points, self.color, self.line_style)
266            };
267            self.vertices = Some(vertices);
268            self.dirty = false;
269        }
270        self.vertices.as_ref().unwrap()
271    }
272
273    pub fn bounds(&mut self) -> BoundingBox {
274        if self.bounds.is_some() && self.x_data.is_empty() {
275            return self.bounds.unwrap();
276        }
277        if self.bounds.is_none() || self.dirty {
278            let points: Vec<Vec3> = self
279                .x_data
280                .iter()
281                .zip(self.y_data.iter())
282                .zip(self.z_data.iter())
283                .map(|((&x, &y), &z)| Vec3::new(x as f32, y as f32, z as f32))
284                .collect();
285            self.bounds = Some(BoundingBox::from_points(&points));
286        }
287        self.bounds.unwrap()
288    }
289
290    pub fn render_data(&mut self) -> RenderData {
291        let single_point = self.x_data.len() == 1 || self.gpu_vertex_count == Some(1);
292        let vertex_count = self
293            .gpu_vertex_count
294            .unwrap_or_else(|| self.generate_vertices().len());
295        let width_px = self.line_width_px();
296        let thick = width_px > 1.0 && !single_point;
297        let indices = if self.gpu_vertices.is_none() && thick {
298            Some((0..vertex_count as u32).collect::<Vec<u32>>())
299        } else {
300            None
301        };
302        RenderData {
303            pipeline_type: if single_point {
304                PipelineType::Scatter3
305            } else if thick {
306                PipelineType::Triangles
307            } else {
308                PipelineType::Lines
309            },
310            vertices: if self.gpu_vertices.is_some() {
311                Vec::new()
312            } else {
313                self.generate_vertices().clone()
314            },
315            indices,
316            gpu_vertices: self.gpu_vertices.clone(),
317            bounds: Some(self.bounds()),
318            material: Material {
319                albedo: self.color,
320                roughness: width_px.max(0.5),
321                ..Default::default()
322            },
323            draw_calls: vec![DrawCall {
324                vertex_offset: 0,
325                vertex_count,
326                index_offset: None,
327                index_count: None,
328                instance_count: 1,
329            }],
330            image: None,
331        }
332    }
333
334    fn pack_gpu_vertices_if_needed(
335        &mut self,
336        gpu: &GpuPackContext<'_>,
337        viewport_px: (u32, u32),
338    ) -> Result<(), String> {
339        if self.gpu_vertices.is_some() {
340            return Ok(());
341        }
342        let Some(inputs) = self.gpu_line_inputs.as_ref() else {
343            return Ok(());
344        };
345        let bounds = self
346            .bounds
347            .as_ref()
348            .ok_or_else(|| "plot3: missing bounds for GPU packing".to_string())?;
349        let width_px = self.line_width_px();
350        let thick_px = width_px > 1.0;
351        let data_per_px = crate::core::data_units_per_px_3d(bounds, viewport_px);
352        let half_width_data = if thick_px {
353            (width_px * 0.5) * data_per_px
354        } else {
355            0.0
356        };
357        let packed = crate::gpu::line3::pack_vertices_from_xyz(
358            gpu.device,
359            gpu.queue,
360            inputs,
361            &Line3GpuParams {
362                color: self.color,
363                half_width_data,
364                thick: thick_px,
365                line_style: self.line_style,
366            },
367        )?;
368        self.gpu_vertex_count =
369            Some((inputs.len.saturating_sub(1) as usize) * if thick_px { 6 } else { 2 });
370        self.gpu_vertices = Some(packed);
371        Ok(())
372    }
373
374    pub fn render_data_with_viewport_gpu(
375        &mut self,
376        viewport_px: Option<(u32, u32)>,
377        view_angles_deg: Option<(f32, f32)>,
378        gpu: Option<&GpuPackContext<'_>>,
379    ) -> RenderData {
380        if self.gpu_line_inputs.is_some() && self.gpu_vertices.is_none() {
381            if let (Some(gpu), Some(vp)) = (gpu, viewport_px) {
382                if let Err(err) = self.pack_gpu_vertices_if_needed(gpu, vp) {
383                    warn!("plot3 gpu pack failed: {err}");
384                }
385            }
386        }
387        self.render_data_with_viewport_and_view(viewport_px, view_angles_deg)
388    }
389
390    pub fn render_data_with_viewport(&mut self, viewport_px: Option<(u32, u32)>) -> RenderData {
391        self.render_data_with_viewport_and_view(viewport_px, None)
392    }
393
394    pub fn render_data_with_viewport_and_view(
395        &mut self,
396        viewport_px: Option<(u32, u32)>,
397        view_angles_deg: Option<(f32, f32)>,
398    ) -> RenderData {
399        if self.gpu_vertices.is_some() {
400            return self.render_data();
401        }
402
403        let single_point = self.x_data.len() == 1;
404        let width_px = self.line_width_px();
405        let (vertices, vertex_count, pipeline) = if !single_point && width_px > 1.0 {
406            let Some(vp) = viewport_px else {
407                return self.render_data();
408            };
409            let points: Vec<Vec3> = self
410                .x_data
411                .iter()
412                .zip(self.y_data.iter())
413                .zip(self.z_data.iter())
414                .map(|((&x, &y), &z)| Vec3::new(x as f32, y as f32, z as f32))
415                .collect();
416            let bounds = self.bounds();
417            let data_per_px =
418                crate::core::data_units_per_px_3d_camera(&bounds, vp, view_angles_deg);
419            let half_width_data = (width_px * 0.5) * data_per_px;
420            let tris = tessellate_polyline_tube(
421                &points,
422                self.color,
423                StrokeStyle3D::new(half_width_data, self.line_style, StrokeCap3D::Square),
424                TUBE_RADIAL_SEGMENTS,
425            );
426            let count = tris.len();
427            (tris, count, PipelineType::Triangles)
428        } else {
429            let verts = self.generate_vertices().clone();
430            let count = verts.len();
431            let pipeline = if single_point {
432                PipelineType::Scatter3
433            } else {
434                PipelineType::Lines
435            };
436            (verts, count, pipeline)
437        };
438
439        let indices = if pipeline == PipelineType::Triangles {
440            Some((0..vertex_count as u32).collect::<Vec<u32>>())
441        } else {
442            None
443        };
444
445        RenderData {
446            pipeline_type: pipeline,
447            vertices,
448            indices,
449            gpu_vertices: None,
450            bounds: Some(self.bounds()),
451            material: Material {
452                albedo: self.color,
453                roughness: width_px.max(0.5),
454                ..Default::default()
455            },
456            draw_calls: vec![DrawCall {
457                vertex_offset: 0,
458                vertex_count,
459                index_offset: None,
460                index_count: None,
461                instance_count: 1,
462            }],
463            image: None,
464        }
465    }
466
467    pub fn estimated_memory_usage(&self) -> usize {
468        self.vertices
469            .as_ref()
470            .map(|v| v.len() * std::mem::size_of::<Vertex>())
471            .unwrap_or(0)
472    }
473}
474
475#[cfg(test)]
476mod tests {
477    use super::*;
478    use pollster::FutureExt;
479    use std::sync::Arc;
480    use wgpu::util::DeviceExt;
481
482    fn maybe_device() -> Option<(Arc<wgpu::Device>, Arc<wgpu::Queue>)> {
483        if std::env::var("RUNMAT_PLOT_SKIP_GPU_TESTS").is_ok()
484            || std::env::var("RUNMAT_PLOT_FORCE_GPU_TESTS").is_err()
485        {
486            return None;
487        }
488        let instance = wgpu::Instance::default();
489        let adapter = instance
490            .request_adapter(&wgpu::RequestAdapterOptions {
491                power_preference: wgpu::PowerPreference::HighPerformance,
492                compatible_surface: None,
493                force_fallback_adapter: false,
494            })
495            .block_on()?;
496        let (device, queue) = adapter
497            .request_device(
498                &crate::wgpu_compat::device_descriptor(
499                    Some("runmat-plot-line3-test-device"),
500                    wgpu::Features::empty(),
501                    adapter.limits(),
502                ),
503                None,
504            )
505            .block_on()
506            .ok()?;
507        Some((Arc::new(device), Arc::new(queue)))
508    }
509
510    #[test]
511    fn line3_creation_and_bounds() {
512        let mut plot = Line3Plot::new(vec![0.0, 1.0], vec![1.0, 2.0], vec![2.0, 3.0]).unwrap();
513        let bounds = plot.bounds();
514        assert_eq!(bounds.min, Vec3::new(0.0, 1.0, 2.0));
515        assert_eq!(bounds.max, Vec3::new(1.0, 2.0, 3.0));
516    }
517
518    #[test]
519    fn line3_dashed_and_thick_generate_geometry() {
520        let mut plot = Line3Plot::new(
521            vec![0.0, 1.0, 2.0],
522            vec![0.0, 1.0, 0.0],
523            vec![0.0, 0.0, 1.0],
524        )
525        .unwrap()
526        .with_style(Vec4::ONE, 3.0, crate::plots::line::LineStyle::Dashed);
527        let render = plot.render_data();
528        assert_eq!(render.pipeline_type, PipelineType::Triangles);
529        assert!(!render.vertices.is_empty());
530        assert!(render.draw_calls[0].vertex_count >= 2);
531    }
532
533    #[test]
534    fn line3_single_point_uses_scatter_pipeline() {
535        let mut plot = Line3Plot::new(vec![1.0], vec![2.0], vec![3.0])
536            .unwrap()
537            .with_style(Vec4::ONE, 2.0, crate::plots::line::LineStyle::Solid);
538        let render = plot.render_data();
539        assert_eq!(render.pipeline_type, PipelineType::Scatter3);
540        assert_eq!(render.vertices.len(), 1);
541        assert!(render.vertices[0].normal[2] >= 6.0);
542    }
543
544    #[test]
545    fn line3_gpu_source_packs_default_width_as_thick_geometry() {
546        let Some((device, queue)) = maybe_device() else {
547            return;
548        };
549        let x_buffer = Arc::new(
550            device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
551                label: Some("line3-test-x"),
552                contents: bytemuck::cast_slice(&[0.0f32, 1.0f32]),
553                usage: wgpu::BufferUsages::STORAGE,
554            }),
555        );
556        let y_buffer = Arc::new(
557            device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
558                label: Some("line3-test-y"),
559                contents: bytemuck::cast_slice(&[0.0f32, 1.0f32]),
560                usage: wgpu::BufferUsages::STORAGE,
561            }),
562        );
563        let z_buffer = Arc::new(
564            device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
565                label: Some("line3-test-z"),
566                contents: bytemuck::cast_slice(&[0.0f32, 1.0f32]),
567                usage: wgpu::BufferUsages::STORAGE,
568            }),
569        );
570        let mut plot = Line3Plot::from_gpu_xyz(
571            Line3GpuInputs {
572                x_buffer,
573                y_buffer,
574                z_buffer,
575                len: 2,
576                scalar: crate::gpu::ScalarType::F32,
577            },
578            Vec4::ONE,
579            1.0,
580            crate::plots::line::LineStyle::Solid,
581            BoundingBox::new(Vec3::ZERO, Vec3::ONE),
582        );
583        let gpu = GpuPackContext {
584            device: &device,
585            queue: &queue,
586        };
587        let render = plot.render_data_with_viewport_gpu(Some((800, 600)), None, Some(&gpu));
588        assert_eq!(render.pipeline_type, PipelineType::Triangles);
589        assert_eq!(render.vertices.len(), 0);
590        assert!(render.gpu_vertices.is_some());
591        assert_eq!(render.draw_calls[0].vertex_count, 6);
592    }
593}