experimental re-embedding (#1186)

* first cut at re-embedding route and back-end support

* update and expand config route tests

* add scanpy_umap

* add reembedding to config route parameters

* front-end support for reembedding fetch and UI

* remove unused imports

* add loading state

* save reembedding in reducer state

* improve withColsFrom

* transmit reembed schema to client; pick unique embedding names

* display embeddings

* format

* lint

* spaces, tab size 2

* lint

* test hack for smoke-test race

* back out hack sleep

* add check for backed mode

* add unit test for reembedding

* lint

* hide re-embedding CLI param from help

client/src/reducers/colors.js

@@ -12,12 +12,12 @@ const ColorsReducer = (
   prevSharedState
 ) => {
   switch (action.type) {
-    case "universe exists, but loading is still in progress":
-    case "reset World to eq Universe": {
+    case "universe exists, but loading is still in progress": {
+      /* initialize everything with default colors, no mode, no color-by accessor */
       const { world } = nextSharedState;
       const colorMode = null;
       const colorAccessor = null;
-      const { rgb, scale } = ColorHelpers.createColors(world, colorMode);
+      const { rgb, scale } = ColorHelpers.createColors(world);
       return {
         ...state,
         colorAccessor,
@@ -27,6 +27,22 @@ const ColorsReducer = (
       };
     }
 
+    case "reset World to eq Universe": {
+      /* need to rebuild colors as world may have changed, but don't switch modes */
+      const { world } = nextSharedState;
+      const { colorMode, colorAccessor } = state;
+      const { rgb, scale } = ColorHelpers.createColors(
+        world,
+        colorMode,
+        colorAccessor
+      );
+      return {
+        ...state,
+        rgb,
+        scale
+      };
+    }
+
     case "set clip quantiles":
     case "set World to current selection": {
       const { world: prevWorld, controls: prevControls } = prevSharedState;

client/src/reducers/index.js

@@ -21,6 +21,7 @@ import autosave from "./autosave";
 import ontology from "./ontology";
 import centroidLabels from "./centroidLabels";
 import pointDialation from "./pointDilation";
+import { reembedController, reembedding } from "./reembed";
 
 import undoableConfig from "./undoableConfig";
 
@@ -31,6 +32,7 @@ const Reducer = undoable(
     ["world", world],
     ["ontology", ontology],
     ["annotations", annotations],
+    ["reembedding", reembedding],
     ["layoutChoice", layoutChoice],
     ["categoricalSelection", categoricalSelection],
     ["continuousSelection", continuousSelection],
@@ -42,6 +44,7 @@ const Reducer = undoable(
     ["responsive", responsive],
     ["centroidLabels", centroidLabels],
     ["pointDilation", pointDialation],
+    ["reembedController", reembedController],
     ["autosave", autosave],
     ["resetCache", resetCache]
   ]),
@@ -57,7 +60,8 @@ const Reducer = undoable(
     "differential",
     "layoutChoice",
     "centroidLabels",
-    "annotations"
+    "annotations",
+    "reembedding"
   ],
   undoableConfig
 );

client/src/reducers/layoutChoice.js

@@ -14,6 +14,14 @@ function bestDefaultLayout(layouts) {
   return layouts[0];
 }
 
+function setToDefaultLayout(world) {
+  const { schema } = world;
+  const available = schema.layout.obs.map(v => v.name).sort();
+  const current = bestDefaultLayout(available);
+  const currentDimNames = schema.layout.obsByName[current].dims;
+  return { available, current, currentDimNames };
+}
+
 const LayoutChoice = (
   state = {
     available: [], // all available choices
@@ -24,14 +32,13 @@ const LayoutChoice = (
   nextSharedState
 ) => {
   switch (action.type) {
-    case "universe exists, but loading is still in progress":
-    case "reset World to eq Universe": {
+    case "universe exists, but loading is still in progress": {
       // set default to default
-      const { schema } = nextSharedState.world;
-      const available = schema.layout.obs.map(v => v.name).sort();
-      const current = bestDefaultLayout(available);
-      const currentDimNames = schema.layout.obsByName[current].dims;
-      return { available, current, currentDimNames };
+      const { universe } = nextSharedState;
+      return {
+        ...state,
+        ...setToDefaultLayout(universe)
+      };
     }
 
     case "set layout choice": {
@@ -41,6 +48,31 @@ const LayoutChoice = (
       return { ...state, current, currentDimNames };
     }
 
+    case "reembed: add reembedding": {
+      const name = action.schema.name;
+      const available = Array.from(new Set(state.available).add(name));
+      return {
+        ...state,
+        available
+      };
+    }
+
+    case "reembed: clear all reembeddings": {
+      const { universe } = nextSharedState;
+      const { current } = state;
+      const dflt = setToDefaultLayout(universe);
+      if (dflt.available.includes(current)) {
+        return {
+          ...state,
+          available: dflt.available
+        };
+      }
+      return {
+        ...state,
+        ...dflt
+      };
+    }
+
     default: {
       return state;
     }

client/src/reducers/reembed.js

@@ -0,0 +1,64 @@
+/*
+controller state is not part of the undo/redo history
+*/
+export const reembedController = (
+  state = {
+    pendingFetch: null
+  },
+  action
+) => {
+  switch (action.type) {
+    case "reembed: request start": {
+      return {
+        ...state,
+        pendingFetch: action.abortableFetch
+      };
+    }
+    case "reembed: request aborted":
+    case "reembed: request cancel":
+    case "reembed: request completed": {
+      return {
+        ...state,
+        pendingFetch: null
+      };
+    }
+    default: {
+      return state;
+    }
+  }
+};
+
+/*
+actual reembedding data is part of the undo/redo history
+*/
+export const reembedding = (
+  state = {
+    reembeddings: new Map()
+  },
+  action
+) => {
+  switch (action.type) {
+    case "reembed: add reembedding": {
+      const { schema, embedding } = action;
+      const { name } = schema.name;
+      const { reembeddings } = state;
+      return {
+        ...state,
+        reembeddings: new Map(reembeddings).set(name, {
+          name,
+          schema,
+          embedding
+        })
+      };
+    }
+    case "reembed: clear all reembeddings": {
+      return {
+        ...state,
+        reembeddings: new Map()
+      };
+    }
+    default: {
+      return state;
+    }
+  }
+};

client/src/reducers/world.js

@@ -4,6 +4,10 @@ import {
   ControlsHelpers,
   AnnotationsHelpers
 } from "../util/stateManager";
+import {
+  addObsLayout,
+  removeObsLayout
+} from "../util/stateManager/schemaHelpers";
 import clip from "../util/clip";
 import quantile from "../util/quantile";
 
@@ -32,6 +36,7 @@ const WorldReducer = (
     }
 
     case "universe: column load success": {
+      /* incremental initial data load - always assumes world == universe */
       const { universe } = nextSharedState;
       const { dim } = action;
       return {
@@ -276,6 +281,52 @@ const WorldReducer = (
       };
     }
 
+    case "reembed: add reembedding": {
+      // new embedding loaded, which *only* affects world's layout.
+      // It may be new, or it may replace a previous re-embedding.
+      const { obsLayout: origObsLayout, schema: origSchema } = state;
+      const { embedding, schema: embeddingSchema } = action;
+
+      const { dims, name } = embeddingSchema;
+      let obsLayout = origObsLayout;
+      let schema = origSchema;
+
+      // alias the names the server sent us, in case they were not the same as the schema
+      const embedingLabels = embedding.colIndex.keys();
+      const labels = {
+        [embedingLabels[0]]: dims[0],
+        [embedingLabels[1]]: dims[1]
+      };
+      obsLayout = obsLayout.withColsFrom(embedding, labels);
+      schema = addObsLayout(schema, embeddingSchema);
+      return {
+        ...state,
+        obsLayout,
+        schema
+      };
+    }
+
+    case "reembed: clear all reembeddings": {
+      // reembedding was cleared -- remove from layout
+      const { obsLayout: origObsLayout, schema: origSchema } = state;
+      const { reembedding } = prevSharedState;
+
+      let schema = origSchema;
+      let obsLayout = origObsLayout;
+
+      reembedding.reembeddings.forEach((emb, name) => {
+        const { dims } = emb.schema;
+        obsLayout = obsLayout.dropCol(dims[0]);
+        obsLayout = obsLayout.dropCol(dims[1]);
+        schema = removeObsLayout(schema, name);
+      });
+      return {
+        ...state,
+        obsLayout,
+        schema
+      };
+    }
+
     default: {
       return state;
     }