subset embedding UI improvement (#1667)

* add user flag feature to annomatrix

* add implicit subsetting for partial embeddings

* lint

* add embedding cell counts to embedding choice menu

* layout

client/src/actions/embedding.js

@@ -0,0 +1,33 @@
+/*
+action creators related to embeddings choice
+*/
+
+import { AnnoMatrixObsCrossfilter } from "../annoMatrix";
+import { _setEmbeddingSubset } from "../util/stateManager/viewStackHelpers";
+
+export const layoutChoiceAction = (newLayoutChoice) => async (
+  dispatch,
+  getState
+) => {
+  /*
+  On layout choice, make sure we have selected all on the previous layout, AND the new
+  layout.
+  */
+  const { annoMatrix: prevAnnoMatrix } = getState();
+
+  const embeddingDf = await prevAnnoMatrix.base().fetch("emb", newLayoutChoice);
+  const annoMatrix = _setEmbeddingSubset(prevAnnoMatrix, embeddingDf);
+  const obsCrossfilter = await new AnnoMatrixObsCrossfilter(annoMatrix).select(
+    "emb",
+    newLayoutChoice,
+    {
+      mode: "all",
+    }
+  );
+  dispatch({
+    type: "set layout choice",
+    layoutChoice: newLayoutChoice,
+    obsCrossfilter,
+    annoMatrix,
+  });
+};

client/src/actions/index.js

@@ -12,6 +12,7 @@ import { loadUserColorConfig } from "../util/stateManager/colorHelpers";
 import * as selnActions from "./selection";
 import * as annoActions from "./annotation";
 import * as viewActions from "./viewStack";
+import * as embActions from "./embedding";
 
 /*
 return promise fetching user-configured colors
@@ -40,6 +41,15 @@ async function configFetch(dispatch) {
   });
 }
 
+function prefetchEmbeddings(annoMatrix) {
+  /*
+  prefetch requests for all embeddings
+  */
+  const { schema } = annoMatrix;
+  const available = schema.layout.obs.map((v) => v.name);
+  available.forEach((embName) => annoMatrix.prefetch("emb", embName));
+}
+
 /*
 Application bootstrap
 */
@@ -57,6 +67,8 @@ const doInitialDataLoad = () =>
       const baseDataUrl = `${globals.API.prefix}${globals.API.version}`;
       const annoMatrix = new AnnoMatrixLoader(baseDataUrl, schema.schema);
       const obsCrossfilter = new AnnoMatrixObsCrossfilter(annoMatrix);
+      prefetchEmbeddings(annoMatrix);
+
       dispatch({
         type: "annoMatrix: init complete",
         annoMatrix,
@@ -210,6 +222,6 @@ export default {
   annotationLabelCurrentSelection: annoActions.annotationLabelCurrentSelection,
   saveObsAnnotationsAction: annoActions.saveObsAnnotationsAction,
   needToSaveObsAnnotations: annoActions.needToSaveObsAnnotations,
-  layoutChoiceAction: selnActions.layoutChoiceAction,
+  layoutChoiceAction: embActions.layoutChoiceAction,
   setCellSetFromSelection: selnActions.setCellSetFromSelection,
 };

client/src/actions/selection.js

@@ -185,31 +185,6 @@ export const graphLassoEndAction = (embName, polygon) => async (
   });
 };
 
-export const layoutChoiceAction = (newLayoutChoice) => async (
-  dispatch,
-  getState
-) => {
-  /*
-  On layout choice, make sure we have selected all on the previous layout, AND the new
-  layout.
-  */
-  const { obsCrossfilter: prevObsCrossfilter, layoutChoice } = getState();
-
-  let obsCrossfilter = await prevObsCrossfilter.select(
-    "emb",
-    layoutChoice.current,
-    { mode: "all" }
-  );
-  obsCrossfilter = await obsCrossfilter.select("emb", newLayoutChoice, {
-    mode: "all",
-  });
-  dispatch({
-    type: "set layout choice",
-    layoutChoice: newLayoutChoice,
-    obsCrossfilter,
-  });
-};
-
 /*
 Differential expression set selection
 */

