fix: render heatmap metadata colour bars

CHANGELOG.md

@@ -37,6 +37,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   The score box/strip plots are now faceted: `facet_rows` adds grid rows and
   `facet_cols` puts the facet on the columns (moving the latent variables onto
   the rows). Setting both is rejected at config-parse time with a clear error.
+- Heatmap metadata colour bars (`--x-meta`/`--y-meta`) were passed to the
+  plotting layer but silently dropped, because `sns.heatmap` cannot draw
+  row/column colour strips. The rank-1 and bootstrap-ratio heatmaps now
+  render the category colour bars via a non-clustered `clustermap`; heatmaps
+  without metadata are unchanged.
 - `plsdo.__version__` was hardcoded to `0.1.0` while the package was `0.1.1`;
   the version is now correct and single-sourced.
 

docs/usage.md

@@ -184,6 +184,10 @@ and `y_bootstrap_ratios.csv` are not filtered.
 
 `--all-plots` generates additional diagnostic figures (scree plot,
 rank-1 heatmaps, bootstrap ratio heatmaps, raw feature distributions).
+When `--x-meta`/`--y-meta` are supplied, the rank-1 and bootstrap-ratio
+heatmaps gain category colour bars alongside their rows (X features) and
+columns (Y features), using the same metadata categories as the loading
+plots.
 When the feature count exceeds 100, heatmaps and distribution plots
 become unreadable and extremely slow, so only the scree plot is
 produced.

plsdo/plotting.py

@@ -14,6 +14,11 @@
 # Threshold above which heatmap annotations are suppressed
 ANNOTATION_THRESHOLD = 30
 
+# clustermap always allocates space for dendrograms; with clustering disabled
+# they are empty, so this near-zero ratio collapses that gutter to (almost)
+# nothing. It cannot be 0 — clustermap requires a positive ratio.
+COLLAPSED_DENDROGRAM_RATIO = 0.02
+
 # Maximum number of features before verbose plots are skipped
 VERBOSE_FEATURE_LIMIT = 100
 
