docs: document array type support (#3895)

Co-authored-by: Ilan Gold <ilanbassgold@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

docs/_static/custom.css

@@ -0,0 +1,5 @@
+code.literal > a > code.literal {
+    border: none;
+    padding: 0;
+    font-size: inherit;
+}

docs/api/index.md

@@ -10,8 +10,19 @@ import scanpy as sc
 Additional functionality is available in the broader {doc}`ecosystem <../ecosystem>`, with some tools being wrapped in the {mod}`scanpy.external` module.
 ```
 
+(array-support)=
+## Array type support
+
+Different APIs have different levels of support for array types,
+and this page lists the supported array types for each function:
+
+```{eval-rst}
+.. array-support:: all
+```
+
 ```{toctree}
 :maxdepth: 2
+:hidden:
 
 preprocessing
 tools

docs/conf.py

@@ -171,6 +171,39 @@
 )
 
 
+array_support: dict[str, tuple[list[str], list[str]]] = {
+    "experimental.pp.highly_variable_genes": (["np", "sp"], []),
+    "get.aggregate": (["np", "sp", "da"], []),
+    "pp.calculate_qc_metrics": (["np", "sp", "da"], []),
+    "pp.combat": (["np"], []),
+    "pp.downsample_counts": (["np", "sp[csr]"], []),
+    "pp.filter_cells": (["np", "sp", "da"], []),
+    "pp.filter_genes": (["np", "sp", "da"], []),
+    "pp.highly_variable_genes": (["np", "sp", "da"], ["da[sp[csc]]"]),
+    "pp.log1p": (["np", "sp", "da"], []),
+    "pp.neighbors": (["np", "sp"], []),
+    "pp.normalize_total": (["np", "sp[csr]", "da"], []),
+    "pp.pca": (["np", "sp", "da"], ["da[sp[csc]]"]),
+    "pp.regress_out": (["np"], []),
+    "pp.sample": (["np", "sp", "da"], []),
+    "pp.scale": (["np", "sp", "da"], []),
+    "pp.scrublet": (["np", "sp"], []),
+    "pp.scrublet_simulate_doublets": (["np", "sp"], []),
+    "tl.dendrogram": (["np", "sp"], []),
+    "tl.diffmap": (["np", "sp"], []),
+    "tl.dpt": (["np", "sp"], []),
+    "tl.draw_graph": (["np", "sp"], []),  # only uses graph in obsp
+    "tl.embedding_density": (["np"], []),
+    "tl.ingest": (["np", "sp"], []),
+    "tl.leiden": (["np", "sp"], []),  # only uses graph in obsp
+    "tl.louvain": (["np", "sp"], []),  # only uses graph in obsp
+    "tl.paga": (["np", "sp"], []),
+    "tl.rank_genes_groups": (["np", "sp"], []),
+    "tl.tsne": (["np", "sp"], []),
+    "tl.umap": (["np", "sp"], []),
+}
+
+
 # -- Options for HTML output ----------------------------------------------
 
 # The theme is sphinx-book-theme, with patches for readthedocs-sphinx-search
@@ -180,12 +213,13 @@
     "use_repository_button": True,
 }
 html_static_path = ["_static"]
+html_css_files = ["custom.css"]
 html_show_sphinx = False
 html_logo = "_static/img/Scanpy_Logo_BrightFG.svg"
 html_title = "scanpy"
 
 
-def setup(app: Sphinx):
+def setup(app: Sphinx) -> None:
     """App setup hook."""
     app.add_generic_role("small", partial(nodes.inline, classes=["small"]))
     app.add_generic_role("smaller", partial(nodes.inline, classes=["smaller"]))

docs/extensions/array_support.py

