func to load cosine sim of localmax

pyproject.toml

@@ -25,6 +25,7 @@ dependencies = [
     "scikit-image>=0.20",
     "scipy>=1.10",
     "seaborn>=0.11",
+    "zarr>=3.0.0b1",
 ]
 classifiers = [
     "Intended Audience :: Science/Research",

sainsc/lazykde/_LazyKDE.py

@@ -36,6 +36,7 @@
     _apply_color,
     _filter_blobs,
     _get_cell_dtype,
+    _load_localmax_cosine,
     _localmax_anndata,
 )
 
@@ -425,6 +426,29 @@ def _load_KDE_maxima(self, genes: list[str]) -> csc_array | csr_array:
             self.counts, genes, self.kernel, self.local_maxima, n_threads=self.n_threads
         )
 
+    def load_local_maxima_cosine_similarity(
+        self, zarr_path: _PathLike, *, celltypes: Iterable[str] | None = None
+    ) -> pd.DataFrame:
+        """
+        Load the cosine similarity for all cell types for the local maxima.
+
+        The cosine similarity can be loaded from a zarr store the was generated by
+        running :py:meth:`sainsc.LazyKDE.assign_celltype`.
+
+        Parameters
+        ----------
+        zarr_path : os.PathLike or str
+            Path to the zarr store.
+        celltypes : Iterable[str] | None, optional
+            If not `None` will only load the provided cell types.
+
+        Returns
+        -------
+        pandas.DataFrame
+        """
+        assert self.local_maxima is not None
+        return _load_localmax_cosine(self.local_maxima, zarr_path, celltypes=celltypes)
+
     ## Celltyping
     def filter_background(
         self,

sainsc/lazykde/_utils.py

@@ -4,12 +4,14 @@
 
 import numpy as np
 import pandas as pd
+import zarr
 from anndata import AnnData
 from numba import njit
 from numpy.typing import NDArray
 from scipy.sparse import csr_matrix, sparray, spmatrix
 from skimage.measure import label, regionprops
 
+from .._typealias import _Local_Max, _PathLike
 from .._utils import _get_coordinate_index
 from .._utils_rust import GridCounts
 
@@ -67,6 +69,23 @@ def _localmax_anndata(
     )
 
 
+def _load_localmax_cosine(
+    coord: _Local_Max, zarr_store: _PathLike, *, celltypes: Iterable[str] | None = None
+) -> pd.DataFrame:
+    cosine_group = zarr.open_group(store=zarr_store, mode="r", path="cosine")
+
+    celltypes = cosine_group.array_keys() if celltypes is None else celltypes
+
+    cosine_df = pd.DataFrame({"x": coord[0], "y": coord[1]})
+
+    for ct in celltypes:
+        cosine_ct = cosine_group[ct]
+        assert isinstance(cosine_ct, zarr.Array)
+        cosine_df[ct] = cosine_ct.get_coordinate_selection(coord)
+
+    return cosine_df
+
+
 class CosineCelltypeCallable(Protocol):
     def __call__(
         self,