dear_implot/plots/
stems.rs

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