dear_implot/plots/
heatmap.rs

1//! Heatmap plot implementation
2
3use super::{
4    Plot, PlotDataLayout, PlotError, PlotItemStyle, plot_spec_with_style, with_plot_str_or_empty,
5};
6use crate::{HeatmapFlags, ItemFlags, sys};
7use dear_imgui_rs::with_scratch_txt_two;
8
9fn validate_grid_counts(
10    caller: &str,
11    rows: usize,
12    cols: usize,
13    values_len: usize,
14) -> Result<(), PlotError> {
15    if rows == 0 || cols == 0 {
16        return Err(PlotError::InvalidData(
17            "Rows and columns must be positive".to_string(),
18        ));
19    }
20
21    let expected_size = rows
22        .checked_mul(cols)
23        .ok_or_else(|| PlotError::InvalidData(format!("{caller} rows * cols overflowed usize")))?;
24    let _ = heatmap_count_to_i32(caller, "rows", rows)?;
25    let _ = heatmap_count_to_i32(caller, "cols", cols)?;
26
27    if values_len != expected_size {
28        return Err(PlotError::DataLengthMismatch {
29            x_len: expected_size,
30            y_len: values_len,
31        });
32    }
33
34    Ok(())
35}
36
37fn heatmap_count_to_i32(caller: &str, name: &str, value: usize) -> Result<i32, PlotError> {
38    i32::try_from(value)
39        .map_err(|_| PlotError::InvalidData(format!("{caller} {name} exceeded ImPlot's i32 range")))
40}
41
42/// Builder for heatmap plots with extensive customization options
43pub struct HeatmapPlot<'a> {
44    label: &'a str,
45    values: &'a [f64],
46    style: PlotItemStyle,
47    rows: usize,
48    cols: usize,
49    scale_min: f64,
50    scale_max: f64,
51    label_fmt: Option<&'a str>,
52    bounds_min: sys::ImPlotPoint,
53    bounds_max: sys::ImPlotPoint,
54    flags: HeatmapFlags,
55    item_flags: ItemFlags,
56}
57
58impl<'a> super::PlotItemStyled for HeatmapPlot<'a> {
59    fn style_mut(&mut self) -> &mut PlotItemStyle {
60        &mut self.style
61    }
62}
63
64impl<'a> HeatmapPlot<'a> {
65    /// Create a new heatmap plot with the given label and data
66    ///
67    /// # Arguments
68    /// * `label` - The label for the heatmap
69    /// * `values` - The data values in row-major order (unless ColMajor flag is set)
70    /// * `rows` - Number of rows in the data
71    /// * `cols` - Number of columns in the data
72    pub fn new(label: &'a str, values: &'a [f64], rows: usize, cols: usize) -> Self {
73        Self {
74            label,
75            values,
76            style: PlotItemStyle::default(),
77            rows,
78            cols,
79            scale_min: 0.0,
80            scale_max: 0.0, // Auto-scale when both are 0
81            label_fmt: Some("%.1f"),
82            bounds_min: sys::ImPlotPoint { x: 0.0, y: 0.0 },
83            bounds_max: sys::ImPlotPoint { x: 1.0, y: 1.0 },
84            flags: HeatmapFlags::NONE,
85            item_flags: ItemFlags::NONE,
86        }
87    }
88
89    /// Set the color scale range (min, max)
90    /// If both are 0.0, auto-scaling will be used
91    pub fn with_scale(mut self, min: f64, max: f64) -> Self {
92        self.scale_min = min;
93        self.scale_max = max;
94        self
95    }
96
97    /// Set the label format for values (e.g., "%.2f", "%.1e")
98    /// Set to None to disable labels
99    pub fn with_label_format(mut self, fmt: Option<&'a str>) -> Self {
100        self.label_fmt = fmt;
101        self
102    }
103
104    /// Set the drawing area bounds in plot coordinates
105    pub fn with_bounds(mut self, min_x: f64, min_y: f64, max_x: f64, max_y: f64) -> Self {
106        self.bounds_min = sys::ImPlotPoint { x: min_x, y: min_y };
107        self.bounds_max = sys::ImPlotPoint { x: max_x, y: max_y };
108        self
109    }
110
111    /// Set the drawing area bounds using ImPlotPoint
112    pub fn with_bounds_points(mut self, min: sys::ImPlotPoint, max: sys::ImPlotPoint) -> Self {
113        self.bounds_min = min;
114        self.bounds_max = max;
115        self
116    }
117
118    /// Set heatmap flags for customization
119    pub fn with_flags(mut self, flags: HeatmapFlags) -> Self {
120        self.flags = flags;
121        self
122    }
123
124    /// Set common item flags for this plot item (applies to all plot types)
125    pub fn with_item_flags(mut self, flags: ItemFlags) -> Self {
126        self.item_flags = flags;
127        self
128    }
129
130    /// Use column-major data ordering instead of row-major
131    pub fn column_major(mut self) -> Self {
132        self.flags |= HeatmapFlags::COL_MAJOR;
133        self
134    }
135
136    /// Validate the plot data
137    pub fn validate(&self) -> Result<(), PlotError> {
138        if self.values.is_empty() {
139            return Err(PlotError::EmptyData);
140        }
141
142        validate_grid_counts(
143            "HeatmapPlot::validate()",
144            self.rows,
145            self.cols,
146            self.values.len(),
147        )
148    }
149}
150
151impl<'a> Plot for HeatmapPlot<'a> {
152    fn plot(&self, plot_ui: &crate::PlotUi<'_>) {
153        if self.validate().is_err() {
154            return; // Skip plotting if data is invalid
155        }
156        let Ok(rows) = heatmap_count_to_i32("HeatmapPlot::plot()", "rows", self.rows) else {
157            return;
158        };
159        let Ok(cols) = heatmap_count_to_i32("HeatmapPlot::plot()", "cols", self.cols) else {
160            return;
161        };
162        let label_fmt = self.label_fmt.filter(|s| !s.contains('\0'));
163        let _guard = plot_ui.bind();
164        match label_fmt {
165            Some(label_fmt) => {
166                let label = if self.label.contains('\0') {
167                    ""
168                } else {
169                    self.label
170                };
171                with_scratch_txt_two(label, label_fmt, |label_ptr, label_fmt_ptr| unsafe {
172                    let spec = plot_spec_with_style(
173                        self.style,
174                        self.flags.bits() | self.item_flags.bits(),
175                        PlotDataLayout::DEFAULT,
176                    );
177                    sys::ImPlot_PlotHeatmap_doublePtr(
178                        label_ptr,
179                        self.values.as_ptr(),
180                        rows,
181                        cols,
182                        self.scale_min,
183                        self.scale_max,
184                        label_fmt_ptr,
185                        self.bounds_min,
186                        self.bounds_max,
187                        spec,
188                    );
189                })
190            }
191            None => with_plot_str_or_empty(self.label, |label_ptr| unsafe {
192                let spec = plot_spec_with_style(
193                    self.style,
194                    self.flags.bits() | self.item_flags.bits(),
195                    PlotDataLayout::DEFAULT,
196                );
197                sys::ImPlot_PlotHeatmap_doublePtr(
198                    label_ptr,
199                    self.values.as_ptr(),
200                    rows,
201                    cols,
202                    self.scale_min,
203                    self.scale_max,
204                    std::ptr::null(),
205                    self.bounds_min,
206                    self.bounds_max,
207                    spec,
208                );
209            }),
210        }
211    }
212
213    fn label(&self) -> &str {
214        self.label
215    }
216}
217
218/// Float version of heatmap for better performance with f32 data
219pub struct HeatmapPlotF32<'a> {
220    label: &'a str,
221    values: &'a [f32],
222    style: PlotItemStyle,
223    rows: usize,
224    cols: usize,
225    scale_min: f64,
226    scale_max: f64,
227    label_fmt: Option<&'a str>,
228    bounds_min: sys::ImPlotPoint,
229    bounds_max: sys::ImPlotPoint,
230    flags: HeatmapFlags,
231    item_flags: ItemFlags,
232}
233
234impl<'a> super::PlotItemStyled for HeatmapPlotF32<'a> {
235    fn style_mut(&mut self) -> &mut PlotItemStyle {
236        &mut self.style
237    }
238}
239
240impl<'a> HeatmapPlotF32<'a> {
241    /// Create a new f32 heatmap plot
242    pub fn new(label: &'a str, values: &'a [f32], rows: usize, cols: usize) -> Self {
243        Self {
244            label,
245            values,
246            style: PlotItemStyle::default(),
247            rows,
248            cols,
249            scale_min: 0.0,
250            scale_max: 0.0,
251            label_fmt: Some("%.1f"),
252            bounds_min: sys::ImPlotPoint { x: 0.0, y: 0.0 },
253            bounds_max: sys::ImPlotPoint { x: 1.0, y: 1.0 },
254            flags: HeatmapFlags::NONE,
255            item_flags: ItemFlags::NONE,
256        }
257    }
258
259    /// Set the color scale range (min, max)
260    pub fn with_scale(mut self, min: f64, max: f64) -> Self {
261        self.scale_min = min;
262        self.scale_max = max;
263        self
264    }
265
266    /// Set the label format for values
267    pub fn with_label_format(mut self, fmt: Option<&'a str>) -> Self {
268        self.label_fmt = fmt;
269        self
270    }
271
272    /// Set the drawing area bounds in plot coordinates
273    pub fn with_bounds(mut self, min_x: f64, min_y: f64, max_x: f64, max_y: f64) -> Self {
274        self.bounds_min = sys::ImPlotPoint { x: min_x, y: min_y };
275        self.bounds_max = sys::ImPlotPoint { x: max_x, y: max_y };
276        self
277    }
278
279    /// Set heatmap flags for customization
280    pub fn with_flags(mut self, flags: HeatmapFlags) -> Self {
281        self.flags = flags;
282        self
283    }
284
285    /// Set common item flags for this plot item (applies to all plot types)
286    pub fn with_item_flags(mut self, flags: ItemFlags) -> Self {
287        self.item_flags = flags;
288        self
289    }
290
291    /// Use column-major data ordering
292    pub fn column_major(mut self) -> Self {
293        self.flags |= HeatmapFlags::COL_MAJOR;
294        self
295    }
296
297    /// Validate the plot data
298    pub fn validate(&self) -> Result<(), PlotError> {
299        if self.values.is_empty() {
300            return Err(PlotError::EmptyData);
301        }
302
303        validate_grid_counts(
304            "HeatmapPlotF32::validate()",
305            self.rows,
306            self.cols,
307            self.values.len(),
308        )
309    }
310}
311
312impl<'a> Plot for HeatmapPlotF32<'a> {
313    fn plot(&self, plot_ui: &crate::PlotUi<'_>) {
314        if self.validate().is_err() {
315            return;
316        }
317        let Ok(rows) = heatmap_count_to_i32("HeatmapPlotF32::plot()", "rows", self.rows) else {
318            return;
319        };
320        let Ok(cols) = heatmap_count_to_i32("HeatmapPlotF32::plot()", "cols", self.cols) else {
321            return;
322        };
323        let label_fmt = self.label_fmt.filter(|s| !s.contains('\0'));
324        let _guard = plot_ui.bind();
325        match label_fmt {
326            Some(label_fmt) => {
327                let label = if self.label.contains('\0') {
328                    ""
329                } else {
330                    self.label
331                };
332                with_scratch_txt_two(label, label_fmt, |label_ptr, label_fmt_ptr| unsafe {
333                    let spec = plot_spec_with_style(
334                        self.style,
335                        self.flags.bits() | self.item_flags.bits(),
336                        PlotDataLayout::DEFAULT,
337                    );
338                    sys::ImPlot_PlotHeatmap_FloatPtr(
339                        label_ptr,
340                        self.values.as_ptr(),
341                        rows,
342                        cols,
343                        self.scale_min,
344                        self.scale_max,
345                        label_fmt_ptr,
346                        self.bounds_min,
347                        self.bounds_max,
348                        spec,
349                    );
350                })
351            }
352            None => with_plot_str_or_empty(self.label, |label_ptr| unsafe {
353                let spec = plot_spec_with_style(
354                    self.style,
355                    self.flags.bits() | self.item_flags.bits(),
356                    PlotDataLayout::DEFAULT,
357                );
358                sys::ImPlot_PlotHeatmap_FloatPtr(
359                    label_ptr,
360                    self.values.as_ptr(),
361                    rows,
362                    cols,
363                    self.scale_min,
364                    self.scale_max,
365                    std::ptr::null(),
366                    self.bounds_min,
367                    self.bounds_max,
368                    spec,
369                );
370            }),
371        }
372    }
373
374    fn label(&self) -> &str {
375        self.label
376    }
377}
378
379/// Convenience functions for quick heatmap plotting
380impl<'ui> crate::PlotUi<'ui> {
381    /// Plot a heatmap with f64 data
382    pub fn heatmap_plot(
383        &self,
384        label: &str,
385        values: &[f64],
386        rows: usize,
387        cols: usize,
388    ) -> Result<(), PlotError> {
389        let plot = HeatmapPlot::new(label, values, rows, cols);
390        plot.validate()?;
391        plot.plot(self);
392        Ok(())
393    }
394
395    /// Plot a heatmap with f32 data
396    pub fn heatmap_plot_f32(
397        &self,
398        label: &str,
399        values: &[f32],
400        rows: usize,
401        cols: usize,
402    ) -> Result<(), PlotError> {
403        let plot = HeatmapPlotF32::new(label, values, rows, cols);
404        plot.validate()?;
405        plot.plot(self);
406        Ok(())
407    }
408
409    /// Plot a heatmap with custom scale and bounds
410    pub fn heatmap_plot_scaled(
411        &self,
412        label: &str,
413        values: &[f64],
414        rows: usize,
415        cols: usize,
416        scale_min: f64,
417        scale_max: f64,
418        bounds_min: sys::ImPlotPoint,
419        bounds_max: sys::ImPlotPoint,
420    ) -> Result<(), PlotError> {
421        let plot = HeatmapPlot::new(label, values, rows, cols)
422            .with_scale(scale_min, scale_max)
423            .with_bounds_points(bounds_min, bounds_max);
424        plot.validate()?;
425        plot.plot(self);
426        Ok(())
427    }
428}
429
430#[cfg(test)]
431mod tests {
432    use super::{HeatmapPlot, HeatmapPlotF32};
433    use crate::PlotError;
434
435    fn invalid_data_message(err: PlotError) -> String {
436        match err {
437            PlotError::InvalidData(message) => message,
438            other => panic!("expected invalid data error, got {other:?}"),
439        }
440    }
441
442    #[test]
443    fn heatmap_rejects_zero_counts_before_ffi() {
444        let values = [1.0];
445        let err = HeatmapPlot::new("heat", &values, 0, 1)
446            .validate()
447            .expect_err("zero row count must be rejected");
448        assert!(invalid_data_message(err).contains("Rows and columns must be positive"));
449    }
450
451    #[test]
452    fn heatmap_rejects_count_multiplication_overflow_before_ffi() {
453        let values = [1.0];
454        let err = HeatmapPlot::new("heat", &values, usize::MAX, 2)
455            .validate()
456            .expect_err("overflowing grid size must be rejected");
457        assert!(invalid_data_message(err).contains("rows * cols overflowed"));
458    }
459
460    #[test]
461    fn heatmap_rejects_i32_count_overflow_before_ffi() {
462        let values = [1.0];
463        let err = HeatmapPlot::new("heat", &values, i32::MAX as usize + 1, 1)
464            .validate()
465            .expect_err("oversized row count must be rejected");
466        assert!(invalid_data_message(err).contains("rows exceeded ImPlot's i32 range"));
467    }
468
469    #[test]
470    fn heatmap_f32_uses_checked_grid_counts() {
471        let values = [1.0f32];
472        let err = HeatmapPlotF32::new("heat", &values, 1, i32::MAX as usize + 1)
473            .validate()
474            .expect_err("oversized column count must be rejected");
475        assert!(invalid_data_message(err).contains("cols exceeded ImPlot's i32 range"));
476    }
477}
dear_implot/plots/heatmap.rs

dear_implot/plots/
heatmap.rs
