[Feat] Add custom color scale for aggregate layers (#2860)

Signed-off-by: Ihor Dykhta <dikhta.igor@gmail.com> Co-authored-by: Ilya Boyandin <iboyandin@foursquare.com>
keplergl · Dec 26, 2024 · a897715 · a897715
1 parent 6bc5946
commit a897715
Show file tree

Hide file tree

Showing 21 changed files with 324 additions and 127 deletions.
diff --git a/src/components/src/common/column-stats-chart.tsx b/src/components/src/common/column-stats-chart.tsx
@@ -7,12 +7,10 @@ import React, {useCallback, useEffect, useMemo, useRef, useState} from 'react';
 import styled from 'styled-components';
 
 import {KeplerTable} from '@kepler.gl/table';
-import {Field} from '@kepler.gl/types';
+import {Bin, Field} from '@kepler.gl/types';
 import {
   ColorBreak,
   ColorBreakOrdinal,
-  histogramFromThreshold,
-  histogramFromValues,
   isNumber,
   isNumericColorBreaks,
   useDimensions
@@ -29,7 +27,6 @@ const COLOR_CHART_TICK_WRAPPER_HEIGHT = 10;
 const COLOR_CHART_TICK_HEIGHT = 8;
 const COLOR_CHART_TICK_WIDTH = 4;
 const COLOR_CHART_TICK_BORDER_COLOR = '#999999';
-const HISTOGRAM_BINS = 30;
 
 const StyledContainer = styled.div.attrs({
   className: 'color-chart-loading'
@@ -208,21 +205,29 @@ export type ColumnStatsChartWLoadingProps = {
   colorField: Field;
   dataset: KeplerTable;
   colorBreaks: ColorBreak[] | ColorBreakOrdinal[] | null;
+  allBins: Bin[];
+  filteredBins: Bin[];
+  isFiltered: boolean;
+  histogramDomain: number[];
   onChangedUpdater: (ticks: ColorBreak[]) => void;
 };
 
 export type ColumnStatsChartProps = {
-  field: Field;
-  dataset: KeplerTable;
+  allBins: Bin[];
+  filteredBins: Bin[];
+  isFiltered: boolean;
+  histogramDomain: number[];
   colorBreaks: ColorBreak[];
   onChangedUpdater: (ticks: ColorBreak[]) => void;
 };
 function ColumnStatsChartFactory(
   HistogramPlot: ReturnType<typeof HistogramPlotFactory>
 ): React.FC<ColumnStatsChartWLoadingProps> {
   const ColumnStatsChart: React.FC<ColumnStatsChartProps> = ({
-    field,
-    dataset,
+    allBins,
+    filteredBins,
+    isFiltered,
+    histogramDomain,
     colorBreaks,
     onChangedUpdater
   }) => {
@@ -241,38 +246,6 @@ function ColumnStatsChartFactory(
       isTickChangingRef.current = false;
     }, [ticks]);
 
-    const valueAccessor = useMemo(() => {
-      return idx => dataset.getValue(field.name, idx);
-    }, [dataset, field.name]);
-
-    const columnStats = field.filterProps?.columnStats;
-
-    // get bins with allIndexes
-    const allBins = useMemo(() => {
-      if (columnStats?.bins) {
-        return columnStats?.bins;
-      }
-      return histogramFromValues(dataset.allIndexes, HISTOGRAM_BINS, valueAccessor);
-    }, [columnStats, dataset.allIndexes, valueAccessor]);
-
-    const isFiltered = dataset.filteredIndexForDomain.length !== dataset.allIndexes.length;
-
-    // get filteredBins
-    const filteredBins = useMemo(() => {
-      if (!isFiltered) {
-        return allBins;
-      }
-      // get threholds
-      const filterEmptyBins = false;
-      const threholds = allBins.map(b => b.x0);
-      return histogramFromThreshold(
-        threholds,
-        dataset.filteredIndexForDomain,
-        valueAccessor,
-        filterEmptyBins
-      );
-    }, [dataset, valueAccessor, allBins, isFiltered]);
-
     // histograms used by histogram-plot.js
     const histogramsByGroup = useMemo(
       () => ({
@@ -282,34 +255,6 @@ function ColumnStatsChartFactory(
       [allBins, filteredBins]
     );
 
-    // get domain (min, max) of histogram
-    const histogramDomain = useMemo(() => {
-      if (columnStats && columnStats.quantiles && columnStats.mean) {
-        // no need to recalcuate min/max/mean if its already in columnStats
-        return [
-          columnStats.quantiles[0].value,
-          columnStats.quantiles[columnStats.quantiles.length - 1].value,
-          columnStats.mean
-        ];
-      }
-      let domainMin = Number.POSITIVE_INFINITY;
-      let domainMax = Number.NEGATIVE_INFINITY;
-      let nValid = 0;
-      let domainSum = 0;
-      dataset.allIndexes.forEach((x, i) => {
-        const val = valueAccessor(x);
-        if (isNumber(val)) {
-          if (val < domainMin) domainMin = val;
-          if (val > domainMax) domainMax = val;
-          domainSum += val;
-          nValid += 1;
-        }
-      });
-      const histogramMean = nValid > 0 ? domainSum / nValid : 0;
-
-      return [domainMin, domainMax, histogramMean];
-    }, [dataset, valueAccessor, columnStats]);
-
     // get colors from colorBreaks
     const domainColors = useMemo(
       () => (colorBreaks ? colorBreaks.map(c => c.data) : []),
@@ -417,19 +362,24 @@ function ColumnStatsChartFactory(
     colorField,
     dataset,
     colorBreaks,
+    allBins,
+    filteredBins,
+    isFiltered,
+    histogramDomain,
     onChangedUpdater
   }) => {
-    const fieldName = colorField.name;
-    const field = useMemo(() => dataset.getColumnField(fieldName), [dataset, fieldName]);
-
-    const isLoading = field?.isLoadingStats;
+    const fieldName = colorField?.name;
+    const field = useMemo(() => (fieldName ? dataset.getColumnField(fieldName) : null), [
+      dataset,
+      fieldName
+    ]);
 
-    if (!isNumericColorBreaks(colorBreaks) || !field) {
+    if (!isNumericColorBreaks(colorBreaks)) {
       // TODO: implement display for ordinal breaks
       return null;
     }
 
-    if (isLoading) {
+    if (field?.isLoadingStats) {
       return (
         <StyledContainer>
           <LoadingSpinner />
@@ -440,8 +390,10 @@ function ColumnStatsChartFactory(
     return (
       <ColumnStatsChart
         colorBreaks={colorBreaks}
-        field={field}
-        dataset={dataset}
+        allBins={allBins}
+        filteredBins={filteredBins}
+        isFiltered={isFiltered}
+        histogramDomain={histogramDomain}
         onChangedUpdater={onChangedUpdater}
       />
     );

diff --git a/src/components/src/map-container.tsx b/src/components/src/map-container.tsx
@@ -30,7 +30,8 @@ import {
   updateMapboxLayers,
   LayerBaseConfig,
   VisualChannelDomain,
-  EditorLayerUtils
+  EditorLayerUtils,
+  AggregatedBin
 } from '@kepler.gl/layers';
 import {MapState, MapControls, Viewport, SplitMap, SplitMapLayers} from '@kepler.gl/types';
 import {
@@ -465,9 +466,13 @@ export default function MapContainerFactory(
       this.props.visStateActions.onLayerHover(info, this.props.index);
     };
 
-    _onLayerSetDomain = (idx: number, colorDomain: {domain: VisualChannelDomain}) => {
+    _onLayerSetDomain = (
+      idx: number,
+      value: {domain: VisualChannelDomain; aggregatedBins: AggregatedBin[]}
+    ) => {
       this.props.visStateActions.layerConfigChange(this.props.visState.layers[idx], {
-        colorDomain: colorDomain.domain
+        colorDomain: value.domain,
+        aggregatedBins: value.aggregatedBins
       } as Partial<LayerBaseConfig>);
     };
 

diff --git a/src/components/src/side-panel/layer-panel/aggr-scale-selector.tsx b/src/components/src/side-panel/layer-panel/aggr-scale-selector.tsx
@@ -6,12 +6,13 @@ import DimensionScaleSelectorFactory from './dimension-scale-selector';
 
 AggrScaleSelectorFactory.deps = [DimensionScaleSelectorFactory];
 export function AggrScaleSelectorFactory(DimensionScaleSelector) {
-  const AggrScaleSelector = ({channel, layer, onChange, setColorUI, label}) => {
+  const AggrScaleSelector = ({channel, dataset, layer, onChange, setColorUI, label}) => {
     const {key} = channel;
     const scaleOptions = layer.getScaleOptions(key);
 
     return Array.isArray(scaleOptions) && scaleOptions.length > 1 ? (
       <DimensionScaleSelector
+        dataset={dataset}
         layer={layer}
         channel={channel}
         label={label || `${key} Scale`}

diff --git a/src/components/src/side-panel/layer-panel/color-breaks-panel.tsx b/src/components/src/side-panel/layer-panel/color-breaks-panel.tsx
@@ -3,7 +3,7 @@
 
 import {SCALE_TYPES} from '@kepler.gl/constants';
 import {KeplerTable} from '@kepler.gl/table';
-import {ColorUI, Field} from '@kepler.gl/types';
+import {Bin, ColorUI, Field} from '@kepler.gl/types';
 import {
   ColorBreak,
   ColorBreakOrdinal,
@@ -83,6 +83,10 @@ export type ColorBreaksPanelProps = {
   dataset: KeplerTable | undefined;
   colorField: Field;
   isCustomBreaks: boolean;
+  allBins: Bin[];
+  filteredBins: Bin[];
+  isFiltered: boolean;
+  histogramDomain: number[];
   setColorUI: SetColorUIFunc;
   onScaleChange: (v: string, visConfg?: Record<string, any>) => void;
   onApply: (e: React.MouseEvent) => void;
@@ -101,6 +105,10 @@ function ColorBreaksPanelFactory(
     dataset,
     colorField,
     isCustomBreaks,
+    allBins,
+    filteredBins,
+    isFiltered,
+    histogramDomain,
     setColorUI,
     onScaleChange,
     onApply,
@@ -168,6 +176,10 @@ function ColorBreaksPanelFactory(
             colorField={colorField}
             dataset={dataset}
             colorBreaks={currentBreaks}
+            allBins={allBins}
+            filteredBins={filteredBins}
+            isFiltered={isFiltered}
+            histogramDomain={histogramDomain}
             onChangedUpdater={onColumnStatsChartChanged}
           />
         ) : null}

diff --git a/src/components/src/side-panel/layer-panel/color-scale-selector.tsx b/src/components/src/side-panel/layer-panel/color-scale-selector.tsx
@@ -4,14 +4,17 @@
 import React, {useCallback, useMemo, useState} from 'react';
 import styled from 'styled-components';
 
-import {SCALE_TYPES} from '@kepler.gl/constants';
-import {Layer, VisualChannelDomain} from '@kepler.gl/layers';
+import {ALL_FIELD_TYPES, SCALE_TYPES} from '@kepler.gl/constants';
+import {AggregatedBin, Layer, VisualChannelDomain} from '@kepler.gl/layers';
 import {KeplerTable} from '@kepler.gl/table';
 import {ColorRange, ColorUI, Field, NestedPartial} from '@kepler.gl/types';
 import {
   colorBreaksToColorMap,
   getLayerColorScale,
   getLegendOfScale,
+  histogramFromValues,
+  histogramFromThreshold,
+  getHistogramDomain,
   hasColorMap
 } from '@kepler.gl/utils';
 
@@ -25,6 +28,8 @@ import Typeahead from '../../common/item-selector/typeahead';
 
 type TippyInstance = any; // 'tippy-js'
 
+const HISTOGRAM_BINS = 30;
+
 export type ScaleOption = {
   label: string;
   value: string;
@@ -50,6 +55,7 @@ export type ColorScaleSelectorProps = {
   searchable: boolean;
   displayOption: string;
   getOptionValue: string;
+  aggregatedBins?: AggregatedBin[];
 };
 
 const DropdownPropContext = React.createContext({});
@@ -126,6 +132,7 @@ function ColorScaleSelectorFactory(
     onSelect,
     scaleType,
     domain,
+    aggregatedBins,
     range,
     setColorUI,
     colorUIConfig,
@@ -151,10 +158,62 @@ function ColorScaleSelectorFactory(
 
     const colorBreaks = useMemo(
       () =>
-        colorScale ? getLegendOfScale({scale: colorScale, scaleType, fieldType: field.type}) : null,
-      [colorScale, scaleType, field.type]
+        colorScale
+          ? getLegendOfScale({
+              scale: colorScale,
+              scaleType,
+              fieldType: field?.type ?? ALL_FIELD_TYPES.real
+            })
+          : null,
+      [colorScale, scaleType, field?.type]
     );
 
+    const columnStats = field?.filterProps?.columnStats;
+
+    const fieldValueAccessor = useMemo(() => {
+      return field
+        ? idx => dataset.getValue(field.name, idx)
+        : idx => dataset.dataContainer.rowAsArray(idx);
+    }, [dataset, field]);
+
+    // aggregatedBins should be the raw data
+    const allBins = useMemo(() => {
+      if (aggregatedBins) {
+        return histogramFromValues(
+          Object.values(aggregatedBins).map(bin => bin.i),
+          HISTOGRAM_BINS,
+          idx => aggregatedBins[idx].value
+        );
+      }
+      return columnStats?.bins
+        ? columnStats?.bins
+        : histogramFromValues(dataset.allIndexes, HISTOGRAM_BINS, fieldValueAccessor);
+    }, [aggregatedBins, columnStats, dataset, fieldValueAccessor]);
+
+    const histogramDomain = useMemo(() => {
+      return getHistogramDomain({aggregatedBins, columnStats, dataset, fieldValueAccessor});
+    }, [dataset, fieldValueAccessor, aggregatedBins, columnStats]);
+
+    const isFiltered = aggregatedBins
+      ? false
+      : dataset.filteredIndexForDomain.length !== dataset.allIndexes.length;
+
+    // get filteredBins (not apply to aggregate layer)
+    const filteredBins = useMemo(() => {
+      if (!isFiltered) {
+        return allBins;
+      }
+      // get threholds
+      const filterEmptyBins = false;
+      const threholds = allBins.map(b => b.x0);
+      return histogramFromThreshold(
+        threholds,
+        dataset.filteredIndexForDomain,
+        fieldValueAccessor,
+        filterEmptyBins
+      );
+    }, [dataset, fieldValueAccessor, allBins, isFiltered]);
+
     const onSelectScale = useCallback(
       val => {
         // highlight selected option
@@ -221,6 +280,10 @@ function ColorScaleSelectorFactory(
           colorUIConfig,
           colorBreaks,
           isCustomBreaks,
+          allBins,
+          filteredBins,
+          isFiltered,
+          histogramDomain,
           onScaleChange: onSelect,
           onApply,
           onCancel

diff --git a/src/components/src/side-panel/layer-panel/dimension-scale-selector.tsx b/src/components/src/side-panel/layer-panel/dimension-scale-selector.tsx
@@ -51,7 +51,9 @@ function DimensionScaleSelectorFactory(
       value: op
     }));
     const disabled = scaleOptions.length < 2;
-    const isColorScale = channelScaleType === CHANNEL_SCALES.color;
+    const isColorScale =
+      channelScaleType === CHANNEL_SCALES.color ||
+      (layer.config.aggregatedBins && channelScaleType === CHANNEL_SCALES.colorAggr);
 
     const onSelect = useCallback(
       (val, newRange) => onChange({[scale]: val}, key, newRange ? {[range]: newRange} : undefined),
@@ -87,6 +89,7 @@ function DimensionScaleSelectorFactory(
             onSelect={onSelect}
             scaleType={scaleType}
             domain={layer.config[domain]}
+            aggregatedBins={layer.config.aggregatedBins}
             range={layer.config.visConfig[range]}
             setColorUI={_setColorUI}
             colorUIConfig={layer.config.colorUI?.[range]}