rich_ext/chart/
heatmap.rs

1//! A labelled grid of values drawn as shades.
2
3use rich::measure::Measurement;
4use rich::{Console, ConsoleOptions, Renderable, Segment, Style};
5
6use super::{
7    Charset, Line, Scale, ValueFormat, cell_units, cells, entries_width, has_colour,
8    lines_to_segments, theme_style, truncate, wrap_entries,
9};
10
11/// Levels for blocks, lowest first.
12const SHADES: [char; 5] = [' ', '░', '▒', '▓', '█'];
13/// Levels for ASCII, lowest first.
14const ASCII_SHADES: [char; 10] = [' ', '.', ':', '-', '=', '+', '*', '#', '%', '@'];
15/// How many `chart.heat.N` styles there are.
16const HEAT_STYLES: usize = 5;
17
18/// A grid of values, one row per label and one column per header, each
19/// cell a shade from lowest to highest.
20///
21/// - Values are scaled from the smallest to the largest (or
22///   [`range`](Self::range)) into equal steps: five in blocks (` ░▒▓█`), ten
23///   in ASCII (` .:-=+*#%@`). The shade is the value, so the grid reads in
24///   black and white; with colour each step also takes a `chart.heat.N`
25///   style, cold to hot.
26/// - A NaN or infinite value is drawn `·` (`?` in ASCII), so a missing
27///   value never passes for a low one.
28/// - The legend under the grid shows the steps between the lowest and
29///   highest value of the scale.
30/// - Each cell is [`cell_width`](Self::cell_width) cells wide (2 by
31///   default). Column headers are written at their column's first cell when
32///   they fit, leaving a space after the one before, so with narrow cells
33///   every second or third header shows.
34/// - Given less width, the cells narrow to one cell, then the row labels
35///   are cut, then neighbouring columns are merged (each shows the mean of
36///   its finite values) so every value still counts.
37///
38/// ```
39/// use rich::Console;
40/// use rich_ext::chart::{Charset, Heatmap};
41///
42/// let console = Console::builder().width(40).color_system(None).build();
43/// let map = Heatmap::new()
44///     .columns(["mon", "tue", "wed", "thu"])
45///     .row("api", [1.0, 4.0, 9.0, 2.0])
46///     .row("web", [0.0, 6.0, 8.0, f64::NAN])
47///     .cell_width(4)
48///     .charset(Charset::Ascii);
49/// assert_eq!(
50///     console.render_export(&map),
51///     concat!(
52///         "    mon tue wed thu        \n",
53///         "api ....====@@@@::::       \n",
54///         "web     ****%%%%????       \n",
55///         "0 [ .:-=+*#%@] 9  ? no data\n",
56///     )
57/// );
58/// ```
59#[derive(Clone, Debug, PartialEq)]
60#[must_use = "builder methods return a new value; it does nothing until it is rendered or used"]
61pub struct Heatmap {
62    columns: Vec<String>,
63    rows: Vec<(String, Vec<f64>)>,
64    min: Option<f64>,
65    max: Option<f64>,
66    charset: Charset,
67    cell_width: usize,
68    legend: bool,
69    format: ValueFormat,
70}
71
72impl Default for Heatmap {
73    fn default() -> Self {
74        Self::new()
75    }
76}
77
78/// How a heatmap fits its width.
79struct Grid {
80    label: usize,
81    cell: usize,
82    /// For each drawn column, the range of data columns it stands for.
83    buckets: Vec<(usize, usize)>,
84}
85
86impl Heatmap {
87    /// An empty heatmap; add rows with [`row`](Self::row).
88    pub fn new() -> Self {
89        Heatmap {
90            columns: Vec::new(),
91            rows: Vec::new(),
92            min: None,
93            max: None,
94            charset: Charset::Auto,
95            cell_width: 2,
96            legend: true,
97            format: ValueFormat::Compact,
98        }
99    }
100
101    /// The column headers.
102    pub fn columns<S: Into<String>>(mut self, headers: impl IntoIterator<Item = S>) -> Self {
103        self.columns = headers.into_iter().map(Into::into).collect();
104        self
105    }
106
107    /// Add a row of values, in column order.
108    pub fn row(mut self, label: impl Into<String>, values: impl IntoIterator<Item = f64>) -> Self {
109        self.rows.push((label.into(), values.into_iter().collect()));
110        self
111    }
112
113    /// Fix the scale; values outside it take the end shades.
114    pub fn range(mut self, min: f64, max: f64) -> Self {
115        self.min = Some(min);
116        self.max = Some(max);
117        self
118    }
119
120    /// Glyphs to draw with: [`Charset::Ascii`] asks for the ASCII shades;
121    /// anything else draws blocks (ASCII on an ASCII-only console).
122    pub fn charset(mut self, charset: Charset) -> Self {
123        self.charset = charset;
124        self
125    }
126
127    /// Cells per value (default 2, at least 1).
128    pub fn cell_width(mut self, width: usize) -> Self {
129        self.cell_width = width.max(1);
130        self
131    }
132
133    /// Show the legend (default on).
134    pub fn legend(mut self, show: bool) -> Self {
135        self.legend = show;
136        self
137    }
138
139    /// How the legend writes the bounds (default [`ValueFormat::Compact`]).
140    pub fn format(mut self, format: ValueFormat) -> Self {
141        self.format = format;
142        self
143    }
144
145    /// The scale every value is shaded on.
146    pub fn scale(&self) -> Scale {
147        let values = self.rows.iter().flat_map(|(_, v)| v.iter().copied());
148        Scale::from_values(values).bounds(self.min, self.max)
149    }
150
151    fn column_count(&self) -> usize {
152        self.rows
153            .iter()
154            .map(|(_, v)| v.len())
155            .max()
156            .unwrap_or(0)
157            .max(self.columns.len())
158    }
159
160    fn label_width(&self) -> usize {
161        self.rows.iter().map(|(l, _)| cells(l)).max().unwrap_or(0)
162    }
163
164    fn has_missing(&self) -> bool {
165        let n = self.column_count();
166        self.rows
167            .iter()
168            .any(|(_, v)| v.len() < n || v.iter().any(|x| !x.is_finite()))
169    }
170
171    fn layout(&self, width: usize) -> Grid {
172        let n = self.column_count();
173        let l = self.label_width();
174        let part = |l: usize| l + usize::from(l > 0);
175        let each = |cell: usize| Grid {
176            label: 0,
177            cell,
178            buckets: (0..n).map(|i| (i, i + 1)).collect(),
179        };
180        let room = width.saturating_sub(part(l));
181        if room >= n {
182            let cell = self.cell_width.min(room / n.max(1)).max(1);
183            return Grid {
184                label: l,
185                ..each(cell)
186            };
187        }
188        let keep = l.min(3);
189        if let Some(label) = width.checked_sub(n + 1)
190            && label >= keep
191            && label > 0
192        {
193            return Grid {
194                label: label.min(l),
195                ..each(1)
196            };
197        }
198        // Merge columns into as many as fit.
199        let label = if width > part(keep) + 1 { keep } else { 0 };
200        let slots = width.saturating_sub(part(label)).max(1);
201        let buckets = (0..slots)
202            .map(|i| {
203                let start = i * n / slots;
204                (start, ((i + 1) * n / slots).max(start + 1))
205            })
206            .collect();
207        Grid {
208            label,
209            cell: 1,
210            buckets,
211        }
212    }
213
214    /// The step of `value` on `scale`, among `levels`.
215    fn level(scale: &Scale, value: f64, levels: usize) -> Option<usize> {
216        let n = scale.normalize(value)?;
217        Some(((n * levels as f64).floor() as usize).min(levels - 1))
218    }
219
220    fn level_style(console: &Console, level: usize, levels: usize) -> Style {
221        let key = format!("chart.heat.{}", level * HEAT_STYLES / levels + 1);
222        theme_style(console, &key)
223    }
224
225    fn lines(&self, console: &Console, options: &ConsoleOptions) -> Vec<Line> {
226        let width = options.max_width;
227        let ascii = self.charset.resolve(console, options, Charset::Blocks) == Charset::Ascii;
228        let colour = has_colour(console);
229        if self.column_count() == 0 {
230            let mut line = Line::new();
231            line.push(&truncate("no data", width, ascii), None);
232            return vec![line];
233        }
234        let shades: &[char] = if ascii { &ASCII_SHADES } else { &SHADES };
235        let missing = if ascii { '?' } else { '·' };
236        let levels = shades.len();
237        let scale = self.scale();
238        let grid = self.layout(width);
239        let label_style = colour.then(|| theme_style(console, "chart.label"));
240        let lead = if grid.label > 0 { grid.label + 1 } else { 0 };
241        let mut out = Vec::new();
242
243        // Headers at their column's first cell, where they fit.
244        if self.columns.iter().any(|c| !c.is_empty()) {
245            let plot = grid.buckets.len() * grid.cell;
246            let mut row = vec![" ".to_string(); plot];
247            let mut free = 0;
248            for (i, (start, _)) in grid.buckets.iter().enumerate() {
249                let Some(header) = self.columns.get(*start) else {
250                    continue;
251                };
252                let header = truncate(header, usize::MAX, ascii);
253                let at = i * grid.cell;
254                let len = cells(&header);
255                if at < free || at + len > plot || len == 0 {
256                    continue;
257                }
258                for (k, unit) in cell_units(&header).into_iter().enumerate() {
259                    row[at + k] = unit;
260                }
261                free = at + len + 1;
262            }
263            if row.iter().any(|c| c != " ") {
264                let mut line = Line::new();
265                line.pad(lead);
266                line.push(&row.concat(), label_style.clone());
267                out.push(line);
268            }
269        }
270
271        for (label, values) in &self.rows {
272            let mut line = Line::new();
273            if grid.label > 0 {
274                let text = truncate(label, grid.label, ascii);
275                let pad = grid.label - cells(&text);
276                line.push(&text, label_style.clone());
277                line.pad(pad + 1);
278            }
279            for &(start, end) in &grid.buckets {
280                let bucket: Vec<f64> = (start..end)
281                    .filter_map(|i| values.get(i).copied())
282                    .filter(|v| v.is_finite())
283                    .collect();
284                let value = if bucket.is_empty() {
285                    f64::NAN
286                } else {
287                    bucket.iter().sum::<f64>() / bucket.len() as f64
288                };
289                let (glyph, style) = match Self::level(&scale, value, levels) {
290                    Some(level) => (
291                        shades[level],
292                        colour.then(|| Self::level_style(console, level, levels)),
293                    ),
294                    None => (missing, None),
295                };
296                line.push(&glyph.to_string().repeat(grid.cell), style);
297            }
298            out.push(line);
299        }
300        if self.legend {
301            out.extend(wrap_entries(
302                self.legend_entries(console, ascii, colour),
303                width,
304            ));
305        }
306        out
307    }
308
309    fn legend_entries(
310        &self,
311        console: &Console,
312        ascii: bool,
313        colour: bool,
314    ) -> Vec<Vec<(String, Option<Style>)>> {
315        let shades: &[char] = if ascii { &ASCII_SHADES } else { &SHADES };
316        let scale = self.scale();
317        let label_style = colour.then(|| theme_style(console, "chart.label"));
318        let mut ramp = vec![(
319            format!("{} [", self.format.format(scale.min())),
320            label_style.clone(),
321        )];
322        for (level, c) in shades.iter().enumerate() {
323            let style = colour.then(|| Self::level_style(console, level, shades.len()));
324            ramp.push((c.to_string(), style));
325        }
326        ramp.push((
327            format!("] {}", self.format.format(scale.max())),
328            label_style.clone(),
329        ));
330        let mut entries = vec![ramp];
331        if self.has_missing() {
332            let missing = if ascii { "?" } else { "·" };
333            entries.push(vec![(format!("{missing} no data"), label_style)]);
334        }
335        entries
336    }
337}
338
339impl crate::surface::Draw for Heatmap {
340    fn draw(
341        &self,
342        console: &Console,
343        options: &ConsoleOptions,
344        screen: &mut crate::surface::Screen,
345        rect: crate::surface::Rect,
346    ) {
347        super::draw_lines(
348            self.lines(console, options),
349            options.max_width,
350            options,
351            screen,
352            rect,
353        );
354    }
355}
356
357impl Renderable for Heatmap {
358    fn rich_render(&self, console: &Console, options: &ConsoleOptions) -> Vec<Segment> {
359        lines_to_segments(self.lines(console, options), options.max_width)
360    }
361
362    fn measure(&self, console: &Console, options: &ConsoleOptions) -> Measurement {
363        let n = self.column_count();
364        if n == 0 {
365            return Measurement::new(7, 7).with_maximum(options.max_width);
366        }
367        let ascii = self.charset.resolve(console, options, Charset::Blocks) == Charset::Ascii;
368        let l = self.label_width();
369        let grid = (l + usize::from(l > 0)).saturating_add(n.saturating_mul(self.cell_width));
370        let legend = if self.legend {
371            entries_width(&self.legend_entries(console, ascii, false))
372        } else {
373            0
374        };
375        let max = grid.max(legend);
376        Measurement::new(n.min(max), max)
377            .with_maximum(options.max_width)
378            .normalize()
379    }
380}
381
382/// `Heatmap, 2 rows by 4 columns, values from 0 to 9, 1 missing`, then
383/// where the highest and lowest values are (`Highest 9 at api, wed; lowest
384/// 0 at web, mon`).
385impl crate::a11y::AccessibleText for Heatmap {
386    fn accessible_text(&self, _width: usize) -> String {
387        use super::describe::{count, counted, range};
388        let columns = self.column_count();
389        if self.rows.is_empty() || columns == 0 {
390            return "Heatmap, no values".to_string();
391        }
392        let mut head = format!(
393            "Heatmap, {} by {}",
394            counted(self.rows.len(), "row", "rows"),
395            counted(columns, "column", "columns")
396        );
397        // Every cell, a short row's missing ones included.
398        let cells: Vec<(&str, usize, f64)> = self
399            .rows
400            .iter()
401            .flat_map(|(label, values)| {
402                (0..columns).map(move |c| {
403                    let value = values.get(c).copied().unwrap_or(f64::NAN);
404                    (label.as_str(), c, value)
405                })
406            })
407            .collect();
408        let finite: Vec<&(&str, usize, f64)> = cells.iter().filter(|c| c.2.is_finite()).collect();
409        let missing = cells.len() - finite.len();
410        if let Some((lo, hi)) = super::describe::extent(finite.iter().map(|c| c.2)) {
411            head.push_str(&format!(", values {}", range(lo, hi, self.format)));
412        }
413        if missing > 0 {
414            head.push_str(&format!(", {} missing", count(missing)));
415        }
416        let header = |c: usize| {
417            self.columns
418                .get(c)
419                .filter(|h| !h.is_empty())
420                .cloned()
421                .unwrap_or_else(|| format!("column {}", c + 1))
422        };
423        let at = |cell: &(&str, usize, f64)| {
424            let row = if cell.0.is_empty() {
425                String::new()
426            } else {
427                format!("{}, ", cell.0)
428            };
429            format!("{} at {row}{}", self.format.format(cell.2), header(cell.1))
430        };
431        // The first highest and the first lowest, in reading order.
432        let mut highest: Option<&(&str, usize, f64)> = None;
433        let mut lowest: Option<&(&str, usize, f64)> = None;
434        for &cell in &finite {
435            if highest.is_none_or(|h| cell.2 > h.2) {
436                highest = Some(cell);
437            }
438            if lowest.is_none_or(|l| cell.2 < l.2) {
439                lowest = Some(cell);
440            }
441        }
442        let mut out = vec![head];
443        if let (Some(hi), Some(lo)) = (highest, lowest) {
444            out.push(format!("Highest {}; lowest {}", at(hi), at(lo)));
445        }
446        out.join("\n")
447    }
448}
449
450#[cfg(test)]
451mod tests {
452    use super::*;
453    use crate::a11y::AccessibleText;
454
455    #[test]
456    fn accessible_text_finds_the_extremes() {
457        let map = Heatmap::new()
458            .columns(["mon", "tue", "wed", "thu"])
459            .row("api", [1.0, 4.0, 9.0, 2.0])
460            .row("web", [0.0, 6.0, 8.0, f64::NAN]);
461        assert_eq!(
462            map.accessible_text(80),
463            "Heatmap, 2 rows by 4 columns, values from 0 to 9, 1 missing\n\
464             Highest 9 at api, wed; lowest 0 at web, mon"
465        );
466        let short = Heatmap::new().row("", [3.0, 3.0]);
467        assert_eq!(
468            short.accessible_text(80),
469            "Heatmap, 1 row by 2 columns, values at 3\n\
470             Highest 3 at column 1; lowest 3 at column 1"
471        );
472        assert_eq!(Heatmap::new().accessible_text(80), "Heatmap, no values");
473    }
474}
rich_ext/chart/heatmap.rs

rich_ext/chart/
heatmap.rs
