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