client/src/actions/viewStack.js

@@ -11,7 +11,12 @@ stack multiple subsets.
 If these conventions change, code elsewhere (eg. menubar/clip.js) will need to
 change as well.
 */
-import { AnnoMatrixObsCrossfilter, clip, isubsetMask } from "../annoMatrix";
+import { AnnoMatrixObsCrossfilter } from "../annoMatrix";
+import {
+  _clipAnnoMatrix,
+  _userSubsetAnnoMatrix,
+  _userResetSubsetAnnoMatrix,
+} from "../util/stateManager/viewStackHelpers";
 
 export const clipAction = (min, max) => (dispatch, getState) => {
   /*
@@ -19,9 +24,7 @@ export const clipAction = (min, max) => (dispatch, getState) => {
   view is ALWAYS the top view.
   */
   const { annoMatrix: prevAnnoMatrix } = getState();
-  const annoMatrix = prevAnnoMatrix.isClipped
+  const annoMatrix = _clipAnnoMatrix(prevAnnoMatrix, min, max);
-    ? clip(prevAnnoMatrix.viewOf, min, max)
-    : clip(prevAnnoMatrix, min, max);
   const obsCrossfilter = new AnnoMatrixObsCrossfilter(annoMatrix);
   dispatch({
     type: "set clip quantiles",
@@ -43,24 +46,10 @@ export const subsetAction = () => (dispatch, getState) => {
     annoMatrix: prevAnnoMatrix,
     obsCrossfilter: prevObsCrossfilter,
   } = getState();
-
+  const annoMatrix = _userSubsetAnnoMatrix(
-  let annoMatrix;
+    prevAnnoMatrix,
-  if (prevAnnoMatrix.isClipped) {
+    prevObsCrossfilter.allSelectedMask()
-    // if there is a clip view, pop it and reapply after we subset
+  );
-    const { clipRange } = prevAnnoMatrix;
-    annoMatrix = isubsetMask(
-      prevAnnoMatrix.viewOf,
-      prevObsCrossfilter.allSelectedMask()
-    );
-    annoMatrix = clip(annoMatrix, ...clipRange);
-  } else {
-    // else just push a subset view.
-    annoMatrix = isubsetMask(
-      prevAnnoMatrix,
-      prevObsCrossfilter.allSelectedMask()
-    );
-  }
-
   const obsCrossfilter = new AnnoMatrixObsCrossfilter(annoMatrix);
   dispatch({
     type: "subset to selection",
@@ -77,20 +66,7 @@ export const resetSubsetAction = () => (dispatch, getState) => {
   */
 
   const { annoMatrix: prevAnnoMatrix } = getState();
-
+  const annoMatrix = _userResetSubsetAnnoMatrix(prevAnnoMatrix);
-  const clipRange = prevAnnoMatrix.isClipped ? prevAnnoMatrix.clipRange : null;
-
-  /* pop all views */
-  let annoMatrix = prevAnnoMatrix;
-  while (annoMatrix.isView) {
-    annoMatrix = annoMatrix.viewOf;
-  }
-
-  /* re-apply the clip, if any */
-  if (clipRange !== null) {
-    annoMatrix = clip(annoMatrix, ...clipRange);
-  }
-
   const obsCrossfilter = new AnnoMatrixObsCrossfilter(annoMatrix);
   dispatch({
     type: "reset subset",

client/src/annoMatrix/annoMatrix.js

@@ -70,6 +70,8 @@ export default class AnnoMatrix {
         The row index labels are as defined by the base dataset from the server.
       * isView - true if this is a view, false if not.
       * viewOf - pointer to parent annomatrix if a view, undefined/null if not a view.
+      * userFlags - container for any additional state a user of this API wants to hang
+        off of an annoMatrix, and have propagated by the (shallow) cloning protocol.
     */
     this.schema = indexEntireSchema(schema);
     this.nObs = nObs;
@@ -77,6 +79,7 @@ export default class AnnoMatrix {
     this.rowIndex = rowIndex || new IdentityInt32Index(nObs);
     this.isView = false;
     this.viewOf = undefined;
+    this.userFlags = {};
 
     /*
 		Private instance variables.

client/src/annoMatrix/viewCreators.js

@@ -33,6 +33,13 @@ export function subset(annoMatrix, obsLabels) {
   return new AnnoMatrixRowSubsetView(annoMatrix, obsIndex);
 }
 
+export function subsetByIndex(annoMatrix, obsIndex) {
+  /*
+  subset based upon the new obs index.
+  */
+  return new AnnoMatrixRowSubsetView(annoMatrix, obsIndex);
+}
+
 export function clip(annoMatrix, qmin, qmax) {
   /*
 		Create a view that clips all continuous data to the [min, max] range.
@@ -59,5 +66,5 @@ function _maskToList(mask) {
       elems += 1;
     }
   }
-  return new Int32Array(list.buffer, 0, elems);
+  return list.subarray(0, elems);
 }

client/src/components/embedding/index.js

@@ -1,5 +1,6 @@
 import React from "react";
 import { connect } from "react-redux";
+import Async from "react-async";
 import {
   ButtonGroup,
   Popover,
@@ -11,10 +12,11 @@ import {
 } from "@blueprintjs/core";
 import * as globals from "../../globals";
 import actions from "../../actions";
+import { getDiscreteCellEmbeddingRowIndex } from "../../util/stateManager/viewStackHelpers";
 
 @connect((state) => {
   return {
-    layoutChoice: state.layoutChoice,
+    layoutChoice: state.layoutChoice, // TODO: really should clean up naming, s/layout/embedding/g
     schema: state.annoMatrix?.schema,
     crossfilter: state.obsCrossfilter,
   };
@@ -32,6 +34,7 @@ class Embedding extends React.PureComponent {
 
   render() {
     const { layoutChoice, schema, crossfilter } = this.props;
+    const { annoMatrix } = crossfilter;
     return (
       <ButtonGroup
         style={{
@@ -61,7 +64,6 @@ class Embedding extends React.PureComponent {
               >
                 {layoutChoice?.current}: {crossfilter.countSelected()} out of{" "}
                 {crossfilter.size()} cells
-                {/* BRUCE to extend 1559 */}
               </Button>
             </Tooltip>
           }
@@ -87,9 +89,9 @@ class Embedding extends React.PureComponent {
                 selectedValue={layoutChoice.current}
               >
                 {layoutChoice.available.map((name) => (
-                  <Radio
+                  <LayoutChoice
-                    label={`${name} ${schema?.dataframe?.nObs} cells`}
+                    annoMatrix={annoMatrix}
-                    value={name}
+                    layoutName={name}
                     key={name}
                   />
                 ))}
@@ -103,3 +105,41 @@ class Embedding extends React.PureComponent {
 }
 
 export default Embedding;
+
+const loadEmbeddingCounts = async ({ annoMatrix, layoutName }) => {
+  const embedding = await annoMatrix.fetch("emb", layoutName);
+  const discreteCellIndex = getDiscreteCellEmbeddingRowIndex(embedding);
+  return { embedding, discreteCellIndex };
+};
+
+const LayoutChoice = ({ annoMatrix, layoutName }) => {
+  return (
+    <Async
+      promiseFn={loadEmbeddingCounts}
+      annoMatrix={annoMatrix}
+      layoutName={layoutName}
+    >
+      {({ data, error, isPending }) => {
+        if (error) {
+          /* log, as this is unexpected */
+          console.error(error);
+        }
+        if (error || isPending) {
+          /* still loading, or errored out - just omit counts (TODO: spinner?) */
+          return <Radio label={`${layoutName}`} value={layoutName} />;
+        }
+        if (data) {
+          const { embedding, discreteCellIndex } = data;
+          const isAllCells = discreteCellIndex.size() === embedding.length;
+          const sizeHint = `${discreteCellIndex.size()} ${
+            isAllCells ? "(all) " : ""
+          }cells`;
+          return (
+            <Radio label={`${layoutName}: ${sizeHint}`} value={layoutName} />
+          );
+        }
+        return null;
+      }}
+    </Async>
+  );
+};

client/src/components/menubar/index.js

@@ -10,6 +10,7 @@ import InformationMenu from "./infoMenu";
 import Subset from "./subset";
 import UndoRedoReset from "./undoRedo";
 import DiffexpButtons from "./diffexpButtons";
+import { getEmbSubsetView } from "../../util/stateManager/viewStackHelpers";
 
 @connect((state) => {
   const { annoMatrix } = state;
@@ -17,9 +18,11 @@ import DiffexpButtons from "./diffexpButtons";
   const selectedCount = crossfilter.countSelected();
 
   const subsetPossible =
-    selectedCount !== 0 && selectedCount !== crossfilter.size(); // ie, not all are selected
+    selectedCount !== 0 && selectedCount !== crossfilter.size(); // ie, not all and not none are selected
-  const subsetResetPossible =
+  const embSubsetView = getEmbSubsetView(annoMatrix);
-    annoMatrix.nObs !== annoMatrix.schema.dataframe.nObs;
+  const subsetResetPossible = !embSubsetView
+    ? annoMatrix.nObs !== annoMatrix.schema.dataframe.nObs
+    : annoMatrix.nObs !== embSubsetView.nObs;
 
   return {
     subsetPossible,

client/src/util/stateManager/viewStackHelpers.js

@@ -0,0 +1,175 @@
+/*
+The annoMatrix view stack has a set of conventions which are assumed elsewhere in the
+application. These helper functions make it simple for action creators to manage
+the stack.
+
+The annoMatrix module does not care about this order, but we maintain it as
+a convention to make it simpler to manipulate the views.
+
+Terminology:
+- clip view: AnnoMatrixClipView
+- subset view: AnnoMatrixRowSubsetView
+- user subset view: create by the user explicitly subsetting by selection
+- embedding subset view: implicitly created by switching the current embedding
+- loader, or base annoMatrix:  the root, which loads data
+
+Rules:
+1. there will be zero or one clip view
+2. there will be zero or more subset views
+3. there will be zero or one embedding view
+4. there will be one loader/base, which is always the bottom view
+5. the view ordering MUST be (top to bottom):
+
+	[clip] -> [user subset] -> [embedding subset] -> loader
+
+There is code elsewhere in the app (eg, menubar/clip.js) which assumes this order.
+
+Views can be interogated for their type with the following:
+
+* is a view:  annoMatrix.isView
+* is the loader:  !anonMatrix.isView (or annoMatrix === annoMatrix.base())
+* is a clip view:  annoMatrix.isClipped (or annoMatrix.clipRange)
+* is a subset view:  (annoMatrix.isView && !annoMatrix.isClipped)
+* is a user subset view:  annoMatrix.userFlags?.isUserSubsetView
+* is an embedding subset view:  annomatrix.userFlags?.isEmbSubsetView
+
+*/
+
+import { clip, isubsetMask, isubset } from "../../annoMatrix";
+import { memoize } from "../dataframe/util";
+
+export function _clipAnnoMatrix(annoMatrix, min, max) {
+  /*
+  clip the annoMatrix.
+  */
+  return annoMatrix.isClipped
+    ? clip(annoMatrix.viewOf, min, max)
+    : clip(annoMatrix, min, max);
+}
+
+export function _userSubsetAnnoMatrix(annoMatrix, mask) {
+  /*
+  user-requested row subset of annoMatrix, to be added on top of any 
+  other previous row subsets.
+  */
+  const { clipRange } = annoMatrix;
+  if (clipRange) {
+    annoMatrix = annoMatrix.viewOf;
+  }
+
+  annoMatrix = isubsetMask(annoMatrix, mask);
+  annoMatrix.userFlags.isUserSubsetView = true;
+
+  if (clipRange) {
+    annoMatrix = clip(annoMatrix, ...clipRange);
+  }
+
+  return annoMatrix;
+}
+
+export function _userResetSubsetAnnoMatrix(annoMatrix) {
+  /*
+  Reset/remove all user-requested subsets.  Do not remove clip or embedding subset.
+  */
+
+  /* stash clipping info, if any */
+  const { clipRange } = annoMatrix;
+  if (clipRange) {
+    annoMatrix = annoMatrix.viewOf;
+  }
+
+  /* pop all views except embedding subset and loader */
+  while (annoMatrix.isView && annoMatrix.userFlags.isUserSubsetView) {
+    annoMatrix = annoMatrix.viewOf;
+  }
+
+  /* re-apply the clip, if any */
+  if (clipRange) {
+    annoMatrix = clip(annoMatrix, ...clipRange);
+  }
+
+  return annoMatrix;
+}
+
+export function _setEmbeddingSubset(annoMatrix, embeddingDf) {
+  /*
+  Set the embedding subset view.  Only create a subset view for the embedding
+  when it is needed, ie, there are NaN values in the embeddings.
+  */
+  const embRowOffsets = _getEmbeddingRowOffsets(
+    annoMatrix.rowIndex,
+    embeddingDf
+  );
+
+  const curEmbSubsetView = getEmbSubsetView(annoMatrix);
+
+  /* if no current embedding subset, and no new embedding subset, just noop */
+  if (!embRowOffsets && !curEmbSubsetView) return annoMatrix;
+
+  // ... otherwise, do the work
+
+  /* stash clipping info, if any */
+  const clipRange = annoMatrix.isClipped ? annoMatrix.clipRange : null;
+
+  /* pop all subsets, user or embedding */
+  while (annoMatrix.isView) {
+    annoMatrix = annoMatrix.viewOf;
+  }
+
+  /* apply new embedding row index, if needed */
+  if (embRowOffsets) {
+    annoMatrix = isubset(annoMatrix, embRowOffsets);
+    annoMatrix.userFlags.isEmbSubsetView = true;
+  }
+
+  /* re-apply clip, if needed */
+  if (clipRange) {
+    annoMatrix = clip(annoMatrix, ...clipRange);
+  }
+
+  return annoMatrix;
+}
+
+function _getEmbeddingRowOffsets(baseRowIndex, embeddingDf) {
+  /*
+  given a dataframe containing an embedding:
+    - if the embedding contains no NaN coordinates, return null
+    - if the embedding contains NaN coordinates, return a rowIndex
+      that contains only the rows with discrete valued coordinates.
+
+  Currently assumes that there will be onl two dimensions in the embedding.
+  */
+  const X = embeddingDf.icol(0).asArray();
+  const Y = embeddingDf.icol(1).asArray();
+  const offsets = new Int32Array(X.length);
+  let numOffsets = 0;
+
+  for (let i = 0, l = X.length; i < l; i += 1) {
+    if (!Number.isNaN(X[i]) && !Number.isNaN(Y[i])) {
+      offsets[numOffsets] = i;
+      numOffsets += 1;
+    }
+  }
+
+  if (numOffsets === X.length) return null;
+  return offsets.subarray(0, numOffsets);
+}
+
+export function _getDiscreteCellEmbeddingRowIndex(embeddingDf) {
+  const idx = _getEmbeddingRowOffsets(embeddingDf.rowIndex, embeddingDf);
+  if (idx === null) return embeddingDf.rowIndex;
+  return embeddingDf.rowIndex.isubset(idx);
+}
+export const getDiscreteCellEmbeddingRowIndex = memoize(
+  _getDiscreteCellEmbeddingRowIndex,
+  (df) => df.__id
+);
+
+export function getEmbSubsetView(annoMatrix) {
+  /* if there is an embedding subset in the view stack, return it.  Falsish if not. */
+  while (annoMatrix.isView) {
+    if (annoMatrix.userFlags.isEmbSubsetView) return annoMatrix;
+    annoMatrix = annoMatrix.viewOf;
+  }
+  return undefined;
+}