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        plot.plot();
217        Ok(())
218    }
219
220    /// Plot a stem plot with custom reference Y value
221    pub fn stem_plot_with_ref(
222        &self,
223        label: &str,
224        x_data: &[f64],
225        y_data: &[f64],
226        y_ref: f64,
227    ) -> Result<(), PlotError> {
228        let plot = StemPlot::new(label, x_data, y_data).with_y_ref(y_ref);
229        plot.validate()?;
230        plot.plot();
231        Ok(())
232    }
233
234    /// Plot a simple stem plot with Y values only (X will be indices)
235    pub fn simple_stem_plot(&self, label: &str, values: &[f64]) -> Result<(), PlotError> {
236        if values.is_empty() {
237            return Err(PlotError::EmptyData);
238        }
239        let plot = SimpleStemPlot::new(label, values);
240        plot.plot();
241        Ok(())
242    }
243
244    /// Plot a simple stem plot with custom reference Y value
245    pub fn simple_stem_plot_with_ref(
246        &self,
247        label: &str,
248        values: &[f64],
249        y_ref: f64,
250    ) -> Result<(), PlotError> {
251        if values.is_empty() {
252            return Err(PlotError::EmptyData);
253        }
254        let plot = SimpleStemPlot::new(label, values).with_y_ref(y_ref);
255        plot.plot();
256        Ok(())
257    }
258}
259
260#[cfg(test)]
261mod tests {
262    use super::*;
263
264    #[test]
265    fn test_stem_plot_creation() {
266        let x_data = [1.0, 2.0, 3.0, 4.0];
267        let y_data = [1.0, 4.0, 2.0, 3.0];
268
269        let plot = StemPlot::new("test", &x_data, &y_data);
270        assert_eq!(plot.label(), "test");
271        assert!(plot.validate().is_ok());
272    }
273
274    #[test]
275    fn test_stem_plot_validation() {
276        let x_data = [1.0, 2.0, 3.0];
277        let y_data = [1.0, 4.0]; // Different length
278
279        let plot = StemPlot::new("test", &x_data, &y_data);
280        assert!(plot.validate().is_err());
281    }
282
283    #[test]
284    fn test_simple_stem_plot() {
285        let values = [1.0, 2.0, 3.0, 4.0];
286        let plot = SimpleStemPlot::new("test", &values)
287            .with_flags(StemsFlags::HORIZONTAL)
288            .with_item_flags(ItemFlags::NO_LEGEND);
289        assert_eq!(plot.label(), "test");
290        assert_eq!(plot.flags.bits(), StemsFlags::HORIZONTAL.bits());
291        assert_eq!(plot.item_flags, ItemFlags::NO_LEGEND);
292    }
293
294    #[test]
295    fn test_stem_plot_with_ref() {
296        let x_data = [1.0, 2.0, 3.0, 4.0];
297        let y_data = [1.0, 4.0, 2.0, 3.0];
298
299        let plot = StemPlot::new("test", &x_data, &y_data).with_y_ref(1.0);
300        assert_eq!(plot.y_ref, 1.0);
301        assert!(plot.validate().is_ok());
302    }
303}