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) {
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        match label_fmt {
164            Some(label_fmt) => {
165                let label = if self.label.contains('\0') {
166                    ""
167                } else {
168                    self.label
169                };
170                with_scratch_txt_two(label, label_fmt, |label_ptr, label_fmt_ptr| unsafe {
171                    let spec = plot_spec_with_style(
172                        self.style,
173                        self.flags.bits() | self.item_flags.bits(),
174                        PlotDataLayout::DEFAULT,
175                    );
176                    sys::ImPlot_PlotHeatmap_doublePtr(
177                        label_ptr,
178                        self.values.as_ptr(),
179                        rows,
180                        cols,
181                        self.scale_min,
182                        self.scale_max,
183                        label_fmt_ptr,
184                        self.bounds_min,
185                        self.bounds_max,
186                        spec,
187                    );
188                })
189            }
190            None => with_plot_str_or_empty(self.label, |label_ptr| unsafe {
191                let spec = plot_spec_with_style(
192                    self.style,
193                    self.flags.bits() | self.item_flags.bits(),
194                    PlotDataLayout::DEFAULT,
195                );
196                sys::ImPlot_PlotHeatmap_doublePtr(
197                    label_ptr,
198                    self.values.as_ptr(),
199                    rows,
200                    cols,
201                    self.scale_min,
202                    self.scale_max,
203                    std::ptr::null(),
204                    self.bounds_min,
205                    self.bounds_max,
206                    spec,
207                );
208            }),
209        }
210    }
211
212    fn label(&self) -> &str {
213        self.label
214    }
215}
216
217/// Float version of heatmap for better performance with f32 data
218pub struct HeatmapPlotF32<'a> {
219    label: &'a str,
220    values: &'a [f32],
221    style: PlotItemStyle,
222    rows: usize,
223    cols: usize,
224    scale_min: f64,
225    scale_max: f64,
226    label_fmt: Option<&'a str>,
227    bounds_min: sys::ImPlotPoint,
228    bounds_max: sys::ImPlotPoint,
229    flags: HeatmapFlags,
230    item_flags: ItemFlags,
231}
232
233impl<'a> super::PlotItemStyled for HeatmapPlotF32<'a> {
234    fn style_mut(&mut self) -> &mut PlotItemStyle {
235        &mut self.style
236    }
237}
238
239impl<'a> HeatmapPlotF32<'a> {
240    /// Create a new f32 heatmap plot
241    pub fn new(label: &'a str, values: &'a [f32], rows: usize, cols: usize) -> Self {
242        Self {
243            label,
244            values,
245            style: PlotItemStyle::default(),
246            rows,
247            cols,
248            scale_min: 0.0,
249            scale_max: 0.0,
250            label_fmt: Some("%.1f"),
251            bounds_min: sys::ImPlotPoint { x: 0.0, y: 0.0 },
252            bounds_max: sys::ImPlotPoint { x: 1.0, y: 1.0 },
253            flags: HeatmapFlags::NONE,
254            item_flags: ItemFlags::NONE,
255        }
256    }
257
258    /// Set the color scale range (min, max)
259    pub fn with_scale(mut self, min: f64, max: f64) -> Self {
260        self.scale_min = min;
261        self.scale_max = max;
262        self
263    }
264
265    /// Set the label format for values
266    pub fn with_label_format(mut self, fmt: Option<&'a str>) -> Self {
267        self.label_fmt = fmt;
268        self
269    }
270
271    /// Set the drawing area bounds in plot coordinates
272    pub fn with_bounds(mut self, min_x: f64, min_y: f64, max_x: f64, max_y: f64) -> Self {
273        self.bounds_min = sys::ImPlotPoint { x: min_x, y: min_y };
274        self.bounds_max = sys::ImPlotPoint { x: max_x, y: max_y };
275        self
276    }
277
278    /// Set heatmap flags for customization
279    pub fn with_flags(mut self, flags: HeatmapFlags) -> Self {
280        self.flags = flags;
281        self
282    }
283
284    /// Set common item flags for this plot item (applies to all plot types)
285    pub fn with_item_flags(mut self, flags: ItemFlags) -> Self {
286        self.item_flags = flags;
287        self
288    }
289
290    /// Use column-major data ordering
291    pub fn column_major(mut self) -> Self {
292        self.flags |= HeatmapFlags::COL_MAJOR;
293        self
294    }
295
296    /// Validate the plot data
297    pub fn validate(&self) -> Result<(), PlotError> {
298        if self.values.is_empty() {
299            return Err(PlotError::EmptyData);
300        }
301
302        validate_grid_counts(
303            "HeatmapPlotF32::validate()",
304            self.rows,
305            self.cols,
306            self.values.len(),
307        )
308    }
309}
310
311impl<'a> Plot for HeatmapPlotF32<'a> {
312    fn plot(&self) {
313        if self.validate().is_err() {
314            return;
315        }
316        let Ok(rows) = heatmap_count_to_i32("HeatmapPlotF32::plot()", "rows", self.rows) else {
317            return;
318        };
319        let Ok(cols) = heatmap_count_to_i32("HeatmapPlotF32::plot()", "cols", self.cols) else {
320            return;
321        };
322        let label_fmt = self.label_fmt.filter(|s| !s.contains('\0'));
323        match label_fmt {
324            Some(label_fmt) => {
325                let label = if self.label.contains('\0') {
326                    ""
327                } else {
328                    self.label
329                };
330                with_scratch_txt_two(label, label_fmt, |label_ptr, label_fmt_ptr| unsafe {
331                    let spec = plot_spec_with_style(
332                        self.style,
333                        self.flags.bits() | self.item_flags.bits(),
334                        PlotDataLayout::DEFAULT,
335                    );
336                    sys::ImPlot_PlotHeatmap_FloatPtr(
337                        label_ptr,
338                        self.values.as_ptr(),
339                        rows,
340                        cols,
341                        self.scale_min,
342                        self.scale_max,
343                        label_fmt_ptr,
344                        self.bounds_min,
345                        self.bounds_max,
346                        spec,
347                    );
348                })
349            }
350            None => with_plot_str_or_empty(self.label, |label_ptr| unsafe {
351                let spec = plot_spec_with_style(
352                    self.style,
353                    self.flags.bits() | self.item_flags.bits(),
354                    PlotDataLayout::DEFAULT,
355                );
356                sys::ImPlot_PlotHeatmap_FloatPtr(
357                    label_ptr,
358                    self.values.as_ptr(),
359                    rows,
360                    cols,
361                    self.scale_min,
362                    self.scale_max,
363                    std::ptr::null(),
364                    self.bounds_min,
365                    self.bounds_max,
366                    spec,
367                );
368            }),
369        }
370    }
371
372    fn label(&self) -> &str {
373        self.label
374    }
375}
376
377/// Convenience functions for quick heatmap plotting
378impl<'ui> crate::PlotUi<'ui> {
379    /// Plot a heatmap with f64 data
380    pub fn heatmap_plot(
381        &self,
382        label: &str,
383        values: &[f64],
384        rows: usize,
385        cols: usize,
386    ) -> Result<(), PlotError> {
387        let plot = HeatmapPlot::new(label, values, rows, cols);
388        plot.validate()?;
389        self.bind();
390        plot.plot();
391        Ok(())
392    }
393
394    /// Plot a heatmap with f32 data
395    pub fn heatmap_plot_f32(
396        &self,
397        label: &str,
398        values: &[f32],
399        rows: usize,
400        cols: usize,
401    ) -> Result<(), PlotError> {
402        let plot = HeatmapPlotF32::new(label, values, rows, cols);
403        plot.validate()?;
404        self.bind();
405        plot.plot();
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        self.bind();
426        plot.plot();
427        Ok(())
428    }
429}
430
431#[cfg(test)]
432mod tests {
433    use super::{HeatmapPlot, HeatmapPlotF32};
434    use crate::PlotError;
435
436    fn invalid_data_message(err: PlotError) -> String {
437        match err {
438            PlotError::InvalidData(message) => message,
439            other => panic!("expected invalid data error, got {other:?}"),
440        }
441    }
442
443    #[test]
444    fn heatmap_rejects_zero_counts_before_ffi() {
445        let values = [1.0];
446        let err = HeatmapPlot::new("heat", &values, 0, 1)
447            .validate()
448            .expect_err("zero row count must be rejected");
449        assert!(invalid_data_message(err).contains("Rows and columns must be positive"));
450    }
451
452    #[test]
453    fn heatmap_rejects_count_multiplication_overflow_before_ffi() {
454        let values = [1.0];
455        let err = HeatmapPlot::new("heat", &values, usize::MAX, 2)
456            .validate()
457            .expect_err("overflowing grid size must be rejected");
458        assert!(invalid_data_message(err).contains("rows * cols overflowed"));
459    }
460
461    #[test]
462    fn heatmap_rejects_i32_count_overflow_before_ffi() {
463        let values = [1.0];
464        let err = HeatmapPlot::new("heat", &values, i32::MAX as usize + 1, 1)
465            .validate()
466            .expect_err("oversized row count must be rejected");
467        assert!(invalid_data_message(err).contains("rows exceeded ImPlot's i32 range"));
468    }
469
470    #[test]
471    fn heatmap_f32_uses_checked_grid_counts() {
472        let values = [1.0f32];
473        let err = HeatmapPlotF32::new("heat", &values, 1, i32::MAX as usize + 1)
474            .validate()
475            .expect_err("oversized column count must be rejected");
476        assert!(invalid_data_message(err).contains("cols exceeded ImPlot's i32 range"));
477    }
478}
dear_implot/plots/heatmap.rs

dear_implot/plots/
heatmap.rs
