load assignment score for localmax

HiDiHlabs · Oct 28, 2024 · 6dee4b7 · 6dee4b7
1 parent dce8532
commit 6dee4b7
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/sainsc/lazykde/_LazyKDE.py b/sainsc/lazykde/_LazyKDE.py
@@ -446,8 +446,18 @@ def load_local_maxima_cosine_similarity(
         -------
         anndata.AnnData
         """
+        # Celltype assignment must have beenn run so the following attributesshoul exist
         assert self.local_maxima is not None
-        return _load_localmax_cosine(self.local_maxima, zarr_path, celltypes=celltypes)
+        assert self.celltypes is not None
+        assert self.celltype_map is not None
+        assert self.assignment_score is not None
+
+        adata = _load_localmax_cosine(self.local_maxima, zarr_path, celltypes=celltypes)
+        adata.obs["celltype"] = pd.Categorical.from_codes(
+            self.celltype_map[self.local_maxima], self.celltypes
+        )
+        adata.obs["assignment_score"] = self.assignment_score[self.local_maxima]
+        return adata
 
     ## Celltyping
     def filter_background(

diff --git a/sainsc/lazykde/_utils.py b/sainsc/lazykde/_utils.py
@@ -84,8 +84,7 @@ def _load_localmax_cosine(
     )
 
     obs = pd.DataFrame(
-        {"celltype": pd.Categorical.from_codes(cosine.argmax(axis=1), celltypes)},
-        index=_get_coordinate_index(*coord, name="local_maxima", n_threads=1),
+        index=_get_coordinate_index(*coord, name="local_maxima", n_threads=1)
     )
     var = pd.DataFrame(index=celltypes)