Fix and test zappy support (scverse#3089)

docs/release-notes/1.10.2.md

@@ -19,6 +19,7 @@
 * Add clear errors where `backed` mode-like matrices (i.e., from `sparse_dataset`) are not supported {pr}`3048` {smaller}`I gold`
 * Write out full pca results when `_choose_representation` is called i.e., `~sc.pp.neighbors` without `sc.pp.pca` {pr}`3078` {smaller}`I gold`
 * Fix deprecated use of `.A` with sparse matrices {pr}`3084` {smaller}`P Angerer`
+* Fix zappy support {pr}`3089` {smaller}`P Angerer`
 
 ```{rubric} Performance
 ```

pyproject.toml

@@ -101,6 +101,8 @@ test = [
 ]
 test-full = [
     "scanpy[test]",
+    # optional storage modes
+    "zappy",
     # additional tested algorithms
     "scanpy[louvain]",
     "scanpy[magic]",

scanpy/_utils/__init__.py

@@ -580,14 +580,34 @@ def check_op(op):
 
 @singledispatch
 def axis_mul_or_truediv(
-    X: sparse.spmatrix,
+    X: np.ndarray,
+    scaling_array: np.ndarray,
+    axis: Literal[0, 1],
+    op: Callable[[Any, Any], Any],
+    *,
+    allow_divide_by_zero: bool = True,
+    out: np.ndarray | None = None,
+) -> np.ndarray:
+    check_op(op)
+    scaling_array = broadcast_axis(scaling_array, axis)
+    if op is mul:
+        return np.multiply(X, scaling_array, out=out)
+    if not allow_divide_by_zero:
+        scaling_array = scaling_array.copy() + (scaling_array == 0)
+    return np.true_divide(X, scaling_array, out=out)
+
+
+@axis_mul_or_truediv.register(sparse.csr_matrix)
+@axis_mul_or_truediv.register(sparse.csc_matrix)
+def _(
+    X: sparse.csr_matrix | sparse.csc_matrix,
     scaling_array,
     axis: Literal[0, 1],
     op: Callable[[Any, Any], Any],
     *,
     allow_divide_by_zero: bool = True,
-    out: sparse.spmatrix | None = None,
-) -> sparse.spmatrix:
+    out: sparse.csr_matrix | sparse.csc_matrix | None = None,
+) -> sparse.csr_matrix | sparse.csc_matrix:
     check_op(op)
     if out is not None:
         if X.data is not out.data:
@@ -629,25 +649,6 @@ def new_data_op(x):
     ).T
 
 
-@axis_mul_or_truediv.register(np.ndarray)
-def _(
-    X: np.ndarray,
-    scaling_array: np.ndarray,
-    axis: Literal[0, 1],
-    op: Callable[[Any, Any], Any],
-    *,
-    allow_divide_by_zero: bool = True,
-    out: np.ndarray | None = None,
-) -> np.ndarray:
-    check_op(op)
-    scaling_array = broadcast_axis(scaling_array, axis)
-    if op is mul:
-        return np.multiply(X, scaling_array, out=out)
-    if not allow_divide_by_zero:
-        scaling_array = scaling_array.copy() + (scaling_array == 0)
-    return np.true_divide(X, scaling_array, out=out)
-
-
 def make_axis_chunks(
     X: DaskArray, axis: Literal[0, 1], pad=True
 ) -> tuple[tuple[int], tuple[int]]: