Add --query-layers-key option for all algorithms

containers/azimuth/context/main.py

@@ -17,6 +17,7 @@ class AzimuthOrganMetadata(t.TypedDict):
 
 class AzimuthOptions(t.TypedDict):
     reference_data_dir: Path
+    query_layers_key: t.Optional[str]
 
 
 class AzimuthAlgorithm(Algorithm[AzimuthOrganMetadata, AzimuthOptions]):
@@ -37,46 +38,92 @@ def do_run(
         # obs columns of dtype 'object'. As a workaround we create a
         # clean matrix without obs columns on which azimuth is run
         # after which the annotations are copied back to the original matrix
+        temp_index = self.create_temp_obs_index(data)
         clean_matrix_path = Path("clean_matrix.h5ad")
-        clean_matrix = self.create_clean_matrix(data)
+        clean_matrix = self.create_clean_matrix(data, temp_index)
+
+        self.set_data_layer(clean_matrix, options["query_layers_key"])
         clean_matrix.write_h5ad(clean_matrix_path)
 
         annotated_matrix_path = self.run_azimuth_scripts(
             clean_matrix_path, reference_data
         )
         annotated_matrix = anndata.read_h5ad(annotated_matrix_path)
-        self.copy_annotations(data, annotated_matrix)
+        self.copy_annotations(data, annotated_matrix, temp_index)
 
         return {
             "data": data,
             "organ_level": metadata["organ_level"],
             "prediction_column": "predicted." + metadata["prediction_column"],
         }
 
-    def create_clean_matrix(self, matrix: anndata.AnnData) -> anndata.AnnData:
+    def create_temp_obs_index(self, matrix: anndata.AnnData) -> pandas.Index:
+        """Creates a new index by adding a prefix to each index name.
+        Used as a workaround for: https://github.com/satijalab/azimuth/issues/178
+        and https://github.com/satijalab/azimuth/issues/138
+
+        Args:
+            matrix (anndata.AnnData): Original data
+
+        Returns:
+            pandas.Index: A new index
+        """
+        return matrix.obs.index.map(lambda name: f"QUERY:{name}")
+
+    def create_clean_matrix(
+        self,
+        matrix: anndata.AnnData,
+        temp_index: pandas.Index,
+    ) -> anndata.AnnData:
         """Creates a copy of the data with all observation columns removed.
 
         Args:
             matrix (anndata.AnnData): Original data
+            temp_index (pandas.Index): Temporary index generated by `create_temp_obs_index`
 
         Returns:
             anndata.AnnData: Cleaned data
         """
-        clean_obs = pandas.DataFrame(index=matrix.obs.index)
+        clean_obs = pandas.DataFrame(index=temp_index)
         clean_matrix = matrix.copy()
         clean_matrix.obs = clean_obs
+
         return clean_matrix
 
+    def set_data_layer(
+        self, matrix: anndata.AnnData, query_layers_key: t.Optional[str]
+    ) -> None:
+        """Set the data layer to use for annotating.
+
+        Args:
+            matrix (anndata.AnnData): Matrix to update
+            query_layers_key (t.Optional[str]): A layer name or 'raw'
+        """
+        if query_layers_key == "raw":
+            matrix.X = matrix.raw.X
+        elif query_layers_key is not None:
+            matrix.X = matrix.layers[query_layers_key].copy()
+
     def copy_annotations(
-        self, matrix: anndata.AnnData, annotated_matrix: anndata.AnnData
+        self,
+        matrix: anndata.AnnData,
+        annotated_matrix: anndata.AnnData,
+        temp_index: pandas.Index,
     ) -> None:
         """Copies annotations from one matrix to another.
 
         Args:
             matrix (anndata.AnnData): Matrix to copy to
             annotated_matrix (anndata.AnnData): Matrix to copy from
+            temp_index (pandas.Index): Temporary index generated by `create_temp_obs_index`
         """
-        matrix.obs = matrix.obs.join(annotated_matrix.obs, rsuffix="_azimuth")
+        matrix.obs = matrix.obs.merge(
+            annotated_matrix.obs,
+            how="left",
+            left_on=temp_index,
+            right_index=True,
+            suffixes=(None, "_azimuth"),
+        )
 
     def run_azimuth_scripts(self, matrix_path: Path, reference_data: Path) -> str:
         """Creates a subprocess running the Azimuth annotation R script.
@@ -156,6 +203,7 @@ def _get_arg_parser():
         required=True,
         help="Path to directory with reference data",
     )
+    parser.add_argument("--query-layers-key", help="Data layer to use")
 
     return parser
 

containers/azimuth/options.yml

@@ -7,3 +7,8 @@ fields:
     label: Directory with reference data directories
     inputBinding:
       prefix: --reference-data-dir
+  queryLayersKey:
+    type: string?
+    label: Data layer to use
+    inputBinding:
+      prefix: --query-layers-key

containers/celltypist/context/main.py

@@ -16,6 +16,7 @@ class CelltypistOrganMetadata(t.TypedDict):
 
 class CelltypistOptions(t.TypedDict):
     ensemble_lookup: Path
+    query_layers_key: t.Optional[str]
 
 
 class CelltypistAlgorithm(Algorithm[CelltypistOrganMetadata, CelltypistOptions]):
@@ -31,6 +32,7 @@ def do_run(
     ) -> RunResult:
         """Annotate data using celltypist."""
         data = scanpy.read_h5ad(matrix)
+        self.set_data_layer(data, options["query_layers_key"])
         data = self.normalize(data)
         data, var_names = self.normalize_var_names(data, options)
         data = celltypist.annotate(
@@ -40,6 +42,20 @@ def do_run(
 
         return {"data": data, "organ_level": metadata["model"].replace(".", "_")}
 
+    def set_data_layer(
+        self, matrix: scanpy.AnnData, query_layers_key: t.Optional[str]
+    ) -> None:
+        """Set the data layer to use for annotating.
+
+        Args:
+            matrix (anndata.AnnData): Matrix to update
+            query_layers_key (t.Optional[str]): A layer name or 'raw'
+        """
+        if query_layers_key == "raw":
+            matrix.X = matrix.raw.X
+        elif query_layers_key is not None:
+            matrix.X = matrix.layers[query_layers_key].copy()
+
     def normalize(self, data: scanpy.AnnData) -> scanpy.AnnData:
         """Normalizes data according to celltypist requirements.
 
@@ -115,6 +131,7 @@ def _get_arg_parser():
         default="/ensemble-lookup.csv",
         help="Ensemble id to gene name csv",
     )
+    parser.add_argument("--query-layers-key", help="Data layer to use")
 
     return parser
 

containers/celltypist/options.yml

@@ -1,4 +1,9 @@
 type: record
 name: options
 label: Celltypist specific options
-fields: {}
+fields:
+  queryLayersKey:
+    type: string?
+    label: Data layer to use
+    inputBinding:
+      prefix: --query-layers-key

containers/popv/context/main.py

@@ -299,9 +299,7 @@ def _get_arg_parser():
         required=True,
         help="Path to models directory",
     )
-    parser.add_argument(
-        "--query-layers-key", required=True, help="Name of layer with raw counts"
-    )
+    parser.add_argument("--query-layers-key", help="Name of layer with raw counts")
     parser.add_argument("--prediction-mode", default="fast", help="Prediction mode")
     parser.add_argument(
         "--cell-ontology-dir",

containers/popv/options.yml

@@ -13,7 +13,7 @@ fields:
     inputBinding:
       prefix: --models-dir
   queryLayersKey:
-    type: string
+    type: string?
     inputBinding:
       prefix: --query-layers-key
   predictionMode: