Update heatmap (#1160)

Add argument `title`.

src/cellrank/pl/_heatmap.py

@@ -77,6 +77,7 @@ def heatmap(
     lineage_height: float = 0.33,
     fontsize: Optional[float] = None,
     xlabel: Optional[str] = None,
+    title: Optional[str] = None,
     cmap: colors.ListedColormap = cm.viridis,
     dendrogram: bool = True,
     return_genes: bool = False,
@@ -147,6 +148,8 @@ def heatmap(
         Size of the title's font.
     xlabel
         Label on the x-axis. If :obj:`None`, it is determined based on ``time_key``.
+    title
+        Title of the figure.
     cmap
         Colormap to use when visualizing the smoothed expression.
     dendrogram
@@ -340,6 +343,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
 
@@ -443,6 +448,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