Skip to main content

dear_implot/plots/
heatmap.rs

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