Skip to main content

dear_implot/plots/
heatmap.rs

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