feat: Add ncols parameter to pl.paga

docs/release-notes/4317.feat.md

@@ -0,0 +1 @@
+Add `ncols` parameter to {func}`scanpy.pl.paga` to wrap multiple panels into a grid {smaller}`Advait Shukla`

src/scanpy/plotting/legacy/_tools/paga.py

@@ -393,6 +393,7 @@ def paga(  # noqa: PLR0912, PLR0913, PLR0915
     plot: bool = True,
     show: bool | None = None,
     ax: Axes | None = None,
+    ncols: int | None = None,
     # deprecated
     save: bool | str | None = None,
 ) -> Axes | list[Axes] | None:
@@ -509,6 +510,10 @@ def paga(  # noqa: PLR0912, PLR0913, PLR0915
         Infer the filetype if ending on {{`'.pdf'`, `'.png'`, `'.svg'`}}.
     ax
         A matplotlib axes object.
+    ncols
+        Number of panels per row. If ``None`` (default), all panels are placed
+        in a single row (the original layout). If set to an integer, panels
+        wrap into a grid with this many columns.
 
     Returns
     -------
@@ -633,9 +638,21 @@ def is_flat(x):
         )
 
     if plot:
-        axs, panel_pos, draw_region_width, _figure_width = _utils.setup_axes(
-            ax, panels=colors, colorbars=colorbars
-        )
+        if ncols is not None and ax is not None:
+            msg = "`ncols` cannot be combined with a pre-supplied `ax`."
+            raise ValueError(msg)
+
+        if ncols is not None:
+            from .scatterplots import _panel_grid
+
+            fig, gs = _panel_grid(
+                hspace=0.25, wspace=0.1, ncols=ncols, num_panels=len(colors)
+            )
+            axs = [fig.add_subplot(gs[i]) for i in range(len(colors))]
+        else:
+            axs, panel_pos, draw_region_width, _figure_width = _utils.setup_axes(
+                ax, panels=colors, colorbars=colorbars
+            )
 
         if len(colors) == 1 and not isinstance(axs, list):
             axs = [axs]
@@ -677,22 +694,33 @@ def is_flat(x):
                 pos=pos,
             )
             if colorbars[icolor]:
-                if cax is None:
+                if cax is not None:
+                    ax_cb = cax[icolor]
+                    _ = plt.colorbar(
+                        sct,
+                        format=ticker.FuncFormatter(_utils.ticks_formatter),
+                        cax=ax_cb,
+                    )
+                elif ncols is not None:
+                    fig = plt.gcf()
+                    _ = fig.colorbar(
+                        sct,
+                        format=ticker.FuncFormatter(_utils.ticks_formatter),
+                        ax=axs[icolor],
+                    )
+                else:
                     bottom = panel_pos[0][0]
                     height = panel_pos[1][0] - bottom
                     width = 0.006 * draw_region_width / len(colors)
                     left = panel_pos[2][2 * icolor + 1] + 0.2 * width
                     rectangle = [left, bottom, width, height]
                     fig = plt.gcf()
                     ax_cb = fig.add_axes(rectangle)
-                else:
-                    ax_cb = cax[icolor]
-
-                _ = plt.colorbar(
-                    sct,
-                    format=ticker.FuncFormatter(_utils.ticks_formatter),
-                    cax=ax_cb,
-                )
+                    _ = plt.colorbar(
+                        sct,
+                        format=ticker.FuncFormatter(_utils.ticks_formatter),
+                        cax=ax_cb,
+                    )
     if add_pos:
         adata.uns["paga"]["pos"] = pos
         logg.hint("added 'pos', the PAGA positions (adata.uns['paga'])")

tests/plotting/legacy/test_paga.py

@@ -3,9 +3,13 @@
 from functools import partial
 from importlib.util import find_spec
 
+import numpy as np
+import pandas as pd
 import pytest
 from matplotlib import colormaps
+from matplotlib import pyplot as plt
 from packaging.version import Version
+from scipy import sparse
 
 import scanpy as sc
 from scanpy._compat import pkg_version
@@ -96,3 +100,48 @@ def test_paga_compare(plot_cmp):
     sc.pl.paga_compare(pbmc, basis="umap", show=False)
 
     plot_cmp("paga_compare_pbmc3k")
+
+
+def test_paga_ncols() -> None:
+    # Tests that https://github.com/scverse/scanpy/issues/1203 is fixed
+    rng = np.random.default_rng(0)
+    adata = sc.AnnData(rng.random((80, 20)))
+    adata.obs["group"] = pd.Categorical(rng.choice(["a", "b", "c", "d", "e"], 80))
+    for i in range(4):
+        adata.obs[f"c{i}"] = pd.Categorical(rng.choice(["x", "y"], 80))
+
+    k = 5
+    rows = np.array([0, 1, 1, 2, 2, 3, 3, 4, 4, 0])
+    cols = np.array([1, 0, 2, 1, 3, 2, 4, 3, 0, 4])
+    connectivities = sparse.csr_matrix(  # noqa: TID251
+        (np.ones(len(rows)), (rows, cols)), shape=(k, k)
+    )
+    adata.uns["paga"] = {
+        "groups": "group",
+        "connectivities": connectivities,
+        "connectivities_tree": connectivities.copy(),
+    }
+    pos = rng.random((k, 2))
+    colors = ["c0", "c1", "c2", "c3"]
+
+    # `ncols` wraps the panels into a grid
+    axs = sc.pl.paga(adata, color=colors, ncols=2, pos=pos, show=False)
+    assert len(axs) == 4
+    gridspec = axs[0].get_subplotspec().get_gridspec()
+    assert (gridspec.nrows, gridspec.ncols) == (2, 2)
+
+    # the default layout keeps all panels in a single row
+    axs = sc.pl.paga(adata, color=colors, pos=pos, show=False)
+    assert len(axs) == 4
+
+    # continuous colors (with colorbars) also wrap into a grid
+    gene_colors = adata.var_names[:3].tolist()
+    axs = sc.pl.paga(adata, color=gene_colors, ncols=2, pos=pos, show=False)
+    assert len(axs) == 3
+    gridspec = axs[0].get_subplotspec().get_gridspec()
+    assert gridspec.ncols == 2
+
+    # `ncols` cannot be combined with a pre-supplied `ax`
+    _, ax = plt.subplots()
+    with pytest.raises(ValueError, match="`ncols` cannot be combined"):
+        sc.pl.paga(adata, color=colors, ncols=2, ax=ax, pos=pos, show=False)