Repository: scverse/scanpy

Issue:
Broken zappy tests
### Please make sure these conditions are met

- [X] I have checked that this issue has not already been reported.
- [X] I have confirmed this bug exists on the latest version of scanpy.
- [ ] (optional) I have confirmed this bug exists on the main branch of scanpy.

### What happened?

Somewhere, we broke zappy tests.

This diff is the start of fixing this: https://github.com/scverse/scanpy/compare/d3a8c7981321e4f0afa3290dd84c06fed3654b0a..fd90edb71b6f5fc4cc58638583ea4ff37b34e660

### Minimal code sample

```python
pip install zappy
pytest scanpy/tests/test_preprocessing_distributed.py
```


### Error output

```pytb
____________________________________________________________________________________________________________________________________ test_normalize_total[direct] ____________________________________________________________________________________________________________________________________

adata = AnnData object with n_obs × n_vars = 10000 × 1000
    var: 'gene_ids', adata_dist = AnnData object with n_obs × n_vars = 10000 × 1000
    var: 'gene_ids'
    uns: 'dist-mode'

    def test_normalize_total(adata: AnnData, adata_dist: AnnData):
>       normalize_total(adata_dist)

scanpy/tests/test_preprocessing_distributed.py:93: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/home/phil/.local/share/hatch/env/virtual/scanpy/q4In3tK-/hatch-test.stable/lib/python3.12/site-packages/legacy_api_wrap/__init__.py:80: in fn_compatible
    return fn(*args_all, **kw)
scanpy/preprocessing/_normalization.py:240: in normalize_total
    adata, _normalize_data(X, counts_per_cell, target_sum), layer=layer
scanpy/preprocessing/_normalization.py:49: in _normalize_data
    return axis_mul_or_truediv(
/usr/lib/python3.12/functools.py:909: in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

X = <zappy.direct.array.DirectZappyArray object at 0x7f06eb607a40>, scaling_array = array([0.9270073, 0.4379562, 1.5985402, ..., 0.8540146, 1.5109489,
       0.8978102], dtype=float32), axis = 0, op = <built-in function truediv>

    @singledispatch
    def axis_mul_or_truediv(
        X: sparse.spmatrix,
        scaling_array,
        axis: Literal[0, 1],
        op: Callable[[Any, Any], Any],
        *,
        allow_divide_by_zero: bool = True,
        out: sparse.spmatrix | None = None,
    ) -> sparse.spmatrix:
        check_op(op)
        if out is not None:
            if X.data is not out.data:
                raise ValueError(
                    "`out` argument provided but not equal to X.  This behavior is not supported for sparse matrix scaling."
                )
        if not allow_divide_by_zero and op is truediv:
            scaling_array = scaling_array.copy() + (scaling_array == 0)
    
        row_scale = axis == 0
        column_scale = axis == 1
        if row_scale:
    
            def new_data_op(x):
                return op(x.data, np.repeat(scaling_array, np.diff(x.indptr)))
    
        elif column_scale:
    
            def new_data_op(x):
                return op(x.data, scaling_array.take(x.indices, mode="clip"))
    
>       if X.format == "csr":
E       AttributeError: 'DirectZappyArray' object has no attribute 'format'

scanpy/_utils/__init__.py:612: AttributeError
```


### Versions

<details>

