plotters_statistical/series/
heatmap.rs

1//! Generic heatmap series — a grid of color-mapped cells with optional numeric
2//! annotations. The building block under the correlation / missingness heatmap
3//! figures, but usable directly on any chart.
4//!
5//! Cells are placed at integer coordinates: column `c` at `x = c`, row `r` at
6//! `y = nrows - 1 - r` (so row 0 is at the top). Build the chart with ranges
7//! `-0.5..ncols as f64 - 0.5` and `-0.5..nrows as f64 - 0.5`.
8
9use plotters::element::{Drawable, PointCollection};
10use plotters::style::text_anchor::{HPos, Pos, VPos};
11use plotters::style::{IntoFont, RGBColor, ShapeStyle, TextStyle};
12use plotters_backend::{BackendCoord, DrawingBackend, DrawingErrorKind};
13
14use crate::colormap::{GradientColorMap, Normalization};
15use crate::style::fill_style;
16
17const MISSING_COLOR: RGBColor = RGBColor(235, 235, 235);
18
19/// Per-cell numeric annotation settings.
20#[derive(Debug, Clone)]
21pub struct HeatmapAnnotation {
22    /// Decimal places shown.
23    pub precision: usize,
24    /// Font size in points.
25    pub font_size: u32,
26    /// Text color; `None` picks black or white automatically for contrast
27    /// against each cell.
28    pub text_color: Option<RGBColor>,
29}
30
31impl Default for HeatmapAnnotation {
32    fn default() -> Self {
33        Self {
34            precision: 2,
35            font_size: 12,
36            text_color: None,
37        }
38    }
39}
40
41/// A grid of color-mapped cells.
42#[derive(Debug, Clone)]
43pub struct Heatmap {
44    // Two corners per cell (row-major): [lower-left, upper-right] in data coords.
45    corners: Vec<(f64, f64)>,
46    values: Vec<f64>,
47    colormap: GradientColorMap,
48    norm: Normalization,
49    annotate: Option<HeatmapAnnotation>,
50    border: Option<ShapeStyle>,
51    cell_gap: i32,
52}
53
54impl Heatmap {
55    /// Build from a row-major matrix (`values[row][col]`). Rows may not be empty
56    /// and must all be the same width; otherwise the shorter rows are padded
57    /// with `NaN` (rendered as the "missing" color). The default normalization
58    /// is linear over the finite value range.
59    pub fn new(values: &[Vec<f64>]) -> Self {
60        let nrows = values.len();
61        let ncols = values.iter().map(|r| r.len()).max().unwrap_or(0);
62        let mut flat = Vec::with_capacity(nrows * ncols);
63        let mut corners = Vec::with_capacity(nrows * ncols * 2);
64        for (r, row) in values.iter().enumerate() {
65            let y = (nrows - 1 - r) as f64;
66            for c in 0..ncols {
67                let v = row.get(c).copied().unwrap_or(f64::NAN);
68                flat.push(v);
69                corners.push((c as f64 - 0.5, y - 0.5));
70                corners.push((c as f64 + 0.5, y + 0.5));
71            }
72        }
73        let (mut lo, mut hi) = (f64::INFINITY, f64::NEG_INFINITY);
74        for &v in &flat {
75            if v.is_finite() {
76                lo = lo.min(v);
77                hi = hi.max(v);
78            }
79        }
80        if !lo.is_finite() {
81            lo = 0.0;
82            hi = 1.0;
83        }
84        Self {
85            corners,
86            values: flat,
87            colormap: GradientColorMap::viridis(),
88            norm: Normalization::Linear { min: lo, max: hi },
89            annotate: None,
90            border: None,
91            cell_gap: 0,
92        }
93    }
94
95    /// Set the color map.
96    pub fn colormap(mut self, colormap: GradientColorMap) -> Self {
97        self.colormap = colormap;
98        self
99    }
100
101    /// Set the value→`[0,1]` normalization.
102    pub fn normalization(mut self, norm: Normalization) -> Self {
103        self.norm = norm;
104        self
105    }
106
107    /// Annotate each cell with its value.
108    pub fn annotate(mut self, annotation: HeatmapAnnotation) -> Self {
109        self.annotate = Some(annotation);
110        self
111    }
112
113    /// Draw a border around each cell.
114    pub fn cell_border(mut self, style: ShapeStyle) -> Self {
115        self.border = Some(style);
116        self
117    }
118
119    /// Shrink each cell by `gap` pixels on every side (visual gutters).
120    pub fn cell_gap(mut self, gap: i32) -> Self {
121        self.cell_gap = gap;
122        self
123    }
124}
125
126fn contrast_text(bg: RGBColor) -> RGBColor {
127    let l = 0.299 * bg.0 as f64 + 0.587 * bg.1 as f64 + 0.114 * bg.2 as f64;
128    if l > 140.0 {
129        RGBColor(0, 0, 0)
130    } else {
131        RGBColor(255, 255, 255)
132    }
133}
134
135impl<'a> PointCollection<'a, (f64, f64)> for &'a Heatmap {
136    type Point = &'a (f64, f64);
137    type IntoIter = &'a [(f64, f64)];
138    fn point_iter(self) -> &'a [(f64, f64)] {
139        &self.corners
140    }
141}
142
143impl<DB: DrawingBackend> Drawable<DB> for Heatmap {
144    fn draw<I: Iterator<Item = BackendCoord>>(
145        &self,
146        points: I,
147        backend: &mut DB,
148        _parent_dim: (u32, u32),
149    ) -> Result<(), DrawingErrorKind<DB::ErrorType>> {
150        let pix: Vec<BackendCoord> = points.collect();
151        let ncells = self.values.len();
152        if pix.len() < ncells * 2 {
153            return Ok(());
154        }
155        for k in 0..ncells {
156            let a = pix[2 * k];
157            let b = pix[2 * k + 1];
158            let (x0, x1) = (a.0.min(b.0) + self.cell_gap, a.0.max(b.0) - self.cell_gap);
159            let (y0, y1) = (a.1.min(b.1) + self.cell_gap, a.1.max(b.1) - self.cell_gap);
160            if x1 <= x0 || y1 <= y0 {
161                continue;
162            }
163            let v = self.values[k];
164            let color = if v.is_finite() {
165                self.colormap.color(self.norm.t(v))
166            } else {
167                MISSING_COLOR
168            };
169            backend.draw_rect((x0, y0), (x1, y1), &fill_style(&color), true)?;
170            if let Some(border) = &self.border {
171                backend.draw_rect((x0, y0), (x1, y1), border, false)?;
172            }
173            if let Some(ann) = &self.annotate {
174                if v.is_finite() {
175                    let tc = ann.text_color.unwrap_or_else(|| contrast_text(color));
176                    let style = TextStyle::from(("sans-serif", ann.font_size as i32).into_font())
177                        .color(&tc)
178                        .pos(Pos::new(HPos::Center, VPos::Center));
179                    let cx = (x0 + x1) / 2;
180                    let cy = (y0 + y1) / 2;
181                    let text = format!("{:.*}", ann.precision, v);
182                    backend.draw_text(&text, &style, (cx, cy))?;
183                }
184            }
185        }
186        Ok(())
187    }
188}
plotters_statistical/series/heatmap.rs

plotters_statistical/series/
heatmap.rs