@@ -0,0 +1,185 @@
+"""Add `array-support` directive."""
+
+from __future__ import annotations
+
+from itertools import groupby
+from typing import TYPE_CHECKING
+
+from docutils import nodes
+from sphinx.util.docutils import SphinxDirective
+
+from scanpy._utils import _docs
+
+if TYPE_CHECKING:
+    from collections.abc import Collection, Generator, Iterable, Sequence
+    from typing import ClassVar
+
+    from sphinx.application import Sphinx
+
+
+ALL_INNER = list(_docs.parse(["np", "sp"], inner=True))
+
+
+class ArraySupport(SphinxDirective):
+    """Document array support."""
+
+    required_arguments: ClassVar = 1
+
+    @property
+    def _array_support(self) -> dict[str, tuple[list[str], list[str]]]:
+        return self.config.array_support
+
+    def run(self) -> list[nodes.Node]:  # noqa: D102
+        if self.arguments[0] == "all":
+            return self._render_overview()
+
+        if not self.arguments[0] not in self._array_support:
+            self.error(
+                f"API not in `array_support`, add it in `docs/conf.py`: {self.arguments[0]}"
+            )
+        array_types = list(_docs.parse(*self._array_support[self.arguments[0]]))
+        headers = (
+            "Array type",
+            "supported",
+            "… experimentally in dask :class:`~dask.array.Array`",
+        )
+        data: list[tuple[_docs.Inner, bool, bool]] = []
+        for array_type in ALL_INNER:
+            dask_array_type = _docs.DaskArray(array_type)
+            data.append((
+                array_type,
+                array_type in array_types,
+                dask_array_type in array_types,
+            ))
+
+        title = nodes.title("", "", *self.parse_inline(":ref:`array-support`")[0])
+        rows = self._render_support_data(data)
+        return self._render_table(headers, rows, title=title)
+
+    def _render_overview(self) -> list[nodes.Node]:
+        headers = ["Function", *(at.rst(short=True) for at in ALL_INNER)]
+        rows: list[nodes.row] = []
+        for fn, (include, exclude) in self._array_support.items():
+            row_header, _ = self.parse_inline(f":func:`scanpy.{fn}`")
+            ats = frozenset(_docs.parse(include, exclude))
+            cells: list[Sequence[nodes.Node]] = [
+                row_header,
+                *(
+                    self._render_support(at in ats, dask=dt in ats)
+                    for at, dt in zip(
+                        ALL_INNER, map(_docs.DaskArray, ALL_INNER), strict=True
+                    )
+                ),
+            ]
+            rows.append(
+                nodes.row(
+                    "",
+                    *(
+                        nodes.entry("", nodes.paragraph("", "", *cell))
+                        for cell in cells
+                    ),
+                )
+            )
+        return self._render_table(headers, rows)
+
+    def _render_support_data(
+        self,
+        data: list[tuple[_docs.Inner, bool, bool]],
+    ) -> Generator[nodes.row, None, None]:
+        for t, group in groupby(data, key=lambda r: type(r[0])):
+            group = list(group)  # noqa: PLW2901
+            if (  # if all sparse types have the same support, just one row
+                t is _docs.ScipySparse
+                and (support := one({s for _, s, _ in group})) is not None
+                and (in_dask := one({d for _, _, d in group})) is not None
+            ):
+                refs: list[nodes.Node] = [
+                    nodes.inline("", "scipy.sparse.{"),
+                    *self.parse_inline(":class:`csr <scipy.sparse.csr_array>`")[0],
+                    nodes.inline("", ","),
+                    *self.parse_inline(":class:`csc <scipy.sparse.csc_matrix>`")[0],
+                    nodes.inline("", "}_{"),
+                    *self.parse_inline(":class:`array <scipy.sparse.csc_array>`")[0],
+                    nodes.inline("", ","),
+                    *self.parse_inline(":class:`matrix <scipy.sparse.csr_matrix>`")[0],
+                    nodes.inline("", "}"),
+                ]
+                header = [nodes.literal("", "", *refs)]
+                yield self._render_row(header, support=support, in_dask=in_dask)
+            else:  # otherwise, show them individually
+                for array_type, support, in_dask in group:
+                    yield self._render_row(
+                        self._render_array_type(array_type),
+                        support=support,
+                        in_dask=in_dask,
+                    )
+
+    def _render_row(
+        self, header: Sequence[nodes.Node], *, support: bool, in_dask: bool
+    ) -> nodes.row:
+        cells: list[Sequence[nodes.Node]] = [
+            header,
+            self._render_support(support),
+            self._render_support(in_dask),
+        ]
+        children = (nodes.entry("", nodes.paragraph("", "", *cell)) for cell in cells)
+        return nodes.row("", *children)
+
+    def _render_table(
+        self,
+        headers: Collection[str],
+        rows: Iterable[nodes.row],
+        *,
+        title: nodes.title | None = None,
+    ) -> list[nodes.Node]:
+        colspecs = [
+            nodes.colspec(stub=True),
+            *(nodes.colspec() for _ in range(len(headers) - 1)),
+        ]
+        header_nodes = [
+            nodes.entry("", nodes.paragraph("", "", *self.parse_inline(t)[0]))
+            for t in headers
+        ]
+        thead = nodes.thead("", nodes.row("", *header_nodes))
+        tbody = nodes.tbody("", *rows)
+        return [
+            nodes.table(
+                "",
+                *([title] if title else []),
+                nodes.tgroup("", *colspecs, thead, tbody, cols=len(colspecs)),
+                ids=["array-support"],
+            )
+        ]
+
+    def _render_support(
+        self,
+        support: bool,  # noqa: FBT001
+        /,
+        *,
+        dask: bool = False,
+    ) -> Sequence[nodes.Node]:
+        dask_expl = "Also supports this type as chunk in a dask Array"
+        return [
+            nodes.Text(("✅" if support else "❌") + " " * dask),
+            *([nodes.abbreviation(text="⚡", explanation=dask_expl)] if dask else []),
+        ]
+
+    def _render_array_type(self, array_type: _docs.ArrayType, /) -> list[nodes.Node]:
+        nodes_, msgs = self.parse_inline(array_type.rst())
+        assert not msgs, msgs
+        return nodes_
+
+
+def one[T](arg: Collection[T]) -> T | None:
+    """Return the only item in `arg` or None if `arg` is not of length 1."""
+    try:
+        [item] = arg
+    except ValueError:
+        return None
+    return item
+
+
+def setup(app: Sphinx) -> None:
+    """App setup hook."""
+    app.add_directive("array-support", ArraySupport)
+    app.add_config_value("array_support", {}, "env")