```
-----
anndata     0.10.7
scanpy      1.10.0rc2.dev61+g874d99b3
-----
PIL                 10.3.0
asciitree           NA
cloudpickle         3.0.0
coverage            7.5.3
cycler              0.12.1
cython_runtime      NA
dask                2024.5.2
dateutil            2.9.0.post0
distutils           3.12.3
h5py                3.11.0
igraph              0.11.5
joblib              1.4.2
kiwisolver          1.4.5
legacy_api_wrap     NA
leidenalg           0.10.2
llvmlite            0.42.0
matplotlib          3.9.0
mpl_toolkits        NA
natsort             8.4.0
numba               0.59.1
numcodecs           0.12.1
numpy               1.26.4
packaging           24.0
pandas              2.2.2
psutil              5.9.8
pyparsing           3.1.2
pytz                2024.1
scipy               1.13.1
session_info        1.0.0
setuptools          70.0.0
setuptools_scm      NA
six                 1.16.0
sklearn             1.5.0
texttable           1.7.0
threadpoolctl       3.5.0
tlz                 0.12.1
toolz               0.12.1
yaml                6.0.1
zappy               NA
zarr                2.18.2
-----
Python 3.12.3 (main, Apr 23 2024, 09:16:07) [GCC 13.2.1 20240417]
Linux-6.9.3-zen1-1-zen-x86_64-with-glibc2.39
-----
Session information updated at 2024-06-03 11:27
```

</details>


Discussion:

Patch:
diff --git a/docs/release-notes/1.10.2.md b/docs/release-notes/1.10.2.md
index 9eaa3a6e..6c7ce0b6 100644
--- a/docs/release-notes/1.10.2.md
+++ b/docs/release-notes/1.10.2.md
@@ -17,7 +17,9 @@
 
 * Compatibility with `matplotlib` 3.9 {pr}`2999` {smaller}`I Virshup`
 * 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
 ```
diff --git a/pyproject.toml b/pyproject.toml
index eaa5cdd5..dd07b621 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -101,6 +101,8 @@ test = [
 ]
 test-full = [
     "scanpy[test]",
+    # optional storage modes
+    "zappy",
     # additional tested algorithms
     "scanpy[louvain]",
     "scanpy[magic]",
diff --git a/scanpy/_utils/__init__.py b/scanpy/_utils/__init__.py
index f28acf2e..f4c21e32 100644
--- a/scanpy/_utils/__init__.py
+++ b/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 axis_mul_or_truediv(
     ).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]]:
diff --git a/scanpy/tools/_utils.py b/scanpy/tools/_utils.py
index 317948ca..b3ffa7d3 100644
--- a/scanpy/tools/_utils.py
+++ b/scanpy/tools/_utils.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import warnings
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -37,14 +38,14 @@ def _choose_representation(
                 X = adata.obsm["X_pca"][:, :n_pcs]
                 logg.info(f"    using 'X_pca' with n_pcs = {X.shape[1]}")
             else:
-                logg.warning(
+                warnings.warn(
                     f"You’re trying to run this on {adata.n_vars} dimensions of `.X`, "
                     "if you really want this, set `use_rep='X'`.\n         "
                     "Falling back to preprocessing with `sc.pp.pca` and default params."
                 )
                 n_pcs_pca = n_pcs if n_pcs is not None else settings.N_PCS
-                X = pca(adata.X, n_comps=n_pcs_pca)
-                adata.obsm["X_pca"] = X
+                pca(adata, n_comps=n_pcs_pca)
+                X = adata.obsm["X_pca"]
         else:
             logg.info("    using data matrix X directly")
             X = adata.X


Test patch:
diff --git a/scanpy/tests/test_neighbors_key_added.py b/scanpy/tests/test_neighbors_key_added.py
index d29470a1..4f30cee0 100644
--- a/scanpy/tests/test_neighbors_key_added.py
+++ b/scanpy/tests/test_neighbors_key_added.py
@@ -32,6 +32,16 @@ def test_neighbors_key_added(adata):
     )
 
 
+def test_neighbors_pca_keys_added_without_previous_pca_run(adata):
+    assert "pca" not in adata.uns and "X_pca" not in adata.obsm
+    with pytest.warns(
+        UserWarning,
+        match=r".*Falling back to preprocessing with `sc.pp.pca` and default params",
+    ):
+        sc.pp.neighbors(adata, n_neighbors=n_neighbors, random_state=0)
+    assert "pca" in adata.uns
+
+
 # test functions with neighbors_key and obsp
 @needs.igraph
 @needs.leidenalg