feat: make use_rep accept anndata.acc

docs/release-notes/4315.feat.md

@@ -0,0 +1 @@
+The `use_rep` parameter of {func}`scanpy.pp.neighbors`, {func}`scanpy.tl.tsne`, and {func}`scanpy.tl.dendrogram` now accepts {mod}`anndata.acc` accessors such as `A.X`, `A.layers["scaled"]`, or `A.obsm["pca"]`, and resolves strings into them if {attr}`scanpy.settings.preset` is {attr}`~scanpy.Preset.ScanpyV2Preview` {smaller}`P Angerer`

pyproject.toml

@@ -99,7 +99,7 @@ scrublet = [ "scikit-image>=0.25" ]
 # highly_variable_genes method 'seurat_v3'
 skmisc = [ "scikit-misc>=0.5.1" ]
 illico = [ "illico>=0.6" ]
-scanpy2 = [ "anndata>=0.13.2", "hv-anndata>=0.0.3a5", "igraph>=0.10.8", "scanpy[illico]", "scikit-misc>=0.5.1" ]
+scanpy2 = [ "anndata>=0.13.3", "hv-anndata>=0.0.3a5", "igraph>=0.10.8", "scanpy[illico]", "scikit-misc>=0.5.1" ]
 
 [dependency-groups]
 dev = [

src/scanpy/get/get.py

@@ -2,7 +2,8 @@
 
 from __future__ import annotations
 
-from collections.abc import Collection
+import json
+from collections.abc import Collection, Sequence
 from importlib.util import find_spec
 from typing import TYPE_CHECKING, TypedDict, overload
 
@@ -19,7 +20,7 @@
     from collections.abc import Iterable
     from typing import Any, Literal, Unpack
 
-    from anndata.acc import Idx2D
+    from anndata.acc import Idx2D, RefAcc
 
     from .._compat import DaskArray
 
@@ -485,6 +486,8 @@ class _Rep(TypedDict, total=False):
 
 
 type ArrAcc = GraphAcc | LayerAcc | MultiAcc
+type RepAcc = LayerAcc | MultiAcc
+"""Accessor usable as a representation (`use_rep`), i.e. a non-graph 2D array."""
 
 
 @overload
@@ -790,3 +793,70 @@ def _get_vec(
         ref = A.resolve(ref, vec=True)
     _ref_dim(ref, dim=dim)
     return adata[ref]
+
+
+def _resolve_rep(rep: RefAcc | str) -> RepAcc:
+    """Resolve a `rep`resentation string into a `LayerAcc`/`MultiAcc` using `anndata.acc`."""
+    if isinstance(rep, str):
+        from anndata.acc import A
+
+        rep = A.resolve(rep, vec=False)
+    if isinstance(rep, MultiAcc) and rep.dim != "obs":
+        msg = (
+            f"Representation must be aligned to `obs`, but {rep!r} is aligned to `var`"
+        )
+        raise ValueError(msg)
+    if isinstance(rep, LayerAcc | MultiAcc):
+        return rep
+    msg = (
+        "Representation must be a `LayerAcc` (e.g. `A.X`, `A.layers[...]`) or a "
+        f"`MultiAcc` (e.g. `A.obsm[...]`), was {rep!r}"
+    )
+    raise TypeError(msg)
+
+
+def _rep_to_json(rep: RepAcc | str | None) -> str | list[str] | None:
+    """Serialize a `rep`resentation for storage in `.uns`.
+
+    v1 strings (`'X'` or an `.obsm` key) are stored unchanged,
+    accessors (and hence v2 strings) as `anndata.acc` JSON inside a 1-element list,
+    e.g. `A.obsm['pca']` as `['["obsm", "pca"]']`.
+
+    TODO: Once AnnData can store a heterogeneous list, store that instead of a 1-element list.
+    See https://github.com/scverse/anndata/issues/1979
+    """
+    from scanpy import settings
+
+    if rep is None or (
+        isinstance(rep, str) and settings.preset is not Preset.ScanpyV2Preview
+    ):
+        return rep
+    from anndata.acc import A
+
+    return [json.dumps(A.to_json(_resolve_rep(rep)))]
+
+
+def _rep_from_json(rep: str | Sequence[str | int | None] | None) -> RepAcc | str | None:
+    """Parse a `rep`resentation stored by `_rep_to_json`."""
+    from scanpy import settings
+
+    if rep is None:
+        return rep
+    if not isinstance(rep, str):
+        from anndata.acc import A
+
+        if (
+            isinstance(rep, Sequence | np.ndarray)
+            and len(rep) == 1
+            and isinstance(rep[0], str)
+        ):
+            # see `_rep_to_json`
+            rep: Sequence[str | int | None] = json.loads(rep[0])
+        return _resolve_rep(A.from_json(rep, vec=False))
+    if settings.preset is Preset.ScanpyV2Preview:
+        from anndata.acc import A
+
+        # a plain string was stored under the v1 preset,
+        # so interpret it as one instead of as an `anndata.acc` spec
+        return A.X if rep == "X" else A.obsm[rep]
+    return rep

src/scanpy/neighbors/__init__.py

@@ -23,6 +23,7 @@
 from .._keys import _EmbeddingKeys, _existing_preset_keys
 from .._utils import NeighborsView, _doc_params, get_literal_vals
 from .._utils.random import _accepts_legacy_random_state, _LegacyRng
+from ..get.get import _rep_to_json
 from . import _connectivity
 from ._common import (
     _get_indices_distances_from_dense_matrix,
@@ -44,6 +45,7 @@
     from numpy.typing import NDArray
 
     from .._utils.random import RNGLike, SeedLike
+    from ..get.get import RepAcc
     from ._types import (
         KnnTransformerLike,
         RPForestDict,
@@ -64,7 +66,7 @@ def neighbors(  # noqa: PLR0913
     n_pcs: int | None = None,
     *,
     distances: np.ndarray | SpBase | None = None,
-    use_rep: str | None = None,
+    use_rep: RepAcc | str | None = None,
     knn: bool = True,
     method: _Method = "umap",
     transformer: KnnTransformerLike | _KnownTransformer | None = None,
@@ -249,7 +251,7 @@ def neighbors(  # noqa: PLR0913
         metric=metric,
         **meta_random_state,
         **({} if not metric_kwds else dict(metric_kwds=metric_kwds)),
-        **({} if use_rep is None else dict(use_rep=use_rep)),
+        **({} if use_rep is None else dict(use_rep=_rep_to_json(use_rep))),
         **({} if n_pcs is None else dict(n_pcs=n_pcs)),
     )
 
@@ -536,7 +538,7 @@ def compute_neighbors(
         n_neighbors: int = 30,
         n_pcs: int | None = None,
         *,
-        use_rep: str | None = None,
+        use_rep: RepAcc | str | None = None,
         knn: bool = True,
         method: _Method | None = "umap",
         transformer: KnnTransformerLike | _KnownTransformer | None = None,
@@ -564,7 +566,7 @@ def compute_neighbors(
         if `method` is not `None`, `.connectivities`.
 
         """
-        from ..tools._utils import _choose_representation
+        from ..tools._utils import _choose_representation_compat
 
         start_neighbors = logg.debug("computing neighbors")
         if transformer is not None and not isinstance(transformer, str):
@@ -590,7 +592,7 @@ def compute_neighbors(
         self._rp_forest = None
         self.n_neighbors = n_neighbors
         self.knn = knn
-        x = _choose_representation(self._adata, use_rep=use_rep, n_pcs=n_pcs)
+        x = _choose_representation_compat(self._adata, use_rep=use_rep, n_pcs=n_pcs)
         self._distances = transformer.fit_transform(x)
         knn_indices, knn_distances = _get_indices_distances_from_sparse_matrix(
             self._distances, n_neighbors

src/scanpy/neighbors/_doc.py

@@ -9,13 +9,19 @@
     ``.obsp[.uns[neighbors_key]['connectivities_key']]`` for connectivities.
 """
 
-doc_use_rep = """\
-use_rep
-    Use the indicated representation. `'X'` or any key for `.obsm` is valid.
+doc_use_rep = r"""use_rep
+    Use the indicated representation:
+    a :class:`~anndata.acc.LayerAcc` (e.g. `A.X`, `A.layers[...]`) or
+    :class:`~anndata.acc.MultiAcc` (e.g. `A.obsm[...]`, `A.varm[...]`).
+    A :class:`str` is :meth:`~anndata.acc.AdAcc.resolve`\ d to one of those
+    if :attr:`scanpy.settings.preset` is :attr:`~scanpy.Preset.ScanpyV2Preview`,
+    otherwise interpreted as `'X'` or a key of `.obsm`.
+
     If `None`, the representation is chosen automatically:
-    For `.n_vars` < :attr:`~scanpy.settings.N_PCS` (default: 50), `.X` is used, otherwise 'X_pca' is used.
-    If 'X_pca' is not present, it’s computed with default parameters or `n_pcs` if present.\
-"""
+    For `.n_vars` < :attr:`~scanpy.settings.N_PCS` (default: 50), `.X` is used, otherwise the PCA
+    representation (`.obsm['X_pca']`, or `.obsm['pca']` if it was computed under
+    :attr:`~scanpy.Preset.ScanpyV2Preview`).
+    If it is not present, it’s computed with default parameters or `n_pcs` if present."""
 
 doc_n_pcs = """\
 n_pcs

src/scanpy/neighbors/_types.py

@@ -98,5 +98,5 @@ class NeighborsParams(TypedDict):
     metric: _Metric | _MetricFn | None
     random_state: NotRequired[_LegacyRandom]
     metric_kwds: NotRequired[Mapping[str, Any]]
-    use_rep: NotRequired[str]
+    use_rep: NotRequired[str | list[str]]  # see `scanpy.get.get._rep_to_json`
     n_pcs: NotRequired[int]

src/scanpy/tools/_dendrogram.py

@@ -9,23 +9,26 @@
 
 from .. import logging as logg
 from .._utils import _doc_params, raise_not_implemented_error_if_backed_type
+from ..get.get import _rep_to_json
 from ..neighbors._doc import doc_n_pcs, doc_use_rep
-from ._utils import _choose_representation
+from ._utils import _choose_representation_compat
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
     from typing import Any
 
     from anndata import AnnData
 
+    from ..get.get import RepAcc
+
 
 @_doc_params(n_pcs=doc_n_pcs, use_rep=doc_use_rep)
 def dendrogram(  # noqa: PLR0913
     adata: AnnData,
     groupby: str | Sequence[str],
     *,
     n_pcs: int | None = None,
-    use_rep: str | None = None,
+    use_rep: RepAcc | str | None = None,
     var_names: Sequence[str] | None = None,
     use_raw: bool | None = None,
     cor_method: str = "pearson",
@@ -125,7 +128,7 @@ def dendrogram(  # noqa: PLR0913
 
     if var_names is None:
         rep_df = pd.DataFrame(
-            _choose_representation(adata, use_rep=use_rep, n_pcs=n_pcs)
+            _choose_representation_compat(adata, use_rep=use_rep, n_pcs=n_pcs)
         )
         categorical = adata.obs[groupby[0]]
         if len(groupby) > 1:
@@ -167,7 +170,7 @@ def dendrogram(  # noqa: PLR0913
     dat = dict(
         linkage=z_var,
         groupby=groupby,
-        use_rep=use_rep,
+        use_rep=_rep_to_json(use_rep),
         cor_method=cor_method,
         linkage_method=linkage_method,
         categories_ordered=dendro_info["ivl"],

src/scanpy/tools/_ingest.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 from collections.abc import MutableMapping
 from typing import TYPE_CHECKING
 
@@ -21,7 +22,9 @@
 from .._utils._doctests import doctest_skipif
 from .._utils.random import _legacy_random_state, _LegacyRng
 from ..get import _check_mask
+from ..get.get import MultiAcc, _rep_from_json
 from ..neighbors import FlatTree
+from ._utils import _choose_representation_compat
 
 if TYPE_CHECKING:
     from collections.abc import Generator, Iterable
@@ -31,6 +34,7 @@
     from umap import UMAP
 
     from .._keys import _EmbeddingKeys
+    from ..get.get import RepAcc
     from ..neighbors import RPForestDict
 
 
@@ -225,7 +229,7 @@ class Ingest:
     _rng: np.random.Generator | None
     # neighbors
     _rep: np.ndarray
-    _use_rep: str
+    _use_rep: RepAcc | str
     _metric: str
     _metric_kwds: dict[str, object]
     _n_neighbors: int
@@ -323,8 +327,10 @@ def _init_neighbors(self, adata: AnnData, neighbors_key: str | None) -> None:
         self._n_neighbors = neighbors["params"]["n_neighbors"]
 
         if "use_rep" in neighbors["params"]:
-            self._use_rep = neighbors["params"]["use_rep"]
-            self._rep = adata.X if self._use_rep == "X" else adata.obsm[self._use_rep]
+            self._use_rep = _rep_from_json(neighbors["params"]["use_rep"])
+            self._rep = _choose_representation_compat(
+                adata, use_rep=self._use_rep, n_pcs=None
+            )
         elif "n_pcs" in neighbors["params"]:
             self._use_rep = "X_pca"
             self._n_pcs = neighbors["params"]["n_pcs"]
@@ -422,10 +428,11 @@ def _same_rep(self):
         adata = self._adata_new
         if self._n_pcs is not None:
             return self._pca(self._n_pcs)
-        if self._use_rep == "X":
-            return adata.X
-        if self._use_rep in adata.obsm:
-            return adata.obsm[self._use_rep]
+        # fall back to `.X` if the representation is missing in the new object
+        with contextlib.suppress(KeyError, ValueError):
+            return _choose_representation_compat(
+                adata, use_rep=self._use_rep, n_pcs=None
+            )
         return adata.X
 
     def fit(self, adata_new: AnnData) -> None:
@@ -558,11 +565,14 @@ def to_adata_joint(
                     self._obsm[key],
                 ))
 
-        if self._use_rep not in ("X_pca", "X"):
-            adata.obsm[self._use_rep] = np.vstack((
-                self._adata_ref.obsm[self._use_rep],
-                self._obsm["rep"],
-            ))
+        pca_keys = _existing_preset_keys(self._adata_ref, "pca")
+        skip = {"X", pca_keys.obsm if pca_keys else "X_pca"}
+        match self._use_rep:
+            case MultiAcc(dim="obs", k=key) | str(key) if key not in skip:
+                adata.obsm[key] = np.vstack((
+                    self._adata_ref.obsm[key],
+                    self._obsm["rep"],
+                ))
 
         if keys := _existing_preset_keys(self._adata_ref, "umap"):
             adata.uns[keys.uns] = self._adata_ref.uns[keys.uns]

src/scanpy/tools/_tsne.py

@@ -9,13 +9,15 @@
 from .._settings import Default, settings
 from .._utils import _doc_params, raise_not_implemented_error_if_backed_type
 from .._utils.random import _accepts_legacy_random_state, _legacy_random_state
+from ..get.get import _rep_to_json
 from ..neighbors._doc import doc_n_pcs, doc_use_rep
-from ._utils import _choose_representation
+from ._utils import _choose_representation_compat
 
 if TYPE_CHECKING:
     from anndata import AnnData
 
     from .._utils.random import RNGLike, SeedLike
+    from ..get.get import RepAcc
 
 
 @_accepts_legacy_random_state(0)
@@ -25,7 +27,7 @@ def tsne(  # noqa: PLR0913
     n_pcs: int | None = None,
     *,
     n_components: int = 2,
-    use_rep: str | None = None,
+    use_rep: RepAcc | str | None = None,
     perplexity: float = 30,
     metric: str = "euclidean",
     early_exaggeration: float = 12,
@@ -105,7 +107,7 @@ def tsne(  # noqa: PLR0913
     start = logg.info("computing tSNE")
     keys = _embedding_keys("tsne", key_added)
     adata = adata.copy() if copy else adata
-    x = _choose_representation(adata, use_rep=use_rep, n_pcs=n_pcs)
+    x = _choose_representation_compat(adata, use_rep=use_rep, n_pcs=n_pcs)
     raise_not_implemented_error_if_backed_type(x, "tsne")
     # params for sklearn
     n_jobs = settings.n_jobs if n_jobs is None else n_jobs
@@ -156,7 +158,7 @@ def tsne(  # noqa: PLR0913
         learning_rate=learning_rate,
         n_jobs=n_jobs,
         metric=metric,
-        use_rep=use_rep,
+        use_rep=_rep_to_json(use_rep),
         n_components=n_components,
     )
     adata.obsm[keys.obsm] = x_tsne  # annotate samples with tSNE coordinates

src/scanpy/tools/_umap.py

@@ -15,7 +15,8 @@
     _legacy_random_state,
     _LegacyRng,
 )
-from ._utils import _choose_representation, get_init_pos_from_paga
+from ..get.get import _rep_from_json
+from ._utils import _choose_representation_compat, get_init_pos_from_paga
 
 if TYPE_CHECKING:
     from typing import Literal
@@ -182,9 +183,9 @@ def umap(  # noqa: PLR0913
         init_coords = check_array(init_coords, dtype=np.float32, accept_sparse=False)
 
     neigh_params = neighbors["params"]
-    x = _choose_representation(
+    x = _choose_representation_compat(
         adata,
-        use_rep=neigh_params.get("use_rep", None),
+        use_rep=_rep_from_json(neigh_params.get("use_rep", None)),
         n_pcs=neigh_params.get("n_pcs", None),
         silent=True,
     )