fix(heatmap): compute nice legend items from color scale

integration/tests/heatmap.test.ts

@@ -0,0 +1,33 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import { common } from '../page_objects/common';
+
+describe('Heatmap color scale', () => {
+  it('quantile', async () => {
+    await common.expectChartAtUrlToMatchScreenshot(
+      'http://localhost:9001/?path=/story/heatmap-alpha--categorical&globals=backgrounds.value:transparent;themes.value:Light&knob-color scale=quantile&knob-ranges=auto&knob-colors=green, yellow, red',
+    );
+  });
+  it('quantize', async () => {
+    await common.expectChartAtUrlToMatchScreenshot(
+      'http://localhost:9001/?path=/story/heatmap-alpha--categorical&globals=backgrounds.value:transparent;themes.value:Light&knob-color scale=quantize&knob-ranges=auto&knob-colors=green, yellow, red',
+    );
+  });
+  it('threshold', async () => {
+    await common.expectChartAtUrlToMatchScreenshot(
+      'http://localhost:9001/?path=/story/heatmap-alpha--categorical&globals=backgrounds.value:transparent;themes.value:Light&knob-color scale=threshold&knob-ranges=10000, 40000&knob-colors=green, yellow, red',
+    );
+  });
+
+  it('linear', async () => {
+    await common.expectChartAtUrlToMatchScreenshot(
+      'http://localhost:9001/?path=/story/heatmap-alpha--categorical&globals=backgrounds.value:transparent;themes.value:Light&knob-color scale=linear&knob-ranges=0, 10000, 25000, 50000, 100000&knob-colors=green, yellow, red, purple',
+    );
+  });
+});

packages/charts/src/chart_types/heatmap/layout/viewmodel/viewmodel.ts

