feat: add configurable columns to PAGA plots

docs/release-notes/1203.feat.md

@@ -0,0 +1 @@
+Add an ``ncols`` parameter to {func}`scanpy.pl.paga` so multi-color PAGA plots can be arranged in a configurable grid instead of a single row. {smaller}`AtomicGlance`

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

@@ -384,6 +384,7 @@ def paga(  # noqa: PLR0912, PLR0913, PLR0915
     cax: Axes | None = None,
     colorbar=None,  # TODO: this seems to be unused
     cb_kwds: Mapping[str, Any] = frozendict({}),
+    ncols: int | None = None,
     frameon: bool | None = None,
     add_pos: bool = True,
     export_to_gexf: bool = False,
@@ -495,6 +496,9 @@ def paga(  # noqa: PLR0912, PLR0913, PLR0915
     cb_kwds
         Keyword arguments for :class:`~matplotlib.colorbar.Colorbar`,
         for instance, `ticks`.
+    ncols
+        Number of panels shown per row when plotting multiple colors. If
+        `None`, all panels are shown in a single row.
     add_pos
         Add the positions to `adata.uns['paga']`.
     title
@@ -547,6 +551,13 @@ def paga(  # noqa: PLR0912, PLR0913, PLR0915
 
     """
     rng = np.random.default_rng(rng)
+    if ncols is not None:
+        if isinstance(ncols, bool) or not isinstance(ncols, int) or ncols < 1:
+            msg = "`ncols` must be a positive integer or `None`."
+            raise ValueError(msg)
+        if ax is not None:
+            msg = "Cannot specify `ncols` when passing a pre-supplied `ax`."
+            raise ValueError(msg)
     if groups is not None:  # backwards compat
         labels = groups
         logg.warning("`groups` is deprecated in `pl.paga`: use `labels` instead")
@@ -633,9 +644,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 None:
+            axs, panel_pos, draw_region_width, _figure_width = _utils.setup_axes(
+                ax, panels=colors, colorbars=colorbars
+            )
+        else:
+            from .scatterplots import _panel_grid
+
+            fig, grid = _panel_grid(
+                hspace=0.25,
+                wspace=0.1,
+                ncols=ncols,
+                num_panels=len(colors),
+            )
+            axs = [fig.add_subplot(grid[i]) for i in range(len(colors))]
+            panel_pos = None
 
         if len(colors) == 1 and not isinstance(axs, list):
             axs = [axs]
@@ -677,7 +700,21 @@ def is_flat(x):
                 pos=pos,
             )
             if colorbars[icolor]:
-                if cax is None:
+                if ncols is not None:
+                    if cax is None:
+                        colorbar_kwds = {
+                            "format": ticker.FuncFormatter(_utils.ticks_formatter)
+                        }
+                        colorbar_kwds.update(cb_kwds)
+                        _ = axs[icolor].figure.colorbar(
+                            sct,
+                            ax=axs[icolor],
+                            use_gridspec=False,
+                            **colorbar_kwds,
+                        )
+                        continue
+                    ax_cb = cax[icolor] if isinstance(cax, Sequence) else cax
+                elif cax is None:
                     bottom = panel_pos[0][0]
                     height = panel_pos[1][0] - bottom
                     width = 0.006 * draw_region_width / len(colors)

tests/plotting/legacy/test_paga.py

@@ -5,6 +5,7 @@
 
 import pytest
 from matplotlib import colormaps
+from matplotlib import pyplot as plt
 from packaging.version import Version
 
 import scanpy as sc
@@ -76,6 +77,44 @@ def test_paga_pie(plot_cmp, pbmc) -> None:
     plot_cmp("paga_pie")
 
 
+@SKIP_IF_OLD_IGRAPH
+@pytest.mark.parametrize(("ncols", "expected_shape"), [(1, (3, 1)), (2, (2, 2))])
+def test_paga_ncols(pbmc, ncols, expected_shape) -> None:
+    axs = sc.pl.paga(
+        pbmc,
+        color=["CST3", "GATA2", "cool_feature"],
+        ncols=ncols,
+        show=False,
+    )
+
+    assert len(axs) == 3
+    assert axs[0].get_subplotspec().get_gridspec().get_geometry() == expected_shape
+    plt.close(axs[0].figure)
+
+
+@pytest.mark.parametrize("ncols", [0, -1, True, 1.5])
+def test_paga_ncols_rejects_invalid(ncols) -> None:
+    with pytest.raises(ValueError, match=r"ncols.*positive integer"):
+        sc.pl.paga(None, ncols=ncols)
+
+
+def test_paga_ncols_rejects_ax() -> None:
+    fig, ax = plt.subplots()
+    with pytest.raises(ValueError, match="Cannot specify `ncols`"):
+        sc.pl.paga(None, ncols=2, ax=ax)
+    plt.close(fig)
+
+
+@SKIP_IF_OLD_IGRAPH
+def test_paga_ncols_custom_colorbar(pbmc) -> None:
+    cax_fig, cax = plt.subplots()
+    ax = sc.pl.paga(pbmc, color="CST3", ncols=1, cax=cax, show=False)
+
+    assert ax.figure is not cax_fig
+    plt.close(ax.figure)
+    plt.close(cax_fig)
+
+
 def test_paga_path(plot_cmp, pbmc) -> None:
     pbmc.uns["iroot"] = 0
     sc.tl.dpt(pbmc)