docs/release-notes/3895.docs.md

@@ -0,0 +1 @@
+Document array type support for most functions in {mod}`~scanpy.pp` and {mod}`~scanpy.tl` {smaller}`P Angerer`

src/scanpy/_utils/_docs.py

@@ -0,0 +1,119 @@
+"""Utilities for `array-support` directive (see `/docs/extensions/array_support.py`)."""
+
+from __future__ import annotations
+
+import re
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, overload
+
+if TYPE_CHECKING:
+    from collections.abc import Collection, Generator
+    from typing import Literal
+
+
+__all__ = ["ArrayType", "DaskArray", "Numpy", "ScipySparse", "parse"]
+
+
+class ArrayType(ABC):
+    def rst(self, *, short: bool = False) -> str:  # pragma: no cover
+        return f":class:`{'~' if short else ''}{self}`"
+
+    @abstractmethod
+    def __hash__(self) -> int: ...
+
+
+@dataclass(unsafe_hash=True, frozen=True)
+class Numpy(ArrayType):
+    def __str__(self) -> str:  # pragma: no cover
+        return "numpy.ndarray"
+
+    def rst(self, *, short: bool = False) -> str:  # pragma: no cover
+        return f":class:`{'~' if short else ''}{self}`"
+
+
+@dataclass(unsafe_hash=True, frozen=True)
+class ScipySparse(ArrayType):
+    format: Literal["csr", "csc"]
+
+    def __str__(self) -> str:  # pragma: no cover
+        return f"scipy.sparse.{self.format}_{{array,matrix}}"
+
+    def rst(self, *, short: bool = False) -> str:  # pragma: no cover
+        return (
+            f":class:`{'~' if short else ''}scipy.sparse.{self.format}_array` / "
+            f":class:`~scipy.sparse.{self.format}_matrix`"
+        )
+
+
+type Inner = Numpy | ScipySparse
+
+
+@dataclass(unsafe_hash=True, frozen=True)
+class DaskArray(ArrayType):
+    chunk: Inner
+
+    def __str__(self) -> str:  # pragma: no cover
+        return f"dask.array.Array[{self.chunk}]"
+
+    def rst(self, *, short: bool = False) -> str:  # pragma: no cover
+        return rf":class:`{'~' if short else ''}dask.array.Array`\ \[{self.chunk.rst(short=short)}\]"
+
+
+@overload
+def parse(
+    include: Collection[str],
+    exclude: Collection[str] = (),
+    *,
+    inner: Literal[False] = False,
+) -> Generator[ArrayType]: ...
+@overload
+def parse(
+    include: Collection[str], exclude: Collection[str] = (), *, inner: Literal[True]
+) -> Generator[Inner]: ...
+def parse(
+    include: Collection[str], exclude: Collection[str] = (), *, inner: bool = False
+) -> Generator[ArrayType]:
+    if exclude:
+        excluded = dict.fromkeys(parse(exclude)).keys()
+        yield from (t for t in parse(include) if t not in excluded)
+        return
+
+    inner_includes = [i for i in include if not i.startswith("da")]
+    for t in include:
+        if (
+            match := re.fullmatch(r"([^\[]+)(?:\[(.+)\])?", t)
+        ) is None:  # pragma: no cover
+            msg = f"invalid {t!r}"
+            raise ValueError(msg)
+        mod, tags = match.groups("")
+        if mod == "da" and inner:  # pragma: no cover
+            msg = "Can’t nest dask arrays"
+            raise ValueError(msg)
+        tags = set(re.split(r",(?![^\[]+\])", tags)) if tags else set()
+        yield from _parse_mod(mod, tags, inner_includes=inner_includes)
+
+
+def _parse_mod(
+    mod: str, tags: set[str], *, inner_includes: Collection[str]
+) -> Generator[ArrayType]:
+    match mod:
+        case "np":
+            if tags:  # pragma: no cover
+                msg = f"`np` takes no tags {tags!r}"
+                raise ValueError(msg)
+            yield Numpy()
+        case "sp":
+            if tags - {"csr", "csc"}:  # pragma: no cover
+                msg = f"invalid tags {tags!r}"
+                raise ValueError(msg)
+            for format in ("csr", "csc"):
+                if tags & {"csr", "csc"} and format not in tags:
+                    continue
+                yield ScipySparse(format=format)
+        case "da":
+            for chunk in parse(tags if tags else inner_includes, inner=True):
+                yield DaskArray(chunk=chunk)
+        case _:  # pragma: no cover
+            msg = f"invalid module {mod!r}"
+            raise ValueError(msg)

src/scanpy/experimental/pp/_highly_variable_genes.py

@@ -320,6 +320,8 @@ def highly_variable_genes(  # noqa: PLR0913
 
     Expects raw count input.
 
+    .. array-support:: experimental.pp.highly_variable_genes
+
     Parameters
     ----------
     {adata}