Fix clustermap (#163)

Author: ahuang11

src/hv_anndata/plotting/clustermap.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 from typing import TYPE_CHECKING, TypedDict
 
 import anndata as ad
@@ -61,8 +62,12 @@ def create_clustermap_plot(
     if hasattr(x, "toarray"):
         x = x.toarray()
 
+    # Take var_names from the same space as x: adata.raw has a different
+    # (usually larger) gene set than adata.var, so mixing them makes the
+    # variance-ranked indices overrun adata.var_names.
+    var_names = adata.raw.var_names if use_raw else adata.var_names
+
     # Filter genes if max_genes is specified
-    var_names = adata.var_names
     if max_genes is not None and len(var_names) > max_genes:
         gene_vars = np.var(x, axis=0)
         top_gene_indices = np.argsort(gene_vars)[-max_genes:]
@@ -97,10 +102,12 @@ def create_clustermap_plot(
         "show_grid": False,
         "tools": ["hover"],
         "colorbar": True,
+        "responsive": True,
     }
 
     return clustered_plot.opts(
-        hv.opts.HeatMap(**main_plot_opts), hv.opts.Dendrogram(xaxis=None, yaxis=None)
+        hv.opts.HeatMap(**main_plot_opts),
+        hv.opts.Dendrogram(xaxis=None, yaxis=None, responsive=True),
     )
 
 
@@ -146,12 +153,19 @@ def __init__(self, adata: AnnData | None = None, **params: object) -> None:
         """Initialize the ClusterMap with the given parameters."""
         if adata is not None:
             params["adata"] = adata
-        super().__init__(**params)
 
-        # Check if raw data exists
-        has_raw = self.adata.raw is not None
-        if self.use_raw is None:
-            self.use_raw = has_raw
+        # Resolve use_raw's default (if not given, or given as None) before
+        # calling super().__init__ so it's set as part of construction rather
+        # than as a later attribute assignment. A later assignment would fire
+        # the async _update_plot watcher below, which relies on a running
+        # event loop to schedule -- not guaranteed outside of a live Panel
+        # session (e.g. plain scripts and tests), where it has been observed
+        # to leak an event loop/socket.
+        has_raw = adata is not None and adata.raw is not None
+        if params.get("use_raw") is None:
+            params["use_raw"] = has_raw
+
+        super().__init__(**params)
 
         use_raw_widget = pmui.widgets.Checkbox.from_param(
             self.param.use_raw,
@@ -160,6 +174,12 @@ def __init__(self, adata: AnnData | None = None, **params: object) -> None:
             disabled=not has_raw,  # Disable if no raw data
         )
 
+        # Shown immediately; later recomputes (triggered by widget changes,
+        # which only ever happen inside a running Panel session) go through
+        # the async _update_plot watcher, which swaps in the finished plot or
+        # an error message behind a loading indicator.
+        self._plot_placeholder = pn.pane.Placeholder(width=300, height=300)
+
         self._widgets = pmui.Column(
             pmui.widgets.Select.from_param(
                 self.param.cmap,
@@ -175,29 +195,49 @@ def __init__(self, adata: AnnData | None = None, **params: object) -> None:
             visible=self.param.show_widgets,
             sx={"border": 1, "borderColor": "#e3e3e3", "borderRadius": 1},
             sizing_mode="stretch_width",
-            max_width=400,
+            max_width=300,
         )
 
-    @param.depends("use_raw", "cmap", "max_genes", "plot_opts")
-    def _plot_view(self) -> pn.viewable.Viewable:
-        """Create the plot view with parameter dependencies."""
+        # Build the first plot directly and synchronously; see the use_raw
+        # comment above for why this doesn't go through _update_plot.
+        self._plot_placeholder.object = self._build_plot_or_error()
+
+    def _build_plot_or_error(self) -> pn.viewable.Viewable:
+        """Build the clustermap pane, or an Alert describing why it failed."""
         config = ClusterMapConfig(
             cmap=self.cmap,
         )
-
-        plot = create_clustermap_plot(
-            self.adata,
-            use_raw=self.use_raw,
-            max_genes=self.max_genes,
-            **config,
-        )
-
-        # Apply user-provided plot options to the HeatMap only
-        if self.plot_opts:
-            plot = plot.opts(hv.opts.HeatMap(**self.plot_opts))
-
-        return plot
+        try:
+            plot = create_clustermap_plot(
+                self.adata,
+                use_raw=self.use_raw,
+                max_genes=self.max_genes,
+                **config,
+            )
+            if self.plot_opts:
+                plot = plot.opts(hv.opts.HeatMap(**self.plot_opts))
+            return pn.pane.HoloViews(
+                plot, sizing_mode="stretch_both", min_height=550, min_width=450
+            )
+        except Exception as e:  # noqa: BLE001
+            msg = f"Could not render clustermap: {e}"
+            if pn.state.notifications is not None:
+                pn.state.notifications.error(msg)
+            return pn.pane.Alert(msg, alert_type="danger", sizing_mode="stretch_width")
+
+    @param.depends("use_raw", "cmap", "max_genes", "plot_opts", watch=True)
+    async def _update_plot(self) -> None:
+        """Recompute the clustermap and swap it into the placeholder."""
+        # Hold the loading overlay across the whole update, including the object
+        # assignment, so it doesn't vanish while the plot is still being built
+        # and serialized. Yield first so it paints, and run the slow clustering
+        # off the event loop.
+        with self._plot_placeholder.param.update(loading=True):
+            await asyncio.sleep(0.01)
+            self._plot_placeholder.object = await asyncio.to_thread(
+                self._build_plot_or_error
+            )
 
     def __panel__(self) -> pn.viewable.Viewable:
         """Create the Panel application layout."""
-        return pmui.Row(self._widgets, self._plot_view)
+        return pmui.Row(self._widgets, self._plot_placeholder)

tests/plotting/test_clustermap.py

@@ -2,14 +2,17 @@
 
 from __future__ import annotations
 
+import asyncio
+
 import anndata as ad
 import numpy as np
 import pandas as pd
+import panel as pn
 import panel_material_ui as pmui
 import pytest
 import scanpy as sc
 
-from hv_anndata import ClusterMap
+from hv_anndata import ClusterMap, create_clustermap_plot
 
 
 @pytest.fixture
@@ -68,8 +71,138 @@ def test_clustermap_no_raw_data() -> None:
     assert cm.use_raw is False  # Should default to False when no raw data
 
 
+@pytest.mark.usefixtures("bokeh_renderer")
+def test_clustermap_builds_plot_synchronously_on_construction(
+    sadata: ad.AnnData,
+) -> None:
+    """Construction alone, with no async machinery involved, must render the plot.
+
+    Building the first plot via the async _update_plot watcher instead would
+    depend on a running event loop to schedule it -- not guaranteed outside a
+    live Panel session (e.g. plain scripts and tests) -- and was observed to
+    leak an event loop/socket when construction happened outside of one;
+    see #163.
+    """
+    cm = ClusterMap(adata=sadata)
+
+    assert isinstance(cm._plot_placeholder.object, pn.pane.HoloViews)
+
+
+@pytest.mark.usefixtures("bokeh_renderer")
+def test_clustermap_explicit_use_raw_builds_plot_synchronously(
+    sadata: ad.AnnData,
+) -> None:
+    """Passing use_raw explicitly must not require a later trigger to render."""
+    cm = ClusterMap(adata=sadata, use_raw=True)
+
+    assert cm.use_raw is True
+    assert isinstance(cm._plot_placeholder.object, pn.pane.HoloViews)
+
+
 @pytest.mark.usefixtures("bokeh_renderer")
 def test_integration() -> None:
     adata = sc.datasets.pbmc68k_reduced()  # errors
 
     assert ClusterMap(adata=adata).__panel__()
+
+
+@pytest.fixture
+def sadata_mismatched_raw() -> ad.AnnData:
+    """AnnData whose raw layer has far more genes than the primary var.
+
+    Used to exercise the max_genes filtering path with use_raw=True, where
+    gene-variance indices are computed over adata.raw.X's (larger) gene axis.
+
+    Returns
+    -------
+    AnnData with 3 main-var genes and 20 raw-var genes.
+
+    """
+    n_obs = 10
+    n_vars = 3
+    n_raw_vars = 20
+
+    rng = np.random.default_rng()
+
+    x = rng.random((n_obs, n_vars))
+    obs = pd.DataFrame(
+        {"cell_type": ["A", "B"] * (n_obs // 2)},
+        index=[f"cell_{i}" for i in range(n_obs)],
+    )
+    var = pd.DataFrame(index=[f"gene_{i}" for i in range(n_vars)])
+    raw_x = rng.random((n_obs, n_raw_vars))
+    raw_var = pd.DataFrame(index=[f"raw_gene_{i}" for i in range(n_raw_vars)])
+    adata = ad.AnnData(X=x, obs=obs, var=var)
+    adata.raw = ad.AnnData(X=raw_x, var=raw_var, obs=obs)
+    return adata
+
+
+@pytest.mark.usefixtures("bokeh_renderer")
+def test_create_clustermap_plot_use_raw_var_names_from_raw(
+    sadata_mismatched_raw: ad.AnnData,
+) -> None:
+    """max_genes filtering with use_raw must index raw.var_names, not var_names.
+
+    Pre-fix this raised an IndexError: the filter indices are computed from
+    adata.raw.X's (larger) gene axis but were used to index the (shorter)
+    adata.var_names; see #163.
+    """
+    plot = create_clustermap_plot(sadata_mismatched_raw, use_raw=True, max_genes=10)
+
+    genes = set(plot.dimension_values("variable", expanded=False))
+    assert len(genes) == 10
+    assert all(gene.startswith("raw_gene_") for gene in genes)
+
+
+@pytest.mark.usefixtures("bokeh_renderer")
+def test_create_clustermap_plot_heatmap_is_responsive(
+    sadata_mismatched_raw: ad.AnnData,
+) -> None:
+    """The HeatMap must be responsive so it doesn't overflow at large gene counts."""
+    plot = create_clustermap_plot(sadata_mismatched_raw, use_raw=True, max_genes=10)
+
+    assert plot.main.opts.get("plot").kwargs["responsive"] is True
+
+
+@pytest.mark.usefixtures("bokeh_renderer")
+# asyncio.run() spins up and tears down a fresh loop; Panel/Bokeh's IOLoop and
+# thread-pool teardown from asyncio.to_thread emit ResourceWarnings on gc that
+# land here (or on whatever test runs next), unrelated to this test's outcome.
+@pytest.mark.filterwarnings("ignore::ResourceWarning")
+def test_clustermap_update_plot_use_raw_with_mismatched_var_names(
+    sadata_mismatched_raw: ad.AnnData,
+) -> None:
+    """The ClusterMap widget must render, not crash, through the same code path."""
+    cm = ClusterMap(adata=sadata_mismatched_raw, use_raw=True, max_genes=10)
+
+    asyncio.run(cm._update_plot())
+
+    assert isinstance(cm._plot_placeholder.object, pn.pane.HoloViews)
+    assert cm._plot_placeholder.loading is False
+
+
+@pytest.mark.usefixtures("bokeh_renderer")
+@pytest.mark.filterwarnings("ignore::ResourceWarning")
+def test_clustermap_update_plot_error_sets_alert(
+    sadata: ad.AnnData, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Errors during recompute must surface as an Alert, not propagate.
+
+    The recompute runs off the main thread via asyncio.to_thread, so an
+    unhandled exception there would otherwise be silently lost; see #163.
+    """
+
+    def _boom(*_args: object, **_kwargs: object) -> None:
+        msg = "kaboom"
+        raise ValueError(msg)
+
+    cm = ClusterMap(adata=sadata)
+    monkeypatch.setattr("hv_anndata.plotting.clustermap.create_clustermap_plot", _boom)
+
+    asyncio.run(cm._update_plot())
+
+    alert = cm._plot_placeholder.object
+    assert isinstance(alert, pn.pane.Alert)
+    assert alert.alert_type == "danger"
+    assert "Could not render clustermap" in alert.object
+    assert cm._plot_placeholder.loading is False