feat: cluster display

app/package.json

@@ -6,8 +6,8 @@
   "license": "None",
   "private": true,
   "dependencies": {
-    "@arizeai/components": "^0.9.6",
-    "@arizeai/point-cloud": "^2.1.1",
+    "@arizeai/components": "^0.9.10",
+    "@arizeai/point-cloud": "^2.1.5",
     "@emotion/react": "^11.10.5",
     "@react-three/drei": "9.5.7",
     "@react-three/fiber": "8.0.12",
@@ -38,7 +38,7 @@
     "lint-staged": "^13.1.0",
     "prettier": "2.8.1",
     "relay-compiler": "^14.1.0",
-    "typescript": "^4.8.4"
+    "typescript": "4.8"
   },
   "scripts": {
     "prebuild": "rm -rf ../src/phoenix/server/static",

app/src/GlobalStyles.tsx

@@ -16,6 +16,21 @@ export function GlobalStyles() {
             height: 100vh;
           }
         }
+
+        /* Remove list styling */
+        ul {
+          display: block;
+          list-style-type: none;
+          margin-block-start: none;
+          margin-block-end: 0;
+          padding-inline-start: 0;
+        }
+
+        :root {
+          --px-primary-color: #9efcfd;
+          --px-primary-color--transparent: rgb(158, 252, 253, 0.2);
+          --px-reference-color: #baa1f9;
+        }
       `}
     />
   );

app/src/components/LoadingMask.tsx

@@ -0,0 +1,23 @@
+import React from "react";
+import { ProgressCircle } from "@arizeai/components";
+import { css } from "@emotion/react";
+
+export const LoadingMask = () => {
+  return (
+    <div
+      css={css`
+        position: absolute;
+        top: 0;
+        left: 0;
+        right: 0;
+        bottom: 0;
+        display: flex;
+        justify-content: center;
+        align-items: center;
+        background-color: rgba(0, 0, 0, 0.2);
+      `}
+    >
+      <ProgressCircle isIndeterminate />
+    </div>
+  );
+};

app/src/components/canvas/PointCloud.tsx

@@ -1,26 +1,31 @@
-import React, { useMemo, useState } from "react";
+import React, { ReactNode, useMemo, useState } from "react";
 import {
   ThreeDimensionalCanvas,
   ThreeDimensionalControls,
   getThreeDimensionalBounds,
   ThreeDimensionalBounds,
   LassoSelect,
   ColorSchemes,
+  Cluster,
 } from "@arizeai/point-cloud";
-import { ErrorBoundary } from "../ErrorBoundary";
 import { theme } from "@arizeai/components";
 import { css } from "@emotion/react";
 import { ControlPanel } from "./ControlPanel";
 import { ColoringStrategyPicker } from "./ColoringStrategyPicker";
 import { CanvasMode, CanvasModeRadioGroup } from "./CanvasModeRadioGroup";
 import { PointCloudPoints } from "./PointCloudPoints";
 import { ThreeDimensionalPointItem } from "./types";
-import { ColoringStrategy } from "./types";
+import { ColoringStrategy, ClusterItem } from "./types";
 import { createColorFn } from "./coloring";
 
 export type PointCloudProps = {
   primaryData: ThreeDimensionalPointItem[];
   referenceData: ThreeDimensionalPointItem[] | null;
+  clusters: readonly ClusterItem[];
+  /**
+   * The id of the cluster that is currently selected
+   */
+  selectedClusterId: string | null;
 };
 
 const CONTROL_PANEL_WIDTH = 300;
@@ -76,7 +81,25 @@ function SelectionControlPanel({ selectedIds }: { selectedIds: Set<string> }) {
   );
 }
 
-export function PointCloud({ primaryData, referenceData }: PointCloudProps) {
+function CanvasWrap({ children }: { children: ReactNode }) {
+  return (
+    <div
+      css={css`
+        flex: 1 1 auto;
+        position: relative;
+      `}
+    >
+      {children}
+    </div>
+  );
+}
+
+export function PointCloud({
+  primaryData,
+  referenceData,
+  clusters,
+  selectedClusterId,
+}: PointCloudProps) {
   // AutoRotate the canvas on initial load
   const [autoRotate, setAutoRotate] = useState<boolean>(true);
   const [coloringStrategy, onColoringStrategyChange] =
@@ -86,9 +109,19 @@ export function PointCloud({ primaryData, referenceData }: PointCloudProps) {
   const allPoints = useMemo(() => {
     return [...primaryData, ...(referenceData || [])];
   }, []);
+
+  // Keep a map of point id to position for fast lookup
+  const pointPositionsMap = useMemo(() => {
+    return allPoints.reduce((acc, point) => {
+      acc[(point.metaData as any).id as string] = point.position;
+      return acc;
+    }, {} as Record<string, ThreeDimensionalPointItem["position"]>);
+  }, [allPoints]);
+
   const bounds = useMemo(() => {
     return getThreeDimensionalBounds(allPoints.map((p) => p.position));
   }, []);
+
   const isMoveMode = canvasMode === CanvasMode.move;
 
   // Determine the color of the points based on the strategy
@@ -100,8 +133,21 @@ export function PointCloud({ primaryData, referenceData }: PointCloudProps) {
     coloringStrategy,
     defaultColor: DEFAULT_COLOR_SCHEME[1],
   });
+
+  // Interleave the cluster point locations with the cluster
+  const clustersWithData = useMemo(() => {
+    return clusters.map((cluster) => {
+      const { pointIds } = cluster;
+      return {
+        ...cluster,
+        data: pointIds.map((pointId) => ({
+          position: pointPositionsMap[pointId],
+        })),
+      };
+    });
+  }, [pointPositionsMap]);
   return (
-    <ErrorBoundary>
+    <CanvasWrap>
       {/* <SelectionControlPanel selectedIds={selectedIds} /> */}
       <CanvasTools
         coloringStrategy={coloringStrategy}
@@ -136,8 +182,16 @@ export function PointCloud({ primaryData, referenceData }: PointCloudProps) {
             primaryColor={primaryColor}
             referenceColor={referenceColor}
           />
+          {clustersWithData.map((cluster) => (
+            <Cluster
+              key={cluster.id}
+              data={cluster.data}
+              opacity={cluster.id === selectedClusterId ? 0.2 : 0}
+              wireframe
+            />
+          ))}
         </ThreeDimensionalBounds>
       </ThreeDimensionalCanvas>
-    </ErrorBoundary>
+    </CanvasWrap>
   );
 }

app/src/components/canvas/types.ts

@@ -5,6 +5,11 @@ export type ThreeDimensionalPointItem = {
   metaData: unknown;
 };
 
+export type ClusterItem = {
+  readonly id: string;
+  readonly pointIds: readonly string[];
+};
+
 export enum ColoringStrategy {
   dataset = "dataset",
   correctness = "correctness",