dear_implot/plots/
heatmap.rs

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