Heatmaps plus

app/components/Heatmap.tsx

@@ -7,10 +7,16 @@
  */
 import { useMemo, useState } from 'react'
 import * as R from 'remeda'
+import { match } from 'ts-pattern'
 import type uPlot from 'uplot'
 
 import { ChartTooltip, type TopBottom, type LeftRight } from '~/components/ChartTooltip'
 import { type UPlotOptions, FramedChart } from '~/components/FramedChart'
+import {
+  getSampleCountMetadata,
+  getScalingFunction,
+  type ColorScale,
+} from '~/components/heatmap-color-scale'
 import {
   timeFormatterForRange,
   useChartTheme,
@@ -40,38 +46,20 @@ const RAMP_STOPS = [
   ['--content-accent', 'oklch(0.77 0.1919 163.7)'],
 ] as const
 
+const CSS_VAR_RAMP: ColorRamp = {
+  stops: RAMP_STOPS.map(([name, fallback]) => `var(${name}, ${fallback})`),
+}
+
 // Interpolates the color between the two closest stops, for 0 <= t <= 1.
 const rampColor = ({ stops }: ColorRamp, t: number): string => {
   const clamped = R.clamp(t, { min: 0, max: 1 })
   const lastIndex = stops.length - 1
   const progress = clamped * lastIndex
   const start = R.clamp(Math.floor(progress), { max: lastIndex - 1 })
-  const lerp = progress - start
+  const localT = progress - start
   const from = stops[start]
   const to = stops[start + 1]
-  return `color-mix(in oklch, ${from} ${(1 - lerp) * 100}%, ${to})`
-}
-
-// grabbing style from the chart root instead of document body
-// that way it can be themed e.g. `yellow-theme`
-const getColorRamp = (el: Element): ColorRamp => {
-  const style = getComputedStyle(el)
-  return {
-    stops: RAMP_STOPS.map(
-      ([name, fallback]) => style.getPropertyValue(name).trim() || fallback
-    ),
-  }
-}
-
-// count -> ramp position. Doing a square root or a log makes for some different
-// smoothing. Nothing is truly intuitive, but linearly showing the fraction is
-// most easily compared to the legend at the bottom
-const rampT = (count: number, maxCount: number) => {
-  const fraction = count / maxCount
-  // TODO: pick!
-  return fraction
-  // return Math.log10(lerp(1, 10, fraction))
-  // return Math.sqrt(fraction)
+  return `color-mix(in oklch, ${from}, ${to} ${localT * 100}%)`
 }
 
 type Hover = {
@@ -93,6 +81,7 @@ type HeatmapProps = {
   unit?: string
   /** Drawn rotated alongside the y-axis tick labels, e.g. a unit */
   yAxisLabel?: string
+  colorScale?: ColorScale
 }
 
 const defaultYAxisTickFormatter = (val: number) => val.toLocaleString()
@@ -105,6 +94,7 @@ export function Heatmap({
   yAxisTickFormatter = defaultYAxisTickFormatter,
   unit,
   yAxisLabel,
+  colorScale = 'ranked',
 }: HeatmapProps) {
   const theme = useChartTheme()
   const [hover, setHover] = useState<Hover | null>(null)
@@ -118,8 +108,8 @@ export function Heatmap({
   const binCount = bins.length
   const colCount = distributions.length
 
-  const maxSampleCount = useMemo(
-    () => distributions.reduce((max, d) => (d ? Math.max(max, ...d.counts) : max), 0),
+  const sampleCountMetadata = useMemo(
+    () => getSampleCountMetadata(distributions.flatMap((d) => d?.counts ?? [])),
     [distributions]
   )
 
@@ -129,10 +119,22 @@ export function Heatmap({
 
   // uPlot doesn't actually have heatmap support, but we can use its position
   // functions to make this (somewhat) easy
-  const drawCells = useMemo(
-    () => (u: uPlot) => {
+  const drawCells = useMemo(() => {
+    const scaleCount = getScalingFunction(sampleCountMetadata, colorScale)
+
+    return (u: uPlot) => {
       if (binCount === 0 || colCount === 0) return
-      const colorRamp = getColorRamp(u.root)
+
+      // since a canvas can't resolve CSS variables, we simulate it by reading
+      // the ramp colors off the element furthest down the tree (the chart's
+      // root element)
+      const style = getComputedStyle(u.root)
+      const colorRamp = {
+        stops: RAMP_STOPS.map(
+          ([name, fallback]) => style.getPropertyValue(name).trim() || fallback
+        ),
+      }
+
       const ctx = u.ctx
       ctx.save()
       ctx.beginPath()
@@ -149,7 +151,7 @@ export function Heatmap({
 
           const topEdge = u.valToPos(rowIndex + 1, 'y', true)
           const bottomEdge = u.valToPos(rowIndex, 'y', true)
-          ctx.fillStyle = rampColor(colorRamp, rampT(count, maxSampleCount))
+          ctx.fillStyle = rampColor(colorRamp, scaleCount(count))
           ctx.fillRect(
             leftEdge - CELL_OVERFLOW,
             topEdge - CELL_OVERFLOW,
@@ -159,9 +161,16 @@ export function Heatmap({
         })
       })
       ctx.restore()
-    },
-    [distributions, leftEdges, rightEdges, colCount, binCount, maxSampleCount]
-  )
+    }
+  }, [
+    distributions,
+    leftEdges,
+    rightEdges,
+    colCount,
+    binCount,
+    colorScale,
+    sampleCountMetadata,
+  ])
 
   const tooltipPlugin = useMemo<uPlot.Plugin>(
     () => ({
@@ -307,7 +316,19 @@ export function Heatmap({
       chartOptions={chartOptions}
       data={data}
       uRef={uRef}
-      legend={<HeatmapLegend maxSampleCount={maxSampleCount} axisText={theme.axisText} />}
+      legend={
+        sampleCountMetadata.ranks.size > 0 && (
+          <HeatmapLegend
+            minSampleCount={match(colorScale)
+              .with('ranked', () => sampleCountMetadata.minSampleCount)
+              .with('linear', () => 0)
+              .with('logarithmic', () => 1)
+              .exhaustive()}
+            maxSampleCount={sampleCountMetadata.maxSampleCount}
+            axisText={theme.axisText}
+          />
+        )
+      }
     >
       {hover && (
         <>
@@ -342,28 +363,69 @@ export function Heatmap({
 }
 
 function HeatmapLegend({
+  minSampleCount,
   maxSampleCount,
   axisText,
 }: {
+  minSampleCount: number
   maxSampleCount: number
   axisText: string
 }) {
   return (
     <div className="mt-3 flex items-center gap-2 pl-5">
-      <span className="text-mono-xs" style={{ color: axisText }}>
-        0
-      </span>
-      <div
-        className="h-3 w-32 rounded-sm"
-        style={{
-          backgroundImage: `linear-gradient(in oklch to right, ${RAMP_STOPS.map(
-            ([name, fallback]) => `var(${name}, ${fallback})`
-          ).join(', ')})`,
-        }}
-      ></div>
+      {minSampleCount === maxSampleCount ? (
+        <div
+          className="h-3 w-3 rounded-sm"
+          style={{
+            backgroundColor: R.last(CSS_VAR_RAMP.stops),
+          }}
+        ></div>
+      ) : (
+        <>
+          <span className="text-mono-xs" style={{ color: axisText }}>
+            {minSampleCount.toLocaleString()}
+          </span>
+          <div
+            className="h-3 w-32 rounded-sm"
+            style={{
+              backgroundImage: `linear-gradient(in oklch to right, ${CSS_VAR_RAMP.stops.join(', ')})`,
+            }}
+          ></div>
+        </>
+      )}
       <span className="text-mono-xs" style={{ color: axisText }}>
         {maxSampleCount.toLocaleString()}
       </span>
     </div>
   )
 }
+
+const SAMPLE_VALUES = [1, 5, 10, 50, 95, 100]
+const SAMPLE_VALUES_METADATA = getSampleCountMetadata(SAMPLE_VALUES)
+const LONGEST_SAMPLE_ITEM = SAMPLE_VALUES.map((n) => n.toLocaleString().length).reduce(
+  (acc, length) => Math.max(acc, length),
+  0
+)
+
+export function SampleColorStrip({ colorScale }: { colorScale: ColorScale }) {
+  const scaleCount = getScalingFunction(SAMPLE_VALUES_METADATA, colorScale)
+
+  return (
+    <div className="flex items-end gap-px">
+      {SAMPLE_VALUES.map((n) => (
+        <div
+          key={n}
+          className="text-secondary text-mono-code flex flex-col items-center text-xs"
+        >
+          <div className="text-center" style={{ width: `${LONGEST_SAMPLE_ITEM + 0.5}ch` }}>
+            {n.toLocaleString()}
+          </div>
+          <div
+            className="h-3 w-full"
+            style={{ backgroundColor: rampColor(CSS_VAR_RAMP, scaleCount(n)) }}
+          />
+        </div>
+      ))}
+    </div>
+  )
+}

app/components/heatmap-color-scale.spec.ts

@@ -0,0 +1,89 @@
+/*
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, you can obtain one at https://mozilla.org/MPL/2.0/.
+ *
+ * Copyright Oxide Computer Company
+ */
+import { describe, expect, it } from 'vitest'
+
+import {
+  getSampleCountMetadata,
+  getScalingFunction,
+  type ColorScale,
+} from './heatmap-color-scale'
+
+const scale = (counts: number[], colorScale: ColorScale) =>
+  getScalingFunction(getSampleCountMetadata(counts), colorScale)
+
+describe('getSampleCountMetadata', () => {
+  it('ignores zeros and duplicates', () => {
+    const { minSampleCount, maxSampleCount, ranks } = getSampleCountMetadata([
+      50, 0, 5, 50, 0, 100,
+    ])
+    expect(minSampleCount).toBe(5)
+    expect(maxSampleCount).toBe(100)
+    expect([...ranks]).toEqual([
+      [5, 0],
+      [50, 1],
+      [100, 2],
+    ])
+  })
+
+  it('handles lack of samples', () => {
+    const { minSampleCount, maxSampleCount, ranks } = getSampleCountMetadata([0, 0])
+    expect(minSampleCount).toBe(0)
+    expect(maxSampleCount).toBe(0)
+    expect(ranks.size).toBe(0)
+  })
+})
+
+describe('linear', () => {
+  it('is the fraction of the max', () => {
+    const f = scale([10, 50, 200], 'linear')
+    expect(f(10)).toBe(0.05)
+    expect(f(50)).toBe(0.25)
+    expect(f(200)).toBe(1)
+  })
+})
+
+describe('logarithmic', () => {
+  it("has 1 as 0 and the collection's greatest value as 1", () => {
+    const f = scale([10, 200], 'logarithmic')
+    expect(f(1)).toBe(0)
+    expect(f(200)).toBe(1)
+  })
+
+  it('gives equal ratios equal steps', () => {
+    const expectEqualSpacing = (values: number[]) => {
+      const f = scale(values, 'logarithmic')
+      values.forEach((value, i) => expect(f(value)).toBeCloseTo(i / (values.length - 1)))
+    }
+
+    expectEqualSpacing([1, 10, 100, 1000, 10000])
+    expectEqualSpacing([1, 7, 49, 343])
+  })
+
+  it("makes an exception when the collection's greatest value is 1", () => {
+    expect(scale([1, 0], 'logarithmic')(1)).toBe(1)
+  })
+})
+
+describe('ranked', () => {
+  it('spaces distinct counts evenly', () => {
+    const f = scale([1, 2, 1000], 'ranked')
+    expect(f(1)).toBe(0)
+    expect(f(2)).toBe(0.5)
+    expect(f(1000)).toBe(1)
+  })
+
+  it('ignores duplicates', () => {
+    const f = scale([3, 3, 7, 9], 'ranked')
+    expect(f(3)).toBe(0)
+    expect(f(7)).toBe(0.5)
+  })
+
+  it('puts a single count at the top', () => {
+    expect(scale([4, 4, 0], 'ranked')(4)).toBe(1)
+  })
+})

app/components/heatmap-color-scale.ts

@@ -0,0 +1,61 @@
+/*
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, you can obtain one at https://mozilla.org/MPL/2.0/.
+ *
+ * Copyright Oxide Computer Company
+ */
+import * as R from 'remeda'
+import { match } from 'ts-pattern'
+
+export type ColorScale = 'linear' | 'logarithmic' | 'ranked'
+
+const scaleLinear =
+  (maxCount: number) =>
+  (count: number): number =>
+    count / maxCount
+const scaleLog = (maxCount: number) =>
+  maxCount <= 1
+    ? // if the max is 1, a logarithmic scaling function is undefined. but
+      // since 1 is then the only valid value, in a charting context it makes
+      // sense to treat that as a scale of 100%
+      R.constant(1)
+    : (count: number): number => Math.log(count) / Math.log(maxCount)
+const scaleRank = (ranks: Map<number, number>) =>
+  ranks.size <= 1
+    ? R.constant(1)
+    : (count: number): number => (ranks.get(count) ?? 0) / (ranks.size - 1)
+
+export type SampleCountMetadata = {
+  minSampleCount: number
+  maxSampleCount: number
+  ranks: Map<number, number>
+}
+
+export const getSampleCountMetadata = (counts: number[]): SampleCountMetadata => {
+  const allCounts = R.sort(
+    R.unique(counts).filter((x) => x !== 0),
+    (a, b) => a - b
+  )
+
+  return {
+    minSampleCount: R.first(allCounts) ?? 0,
+    maxSampleCount: R.last(allCounts) ?? 0,
+    ranks: new Map(allCounts.map((n, i) => [n, i])),
+  }
+}
+
+/**
+ * Maps a sample count to a position for a color ramp, from 0 to 1.
+ * Expects that the count given to the final scaling function is a nonzero
+ * natural number.
+ */
+export const getScalingFunction = (
+  metadata: SampleCountMetadata,
+  colorScale: ColorScale
+): ((n: number) => number) =>
+  match(colorScale)
+    .with('linear', () => scaleLinear(metadata.maxSampleCount))
+    .with('logarithmic', () => scaleLog(metadata.maxSampleCount))
+    .with('ranked', () => scaleRank(metadata.ranks))
+    .exhaustive()