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,8 +81,8 @@ export function shapeViewModel(
   settingsSpec: SettingsSpec,
   chartDimensions: Dimensions,
   heatmapTable: HeatmapTable,
-  colorScale: ColorScaleType,
-  filterRanges: Array<[number, number | null]>,
+  colorScale: ColorScaleType['scale'],
+  bandsToHide: Array<[number, number]>,
   { height, pageSize }: GridHeightParams,
 ): ShapeViewModel {
   const gridStrokeWidth = config.grid.stroke.width ?? 1;
@@ -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;
     }
@@ -208,7 +208,7 @@ export function shapeViewModel(
         width: config.cell.border.strokeWidth,
       },
       value: d.value,
-      visible: !isFilteredValue(filterRanges, d.value),
+      visible: !isValueHidden(d.value, bandsToHide),
       formatted: spec.valueFormatter(d.value),
     };
     return acc;
@@ -391,11 +391,6 @@ function getCellKey(x: NonNullable<PrimitiveValue>, y: NonNullable<PrimitiveValu
   return [String(x), String(y)].join('&_&');
 }
 
-function isFilteredValue(filterRanges: Array<[number, number | null]>, value: number) {
-  return filterRanges.some(([min, max]) => {
-    if (max !== null && value > min && value < max) {
-      return true;
-    }
-    return max === null && value > min;
-  });
+function isValueHidden(value: number, rangesToHide: Array<[number, number]>) {
+  return rangesToHide.some(([min, max]) => min <= value && value < max);
 }

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

@@ -12,30 +12,29 @@ 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, { bands }, deselectedDataSeries): LegendItem[] => {
+    if (spec === null) {
+      return EMPTY_LEGEND;
     }
 
-    return colorScale.ticks.map((tick) => {
-      const color = colorScale.config(tick);
-      const seriesIdentifier = {
-        key: String(tick),
-        specId: String(tick),
-      };
-
+    return bands.map(({ label, color }) => {
       return {
+        // the band label is considered unique by construction
+        seriesIdentifiers: [
+          {
+            key: label,
+            specId: label,
+          },
+        ],
         color,
-        label: `> ${spec.valueFormatter ? spec.valueFormatter(tick) : tick}`,
-        seriesIdentifiers: [seriesIdentifier],
-        isSeriesHidden: deselectedDataSeries.some((dataSeries) => dataSeries.key === seriesIdentifier.key),
+        label,
+        isSeriesHidden: deselectedDataSeries.some((dataSeries) => dataSeries.key === label),
         isToggleable: true,
-        path: [{ index: 0, value: seriesIdentifier.key }],
+        path: [{ index: 0, value: label }],
         keys: [],
       };
     });

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,27 +35,25 @@ export const geometries = createCustomCachedSelector(
     chartDimensions,
     settingSpec,
     heatmapTable,
-    colorScale,
-    legendItems,
+    { bands, scale: colorScale },
     deselectedSeries,
     gridHeightParams,
   ): ShapeViewModel => {
-    const deselectedTicks = new Set(
+    // instead of using the specId, each legend item is associated with an unique band label
+    const disabledBandLabels = new Set(
       deselectedSeries.map(({ specId }) => {
-        return Number(specId);
+        return 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]);
-      }
-      return acc;
-    }, []);
+
+    const bandsToHide: Array<[number, number]> = bands
+      .filter(({ label }) => {
+        return disabledBandLabels.has(label);
+      })
+      .map(({ start, end }) => [start, end]);
 
     return heatmapSpec
-      ? render(heatmapSpec, settingSpec, chartDimensions, heatmapTable, colorScale, ranges, gridHeightParams)
+      ? render(heatmapSpec, settingSpec, chartDimensions, heatmapTable, colorScale, bandsToHide, gridHeightParams)
       : nullShapeViewModel();
   },
 );

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,117 @@ import {
 
 import { ScaleType } from '../../../../scales/constants';
 import { createCustomCachedSelector } from '../../../../state/create_selector';
+import { Color, identity } from '../../../../utils/common';
+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;
-  ticks: number[];
+type ScaleModelType<S> = {
+  scale: S;
+  bands: Array<{ start: 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 Band = { start: number; end: number; label: string; color: Color };
+
+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,
+};
+const DEFAULT_COLOR_SCALE_TYPE = ScaleType.Linear;
+
 /**
  * @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, bands } = SCALE_TYPE_TO_SCALE_FN[spec.colorScale ?? DEFAULT_COLOR_SCALE_TYPE](spec, heatmapTable);
+    return {
+      scale,
+      bands: dedupBands(bands, spec, scale),
+    };
   },
 );
 
-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 dedupBands(
+  bands: Array<{ start: number }>,
+  spec: HeatmapSpec,
+  scale: ScaleLinear<string, string> | ScaleQuantile<string> | ScaleQuantize<string> | ScaleThreshold<number, string>,
+) {
+  const formatter = spec.valueFormatter ?? identity;
+  const bandsWithFormattedStarts = bands.reduce<Map<string, Band>>((acc, { start }, index, array) => {
+    const label = `≥ ${formatter(start)}`;
+    const end = array.length === index + 1 ? Infinity : array[index + 1].start;
+    acc.set(label, { start, end, label, color: scale(start) });
+    return acc;
+  }, new Map());
+  return [...bandsWithFormattedStarts.values()];
+}
+
+function getQuantizedScale(spec: HeatmapSpec, heatmapTable: HeatmapTable): ScaleQuantizeType {
+  const domain =
+    Array.isArray(spec.ranges) && spec.ranges.length > 1 ? (spec.ranges as [number, number]) : heatmapTable.extent;
+  const colors = spec.colors?.length > 0 ? spec.colors : DEFAULT_COLORS;
+  // we use the data extent or only the first two values in the `ranges` prop
+  const scale = scaleQuantize<string>().domain(domain).range(colors);
+  // quantize scale works as the linear one, we should manually
+  // compute the start of each band corresponding to the quantized segments
+  const numOfSegments = colors.length;
+  const interval = (domain[1] - domain[0]) / numOfSegments;
+  const bands = colors.map((color, i) => ({ start: domain[0] + interval * i }));
+
+  return {
+    scale,
+    bands,
+  };
+}
+
+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 bands array should contain all quantiles + the minimum value
+  const bands = [...new Set([heatmapTable.extent[0], ...scale.quantiles()])].map((start) => ({ start }));
+  return {
+    scale,
+    bands,
+  };
+}
+
+function getThresholdScale(spec: HeatmapSpec, heatmapTable: HeatmapTable): ScaleThresholdType {
+  const colors = spec.colors ?? DEFAULT_COLORS;
+  const domain = spec.ranges ?? heatmapTable.extent;
+  const scale = scaleThreshold<number, string>().domain(domain).range(colors);
+  // the start band array should contain all the thresholds + the minimum value
+  const bands = [...new Set([heatmapTable.extent[0], ...domain])].map((start) => ({ start }));
+  return {
+    scale,
+    bands,
+  };
+}
 
-  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 bands = [...new Set([domain[0], ...scale.ticks(6)])].map((start) => ({ start }));
+  return {
+    scale,
+    bands,
+  };
 }

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

@@ -25,8 +25,8 @@ export function render(
   settingsSpec: SettingsSpec,
   chartDimensions: Dimensions,
   heatmapTable: HeatmapTable,
-  colorScale: ColorScaleType,
-  filterRanges: Array<[number, number | null]>,
+  colorScale: ColorScaleType['scale'],
+  bandsToHide: Array<[number, number]>,
   gridHeightParams: GridHeightParams,
 ): ShapeViewModel {
   const textMeasurer = document.createElement('canvas');
@@ -46,7 +46,7 @@ export function render(
     chartDimensions,
     heatmapTable,
     colorScale,
-    filterRanges,
+    bandsToHide,
     gridHeightParams,
   );
 }