Merge pull request #35 from HiDiHlabs/shape-typing

Shape typing for numpy

Author: niklasmueboe

sainsc/_typealias.py

@@ -2,7 +2,6 @@
 from typing import TypeAlias
 
 import numpy as np
-from numpy.typing import NDArray
 from scipy.sparse import csc_array, csc_matrix, csr_array, csr_matrix
 
 _PathLike: TypeAlias = os.PathLike[str] | str
@@ -15,7 +14,20 @@
 _RangeTuple: TypeAlias = tuple[int, int]
 _RangeTuple2D: TypeAlias = tuple[_RangeTuple, _RangeTuple]
 
-_Local_Max: TypeAlias = tuple[NDArray[np.int_], NDArray[np.int_]]
+
+_Local_Max: TypeAlias = tuple[
+    np.ndarray[tuple[int], np.dtype[np.int_]], np.ndarray[tuple[int], np.dtype[np.int_]]
+]
 
 _Color: TypeAlias = str | tuple[float, float, float]
 _Cmap: TypeAlias = str | list[_Color] | dict[str, _Color]
+
+
+_Coord: TypeAlias = np.ndarray[tuple[int], np.dtype[np.int32]]
+_CosineMap: TypeAlias = np.ndarray[tuple[int, int], np.dtype[np.float32]]
+_AssignmentScoreMap: TypeAlias = np.ndarray[tuple[int, int], np.dtype[np.float32]]
+_Kernel: TypeAlias = np.ndarray[tuple[int, int], np.dtype[np.float32]]
+_SignatureArray: TypeAlias = np.ndarray[tuple[int, int], np.dtype[np.float32]]
+_Background: TypeAlias = np.ndarray[tuple[int, int], np.dtype[np.bool_]]
+_CountMap: TypeAlias = np.ndarray[tuple[int, int], np.dtype[np.unsignedinteger]]
+_KDE: TypeAlias = np.ndarray[tuple[int, int], np.dtype[np.single]]

sainsc/_utils.py

@@ -4,7 +4,6 @@
 
 import numpy as np
 import pandas as pd
-from numpy.typing import NDArray
 
 from ._utils_rust import coordinate_as_string
 
@@ -38,15 +37,18 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
     return wrapper
 
 
+N = TypeVar("N", bound=int)
+
+
 def _get_coordinate_index(
-    x: NDArray[np.integer],
-    y: NDArray[np.integer],
+    x: np.ndarray[tuple[N], np.dtype[np.integer]],
+    y: np.ndarray[tuple[N], np.dtype[np.integer]],
     *,
     name: str | None = None,
     n_threads: int | None = None,
 ) -> pd.Index:
-    x_i32: NDArray[np.int32] = x.astype(np.int32, copy=False)
-    y_i32: NDArray[np.int32] = y.astype(np.int32, copy=False)
+    x_i32: np.ndarray[tuple[N], np.dtype[np.int32]] = x.astype(np.int32, copy=False)
+    y_i32: np.ndarray[tuple[N], np.dtype[np.int32]] = y.astype(np.int32, copy=False)
 
     return pd.Index(
         coordinate_as_string(x_i32, y_i32, n_threads=n_threads), dtype=str, name=name

sainsc/_utils_rust.pyi

@@ -1,13 +1,20 @@
 from typing import Self
 
 import numpy as np
-from numpy.typing import NDArray
 from polars import DataFrame
 
-from ._typealias import _Csx, _CsxArray
+from ._typealias import (
+    _AssignmentScoreMap,
+    _Coord,
+    _CosineMap,
+    _Csx,
+    _CsxArray,
+    _Kernel,
+    _SignatureArray,
+)
 
 def sparse_kde_csx_py(
-    counts: _Csx, kernel: NDArray[np.float32], *, threshold: float = 0
+    counts: _Csx, kernel: _Kernel, *, threshold: float = 0
 ) -> _CsxArray:
     """
     Calculate the KDE for each spot with counts as uint16.
@@ -16,8 +23,11 @@ def sparse_kde_csx_py(
 def kde_at_coord(
     counts: GridCounts,
     genes: list[str],
-    kernel: NDArray[np.float32],
-    coordinates: tuple[NDArray[np.int_], NDArray[np.int_]],
+    kernel: _Kernel,
+    coordinates: tuple[
+        np.ndarray[tuple[int], np.dtype[np.int_]],
+        np.ndarray[tuple[int], np.dtype[np.int_]],
+    ],
     *,
     n_threads: int | None = None,
 ) -> _CsxArray:
@@ -26,29 +36,34 @@ def kde_at_coord(
     """
 
 def categorical_coordinate(
-    x: NDArray[np.int32], y: NDArray[np.int32], *, n_threads: int | None = None
-) -> tuple[NDArray[np.int32], NDArray[np.int32]]:
+    x: _Coord, y: _Coord, *, n_threads: int | None = None
+) -> tuple[
+    np.ndarray[tuple[int], np.dtype[np.int32]],
+    np.ndarray[tuple[int, int], np.dtype[np.int32]],
+]:
     """
     Get the codes and the coordinates (comparable to a pandas.Categorical)
     """
 
 def coordinate_as_string(
-    x: NDArray[np.int32], y: NDArray[np.int32], *, n_threads: int | None = None
-) -> NDArray[np.str_]:
+    x: _Coord, y: _Coord, *, n_threads: int | None = None
+) -> np.ndarray[tuple[int], np.dtype[np.str_]]:
     """
     Concatenate two int arrays elementwise into a string representation (i.e. 'x_y').
     """
 
 def cosinef32_and_celltypei8(
     counts: GridCounts,
     genes: list[str],
-    signatures: NDArray[np.float32],
-    kernel: NDArray[np.float32],
+    signatures: _SignatureArray,
+    kernel: _Kernel,
     *,
     log: bool = False,
     chunk_size: tuple[int, int] = (500, 500),
     n_threads: int | None = None,
-) -> tuple[NDArray[np.float32], NDArray[np.float32], NDArray[np.int8]]:
+) -> tuple[
+    _CosineMap, _AssignmentScoreMap, np.ndarray[tuple[int, int], np.dtype[np.int8]]
+]:
     """
     Calculate the cosine similarity given counts and signatures and assign the most
     similar celltype.
@@ -57,13 +72,15 @@ def cosinef32_and_celltypei8(
 def cosinef32_and_celltypei16(
     counts: GridCounts,
     genes: list[str],
-    signatures: NDArray[np.float32],
-    kernel: NDArray[np.float32],
+    signatures: _SignatureArray,
+    kernel: _Kernel,
     *,
     log: bool = False,
     chunk_size: tuple[int, int] = (500, 500),
     n_threads: int | None = None,
-) -> tuple[NDArray[np.float32], NDArray[np.float32], NDArray[np.int16]]:
+) -> tuple[
+    _CosineMap, _AssignmentScoreMap, np.ndarray[tuple[int, int], np.dtype[np.int16]]
+]:
     """
     Calculate the cosine similarity given counts and signatures and assign the most
     similar celltype.