@@ -81,7 +81,7 @@ export function shapeViewModel(
   settingsSpec: SettingsSpec,
   chartDimensions: Dimensions,
   heatmapTable: HeatmapTable,
-  colorScale: ColorScaleType,
+  colorScale: ColorScaleType['scale'],
   filterRanges: Array<[number, number | null]>,
   { height, pageSize }: GridHeightParams,
 ): ShapeViewModel {
@@ -187,7 +187,7 @@ export function shapeViewModel(
     const x = xScale(String(d.x));
     const y = yScale(String(d.y))! + gridStrokeWidth;
     const yIndex = yValues.indexOf(d.y);
-    const color = colorScale.config(d.value);
+    const color = colorScale(d.value);
     if (x === undefined || y === undefined || yIndex === -1) {
       return acc;
     }
@@ -393,9 +393,9 @@ function getCellKey(x: NonNullable<PrimitiveValue>, y: NonNullable<PrimitiveValu
 
 function isFilteredValue(filterRanges: Array<[number, number | null]>, value: number) {
   return filterRanges.some(([min, max]) => {
-    if (max !== null && value > min && value < max) {
+    if (max !== null && value >= min && value < max) {
       return true;
     }
-    return max === null && value > min;
+    return max === null && value >= min;
   });
 }

packages/charts/src/chart_types/heatmap/state/selectors/compute_legend.ts

@@ -12,26 +12,25 @@ import { getDeselectedSeriesSelector } from '../../../../state/selectors/get_des
 import { getColorScale } from './get_color_scale';
 import { getSpecOrNull } from './heatmap_spec';
 
+const EMPTY_LEGEND: LegendItem[] = [];
 /** @internal */
 export const computeLegendSelector = createCustomCachedSelector(
   [getSpecOrNull, getColorScale, getDeselectedSeriesSelector],
-  (spec, colorScale, deselectedDataSeries): LegendItem[] => {
-    const legendItems: LegendItem[] = [];
-
-    if (colorScale === null || spec === null) {
-      return legendItems;
+  (spec, { ticks, scale }, deselectedDataSeries): LegendItem[] => {
+    if (spec === null) {
+      return EMPTY_LEGEND;
     }
 
-    return colorScale.ticks.map((tick) => {
-      const color = colorScale.config(tick);
+    return ticks.map(({ tick, formattedTick }, i) => {
+      const color = scale(tick);
       const seriesIdentifier = {
         key: String(tick),
         specId: String(tick),
       };
 
       return {
         color,
-        label: `> ${spec.valueFormatter ? spec.valueFormatter(tick) : tick}`,
+        label: `>${i === 0 ? '=' : ''} ${formattedTick}`,
         seriesIdentifiers: [seriesIdentifier],
         isSeriesHidden: deselectedDataSeries.some((dataSeries) => dataSeries.key === seriesIdentifier.key),
         isToggleable: true,

packages/charts/src/chart_types/heatmap/state/selectors/geometries.ts

@@ -15,7 +15,6 @@ import { getColorScale } from './get_color_scale';
 import { getGridHeightParamsSelector } from './get_grid_full_height';
 import { getHeatmapSpecSelector } from './get_heatmap_spec';
 import { getHeatmapTableSelector } from './get_heatmap_table';
-import { getLegendItemsLabelsSelector } from './get_legend_items_labels';
 import { render } from './scenegraph';
 
 const getDeselectedSeriesSelector = (state: GlobalChartState) => state.interactions.deselectedDataSeries;
@@ -28,7 +27,6 @@ export const geometries = createCustomCachedSelector(
     getSettingsSpecSelector,
     getHeatmapTableSelector,
     getColorScale,
-    getLegendItemsLabelsSelector,
     getDeselectedSeriesSelector,
     getGridHeightParamsSelector,
   ],
@@ -37,8 +35,7 @@ export const geometries = createCustomCachedSelector(
     chartDimensions,
     settingSpec,
     heatmapTable,
-    colorScale,
-    legendItems,
+    { ticks, scale: colorScale },
     deselectedSeries,
     gridHeightParams,
   ): ShapeViewModel => {
@@ -47,11 +44,10 @@ export const geometries = createCustomCachedSelector(
         return Number(specId);
       }),
     );
-    const { ticks } = colorScale;
-    const ranges = ticks.reduce<Array<[number, number | null]>>((acc, d, i) => {
-      if (deselectedTicks.has(d)) {
-        const rangeEnd = i + 1 === ticks.length ? null : ticks[i + 1];
-        acc.push([d, rangeEnd]);
+    const ranges = ticks.reduce<Array<[number, number | null]>>((acc, { tick }, i) => {
+      if (deselectedTicks.has(tick)) {
+        const rangeEnd = i + 1 === ticks.length ? null : ticks[i + 1].tick;
+        acc.push([tick, rangeEnd]);
       }
       return acc;
     }, []);

packages/charts/src/chart_types/heatmap/state/selectors/get_color_scale.ts

@@ -6,7 +6,6 @@
  * Side Public License, v 1.
  */
 
-import { extent as d3Extent } from 'd3-array';
 import { interpolateHcl } from 'd3-interpolate';
 import {
   ScaleLinear,
@@ -21,63 +20,115 @@ import {
 
 import { ScaleType } from '../../../../scales/constants';
 import { createCustomCachedSelector } from '../../../../state/create_selector';
+import { HeatmapSpec } from '../../specs/heatmap';
+import { HeatmapTable } from './compute_chart_dimensions';
 import { getHeatmapSpecSelector } from './get_heatmap_spec';
 import { getHeatmapTableSelector } from './get_heatmap_table';
 
-type ScaleModelType<Type, Config> = {
-  type: Type;
-  config: Config;
+type ScaleModelType<S> = {
+  scale: S;
   ticks: number[];
 };
-type ScaleLinearType = ScaleModelType<typeof ScaleType.Linear, ScaleLinear<string, string>>;
-type ScaleQuantizeType = ScaleModelType<typeof ScaleType.Quantize, ScaleQuantize<string>>;
-type ScaleQuantileType = ScaleModelType<typeof ScaleType.Quantile, ScaleQuantile<string>>;
-type ScaleThresholdType = ScaleModelType<typeof ScaleType.Threshold, ScaleThreshold<number, string>>;
+
+type ScaleLinearType = ScaleModelType<ScaleLinear<string, string>>;
+type ScaleQuantizeType = ScaleModelType<ScaleQuantize<string>>;
+type ScaleQuantileType = ScaleModelType<ScaleQuantile<string>>;
+type ScaleThresholdType = ScaleModelType<ScaleThreshold<number, string>>;
+
 /** @internal */
 export type ColorScaleType = ScaleLinearType | ScaleQuantizeType | ScaleQuantileType | ScaleThresholdType;
 
+const DEFAULT_COLORS = ['green', 'red'];
+
+const SCALE_TYPE_TO_SCALE_FN = {
+  [ScaleType.Linear]: getLinearScale,
+  [ScaleType.Quantile]: getQuantileScale,
+  [ScaleType.Quantize]: getQuantizedScale,
+  [ScaleType.Threshold]: getThresholdScale,
+};
+
 /**
  * @internal
  * Gets color scale based on specification and values range.
  */
 export const getColorScale = createCustomCachedSelector(
   [getHeatmapSpecSelector, getHeatmapTableSelector],
   (spec, heatmapTable) => {
-    const { colors, colorScale: colorScaleSpec } = spec;
-
-    // compute the color scale based domain and colors
-    const { ranges = heatmapTable.extent } = spec;
-    const colorRange = colors ?? ['green', 'red'];
-
-    const colorScale = {
-      type: colorScaleSpec,
-    } as ColorScaleType;
-    if (colorScale.type === ScaleType.Quantize) {
-      colorScale.config = scaleQuantize<string>()
-        .domain(d3Extent(ranges) as [number, number])
-        .range(colorRange);
-      colorScale.ticks = colorScale.config.ticks(spec.colors.length);
-    } else if (colorScale.type === ScaleType.Quantile) {
-      colorScale.config = scaleQuantile<string>().domain(ranges).range(colorRange);
-      colorScale.ticks = colorScale.config.quantiles();
-    } else if (colorScale.type === ScaleType.Threshold) {
-      colorScale.config = scaleThreshold<number, string>().domain(ranges).range(colorRange);
-      colorScale.ticks = colorScale.config.domain();
-    } else {
-      colorScale.config = scaleLinear<string>().domain(ranges).interpolate(interpolateHcl).range(colorRange);
-      colorScale.ticks = addBaselineOnLinearScale(ranges[0], ranges[1], colorScale.config.ticks(6));
-    }
-    return colorScale;
+    const { scale, ticks } = SCALE_TYPE_TO_SCALE_FN[spec.colorScale ?? ScaleType.Linear](spec, heatmapTable);
+    return {
+      scale,
+      ...dedupTicks(ticks, spec),
+    };
   },
 );
 
-function addBaselineOnLinearScale(min: number, max: number, ticks: Array<number>): Array<number> {
-  if (min < 0 && max < 0) {
-    return [...ticks, 0];
-  }
-  if (min >= 0 && max >= 0) {
-    return [0, ...ticks];
-  }
+function dedupTicks(ticks: number[], spec: HeatmapSpec) {
+  return ticks.reduce<{ uniqueTicks: string[]; ticks: Array<{ tick: number; formattedTick: string }> }>(
+    (acc, curr) => {
+      const formattedTick = `${spec.valueFormatter ? spec.valueFormatter(curr) : curr}`;
+      if (acc.uniqueTicks.includes(formattedTick)) {
+        return acc;
+      }
+      return {
+        uniqueTicks: [...acc.uniqueTicks, formattedTick],
+        ticks: [...acc.ticks, { tick: curr, formattedTick }],
+      };
+    },
+    { uniqueTicks: [], ticks: [] },
+  );
+}
+
+function getQuantizedScale(spec: HeatmapSpec, heatmapTable: HeatmapTable): ScaleQuantizeType {
+  const dataExtent = spec.ranges ?? heatmapTable.extent;
+  const colors = spec.colors ?? DEFAULT_COLORS;
+  // we use the data extent or only the first two values in the `ranges` prop
+  const domain: [number, number] = [dataExtent[0], dataExtent[1]];
+  const scale = scaleQuantize<string>().domain(domain).range(colors);
+  // quantize scale works as the linear one, we should manually
+  // compute the ticks corresponding to the quantized segments
+  const numOfSegments = colors.length;
+  const interval = (domain[1] - domain[0]) / numOfSegments;
+  const ticks = colors.map((d, i) => domain[0] + interval * i);
+
+  return {
+    scale,
+    ticks,
+  };
+}
+
+function getQuantileScale(spec: HeatmapSpec, heatmapTable: HeatmapTable): ScaleQuantileType {
+  const colors = spec.colors ?? DEFAULT_COLORS;
+  const domain = heatmapTable.table.map(({ value }) => value);
+  const scale = scaleQuantile<string>().domain(domain).range(colors);
+  // the ticks array should contain all quantiles + the minimum value
+  const ticks = [...new Set([heatmapTable.extent[0], ...scale.quantiles()])];
+  return {
+    scale,
+    ticks,
+  };
+}
+
+function getThresholdScale(spec: HeatmapSpec, heatmapTable: HeatmapTable): ScaleThresholdType {
+  const colors = spec.colors ?? DEFAULT_COLORS;
+  const scale = scaleThreshold<number, string>()
+    .domain(spec.ranges ?? heatmapTable.extent)
+    .range(colors);
+  // the ticks array should contain all the thresholds + the minimum value
+  const ticks = [...new Set([heatmapTable.extent[0], ...scale.domain()])];
+  return {
+    scale,
+    ticks,
+  };
+}
 
-  return ticks;
+function getLinearScale(spec: HeatmapSpec, heatmapTable: HeatmapTable): ScaleLinearType {
+  const domain = spec.ranges ?? heatmapTable.extent;
+  const colors = spec.colors ?? DEFAULT_COLORS;
+  const scale = scaleLinear<string>().domain(domain).interpolate(interpolateHcl).range(colors).clamp(true);
+  // adding initial and final range/extent value if they are rounded values.
+  const ticks = [...new Set([domain[0], ...scale.ticks(6)])];
+  return {
+    scale,
+    ticks,
+  };
 }

packages/charts/src/chart_types/heatmap/state/selectors/scenegraph.ts

@@ -25,7 +25,7 @@ export function render(
   settingsSpec: SettingsSpec,
   chartDimensions: Dimensions,
   heatmapTable: HeatmapTable,
-  colorScale: ColorScaleType,
+  colorScale: ColorScaleType['scale'],
   filterRanges: Array<[number, number | null]>,
   gridHeightParams: GridHeightParams,
 ): ShapeViewModel {

storybook/stories/heatmap/2_categorical.tsx

@@ -7,7 +7,7 @@
  */
 
 import { action } from '@storybook/addon-actions';
-import { extent } from 'd3-array';
+import { select, text } from '@storybook/addon-knobs';
 import React from 'react';
 
 import { Chart, Heatmap, ScaleType, Settings } from '@elastic/charts';
@@ -16,9 +16,20 @@ import { BABYNAME_DATA } from '@elastic/charts/src/utils/data_samples/babynames'
 import { useBaseTheme } from '../../use_base_theme';
 
 export const Example = () => {
-  const data = BABYNAME_DATA.filter(([year]) => year > 1950);
-  const values = data.map((d) => +d[3]);
-  const [min, max] = extent(values);
+  const data = BABYNAME_DATA.filter(([year]) => year > 1950 && year < 1960);
+  const colorScale = select(
+    'color scale',
+    {
+      [ScaleType.Linear]: ScaleType.Linear,
+      [ScaleType.Quantile]: ScaleType.Quantile,
+      [ScaleType.Quantize]: ScaleType.Quantize,
+      [ScaleType.Threshold]: ScaleType.Threshold,
+    },
+    ScaleType.Linear,
+  );
+  const ranges = text('ranges', 'auto');
+  const colors = text('colors', 'green, yellow, red');
+
   return (
     <Chart>
       <Settings
@@ -30,10 +41,10 @@ export const Example = () => {
       />
       <Heatmap
         id="heatmap2"
-        colorScale={ScaleType.Linear}
-        ranges={[min!, (max! - min!) / 2, max!]}
-        colors={['green', 'yellow', 'red']}
-        data={BABYNAME_DATA.filter(([year]) => year > 1950)}
+        colorScale={colorScale}
+        ranges={ranges === 'auto' ? undefined : ranges.split(',').map((d) => Number(d.trim()))}
+        colors={colors.split(',').map((d) => d.trim())}
+        data={data}
         xAccessor={(d) => d[2]}
         yAccessor={(d) => d[0]}
         valueAccessor={(d) => d[3]}