feat: allow adjusting embedding view colors and ordinal colors in theme

packages/viewer/src/charts/common/theme.ts

@@ -6,47 +6,6 @@ import * as d3 from "d3";
 import colors from "tailwindcss/colors";
 export { defaultCategoryColors } from "@embedding-atlas/component";
 
-const interpolates: Record<string, (v: number) => string> = {
-  blues: d3.interpolateBlues,
-  brbg: d3.interpolateBrBG,
-  bugn: d3.interpolateBuGn,
-  bupu: d3.interpolateBuPu,
-  cividis: d3.interpolateCividis,
-  cool: d3.interpolateCool,
-  cubehelixdefault: d3.interpolateCubehelixDefault,
-  gnbu: d3.interpolateGnBu,
-  greens: d3.interpolateGreens,
-  greys: d3.interpolateGreys,
-  inferno: d3.interpolateInferno,
-  magma: d3.interpolateMagma,
-  oranges: d3.interpolateOranges,
-  orrd: d3.interpolateOrRd,
-  piyg: d3.interpolatePiYG,
-  plasma: d3.interpolatePlasma,
-  prgn: d3.interpolatePRGn,
-  pubu: d3.interpolatePuBu,
-  pubugn: d3.interpolatePuBuGn,
-  puor: d3.interpolatePuOr,
-  purd: d3.interpolatePuRd,
-  purples: d3.interpolatePurples,
-  rainbow: d3.interpolateRainbow,
-  rdbu: d3.interpolateRdBu,
-  rdgy: d3.interpolateRdGy,
-  rdpu: d3.interpolateRdPu,
-  rdylbu: d3.interpolateRdYlBu,
-  rdylgn: d3.interpolateRdYlGn,
-  reds: d3.interpolateReds,
-  sinebow: d3.interpolateSinebow,
-  spectral: d3.interpolateSpectral,
-  turbo: d3.interpolateTurbo,
-  viridis: d3.interpolateViridis,
-  warm: d3.interpolateWarm,
-  ylgn: d3.interpolateYlGn,
-  ylgnbu: d3.interpolateYlGnBu,
-  ylorbr: d3.interpolateYlOrBr,
-  ylorrd: d3.interpolateYlOrRd,
-};
-
 export interface ChartTheme {
   scheme: "light" | "dark";
 
@@ -56,13 +15,25 @@ export interface ChartTheme {
   /** Category color scheme */
   categoryColors: string[] | ((count: number) => string[]);
 
+  /** Ordinal color scheme */
+  ordinalColors: string[] | ((count: number) => string[]);
+
+  /** Color for the '(other)' category */
+  otherColor: string;
+
+  /** Color for the '(null)' category */
+  nullColor: string;
+
   /** Mark color */
   markColor: string;
   markColorFade: string;
   markColorGray: string;
   markColorGrayFade: string;
   ruleColor: string;
 
+  /** Embedding view point / contour color when there is no color encoding */
+  embeddingColor: string;
+
   /** Grid color */
   gridColor: string;
 
@@ -97,10 +68,14 @@ const defaultTheme: { light: ChartTheme; dark: ChartTheme } = {
     scheme: "light",
     interpolate: "pubugn",
     categoryColors: defaultCategoryColors,
+    ordinalColors: defaultOrdinalColors,
+    otherColor: "#9eabc2",
+    nullColor: "#aaaaaa",
     markColor: "#3b82f6",
     markColorFade: "#dbeafe",
     markColorGray: adjustForGray("#3b82f6", 20),
     markColorGrayFade: adjustForGray("#dbeafe"),
+    embeddingColor: "#1f77b4",
     ruleColor: "#000",
     gridColor: colors.slate[300],
     labelColor: colors.slate[400],
@@ -115,10 +90,14 @@ const defaultTheme: { light: ChartTheme; dark: ChartTheme } = {
     scheme: "dark",
     interpolate: "inferno",
     categoryColors: defaultCategoryColors,
+    ordinalColors: defaultOrdinalColors,
+    otherColor: "#9eabc2",
+    nullColor: "#aaaaaa",
     markColor: "#3b82f6",
     markColorFade: "#3b4d7f",
     markColorGray: adjustForGray("#3b82f6", -20),
     markColorGrayFade: adjustForGray("#1f398a"),
+    embeddingColor: "#1f77b4",
     ruleColor: "#fff",
     gridColor: colors.slate[700],
     labelColor: colors.slate[500],
@@ -131,6 +110,47 @@ const defaultTheme: { light: ChartTheme; dark: ChartTheme } = {
   },
 };
 
+const interpolates: Record<string, (v: number) => string> = {
+  blues: d3.interpolateBlues,
+  brbg: d3.interpolateBrBG,
+  bugn: d3.interpolateBuGn,
+  bupu: d3.interpolateBuPu,
+  cividis: d3.interpolateCividis,
+  cool: d3.interpolateCool,
+  cubehelixdefault: d3.interpolateCubehelixDefault,
+  gnbu: d3.interpolateGnBu,
+  greens: d3.interpolateGreens,
+  greys: d3.interpolateGreys,
+  inferno: d3.interpolateInferno,
+  magma: d3.interpolateMagma,
+  oranges: d3.interpolateOranges,
+  orrd: d3.interpolateOrRd,
+  piyg: d3.interpolatePiYG,
+  plasma: d3.interpolatePlasma,
+  prgn: d3.interpolatePRGn,
+  pubu: d3.interpolatePuBu,
+  pubugn: d3.interpolatePuBuGn,
+  puor: d3.interpolatePuOr,
+  purd: d3.interpolatePuRd,
+  purples: d3.interpolatePurples,
+  rainbow: d3.interpolateRainbow,
+  rdbu: d3.interpolateRdBu,
+  rdgy: d3.interpolateRdGy,
+  rdpu: d3.interpolateRdPu,
+  rdylbu: d3.interpolateRdYlBu,
+  rdylgn: d3.interpolateRdYlGn,
+  reds: d3.interpolateReds,
+  sinebow: d3.interpolateSinebow,
+  spectral: d3.interpolateSpectral,
+  turbo: d3.interpolateTurbo,
+  viridis: d3.interpolateViridis,
+  warm: d3.interpolateWarm,
+  ylgn: d3.interpolateYlGn,
+  ylgnbu: d3.interpolateYlGnBu,
+  ylorbr: d3.interpolateYlOrBr,
+  ylorrd: d3.interpolateYlOrRd,
+};
+
 export function resolveChartTheme(scheme: "light" | "dark", config?: ChartThemeConfig): ChartTheme {
   let theme = defaultTheme[scheme];
   if (config != undefined) {

packages/viewer/src/charts/embedding/Embedding.svelte

@@ -54,6 +54,7 @@
   import { IconSettings } from "../../assets/icons.js";
   import { isolatedWritable } from "../../utils/store.js";
   import type { ChartViewProps, RowID } from "../chart.js";
+  import { resolveChartTheme } from "../common/theme.js";
   import { makeCategoryColumn } from "./category_column.js";
   import type { EmbeddingSpec } from "./types.js";
   import { interpolateViewport } from "./viewport_animation.js";
@@ -71,7 +72,9 @@
     onSpecChange,
   }: ChartViewProps<EmbeddingSpec, State> = $props();
 
-  let { colorScheme, columnStyles, searchResult } = context;
+  let { colorScheme, columnStyles, searchResult, theme: themeConfig } = context;
+  let theme = $derived(resolveChartTheme($colorScheme, $themeConfig));
+
   let highlightStore = isolatedWritable(context.highlight);
 
   let categoryColumn = $derived(spec.data.category);
@@ -85,7 +88,7 @@
   // Update the category mapping and legend.
   $effect.pre(() => {
     let promise = context.cache.value(`embedding/category/${categoryColumn}`, () =>
-      makeCategoryColumn(context.coordinator, context.table, categoryColumn),
+      makeCategoryColumn(context.coordinator, context.table, categoryColumn, theme),
     );
     promise.then((v) => {
       categoryLegend = v;
@@ -200,7 +203,7 @@
     y={spec.data.y}
     text={spec.data.text}
     category={categoryLegend?.indexColumn}
-    categoryColors={categoryLegend?.legend.map((x) => x.color)}
+    categoryColors={categoryLegend?.legend.map((x) => x.color) ?? [theme.embeddingColor]}
     config={{
       colorScheme: $colorScheme,
       ...context.embeddingViewConfig,

packages/viewer/src/charts/embedding/category_column.ts

@@ -1,13 +1,12 @@
 // Copyright (c) 2025 Apple Inc. Licensed under MIT License.
 
-import { defaultCategoryColors } from "@embedding-atlas/component";
 import type { Coordinator } from "@uwdata/mosaic-core";
 import * as SQL from "@uwdata/mosaic-sql";
 import * as d3 from "d3";
 
 import { distinctCount, jsTypeFromDBType } from "../../utils/database.js";
 import { inferBinning } from "../common/binning.js";
-import { defaultOrdinalColors } from "../common/theme.js";
+import { type ChartTheme } from "../common/theme.js";
 
 export interface EmbeddingLegend {
   indexColumn: string;
@@ -23,6 +22,7 @@ export async function makeCategoryColumn(
   coordinator: Coordinator,
   table: string,
   column: string | null | undefined,
+  theme: ChartTheme,
 ): Promise<EmbeddingLegend | null> {
   if (column == null) {
     return null;
@@ -33,13 +33,13 @@ export async function makeCategoryColumn(
   }
   let jsType = jsTypeFromDBType(desc.column_type);
   if (jsType == "string") {
-    return await makeDiscreteCategoryColumn(coordinator, table, column, 10);
+    return await makeDiscreteCategoryColumn(coordinator, table, column, 10, theme);
   } else if (jsType == "number") {
     let distinct = await distinctCount(coordinator, table, column);
     if (distinct <= 10) {
-      return await makeDiscreteCategoryColumn(coordinator, table, column, 10);
+      return await makeDiscreteCategoryColumn(coordinator, table, column, 10, theme);
     } else {
-      return await makeBinnedNumericColumn(coordinator, table, column);
+      return await makeBinnedNumericColumn(coordinator, table, column, theme);
     }
   }
   return null;
@@ -50,6 +50,7 @@ async function makeDiscreteCategoryColumn(
   table: string,
   column: string,
   maxCategories: number,
+  theme: ChartTheme,
 ): Promise<EmbeddingLegend> {
   let indexColumnName = `_ev_${column}_id`;
   let values = Array.from(
@@ -90,7 +91,7 @@ async function makeDiscreteCategoryColumn(
   let otherCount = countMap.get(otherIndex) ?? 0;
   let nullCount = countMap.get(nullIndex) ?? 0;
 
-  let colors = defaultCategoryColors(values.length);
+  let colors = resolveCategoryColors(theme, values.length);
 
   let legend: EmbeddingLegend["legend"] = values.map(({ value }, i) => ({
     label: value,
@@ -109,7 +110,7 @@ async function makeDiscreteCategoryColumn(
     ).get(0);
     legend.push({
       label: `(other ${otherCategoryCount.toLocaleString()})`,
-      color: "#9eabc2",
+      color: theme.otherColor,
       predicate:
         values.length > 0
           ? SQL.sql`${SQL.column(column)} IS NOT NULL AND ${SQL.column(column)}::TEXT NOT IN (${values.map((x) => SQL.literal(x.value)).join(",")})`
@@ -128,7 +129,7 @@ async function makeDiscreteCategoryColumn(
     }
     legend.push({
       label: "(null)",
-      color: "#aaaaaa",
+      color: theme.nullColor,
       predicate: SQL.isNull(SQL.column(column)),
       count: nullCount,
     });
@@ -144,6 +145,7 @@ async function makeBinnedNumericColumn(
   coordinator: Coordinator,
   table: string,
   column: string,
+  theme: ChartTheme,
 ): Promise<EmbeddingLegend> {
   let stats = (
     await coordinator.query(
@@ -207,7 +209,7 @@ async function makeBinnedNumericColumn(
   let fmt = d3.format(".6");
 
   if (minIndex != null && maxIndex != null) {
-    let colors = defaultOrdinalColors(maxIndex - minIndex + 1);
+    let colors = resolveOrdinalColors(theme, maxIndex - minIndex + 1);
     for (let index = minIndex; index <= maxIndex; index++) {
       let lowerBound = reverse(index * binning.binSize + binning.binStart);
       let upperBound = reverse((index + 1) * binning.binSize + binning.binStart);
@@ -229,7 +231,7 @@ async function makeBinnedNumericColumn(
       `);
     legend.push({
       label: "(null / nan / inf)",
-      color: "#aaaaaa",
+      color: theme.nullColor,
       predicate: SQL.isNull(binIndexExpr),
       count: index2Count.get(null) ?? 0,
     });
@@ -240,3 +242,29 @@ async function makeBinnedNumericColumn(
     legend: legend,
   };
 }
+
+function resolveCategoryColors(theme: ChartTheme, length: number): string[] {
+  if (typeof theme.categoryColors == "function") {
+    return theme.categoryColors(length);
+  } else {
+    let result: string[] = [];
+    for (let i = 0; i < length; i++) {
+      result.push(theme.categoryColors[i % theme.categoryColors.length]);
+    }
+    return result;
+  }
+}
+
+function resolveOrdinalColors(theme: ChartTheme, length: number): string[] {
+  if (typeof theme.ordinalColors == "function") {
+    return theme.ordinalColors(length);
+  } else {
+    if (length == theme.ordinalColors.length) {
+      return theme.ordinalColors.slice();
+    } else {
+      // Re-interpolate
+      let interp = d3.interpolateRgbBasis(theme.ordinalColors);
+      return Array.from({ length: length }).map((_, i) => interp(i / (length - 1)));
+    }
+  }
+}