feat: globally switch between continuous and discrete coloring of ints

src/components/AppBar.tsx

@@ -19,7 +19,6 @@ import MainWalkthrough, {
     Handle as MainWalkthroughRef,
 } from './walkthrough/MainWalkthrough';
 import { useColors } from '../stores/colors';
-import type { ColorsState } from '../stores/colors';
 import ColorPaletteSelect from './ui/ColorPaletteSelect';
 import { categoricalPalettes, continuousPalettes } from '../palettes';
 
@@ -137,16 +136,8 @@ const HelpMenu = (): JSX.Element => {
     );
 };
 
-const useRobustColorScalesSelector = (c: ColorsState) => ({
-    useRobustColorScales: c.useRobustColorScales,
-    setUseRobustColorScales: c.setUseRobustColorScales,
-});
-
 const ColorMenu = () => {
     const colors = useColors();
-    const { useRobustColorScales, setUseRobustColorScales } = useColors(
-        useRobustColorScalesSelector
-    );
 
     const content = (
         <div tw="flex flex-col w-72 pb-1">
@@ -167,13 +158,27 @@ const ColorMenu = () => {
                         onChangeColorPalette={colors.setCategoricalPalette}
                     />
                 </Menu.Item>
-                <Menu.Title>Robust Coloring</Menu.Title>
+
+                <Menu.Title>Continuous Ints</Menu.Title>
                 <Menu.Switch
-                    value={useRobustColorScales}
-                    onChange={setUseRobustColorScales}
+                    value={colors.continuousInts}
+                    onChange={colors.setContinuousInts}
                 >
                     Enable
                 </Menu.Switch>
+
+                <Menu.Title>Continuous Categories</Menu.Title>
+                <Menu.Switch
+                    value={colors.continuousCategories}
+                    onChange={colors.setContinuousCategories}
+                >
+                    Enable
+                </Menu.Switch>
+
+                <Menu.Title>Robust Coloring</Menu.Title>
+                <Menu.Switch value={colors.robust} onChange={colors.setRobust}>
+                    Enable
+                </Menu.Switch>
             </Menu>
         </div>
     );

src/components/ScalarValue.tsx

@@ -41,7 +41,7 @@ const ScalarValue: FunctionComponent<Props> = ({
     // eslint-disable-next-line react-hooks/exhaustive-deps
     const colorTransferFunctionSelector = useCallback(
         (d: Dataset) =>
-            d.colorTransferFunctions[column.key]?.[filtered ? 'filtered' : 'full'][0],
+            d.colorTransferFunctions[column.key]?.[filtered ? 'filtered' : 'full'],
         [column.key, filtered]
     );
 

src/hooks/useColorTransferFunction.ts

@@ -109,36 +109,40 @@ export const createConstantTransferFunction = (
     return tf as ConstantTransferFunction;
 };
 
-const createColorTransferFunction = (
+export const createColorTransferFunction = (
     data: ColumnData | undefined,
     dType: DataType | undefined,
+    robust = false,
+    continuousInts = false,
+    continuousCategories = false,
     classBreaks?: number[]
 ): TransferFunction => {
-    const robustColoring = useColors.getState().useRobustColorScales;
-
     if (dType === undefined) return createConstantTransferFunction(unknownDataType);
     if (data === undefined) return createConstantTransferFunction(dType);
 
-    if (['int', 'bool', 'Category', 'str'].includes(dType.kind)) {
-        const uniqValues = _.uniq(data);
+    if (['bool', 'str'].includes(dType.kind)) {
+        return createCategoricalTransferFunction(_.uniq(data), dType);
+    }
 
+    if (dType.kind === 'int' && !continuousInts) {
+        const uniqValues = _.uniq(data);
         const tooManyInts =
             dType.kind === 'int' && uniqValues.length > MAX_VALUES_FOR_INT_CATEGORY;
 
         if (!tooManyInts) {
-            const transferFunction = createCategoricalTransferFunction(
-                uniqValues,
-                dType
-            );
-            return transferFunction;
+            return createCategoricalTransferFunction(uniqValues, dType);
         }
     }
 
-    if (['int', 'float'].includes(dType.kind)) {
+    if (dType.kind === 'Category' && !continuousCategories) {
+        return createCategoricalTransferFunction(_.uniq(data), dType);
+    }
+
+    if (['int', 'float', 'Category'].includes(dType.kind)) {
         const stats = makeStats(dType, data);
         return createContinuousTransferFunction(
-            (robustColoring ? stats?.p5 : stats?.min) ?? 0,
-            (robustColoring ? stats?.p95 : stats?.max) ?? 1,
+            (robust ? stats?.p5 : stats?.min) ?? 0,
+            (robust ? stats?.p95 : stats?.max) ?? 1,
             dType,
             classBreaks
         );
@@ -148,7 +152,19 @@ const createColorTransferFunction = (
 };
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
-export const useColorTransferFunction = (data: any[], dtype: DataType) =>
-    useMemo(() => createColorTransferFunction(data, dtype), [dtype, data]);
+export const useColorTransferFunction = (data: any[], dtype: DataType) => {
+    const colors = useColors();
+    return useMemo(
+        () =>
+            createColorTransferFunction(
+                data,
+                dtype,
+                colors.robust,
+                colors.continuousInts,
+                colors.continuousCategories
+            ),
+        [dtype, data, colors.robust, colors.continuousInts, colors.continuousCategories]
+    );
+};
 
 export default useColorTransferFunction;

src/stores/colors.ts

@@ -17,11 +17,15 @@ export interface ColorsState {
     constantPalette: ConstantPalette;
     categoricalPalette: CategoricalPalette;
     continuousPalette: ContinuousPalette;
-    useRobustColorScales: boolean;
+    robust: boolean;
+    continuousInts: boolean;
+    continuousCategories: boolean;
     setConstantPalette: (palette?: ConstantPalette) => void;
     setCategoricalPalette: (palette?: CategoricalPalette) => void;
     setContinuousPalette: (palette?: ContinuousPalette) => void;
-    setUseRobustColorScales: (useRobust: boolean) => void;
+    setRobust: (robust: boolean) => void;
+    setContinuousInts: (continuous: boolean) => void;
+    setContinuousCategories: (continuous: boolean) => void;
 }
 
 export const useColors = create<ColorsState>()(
@@ -30,35 +34,26 @@ export const useColors = create<ColorsState>()(
             constantPalette: defaultConstantPalette,
             categoricalPalette: defaultCategoricalPalette,
             continuousPalette: defaultContinuousPalette,
-            useRobustColorScales: false,
+            robust: false,
+            continuousInts: false,
+            continuousCategories: false,
             setConstantPalette: (palette) => {
-                set((state) => {
-                    return {
-                        ...state,
-                        constantPalette: palette ?? defaultConstantPalette,
-                    };
-                });
+                set({ constantPalette: palette ?? defaultConstantPalette });
             },
             setCategoricalPalette: (palette) => {
-                set((state) => {
-                    return {
-                        ...state,
-                        categoricalPalette: palette ?? defaultCategoricalPalette,
-                    };
-                });
+                set({ categoricalPalette: palette ?? defaultCategoricalPalette });
             },
             setContinuousPalette: (palette) => {
-                set((state) => {
-                    return {
-                        ...state,
-                        continuousPalette: palette ?? defaultContinuousPalette,
-                    };
-                });
+                set({ continuousPalette: palette ?? defaultContinuousPalette });
+            },
+            setRobust: (robust: boolean) => {
+                set({ robust });
+            },
+            setContinuousInts: (continuousInts: boolean) => {
+                set({ continuousInts });
             },
-            setUseRobustColorScales: (useRobustColorScales: boolean) => {
-                set((state) => {
-                    return { ...state, useRobustColorScales };
-                });
+            setContinuousCategories: (continuousCategories: boolean) => {
+                set({ continuousCategories });
             },
         }),
         {

src/stores/dataset/colorTransferFunctionFactory.tsx

@@ -1,83 +1,39 @@
 import {
-    DataType,
-    isCategorical,
-    isFloat,
-    isNumerical,
-    isScalar,
-} from '../../datatypes';
-import {
-    createCategoricalTransferFunction,
-    createConstantTransferFunction,
-    createContinuousTransferFunction,
+    createColorTransferFunction,
     TransferFunction,
 } from '../../hooks/useColorTransferFunction';
-import _ from 'lodash';
-import { useColors } from '../../stores/colors';
-import {
-    ColumnData,
-    DataColumn,
-    DataStatistics,
-    isCategoricalColumn,
-    isScalarColumn,
-    TableData,
-} from '../../types';
-import { Dataset } from './dataset';
-
-export const makeApplicableColorTransferFunctions = (
-    type: DataType,
-    data: ColumnData,
-    stats?: DataStatistics
-): TransferFunction[] => {
-    const transferFunctions: TransferFunction[] = [];
-
-    if ((isCategorical(type) || isScalar(type)) && !isFloat(type)) {
-        const uniqueValues = _.uniq(data);
-        const transFn = createCategoricalTransferFunction(uniqueValues, type);
-        transferFunctions.push(transFn);
-    }
-
-    if (isNumerical(type)) {
-        const useRobustColoring = useColors.getState().useRobustColorScales;
-
-        const min = useRobustColoring ? stats?.p5 : stats?.min;
-        const max = useRobustColoring ? stats?.p95 : stats?.max;
-
-        transferFunctions.push(
-            createContinuousTransferFunction(min || 0, max || 1, type)
-        );
-    }
-
-    transferFunctions.push(createConstantTransferFunction(type));
-
-    return transferFunctions;
-};
+import { DataColumn, TableData } from '../../types';
+import { useColors } from '../colors';
 
 type ColumnsTransferFunctions = Record<
     string,
-    { full: TransferFunction[]; filtered: TransferFunction[] }
+    { full: TransferFunction; filtered: TransferFunction }
 >;
 
 export const makeColumnsColorTransferFunctions = (
     columns: DataColumn[],
     data: TableData,
-    stats: Dataset['columnStats'],
-    filteredMask: boolean[]
+    filteredIndices: Int32Array
 ): ColumnsTransferFunctions => {
-    return columns
-        .filter((column) => isScalarColumn(column) || isCategoricalColumn(column))
-        .reduce((a, column) => {
-            a[column.key] = {
-                full: makeApplicableColorTransferFunctions(
-                    column.type,
-                    data[column.key],
-                    stats.full[column.key]
-                ),
-                filtered: makeApplicableColorTransferFunctions(
-                    column.type,
-                    data[column.key].filter((_, i) => filteredMask[i]),
-                    stats.filtered[column.key]
-                ),
-            };
-            return a;
-        }, {} as ColumnsTransferFunctions);
+    const colors = useColors.getState();
+
+    return columns.reduce((a, column) => {
+        a[column.key] = {
+            full: createColorTransferFunction(
+                data[column.key],
+                column.type,
+                colors.robust,
+                colors.continuousInts,
+                colors.continuousCategories
+            ),
+            filtered: createColorTransferFunction(
+                filteredIndices.map((i) => data[column.key][i]),
+                column.type,
+                colors.robust,
+                colors.continuousInts,
+                colors.continuousCategories
+            ),
+        };
+        return a;
+    }, {} as ColumnsTransferFunctions);
 };

src/stores/dataset/dataset.ts

@@ -44,8 +44,8 @@ export interface Dataset {
     colorTransferFunctions: Record<
         string,
         {
-            full: TransferFunction[];
-            filtered: TransferFunction[];
+            full: TransferFunction;
+            filtered: TransferFunction;
         }
     >;
     recomputeColorTransferFunctions: () => void;
@@ -430,8 +430,7 @@ export const useDataset = create(
                 const newTransferFunctions = makeColumnsColorTransferFunctions(
                     get().columns.filter(({ key }) => columnsToCompute.includes(key)),
                     get().columnData,
-                    get().columnStats,
-                    get().isIndexFiltered
+                    get().filteredIndices
                 );
 
                 set({
@@ -535,6 +534,8 @@ useDataset.subscribe(
     }
 );
 
+useColors.subscribe(useDataset.getState().recomputeColorTransferFunctions);
+
 useDataset.subscribe(
     (state) => state.selectedIndices,
     useDataset.getState().recomputeColumnRelevance

src/stores/dataset/statisticsFactory.tsx

@@ -1,4 +1,4 @@
-import { DataType, isNumerical } from '../../datatypes';
+import { DataType, isCategorical, isNumerical } from '../../datatypes';
 import { max, mean, min, quantile, standardDeviation } from 'simple-statistics';
 import {
     ColumnData,
@@ -13,7 +13,7 @@ export const makeStats = (
     data: ColumnData,
     mask?: boolean[]
 ): DataStatistics | undefined => {
-    if (!isNumerical(type)) {
+    if (!isNumerical(type) && !isCategorical(type)) {
         return;
     }
 

src/widgets/DataGrid/Cell/CategoricalCell.tsx

@@ -22,7 +22,7 @@ const CategoricalCell: FunctionComponent<Props> = ({ value, column }) => {
         (d: Dataset) =>
             d.colorTransferFunctions[column.key]?.[
                 tableView !== 'full' ? 'filtered' : 'full'
-            ][0],
+            ],
         [column.key, tableView]
     );
 

src/widgets/Histogram/Histogram.tsx

@@ -81,7 +81,7 @@ const Histogram: Widget = () => {
             stackByColumnKey
                 ? d.colorTransferFunctions[stackByColumnKey]?.[
                       filter ? 'filtered' : 'full'
-                  ][0]
+                  ]
                 : createConstantTransferFunction(),
         [filter, stackByColumnKey]
     );

src/widgets/ScatterplotView/ScatterplotView.tsx

@@ -139,9 +139,7 @@ const ScatterplotView: Widget = () => {
     const transferFunctionSelector = useCallback(
         (d: Dataset) =>
             colorByKey !== undefined && colorByKey.length > 0
-                ? d.colorTransferFunctions[colorByKey]?.[
-                      filter ? 'filtered' : 'full'
-                  ][0]
+                ? d.colorTransferFunctions[colorByKey]?.[filter ? 'filtered' : 'full']
                 : createConstantTransferFunction(),
         [colorByKey, filter]
     );