drop numba dependency

Author: niklasmueboe

docs/source/tutorials/VisiumHD_CRC.ipynb

@@ -647,7 +647,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": null,
    "id": "672b4258-dc91-4fb0-9158-eeb0870c83c9",
    "metadata": {
     "editable": true,
@@ -664,14 +664,11 @@
     "from matplotlib.colors import to_rgb\n",
     "from scipy.sparse import coo_array\n",
     "\n",
-    "from sainsc.lazykde._utils import _apply_color\n",
-    "\n",
     "background = \"black\"\n",
     "\n",
     "cmap_ls = [cmap[celltype] for celltype in adata_16um.obs[\"leiden\"].cat.categories]\n",
-    "color_map = tuple(\n",
-    "    (np.array(to_rgb(c)) * 255).round().astype(np.uint8) for c in [background] + cmap_ls\n",
-    ")\n",
+    "color_map = np.array([to_rgb(c) for c in [background] + cmap_ls])\n",
+    "\n",
     "\n",
     "celltype = adata_16um.obs[\"leiden\"].cat.codes\n",
     "x = adata_16um.obs[\"array_row\"]\n",
@@ -681,7 +678,7 @@
     "ctmap_16um = coo_array(\n",
     "    (celltype + 1, (x, y)), shape=[ceil(i / 8) for i in visium_hd.shape]\n",
     ")\n",
-    "ctmap_16um = _apply_color((ctmap_16um).toarray(), color_map)"
+    "ctmap_16um = np.take(color_map, ctmap_16um.toarray(), axis=0)"
    ]
   },
   {

pyproject.toml

@@ -19,7 +19,6 @@ dependencies = [
     "h5py>=3",
     "matplotlib>=3.7",
     "matplotlib-scalebar",
-    "numba>=0.44",
     "numpy>=1.24",
     "pandas>=2.0",
     "polars[pandas]>=1",

sainsc/lazykde/_LazyKDE.py

@@ -1,7 +1,6 @@
 from collections.abc import Iterable
-from itertools import chain
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Self
+from typing import TYPE_CHECKING, Any, Self, TypeVar
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -15,7 +14,6 @@
 from matplotlib.patches import Patch
 from matplotlib_scalebar.scalebar import ScaleBar
 from mpl_toolkits import axes_grid1
-from numba import njit
 from numpy.typing import NDArray
 from scipy.sparse import coo_array, csc_array, csr_array
 from skimage.feature import peak_local_max
@@ -33,7 +31,6 @@
 from ._utils import (
     SCALEBAR_PARAMS,
     CosineCelltypeCallable,
-    _apply_color,
     _filter_blobs,
     _get_cell_dtype,
     _localmax_anndata,
@@ -358,7 +355,6 @@ def load_local_maxima(
                 }
 
                 if self.total_mRNA_KDE is not None:
-
                     sdata_dict["total_mRNA"] = Image2DModel.parse(
                         np.atleast_3d(self.total_mRNA_KDE).T, dims=("c", "y", "x")
                     )
@@ -416,7 +412,6 @@ def load_local_maxima(
             return adata
 
     def _load_KDE_maxima(self, genes: list[str]) -> csc_array | csr_array:
-
         assert self.local_maxima is not None
         if self.kernel is None:
             raise ValueError("`kernel` must be set before running KDE")
@@ -461,17 +456,13 @@ def filter_background(
             If cell type-specific thresholds do not include all cell types or if
             using cell type-specific thresholds before cell type assignment.
         """
+        T = TypeVar("T")
 
-        @njit
         def _map_celltype_to_value(
-            ct_map: NDArray[np.integer], thresholds: tuple[float, ...]
+            ct_map: NDArray[np.integer], thresholds: dict[T, float], classes: list[T]
         ) -> NDArray[np.floating]:
-            values = np.zeros(shape=ct_map.shape, dtype=float)
-            for i in range(ct_map.shape[0]):
-                for j in range(ct_map.shape[1]):
-                    if ct_map[i, j] >= 0:
-                        values[i, j] = thresholds[ct_map[i, j]]
-            return values
+            ordered_thresholds = np.array([0] + [thresholds[ct] for ct in classes])
+            return np.take(ordered_thresholds, ct_map + 1)
 
         if self.total_mRNA_KDE is None:
             raise ValueError(
@@ -485,8 +476,9 @@ def _map_celltype_to_value(
                 )
             elif not all([ct in min_norm.keys() for ct in self.celltypes]):
                 raise ValueError("'min_norm' does not contain all celltypes.")
-            idx2threshold = tuple(min_norm[ct] for ct in self.celltypes)
-            threshold = _map_celltype_to_value(self.celltype_map, idx2threshold)
+            threshold = _map_celltype_to_value(
+                self.celltype_map, min_norm, self.celltypes
+            )
             background = self.total_mRNA_KDE < threshold
         else:
             background = self.total_mRNA_KDE < min_norm
@@ -503,8 +495,9 @@ def _map_celltype_to_value(
                     )
                 elif not all([ct in min_cosine.keys() for ct in self.celltypes]):
                     raise ValueError("'min_cosine' does not contain all celltypes.")
-                idx2threshold = tuple(min_cosine[ct] for ct in self.celltypes)
-                threshold = _map_celltype_to_value(self.celltype_map, idx2threshold)
+                threshold = _map_celltype_to_value(
+                    self.celltype_map, min_cosine, self.celltypes
+                )
                 background |= self.cosine_similarity <= threshold
             else:
                 background |= self.cosine_similarity <= min_cosine
@@ -521,8 +514,9 @@ def _map_celltype_to_value(
                     )
                 elif not all([ct in min_assignment.keys() for ct in self.celltypes]):
                     raise ValueError("'min_assignment' does not contain all celltypes.")
-                idx2threshold = tuple(min_assignment[ct] for ct in self.celltypes)
-                threshold = _map_celltype_to_value(self.celltype_map, idx2threshold)
+                threshold = _map_celltype_to_value(
+                    self.celltype_map, min_assignment, self.celltypes
+                )
                 background |= self.assignment_score <= threshold
             else:
                 background |= self.assignment_score <= min_assignment
@@ -1006,11 +1000,11 @@ def plot_celltype_map(
             color_map = [to_rgb(c) if isinstance(c, str) else c for c in cmap]
 
         # convert to uint8 to reduce memory of final image
-        color_map_int = tuple(
-            (np.array(c) * 255).round().astype(np.uint8)
-            for c in chain([to_rgb(background)], color_map)
+        color_map_int = (
+            (np.array([to_rgb(background)] + color_map) * 255).round().astype(np.uint8)
         )
-        img = _apply_color(celltype_map.T, color_map_int)
+
+        img = np.take(color_map_int, celltype_map.T, axis=0)
 
         if return_img:
             return img

sainsc/lazykde/_utils.py

@@ -5,7 +5,6 @@
 import numpy as np
 import pandas as pd
 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
@@ -20,17 +19,6 @@
 """Default scalebar parameters"""
 
 
-@njit
-def _apply_color(
-    img_in: NDArray[np.integer], cmap: tuple[NDArray[T], ...]
-) -> NDArray[T]:
-    img = np.empty(shape=(*img_in.shape, 3), dtype=cmap[0].dtype)
-    for i in range(img_in.shape[0]):
-        for j in range(img_in.shape[1]):
-            img[i, j, :] = cmap[img_in[i, j]]
-    return img
-
-
 def _get_cell_dtype(n: int) -> np.dtype:
     return np.result_type("int8", n)