load embeddings in parallel (#1352)

* load embeddings in parallel

* correctly capture unclipped

* test

* another test

* add convenient copy assets target

* cleanup

server/common/rest.py

@@ -10,6 +10,7 @@ from server.common.errors import (
     PrepareError,
     DisabledFeatureError,
     ExceedsLimitError,
+    DatasetAccessError,
 )
 
 import json
@@ -179,14 +180,21 @@ def diffexp_obs_post(request, data_adaptor):
 
 
 def layout_obs_get(request, data_adaptor):
+    fields = request.args.getlist("layout-name", None)
+    num_columns_requested = len(data_adaptor.get_embedding_names()) if len(fields) == 0 else len(fields)
+    if data_adaptor.config.exceeds_limit("column_request_max", num_columns_requested):
+        return abort(HTTPStatus.BAD_REQUEST)
+
     preferred_mimetype = request.accept_mimetypes.best_match(["application/octet-stream"])
     if preferred_mimetype != "application/octet-stream":
         return abort(HTTPStatus.NOT_ACCEPTABLE)
 
     try:
         return make_response(
-            data_adaptor.layout_to_fbs_matrix(), HTTPStatus.OK, {"Content-Type": "application/octet-stream"}
+            data_adaptor.layout_to_fbs_matrix(fields), HTTPStatus.OK, {"Content-Type": "application/octet-stream"}
         )
+    except (KeyError, DatasetAccessError) as e:
+        return abort_and_log(HTTPStatus.BAD_REQUEST, str(e), include_exc_info=True)
     except PrepareError:
         return abort_and_log(
             HTTPStatus.NOT_IMPLEMENTED,

server/data_common/data_adaptor.py

@@ -331,10 +331,9 @@ class DataAdaptor(metaclass=ABCMeta):
         normalized_layout = normalized_layout.astype(dtype=np.float32)
         return normalized_layout
 
-    def layout_to_fbs_matrix(self):
-        """ same as layout, except returns a flatbuffer """
+    def layout_to_fbs_matrix(self, fields):
         """
-        return all embeddings as a flatbuffer, using the cellxgene matrix fbs encoding.
+        return specified embeddings as a flatbuffer, using the cellxgene matrix fbs encoding.
 
         * returns only first two dimensions, with name {ename}_0 and {ename}_1,
           where {ename} is the embedding name.
@@ -343,8 +342,7 @@ class DataAdaptor(metaclass=ABCMeta):
         * does not support filtering
 
         """
-
-        embeddings = self.get_embedding_names()
+        embeddings = self.get_embedding_names() if fields is None or len(fields) == 0 else fields
         layout_data = []
         with ServerTiming.time(f"layout.query"):
             for ename in embeddings:

server/test/schema.json

@@ -67,6 +67,16 @@
         "name": "umap",
         "type": "float32",
         "dims": ["umap_0", "umap_1"]
+      },
+      {
+        "name": "tsne",
+        "type": "float32",
+        "dims": ["tsne_0", "tsne_1"]
+      },
+      {
+        "name": "pca",
+        "type": "float32",
+        "dims": ["pca_0", "pca_1"]
       }
     ]
   }

server/test/test_anndata_adaptor.py

@@ -34,7 +34,7 @@ Test the anndata adaptor using the pbmc3k data set.
 class AdaptorTest(unittest.TestCase):
     def setUp(self):
         args = {
-            "embeddings__names": ["umap"],
+            "embeddings__names": ["umap", "tsne", "pca"],
             "presentation__max_categories": 100,
             "single_dataset__obs_names": None,
             "single_dataset__var_names": None,
@@ -122,9 +122,9 @@ class AdaptorTest(unittest.TestCase):
         check_feature("PUT", "/annotations/obs", False)
 
     def test_layout(self):
-        fbs = self.data.layout_to_fbs_matrix()
+        fbs = self.data.layout_to_fbs_matrix(fields=None)
         layout = decode_fbs.decode_matrix_FBS(fbs)
-        self.assertEqual(layout["n_cols"], 2)
+        self.assertEqual(layout["n_cols"], 6)
         self.assertEqual(layout["n_rows"], 2638)
 
         X = layout["columns"][0]
@@ -132,6 +132,20 @@ class AdaptorTest(unittest.TestCase):
         Y = layout["columns"][1]
         self.assertTrue((Y >= 0).all() and (Y <= 1).all())
 
+    def test_layout_fields(self):
+        """ X_pca, X_tsne, X_umap are available """
+        fbs = self.data.layout_to_fbs_matrix(["pca"])
+        layout = decode_fbs.decode_matrix_FBS(fbs)
+        self.assertEqual(layout["n_cols"], 2)
+        self.assertEqual(layout["n_rows"], 2638)
+        self.assertCountEqual(layout["col_idx"], ["pca_0", "pca_1"])
+
+        fbs = self.data.layout_to_fbs_matrix(["tsne", "pca"])
+        layout = decode_fbs.decode_matrix_FBS(fbs)
+        self.assertEqual(layout["n_cols"], 4)
+        self.assertEqual(layout["n_rows"], 2638)
+        self.assertCountEqual(layout["col_idx"], ["tsne_0", "tsne_1", "pca_0", "pca_1"])
+
     def test_annotations(self):
         fbs = self.data.annotation_to_fbs_matrix("obs")
         annotations = decode_fbs.decode_matrix_FBS(fbs)