Skip to main content

dear_implot/plots/
stems.rs

1//! Stem plot implementation
2
3use super::{
4    Plot, PlotError, PlotItemStyle, plot_spec_with_style, validate_data_lengths,
5    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    offset: i32,
19    stride: i32,
20}
21
22impl<'a> super::PlotItemStyled for StemPlot<'a> {
23    fn style_mut(&mut self) -> &mut PlotItemStyle {
24        &mut self.style
25    }
26}
27
28impl<'a> StemPlot<'a> {
29    /// Create a new stem plot with the given label and data
30    pub fn new(label: &'a str, x_data: &'a [f64], y_data: &'a [f64]) -> Self {
31        Self {
32            label,
33            x_data,
34            y_data,
35            style: PlotItemStyle::default(),
36            y_ref: 0.0, // Default reference line at Y=0
37            flags: StemsFlags::NONE,
38            item_flags: ItemFlags::NONE,
39            offset: 0,
40            stride: std::mem::size_of::<f64>() as i32,
41        }
42    }
43
44    /// Set the reference Y value for stems
45    /// Stems will be drawn from this Y value to the data points
46    pub fn with_y_ref(mut self, y_ref: f64) -> Self {
47        self.y_ref = y_ref;
48        self
49    }
50
51    /// Set stem flags for customization
52    pub fn with_flags(mut self, flags: StemsFlags) -> Self {
53        self.flags = flags;
54        self
55    }
56
57    /// Set common item flags for this plot item (applies to all plot types)
58    pub fn with_item_flags(mut self, flags: ItemFlags) -> Self {
59        self.item_flags = flags;
60        self
61    }
62
63    /// Set data offset for partial plotting
64    pub fn with_offset(mut self, offset: i32) -> Self {
65        self.offset = offset;
66        self
67    }
68
69    /// Set data stride for non-contiguous data
70    pub fn with_stride(mut self, stride: i32) -> Self {
71        self.stride = stride;
72        self
73    }
74
75    /// Validate the plot data
76    pub fn validate(&self) -> Result<(), PlotError> {
77        validate_data_lengths(self.x_data, self.y_data)
78    }
79}
80
81impl<'a> Plot for StemPlot<'a> {
82    fn plot(&self) {
83        if self.validate().is_err() {
84            return;
85        }
86        let Ok(count) = i32::try_from(self.x_data.len()) else {
87            return;
88        };
89
90        with_plot_str_or_empty(self.label, |label_ptr| unsafe {
91            let spec = plot_spec_with_style(
92                self.style,
93                self.flags.bits() | self.item_flags.bits(),
94                self.offset,
95                self.stride,
96            );
97            sys::ImPlot_PlotStems_doublePtrdoublePtr(
98                label_ptr,
99                self.x_data.as_ptr(),
100                self.y_data.as_ptr(),
101                count,
102                self.y_ref,
103                spec,
104            );
105        })
106    }
107
108    fn label(&self) -> &str {
109        self.label
110    }
111}
112
113/// Simple stem plot for quick plotting without builder pattern
114pub struct SimpleStemPlot<'a> {
115    label: &'a str,
116    values: &'a [f64],
117    style: PlotItemStyle,
118    y_ref: f64,
119    flags: StemsFlags,
120    item_flags: ItemFlags,
121    x_scale: f64,
122    x_start: f64,
123}
124
125impl<'a> super::PlotItemStyled for SimpleStemPlot<'a> {
126    fn style_mut(&mut self) -> &mut PlotItemStyle {
127        &mut self.style
128    }
129}
130
131impl<'a> SimpleStemPlot<'a> {
132    /// Create a simple stem plot with Y values only (X will be indices)
133    pub fn new(label: &'a str, values: &'a [f64]) -> Self {
134        Self {
135            label,
136            values,
137            style: PlotItemStyle::default(),
138            y_ref: 0.0,
139            flags: StemsFlags::NONE,
140            item_flags: ItemFlags::NONE,
141            x_scale: 1.0,
142            x_start: 0.0,
143        }
144    }
145
146    /// Set the reference Y value for stems
147    pub fn with_y_ref(mut self, y_ref: f64) -> Self {
148        self.y_ref = y_ref;
149        self
150    }
151
152    /// Set stem flags for customization
153    pub fn with_flags(mut self, flags: StemsFlags) -> Self {
154        self.flags = flags;
155        self
156    }
157
158    /// Set common item flags for this plot item (applies to all plot types)
159    pub fn with_item_flags(mut self, flags: ItemFlags) -> Self {
160        self.item_flags = flags;
161        self
162    }
163
164    /// Set X scale factor
165    pub fn with_x_scale(mut self, scale: f64) -> Self {
166        self.x_scale = scale;
167        self
168    }
169
170    /// Set X start value
171    pub fn with_x_start(mut self, start: f64) -> Self {
172        self.x_start = start;
173        self
174    }
175}
176
177impl<'a> Plot for SimpleStemPlot<'a> {
178    fn plot(&self) {
179        if self.values.is_empty() {
180            return;
181        }
182        let Ok(count) = i32::try_from(self.values.len()) else {
183            return;
184        };
185
186        with_plot_str_or_empty(self.label, |label_ptr| unsafe {
187            let spec = plot_spec_with_style(
188                self.style,
189                self.flags.bits() | self.item_flags.bits(),
190                0,
191                std::mem::size_of::<f64>() as i32,
192            );
193            sys::ImPlot_PlotStems_doublePtrInt(
194                label_ptr,
195                self.values.as_ptr(),
196                count,
197                self.y_ref,
198                self.x_scale,
199                self.x_start,
200                spec,
201            );
202        })
203    }
204
205    fn label(&self) -> &str {
206        self.label
207    }
208}
209
210/// Convenience functions for quick stem plotting
211impl<'ui> crate::PlotUi<'ui> {
212    /// Plot a stem plot with X and Y data
213    pub fn stem_plot(&self, label: &str, x_data: &[f64], y_data: &[f64]) -> Result<(), PlotError> {
214        let plot = StemPlot::new(label, x_data, y_data);
215        plot.validate()?;
216        self.bind();
217        plot.plot();
218        Ok(())
219    }
220
221    /// Plot a stem plot with custom reference Y value
222    pub fn stem_plot_with_ref(
223        &self,
224        label: &str,
225        x_data: &[f64],
226        y_data: &[f64],
227        y_ref: f64,
228    ) -> Result<(), PlotError> {
229        let plot = StemPlot::new(label, x_data, y_data).with_y_ref(y_ref);
230        plot.validate()?;
231        self.bind();
232        plot.plot();
233        Ok(())
234    }
235
236    /// Plot a simple stem plot with Y values only (X will be indices)
237    pub fn simple_stem_plot(&self, label: &str, values: &[f64]) -> Result<(), PlotError> {
238        if values.is_empty() {
239            return Err(PlotError::EmptyData);
240        }
241        let plot = SimpleStemPlot::new(label, values);
242        self.bind();
243        plot.plot();
244        Ok(())
245    }
246
247    /// Plot a simple stem plot with custom reference Y value
248    pub fn simple_stem_plot_with_ref(
249        &self,
250        label: &str,
251        values: &[f64],
252        y_ref: f64,
253    ) -> Result<(), PlotError> {
254        if values.is_empty() {
255            return Err(PlotError::EmptyData);
256        }
257        let plot = SimpleStemPlot::new(label, values).with_y_ref(y_ref);
258        self.bind();
259        plot.plot();
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}