@@ -186,7 +203,7 @@ class GridCounts:
             Mapping from gene to number of counts.
         """
 
-    def grid_counts(self) -> NDArray[np.uintc]:
+    def grid_counts(self) -> np.ndarray[tuple[int, int], np.dtype[np.uintc]]:
         """
         Counts per pixel.
 
@@ -231,7 +248,7 @@ class GridCounts:
             Range to crop as `(ymin, ymax)`
         """
 
-    def filter_mask(self, mask: NDArray[np.bool_]):
+    def filter_mask(self, mask: np.ndarray[tuple[int, int], np.dtype[np.bool_]]):
         """
         Filter all genes with a binary mask.
 

sainsc/datasets.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import TYPE_CHECKING
 
 import pandas as pd
@@ -13,7 +15,7 @@
 version = "v" + __version__
 
 
-def _get_signature_pooch() -> "Pooch":
+def _get_signature_pooch() -> Pooch:
     # use indirection to enable pooch as optional dependency w/o lazy loading
     try:
         import pooch

sainsc/io/_io.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from collections.abc import Collection
 from pathlib import Path
 from typing import TYPE_CHECKING, Literal, get_args
@@ -457,7 +459,7 @@ def read_StereoSeq_bins(
     sep: str = "\t",
     n_threads: int | None = None,
     **kwargs,
-) -> "AnnData | SpatialData":
+) -> AnnData | SpatialData:
     """
     Read a Stereo-seq GEM file into bins.
 
@@ -521,7 +523,10 @@ def read_StereoSeq_bins(
 
     obs = pd.DataFrame(
         index=_get_coordinate_index(
-            coordinates[:, 0], coordinates[:, 1], name="bin", n_threads=n_threads
+            coordinates[:, 0],  # type: ignore
+            coordinates[:, 1],  # type: ignore
+            name="bin",
+            n_threads=n_threads,
         )
     )
 

sainsc/io/_io_utils.py

@@ -5,9 +5,8 @@
 
 import numpy as np
 import polars as pl
-from numpy.typing import NDArray
 
-from .._typealias import _PathLike
+from .._typealias import _Coord, _PathLike
 from .._utils_rust import categorical_coordinate
 
 
@@ -20,8 +19,11 @@ def _bin_coordinates(df: pl.DataFrame, bin_size: float) -> pl.DataFrame:
 
 
 def _categorical_coordinate(
-    x: NDArray[np.int32], y: NDArray[np.int32], *, n_threads: int | None = None
-) -> tuple[NDArray[np.int32], NDArray[np.int32]]:
+    x: _Coord, y: _Coord, *, n_threads: int | None = None
+) -> tuple[
+    np.ndarray[tuple[int], np.dtype[np.int32]],
+    np.ndarray[tuple[int, int], np.dtype[np.int32]],
+]:
     assert len(x) == len(y)
 
     return categorical_coordinate(x, y, n_threads=n_threads)