Skip to main content

dear_implot/plots/
stems.rs

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