Animate embedding switching via D3Map pipeline system

Reuse D3Map's existing PipelineSelector and animation system for h5ad
embedding switching instead of re-creating preloadedData on each change.
All embeddings are now parsed once and injected via preloadedPipelineData
prop, preserving painted groups and providing smooth 1s transitions.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: patcon

src/components/convo-explorer/App.h5ad.stories.tsx

@@ -3,20 +3,16 @@ import type { Meta, StoryObj } from '@storybook/react';
 import { App } from './App';
 import type { PreloadedData } from './App';
 import { loadH5adFile } from '../../lib/h5ad-loader';
-import type { H5adData } from '../../lib/h5ad-loader';
 
 /**
  * Wrapper component that provides a file picker for loading h5ad files,
  * parses them with h5wasm, and renders the App with preloaded data.
+ * Embedding switching is handled by D3Map's PipelineSelector with animation.
  */
 function H5adFileLoader() {
-  const [data, setData] = React.useState<H5adData | null>(null);
+  const [preloadedData, setPreloadedData] = React.useState<PreloadedData | null>(null);
   const [loading, setLoading] = React.useState(false);
   const [error, setError] = React.useState<string | null>(null);
-  const [selectedEmbedding, setSelectedEmbedding] = React.useState<string | null>(null);
-  // Store the file buffer so we can re-parse with a different embedding
-  // without needing the file input element (which is removed from DOM after loading)
-  const bufferRef = React.useRef<ArrayBuffer | null>(null);
 
   const handleFileChange = React.useCallback(async (event: React.ChangeEvent<HTMLInputElement>) => {
     const file = event.target.files?.[0];
@@ -27,14 +23,13 @@ function H5adFileLoader() {
 
     try {
       const buffer = await file.arrayBuffer();
-      bufferRef.current = buffer;
       const parsed = await loadH5adFile(buffer);
-      setData(parsed);
-      // Set the initially-selected embedding (first in preferred order)
-      const preferredOrder = ['X_localmap', 'X_umap', 'X_pacmap'];
-      const defaultEmbedding = preferredOrder.find(k => parsed.availableEmbeddings.includes(k))
-        ?? parsed.availableEmbeddings[0];
-      setSelectedEmbedding(defaultEmbedding);
+      setPreloadedData({
+        dataset: parsed.dataset,
+        statements: parsed.statements,
+        votesRows: parsed.votesRows,
+        pipelineData: parsed.allEmbeddings,
+      });
     } catch (err) {
       console.error('Failed to load h5ad file:', err);
       setError(err instanceof Error ? err.message : 'Failed to parse h5ad file');
@@ -43,25 +38,8 @@ function H5adFileLoader() {
     }
   }, []);
 
-  const handleEmbeddingChange = React.useCallback(async (newEmbedding: string) => {
-    if (!bufferRef.current) return;
-    setSelectedEmbedding(newEmbedding);
-    setLoading(true);
-    setError(null);
-
-    try {
-      const parsed = await loadH5adFile(bufferRef.current, newEmbedding);
-      setData(parsed);
-    } catch (err) {
-      console.error('Failed to reload with new embedding:', err);
-      setError(err instanceof Error ? err.message : 'Failed to reload embedding');
-    } finally {
-      setLoading(false);
-    }
-  }, []);
-
   // File picker screen
-  if (!data) {
+  if (!preloadedData) {
     return (
       <div className="h-screen w-screen flex items-center justify-center bg-background">
         <div className="flex flex-col items-center gap-4 max-w-md text-center p-8">
@@ -93,39 +71,7 @@ function H5adFileLoader() {
     );
   }
 
-  const preloadedData: PreloadedData = {
-    dataset: data.dataset,
-    statements: data.statements,
-    votesRows: data.votesRows,
-  };
-
-  return (
-    <div className="relative h-screen w-screen">
-      {/* Embedding selector overlay */}
-      {data.availableEmbeddings.length > 1 && (
-        <div className="absolute top-2 left-2 z-[100] pointer-events-auto">
-          <select
-            value={selectedEmbedding ?? ''}
-            onChange={(e) => handleEmbeddingChange(e.target.value)}
-            disabled={loading}
-            className="text-xs px-2 py-1 rounded border bg-background/90 backdrop-blur-sm shadow-sm"
-          >
-            {data.availableEmbeddings.map((key) => (
-              <option key={key} value={key}>
-                {key.replace(/^X_/, '')}
-              </option>
-            ))}
-          </select>
-        </div>
-      )}
-      {loading && (
-        <div className="absolute inset-0 z-[200] flex items-center justify-center bg-background/50">
-          <p className="text-sm text-muted-foreground">Switching embedding...</p>
-        </div>
-      )}
-      <App preloadedData={preloadedData} />
-    </div>
-  );
+  return <App preloadedData={preloadedData} />;
 }
 
 const meta: Meta = {

src/components/convo-explorer/App.tsx

@@ -35,6 +35,7 @@ export type PreloadedData = {
   dataset: [string, [number, number]][];
   statements: { statement_id: string; txt: string; moderated: number }[];
   votesRows: { participant_id: string; comment_id: string; vote: number }[];
+  pipelineData?: Record<string, [string, [number, number]][]>;
 };
 
 type AppProps = {
@@ -691,11 +692,12 @@ export const App: React.FC<AppProps> = ({ testAnimation = false, kedroBaseUrl, i
             flipX={currentDisplayState.flipX}
             flipY={currentDisplayState.flipY}
             colorsToFront={colorsToFront}
-            testAnimation={testAnimation}
+            testAnimation={testAnimation || !!preloadedData?.pipelineData}
             kedroBaseUrl={kedroBaseUrl}
             pipelineFilter={pipelineFilter}
             availablePipelines={kedroBaseUrl ? [] : undefined} // Will be populated by D3Map's usePipelineOptions
             onPipelineChange={handlePipelineChange}
+            preloadedPipelineData={preloadedData?.pipelineData}
           />
         </div>
       </div>

src/components/convo-explorer/D3Map.tsx

@@ -42,6 +42,8 @@ type D3MapProps = {
   availablePipelines?: Array<{id: string, name: string}>;
   /** Called when pipeline changes in the selector */
   onPipelineChange?: (pipelineId: string) => void;
+  /** Preloaded pipeline data (e.g. from h5ad file) — bypasses fetch-based loading */
+  preloadedPipelineData?: Record<string, [string, [number, number]][] | null>;
 };
 
 const PREFERRED_KEDRO_PIPELINE = 'mean_localmap_bestkmeans';
@@ -64,6 +66,7 @@ export const D3Map: React.FC<D3MapProps> = ({
   pipelineFilter,
   availablePipelines = [],
   onPipelineChange,
+  preloadedPipelineData,
 }) => {
   const svgRef = React.useRef<SVGSVGElement>(null);
   const containerRef = React.useRef<d3.Selection<SVGGElement, unknown, null, undefined> | null>(null);
@@ -103,8 +106,14 @@ export const D3Map: React.FC<D3MapProps> = ({
   );
   const kedroOptions = availablePipelines?.length ? availablePipelines : fetchedKedroOptions;
 
+  // Preloaded pipeline options derived from preloadedPipelineData keys
+  const preloadedPipelineOptions = React.useMemo(() => {
+    if (!preloadedPipelineData) return [];
+    return Object.keys(preloadedPipelineData).map(id => ({ id, name: id }));
+  }, [preloadedPipelineData]);
+
   // Current pipeline options based on mode
-  const currentPipelineOptions = isKedroMode ? kedroOptions : staticPipelines;
+  const currentPipelineOptions = preloadedPipelineData ? preloadedPipelineOptions : isKedroMode ? kedroOptions : staticPipelines;
 
   // Auto-cycling state
   const [isAutoCycling, setIsAutoCycling] = React.useState(false);
@@ -115,7 +124,12 @@ export const D3Map: React.FC<D3MapProps> = ({
   // Initialize selectedPipeline when pipeline options become available
   React.useEffect(() => {
     if (currentPipelineOptions.length > 0 && !selectedPipeline) {
-      if (isKedroMode) {
+      if (preloadedPipelineData) {
+        // For preloaded data, use preferred order: localmap > umap > pacmap > first key
+        const preferredOrder = ['localmap', 'umap', 'pacmap'];
+        const preferred = preferredOrder.find(id => id in preloadedPipelineData);
+        setSelectedPipeline(preferred || currentPipelineOptions[0].id);
+      } else if (isKedroMode) {
         // Prioritize preferred Kedro pipeline if available, otherwise use first
         const preferredPipeline = currentPipelineOptions.find(p => p.id === PREFERRED_KEDRO_PIPELINE);
         const defaultPipeline = preferredPipeline || currentPipelineOptions[0];
@@ -125,7 +139,7 @@ export const D3Map: React.FC<D3MapProps> = ({
         setSelectedPipeline('localmap');
       }
     }
-  }, [currentPipelineOptions, selectedPipeline, isKedroMode, testAnimation]);
+  }, [currentPipelineOptions, selectedPipeline, isKedroMode, testAnimation, preloadedPipelineData]);
 
   // Handle window resize to update radius (with throttling)
   React.useEffect(() => {
@@ -149,6 +163,12 @@ export const D3Map: React.FC<D3MapProps> = ({
   React.useEffect(() => {
     if (!testAnimation) return;
 
+    // If preloaded pipeline data is provided, use it directly
+    if (preloadedPipelineData) {
+      setPipelineData(preloadedPipelineData);
+      return;
+    }
+
     const loadProjections = async () => {
       try {
         if (isKedroMode && kedroOptions.length > 0) {
@@ -197,7 +217,7 @@ export const D3Map: React.FC<D3MapProps> = ({
     };
 
     loadProjections();
-  }, [testAnimation, isKedroMode, kedroOptions]);
+  }, [testAnimation, isKedroMode, kedroOptions, preloadedPipelineData]);
 
   // Calculate responsive base radius directly in JavaScript
   const BASE_RADIUS = React.useMemo(() => {

src/lib/h5ad-loader.ts

@@ -6,6 +6,8 @@ export type H5adData = {
   statements: { statement_id: string; txt: string; moderated: number }[];
   votesRows: { participant_id: string; comment_id: string; vote: number }[];
   availableEmbeddings: string[];
+  /** All 2D embeddings keyed by pipeline-style ID (X_ prefix stripped) */
+  allEmbeddings: Record<string, [string, [number, number]][]>;
 };
 
 /**
@@ -148,35 +150,47 @@ export async function loadH5adFile(
     if (!obsGroup) throw new Error('Missing /obs group');
     const obsNames = readIndex(obsGroup);
 
-    // --- Read embedding coordinates ---
+    // --- Read embedding coordinates for all available embeddings ---
     const obsmGroup = file.get('obsm') as Group;
     if (!obsmGroup) throw new Error('Missing /obsm group');
-    const embeddingDs = obsmGroup.get(selectedEmbedding) as Dataset | null;
-    if (!embeddingDs) throw new Error(`Embedding "${selectedEmbedding}" not found in /obsm`);
 
-    const shape = embeddingDs.shape;
-    if (!shape || shape.length !== 2) {
-      throw new Error(`Embedding "${selectedEmbedding}" has unexpected shape`);
-    }
-    const nObs = shape[0];
-    const nDims = shape[1];
-
-    // Read the flat typed array and reshape to [n_obs, 2]
-    const rawValue = embeddingDs.value;
-    let flatCoords: number[];
-    if (ArrayBuffer.isView(rawValue)) {
-      flatCoords = Array.from(rawValue as Float64Array | Float32Array);
-    } else if (Array.isArray(rawValue)) {
-      flatCoords = (rawValue as number[][]).flat();
-    } else {
-      throw new Error(`Unexpected embedding data format`);
+    const allEmbeddings: Record<string, [string, [number, number]][]> = {};
+    for (const embKey of availableEmbeddings) {
+      const ds = obsmGroup.get(embKey) as Dataset | null;
+      if (!ds) continue;
+
+      const shape = ds.shape;
+      if (!shape || shape.length !== 2) continue;
+      const nObs = shape[0];
+      const nDims = shape[1];
+
+      const rawValue = ds.value;
+      let flatCoords: number[];
+      if (ArrayBuffer.isView(rawValue)) {
+        flatCoords = Array.from(rawValue as Float64Array | Float32Array);
+      } else if (Array.isArray(rawValue)) {
+        flatCoords = (rawValue as number[][]).flat();
+      } else {
+        continue;
+      }
+
+      const embDataset: [string, [number, number]][] = [];
+      for (let i = 0; i < nObs; i++) {
+        const x = flatCoords[i * nDims];
+        const y = flatCoords[i * nDims + 1];
+        embDataset.push([obsNames[i], [x, y]]);
+      }
+
+      // Strip X_ prefix for pipeline-style IDs
+      const pipelineId = embKey.replace(/^X_/, '');
+      allEmbeddings[pipelineId] = embDataset;
     }
 
-    const dataset: [string, [number, number]][] = [];
-    for (let i = 0; i < nObs; i++) {
-      const x = flatCoords[i * nDims];
-      const y = flatCoords[i * nDims + 1];
-      dataset.push([obsNames[i], [x, y]]);
+    // Use the selected embedding as the default dataset
+    const pipelineId = selectedEmbedding.replace(/^X_/, '');
+    const dataset = allEmbeddings[pipelineId];
+    if (!dataset) {
+      throw new Error(`Embedding "${selectedEmbedding}" could not be read from /obsm`);
     }
 
     // --- Read statements (var) ---
@@ -224,7 +238,7 @@ export async function loadH5adFile(
     // --- Read votes from uns/votes ---
     const votesRows = readVotes(file);
 
-    return { dataset, statements, votesRows, availableEmbeddings };
+    return { dataset, statements, votesRows, availableEmbeddings, allEmbeddings };
   } finally {
     if (file) {
       file.close();