re-implement re-embeddings (#1679)

* fix mispelling

* re-implement re-embedding

* always load base embedding to fetch counts

* format

* lint

* fix tests

* lint

* fix accept handling

* test log

* more debug

* more

* more

* more

* more

* remove logging

* logging

* jsonify

* remove debugging logs

* lint

* clean up errors a bit

* fix issue found in PR review

* PR review changes

server/common/rest.py

@@ -304,10 +304,6 @@ def layout_obs_put(request, data_adaptor):
     if not data_adaptor.dataset_config.embeddings__enable_reembedding:
         return abort(HTTPStatus.NOT_IMPLEMENTED)
 
-    preferred_mimetype = request.accept_mimetypes.best_match(["application/octet-stream"])
-    if preferred_mimetype != "application/octet-stream":
-        return abort(HTTPStatus.NOT_ACCEPTABLE)
-
     args = request.get_json()
     filter = args["filter"] if args else None
     if not filter:
@@ -315,17 +311,9 @@ def layout_obs_put(request, data_adaptor):
     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",
-            },
-        )
+        schema = data_adaptor.compute_embedding(method, filter)
+        return make_response(jsonify(schema), HTTPStatus.OK, {"Content-Type": "application/json"})
     except NotImplementedError as e:
-        return abort_and_log(HTTPStatus.NOT_IMPLEMENTED, str(e), include_exc_info=True)
+        return abort_and_log(HTTPStatus.NOT_IMPLEMENTED, str(e))
     except (ValueError, DisabledFeatureError, FilterError) as e:
         return abort_and_log(HTTPStatus.BAD_REQUEST, str(e), include_exc_info=True)

server/compute/scanpy.py

@@ -1,4 +1,5 @@
 import importlib
+import numpy as np
 
 """
 Wrapper for various scanpy modules.  Will raise NotImplementedError if the scanpy
@@ -11,8 +12,8 @@ def get_scanpy_module():
         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 ModuleNotFoundError as e:
+        raise NotImplementedError("Please install scanpy to enable UMAP re-embedding") from e
     except Exception as e:
         # will capture other ImportError corner cases
         raise NotImplementedError() from e
@@ -46,4 +47,7 @@ def scanpy_umap(adata, obs_mask=None, pca_options={}, neighbors_options={}, umap
     sc.pp.neighbors(adata, **neighbors_options)
     sc.tl.umap(adata, **umap_options)
 
-    return adata.obsm["X_umap"]
+    umap = adata.obsm["X_umap"]
+    result = np.full((obs_mask.shape[0], umap.shape[1]), np.NaN)
+    result[obs_mask] = umap
+    return result

server/data_anndata/anndata_adaptor.py

@@ -1,7 +1,6 @@
 import warnings
 
 import numpy as np
-import pandas as pd
 from pandas.core.dtypes.dtypes import CategoricalDtype
 import anndata
 from scipy import sparse
@@ -314,16 +313,15 @@ class AnndataAdaptor(DataAdaptor):
             raise FilterError("Error parsing filter")
         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.
+        # embedding name generated by this backend.
         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)
+        layout_schema = {"name": name, "type": "float32", "dims": dims}
+        self.schema["layout"]["obs"].append(layout_schema)
+        self.data.obsm[f"X_{name}"] = X_umap
+        return layout_schema
 
     def compute_diffexp_ttest(self, maskA, maskB, top_n=None, lfc_cutoff=None):
         if top_n is None:

server/data_common/data_adaptor.py

@@ -71,8 +71,7 @@ class DataAdaptor(metaclass=ABCMeta):
 
     @abstractmethod
     def compute_embedding(self, method, filter):
-        """compute a new embedding on the specified obs subset, and return a
-           tuple of (schema, fbs)."""
+        """compute a new embedding on the specified obs subset, and return the embedding schema. """
         pass
 
     @abstractmethod

server/test/test_anndata_adaptor.py

@@ -237,14 +237,14 @@ class AdaptorTest(unittest.TestCase):
                 self.data.compute_embedding("umap", filter)
             return
 
-        (schema, fbs) = self.data.compute_embedding("umap", filter)
+        schema = 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"])
+        emb = self.data.data.obsm[f"X_{name}"]
+        self.assertEqual(emb.shape, (2638, 2))
+        self.assertTrue(np.isfinite(emb[0:100]).all())
+        self.assertTrue(np.isnan(emb[100:]).all())

server/test/test_api.py

@@ -73,21 +73,23 @@ class EndPoints(object):
         # attempt to reembed with umap over 100 cells.
         endpoint = "layout/obs"
         url = f"{self.URL_BASE}{endpoint}"
-        header = {"Accept": "application/octet-stream"}
         data = {}
         data["filter"] = {}
         data["filter"]["obs"] = {}
         data["filter"]["obs"]["index"] = list(range(100))
         data["method"] = "umap"
-        result = self.session.put(url, headers=header, json=data)
+        result = self.session.put(url, json=data)
 
         self.assertEqual(result.status_code, HTTPStatus.OK)
-        df = decode_fbs.decode_matrix_FBS(result.content)
-        self.assertEqual(df["n_rows"], 100)
-        self.assertEqual(df["n_cols"], 2)
-        cols = list(df["col_idx"])
-        self.assertTrue(cols[0].startswith("reembed:umap_") and cols[0].endswith("_0"))
-        self.assertTrue(cols[1].startswith("reembed:umap_") and cols[1].endswith("_1"))
+        result_data = result.json()
+        self.assertIsInstance(result_data, dict)
+        self.assertEqual(result_data["type"], "float32")
+        self.assertTrue(result_data["name"].startswith("reembed:umap_"))
+        self.assertIsInstance(result_data["dims"], list)
+        self.assertEqual(len(result_data["dims"]), 2)
+        dims = result_data["dims"]
+        self.assertTrue(dims[0].startswith("reembed:umap_") and dims[0].endswith("_0"))
+        self.assertTrue(dims[1].startswith("reembed:umap_") and dims[1].endswith("_1"))
 
     def test_bad_filter(self):
         endpoint = "data/var"