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

server/app/app.py

@@ -163,6 +163,10 @@ class LayoutObsAPI(Resource):
     def get(self, data_adaptor):
         return common_rest.layout_obs_get(request, data_adaptor)
 
+    @rest_get_data_adaptor
+    def put(self, data_adaptor):
+        return common_rest.layout_obs_put(request, data_adaptor)
+
 
 def get_api_resources(bp_api):
     api = Api(bp_api)

server/cli/launch.py

@@ -13,7 +13,7 @@ import click
 from server.common.utils import custom_format_warning
 from server.common.utils import find_available_port, is_port_available, sort_options
 from server.common.errors import DatasetAccessError
-from server.data_common.matrix_loader import MatrixDataLoader, MatrixDataCacheManager
+from server.data_common.matrix_loader import MatrixDataLoader, MatrixDataCacheManager, MatrixDataType
 from server.common.annotations import AnnotationsLocalFile
 from server.common.app_config import AppConfig
 
@@ -103,6 +103,14 @@ def config_args(func):
         metavar="<text>",
         help="Embedding name, eg, 'umap'. Repeat option for multiple embeddings. Defaults to all.",
     )
+    @click.option(
+        "--experimental-enable-reembedding",
+        is_flag=True,
+        default=False,
+        show_default=False,
+        hidden=True,
+        help="Enable experimental on-demand re-embedding using UMAP. WARNING: may be very slow.",
+    )
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
         return func(*args, **kwargs)
