From 2c080cfea6e528c8eda4df3aeaddd8c1269af56a Mon Sep 17 00:00:00 2001
From: Philipp Weiler <weiler.philipp@gmail.com>
Date: Thu, 22 Feb 2024 08:03:18 +0000
Subject: [PATCH] Update `heatmap`

Add argument `title`.
---
 src/cellrank/pl/_heatmap.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/cellrank/pl/_heatmap.py b/src/cellrank/pl/_heatmap.py
index be49a7b84..6387cafa8 100644
--- a/src/cellrank/pl/_heatmap.py
+++ b/src/cellrank/pl/_heatmap.py
@@ -89,6 +89,7 @@ def heatmap(
     dpi: Optional[int] = None,
     save: Optional[Union[str, pathlib.Path]] = None,
     gene_order: Optional[Sequence[str]] = None,
+    title: Optional[str] = None,
     **kwargs: Any,
 ) -> Optional[Union[Dict[str, pd.DataFrame], Tuple[_return_model_type, Dict[str, pd.DataFrame]]]]:
     """Plot a heatmap of smoothed gene expression along specified lineages.
@@ -340,6 +341,8 @@ def color_fill_rec(ax, xs, y1, y2, colors=None, cmap=cmap, **kwargs) -> None:
             labelbottom=True,
         )
         ax.set_xlabel(xlabel)
+        if title is not None:
+            ax.set_title(title)
 
         return fig, None
 
@@ -442,6 +445,8 @@ def _(gene_order: Optional[Sequence[str]] = None) -> Tuple[List[plt.Figure], pd.
             g.ax_heatmap.set_xlabel(xlabel)
             g.ax_heatmap.set_xticks(np.linspace(0, len(df.columns), _N_XTICKS))
             g.ax_heatmap.set_xticklabels([round(n, 3) for n in np.linspace(x_min, x_max, _N_XTICKS)])
+            if title is not None:
+                g.ax_heatmap.set_title(title)
             # fmt: on
             if show_clust:
                 # robustly show dendrogram, because gene names can be long