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    /// Sets the legend label.
36    pub fn label(&mut self, label: &str) -> &mut Self {
37        self.label = Some(label.to_string());
38        self
39    }
40
41    /// Sets the x-axis labels for the heatmap columns.
42    pub fn x_labels(&mut self, labels: Vec<String>) -> &mut Self {
43        self.x_labels = Some(labels);
44        self
45    }
46
47    /// Sets the y-axis labels for the heatmap rows.
48    pub fn y_labels(&mut self, labels: Vec<String>) -> &mut Self {
49        self.y_labels = Some(labels);
50        self
51    }
52
53    /// Returns the effective minimum value for colormap normalisation.
54    pub fn effective_vmin(&self) -> f64 {
55        if let Some(v) = self.vmin {
56            return v;
57        }
58        let mut lo = f64::INFINITY;
59        for row in &self.data {
60            for &val in row {
61                if val.is_finite() && val < lo {
62                    lo = val;
63                }
64            }
65        }
66        if lo.is_finite() { lo } else { 0.0 }
67    }
68
69    /// Returns the effective maximum value for colormap normalisation.
70    pub fn effective_vmax(&self) -> f64 {
71        if let Some(v) = self.vmax {
72            return v;
73        }
74        let mut hi = f64::NEG_INFINITY;
75        for row in &self.data {
76            for &val in row {
77                if val.is_finite() && val > hi {
78                    hi = val;
79                }
80            }
81        }
82        if hi.is_finite() { hi } else { 1.0 }
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use crate::artist::HeatmapArtist;
89    use crate::colormap::Colormap;
90    use crate::primitives::Color;
91
92    fn sample_heatmap() -> HeatmapArtist {
93        HeatmapArtist {
94            data: vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]],
95            x_labels: None, y_labels: None,
96            cmap: Colormap::Viridis, vmin: None, vmax: None,
97            show_values: false, color: Color::TAB_BLUE, label: None,
98        }
99    }
100
101    #[test]
102    fn builder_colormap() {
103        let mut h = sample_heatmap();
104        h.colormap(Colormap::Plasma);
105        assert_eq!(h.cmap, Colormap::Plasma);
106    }
107
108    #[test]
109    fn builder_vmin() {
110        let mut h = sample_heatmap();
111        h.vmin(-10.0);
112        assert_eq!(h.vmin, Some(-10.0));
113    }
114
115    #[test]
116    fn builder_vmax() {
117        let mut h = sample_heatmap();
118        h.vmax(100.0);
119        assert_eq!(h.vmax, Some(100.0));
120    }
121
122    #[test]
123    fn builder_show_values() {
124        let mut h = sample_heatmap();
125        assert!(!h.show_values);
126        h.show_values(true);
127        assert!(h.show_values);
128    }
129
130    #[test]
131    fn builder_label() {
132        let mut h = sample_heatmap();
133        h.label("my heatmap");
134        assert_eq!(h.label.as_deref(), Some("my heatmap"));
135    }
136
137    #[test]
138    fn builder_x_labels() {
139        let mut h = sample_heatmap();
140        h.x_labels(vec!["A".into(), "B".into(), "C".into()]);
141        assert_eq!(h.x_labels.as_ref().unwrap().len(), 3);
142    }
143
144    #[test]
145    fn builder_y_labels() {
146        let mut h = sample_heatmap();
147        h.y_labels(vec!["row1".into(), "row2".into()]);
148        assert_eq!(h.y_labels.as_ref().unwrap().len(), 2);
149    }
150
151    #[test]
152    fn effective_vmin_auto() {
153        let h = sample_heatmap();
154        assert!((h.effective_vmin() - 1.0).abs() < f64::EPSILON);
155    }
156
157    #[test]
158    fn effective_vmax_auto() {
159        let h = sample_heatmap();
160        assert!((h.effective_vmax() - 6.0).abs() < f64::EPSILON);
161    }
162
163    #[test]
164    fn effective_vmin_explicit() {
165        let mut h = sample_heatmap();
166        h.vmin(-5.0);
167        assert!((h.effective_vmin() - (-5.0)).abs() < f64::EPSILON);
168    }
169
170    #[test]
171    fn effective_vmax_explicit() {
172        let mut h = sample_heatmap();
173        h.vmax(50.0);
174        assert!((h.effective_vmax() - 50.0).abs() < f64::EPSILON);
175    }
176
177    #[test]
178    fn effective_vmin_empty_data() {
179        let h = HeatmapArtist {
180            data: vec![], x_labels: None, y_labels: None,
181            cmap: Colormap::Viridis, vmin: None, vmax: None,
182            show_values: false, color: Color::TAB_BLUE, label: None,
183        };
184        assert!((h.effective_vmin() - 0.0).abs() < f64::EPSILON);
185    }
186
187    #[test]
188    fn effective_vmax_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        };
194        assert!((h.effective_vmax() - 1.0).abs() < f64::EPSILON);
195    }
196
197    #[test]
198    fn data_bounds_basic() {
199        let h = sample_heatmap();
200        let (xmin, xmax, ymin, ymax) = h.data_bounds();
201        assert!((xmin - 0.0).abs() < f64::EPSILON);
202        assert!((xmax - 3.0).abs() < f64::EPSILON);
203        assert!((ymin - 0.0).abs() < f64::EPSILON);
204        assert!((ymax - 2.0).abs() < f64::EPSILON);
205    }
206
207    #[test]
208    fn data_bounds_empty() {
209        let h = HeatmapArtist {
210            data: vec![], x_labels: None, y_labels: None,
211            cmap: Colormap::Viridis, vmin: None, vmax: None,
212            show_values: false, color: Color::TAB_BLUE, label: None,
213        };
214        assert_eq!(h.data_bounds(), (0.0, 1.0, 0.0, 1.0));
215    }
216
217    #[test]
218    fn data_bounds_single_cell() {
219        let h = HeatmapArtist {
220            data: vec![vec![42.0]], x_labels: None, y_labels: None,
221            cmap: Colormap::Viridis, vmin: None, vmax: None,
222            show_values: false, color: Color::TAB_BLUE, label: None,
223        };
224        let (xmin, xmax, ymin, ymax) = h.data_bounds();
225        assert!((xmin - 0.0).abs() < f64::EPSILON);
226        assert!((xmax - 1.0).abs() < f64::EPSILON);
227        assert!((ymin - 0.0).abs() < f64::EPSILON);
228        assert!((ymax - 1.0).abs() < f64::EPSILON);
229    }
230
231    #[test]
232    fn builder_chaining() {
233        let mut h = sample_heatmap();
234        h.colormap(Colormap::Plasma).vmin(0.0).vmax(10.0)
235            .show_values(true).label("chained");
236        assert_eq!(h.cmap, Colormap::Plasma);
237        assert_eq!(h.vmin, Some(0.0));
238        assert_eq!(h.vmax, Some(10.0));
239        assert!(h.show_values);
240        assert_eq!(h.label.as_deref(), Some("chained"));
241    }
242
243    #[test]
244    fn effective_bounds_with_nan() {
245        let h = HeatmapArtist {
246            data: vec![vec![f64::NAN, 3.0], vec![1.0, f64::NAN]],
247            x_labels: None, y_labels: None,
248            cmap: Colormap::Viridis, vmin: None, vmax: None,
249            show_values: false, color: Color::TAB_BLUE, label: None,
250        };
251        assert!((h.effective_vmin() - 1.0).abs() < f64::EPSILON);
252        assert!((h.effective_vmax() - 3.0).abs() < f64::EPSILON);
253    }
254}