@@ -278,6 +286,7 @@ def launch(
     disable_diffexp,
     experimental_annotations_ontology,
     experimental_annotations_ontology_obo,
+    experimental_enable_reembedding,
 ):
     """Launch the cellxgene data viewer.
     This web app lets you explore single-cell expression data.
@@ -317,6 +326,14 @@ def launch(
         except DatasetAccessError as e:
             raise click.ClickException(str(e))
 
+        if experimental_enable_reembedding:
+            if matrix_data_loader.matrix_data_type() != MatrixDataType.H5AD:
+                raise click.ClickException("--experimental-enable-reembedding is only supported with H5AD files.")
+            if backed:
+                raise click.ClickException(
+                    "--experimental-enable-reembedding is not supported when run in --backed mode."
+                )
+
         file_size = matrix_data_loader.file_size()
         if file_size > BIG_FILE_SIZE_THRESHOLD:
             click.echo(f"[cellxgene] Loading data from {basename(datapath)}, this may take a while...")
@@ -402,6 +419,7 @@ def launch(
         var_names=var_names,
         anndata_backed=backed,
         disable_diffexp=disable_diffexp,
+        enable_reembedding=experimental_enable_reembedding,
     )
 
     matrix_data_cache_manager = MatrixDataCacheManager()

server/common/app_config.py

@@ -32,6 +32,7 @@ class AppConfig(object):
         self.max_category_items = 100
         self.diffexp_lfc_cutoff = 0.01
         self.disable_diffexp = False
+        self.enable_reembedding = False
         self.anndata_backed = False
 
         # TODO these options may not apply to all datasets in the multi dataset.
@@ -56,6 +57,7 @@ class AppConfig(object):
             "var_names",
             "anndata_backed",
             "disable_diffexp",
+            "enable_reembedding",
         ]
 
         self.update(inputs, kw)
@@ -80,7 +82,7 @@ class AppConfig(object):
         # we have camalCase, hyphen-text, and underscore_text
 
         # features
-        features = [f.todict() for f in data_adaptor.get_features().values()]
+        features = [f.todict() for f in data_adaptor.get_features(annotation)]
 
         # display_names
         title = self.get_title(data_adaptor)
@@ -105,6 +107,7 @@ class AppConfig(object):
             "diffexp_lfc_cutoff": self.diffexp_lfc_cutoff,
             "backed": self.anndata_backed,
             "disable-diffexp": self.disable_diffexp,
+            "enable-reembedding": self.enable_reembedding,
             "annotations": False,
             "annotations_file": None,
             "annotations_output_dir": None,

server/common/rest.py

@@ -165,7 +165,7 @@ def diffexp_obs_post(request, data_adaptor):
     try:
         diffexp = data_adaptor.diffexp_topN(set1_filter, set2_filter, count)
         return make_response(diffexp, HTTPStatus.OK, {"Content-Type": "application/json"})
-    except (ValueError, FilterError) as e:
+    except (ValueError, DisabledFeatureError, FilterError) as e:
         return make_response(str(e), HTTPStatus.BAD_REQUEST)
     except JSONEncodingValueError as e:
         # JSON encoding failure, usually due to bad data
@@ -188,3 +188,33 @@ def layout_obs_get(request, data_adaptor):
         return make_response(str(e), HTTPStatus.INTERNAL_SERVER_ERROR)
     except ValueError as e:
         return make_response(str(e), HTTPStatus.INTERNAL_SERVER_ERROR)
+
+
+def layout_obs_put(request, data_adaptor):
+    preferred_mimetype = request.accept_mimetypes.best_match(["application/octet-stream"])
+    if preferred_mimetype != "application/octet-stream":
+        return make_response(f"Unsupported MIME type '{request.accept_mimetypes}'", HTTPStatus.NOT_ACCEPTABLE)
+    if not data_adaptor.config.enable_reembedding:
+        return make_response(f"Computed embedding not supported.", HTTPStatus.BAD_REQUEST)
+
+    args = request.get_json()
+    filter = args["filter"] if args else None
+    if not filter:
+        return make_response("Error: obs filter is required", HTTPStatus.BAD_REQUEST)
+    method = args["method"] if args else "umap"
+
+    try:
+        schema, fbs = data_adaptor.compute_embedding(method, filter)
+        return make_response(
+            fbs,
+            HTTPStatus.OK,
+            {
+                "Content-Type": "application/octet-stream",
+                "CxG-Schema": json.dumps(schema),
+                "Access-Control-Expose-Headers": "CxG-Schema",
+            },
+        )
+    except NotImplementedError as e:
+        return make_response(str(e), HTTPStatus.NOT_IMPLEMENTED)
+    except (ValueError, DisabledFeatureError, FilterError) as e:
+        return make_response(str(e), HTTPStatus.BAD_REQUEST)

server/compute/scanpy.py

@@ -0,0 +1,49 @@
+import importlib
+
+"""
+Wrapper for various scanpy modules.  Will raise NotImplementedError if the scanpy
+module is not installed/available
+"""
+
+
+def get_scanpy_module():
+    try:
+        sc = importlib.import_module("scanpy")
+        # Future: we could enforce versions here, eg, lookat sc.__version__
+        return sc
+    except ModuleNotFoundError:
+        raise NotImplementedError("Please install scanpy to enable UMAP re-embedding")
+    except Exception as e:
+        # will capture other ImportError corner cases
+        raise NotImplementedError(str(e))
+
+
+def scanpy_umap(adata, obs_mask=None, pca_options={}, neighbors_options={}, umap_options={}):
+    """
+    Given adata and an obs mask, return a new embedding for adata[obs_mask, :]
+    as an ndarray of shape (len(obs_mask), N), where N>=2.
+
+    Do NOT mutate adata.
+    """
+
+    # backed mode is incompatible with the current implementation
+    if adata.isbacked:
+        raise NotImplementedError("Backed mode is incompatible with re-embedding")
+
+    # safely get scanpy module, which may not be present.
+    sc = get_scanpy_module()
+
+    # https://github.com/theislab/anndata/issues/311
+    obs_mask = slice(None) if obs_mask is None else obs_mask
+    adata = adata[obs_mask, :].copy()
+
+    for k in list(adata.obsm.keys()):
+        del adata.obsm[k]
+    for k in list(adata.uns.keys()):
+        del adata.uns[k]
+
+    sc.pp.pca(adata, zero_center=None, n_comps=min(adata.n_obs - 1, 50), **pca_options)
+    sc.pp.neighbors(adata, **neighbors_options)
+    sc.tl.umap(adata, **umap_options)
+
+    return adata.obsm["X_umap"]

server/data_anndata/anndata_adaptor.py

@@ -1,17 +1,21 @@
 import warnings
 
 import numpy as np
+import pandas as pd
 from pandas.core.dtypes.dtypes import CategoricalDtype
 import anndata
 from scipy import sparse
 from packaging import version
+from datetime import datetime
+from server_timing import Timing as ServerTiming
 
 from server.data_common.data_adaptor import DataAdaptor
 from server.data_common.fbs.matrix import encode_matrix_fbs
 from server.common.utils import series_to_schema
 from server.common.constants import Axis, MAX_LAYOUTS
-from server.common.errors import PrepareError, DatasetAccessError
+from server.common.errors import PrepareError, DatasetAccessError, FilterError
 from server.common.data_locator import DataLocator
+from server.compute.scanpy import scanpy_umap
 
 anndata_version = version.parse(str(anndata.__version__)).release
 
@@ -261,7 +265,10 @@ class AnndataAdaptor(DataAdaptor):
         return encode_matrix_fbs(df, col_idx=df.columns)
 
     def get_embedding_names(self):
-        """ function:
+        """
+        Return pre-computed embeddings.
+
+        function:
             a) generate list of default layouts
             b) validate layouts are legal.  remove/warn on any that are not
             c) cap total list of layouts at global const MAX_LAYOUTS
@@ -294,6 +301,30 @@ class AnndataAdaptor(DataAdaptor):
         full_embedding = self.data.obsm[f"X_{ename}"]
         return full_embedding[:, 0:dims]
 
+    def compute_embedding(self, method, obsFilter):
+        if Axis.VAR in obsFilter:
+            raise FilterError("Observation filters may not contain variable conditions")
+        if method != "umap":
+            raise NotImplementedError(f"re-embedding method {method} is not available.")
+        try:
+            shape = self.get_shape()
+            obs_mask = self._axis_filter_to_mask(Axis.OBS, obsFilter["obs"], shape[0])
+        except (KeyError, IndexError) as e:
+            raise FilterError(f"Error parsing filter: {e}") from e
+
+        with ServerTiming.time("layout.compute"):
+            X_umap = scanpy_umap(self.data, obs_mask)
+            normalized_layout = DataAdaptor.normalize_embedding(X_umap)
+
+        # Server picks reemedding name, which must not collide with any other
+        # embedding name generated by this backed.
+        name = f"reembed:{method}_{datetime.now().isoformat(timespec='milliseconds')}"
+        dims = [f"{name}_0", f"{name}_1"]
+        df = pd.DataFrame(normalized_layout, columns=dims)
+        fbs = encode_matrix_fbs(df, col_idx=df.columns, row_idx=None)
+        schema = {"name": name, "type": "float32", "dims": dims}
+        return (schema, fbs)
+
     def get_X_array(self, obs_mask=None, var_mask=None):
         if obs_mask is None:
             obs_mask = slice(None)

server/data_common/data_adaptor.py

@@ -57,12 +57,18 @@ class DataAdaptor(metaclass=ABCMeta):
 
     @abstractmethod
     def get_embedding_names(self):
-        """return a list of embedding names"""
+        """return a list of pre-computed embedding names"""
         pass
 
     @abstractmethod
     def get_embedding_array(self, ename, dims=2):
-        """return an numpy array for the given embedding name."""
+        """return an numpy array for the given pre-computed embedding name."""
+        pass
+
+    @abstractmethod
+    def compute_embedding(self, method, filter):
+        """compute a new embedding on the specified obs subset, and return a
+           tuple of (schema, fbs)."""
         pass
 
     @abstractmethod
@@ -126,21 +132,15 @@ class DataAdaptor(metaclass=ABCMeta):
         """
         pass
 
