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() { lo } else { 0.0 }
76    }
77
78    /// Returns the effective maximum value for colormap normalisation.
79    pub fn effective_vmax(&self) -> f64 {
80        if let Some(v) = self.vmax {
81            return v;
82        }
83        let mut hi = f64::NEG_INFINITY;
84        for row in &self.data {
85            for &val in row {
86                if val.is_finite() && val > hi {
87                    hi = val;
88                }
89            }
90        }
91        if hi.is_finite() { hi } else { 1.0 }
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use crate::artist::HeatmapArtist;
98    use crate::colormap::Colormap;
99    use crate::primitives::Color;
100
101    fn sample_heatmap() -> HeatmapArtist {
102        HeatmapArtist {
103            data: vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]],
104            x_labels: None, y_labels: None,
105            cmap: Colormap::Viridis, vmin: None, vmax: None,
106            show_values: false, color: Color::TAB_BLUE, label: None,
107            show_colorbar: false,
108        }
109    }
110
111    #[test]
112    fn builder_colormap() {
113        let mut h = sample_heatmap();
114        h.colormap(Colormap::Plasma);
115        assert_eq!(h.cmap, Colormap::Plasma);
116    }
117
118    #[test]
119    fn builder_vmin() {
120        let mut h = sample_heatmap();
121        h.vmin(-10.0);
122        assert_eq!(h.vmin, Some(-10.0));
123    }
124
125    #[test]
126    fn builder_vmax() {
127        let mut h = sample_heatmap();
128        h.vmax(100.0);
129        assert_eq!(h.vmax, Some(100.0));
130    }
131
132    #[test]
133    fn builder_show_values() {
134        let mut h = sample_heatmap();
135        assert!(!h.show_values);
136        h.show_values(true);
137        assert!(h.show_values);
138    }
139
140    #[test]
141    fn builder_label() {
142        let mut h = sample_heatmap();
143        h.label("my heatmap");
144        assert_eq!(h.label.as_deref(), Some("my heatmap"));
145    }
146
147    #[test]
148    fn builder_x_labels() {
149        let mut h = sample_heatmap();
150        h.x_labels(vec!["A".into(), "B".into(), "C".into()]);
151        assert_eq!(h.x_labels.as_ref().unwrap().len(), 3);
152    }
153
154    #[test]
155    fn builder_y_labels() {
156        let mut h = sample_heatmap();
157        h.y_labels(vec!["row1".into(), "row2".into()]);
158        assert_eq!(h.y_labels.as_ref().unwrap().len(), 2);
159    }
160
161    #[test]
162    fn effective_vmin_auto() {
163        let h = sample_heatmap();
164        assert!((h.effective_vmin() - 1.0).abs() < f64::EPSILON);
165    }
166
167    #[test]
168    fn effective_vmax_auto() {
169        let h = sample_heatmap();
170        assert!((h.effective_vmax() - 6.0).abs() < f64::EPSILON);
171    }
172
173    #[test]
174    fn effective_vmin_explicit() {
175        let mut h = sample_heatmap();
176        h.vmin(-5.0);
177        assert!((h.effective_vmin() - (-5.0)).abs() < f64::EPSILON);
178    }
179
180    #[test]
181    fn effective_vmax_explicit() {
182        let mut h = sample_heatmap();
183        h.vmax(50.0);
184        assert!((h.effective_vmax() - 50.0).abs() < f64::EPSILON);
185    }
186
187    #[test]
188    fn effective_vmin_empty_data() {
189        let h = HeatmapArtist {
190            data: vec![], x_labels: None, y_labels: None,
191            cmap: Colormap::Viridis, vmin: None, vmax: None,
192            show_values: false, color: Color::TAB_BLUE, label: None,
193            show_colorbar: false,
194        };
195        assert!((h.effective_vmin() - 0.0).abs() < f64::EPSILON);
196    }
197
198    #[test]
199    fn effective_vmax_empty_data() {
200        let h = HeatmapArtist {
201            data: vec![], x_labels: None, y_labels: None,
202            cmap: Colormap::Viridis, vmin: None, vmax: None,
203            show_values: false, color: Color::TAB_BLUE, label: None,
204            show_colorbar: false,
205        };
206        assert!((h.effective_vmax() - 1.0).abs() < f64::EPSILON);
207    }
208
209    #[test]
210    fn data_bounds_basic() {
211        let h = sample_heatmap();
212        let (xmin, xmax, ymin, ymax) = h.data_bounds();
213        assert!((xmin - 0.0).abs() < f64::EPSILON);
214        assert!((xmax - 3.0).abs() < f64::EPSILON);
215        assert!((ymin - 0.0).abs() < f64::EPSILON);
216        assert!((ymax - 2.0).abs() < f64::EPSILON);
217    }
218
219    #[test]
220    fn data_bounds_empty() {
221        let h = HeatmapArtist {
222            data: vec![], x_labels: None, y_labels: None,
223            cmap: Colormap::Viridis, vmin: None, vmax: None,
224            show_values: false, color: Color::TAB_BLUE, label: None,
225            show_colorbar: false,
226        };
227        assert_eq!(h.data_bounds(), (0.0, 1.0, 0.0, 1.0));
228    }
229
230    #[test]
231    fn data_bounds_single_cell() {
232        let h = HeatmapArtist {
233            data: vec![vec![42.0]], x_labels: None, y_labels: None,
234            cmap: Colormap::Viridis, vmin: None, vmax: None,
235            show_values: false, color: Color::TAB_BLUE, label: None,
236            show_colorbar: false,
237        };
238        let (xmin, xmax, ymin, ymax) = h.data_bounds();
239        assert!((xmin - 0.0).abs() < f64::EPSILON);
240        assert!((xmax - 1.0).abs() < f64::EPSILON);
241        assert!((ymin - 0.0).abs() < f64::EPSILON);
242        assert!((ymax - 1.0).abs() < f64::EPSILON);
243    }
244
245    #[test]
246    fn builder_chaining() {
247        let mut h = sample_heatmap();
248        h.colormap(Colormap::Plasma).vmin(0.0).vmax(10.0)
249            .show_values(true).label("chained");
250        assert_eq!(h.cmap, Colormap::Plasma);
251        assert_eq!(h.vmin, Some(0.0));
252        assert_eq!(h.vmax, Some(10.0));
253        assert!(h.show_values);
254        assert_eq!(h.label.as_deref(), Some("chained"));
255    }
256
257    #[test]
258    fn effective_bounds_with_nan() {
259        let h = HeatmapArtist {
260            data: vec![vec![f64::NAN, 3.0], vec![1.0, f64::NAN]],
261            x_labels: None, y_labels: None,
262            cmap: Colormap::Viridis, vmin: None, vmax: None,
263            show_values: false, color: Color::TAB_BLUE, label: None,
264            show_colorbar: false,
265        };
266        assert!((h.effective_vmin() - 1.0).abs() < f64::EPSILON);
267        assert!((h.effective_vmax() - 3.0).abs() < f64::EPSILON);
268    }
269}