Skip to main content

dear_implot/plots/
stems.rs

1//! Stem plot implementation
2
3use super::{
4    Plot, PlotDataLayout, PlotDataOffset, PlotDataStride, PlotError, PlotItemStyle,
5    plot_spec_with_style, validate_data_lengths, with_plot_str_or_empty,
6};
7use crate::{ItemFlags, StemsFlags, sys};
8
9/// Builder for stem plots (lollipop charts)
10pub struct StemPlot<'a> {
11    label: &'a str,
12    x_data: &'a [f64],
13    y_data: &'a [f64],
14    style: PlotItemStyle,
15    y_ref: f64,
16    flags: StemsFlags,
17    item_flags: ItemFlags,
18    layout: PlotDataLayout,
19}
20
21impl<'a> super::PlotItemStyled for StemPlot<'a> {
22    fn style_mut(&mut self) -> &mut PlotItemStyle {
23        &mut self.style
24    }
25}
26
27impl<'a> StemPlot<'a> {
28    /// Create a new stem plot with the given label and data
29    pub fn new(label: &'a str, x_data: &'a [f64], y_data: &'a [f64]) -> Self {
30        Self {
31            label,
32            x_data,
33            y_data,
34            style: PlotItemStyle::default(),
35            y_ref: 0.0, // Default reference line at Y=0
36            flags: StemsFlags::NONE,
37            item_flags: ItemFlags::NONE,
38            layout: PlotDataLayout::DEFAULT,
39        }
40    }
41
42    /// Set the reference Y value for stems
43    /// Stems will be drawn from this Y value to the data points
44    pub fn with_y_ref(mut self, y_ref: f64) -> Self {
45        self.y_ref = y_ref;
46        self
47    }
48
49    /// Set stem flags for customization
50    pub fn with_flags(mut self, flags: StemsFlags) -> Self {
51        self.flags = flags;
52        self
53    }
54
55    /// Set common item flags for this plot item (applies to all plot types)
56    pub fn with_item_flags(mut self, flags: ItemFlags) -> Self {
57        self.item_flags = flags;
58        self
59    }
60
61    /// Set the data layout used to read X/Y samples.
62    pub fn with_data_layout(mut self, layout: PlotDataLayout) -> Self {
63        self.layout = layout;
64        self
65    }
66
67    /// Set the sample-index offset used to read X/Y samples.
68    pub fn with_offset(mut self, offset: PlotDataOffset) -> Self {
69        self.layout = self.layout.with_offset(offset);
70        self
71    }
72
73    /// Set the byte stride used to read X/Y samples.
74    pub fn with_stride(mut self, stride: PlotDataStride) -> Self {
75        self.layout = self.layout.with_stride(stride);
76        self
77    }
78
79    /// Validate the plot data
80    pub fn validate(&self) -> Result<(), PlotError> {
81        validate_data_lengths(self.x_data, self.y_data)
82    }
83}
84
85impl<'a> Plot for StemPlot<'a> {
86    fn plot(&self, plot_ui: &crate::PlotUi<'_>) {
87        if self.validate().is_err() {
88            return;
89        }
90        let Ok(count) = i32::try_from(self.x_data.len()) else {
91            return;
92        };
93
94        let _guard = plot_ui.bind();
95        with_plot_str_or_empty(self.label, |label_ptr| unsafe {
96            let spec = plot_spec_with_style(
97                self.style,
98                self.flags.bits() | self.item_flags.bits(),
99                self.layout,
100            );
101            sys::ImPlot_PlotStems_doublePtrdoublePtr(
102                label_ptr,
103                self.x_data.as_ptr(),
104                self.y_data.as_ptr(),
105                count,
106                self.y_ref,
107                spec,
108            );
109        })
110    }
111
112    fn label(&self) -> &str {
113        self.label
114    }
115}
116
117/// Simple stem plot for quick plotting without builder pattern
118pub struct SimpleStemPlot<'a> {
119    label: &'a str,
120    values: &'a [f64],
121    style: PlotItemStyle,
122    y_ref: f64,
123    flags: StemsFlags,
124    item_flags: ItemFlags,
125    x_scale: f64,
126    x_start: f64,
127}
128
129impl<'a> super::PlotItemStyled for SimpleStemPlot<'a> {
130    fn style_mut(&mut self) -> &mut PlotItemStyle {
131        &mut self.style
132    }
133}
134
135impl<'a> SimpleStemPlot<'a> {
136    /// Create a simple stem plot with Y values only (X will be indices)
137    pub fn new(label: &'a str, values: &'a [f64]) -> Self {
138        Self {
139            label,
140            values,
141            style: PlotItemStyle::default(),
142            y_ref: 0.0,
143            flags: StemsFlags::NONE,
144            item_flags: ItemFlags::NONE,
145            x_scale: 1.0,
146            x_start: 0.0,
147        }
148    }
149
150    /// Set the reference Y value for stems
151    pub fn with_y_ref(mut self, y_ref: f64) -> Self {
152        self.y_ref = y_ref;
153        self
154    }
155
156    /// Set stem flags for customization
157    pub fn with_flags(mut self, flags: StemsFlags) -> Self {
158        self.flags = flags;
159        self
160    }
161
162    /// Set common item flags for this plot item (applies to all plot types)
163    pub fn with_item_flags(mut self, flags: ItemFlags) -> Self {
164        self.item_flags = flags;
165        self
166    }
167
168    /// Set X scale factor
169    pub fn with_x_scale(mut self, scale: f64) -> Self {
170        self.x_scale = scale;
171        self
172    }
173
174    /// Set X start value
175    pub fn with_x_start(mut self, start: f64) -> Self {
176        self.x_start = start;
177        self
178    }
179}
180
181impl<'a> Plot for SimpleStemPlot<'a> {
182    fn plot(&self, plot_ui: &crate::PlotUi<'_>) {
183        if self.values.is_empty() {
184            return;
185        }
186        let Ok(count) = i32::try_from(self.values.len()) else {
187            return;
188        };
189
190        let _guard = plot_ui.bind();
191        with_plot_str_or_empty(self.label, |label_ptr| unsafe {
192            let spec = plot_spec_with_style(
193                self.style,
194                self.flags.bits() | self.item_flags.bits(),
195                PlotDataLayout::DEFAULT,
196            );
197            sys::ImPlot_PlotStems_doublePtrInt(
198                label_ptr,
199                self.values.as_ptr(),
200                count,
201                self.y_ref,
202                self.x_scale,
203                self.x_start,
204                spec,
205            );
206        })
207    }
208
209    fn label(&self) -> &str {
210        self.label
211    }
212}
213
214/// Convenience functions for quick stem plotting
215impl<'ui> crate::PlotUi<'ui> {
216    /// Plot a stem plot with X and Y data
217    pub fn stem_plot(&self, label: &str, x_data: &[f64], y_data: &[f64]) -> Result<(), PlotError> {
218        let plot = StemPlot::new(label, x_data, y_data);
219        plot.validate()?;
220        plot.plot(self);
221        Ok(())
222    }
223
224    /// Plot a stem plot with custom reference Y value
225    pub fn stem_plot_with_ref(
226        &self,
227        label: &str,
228        x_data: &[f64],
229        y_data: &[f64],
230        y_ref: f64,
231    ) -> Result<(), PlotError> {
232        let plot = StemPlot::new(label, x_data, y_data).with_y_ref(y_ref);
233        plot.validate()?;
234        plot.plot(self);
235        Ok(())
236    }
237
238    /// Plot a simple stem plot with Y values only (X will be indices)
239    pub fn simple_stem_plot(&self, label: &str, values: &[f64]) -> Result<(), PlotError> {
240        if values.is_empty() {
241            return Err(PlotError::EmptyData);
242        }
243        let plot = SimpleStemPlot::new(label, values);
244        plot.plot(self);
245        Ok(())
246    }
247
248    /// Plot a simple stem plot with custom reference Y value
249    pub fn simple_stem_plot_with_ref(
250        &self,
251        label: &str,
252        values: &[f64],
253        y_ref: f64,
254    ) -> Result<(), PlotError> {
255        if values.is_empty() {
256            return Err(PlotError::EmptyData);
257        }
258        let plot = SimpleStemPlot::new(label, values).with_y_ref(y_ref);
259        plot.plot(self);
260        Ok(())
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267
268    #[test]
269    fn test_stem_plot_creation() {
270        let x_data = [1.0, 2.0, 3.0, 4.0];
271        let y_data = [1.0, 4.0, 2.0, 3.0];
272
273        let plot = StemPlot::new("test", &x_data, &y_data);
274        assert_eq!(plot.label(), "test");
275        assert!(plot.validate().is_ok());
276    }
277
278    #[test]
279    fn test_stem_plot_validation() {
280        let x_data = [1.0, 2.0, 3.0];
281        let y_data = [1.0, 4.0]; // Different length
282
283        let plot = StemPlot::new("test", &x_data, &y_data);
284        assert!(plot.validate().is_err());
285    }
286
287    #[test]
288    fn test_simple_stem_plot() {
289        let values = [1.0, 2.0, 3.0, 4.0];
290        let plot = SimpleStemPlot::new("test", &values)
291            .with_flags(StemsFlags::HORIZONTAL)
292            .with_item_flags(ItemFlags::NO_LEGEND);
293        assert_eq!(plot.label(), "test");
294        assert_eq!(plot.flags.bits(), StemsFlags::HORIZONTAL.bits());
295        assert_eq!(plot.item_flags, ItemFlags::NO_LEGEND);
296    }
297
298    #[test]
299    fn test_stem_plot_with_ref() {
300        let x_data = [1.0, 2.0, 3.0, 4.0];
301        let y_data = [1.0, 4.0, 2.0, 3.0];
302
303        let plot = StemPlot::new("test", &x_data, &y_data).with_y_ref(1.0);
304        assert_eq!(plot.y_ref, 1.0);
305        assert!(plot.validate().is_ok());
306    }
307}