-    def get_features(self):
-        features = {}
-        features["cluster"] = AppFeature("/cluster/")
-
-        if self.get_embedding_names():
-            # TODO handle "var" when gene layout becomes available
-            features["layout_obs"] = AppFeature("/layout/obs", available=True)
-        else:
-            features["layout_obs"] = AppFeature("/layout/obs")
-
-        if self.config.disable_diffexp:
-            features["diffexp"] = AppFeature("/diffexp/")
-        else:
-            features["diffexp"] = AppFeature("/diffexp/", available=True)
-
+    def get_features(self, annotations=None):
+        """Return list of features, to return as part of the config route"""
+        features = [
+            AppFeature("/cluster/", method="POST", available=False),
+            AppFeature("/layout/obs", method="GET", available=self.get_embedding_names() is not None),
+            AppFeature("/layout/obs", method="PUT", available=self.config.enable_reembedding),
+            AppFeature("/diffexp/", method="POST", available=not self.config.disable_diffexp),
+            AppFeature("/annotations/obs", method="PUT", available=annotations is not None),
+        ]
         return features
 
     def update_parameters(self, parameters):
@@ -294,6 +294,25 @@ class DataAdaptor(metaclass=ABCMeta):
         except ValueError:
             raise JSONEncodingValueError("Error encoding differential expression to JSON")
 
