@vitessce/scatterplot-embedding 4.0.4 → 4.0.5

package/src/EmbeddingScatterplotSubscriber.js

@@ -1,5 +1,5 @@
 import React, {
-  useState, useEffect, useCallback, useMemo,
+  useState, useEffect, useCallback, useMemo, useDeferredValue,
 } from 'react';
 import { extent, quantileSorted } from 'd3-array';
 import { isEqual } from 'lodash-es';
@@ -28,8 +28,8 @@ import {
   useCoordinationScopes,
 } from '@vitessce/vit-s';
 import {
-  setObsSelection, mergeObsSets, getCellSetPolygons, getCellColors,
-  stratifyArrays,
+  setObsSelection, mergeObsSets, getCellSetPolygons, treeToColorIndicesArray,
+  colorIndicesFromCodes, getObsIndexMap, stratifyArrays,
 } from '@vitessce/sets-utils';
 import { pluralize as plur, commaNumber, aggregateFeatureArrays } from '@vitessce/utils';
 import {
@@ -187,7 +187,8 @@ export function EmbeddingScatterplotSubscriber(props) {
   );
   const cellsCount = obsEmbeddingIndex?.length || 0;
   const [
-    { obsSets: cellSets, obsSetsMembership }, obsSetsStatus, obsSetsUrls, obsSetsError,
+    { obsSets: cellSets, obsSetsMembership, obsSetsColumns },
+    obsSetsStatus, obsSetsUrls, obsSetsError,
   ] = useObsSetsData(
     loaders, dataset, false,
     { setObsSetSelection: setCellSetSelection, setObsSetColor: setCellSetColor },
@@ -261,6 +262,7 @@ export function EmbeddingScatterplotSubscriber(props) {
   ]);
 
   const [dynamicCellRadius, setDynamicCellRadius] = useState(cellRadiusFixed);
+  const [isSelectionPending, setIsSelectionPending] = useState(false);
   const [dynamicCellOpacity, setDynamicCellOpacity] = useState(cellOpacityFixed);
 
   const [originalViewState, setOriginalViewState] = useState(null);
@@ -278,14 +280,38 @@ export function EmbeddingScatterplotSubscriber(props) {
   }, [additionalCellSets, cellSetColor, setCellColorEncoding,
     setAdditionalCellSets, setCellSetColor, setCellSetSelection]);
 
-  const cellColors = useMemo(() => getCellColors({
-    cellSets: mergedCellSets,
-    cellSetSelection,
-    cellSetColor,
-    obsIndex: matrixObsIndex,
-    theme,
-  }), [mergedCellSets, theme,
-    cellSetSelection, cellSetColor, matrixObsIndex]);
+  // A set selection or color change re-encodes every observation in the memos
+  // below, which at atlas scale takes longer than a frame. Deferring these values
+  // lets React commit the urgent update first — the sets manager checkbox that
+  // initiated the change paints immediately — and re-render this view afterwards.
+  // The urgent render sees the previous values, so the memos below keep their
+  // cached results and cost nothing in that first commit.
+  const deferredCellSetSelection = useDeferredValue(cellSetSelection);
+  const deferredCellSetColor = useDeferredValue(cellSetColor);
+  const deferredSampleSetSelection = useDeferredValue(sampleSetSelection);
+
+  // Positional rather than keyed by observation ID: at atlas scale an ID-keyed color
+  // Map costs one string hash lookup per point per render, plus a per-observation
+  // color array whenever the selected sets carry confidence scores.
+  const obsColorIndices = useMemo(() => (
+    // When the loader provides raw categorical codes over this same observation
+    // axis (checked by reference), the encoding is one typed-array pass with no
+    // tree walk. colorIndicesFromCodes returns null for selections it cannot
+    // resolve (e.g. user-defined lasso selections), falling back to the tree.
+    (obsSetsColumns && obsSetsColumns.obsIndex === obsEmbeddingIndex
+      && colorIndicesFromCodes({
+        columns: obsSetsColumns.columns,
+        obsIndex: obsEmbeddingIndex,
+        selectedNamePaths: deferredCellSetSelection,
+        cellSetColor: deferredCellSetColor,
+        theme,
+      }))
+    || treeToColorIndicesArray(
+      mergedCellSets, deferredCellSetSelection, deferredCellSetColor,
+      obsEmbeddingIndex, theme,
+    )
+  ), [mergedCellSets, deferredCellSetSelection, deferredCellSetColor,
+    obsEmbeddingIndex, theme, obsSetsColumns]);
 
   // cellSetPolygonCache is an array of tuples like [(key0, val0), (key1, val1), ...],
   // where the keys are cellSetSelection arrays.
