Skip to main content

esoc_chart/chart/
heatmap.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! Heatmap series with color-mapped cells.
3
4use esoc_gfx::canvas::Canvas;
5use esoc_gfx::color::Color;
6use esoc_gfx::element::{DrawElement, Element};
7use esoc_gfx::geom::Rect;
8use esoc_gfx::layer::Layer;
9use esoc_gfx::palette::Palette;
10use esoc_gfx::style::{Fill, FontStyle, TextAnchor};
11use esoc_gfx::transform::CoordinateTransform;
12
13use crate::series::{DataBounds, SeriesRenderer};
14use crate::theme::Theme;
15
16/// A heatmap series rendering a 2D matrix as colored cells.
17#[derive(Clone, Debug)]
18pub struct HeatmapSeries {
19    /// 2D data: `data[row][col]`.
20    pub data: Vec<Vec<f64>>,
21    /// Optional series label.
22    pub label: Option<String>,
23    /// Whether to annotate cells with values.
24    pub annotate: bool,
25    /// Optional row labels.
26    pub row_labels: Option<Vec<String>>,
27    /// Optional column labels.
28    pub col_labels: Option<Vec<String>>,
29    /// Color palette for mapping values.
30    pub palette: Option<Palette>,
31}
32
33impl HeatmapSeries {
34    /// Create a new heatmap series.
35    pub fn new(data: Vec<Vec<f64>>) -> Self {
36        Self {
37            data,
38            label: None,
39            annotate: false,
40            row_labels: None,
41            col_labels: None,
42            palette: None,
43        }
44    }
45
46    fn value_range(&self) -> (f64, f64) {
47        let min = self
48            .data
49            .iter()
50            .flat_map(|row| row.iter().copied())
51            .fold(f64::INFINITY, f64::min);
52        let max = self
53            .data
54            .iter()
55            .flat_map(|row| row.iter().copied())
56            .fold(f64::NEG_INFINITY, f64::max);
57        (min, max)
58    }
59}
60
61impl SeriesRenderer for HeatmapSeries {
62    fn data_bounds(&self) -> DataBounds {
63        let rows = self.data.len();
64        let cols = self.data.first().map_or(0, Vec::len);
65        DataBounds::new(-0.5, cols as f64 - 0.5, -0.5, rows as f64 - 0.5)
66    }
67
68    fn render(
69        &self,
70        canvas: &mut Canvas,
71        transform: &CoordinateTransform,
72        theme: &Theme,
73        _series_index: usize,
74    ) {
75        let rows = self.data.len();
76        if rows == 0 {
77            return;
78        }
79        let _cols = self.data[0].len();
80        let (vmin, vmax) = self.value_range();
81        let palette = self.palette.clone().unwrap_or_else(Palette::viridis);
82
83        for (r, row) in self.data.iter().enumerate() {
84            for (c, &val) in row.iter().enumerate() {
85                let t = if (vmax - vmin).abs() < 1e-15 {
86                    0.5
87                } else {
88                    (val - vmin) / (vmax - vmin)
89                };
90                let color = palette.sample(t);
91
92                // Cell rectangle — heatmap uses y-axis inverted from data row order
93                let y = (rows - 1 - r) as f64;
94                let p_tl = transform.to_pixel(c as f64 - 0.5, y + 0.5);
95                let p_br = transform.to_pixel(c as f64 + 0.5, y - 0.5);
96                let rx = p_tl.x.min(p_br.x);
97                let ry = p_tl.y.min(p_br.y);
98                let rw = (p_br.x - p_tl.x).abs();
99                let rh = (p_br.y - p_tl.y).abs();
100
101                canvas.add(DrawElement::new(
102                    Element::Rect {
103                        rect: Rect::new(rx, ry, rw, rh),
104                        fill: Fill::Solid(color),
105                        stroke: None,
106                        rx: 0.0,
107                    },
108                    Layer::Data,
109                ));
110
111                // Annotation
112                if self.annotate {
113                    let center = transform.to_pixel(c as f64, y);
114                    let text_color = if t > 0.5 { Color::BLACK } else { Color::WHITE };
115                    let font = FontStyle {
116                        family: theme.font_family.clone(),
117                        size: theme.tick_font_size,
118                        weight: 400,
119                        color: text_color,
120                        anchor: TextAnchor::Middle,
121                    };
122                    let text = if (val - val.round()).abs() < 1e-9 {
123                        format!("{}", val as i64)
124                    } else {
125                        format!("{val:.2}")
126                    };
127                    canvas.add(DrawElement::new(
128                        Element::text(center.x, center.y + theme.tick_font_size * 0.35, text, font),
129                        Layer::Annotations,
130                    ));
131                }
132            }
133        }
134    }
135
136    fn label(&self) -> Option<&str> {
137        self.label.as_deref()
138    }
139}