@@ -94,6 +99,24 @@ def plot_heatmap(
     n_rows, n_cols = data.shape
     figsize = figure_size(n_rows, n_cols)
     annotate = n_rows <= ANNOTATION_THRESHOLD and n_cols <= ANNOTATION_THRESHOLD
+    tick_fontsize = max(4.0, min(10.0, 200 / max(n_rows, n_cols)))
+
+    if row_colors is not None or col_colors is not None:
+        _heatmap_with_colour_bars(
+            data,
+            v=v,
+            xticklabels=xticklabels,
+            yticklabels=yticklabels,
+            out_path=out_path,
+            subtitle=subtitle,
+            row_colors=row_colors,
+            col_colors=col_colors,
+            figsize=figsize,
+            annotate=annotate,
+            tick_fontsize=tick_fontsize,
+            dpi=dpi,
+        )
+        return
 
     fig, ax = plt.subplots(figsize=figsize)
     sns.heatmap(
@@ -109,7 +132,6 @@ def plot_heatmap(
         annot=annotate,
         fmt=".2f" if annotate else "",
     )
-    tick_fontsize = max(4.0, min(10.0, 200 / max(n_rows, n_cols)))
     ax.set_xticklabels(
         ax.get_xticklabels(), rotation=45, ha="right",
         fontsize=tick_fontsize,
@@ -122,6 +144,66 @@ def plot_heatmap(
     _finalise(fig, out_path, dpi)
 
 
+def _heatmap_with_colour_bars(
+    data: np.ndarray,
+    *,
+    v: float,
+    xticklabels: list[str],
+    yticklabels: list[str],
+    out_path: Path,
+    subtitle: Optional[str],
+    row_colors: Optional[list],
+    col_colors: Optional[list],
+    figsize: tuple[float, float],
+    annotate: bool,
+    tick_fontsize: float,
+    dpi: int,
+) -> None:
+    """Heatmap with metadata colour bars beside the rows and/or columns.
+
+    ``sns.heatmap`` cannot draw row/column colour strips, so a ``clustermap``
+    with both dendrograms disabled is used instead: it lays out the colour
+    bars alongside the heatmap while preserving the original row/column order
+    (``row_cluster=False``, ``col_cluster=False``). The diverging ``vlag``
+    scale, symmetric range, and annotation behaviour match the plain heatmap.
+    The colour bar itself is positioned and sized by ``clustermap`` (which
+    manages its own layout), so it will not match the plain heatmap's
+    ``shrink``-ed colour bar exactly.
+    """
+    g = sns.clustermap(
+        data,
+        row_cluster=False,
+        col_cluster=False,
+        row_colors=row_colors,
+        col_colors=col_colors,
+        vmin=-v,
+        vmax=v,
+        center=0,
+        cmap="vlag",
+        xticklabels=xticklabels,
+        yticklabels=yticklabels,
+        annot=annotate,
+        fmt=".2f" if annotate else "",
+        figsize=figsize,
+        dendrogram_ratio=COLLAPSED_DENDROGRAM_RATIO,
+    )
+    # No clustering, so the (empty) dendrogram axes are just wasted space.
+    g.ax_row_dendrogram.set_visible(False)
+    g.ax_col_dendrogram.set_visible(False)
+
+    ax = g.ax_heatmap
+    ax.set_xticklabels(
+        ax.get_xticklabels(), rotation=45, ha="right", fontsize=tick_fontsize
+    )
+    ax.set_yticklabels(
+        ax.get_yticklabels(), rotation=0, fontsize=tick_fontsize
+    )
+    if subtitle:
+        g.figure.suptitle(subtitle)
+    g.savefig(out_path, transparent=False, dpi=dpi)
+    plt.close(g.figure)
+
+
 def plot_permutation(
     observed_s: np.ndarray,
     permuted_s: np.ndarray,

tests/test_plotting.py

@@ -41,6 +41,25 @@ def test_tiny_data_floored(self):
         assert h >= 4
 
 
+def _spy_clustermap(plotting_mod, monkeypatch):
+    """Wrap sns.clustermap so tests can inspect the ClusterGrid it returns."""
+    captured = {}
+    real = plotting_mod.sns.clustermap
+
+    def spy(*args, **kwargs):
+        grid = real(*args, **kwargs)
+        captured["grid"] = grid
+        return grid
+
+    monkeypatch.setattr(plotting_mod.sns, "clustermap", spy)
+    return captured
+
+
+def _axis_has_artists(ax):
+    """True if a colour-bar axis exists and actually drew something."""
+    return ax is not None and (len(ax.collections) + len(ax.images)) > 0
+
+
 class TestPlotHeatmap:
     def test_saves_file(self, tmp_output):
         data = np.random.default_rng(0).standard_normal((5, 4))
@@ -55,6 +74,24 @@ def test_saves_file(self, tmp_output):
         assert out.exists()
         assert out.stat().st_size > 0
 
+    def test_no_colour_bars_uses_plain_heatmap(self, tmp_output, monkeypatch):
+        """Without colour bars the plain sns.heatmap path is used unchanged."""
+        from plsdo import plotting as plotting_mod
+
+        def fail(*args, **kwargs):
+            raise AssertionError("clustermap must not be used without colour bars")
+
+        monkeypatch.setattr(plotting_mod.sns, "clustermap", fail)
+        out = tmp_output / "plain_heatmap.svg"
+        plot_heatmap(
+            np.zeros((3, 3)),
+            v=1.0,
+            xticklabels=["a", "b", "c"],
+            yticklabels=["x", "y", "z"],
+            out_path=out,
+        )
+        assert out.exists()
+
     def test_annotations_suppressed_for_large_data(self, tmp_output, monkeypatch):
         from plsdo import plotting as plotting_mod
 
@@ -455,6 +492,29 @@ def test_each_lv_different_output(self, tmp_output):
             )
             assert out.exists()
 
+    def test_metadata_colours_rendered_on_both_axes(self, tmp_output, monkeypatch):
+        """X/Y metadata colours must appear as row and column colour bars."""
+        from plsdo import plotting as plotting_mod
+
+        captured = _spy_clustermap(plotting_mod, monkeypatch)
+        rng = np.random.default_rng(0)
+        out = tmp_output / "lv_heatmap_meta.svg"
+        plot_lv_heatmap(
+            lv_idx=0,
+            u=rng.standard_normal((5, 3)),
+            s=np.array([2.0, 1.5, 0.5]),
+            vt=rng.standard_normal((3, 4)),
+            x_feature_names=["x1", "x2", "x3", "x4", "x5"],
+            y_feature_names=["y1", "y2", "y3", "y4"],
+            x_colours=["red", "red", "blue", "blue", "green"],
+            y_colours=["orange", "orange", "purple", "purple"],
+            out_path=out,
+        )
+        assert out.exists()
+        g = captured["grid"]
+        assert _axis_has_artists(g.ax_row_colors)
+        assert _axis_has_artists(g.ax_col_colors)
+
 
 class TestPlotBootstrapHeatmap:
     def test_saves_file(self, tmp_output):
@@ -482,6 +542,26 @@ def test_single_lv(self, tmp_output):
         )
         assert out.exists()
 
+    def test_metadata_colours_rendered_on_rows_only(self, tmp_output, monkeypatch):
+        """Feature metadata colours appear as a row colour bar; columns are LVs
+        (no metadata), so there is no column colour bar."""
+        from plsdo import plotting as plotting_mod
+
+        captured = _spy_clustermap(plotting_mod, monkeypatch)
+        rng = np.random.default_rng(0)
+        out = tmp_output / "bsr_meta.svg"
+        plot_bootstrap_heatmap(
+            bootstrap_ratios=rng.standard_normal((5, 2)),
+            feature_names=["f1", "f2", "f3", "f4", "f5"],
+            lv_names=["LV1", "LV2"],
+            colours=["red", "red", "blue", "blue", "green"],
+            out_path=out,
+        )
+        assert out.exists()
+        g = captured["grid"]
+        assert _axis_has_artists(g.ax_row_colors)
+        assert g.ax_col_colors is None
+
 
 class TestPlotRawDistributions:
     def test_saves_file(self, tmp_output):