1use crate::context::shared_wgpu_context;
4use crate::core::{
5 marker_shape_code, vertex_utils, AlphaMode, BoundingBox, DrawCall, GpuVertexBuffer, Material,
6 PipelineType, RenderData, Vertex,
7};
8use crate::gpu::stem::StemGpuInputs;
9use crate::gpu::util::readback_scalar_buffer_f64;
10use crate::plots::line::{LineMarkerAppearance, LineStyle};
11use glam::{Vec3, Vec4};
12
13#[derive(Debug, Clone)]
14pub struct StemPlot {
15 pub x: Vec<f64>,
16 pub y: Vec<f64>,
17 pub baseline: f64,
18 pub color: Vec4,
19 pub line_width: f32,
20 pub line_style: LineStyle,
21 pub baseline_color: Vec4,
22 pub baseline_visible: bool,
23 pub marker: Option<LineMarkerAppearance>,
24 pub label: Option<String>,
25 pub visible: bool,
26 vertices: Option<Vec<Vertex>>,
27 bounds: Option<BoundingBox>,
28 dirty: bool,
29 gpu_vertices: Option<GpuVertexBuffer>,
30 gpu_vertex_count: Option<usize>,
31 gpu_bounds: Option<BoundingBox>,
32 gpu_inputs: Option<StemGpuInputs>,
33 marker_vertices: Option<Vec<Vertex>>,
34 marker_gpu_vertices: Option<GpuVertexBuffer>,
35 marker_dirty: bool,
36}
37
38impl StemPlot {
39 pub async fn export_scene_xy_data(&self) -> Result<(Vec<f64>, Vec<f64>), String> {
40 if !self.x.is_empty() && self.x.len() == self.y.len() {
41 return Ok((self.x.clone(), self.y.clone()));
42 }
43 if !self.x.is_empty() || !self.y.is_empty() {
44 return Err(format!(
45 "stem plot has partial CPU source data: x has {} values, y has {} values",
46 self.x.len(),
47 self.y.len()
48 ));
49 }
50
51 if let Some(inputs) = &self.gpu_inputs {
52 let context = shared_wgpu_context().ok_or_else(|| {
53 "stem plot has GPU source data but no shared WGPU context is installed".to_string()
54 })?;
55 let len = inputs.len as usize;
56 let x = readback_scalar_buffer_f64(
57 &context.device,
58 &context.queue,
59 &inputs.x_buffer,
60 len,
61 inputs.scalar,
62 )
63 .await?;
64 let y = readback_scalar_buffer_f64(
65 &context.device,
66 &context.queue,
67 &inputs.y_buffer,
68 len,
69 inputs.scalar,
70 )
71 .await?;
72 return Ok((x, y));
73 }
74
75 if self.gpu_vertices.is_some() {
76 return Err(
77 "stem plot has GPU render vertices but no exportable source data".to_string(),
78 );
79 }
80
81 Ok((Vec::new(), Vec::new()))
82 }
83
84 pub fn new(x: Vec<f64>, y: Vec<f64>) -> Result<Self, String> {
85 if x.len() != y.len() || x.is_empty() {
86 return Err("stem: X and Y must be same non-zero length".to_string());
87 }
88 Ok(Self {
89 x,
90 y,
91 baseline: 0.0,
92 color: Vec4::new(0.0, 0.447, 0.741, 1.0),
93 line_width: 1.0,
94 line_style: LineStyle::Solid,
95 baseline_color: Vec4::new(0.15, 0.15, 0.15, 1.0),
96 baseline_visible: true,
97 marker: Some(LineMarkerAppearance {
98 kind: crate::plots::scatter::MarkerStyle::Circle,
99 size: 6.0,
100 edge_color: Vec4::new(0.0, 0.447, 0.741, 1.0),
101 face_color: Vec4::new(0.0, 0.447, 0.741, 1.0),
102 filled: false,
103 }),
104 label: None,
105 visible: true,
106 vertices: None,
107 bounds: None,
108 dirty: true,
109 gpu_vertices: None,
110 gpu_vertex_count: None,
111 gpu_bounds: None,
112 gpu_inputs: None,
113 marker_vertices: None,
114 marker_gpu_vertices: None,
115 marker_dirty: true,
116 })
117 }
118
119 #[allow(clippy::too_many_arguments)]
120 pub fn from_gpu_buffer(
121 color: Vec4,
122 line_width: f32,
123 line_style: LineStyle,
124 baseline: f64,
125 baseline_color: Vec4,
126 baseline_visible: bool,
127 buffer: GpuVertexBuffer,
128 vertex_count: usize,
129 bounds: BoundingBox,
130 ) -> Self {
131 Self {
132 x: Vec::new(),
133 y: Vec::new(),
134 baseline,
135 color,
136 line_width,
137 line_style,
138 baseline_color,
139 baseline_visible,
140 marker: None,
141 label: None,
142 visible: true,
143 vertices: None,
144 bounds: None,
145 dirty: false,
146 gpu_vertices: Some(buffer),
147 gpu_vertex_count: Some(vertex_count),
148 gpu_bounds: Some(bounds),
149 gpu_inputs: None,
150 marker_vertices: None,
151 marker_gpu_vertices: None,
152 marker_dirty: true,
153 }
154 }
155
156 pub fn with_style(
157 mut self,
158 color: Vec4,
159 line_width: f32,
160 line_style: LineStyle,
161 baseline: f64,
162 ) -> Self {
163 self.color = color;
164 self.line_width = line_width.max(0.5);
165 self.line_style = line_style;
166 self.baseline = baseline;
167 self.dirty = true;
168 self.marker_dirty = true;
169 self.gpu_vertices = None;
170 self.gpu_vertex_count = None;
171 self.gpu_bounds = None;
172 self.marker_gpu_vertices = None;
173 self
174 }
175
176 pub fn with_gpu_source_inputs(mut self, inputs: StemGpuInputs) -> Self {
177 self.gpu_inputs = Some(inputs);
178 self
179 }
180
181 pub fn with_baseline_style(mut self, color: Vec4, visible: bool) -> Self {
182 self.baseline_color = color;
183 self.baseline_visible = visible;
184 self.dirty = true;
185 self
186 }
187
188 pub fn with_label<S: Into<String>>(mut self, label: S) -> Self {
189 self.label = Some(label.into());
190 self
191 }
192
193 pub fn set_visible(&mut self, visible: bool) {
194 self.visible = visible;
195 }
196
197 pub fn set_marker(&mut self, marker: Option<LineMarkerAppearance>) {
198 self.marker = marker;
199 self.marker_dirty = true;
200 if self.marker.is_none() {
201 self.marker_vertices = None;
202 self.marker_gpu_vertices = None;
203 }
204 }
205
206 pub fn set_marker_gpu_vertices(&mut self, buffer: Option<GpuVertexBuffer>) {
207 let has_gpu = buffer.is_some();
208 self.marker_gpu_vertices = buffer;
209 if has_gpu {
210 self.marker_vertices = None;
211 }
212 }
213
214 pub fn generate_vertices(&mut self) -> &Vec<Vertex> {
215 if self.gpu_vertices.is_some() {
216 if self.vertices.is_none() {
217 self.vertices = Some(Vec::new());
218 }
219 return self.vertices.as_ref().unwrap();
220 }
221 if self.dirty || self.vertices.is_none() {
222 let mut vertices = Vec::new();
223 let finite_x: Vec<f32> = self
224 .x
225 .iter()
226 .map(|v| *v as f32)
227 .filter(|v| v.is_finite())
228 .collect();
229 if self.baseline_visible && !finite_x.is_empty() {
230 let min_x = finite_x.iter().copied().fold(f32::INFINITY, f32::min);
231 let max_x = finite_x.iter().copied().fold(f32::NEG_INFINITY, f32::max);
232 vertices.push(Vertex::new(
233 Vec3::new(min_x, self.baseline as f32, 0.0),
234 self.baseline_color,
235 ));
236 vertices.push(Vertex::new(
237 Vec3::new(max_x, self.baseline as f32, 0.0),
238 self.baseline_color,
239 ));
240 }
241 for i in 0..self.x.len() {
242 let x = self.x[i] as f32;
243 let y = self.y[i] as f32;
244 let b = self.baseline as f32;
245 if !x.is_finite() || !y.is_finite() {
246 continue;
247 }
248 if include_segment(i, self.line_style) {
249 vertices.push(Vertex::new(Vec3::new(x, b, 0.0), self.color));
250 vertices.push(Vertex::new(Vec3::new(x, y, 0.0), self.color));
251 }
252 }
253 self.vertices = Some(vertices);
254 self.dirty = false;
255 }
256 self.vertices.as_ref().unwrap()
257 }
258
259 pub fn marker_render_data(&mut self) -> Option<RenderData> {
260 let marker = self.marker.clone()?;
261 if let Some(gpu_vertices) = self.marker_gpu_vertices.clone() {
262 let vertex_count = gpu_vertices.vertex_count;
263 if vertex_count == 0 {
264 return None;
265 }
266 return Some(RenderData {
267 pipeline_type: PipelineType::Points,
268 vertices: Vec::new(),
269 indices: None,
270 gpu_vertices: Some(gpu_vertices),
271 bounds: None,
272 material: Material {
273 albedo: marker.face_color,
274 emissive: marker.edge_color,
275 roughness: 1.0,
276 metallic: marker_shape_code(marker.kind) as f32,
277 alpha_mode: if marker.face_color.w < 0.999 {
278 AlphaMode::Blend
279 } else {
280 AlphaMode::Opaque
281 },
282 ..Default::default()
283 },
284 draw_calls: vec![DrawCall {
285 vertex_offset: 0,
286 vertex_count,
287 index_offset: None,
288 index_count: None,
289 instance_count: 1,
290 }],
291 image: None,
292 });
293 }
294 if self.marker_dirty || self.marker_vertices.is_none() {
295 let mut vertices = Vec::new();
296 for (&x, &y) in self.x.iter().zip(self.y.iter()) {
297 let x = x as f32;
298 let y = y as f32;
299 if !x.is_finite() || !y.is_finite() {
300 continue;
301 }
302 vertices.push(Vertex {
303 position: [x, y, 0.0],
304 color: marker.face_color.to_array(),
305 normal: [0.0, 0.0, marker.size],
306 tex_coords: [0.0, 0.0],
307 });
308 }
309 self.marker_vertices = Some(vertices);
310 self.marker_dirty = false;
311 }
312 let vertices = self.marker_vertices.as_ref()?;
313 if vertices.is_empty() {
314 return None;
315 }
316 Some(RenderData {
317 pipeline_type: PipelineType::Points,
318 vertices: vertices.clone(),
319 indices: None,
320 gpu_vertices: None,
321 bounds: None,
322 material: Material {
323 albedo: marker.face_color,
324 emissive: marker.edge_color,
325 roughness: 1.0,
326 metallic: marker_shape_code(marker.kind) as f32,
327 alpha_mode: if marker.face_color.w < 0.999 {
328 AlphaMode::Blend
329 } else {
330 AlphaMode::Opaque
331 },
332 ..Default::default()
333 },
334 draw_calls: vec![DrawCall {
335 vertex_offset: 0,
336 vertex_count: vertices.len(),
337 index_offset: None,
338 index_count: None,
339 instance_count: 1,
340 }],
341 image: None,
342 })
343 }
344
345 pub fn bounds(&mut self) -> BoundingBox {
346 if let Some(bounds) = self.gpu_bounds {
347 return bounds;
348 }
349 if self.dirty || self.bounds.is_none() {
350 let mut min = Vec3::new(f32::INFINITY, f32::INFINITY, 0.0);
351 let mut max = Vec3::new(f32::NEG_INFINITY, f32::NEG_INFINITY, 0.0);
352 for (&x, &y) in self.x.iter().zip(self.y.iter()) {
353 let (x, y) = (x as f32, y as f32);
354 if !x.is_finite() || !y.is_finite() {
355 continue;
356 }
357 min.x = min.x.min(x);
358 max.x = max.x.max(x);
359 min.y = min.y.min(y.min(self.baseline as f32));
360 max.y = max.y.max(y.max(self.baseline as f32));
361 }
362 if !min.x.is_finite() {
363 min = Vec3::ZERO;
364 max = Vec3::ZERO;
365 }
366 self.bounds = Some(BoundingBox::new(min, max));
367 }
368 self.bounds.unwrap()
369 }
370
371 pub fn render_data(&mut self) -> RenderData {
372 let bounds = self.bounds();
373 let (vertices, vertex_count, gpu_vertices) = if self.gpu_vertices.is_some() {
374 (
375 Vec::new(),
376 self.gpu_vertex_count.unwrap_or(0),
377 self.gpu_vertices.clone(),
378 )
379 } else {
380 let vertices = self.generate_vertices().clone();
381 let count = vertices.len();
382 (vertices, count, None)
383 };
384 RenderData {
385 pipeline_type: PipelineType::Lines,
386 vertices,
387 indices: None,
388 gpu_vertices,
389 bounds: Some(bounds),
390 material: Material {
391 albedo: self.color,
392 roughness: self.line_width,
393 ..Default::default()
394 },
395 draw_calls: vec![DrawCall {
396 vertex_offset: 0,
397 vertex_count,
398 index_offset: None,
399 index_count: None,
400 instance_count: 1,
401 }],
402 image: None,
403 }
404 }
405
406 pub fn render_data_with_viewport(&mut self, viewport_px: Option<(u32, u32)>) -> RenderData {
407 if self.gpu_vertices.is_some() {
408 return self.render_data();
409 }
410
411 let bounds = self.bounds();
412 let (vertices, vertex_count, pipeline_type) = if self.line_width > 1.0 {
413 let Some(viewport_px) = viewport_px else {
414 return self.render_data();
415 };
416 let data_per_px = crate::core::data_units_per_px(&bounds, viewport_px);
417 let width_data = self.line_width.max(0.1) * data_per_px;
418 let verts = self.generate_vertices().clone();
419 let mut thick = Vec::new();
420 for segment in verts.chunks_exact(2) {
421 let x = [segment[0].position[0] as f64, segment[1].position[0] as f64];
422 let y = [segment[0].position[1] as f64, segment[1].position[1] as f64];
423 let color = Vec4::from_array(segment[0].color);
424 thick.extend(vertex_utils::create_thick_polyline(
425 &x, &y, color, width_data,
426 ));
427 }
428 let count = thick.len();
429 (thick, count, PipelineType::Triangles)
430 } else {
431 let verts = self.generate_vertices().clone();
432 let count = verts.len();
433 (verts, count, PipelineType::Lines)
434 };
435 RenderData {
436 pipeline_type,
437 vertices,
438 indices: None,
439 gpu_vertices: None,
440 bounds: Some(bounds),
441 material: Material {
442 albedo: self.color,
443 roughness: self.line_width.max(0.0),
444 ..Default::default()
445 },
446 draw_calls: vec![DrawCall {
447 vertex_offset: 0,
448 vertex_count,
449 index_offset: None,
450 index_count: None,
451 instance_count: 1,
452 }],
453 image: None,
454 }
455 }
456
457 pub fn estimated_memory_usage(&self) -> usize {
458 self.vertices
459 .as_ref()
460 .map_or(0, |v| v.len() * std::mem::size_of::<Vertex>())
461 + self
462 .marker_vertices
463 .as_ref()
464 .map_or(0, |v| v.len() * std::mem::size_of::<Vertex>())
465 }
466}
467
468fn include_segment(index: usize, style: LineStyle) -> bool {
469 match style {
470 LineStyle::Solid => true,
471 LineStyle::Dashed => (index % 4) < 2,
472 LineStyle::Dotted => index.is_multiple_of(4),
473 LineStyle::DashDot => {
474 let m = index % 6;
475 m < 2 || m == 3
476 }
477 }
478}
479
480#[cfg(test)]
481mod tests {
482 use super::*;
483
484 #[test]
485 fn stem_bounds_include_baseline() {
486 let mut plot = StemPlot::new(vec![0.0, 1.0], vec![1.0, -2.0])
487 .unwrap()
488 .with_style(Vec4::ONE, 1.0, LineStyle::Solid, -1.0);
489 let bounds = plot.bounds();
490 assert_eq!(bounds.min.y, -2.0);
491 assert_eq!(bounds.max.y, 1.0);
492 }
493
494 #[test]
495 fn thick_stem_use_viewport_aware_triangles() {
496 let mut plot = StemPlot::new(vec![0.0, 1.0], vec![1.0, 2.0])
497 .unwrap()
498 .with_style(Vec4::ONE, 2.0, LineStyle::Solid, 0.0);
499 let render = plot.render_data_with_viewport(Some((600, 400)));
500 assert_eq!(render.pipeline_type, PipelineType::Triangles);
501 assert!(!render.vertices.is_empty());
502 }
503}