@@ -294,28 +320,29 @@ export function EmbeddingScatterplotSubscriber(props) {
   const cacheGet = (cache, key) => cache.find(el => isEqual(el[0], key))?.[1];
   const cellSetPolygons = useMemo(() => {
     if ((cellSetLabelsVisible || cellSetPolygonsVisible)
-      && !cacheHas(cellSetPolygonCache, cellSetSelection)
+      && !cacheHas(cellSetPolygonCache, deferredCellSetSelection)
       && mergedCellSets?.tree?.length
       && obsEmbedding
       && obsEmbeddingIndex
-      && cellSetColor?.length) {
+      && deferredCellSetColor?.length) {
       const newCellSetPolygons = getCellSetPolygons({
         obsIndex: obsEmbeddingIndex,
         obsEmbedding,
         cellSets: mergedCellSets,
-        cellSetSelection,
-        cellSetColor,
+        cellSetSelection: deferredCellSetSelection,
+        cellSetColor: deferredCellSetColor,
         theme,
       });
-      setCellSetPolygonCache(cache => [...cache, [cellSetSelection, newCellSetPolygons]]);
+      setCellSetPolygonCache(
+        cache => [...cache, [deferredCellSetSelection, newCellSetPolygons]],
+      );
       return newCellSetPolygons;
     }
-    return cacheGet(cellSetPolygonCache, cellSetSelection) || [];
+    return cacheGet(cellSetPolygonCache, deferredCellSetSelection) || [];
   }, [cellSetPolygonsVisible, cellSetPolygonCache, cellSetLabelsVisible, theme,
-    obsEmbeddingIndex, obsEmbedding, mergedCellSets, cellSetSelection, cellSetColor]);
-
+    obsEmbeddingIndex, obsEmbedding, mergedCellSets, deferredCellSetSelection,
+    deferredCellSetColor]);
 
