Skip to main content

plotkit_core/charts/
heatmap.rs

1//! Heatmap chart builder methods.
2//!
3//! Provides a fluent builder API for configuring [`HeatmapArtist`] instances.
4//! Each method returns `&mut Self`, allowing calls to be chained together
5//! for concise, readable chart construction.
6
7use crate::artist::HeatmapArtist;
8use crate::colormap::Colormap;
9
10impl HeatmapArtist {
11    /// Sets the colormap used to map cell values to colors.
12    pub fn colormap(&mut self, cmap: Colormap) -> &mut Self {
13        self.cmap = cmap;
14        self
15    }
16
17    /// Sets the minimum value for colormap normalisation.
18    pub fn vmin(&mut self, min: f64) -> &mut Self {
19        self.vmin = Some(min);
20        self
21    }
22
23    /// Sets the maximum value for colormap normalisation.
24    pub fn vmax(&mut self, max: f64) -> &mut Self {
25        self.vmax = Some(max);
26        self
27    }
28
29    /// Enables or disables drawing cell values as text in each cell.
30    pub fn show_values(&mut self, show: bool) -> &mut Self {
31        self.show_values = show;
32        self
33    }
34
35    /// Enables or disables auto-attaching a colorbar when this heatmap is drawn.
36    ///
37    /// When `true`, the parent axes will automatically add a colorbar showing
38    /// the color-to-value mapping used by this heatmap.
39    pub fn colorbar(&mut self, show: bool) -> &mut Self {
40        self.show_colorbar = show;
41        self
42    }
43
44    /// Sets the legend label.
45    pub fn label(&mut self, label: &str) -> &mut Self {
46        self.label = Some(label.to_string());
47        self
48    }
49
50    /// Sets the x-axis labels for the heatmap columns.
51    pub fn x_labels(&mut self, labels: Vec<String>) -> &mut Self {
52        self.x_labels = Some(labels);
53        self
54    }
55
56    /// Sets the y-axis labels for the heatmap rows.
57    pub fn y_labels(&mut self, labels: Vec<String>) -> &mut Self {
58        self.y_labels = Some(labels);
59        self
60    }
61
62    /// Returns the effective minimum value for colormap normalisation.
63    pub fn effective_vmin(&self) -> f64 {
64        if let Some(v) = self.vmin {
65            return v;
66        }
67        let mut lo = f64::INFINITY;
68        for row in &self.data {
69            for &val in row {
70                if val.is_finite() && val < lo {
71                    lo = val;
72                }
73            }
74        }
75        if lo.is_finite() {
76            lo
77        } else {
78            0.0
79        }
80    }
81
82    /// Returns the effective maximum value for colormap normalisation.
83    pub fn effective_vmax(&self) -> f64 {
84        if let Some(v) = self.vmax {
85            return v;
86        }
87        let mut hi = f64::NEG_INFINITY;
88        for row in &self.data {
89            for &val in row {
90                if val.is_finite() && val > hi {
91                    hi = val;
92                }
93            }
94        }
95        if hi.is_finite() {
96            hi
97        } else {
98            1.0
99        }
100    }
101}
102
103#[cfg(test)]
104mod tests {
105    use crate::artist::HeatmapArtist;
106    use crate::colormap::Colormap;
107    use crate::primitives::Color;
108
109    fn sample_heatmap() -> HeatmapArtist {
110        HeatmapArtist {
111            data: vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]],
112            x_labels: None,
113            y_labels: None,
114            cmap: Colormap::Viridis,
115            vmin: None,
116            vmax: None,
117            show_values: false,
118            color: Color::TAB_BLUE,
119            label: None,
120            show_colorbar: false,
121        }
122    }
123
124    #[test]
125    fn builder_colormap() {
126        let mut h = sample_heatmap();
127        h.colormap(Colormap::Plasma);
128        assert_eq!(h.cmap, Colormap::Plasma);
129    }
130
131    #[test]
132    fn builder_vmin() {
133        let mut h = sample_heatmap();
134        h.vmin(-10.0);
135        assert_eq!(h.vmin, Some(-10.0));
136    }
137
138    #[test]
139    fn builder_vmax() {
140        let mut h = sample_heatmap();
141        h.vmax(100.0);
142        assert_eq!(h.vmax, Some(100.0));
143    }
144
145    #[test]
146    fn builder_show_values() {
147        let mut h = sample_heatmap();
148        assert!(!h.show_values);
149        h.show_values(true);
150        assert!(h.show_values);
151    }
152
153    #[test]
154    fn builder_label() {
155        let mut h = sample_heatmap();
156        h.label("my heatmap");
157        assert_eq!(h.label.as_deref(), Some("my heatmap"));
158    }
159
160    #[test]
161    fn builder_x_labels() {
162        let mut h = sample_heatmap();
163        h.x_labels(vec!["A".into(), "B".into(), "C".into()]);
164        assert_eq!(h.x_labels.as_ref().unwrap().len(), 3);
165    }
166
167    #[test]
168    fn builder_y_labels() {
169        let mut h = sample_heatmap();
170        h.y_labels(vec!["row1".into(), "row2".into()]);
171        assert_eq!(h.y_labels.as_ref().unwrap().len(), 2);
172    }
173
174    #[test]
175    fn effective_vmin_auto() {
176        let h = sample_heatmap();
177        assert!((h.effective_vmin() - 1.0).abs() < f64::EPSILON);
178    }
179
180    #[test]
181    fn effective_vmax_auto() {
182        let h = sample_heatmap();
183        assert!((h.effective_vmax() - 6.0).abs() < f64::EPSILON);
184    }
185
186    #[test]
187    fn effective_vmin_explicit() {
188        let mut h = sample_heatmap();
189        h.vmin(-5.0);
190        assert!((h.effective_vmin() - (-5.0)).abs() < f64::EPSILON);
191    }
192
193    #[test]
194    fn effective_vmax_explicit() {
195        let mut h = sample_heatmap();
196        h.vmax(50.0);
197        assert!((h.effective_vmax() - 50.0).abs() < f64::EPSILON);
198    }
199
200    #[test]
201    fn effective_vmin_empty_data() {
202        let h = HeatmapArtist {
203            data: vec![],
204            x_labels: None,
205            y_labels: None,
206            cmap: Colormap::Viridis,
207            vmin: None,
208            vmax: None,
209            show_values: false,
210            color: Color::TAB_BLUE,
211            label: None,
212            show_colorbar: false,
213        };
214        assert!((h.effective_vmin() - 0.0).abs() < f64::EPSILON);
215    }
216
217    #[test]
218    fn effective_vmax_empty_data() {
219        let h = HeatmapArtist {
220            data: vec![],
221            x_labels: None,
222            y_labels: None,
223            cmap: Colormap::Viridis,
224            vmin: None,
225            vmax: None,
226            show_values: false,
227            color: Color::TAB_BLUE,
228            label: None,
229            show_colorbar: false,
230        };
231        assert!((h.effective_vmax() - 1.0).abs() < f64::EPSILON);
232    }
233
234    #[test]
235    fn data_bounds_basic() {
236        let h = sample_heatmap();
237        let (xmin, xmax, ymin, ymax) = h.data_bounds();
238        assert!((xmin - 0.0).abs() < f64::EPSILON);
239        assert!((xmax - 3.0).abs() < f64::EPSILON);
240        assert!((ymin - 0.0).abs() < f64::EPSILON);
241        assert!((ymax - 2.0).abs() < f64::EPSILON);
242    }
243
244    #[test]
245    fn data_bounds_empty() {
246        let h = HeatmapArtist {
247            data: vec![],
248            x_labels: None,
249            y_labels: None,
250            cmap: Colormap::Viridis,
251            vmin: None,
252            vmax: None,
253            show_values: false,
254            color: Color::TAB_BLUE,
255            label: None,
256            show_colorbar: false,
257        };
258        assert_eq!(h.data_bounds(), (0.0, 1.0, 0.0, 1.0));
259    }
260
261    #[test]
262    fn data_bounds_single_cell() {
263        let h = HeatmapArtist {
264            data: vec![vec![42.0]],
265            x_labels: None,
266            y_labels: None,
267            cmap: Colormap::Viridis,
268            vmin: None,
269            vmax: None,
270            show_values: false,
271            color: Color::TAB_BLUE,
272            label: None,
273            show_colorbar: false,
274        };
275        let (xmin, xmax, ymin, ymax) = h.data_bounds();
276        assert!((xmin - 0.0).abs() < f64::EPSILON);
277        assert!((xmax - 1.0).abs() < f64::EPSILON);
278        assert!((ymin - 0.0).abs() < f64::EPSILON);
279        assert!((ymax - 1.0).abs() < f64::EPSILON);
280    }
281
282    #[test]
283    fn builder_chaining() {
284        let mut h = sample_heatmap();
285        h.colormap(Colormap::Plasma)
286            .vmin(0.0)
287            .vmax(10.0)
288            .show_values(true)
289            .label("chained");
290        assert_eq!(h.cmap, Colormap::Plasma);
291        assert_eq!(h.vmin, Some(0.0));
292        assert_eq!(h.vmax, Some(10.0));
293        assert!(h.show_values);
294        assert_eq!(h.label.as_deref(), Some("chained"));
295    }
296
297    #[test]
298    fn effective_bounds_with_nan() {
299        let h = HeatmapArtist {
300            data: vec![vec![f64::NAN, 3.0], vec![1.0, f64::NAN]],
301            x_labels: None,
302            y_labels: None,
303            cmap: Colormap::Viridis,
304            vmin: None,
305            vmax: None,
306            show_values: false,
307            color: Color::TAB_BLUE,
308            label: None,
309            show_colorbar: false,
310        };
311        assert!((h.effective_vmin() - 1.0).abs() < f64::EPSILON);
312        assert!((h.effective_vmax() - 3.0).abs() < f64::EPSILON);
313    }
314}