dear_implot/plots/
heatmap.rs

1//! Heatmap plot implementation
2
3use super::{Plot, PlotError, with_plot_str_or_empty};
4use crate::{HeatmapFlags, sys};
5use dear_imgui_rs::with_scratch_txt_two;
6
7/// Builder for heatmap plots with extensive customization options
8pub struct HeatmapPlot<'a> {
9    label: &'a str,
10    values: &'a [f64],
11    rows: i32,
12    cols: i32,
13    scale_min: f64,
14    scale_max: f64,
15    label_fmt: Option<&'a str>,
16    bounds_min: sys::ImPlotPoint,
17    bounds_max: sys::ImPlotPoint,
18    flags: HeatmapFlags,
19}
20
21impl<'a> HeatmapPlot<'a> {
22    /// Create a new heatmap plot with the given label and data
23    ///
24    /// # Arguments
25    /// * `label` - The label for the heatmap
26    /// * `values` - The data values in row-major order (unless ColMajor flag is set)
27    /// * `rows` - Number of rows in the data
28    /// * `cols` - Number of columns in the data
29    pub fn new(label: &'a str, values: &'a [f64], rows: usize, cols: usize) -> Self {
30        Self {
31            label,
32            values,
33            rows: rows as i32,
34            cols: cols as i32,
35            scale_min: 0.0,
36            scale_max: 0.0, // Auto-scale when both are 0
37            label_fmt: Some("%.1f"),
38            bounds_min: sys::ImPlotPoint { x: 0.0, y: 0.0 },
39            bounds_max: sys::ImPlotPoint { x: 1.0, y: 1.0 },
40            flags: HeatmapFlags::NONE,
41        }
42    }
43
44    /// Set the color scale range (min, max)
45    /// If both are 0.0, auto-scaling will be used
46    pub fn with_scale(mut self, min: f64, max: f64) -> Self {
47        self.scale_min = min;
48        self.scale_max = max;
49        self
50    }
51
52    /// Set the label format for values (e.g., "%.2f", "%.1e")
53    /// Set to None to disable labels
54    pub fn with_label_format(mut self, fmt: Option<&'a str>) -> Self {
55        self.label_fmt = fmt;
56        self
57    }
58
59    /// Set the drawing area bounds in plot coordinates
60    pub fn with_bounds(mut self, min_x: f64, min_y: f64, max_x: f64, max_y: f64) -> Self {
61        self.bounds_min = sys::ImPlotPoint { x: min_x, y: min_y };
62        self.bounds_max = sys::ImPlotPoint { x: max_x, y: max_y };
63        self
64    }
65
66    /// Set the drawing area bounds using ImPlotPoint
67    pub fn with_bounds_points(mut self, min: sys::ImPlotPoint, max: sys::ImPlotPoint) -> Self {
68        self.bounds_min = min;
69        self.bounds_max = max;
70        self
71    }
72
73    /// Set heatmap flags for customization
74    pub fn with_flags(mut self, flags: HeatmapFlags) -> Self {
75        self.flags = flags;
76        self
77    }
78
79    /// Use column-major data ordering instead of row-major
80    pub fn column_major(mut self) -> Self {
81        self.flags |= HeatmapFlags::COL_MAJOR;
82        self
83    }
84
85    /// Validate the plot data
86    pub fn validate(&self) -> Result<(), PlotError> {
87        if self.values.is_empty() {
88            return Err(PlotError::EmptyData);
89        }
90
91        let expected_size = (self.rows * self.cols) as usize;
92        if self.values.len() != expected_size {
93            return Err(PlotError::DataLengthMismatch {
94                x_len: expected_size,
95                y_len: self.values.len(),
96            });
97        }
98
99        if self.rows <= 0 || self.cols <= 0 {
100            return Err(PlotError::InvalidData(
101                "Rows and columns must be positive".to_string(),
102            ));
103        }
104
105        Ok(())
106    }
107}
108
109impl<'a> Plot for HeatmapPlot<'a> {
110    fn plot(&self) {
111        if self.validate().is_err() {
112            return; // Skip plotting if data is invalid
113        }
114        let label_fmt = self.label_fmt.filter(|s| !s.contains('\0'));
115        match label_fmt {
116            Some(label_fmt) => {
117                let label = if self.label.contains('\0') {
118                    ""
119                } else {
120                    self.label
121                };
122                with_scratch_txt_two(label, label_fmt, |label_ptr, label_fmt_ptr| unsafe {
123                    sys::ImPlot_PlotHeatmap_doublePtr(
124                        label_ptr,
125                        self.values.as_ptr(),
126                        self.rows,
127                        self.cols,
128                        self.scale_min,
129                        self.scale_max,
130                        label_fmt_ptr,
131                        self.bounds_min,
132                        self.bounds_max,
133                        self.flags.bits() as i32,
134                    );
135                })
136            }
137            None => with_plot_str_or_empty(self.label, |label_ptr| unsafe {
138                sys::ImPlot_PlotHeatmap_doublePtr(
139                    label_ptr,
140                    self.values.as_ptr(),
141                    self.rows,
142                    self.cols,
143                    self.scale_min,
144                    self.scale_max,
145                    std::ptr::null(),
146                    self.bounds_min,
147                    self.bounds_max,
148                    self.flags.bits() as i32,
149                );
150            }),
151        }
152    }
153
154    fn label(&self) -> &str {
155        self.label
156    }
157}
158
159/// Float version of heatmap for better performance with f32 data
160pub struct HeatmapPlotF32<'a> {
161    label: &'a str,
162    values: &'a [f32],
163    rows: i32,
164    cols: i32,
165    scale_min: f64,
166    scale_max: f64,
167    label_fmt: Option<&'a str>,
168    bounds_min: sys::ImPlotPoint,
169    bounds_max: sys::ImPlotPoint,
170    flags: HeatmapFlags,
171}
172
173impl<'a> HeatmapPlotF32<'a> {
174    /// Create a new f32 heatmap plot
175    pub fn new(label: &'a str, values: &'a [f32], rows: usize, cols: usize) -> Self {
176        Self {
177            label,
178            values,
179            rows: rows as i32,
180            cols: cols as i32,
181            scale_min: 0.0,
182            scale_max: 0.0,
183            label_fmt: Some("%.1f"),
184            bounds_min: sys::ImPlotPoint { x: 0.0, y: 0.0 },
185            bounds_max: sys::ImPlotPoint { x: 1.0, y: 1.0 },
186            flags: HeatmapFlags::NONE,
187        }
188    }
189
190    /// Set the color scale range (min, max)
191    pub fn with_scale(mut self, min: f64, max: f64) -> Self {
192        self.scale_min = min;
193        self.scale_max = max;
194        self
195    }
196
197    /// Set the label format for values
198    pub fn with_label_format(mut self, fmt: Option<&'a str>) -> Self {
199        self.label_fmt = fmt;
200        self
201    }
202
203    /// Set the drawing area bounds in plot coordinates
204    pub fn with_bounds(mut self, min_x: f64, min_y: f64, max_x: f64, max_y: f64) -> Self {
205        self.bounds_min = sys::ImPlotPoint { x: min_x, y: min_y };
206        self.bounds_max = sys::ImPlotPoint { x: max_x, y: max_y };
207        self
208    }
209
210    /// Set heatmap flags for customization
211    pub fn with_flags(mut self, flags: HeatmapFlags) -> Self {
212        self.flags = flags;
213        self
214    }
215
216    /// Use column-major data ordering
217    pub fn column_major(mut self) -> Self {
218        self.flags |= HeatmapFlags::COL_MAJOR;
219        self
220    }
221
222    /// Validate the plot data
223    pub fn validate(&self) -> Result<(), PlotError> {
224        if self.values.is_empty() {
225            return Err(PlotError::EmptyData);
226        }
227
228        let expected_size = (self.rows * self.cols) as usize;
229        if self.values.len() != expected_size {
230            return Err(PlotError::DataLengthMismatch {
231                x_len: expected_size,
232                y_len: self.values.len(),
233            });
234        }
235
236        if self.rows <= 0 || self.cols <= 0 {
237            return Err(PlotError::InvalidData(
238                "Rows and columns must be positive".to_string(),
239            ));
240        }
241
242        Ok(())
243    }
244}
245
246impl<'a> Plot for HeatmapPlotF32<'a> {
247    fn plot(&self) {
248        if self.validate().is_err() {
249            return;
250        }
251        let label_fmt = self.label_fmt.filter(|s| !s.contains('\0'));
252        match label_fmt {
253            Some(label_fmt) => {
254                let label = if self.label.contains('\0') {
255                    ""
256                } else {
257                    self.label
258                };
259                with_scratch_txt_two(label, label_fmt, |label_ptr, label_fmt_ptr| unsafe {
260                    sys::ImPlot_PlotHeatmap_FloatPtr(
261                        label_ptr,
262                        self.values.as_ptr(),
263                        self.rows,
264                        self.cols,
265                        self.scale_min,
266                        self.scale_max,
267                        label_fmt_ptr,
268                        self.bounds_min,
269                        self.bounds_max,
270                        self.flags.bits() as i32,
271                    );
272                })
273            }
274            None => with_plot_str_or_empty(self.label, |label_ptr| unsafe {
275                sys::ImPlot_PlotHeatmap_FloatPtr(
276                    label_ptr,
277                    self.values.as_ptr(),
278                    self.rows,
279                    self.cols,
280                    self.scale_min,
281                    self.scale_max,
282                    std::ptr::null(),
283                    self.bounds_min,
284                    self.bounds_max,
285                    self.flags.bits() as i32,
286                );
287            }),
288        }
289    }
290
291    fn label(&self) -> &str {
292        self.label
293    }
294}
295
296/// Convenience functions for quick heatmap plotting
297impl<'ui> crate::PlotUi<'ui> {
298    /// Plot a heatmap with f64 data
299    pub fn heatmap_plot(
300        &self,
301        label: &str,
302        values: &[f64],
303        rows: usize,
304        cols: usize,
305    ) -> Result<(), PlotError> {
306        let plot = HeatmapPlot::new(label, values, rows, cols);
307        plot.validate()?;
308        plot.plot();
309        Ok(())
310    }
311
312    /// Plot a heatmap with f32 data
313    pub fn heatmap_plot_f32(
314        &self,
315        label: &str,
316        values: &[f32],
317        rows: usize,
318        cols: usize,
319    ) -> Result<(), PlotError> {
320        let plot = HeatmapPlotF32::new(label, values, rows, cols);
321        plot.validate()?;
322        plot.plot();
323        Ok(())
324    }
325
326    /// Plot a heatmap with custom scale and bounds
327    pub fn heatmap_plot_scaled(
328        &self,
329        label: &str,
330        values: &[f64],
331        rows: usize,
332        cols: usize,
333        scale_min: f64,
334        scale_max: f64,
335        bounds_min: sys::ImPlotPoint,
336        bounds_max: sys::ImPlotPoint,
337    ) -> Result<(), PlotError> {
338        let plot = HeatmapPlot::new(label, values, rows, cols)
339            .with_scale(scale_min, scale_max)
340            .with_bounds_points(bounds_min, bounds_max);
341        plot.validate()?;
342        plot.plot();
343        Ok(())
344    }
345}