Multi-year drought heatmap

src/spei/plot.py

@@ -5,7 +5,7 @@
 import matplotlib as mpl
 import matplotlib.pyplot as plt
 from matplotlib.dates import date2num
-from numpy import array, concatenate, linspace, meshgrid, reshape
+from numpy import arange, array, concatenate, linspace, meshgrid, reshape
 from pandas import DatetimeIndex, Series
 from scipy.stats import gaussian_kde
 
@@ -80,6 +80,7 @@ def si(
         lc.set_linewidth(1.2)
         _ = ax.add_collection(lc)
 
+    ax.yaxis.set_major_locator(mpl.ticker.MultipleLocator(1))
     ax.set_ylim(nmin, nmax)
 
     return ax
@@ -150,6 +151,86 @@ def monthly_density(
     return ax
 
 
+def heatmap(
+    sis: List[Series],
+    cmap: str = "Reds_r",
+    vmin: float = -3.0,
+    vmax: float = -1.0,
+    yticklabels: List[str] | None = None,
+    ax: Optional[Axes] = None,
+) -> Axes:
+    """
+    Plots multiple standardized indices on a heatmap from [mourik_2024]_
+
+    Parameters
+    ----------
+    sis : List[Series]
+        A list of pandas Series objects, each representing a time series of SI values.
+    cmap : str, optional
+        The colormap to use for the heatmap. Default is "Reds_r".
+    vmin : float, optional
+        The minimum value for color normalization. Default is -3.0.
+    vmax : float, optional
+        The maximum value for color normalization. Default is -1.0.
+    yticklabels : List[str] or None, optional
+        Custom labels for the y-axis ticks. If None, the names of the Series objects are used. Default is None.
+    ax : Axes, optional
+        A matplotlib Axes object to plot on. If None, a new figure and axes are created. Default is None.
+    Returns
+    -------
+    Axes
+        The matplotlib Axes object with the heatmap.
+
+    References
+    ----------
+    .. [mourik_2024] van Mourik, J., Ruijsch, D., van der Wiel, K., Hazeleger,
+    W., Wanders, N.: Regional drivers and characteristics of multi-year
+    droughts. 2024
+    """
+    if ax is None:
+        fig, ax = plt.subplots(figsize=(6.5, 4))
+    else:
+        fig = ax.get_figure()
+
+    if cmap in Crameri._available_cmaps:
+        colormap = Crameri(cmap).cmap
+    else:
+        colormap = plt.get_cmap(cmap)
+
+    norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
+    for i, s in enumerate(sis):
+        _ = ax.pcolormesh(
+            [s.index, s.index],
+            [i, i + 1],
+            [s.values[:-1]],
+            norm=norm,
+            cmap=colormap,
+            shading="flat",
+        )
+
+    ax.set_yticks(arange(0.5, len(sis) + 0.5, 1.0), minor=False)
+    ax.set_yticks(arange(0.0, len(sis) + 1.5, 1.0), minor=True)
+    yticklabels = (
+        [getattr(s, "name") for s in sis] if yticklabels is None else yticklabels
+    )
+    ax.set_yticklabels(yticklabels)
+    for tick in ax.yaxis.get_major_ticks():  # don't show major ytick marker
+        tick.tick1line.set_visible(False)
+
+    ax.set_ylim(0, len(sis))
+    scm = mpl.cm.ScalarMappable(norm=norm, cmap=colormap)
+    cax, cbar_kw = mpl.colorbar.make_axes(
+        ax, fraction=0.05, pad=0.01, orientation="vertical"
+    )
+    _ = fig.colorbar(
+        scm,
+        cax=cax,
+        **cbar_kw,
+    )
+
+    return ax
+
+
 class Crameri:
     _available_cmaps = ("roma", "roma_r", "vik", "vik_r", "lajolla", "lajolla_r")
 

tests/test_plots.py

@@ -1,6 +1,6 @@
 import matplotlib as mpl
 from pandas import Series
-from spei.plot import monthly_density
+from spei.plot import heatmap, monthly_density
 from spei.plot import si as plot_si
 
 mpl.use("Agg")  # prevent _tkinter.TclError: Can't find a usable tk.tcl error
@@ -16,3 +16,7 @@ def test_plot_si_no_background(si: Series) -> None:
 
 def test_plot_monthly_density(si: Series) -> None:
     _ = monthly_density(si, years=[2011], months=[1, 2, 3, 4, 5])
+
+
+def test_plot_heatmap(si: Series) -> None:
+    _ = heatmap([si], cmap="vik", vmin=-3.0, vmax=3.0)