Write predictions to common column across algorithms

hubmapconsortium · Nov 6, 2023 · d30bc0b · d30bc0b
1 parent b545052
commit d30bc0b
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 5 deletions.
diff --git a/containers/azimuth/context/main.py b/containers/azimuth/context/main.py
@@ -37,7 +37,7 @@ def get_builtin_options(self):
 
 class AzimuthAlgorithm(Algorithm[str, AzimuthOptions]):
     def __init__(self):
-        super().__init__(AzimuthOrganLookup)
+        super().__init__(AzimuthOrganLookup, "azimuth_label")
 
     def do_run(self, matrix: Path, organ: str, options: AzimuthOptions):
         data = anndata.read_h5ad(matrix)

diff --git a/containers/celltypist/context/main.py b/containers/celltypist/context/main.py
@@ -28,7 +28,7 @@ def from_raw(self, id: str):
 
 class CelltypistAlgorithm(Algorithm[celltypist.Model, CelltypistOptions]):
     def __init__(self):
-        super().__init__(CelltypistOrganLookup)
+        super().__init__(CelltypistOrganLookup, "predicted_labels")
 
     def do_run(self, matrix: Path, organ: celltypist.Model, options: CelltypistOptions):
         data = scanpy.read_h5ad(matrix)

diff --git a/containers/popv/context/main.py b/containers/popv/context/main.py
@@ -30,7 +30,7 @@ class PopvOptions(t.TypedDict):
 
 class PopvAlgorithm(Algorithm[str, PopvOptions]):
     def __init__(self):
-        super().__init__(OrganLookup)
+        super().__init__(OrganLookup, "popv_prediction")
 
     def do_run(self, matrix: Path, organ: str, options: PopvOptions):
         data = scanpy.read_h5ad(matrix)

diff --git a/src/algorithm/algorithm.py b/src/algorithm/algorithm.py
@@ -12,8 +12,13 @@
 
 
 class Algorithm(t.Generic[Organ, Options], abc.ABC):
-    def __init__(self, organ_lookup: t.Callable[[Path], OrganLookup[Organ]]):
+    def __init__(
+        self,
+        organ_lookup: t.Callable[[Path], OrganLookup[Organ]],
+        prediction_column: str,
+    ):
         self.organ_lookup = organ_lookup
+        self.prediction_column = prediction_column
 
     def run(
         self,
@@ -26,7 +31,11 @@ def run(
         **options,
     ) -> AlgorithmReport:
         report = AlgorithmReport(
-            output_matrix, output_annotations, output_report, matrix
+            output_matrix,
+            output_annotations,
+            output_report,
+            matrix,
+            self.prediction_column,
         )
         try:
             lookup = self.organ_lookup(organ_mapping)

diff --git a/src/algorithm/report.py b/src/algorithm/report.py
@@ -23,6 +23,7 @@ class AlgorithmReport:
     annotations: Path
     report: Path
     data: t.Union[Path, anndata.AnnData]
+    prediction_column: str
     status = Status.SUCCESS
     failure_cause: t.Any = None
 
@@ -53,6 +54,7 @@ def save_matrix(self):
         if not self.is_success():
             # Create an empty observations frame with the same columns as the original
             obs = pandas.DataFrame(columns=obs.columns)
+        obs['hra_prediction'] = obs[self.prediction_column]
         obs.to_csv(self.annotations)
 
     def save_report(self):