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

dear_implot/plots/
heatmap.rs
