feat(embeddings): color by dimension

app/package.json

@@ -6,13 +6,14 @@
   "license": "None",
   "private": true,
   "dependencies": {
-    "@arizeai/components": "^0.9.15",
+    "@arizeai/components": "^0.9.19",
     "@arizeai/point-cloud": "^2.1.5",
     "@emotion/react": "^11.10.5",
     "@react-three/drei": "9.5.7",
     "@react-three/fiber": "8.0.12",
     "@types/react-table": "^7.7.14",
     "@types/recharts": "^1.8.24",
+    "d3-scale-chromatic": "^3.0.0",
     "d3-time-format": "^4.1.0",
     "date-fns": "^2.29.3",
     "normalize.css": "^8.0.1",
@@ -32,6 +33,7 @@
     "zustand": "^4.3.5"
   },
   "devDependencies": {
+    "@types/d3-scale-chromatic": "^3.0.0",
     "@types/d3-time-format": "^4.0.0",
     "@types/react-dom": "^18.0.8",
     "@types/react-relay": "^14.1.2",
@@ -56,7 +58,7 @@
     "build:ts": "node ./esbuild.config.mjs",
     "build:relay": "relay-compiler",
     "watch": "./esbuild.config.mjs dev",
-    "dev": "npm run dev:server:mnist & npm run build:static && npm run watch",
+    "dev": "npm run dev:server:sentiment & npm run build:static && npm run watch",
     "dev:server:mnist": "python3 -m phoenix.server.main fixture fashion_mnist",
     "dev:server:sentiment": "python3 -m phoenix.server.main fixture sentiment_classification_language_drift",
     "dev:server:ner": "python3 -m phoenix.server.main fixture ner_token_drift",

app/src/components/Empty.tsx

@@ -0,0 +1,34 @@
+import React from "react";
+import { css } from "@emotion/react";
+
+import { EmptyGraphic, EmptyGraphicProps, Text } from "@arizeai/components";
+
+interface EmptyProps extends EmptyGraphicProps {
+  message?: string;
+}
+export function Empty(props: EmptyProps) {
+  const { message, ...graphicsProps } = props;
+  return (
+    <div
+      css={css`
+        width: 100%;
+        display: flex;
+        justify-content: center;
+      `}
+    >
+      <div
+        css={(theme) =>
+          css`
+            margin: ${theme.spacing.margin24}px;
+            display: flex;
+            flex-direction: column;
+            align-items: center;
+          `
+        }
+      >
+        <EmptyGraphic {...graphicsProps} />
+        {message && <Text>{message}</Text>}
+      </div>
+    </div>
+  );
+}

app/src/components/Loading.tsx

@@ -1,20 +1,24 @@
 import React from "react";
 import { css } from "@emotion/react";
 
-import { ProgressCircle } from "@arizeai/components";
+import { ProgressCircle, Text } from "@arizeai/components";
 
-export const Loading = () => {
+type LoadingProps = { message?: string };
+export const Loading = ({ message }: LoadingProps) => {
   return (
     <div
       css={css`
         display: flex;
+        flex-direction: column;
         justify-content: center;
         align-items: center;
         width: 100%;
         height: 100%;
+        gap: var(--px-spacing-med);
       `}
     >
       <ProgressCircle isIndeterminate aria-label="loading" />
+      {message != null ? <Text>{message}</Text> : null}
     </div>
   );
 };

app/src/components/form/DimensionPicker.tsx

@@ -0,0 +1,159 @@
+import React, { startTransition, useEffect, useState } from "react";
+import { fetchQuery, graphql } from "react-relay";
+import { css } from "@emotion/react";
+
+import {
+  Item,
+  Label,
+  LabelProps,
+  Picker,
+  PickerProps,
+} from "@arizeai/components";
+
+import RelayEnvironment from "@phoenix/RelayEnvironment";
+import { Dimension } from "@phoenix/types";
+import { assertUnreachable } from "@phoenix/typeUtils";
+
+import { DimensionPickerQuery } from "./__generated__/DimensionPickerQuery.graphql";
+
+type DimensionPickerProps<T> = Omit<
+  PickerProps<T>,
+  "onSelectionChange" | "children"
+> & {
+  selectedDimension: Dimension | null;
+  onChange: (dimension: Dimension) => void;
+  dimensions: Dimension[];
+  /**
+   * Boolean flag to indicate if the picker is loading data
+   * @default false
+   */
+  isLoading?: boolean;
+};
+
+function DimensionTypeLabel(props: { type: Dimension["type"] }) {
+  const { type } = props;
+  let labelColor: LabelProps["color"] = "gray";
+  let text = "";
+  switch (type) {
+    case "feature":
+      labelColor = "blue";
+      text = "FEA";
+      break;
+    case "tag":
+      labelColor = "purple";
+      text = "TAG";
+      break;
+    case "prediction":
+      labelColor = "white";
+      text = "PRE";
+      break;
+    case "actual":
+      labelColor = "orange";
+      text = "ACT";
+      break;
+    default:
+      assertUnreachable(type);
+  }
+  return (
+    <Label color={labelColor} aria-label={type} title="type">
+      {text}
+    </Label>
+  );
+}
+
+export function DimensionPicker<T>(props: DimensionPickerProps<T>) {
+  const { selectedDimension, dimensions, onChange, isLoading, ...restProps } =
+    props;
+  return (
+    <Picker
+      {...restProps}
+      defaultSelectedKey={
+        selectedDimension ? selectedDimension.name : undefined
+      }
+      aria-label="Select a dimension"
+      onSelectionChange={(key) => {
+        // Find the dimension in the list
+        const dimension = dimensions.find((d) => d.name === key);
+        if (dimension) {
+          startTransition(() => onChange(dimension));
+        }
+      }}
+      label="Dimension"
+      isDisabled={isLoading}
+      placeholder={isLoading ? "Loading..." : "Select a dimension"}
+    >
+      {dimensions.map((dimension) => (
+        <Item key={dimension.name}>
+          <div
+            css={css`
+              .ac-label {
+                margin-right: var(--px-spacing-med);
+              }
+            `}
+          >
+            <DimensionTypeLabel type={dimension.type} />
+            {dimension.name}
+          </div>
+        </Item>
+      ))}
+    </Picker>
+  );
+}
+
+type ConnectedDimensionPickerProps<T> = Omit<
+  DimensionPickerProps<T>,
+  "dimensions"
+>;
+
+export function ConnectedDimensionPicker<T>(
+  props: ConnectedDimensionPickerProps<T>
+) {
+  const [dimensions, setDimensions] = useState<Dimension[]>([]);
+  const [isLoading, setIsLoading] = useState<boolean>(true);
+  const { selectedDimension, onChange, ...restProps } = props;
+
+  // Async load the dimensions
+  useEffect(() => {
+    fetchQuery<DimensionPickerQuery>(
+      RelayEnvironment,
+      graphql`
+        query DimensionPickerQuery {
+          model {
+            dimensions {
+              edges {
+                node {
+                  id
+                  name
+                  type
+                  dataType
+                }
+              }
+            }
+          }
+        }
+      `,
+      {},
+      {
+        fetchPolicy: "store-or-network",
+      }
+    )
+      .toPromise()
+      .then((data) => {
+        const dims: Dimension[] =
+          data?.model.dimensions.edges.map((edge) => edge.node) ?? [];
+        setDimensions(dims);
+        setIsLoading(false);
+      });
+  }, []);
+
+  return (
+    <DimensionPicker
+      {...restProps}
+      onChange={onChange}
+      dimensions={dimensions}
+      label="Dimension"
+      selectedDimension={selectedDimension}
+      isLoading={isLoading}
+    />
+  );
+}