(fix): correct typing of AnnData.X (#1616)

Co-authored-by: Philipp A. <flying-sheep@web.de>

Author: ilan-gold

docs/concatenation.rst

@@ -54,8 +54,8 @@ When the variables present in the objects to be concatenated aren't exactly the
 This is otherwise called taking the `"inner"` (intersection) or `"outer"` (union) join.
 For example, given two anndata objects with differing variables:
 
-    >>> a = AnnData(sparse.eye(3), var=pd.DataFrame(index=list("abc")))
-    >>> b = AnnData(sparse.eye(2), var=pd.DataFrame(index=list("ba")))
+    >>> a = AnnData(sparse.eye(3, format="csr"), var=pd.DataFrame(index=list("abc")))
+    >>> b = AnnData(sparse.eye(2, format="csr"), var=pd.DataFrame(index=list("ba")))
     >>> ad.concat([a, b], join="inner").X.toarray()
     array([[1., 0.],
            [0., 1.],
@@ -208,11 +208,11 @@ Note that comparisons are made after indices are aligned.
 That is, if the objects only share a subset of indices on the alternative axis, it's only required that values for those indices match when using a strategy like `"same"`.
 
     >>> a = AnnData(
-    ...     sparse.eye(3),
+    ...     sparse.eye(3, format="csr"),
     ...     var=pd.DataFrame({"nums": [1, 2, 3]}, index=list("abc"))
     ... )
     >>> b = AnnData(
-    ...     sparse.eye(2),
+    ...     sparse.eye(2, format="csr"),
     ...     var=pd.DataFrame({"nums": [2, 1]}, index=list("ba"))
     ... )
     >>> ad.concat([a, b], merge="same").var

docs/release-notes/1616.doc.md

@@ -0,0 +1 @@
+Correct {attr}`anndata.AnnData.X` type to include {class}`~anndata.experimental.CSRDataset` and {class}`~anndata.experimental.CSCDataset` as possible types {user}`ilan-gold`

src/anndata/_core/anndata.py

@@ -52,9 +52,9 @@
     from os import PathLike
     from typing import Any, Literal
 
+    from .._types import ArrayDataStructureType
     from .aligned_mapping import AxisArraysView, LayersView, PairwiseArraysView
     from .index import Index, Index1D
-    from .views import ArrayView
 
 
 # for backwards compat
@@ -541,7 +541,7 @@ def shape(self) -> tuple[int, int]:
         return self.n_obs, self.n_vars
 
     @property
-    def X(self) -> np.ndarray | sparse.spmatrix | SpArray | ArrayView | None:
+    def X(self) -> ArrayDataStructureType | None:
         """Data matrix of shape :attr:`n_obs` × :attr:`n_vars`."""
         if self.isbacked:
             if not self.file.is_open:

src/anndata/_core/merge.py

@@ -737,8 +737,8 @@ def gen_reindexer(new_var: pd.Index, cur_var: pd.Index):
     Usage
     -----
 
-    >>> a = AnnData(sparse.eye(3), var=pd.DataFrame(index=list("abc")))
-    >>> b = AnnData(sparse.eye(2), var=pd.DataFrame(index=list("ba")))
+    >>> a = AnnData(sparse.eye(3, format="csr"), var=pd.DataFrame(index=list("abc")))
+    >>> b = AnnData(sparse.eye(2, format="csr"), var=pd.DataFrame(index=list("ba")))
     >>> reindexer = gen_reindexer(a.var_names, b.var_names)
     >>> sparse.vstack([a.X, reindexer(b.X)]).toarray()
     array([[1., 0., 0.],

src/anndata/_core/storage.py

@@ -1,8 +1,7 @@
 from __future__ import annotations
 
 import warnings
-from enum import Enum
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Union, get_args
 
 import numpy as np
 import pandas as pd
@@ -25,48 +24,27 @@
     join_english,
     raise_value_error_if_multiindex_columns,
 )
-from .sparse_dataset import BaseCompressedSparseDataset
+from .sparse_dataset import CSCDataset, CSRDataset
 
 if TYPE_CHECKING:
-    from collections.abc import Generator
-    from typing import Any
+    from typing import Any, TypeAlias
 
-
-class ArrayDataStructureType(Enum):
-    # Memory
-    Array = (np.ndarray, "np.ndarray")
-    Masked = (ma.MaskedArray, "numpy.ma.core.MaskedArray")
-    Sparse = (sparse.spmatrix, "scipy.sparse.spmatrix")
-    SparseArray = (SpArray, "scipy.sparse.sparray")
-    AwkArray = (AwkArray, "awkward.Array")
-    # Backed
-    HDF5Dataset = (H5Array, "h5py.Dataset")
-    ZarrArray = (ZarrArray, "zarr.Array")
-    ZappyArray = (ZappyArray, "zappy.base.ZappyArray")
-    BackedSparseMatrix = (
-        BaseCompressedSparseDataset,
-        "anndata.experimental.[CSC,CSR]Dataset",
-    )
-    # Distributed
-    DaskArray = (DaskArray, "dask.array.Array")
-    CupyArray = (CupyArray, "cupy.ndarray")
-    CupySparseMatrix = (CupySparseMatrix, "cupyx.scipy.sparse.spmatrix")
-
-    @property
-    def cls(self):
-        return self.value[0]
-
-    @property
-    def qualname(self):
-        return self.value[1]
-
-    @classmethod
-    def classes(cls) -> tuple[type, ...]:
-        return tuple(v.cls for v in cls)
-
-    @classmethod
-    def qualnames(cls) -> Generator[str, None, None]:
-        yield from (v.qualname for v in cls)
+ArrayDataStructureType: TypeAlias = Union[
+    np.ndarray,
+    ma.MaskedArray,
+    sparse.csr_matrix,
+    sparse.csc_matrix,
+    SpArray,
+    AwkArray,
+    H5Array,
+    ZarrArray,
+    ZappyArray,
+    CSRDataset,
+    CSCDataset,
+    DaskArray,
+    CupyArray,
+    CupySparseMatrix,
+]
 
 
 def coerce_array(
@@ -81,12 +59,21 @@ def coerce_array(
     if allow_array_like and np.isscalar(value):
         return value
     # If value is one of the allowed types, return it
-    if isinstance(value, ArrayDataStructureType.classes()):
+    array_data_structure_types = get_args(ArrayDataStructureType)
+    if isinstance(value, array_data_structure_types):
         if isinstance(value, np.matrix):
             msg = f"{name} should not be a np.matrix, use np.ndarray instead."
             warnings.warn(msg, ImplicitModificationWarning)
             value = value.A
         return value
+    elif isinstance(value, sparse.spmatrix):
+        msg = (
+            f"AnnData previously had undefined behavior around matrices of type {type(value)}."
+            "In 0.12, passing in this type will throw an error. Please convert to a supported type."
+            "Continue using for this minor version at your own risk."
+        )
+        warnings.warn(msg, FutureWarning)
+        return value
     if isinstance(value, pd.DataFrame):
         if allow_df:
             raise_value_error_if_multiindex_columns(value, name)
@@ -100,7 +87,7 @@ def coerce_array(
         except (ValueError, TypeError) as _e:
             e = _e
     # if value isn’t the right type or convertible, raise an error
-    msg = f"{name} needs to be of one of {join_english(ArrayDataStructureType.qualnames())}, not {type(value)}."
+    msg = f"{name} needs to be of one of {join_english(map(str, array_data_structure_types))}, not {type(value)}."
     if e is not None:
         msg += " (Failed to convert it to an array, see above for details.)"
     raise ValueError(msg) from e

src/anndata/_types.py

@@ -8,21 +8,13 @@
 
 import numpy as np
 import pandas as pd
-from numpy.typing import NDArray
-from scipy import sparse
 
 from anndata._core.anndata import AnnData
 
-from ._core.sparse_dataset import BaseCompressedSparseDataset
+from ._core.storage import ArrayDataStructureType
 from .compat import (
-    AwkArray,
-    CupyArray,
-    CupySparseMatrix,
-    DaskArray,
     H5Array,
     H5Group,
-    SpArray,
-    ZappyArray,
     ZarrArray,
     ZarrGroup,
 )
@@ -34,6 +26,7 @@
     from anndata._io.specs.registry import DaskReader
 
     from ._io.specs.registry import IOSpec, Reader, Writer
+    from .compat import DaskArray
 
 __all__ = [
     "ArrayStorageType",
@@ -42,21 +35,7 @@
 ]
 
 InMemoryArrayOrScalarType: TypeAlias = Union[
-    NDArray,
-    np.ma.MaskedArray,
-    sparse.spmatrix,
-    SpArray,
-    H5Array,
-    ZarrArray,
-    ZappyArray,
-    BaseCompressedSparseDataset,
-    DaskArray,
-    CupyArray,
-    CupySparseMatrix,
-    AwkArray,
-    pd.DataFrame,
-    np.number,
-    str,
+    pd.DataFrame, np.number, str, ArrayDataStructureType
 ]
 RWAble: TypeAlias = Union[
     InMemoryArrayOrScalarType, dict[str, "RWAble"], list["RWAble"]

tests/test_base.py

@@ -30,7 +30,7 @@ def test_creation():
     AnnData(np.array([[1, 2], [3, 4]]))
     AnnData(np.array([[1, 2], [3, 4]]), {}, {})
     AnnData(ma.array([[1, 2], [3, 4]]), uns=dict(mask=[0, 1, 1, 0]))
-    AnnData(sp.eye(2))
+    AnnData(sp.eye(2, format="csr"))
     if CAN_USE_SPARSE_ARRAY:
         AnnData(sp.eye_array(2))
     X = np.array([[1, 2, 3], [4, 5, 6]])
@@ -95,7 +95,7 @@ def test_creation_error(src, src_arg, dim_msg, dim, dim_arg, msg: str | None):
 def test_invalid_X():
     with pytest.raises(
         ValueError,
-        match=r"X needs to be of one of np\.ndarray.*not <class 'str'>\.",
+        match=r"X needs to be of one of <class 'numpy.ndarray'>.*not <class 'str'>\.",
     ):
         AnnData("string is not a valid X")
 
@@ -126,7 +126,7 @@ def test_error_create_from_multiindex_df(attr):
 
 
 def test_create_from_sparse_df():
-    s = sp.random(20, 30, density=0.2)
+    s = sp.random(20, 30, density=0.2, format="csr")
     obs_names = [f"obs{i}" for i in range(20)]
     var_names = [f"var{i}" for i in range(30)]
     df = pd.DataFrame.sparse.from_spmatrix(s, index=obs_names, columns=var_names)

tests/test_obsmvarm.py

@@ -85,21 +85,21 @@ def test_setting_dataframe(adata: AnnData):
 
 
 def test_setting_sparse(adata: AnnData):
-    obsm_sparse = sparse.random(M, 100)
+    obsm_sparse = sparse.random(M, 100, format="csr")
     adata.obsm["a"] = obsm_sparse
     assert not np.any((adata.obsm["a"] != obsm_sparse).data)
 
-    varm_sparse = sparse.random(N, 100)
+    varm_sparse = sparse.random(N, 100, format="csr")
     adata.varm["a"] = varm_sparse
     assert not np.any((adata.varm["a"] != varm_sparse).data)
 
     h = joblib.hash(adata)
 
-    bad_obsm_sparse = sparse.random(M * 2, M)
+    bad_obsm_sparse = sparse.random(M * 2, M, format="csr")
     with pytest.raises(ValueError, match=r"incorrect shape"):
         adata.obsm["b"] = bad_obsm_sparse
 
-    bad_varm_sparse = sparse.random(N * 2, N)
+    bad_varm_sparse = sparse.random(N * 2, N, format="csr")
     with pytest.raises(ValueError, match=r"incorrect shape"):
         adata.varm["b"] = bad_varm_sparse
 

tests/test_obspvarp.py

@@ -65,21 +65,21 @@ def test_setting_ndarray(adata: AnnData):
 
 
 def test_setting_sparse(adata: AnnData):
-    obsp_sparse = sparse.random(M, M)
+    obsp_sparse = sparse.random(M, M, format="csr")
     adata.obsp["a"] = obsp_sparse
     assert not np.any((adata.obsp["a"] != obsp_sparse).data)
 
-    varp_sparse = sparse.random(N, N)
+    varp_sparse = sparse.random(N, N, format="csr")
     adata.varp["a"] = varp_sparse
     assert not np.any((adata.varp["a"] != varp_sparse).data)
 
     h = joblib.hash(adata)
 
-    bad_obsp_sparse = sparse.random(M * 2, M)
+    bad_obsp_sparse = sparse.random(M * 2, M, format="csr")
     with pytest.raises(ValueError, match=r"incorrect shape"):
         adata.obsp["b"] = bad_obsp_sparse
 
-    bad_varp_sparse = sparse.random(N * 2, N)
+    bad_varp_sparse = sparse.random(N * 2, N, format="csr")
     with pytest.raises(ValueError, match=r"incorrect shape"):
         adata.varp["b"] = bad_varp_sparse
 

tests/test_x.py

@@ -182,3 +182,12 @@ def test_set_dense_x_view_from_sparse():
     assert_equal(view.X, x1[:30])
     assert_equal(orig.X[:30], x1[:30])  # change propagates through
     assert_equal(orig.X[30:], x[30:])  # change propagates through
+
+
+def test_warn_on_non_csr_csc_matrix():
+    X = sparse.eye(100)
+    with pytest.warns(
+        FutureWarning,
+        match=rf"AnnData previously had undefined behavior around matrices of type {type(X)}.*",
+    ):
+        ad.AnnData(X=X)