Backport CS* tightening (#1889)

Author: flying-sheep

src/anndata/_core/aligned_mapping.py

@@ -9,10 +9,9 @@
 
 import numpy as np
 import pandas as pd
-from scipy.sparse import spmatrix
 
 from .._warnings import ExperimentalFeatureWarning, ImplicitModificationWarning
-from ..compat import AwkArray
+from ..compat import AwkArray, CSArray, CSMatrix
 from ..utils import (
     axis_len,
     convert_to_dict,
@@ -36,7 +35,7 @@
 OneDIdx = Sequence[int] | Sequence[bool] | slice
 TwoDIdx = tuple[OneDIdx, OneDIdx]
 # TODO: pd.DataFrame only allowed in AxisArrays?
-Value = pd.DataFrame | spmatrix | np.ndarray
+Value = pd.DataFrame | CSArray | CSMatrix | np.ndarray
 
 P = TypeVar("P", bound="AlignedMappingBase")
 """Parent mapping an AlignedView is based on."""

src/anndata/_core/anndata.py

@@ -195,13 +195,13 @@ class AnnData(metaclass=utils.DeprecationMixinMeta):
 
     def __init__(
         self,
-        X: np.ndarray | sparse.spmatrix | pd.DataFrame | None = None,
+        X: ArrayDataStructureType | pd.DataFrame | None = None,
         obs: pd.DataFrame | Mapping[str, Iterable[Any]] | None = None,
         var: pd.DataFrame | Mapping[str, Iterable[Any]] | None = None,
         uns: Mapping[str, Any] | None = None,
         obsm: np.ndarray | Mapping[str, Sequence[Any]] | None = None,
         varm: np.ndarray | Mapping[str, Sequence[Any]] | None = None,
-        layers: Mapping[str, np.ndarray | sparse.spmatrix] | None = None,
+        layers: Mapping[str, ArrayDataStructureType] | None = None,
         raw: Mapping[str, Any] | None = None,
         dtype: np.dtype | type | str | None = None,
         shape: tuple[int, int] | None = None,
@@ -557,7 +557,7 @@ def X(self) -> ArrayDataStructureType | None:
         #     return X
 
     @X.setter
-    def X(self, value: np.ndarray | sparse.spmatrix | SpArray | None):
+    def X(self, value: ArrayDataStructureType | None):
         if value is None:
             if self.isbacked:
                 msg = "Cannot currently remove data matrix from backed object."
@@ -1159,7 +1159,7 @@ def _inplace_subset_obs(self, index: Index1D):
         self._init_as_actual(adata_subset)
 
     # TODO: Update, possibly remove
-    def __setitem__(self, index: Index, val: float | np.ndarray | sparse.spmatrix):
+    def __setitem__(self, index: Index, val: float | ArrayDataStructureType):
         if self.is_view:
             msg = "Object is view and cannot be accessed with `[]`."
             raise ValueError(msg)

src/anndata/_core/index.py

@@ -13,7 +13,7 @@
 from ..compat import AwkArray, DaskArray, SpArray
 
 if TYPE_CHECKING:
-    from ..compat import Index, Index1D
+    from ..compat import CSArray, CSMatrix, Index, Index1D
 
 
 def _normalize_indices(
@@ -182,7 +182,7 @@ def _subset_dask(a: DaskArray, subset_idx: Index):
 
 @_subset.register(spmatrix)
 @_subset.register(SpArray)
-def _subset_sparse(a: spmatrix | SpArray, subset_idx: Index):
+def _subset_sparse(a: CSMatrix | CSArray, subset_idx: Index):
     # Correcting for indexing behaviour of sparse.spmatrix
     if len(subset_idx) > 1 and all(isinstance(x, Iterable) for x in subset_idx):
         first_idx = subset_idx[0]

src/anndata/_core/merge.py

@@ -41,6 +41,8 @@
 
     from pandas.api.extensions import ExtensionDtype
 
+    from ..compat import CSArray, CSMatrix
+
 T = TypeVar("T")
 
 ###################
@@ -174,7 +176,7 @@ def equal_sparse(a, b) -> bool:
 
     xp = array_api_compat.array_namespace(a.data)
 
-    if isinstance(b, CupySparseMatrix | sparse.spmatrix | SpArray):
+    if isinstance(b, CupySparseMatrix | spmatrix | SpArray):
         if isinstance(a, CupySparseMatrix):
             # Comparison broken for CSC matrices
             # https://github.com/cupy/cupy/issues/7757
@@ -205,8 +207,8 @@ def equal_awkward(a, b) -> bool:
     return ak.almost_equal(a, b)
 
 
-def as_sparse(x, use_sparse_array=False):
-    if not isinstance(x, sparse.spmatrix | SpArray):
+def as_sparse(x, *, use_sparse_array: bool = False) -> CSMatrix | CSArray:
+    if not isinstance(x, spmatrix | SpArray):
         if CAN_USE_SPARSE_ARRAY and use_sparse_array:
             return sparse.csr_array(x)
         return sparse.csr_matrix(x)
@@ -537,7 +539,7 @@ def apply(self, el, *, axis, fill_value=None):
             return el
         if isinstance(el, pd.DataFrame):
             return self._apply_to_df(el, axis=axis, fill_value=fill_value)
-        elif isinstance(el, sparse.spmatrix | SpArray | CupySparseMatrix):
+        elif isinstance(el, spmatrix | SpArray | CupySparseMatrix):
             return self._apply_to_sparse(el, axis=axis, fill_value=fill_value)
         elif isinstance(el, AwkArray):
             return self._apply_to_awkward(el, axis=axis, fill_value=fill_value)
@@ -615,7 +617,7 @@ def _apply_to_array(self, el, *, axis, fill_value=None):
         )
 
     def _apply_to_sparse(
-        self, el: sparse.spmatrix | SpArray, *, axis, fill_value=None
+        self, el: CSMatrix | CSArray, *, axis, fill_value=None
     ) -> spmatrix:
         if isinstance(el, CupySparseMatrix):
             from cupyx.scipy import sparse
@@ -730,11 +732,8 @@ def default_fill_value(els):
     This is largely due to backwards compat, and might not be the ideal solution.
     """
     if any(
-        isinstance(el, sparse.spmatrix | SpArray)
-        or (
-            isinstance(el, DaskArray)
-            and isinstance(el._meta, sparse.spmatrix | SpArray)
-        )
+        isinstance(el, spmatrix | SpArray)
+        or (isinstance(el, DaskArray) and isinstance(el._meta, spmatrix | SpArray))
         for el in els
     ):
         return 0
@@ -830,7 +829,7 @@ def concat_arrays(arrays, reindexers, axis=0, index=None, fill_value=None):
             ],
             axis=axis,
         )
-    elif any(isinstance(a, sparse.spmatrix | SpArray) for a in arrays):
+    elif any(isinstance(a, spmatrix | SpArray) for a in arrays):
         sparse_stack = (sparse.vstack, sparse.hstack)[axis]
         use_sparse_array = any(issubclass(type(a), SpArray) for a in arrays)
         return sparse_stack(
@@ -941,7 +940,7 @@ def gen_outer_reindexers(els, shapes, new_index: pd.Index, *, axis=0):
 
 def missing_element(
     n: int,
-    els: list[SpArray | sparse.csr_matrix | sparse.csc_matrix | np.ndarray | DaskArray],
+    els: list[CSArray | CSMatrix | np.ndarray | DaskArray],
     axis: Literal[0, 1] = 0,
     fill_value: Any | None = None,
     off_axis_size: int = 0,

src/anndata/_core/raw.py

@@ -17,8 +17,7 @@
     from collections.abc import Mapping, Sequence
     from typing import ClassVar
 
-    from scipy import sparse
-
+    from ..compat import CSMatrix
     from .aligned_mapping import AxisArraysView
     from .anndata import AnnData
     from .sparse_dataset import BaseCompressedSparseDataset
@@ -31,7 +30,7 @@ class Raw:
     def __init__(
         self,
         adata: AnnData,
-        X: np.ndarray | sparse.spmatrix | None = None,
+        X: np.ndarray | CSMatrix | None = None,
         var: pd.DataFrame | Mapping[str, Sequence] | None = None,
         varm: AxisArrays | Mapping[str, np.ndarray] | None = None,
     ):
@@ -67,7 +66,7 @@ def _get_X(self, layer=None):
         return self.X
 
     @property
-    def X(self) -> BaseCompressedSparseDataset | np.ndarray | sparse.spmatrix:
+    def X(self) -> BaseCompressedSparseDataset | np.ndarray | CSMatrix:
         # TODO: Handle unsorted array of integer indices for h5py.Datasets
         if not self._adata.isbacked:
             return self._X

src/anndata/_core/sparse_dataset.py

@@ -40,7 +40,7 @@
     from scipy.sparse._compressed import _cs_matrix
 
     from .._types import GroupStorageType
-    from ..compat import H5Array
+    from ..compat import CSArray, CSMatrix, H5Array
     from .index import Index, Index1D
 else:
     from scipy.sparse import spmatrix as _cs_matrix
@@ -67,7 +67,7 @@ class BackedSparseMatrix(_cs_matrix):
     indices: GroupStorageType
     indptr: np.ndarray
 
-    def copy(self) -> ss.csr_matrix | ss.csc_matrix:
+    def copy(self) -> CSMatrix:
         if isinstance(self.data, h5py.Dataset):
             return sparse_dataset(self.data.parent).to_memory()
         if isinstance(self.data, ZarrArray):
@@ -433,9 +433,7 @@ def __repr__(self) -> str:
         name = type(self).__name__.removeprefix("_")
         return f"{name}: backend {self.backend}, shape {self.shape}, data_dtype {self.dtype}"
 
-    def __getitem__(
-        self, index: Index | tuple[()]
-    ) -> float | ss.csr_matrix | ss.csc_matrix | SpArray:
+    def __getitem__(self, index: Index | tuple[()]) -> float | CSMatrix | CSArray:
         indices = self._normalize_index(index)
         row, col = indices
         mtx = self._to_backed()
@@ -494,7 +492,7 @@ def __setitem__(self, index: Index | tuple[()], value) -> None:
         mock_matrix[row, col] = value
 
     # TODO: split to other classes?
-    def append(self, sparse_matrix: ss.csr_matrix | ss.csc_matrix | SpArray) -> None:
+    def append(self, sparse_matrix: CSMatrix | CSArray) -> None:
         """Append an in-memory or on-disk sparse matrix to the current object's store.
 
         Parameters
@@ -620,7 +618,7 @@ def _to_backed(self) -> BackedSparseMatrix:
         mtx.indptr = self._indptr
         return mtx
 
-    def to_memory(self) -> ss.csr_matrix | ss.csc_matrix | SpArray:
+    def to_memory(self) -> CSMatrix | CSArray:
         format_class = get_memory_class(
             self.format, use_sparray_in_io=settings.use_sparse_array_on_read
         )

src/anndata/_io/h5ad.py

@@ -38,6 +38,7 @@
     from typing import Any, Literal
 
     from .._core.file_backing import AnnDataFileManager
+    from ..compat import CSMatrix
 
 T = TypeVar("T")
 
@@ -115,7 +116,7 @@ def write_h5ad(
 def write_sparse_as_dense(
     f: h5py.Group,
     key: str,
-    value: sparse.spmatrix | BaseCompressedSparseDataset,
+    value: CSMatrix | BaseCompressedSparseDataset,
     *,
     dataset_kwargs: Mapping[str, Any] = MappingProxyType({}),
 ):
@@ -172,7 +173,7 @@ def read_h5ad(
     backed: Literal["r", "r+"] | bool | None = None,
     *,
     as_sparse: Sequence[str] = (),
-    as_sparse_fmt: type[sparse.spmatrix] = sparse.csr_matrix,
+    as_sparse_fmt: type[CSMatrix] = sparse.csr_matrix,
     chunk_size: int = 6000,  # TODO, probably make this 2d chunks
 ) -> AnnData:
     """\
@@ -273,7 +274,7 @@ def callback(func, elem_name: str, elem, iospec):
 def _read_raw(
     f: h5py.File | AnnDataFileManager,
     as_sparse: Collection[str] = (),
-    rdasp: Callable[[h5py.Dataset], sparse.spmatrix] | None = None,
+    rdasp: Callable[[h5py.Dataset], CSMatrix] | None = None,
     *,
     attrs: Collection[str] = ("X", "var", "varm"),
 ) -> dict:
@@ -346,7 +347,7 @@ def read_dataset(dataset: h5py.Dataset):
 
 @report_read_key_on_error
 def read_dense_as_sparse(
-    dataset: h5py.Dataset, sparse_format: sparse.spmatrix, axis_chunk: int
+    dataset: h5py.Dataset, sparse_format: CSMatrix, axis_chunk: int
 ):
     if sparse_format == sparse.csr_matrix:
         return read_dense_as_csr(dataset, axis_chunk)

src/anndata/_io/specs/lazy_methods.py

@@ -20,7 +20,7 @@
     from collections.abc import Generator, Mapping, Sequence
     from typing import Literal, ParamSpec, TypeVar
 
-    from ...compat import DaskArray, H5File, SpArray
+    from ...compat import CSArray, CSMatrix, DaskArray, H5File
     from .registry import DaskReader
 
     BlockInfo = Mapping[
@@ -72,7 +72,7 @@ def make_dask_chunk(
     path_or_sparse_dataset: Path | D,
     elem_name: str,
     block_info: BlockInfo | None = None,
-) -> sparse.csr_matrix | sparse.csc_matrix | SpArray:
+) -> CSMatrix | CSArray:
     if block_info is None:
         msg = "Block info is required"
         raise ValueError(msg)

src/anndata/_io/specs/methods.py

@@ -51,10 +51,9 @@
     from numpy import typing as npt
     from numpy.typing import NDArray
 
-    from anndata._types import ArrayStorageType, GroupStorageType
-    from anndata.compat import SpArray
-    from anndata.typing import AxisStorable, InMemoryArrayOrScalarType
-
+    from ..._types import ArrayStorageType, GroupStorageType
+    from ...compat import CSArray, CSMatrix
+    from ...typing import AxisStorable, InMemoryArrayOrScalarType
     from .registry import Reader, Writer
 
 ####################
@@ -127,7 +126,7 @@ def wrapper(
 @_REGISTRY.register_read(H5Array, IOSpec("", ""))
 def read_basic(
     elem: H5File | H5Group | H5Array, *, _reader: Reader
-) -> dict[str, InMemoryArrayOrScalarType] | npt.NDArray | sparse.spmatrix | SpArray:
+) -> dict[str, InMemoryArrayOrScalarType] | npt.NDArray | CSMatrix | CSArray:
     from anndata._io import h5ad
 
     warn(
@@ -149,7 +148,7 @@ def read_basic(
 @_REGISTRY.register_read(ZarrArray, IOSpec("", ""))
 def read_basic_zarr(
     elem: ZarrGroup | ZarrArray, *, _reader: Reader
-) -> dict[str, InMemoryArrayOrScalarType] | npt.NDArray | sparse.spmatrix | SpArray:
+) -> dict[str, InMemoryArrayOrScalarType] | npt.NDArray | CSMatrix | CSArray:
     from anndata._io import zarr
 
     warn(
@@ -590,7 +589,7 @@ def write_recarray_zarr(
 def write_sparse_compressed(
     f: GroupStorageType,
     key: str,
-    value: sparse.spmatrix | SpArray,
+    value: CSMatrix | CSArray,
     *,
     _writer: Writer,
     fmt: Literal["csr", "csc"],
@@ -756,9 +755,7 @@ def chunk_slice(start: int, stop: int) -> tuple[slice | None, slice | None]:
 @_REGISTRY.register_read(H5Group, IOSpec("csr_matrix", "0.1.0"))
 @_REGISTRY.register_read(ZarrGroup, IOSpec("csc_matrix", "0.1.0"))
 @_REGISTRY.register_read(ZarrGroup, IOSpec("csr_matrix", "0.1.0"))
-def read_sparse(
-    elem: GroupStorageType, *, _reader: Reader
-) -> sparse.spmatrix | SpArray:
+def read_sparse(elem: GroupStorageType, *, _reader: Reader) -> CSMatrix | CSArray:
     return sparse_dataset(elem).to_memory()
 
 

src/anndata/abc.py

@@ -7,9 +7,8 @@
     from typing import ClassVar, Literal
 
     import numpy as np
-    from scipy.sparse import csc_matrix, csr_matrix
 
-    from .compat import Index, SpArray
+    from .compat import CSArray, CSMatrix, Index
 
 
 __all__ = ["CSRDataset", "CSCDataset"]
@@ -31,7 +30,7 @@ class _AbstractCSDataset(ABC):
     """Which file type is used on-disk."""
 
     @abstractmethod
-    def __getitem__(self, index: Index) -> float | csr_matrix | csc_matrix | SpArray:
+    def __getitem__(self, index: Index) -> float | CSMatrix | CSArray:
         """Load a slice or an element from the sparse dataset into memory.
 
         Parameters
@@ -45,7 +44,7 @@ def __getitem__(self, index: Index) -> float | csr_matrix | csc_matrix | SpArray
         """
 
     @abstractmethod
-    def to_memory(self) -> csr_matrix | csc_matrix | SpArray:
+    def to_memory(self) -> CSMatrix | CSArray:
         """Load the sparse dataset into memory.
 
         Returns