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) {
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        with_plot_str_or_empty(self.label, |label_ptr| unsafe {
95            let spec = plot_spec_with_style(
96                self.style,
97                self.flags.bits() | self.item_flags.bits(),
98                self.layout,
99            );
100            sys::ImPlot_PlotStems_doublePtrdoublePtr(
101                label_ptr,
102                self.x_data.as_ptr(),
103                self.y_data.as_ptr(),
104                count,
105                self.y_ref,
106                spec,
107            );
108        })
109    }
110
111    fn label(&self) -> &str {
112        self.label
113    }
114}
115
116/// Simple stem plot for quick plotting without builder pattern
117pub struct SimpleStemPlot<'a> {
118    label: &'a str,
119    values: &'a [f64],
120    style: PlotItemStyle,
121    y_ref: f64,
122    flags: StemsFlags,
123    item_flags: ItemFlags,
124    x_scale: f64,
125    x_start: f64,
126}
127
128impl<'a> super::PlotItemStyled for SimpleStemPlot<'a> {
129    fn style_mut(&mut self) -> &mut PlotItemStyle {
130        &mut self.style
131    }
132}
133
134impl<'a> SimpleStemPlot<'a> {
135    /// Create a simple stem plot with Y values only (X will be indices)
136    pub fn new(label: &'a str, values: &'a [f64]) -> Self {
137        Self {
138            label,
139            values,
140            style: PlotItemStyle::default(),
141            y_ref: 0.0,
142            flags: StemsFlags::NONE,
143            item_flags: ItemFlags::NONE,
144            x_scale: 1.0,
145            x_start: 0.0,
146        }
147    }
148
149    /// Set the reference Y value for stems
150    pub fn with_y_ref(mut self, y_ref: f64) -> Self {
151        self.y_ref = y_ref;
152        self
153    }
154
155    /// Set stem flags for customization
156    pub fn with_flags(mut self, flags: StemsFlags) -> Self {
157        self.flags = flags;
158        self
159    }
160
161    /// Set common item flags for this plot item (applies to all plot types)
162    pub fn with_item_flags(mut self, flags: ItemFlags) -> Self {
163        self.item_flags = flags;
164        self
165    }
166
167    /// Set X scale factor
168    pub fn with_x_scale(mut self, scale: f64) -> Self {
169        self.x_scale = scale;
170        self
171    }
172
173    /// Set X start value
174    pub fn with_x_start(mut self, start: f64) -> Self {
175        self.x_start = start;
176        self
177    }
178}
179
180impl<'a> Plot for SimpleStemPlot<'a> {
181    fn plot(&self) {
182        if self.values.is_empty() {
183            return;
184        }
185        let Ok(count) = i32::try_from(self.values.len()) else {
186            return;
187        };
188
189        with_plot_str_or_empty(self.label, |label_ptr| unsafe {
190            let spec = plot_spec_with_style(
191                self.style,
192                self.flags.bits() | self.item_flags.bits(),
193                PlotDataLayout::DEFAULT,
194            );
195            sys::ImPlot_PlotStems_doublePtrInt(
196                label_ptr,
197                self.values.as_ptr(),
198                count,
199                self.y_ref,
200                self.x_scale,
201                self.x_start,
202                spec,
203            );
204        })
205    }
206
207    fn label(&self) -> &str {
208        self.label
209    }
210}
211
212/// Convenience functions for quick stem plotting
213impl<'ui> crate::PlotUi<'ui> {
214    /// Plot a stem plot with X and Y data
215    pub fn stem_plot(&self, label: &str, x_data: &[f64], y_data: &[f64]) -> Result<(), PlotError> {
216        let plot = StemPlot::new(label, x_data, y_data);
217        plot.validate()?;
218        self.bind();
219        plot.plot();
220        Ok(())
221    }
222
223    /// Plot a stem plot with custom reference Y value
224    pub fn stem_plot_with_ref(
225        &self,
226        label: &str,
227        x_data: &[f64],
228        y_data: &[f64],
229        y_ref: f64,
230    ) -> Result<(), PlotError> {
231        let plot = StemPlot::new(label, x_data, y_data).with_y_ref(y_ref);
232        plot.validate()?;
233        self.bind();
234        plot.plot();
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        self.bind();
245        plot.plot();
246        Ok(())
247    }
248
249    /// Plot a simple stem plot with custom reference Y value
250    pub fn simple_stem_plot_with_ref(
251        &self,
252        label: &str,
253        values: &[f64],
254        y_ref: f64,
255    ) -> Result<(), PlotError> {
256        if values.is_empty() {
257            return Err(PlotError::EmptyData);
258        }
259        let plot = SimpleStemPlot::new(label, values).with_y_ref(y_ref);
260        self.bind();
261        plot.plot();
262        Ok(())
263    }
264}
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269
270    #[test]
271    fn test_stem_plot_creation() {
272        let x_data = [1.0, 2.0, 3.0, 4.0];
273        let y_data = [1.0, 4.0, 2.0, 3.0];
274
275        let plot = StemPlot::new("test", &x_data, &y_data);
276        assert_eq!(plot.label(), "test");
277        assert!(plot.validate().is_ok());
278    }
279
280    #[test]
281    fn test_stem_plot_validation() {
282        let x_data = [1.0, 2.0, 3.0];
283        let y_data = [1.0, 4.0]; // Different length
284
285        let plot = StemPlot::new("test", &x_data, &y_data);
286        assert!(plot.validate().is_err());
287    }
288
289    #[test]
290    fn test_simple_stem_plot() {
291        let values = [1.0, 2.0, 3.0, 4.0];
292        let plot = SimpleStemPlot::new("test", &values)
293            .with_flags(StemsFlags::HORIZONTAL)
294            .with_item_flags(ItemFlags::NO_LEGEND);
295        assert_eq!(plot.label(), "test");
296        assert_eq!(plot.flags.bits(), StemsFlags::HORIZONTAL.bits());
297        assert_eq!(plot.item_flags, ItemFlags::NO_LEGEND);
298    }
299
300    #[test]
301    fn test_stem_plot_with_ref() {
302        let x_data = [1.0, 2.0, 3.0, 4.0];
303        let y_data = [1.0, 4.0, 2.0, 3.0];
304
305        let plot = StemPlot::new("test", &x_data, &y_data).with_y_ref(1.0);
306        assert_eq!(plot.y_ref, 1.0);
307        assert!(plot.validate().is_ok());
308    }
309}