src/plot/heatmap.ts

Functions

Imports

It imports src/render/data.ts, src/render/figure.ts, src/render/plot.ts, src/render/size.ts. Imported by src/plot/index.ts.

Source

1import * as Plot from "@observablehq/plot";
2import {
3  defaultFormat,
4  domainOf,
5  type KeysOfType,
6  markKey,
7  maxOf,
8  minOf,
9  rampStep,
10  rampThresholds,
11  readCategory,
12  readNumberOrNull,
13} from "../render/data.js";
14import { type FigureOptions, figure, type RampItem, validateFigure } from "../render/figure.js";
15import { decorate, keyed, renderPlot } from "../render/plot.js";
16import { DEFAULT_WIDTH } from "../render/size.js";
17
18/** Options for {@link heatmap}. */
19export interface HeatmapOptions<T extends object> extends FigureOptions {
20  readonly data: readonly T[];
21  /** The column field. */
22  readonly x: KeysOfType<T, string | number>;
23  /** The row field. */
24  readonly y: KeysOfType<T, string | number>;
25  /** The field colored by the theme's sequential ramp. `null` is an empty cell, outlined dashed. */
26  readonly value: KeysOfType<T, number | null | undefined>;
27  readonly xDomain?: readonly (string | number)[];
28  readonly yDomain?: readonly (string | number)[];
29  /**
30   * Four ascending thresholds splitting values into the ramp's five steps. Defaults to equal
31   * intervals between the lowest and highest value.
32   */
33  readonly thresholds?: readonly [number, number, number, number];
34  /** Print each value in its cell (default true when there are 100 cells or fewer). */
35  readonly cellLabels?: boolean;
36  readonly xLabel?: string | null;
37  readonly yLabel?: string | null;
38  /** What the value measures, for hover details and the table header. */
39  readonly valueLabel?: string;
40  readonly formatValue?: (value: number) => string;
41}
42
43interface Cell {
44  readonly x: string;
45  readonly y: string;
46  readonly value: number | null;
47  readonly step: number;
48}
49
50const KIND = "heatmap";
51
52/**
53 * A grid of cells colored by value on the theme's sequential ramp: cohort retention, weekly
54 * activity, plate layouts. Colors come in five labelled steps, printed values keep the chart
55 * readable without color, and empty cells are outlined with a dashed line rather than colored as zero.
56 *
57 * @example
58 * heatmap({ data: cohorts, x: "month", y: "cohort", value: "retained", alt: "..." });
59 */
60export function heatmap<T extends object>(options: HeatmapOptions<T>): string {
61  const { data, x, y, value } = options;
62  validateFigure(KIND, options);
63  if (data.length === 0) throw new Error(`${KIND}: data is empty`);
64  const format = options.formatValue ?? defaultFormat;
65  const valueLabel = options.valueLabel ?? value;
66  const xs = data.map((row, i) => readCategory(KIND, row, x, i));
67  const ys = data.map((row, i) => readCategory(KIND, row, y, i));
68  const xNames = domainOf(KIND, x, xs, options.xDomain);
69  const yNames = domainOf(KIND, y, ys, options.yDomain);
70  const values = data.map((row, i) => readNumberOrNull(KIND, row, value, i));
71  const present = values.filter((v): v is number => v !== null);
72  if (present.length === 0) throw new Error(`${KIND}: every value is empty`);
73
74  const seen = new Set<string>();
75  data.forEach((_, i) => {
76    const id = `${xs[i]}|${ys[i]}`;
77    if (seen.has(id)) throw new Error(`${KIND}: two rows for cell ${xs[i]}, ${ys[i]}`);
78    seen.add(id);
79  });
80
81  const min = minOf(present);
82  const max = maxOf(present);
83  const thresholds = rampThresholds(KIND, min, max, options.thresholds);
84  const cells = data.map((_, i): Cell => {
85    const v = values[i] ?? null;
86    return {
87      x: String(xs[i]),
88      y: String(ys[i]),
89      value: v,
90      step: v === null ? 0 : rampStep(v, thresholds),
91    };
92  });
93
94  const filledCells = cells.filter((c) => c.value !== null);
95  const emptyCells = cells.filter((c) => c.value === null);
96  const xColumn = options.xLabel ?? x;
97  const yColumn = options.yLabel ?? y;
98  const describe = (c: Cell): string =>
99    `${yColumn} ${c.y}, ${xColumn} ${c.x}: ${c.value === null ? "no data" : `${format(c.value)} ${valueLabel}`}`;
100  const showLabels = options.cellLabels ?? cells.length <= 100;
101
102  const longest = Math.max(maxOf(yNames.map((n) => n.length)), 1);
103  const marks: Plot.Markish[] = [
104    Plot.cell(emptyCells, {
105      x: "x",
106      y: "y",
107      fill: "none",
108      stroke: "var(--enarratio-status-unknown)",
109      strokeDasharray: "3,2",
110      inset: 1.5,
111      title: describe,
112    }),
113    Plot.cell(filledCells, {
114      x: "x",
115      y: "y",
116      inset: 1,
117      title: describe,
118      render: keyed(
119        filledCells,
120        (c) => ({
121          key: markKey(c.x, c.y),
122          filter: { field: y, value: c.y },
123          x: c.x,
124          column: xNames.indexOf(c.x),
125          row: yNames.indexOf(c.y),
126        }),
127        (el, c) => el.setAttribute("fill", `var(--enarratio-sequential-${c.step})`),
128      ),
129    }),
130  ];
131  if (showLabels) {
132    marks.push(
133      Plot.text(filledCells, {
134        x: "x",
135        y: "y",
136        text: (c: Cell) => (c.value === null ? "" : format(c.value)),
137        render: decorate(filledCells, (el, c) =>
138          el.setAttribute("fill", `var(--enarratio-sequential-text-${c.step})`),
139        ),
140      }),
141    );
142  }
143
144  const { svg } = renderPlot(
145    {
146      width: options.width ?? DEFAULT_WIDTH,
147      height: options.height ?? Math.max(160, yNames.length * 32 + 64),
148      marginTop: 20,
149      marginLeft: Math.min(200, longest * 7 + 24) + (options.yLabel === null ? 0 : 20),
150      marginBottom: 44,
151      x: {
152        type: "band",
153        domain: xNames,
154        label: options.xLabel === undefined ? x : options.xLabel,
155        labelAnchor: "center",
156        labelArrow: "none",
157      },
158      y: {
159        type: "band",
160        domain: yNames,
161        label: options.yLabel === undefined ? y : options.yLabel,
162        labelAnchor: "center",
163        labelArrow: "none",
164      },
165      marks,
166    },
167    options.alt,
168  );
169
170  const edges = [min, ...thresholds, max];
171  const ramp: RampItem[] = [1, 2, 3, 4, 5].map((step) => ({
172    step,
173    label:
174      step === 1
175        ? `below ${format(edges[1] as number)}`
176        : step === 5
177          ? `${format(edges[4] as number)} or more`
178          : `${format(edges[step - 1] as number)} to ${format(edges[step] as number)}`,
179  }));
180  const cellOf = new Map(cells.map((c) => [`${c.x}|${c.y}`, c]));
181
182  return figure(options, {
183    kind: "heatmap",
184    svg,
185    ramp,
186    table: {
187      columns: [`${yColumn} \\ ${xColumn}`, ...xNames],
188      rows: yNames.map((yn) => [
189        yn,
190        ...xNames.map((xn) => {
191          const c = cellOf.get(`${xn}|${yn}`);
192          return c?.value === null || c === undefined ? "" : format(c.value);
193        }),
194      ]),
195    },
196  });
197}