+    @staticmethod
+    def normalize_embedding(embedding):
+        """Normalize embedding layout to meet client assumptions.
+           Embedding is an ndarray, shape (n_obs, n)., where n is normally 2
+        """
+
+        # scale isotropically
+        min = embedding.min(axis=0)
+        max = embedding.max(axis=0)
+        scale = np.amax(max - min)
+        normalized_layout = (embedding - min) / scale
+
+        # translate to center on both axis
+        translate = 0.5 - ((max - min) / scale / 2)
+        normalized_layout = normalized_layout + translate
+
+        normalized_layout = normalized_layout.astype(dtype=np.float32)
+        return normalized_layout
+
     def layout_to_fbs_matrix(self):
         """ same as layout, except returns a flatbuffer """
         """
@@ -312,18 +331,7 @@ class DataAdaptor(metaclass=ABCMeta):
         with ServerTiming.time(f"layout.query"):
             for ename in embeddings:
                 embedding = self.get_embedding_array(ename, 2)
-
-                # scale isotropically
-                min = embedding.min(axis=0)
-                max = embedding.max(axis=0)
-                scale = np.amax(max - min)
-                normalized_layout = (embedding - min) / scale
-
-                # translate to center on both axis
-                translate = 0.5 - ((max - min) / scale / 2)
-                normalized_layout = normalized_layout + translate
-
-                normalized_layout = normalized_layout.astype(dtype=np.float32)
+                normalized_layout = DataAdaptor.normalize_embedding(embedding)
                 layout_data.append(pd.DataFrame(normalized_layout, columns=[f"{ename}_0", f"{ename}_1"]))
 
         with ServerTiming.time(f"layout.encode"):

server/data_cxg/cxg_adaptor.py

@@ -164,6 +164,9 @@ class CxgAdaptor(DataAdaptor):
         array = self.open_array(f"emb/{ename}")
         return array[:, 0:dims]
 
+    def compute_embedding(self, method, filter):
+        raise NotImplementedError("CXG does not yet support re-embedding")
+
     def get_X_array(self, obs_mask=None, var_mask=None):
         obs_items = self._convert_mask(obs_mask)
         var_items = self._convert_mask(var_mask)

server/test/test_anndata_adaptor.py

@@ -3,6 +3,7 @@ from os import path
 import pytest
 import time
 import unittest
+import sys
 import server.test.decode_fbs as decode_fbs
 from parameterized import parameterized_class
 
@@ -97,7 +98,21 @@ class AdaptorTest(unittest.TestCase):
             self.data._create_schema()
 
     def test_config(self):
-        self.assertEqual(self.data.get_features()["layout_obs"].available, True)
+        features = self.data.get_features(annotations=None)
+
+        # test each for singular presence and accuracy of available flag
+        def check_feature(method, path, available):
+            feature = list(
+                filter(lambda f: f.method == method and f.path == path and f.available == available, features)
+            )
+            self.assertIsNotNone(feature)
+            self.assertEqual(len(feature), 1)
+
+        check_feature("POST", "/cluster/", False)
+        check_feature("POST", "/diffexp/", not self.data.config.disable_diffexp)
+        check_feature("GET", "/layout/obs", True)
+        check_feature("PUT", "/layout/obs", self.data.config.enable_reembedding)
+        check_feature("PUT", "/annotations/obs", False)
 
     def test_layout(self):
         fbs = self.data.layout_to_fbs_matrix()
@@ -185,3 +200,39 @@ class AdaptorTest(unittest.TestCase):
         self.assertEqual(data["n_rows"], 2638)
         self.assertEqual(data["n_cols"], 3)
         self.assertTrue((data["col_idx"] == [15, 1818, 1837]).all())
+
+    def test_compute_embedding(self):
+        filter = {"obs": {"index": [[0, 100]]}}
+
+        # Verify that we correctly handle the case where we lack scanpy
+        import unittest.mock
+
+        with unittest.mock.patch.dict(sys.modules, {"scanpy": None}):
+            with self.assertRaises(NotImplementedError):
+                self.data.compute_embedding("umap", filter)
+
+        # if we happen to have scanpy, test the full API, else punt
+        import importlib
+
+        scanpy_spec = importlib.util.find_spec("scanpy")
+        if scanpy_spec is None:
+            print("Skipping compute_embedding test as ScanPy not installed")
+            return
+
+        # this feature is unsupported in backed mode, and we expect an error
+        if self.data.data.isbacked:
+            with self.assertRaises(NotImplementedError):
+                self.data.compute_embedding("umap", filter)
+            return
+
+        (schema, fbs) = self.data.compute_embedding("umap", filter)
+
+        self.assertIsInstance(schema["name"], str)
+        name = schema["name"]
+        self.assertEqual(schema["type"], "float32")
+        self.assertEqual(schema["dims"], [f"{name}_0", f"{name}_1"])
+
+        emb = decode_fbs.decode_matrix_FBS(fbs)
+        self.assertEqual(emb["n_rows"], 100)
+        self.assertEqual(emb["n_cols"], 2)
+        self.assertEqual(emb["col_idx"], [f"{name}_0", f"{name}_1"])

server/test/test_api.py

@@ -44,7 +44,7 @@ class EndPoints(object):
         result_data = result.json()
         self.assertIn("library_versions", result_data["config"])
         self.assertEqual(result_data["config"]["displayNames"]["dataset"], "pbmc3k")
-        self.assertEqual(len(result_data["config"]["features"]), 3)
+        self.assertEqual(len(result_data["config"]["features"]), 5)
 
     def test_get_layout_fbs(self):
         endpoint = "layout/obs"

server/test/test_writable_annotation.py

@@ -136,3 +136,20 @@ class WritableAnnotationTest(unittest.TestCase):
             all_col_schema["cat_B"],
             {"name": "cat_B", "type": "categorical", "categories": ["label_B"], "writable": True},
         )
+
+    def test_config(self):
+        features = self.data.get_features(self.annotations)
+
+        # test each for singular presence and accuracy of available flag
+        def check_feature(method, path, available):
+            feature = list(
+                filter(lambda f: f.method == method and f.path == path and f.available == available, features)
+            )
+            self.assertIsNotNone(feature)
+            self.assertEqual(len(feature), 1)
+
+        check_feature("POST", "/cluster/", False)
+        check_feature("POST", "/diffexp/", not self.data.config.disable_diffexp)
+        check_feature("GET", "/layout/obs", True)
+        check_feature("PUT", "/layout/obs", self.data.config.enable_reembedding)
+        check_feature("PUT", "/annotations/obs", True)