-  const cellSelection = useMemo(() => Array.from(cellColors.keys()), [cellColors]);
 
   const [xRange, yRange, xExtent, yExtent, numCells] = useMemo(() => {
     if (obsEmbedding && obsEmbedding.data && obsEmbedding.shape) {
@@ -373,10 +400,13 @@ export function EmbeddingScatterplotSubscriber(props) {
     observationsLabel, obsLabelsTypes, obsLabelsData, obsSetsMembership,
   );
 
-  const cellSelectionSet = useMemo(() => new Set(cellSelection), [cellSelection]);
+  // With no set selection every observation counts as selected, matching the
+  // behavior of the ID-keyed color map this replaced. Uses the deferred selection
+  // so that it stays consistent with obsColorIndices within each commit.
+  const allCellsSelected = !(deferredCellSetSelection && mergedCellSets);
   const getCellIsSelected = useCallback((object, { index }) => (
-    (cellSelectionSet || new Set([])).has(obsEmbeddingIndex[index]) ? 1.0 : 0.0
-  ), [cellSelectionSet, obsEmbeddingIndex]);
+    (allCellsSelected || obsColorIndices.colorIndices[index] !== 0) ? 1.0 : 0.0
+  ), [allCellsSelected, obsColorIndices]);
 
   const cellRadius = (cellRadiusMode === 'manual' ? cellRadiusFixed : dynamicCellRadius);
   const cellOpacity = (cellOpacityMode === 'manual' ? cellOpacityFixed : dynamicCellOpacity);
@@ -468,8 +498,11 @@ export function EmbeddingScatterplotSubscriber(props) {
   // useExpressionValueGetter hook and getter function.
   const [alignedEmbeddingIndex, alignedEmbeddingData] = useMemo(() => {
     // Sort the embedding data according to the matrix obsIndex.
-    if (obsEmbedding?.data && obsEmbeddingIndex && matrixObsIndex) {
-      const matrixIndexMap = new Map(matrixObsIndex.map((key, i) => ([key, i])));
+    if (obsEmbedding?.data && obsEmbeddingIndex && matrixObsIndex
+      && obsEmbeddingIndex !== matrixObsIndex) {
+      // (The same index array, as when both come from one data source,
+      // is already aligned and skips this copy.)
+      const matrixIndexMap = getObsIndexMap(matrixObsIndex);
       const toMatrixIndex = obsEmbeddingIndex.map(key => matrixIndexMap.get(key));
 
       const newEmbeddingIndex = new Array(obsEmbeddingIndex.length);
@@ -489,42 +522,32 @@ export function EmbeddingScatterplotSubscriber(props) {
     return [obsEmbeddingIndex, obsEmbedding];
   }, [matrixObsIndex, obsEmbeddingIndex, obsEmbedding]);
 
-  const sampleIdToObsIdsMap = useMemo(() => {
-    // sampleEdges maps obsId -> sampleId.
-    // However when we stratify we want to map sampleId -> [obsId1, obsId2, ...].
-    // Here we create this reverse mapping.
-    if (sampleEdges) {
-      const result = new Map();
-      Array.from(sampleEdges.entries()).forEach(([obsId, sampleId]) => {
-        if (!result.has(sampleId)) {
-          result.set(sampleId, [obsId]);
-        } else {
-          result.get(sampleId).push(obsId);
-        }
-      });
-      return result;
-    }
-    return null;
-  }, [sampleEdges]);
-
   // Stratify multiple arrays: per-cellSet and per-sampleSet.
+  // The strata feed only the contour layers and, when the points are hidden, the
+  // observation count shown in the title — so with contours off and points on,
+  // this per-observation pass is skipped entirely rather than recomputed on
+  // every selection change.
+  const isStratifiedDataNeeded = embeddingContoursVisible || !embeddingPointsVisible;
   const [stratifiedData, stratifiedDataCount] = useMemo(() => {
-    if (alignedEmbeddingData?.data) {
+    if (isStratifiedDataNeeded && alignedEmbeddingData?.data) {
       const [result, cellCountResult] = stratifyArrays(
-        sampleEdges, sampleIdToObsIdsMap,
-        sampleSets, sampleSetSelection,
-        alignedEmbeddingIndex, mergedCellSets, cellSetSelection, {
+        sampleEdges,
+        sampleSets, deferredSampleSetSelection,
+        alignedEmbeddingIndex, mergedCellSets, deferredCellSetSelection, {
           obsEmbeddingX: alignedEmbeddingData.data[0],
           obsEmbeddingY: alignedEmbeddingData.data[1],
           ...(uint8ExpressionData?.[0] ? { featureValue: uint8ExpressionData } : {}),
         }, featureAggregationStrategyToUse,
+        // Raw codes let the strata come from typed arrays when the indices match.
+        { obsSetsColumns },
       );
       return [result, cellCountResult];
     }
     return [null, null];
-  }, [alignedEmbeddingIndex, alignedEmbeddingData, uint8ExpressionData,
-    sampleEdges, sampleIdToObsIdsMap, sampleSets, sampleSetSelection,
-    cellSetSelection, mergedCellSets, featureAggregationStrategyToUse,
+  }, [isStratifiedDataNeeded, alignedEmbeddingIndex, alignedEmbeddingData,
+    uint8ExpressionData, sampleEdges, sampleSets, deferredSampleSetSelection,
+    deferredCellSetSelection, mergedCellSets, featureAggregationStrategyToUse,
+    obsSetsColumns,
   ]);
 
   const setViewState = ({ zoom: newZoom, target }) => {
@@ -548,7 +571,7 @@ export function EmbeddingScatterplotSubscriber(props) {
       removeGridComponent={removeGridComponent}
       urls={urls}
       theme={theme}
-      isReady={isReady}
+      isReady={isReady && !isSelectionPending}
       helpText={helpText}
       errors={errors}
       options={(
@@ -595,6 +618,7 @@ export function EmbeddingScatterplotSubscriber(props) {
       <Scatterplot
         ref={deckRef}
         uuid={uuid}
+        onSelectionBusy={setIsSelectionPending}
         theme={theme}
         viewState={{ zoom, target: [targetX, targetY, targetZ] }}
         setViewState={setViewState}
@@ -602,9 +626,8 @@ export function EmbeddingScatterplotSubscriber(props) {
         obsEmbeddingIndex={obsEmbeddingIndex}
         obsEmbedding={obsEmbedding}
         cellFilter={cellFilter}
-        cellSelection={cellSelection}
         cellHighlight={cellHighlight}
-        cellColors={cellColors}
+        obsColorIndices={obsColorIndices}
         cellSetPolygons={cellSetPolygons}
         cellSetLabelSize={cellSetLabelSize}
         cellSetLabelsVisible={cellSetLabelsVisible}
@@ -624,12 +647,12 @@ export function EmbeddingScatterplotSubscriber(props) {
         getExpressionValue={getExpressionValue}
         getCellIsSelected={getCellIsSelected}
 
-        obsSetSelection={cellSetSelection}
-        sampleSetSelection={sampleSetSelection}
+        obsSetSelection={deferredCellSetSelection}
+        sampleSetSelection={deferredSampleSetSelection}
         // InternMap data structures where keys are
         // obsSet -> sampleSet -> arrayKey -> [].
         stratifiedData={stratifiedData}
-        obsSetColor={cellSetColor}
+        obsSetColor={deferredCellSetColor}
         sampleSetColor={sampleSetColor}
         contourThresholds={contourThresholds}
         contourColorEncoding={contourColorEncoding}