add multi-layout support to back-end (#766)

* add multi-layout support to back-end

* remove obsolete code

* temporary code to apply heuristic choice of default layout

* fix tests

* update python tests

* more py lint

* PR review changes

* more PR lint

* PR lint

client/__tests__/util/stateManager/sampleResponses.js

@@ -162,29 +162,7 @@ const aLayoutFBSResponse = (() => {
     new Float32Array(nObs).fill(Math.random()),
     new Float32Array(nObs).fill(Math.random())
   ];
-  const builder = new flatbuffers.Builder(1024);
-
-  const cols = _.map(coords, carr => {
-    const cdv = NetEncoding.Float32Array.createDataVector(builder, carr);
-    NetEncoding.Float32Array.startFloat32Array(builder);
-    NetEncoding.Float32Array.addData(builder, cdv);
-    const floatArr = NetEncoding.Float32Array.endFloat32Array(builder);
-
-    NetEncoding.Column.startColumn(builder);
-    NetEncoding.Column.addUType(builder, NetEncoding.TypedArray.Float32Array);
-    NetEncoding.Column.addU(builder, floatArr);
-    return NetEncoding.Column.endColumn(builder);
-  });
-
-  const columns = NetEncoding.Matrix.createColumnsVector(builder, cols);
-
-  NetEncoding.Matrix.startMatrix(builder);
-  NetEncoding.Matrix.addNRows(builder, nObs);
-  NetEncoding.Matrix.addNCols(builder, coords.length);
-  NetEncoding.Matrix.addColumns(builder, columns);
-  const matrix = NetEncoding.Matrix.endMatrix(builder);
-  builder.finish(matrix);
-  return builder.asUint8Array();
+  return encodeMatrix(coords, ["umap_0", "umap_1"]);
 })();
 
 const aDataObsResponse = {

client/src/util/stateManager/universe.js

@@ -78,6 +78,7 @@ function AnnotationsFBSToDataframe(arrayBuffer) {
   The application has strong assumptions that all scalar data will be
   stored as a float32 or float64 (regardless of underlying data types).
   For example, clipping of value ranges (eg, user-selected percentiles)
+  depends on the ability to use NaN in any numeric type.
 
   All float data from the server is left as is.  All non-float is promoted
   to an appropriate float.
@@ -98,13 +99,30 @@ function AnnotationsFBSToDataframe(arrayBuffer) {
 
 function LayoutFBSToDataframe(arrayBuffer) {
   const fbs = decodeMatrixFBS(arrayBuffer, true);
-  if (fbs.columns.length !== 2 || !fbs.columns.every(isFpTypedArray)) {
+  if (fbs.columns.length < 2 || !fbs.columns.every(isFpTypedArray)) {
     // We have strong assumptions about the shape & type of layout data.
     throw new Error("Unexpected layout data type returned from server");
   }
+
+  /*
+  TODO: XXX
+
+  TEMPORARY CODE AND COMMENT to support the progressive implementation
+  of multi-layout support.  For now, we search for one of the following 
+  in the layouts and use it if we find it: umap, then tsne, then pca, 
+  then whatever is first in the list.
+  */
+  let layoutIndex = 0;
+  ["umap", "tsne", "pca"].some(name => {
+    const idx = fbs.colIdx.indexOf(`${name}_0`);
+    if (idx !== -1) {
+      layoutIndex = idx;
+    }
+    return idx !== -1;
+  });
   const df = new Dataframe.Dataframe(
-    [fbs.nRows, fbs.nCols],
-    fbs.columns,
+    [fbs.nRows, 2],
+    [fbs.columns[layoutIndex], fbs.columns[layoutIndex + 1]],
     null,
     new Dataframe.KeyIndex(["X", "Y"])
   );

docs/faq.md

@@ -38,7 +38,7 @@ Currently this is not supported directly, but you should be able to do this your
 
 - `.obs` and `.var` annotations are use to extract metadata for filtering
 - `.X` is used to display expression (histograms, scatterplot & colorscale) and to compute differential expression
-- `.obsm` is used for layout
+- `.obsm` is used for layout. If an embedding has more than two components, the first two will be used for visualization.
 
 #### I have a BIG dataset - how can I make cellxgene run as fast as possible?
 

server/app/scanpy_engine/scanpy_engine.py

@@ -1,12 +1,13 @@
 import warnings
 
 import numpy as np
+import pandas
 from pandas.core.dtypes.dtypes import CategoricalDtype
 import anndata
 from scipy import sparse
 
 from server.app.driver.driver import CXGDriver
-from server.app.util.constants import Axis, DEFAULT_TOP_N
+from server.app.util.constants import Axis, DEFAULT_TOP_N, MAX_LAYOUTS
 from server.app.util.errors import (
     FilterError,
     JSONEncodingValueError,
@@ -41,7 +42,7 @@ class ScanpyEngine(CXGDriver):
     @staticmethod
     def _get_default_config():
         return {
-            "layout": "umap",
+            "layout": [],
             "diffexp": "ttest",
             "max_category_items": 100,
             "obs_names": None,
@@ -166,11 +167,55 @@ class ScanpyEngine(CXGDriver):
         self._alias_annotation_names(Axis.OBS, self.config["obs_names"])
         self._alias_annotation_names(Axis.VAR, self.config["var_names"])
         self._validate_data_types()
-        self._validate_data_calculations()
         self.cell_count = self.data.shape[0]
         self.gene_count = self.data.shape[1]
+        self._default_and_validate_layouts()
         self._create_schema()
 
+    @requires_data
+    def _default_and_validate_layouts(self):
+        """ function:
+            a) generate list of default layouts, if not already user specified
+            b) validate layouts are legal.  remove/warn on any that are not
+            c) cap total list of layouts at global const MAX_LAYOUTS
+        """
+        layouts = self.config['layout']
+        # handle default
+        if layouts is None or len(layouts) == 0:
+            # load default layouts from the data.
+            layouts = [key[2:] for key in self.data.obsm_keys() if type(key) == str and key.startswith("X_")]
+            if len(layouts) == 0:
+                raise PrepareError(f"Unable to find any precomputed layouts within the dataset.")
+
+        # remove invalid layouts
+        valid_layouts = []
+        obsm_keys = self.data.obsm_keys()
+        for layout in layouts:
+            layout_name = f"X_{layout}"
+            if layout_name not in obsm_keys:
+                warnings.warn(f"Ignoring unknown layout name: {layout}.")
+            elif not self._is_valid_layout(self.data.obsm[layout_name]):
+                warnings.warn(f"Ignoring layout due to malformed shape or data type: {layout}")
+            else:
+                valid_layouts.append(layout)
+
+        if len(valid_layouts) == 0:
+            raise PrepareError(f"No valid layout data.")
+
+        # cap layouts to MAX_LAYOUTS
+        self.config['layout'] = valid_layouts[0:MAX_LAYOUTS]
+
+    @requires_data
+    def _is_valid_layout(self, arr):
+        """ return True if this layout data is a valid array for front-end presentation:
+            * ndarray, with shape (n_obs, >= 2), dtype float/int/uint
+            * contains only finite values
+        """
+        is_valid = type(arr) == np.ndarray and arr.dtype.kind in "fiu"
+        is_valid = is_valid and arr.shape[0] == self.data.n_obs and arr.shape[1] >= 2
+        is_valid = is_valid and np.all(np.isfinite(arr))
+        return is_valid
+
     @requires_data
     def _validate_data_types(self):
         if sparse.isspmatrix(self.data.X) and not sparse.isspmatrix_csc(self.data.X):
@@ -208,20 +253,6 @@ class ScanpyEngine(CXGDriver):
                             f"annotations with more than 500 categories in the UI"
                         )
 
-    @requires_data
-    def _validate_data_calculations(self):
-        layout_key = f"X_{self.config['layout']}"
-        try:
-            assert layout_key in self.data.obsm_keys()
-        except AssertionError:
-            raise PrepareError(
-                f"Cannot find a field with coordinates for the {self.config['layout']} layout requested. A different"
-                f" layout may have been computed. The requested layout must be pre-calculated and saved "
-                f"back in the h5ad file. You can run "
-                f"`cellxgene prepare --layout {self.config['layout']} <datafile>` "
-                f"to solve this problem. "
-            )
-
     @staticmethod
     def _annotation_filter_to_mask(filter, d_axis, count):
         mask = np.ones((count,), dtype=bool)
@@ -372,15 +403,18 @@ class ScanpyEngine(CXGDriver):
         * only returns Matrix in columnar layout
         """
         try:
-            full_embedding = self.data.obsm[f"X_{self.config['layout']}"]
-            if full_embedding.shape[1] > 2:
-                warnings.warn(f"Warning: found {full_embedding.shape[1]} \
-                components of embedding. Using the first two for layout display.")
-            df_layout = full_embedding[:, :2]
+            layout_data = []
+            for layout in self.config["layout"]:
+                full_embedding = self.data.obsm[f"X_{layout}"]
+                embedding = full_embedding[:, :2]
+                normalized_layout = (embedding - embedding.min()) / (embedding.max() - embedding.min())
+                normalized_layout = normalized_layout.astype(dtype=np.float32)
+                layout_data.append(pandas.DataFrame(normalized_layout, columns=[f"{layout}_0", f"{layout}_1"]))
+
         except ValueError as e:
             raise PrepareError(
                 f"Layout has not been calculated using {self.config['layout']}, "
                 f"please prepare your datafile and relaunch cellxgene") from e
 
-        normalized_layout = (df_layout - df_layout.min()) / (df_layout.max() - df_layout.min())
-        return encode_matrix_fbs(normalized_layout.astype(dtype=np.float32), col_idx=None, row_idx=None)
+        df = pandas.concat(layout_data, axis=1, copy=False)
+        return encode_matrix_fbs(df, col_idx=df.columns, row_idx=None)

server/app/util/constants.py

@@ -31,3 +31,5 @@ JSON_NaN_to_num_warning_msg = (
     "JSON encoding failure - please verify all data are finite values (no NaN or Infinities)"
 )
 REACTIVE_LIMIT = 1_000_000
+
+MAX_LAYOUTS = 30

server/cli/launch.py

@@ -12,7 +12,6 @@ import psutil
 from server.app.app import Server
 from server.app.util.errors import ScanpyFileError
 from server.app.util.utils import custom_format_warning
-from server.utils.constants import MODES
 from server.utils.utils import find_available_port
 
 
@@ -25,10 +24,10 @@ BIG_FILE_SIZE_THRESHOLD = 100 * 2**20  # 100MB
 @click.option(
     "--layout",
     "-l",
-    type=click.Choice(MODES),
-    default="umap",
+    default=[],
+    multiple=True,
     show_default=True,
-    help="Method for layout."
+    help="Layout name, eg, 'umap'."
 )
 @click.option(
     "--diffexp",

server/test/test_api.py

@@ -67,9 +67,11 @@ class EndPoints(unittest.TestCase):
         self.assertEqual(result.headers["Content-Type"], "application/octet-stream")
         df = decode_fbs.decode_matrix_FBS(result.content)
         self.assertEqual(df['n_rows'], 2638)
-        self.assertEqual(df['n_cols'], 2)
+        self.assertEqual(df['n_cols'], 8)
         self.assertIsNotNone(df['columns'])
-        self.assertIsNone(df['col_idx'])
+        self.assertListEqual(df['col_idx'], [
+            'pca_0', 'pca_1', 'tsne_0', 'tsne_1', 'umap_0', 'umap_1', 'draw_graph_fr_0', 'draw_graph_fr_1'
+        ])
         self.assertIsNone(df['row_idx'])
         self.assertEqual(len(df['columns']), df['n_cols'])
 

server/test/test_nan_scanpy_engine.py

@@ -12,7 +12,7 @@ from server.app.util.errors import FilterError
 class NaNTest(unittest.TestCase):
     def setUp(self):
         self.args = {
-            "layout": "umap",
+            "layout": ["umap"],
             "diffexp": "ttest",
             "max_category_items": 100,
             "obs_names": None,

server/test/test_scanpy_engine.py

@@ -15,7 +15,7 @@ from server.app.util.errors import FilterError
 class EngineTest(unittest.TestCase):
     def setUp(self):
         args = {
-            "layout": "umap",
+            "layout": ["umap"],
             "diffexp": "ttest",
             "max_category_items": 100,
             "obs_names": None,

server/test/test_scanpy_engine_data_load.py

@@ -15,7 +15,7 @@ class DataLoadEngineTest(unittest.TestCase):
 
     def test_delayed_load_args(self):
         args = {
-            "layout": "tsne",
+            "layout": ["tsne"],
             "diffexp": "ttest",
             "max_category_items": 1000,
             "obs_names": "foo",