Add ability to specify median and mean lines for diurnal plots

Python_Engine/Python/src/python_toolkit/plot/diurnal.py

@@ -2,6 +2,7 @@
 
 import calendar
 import textwrap
+from typing import Tuple
 
 import matplotlib.collections as mcollections
 import matplotlib.lines as mlines
@@ -20,6 +21,9 @@ def diurnal(
     series: pd.Series,
     ax: plt.Axes = None,
     period: str = "daily",
+    quantile_range: Tuple[float, float] = (0.05, 0.95),
+    median: bool = True,
+    mean: bool = True,
     **kwargs,
 ) -> plt.Axes:
     """Plot a profile aggregated across days in the specified timeframe.
@@ -31,6 +35,12 @@ def diurnal(
             A matplotlib Axes object. Defaults to None.
         period (str, optional):
             The period to aggregate over. Must be one of "dailyy", "weekly", or "monthly". Defaults to "daily".
+        quantile_range (Tuple[float, float]):
+            The quantile range to display in a lighter (30% alpha) colour on the plot. Defaults to (0.05, 0.95).
+        median (bool, optional):
+            Whether to plot the median line. Default `True`.
+        mean (bool, optional):
+            Whether to plot the mean line. Default `True`.
         **kwargs (Dict[str, Any], optional):
             Additional keyword arguments to pass to the matplotlib plotting function.
             legend (bool, optional):
@@ -43,7 +53,6 @@ def diurnal(
             A matplotlib Axes object.
     """
 
-
     if not isinstance(series.index, pd.DatetimeIndex):
         raise ValueError("Series passed is not datetime indexed.")
 
@@ -61,7 +70,6 @@ def diurnal(
             raise ValueError("minmax_range must be increasing.")
         minmax_alpha = kwargs.pop("minmax_alpha", 0.1)
 
-        quantile_range = kwargs.pop("quantile_range", [0.05, 0.95])
         if quantile_range[0] > quantile_range[1]:
             raise ValueError("quantile_range must be increasing.")
         if quantile_range[0] < minmax_range[0] or quantile_range[1] > minmax_range[1]:
@@ -183,24 +191,26 @@ def diurnal(
                 label="_nolegend_",
             )
             # mean/median
-            ax.plot(
-                range(len(df) + 1)[i : i + 25],
-                (df["mean"].tolist() + [df["mean"].values[0]])[i : i + 24]
-                + [(df["mean"].tolist() + [df["mean"].values[0]])[i : i + 24][0]],
-                c=color,
-                ls="-",
-                lw=1,
-                label="Average" if n == 0 else "_nolegend_",
-            )
-            ax.plot(
-                range(len(df) + 1)[i : i + 25],
-                (df["median"].tolist() + [df["median"].values[0]])[i : i + 24]
-                + [(df["median"].tolist() + [df["median"].values[0]])[i : i + 24][0]],
-                c=color,
-                ls="--",
-                lw=1,
-                label="Median" if n == 0 else "_nolegend_",
-            )
+            if mean:
+                ax.plot(
+                    range(len(df) + 1)[i : i + 25],
+                    (df["mean"].tolist() + [df["mean"].values[0]])[i : i + 24]
+                    + [(df["mean"].tolist() + [df["mean"].values[0]])[i : i + 24][0]],
+                    c=color,
+                    ls="-",
+                    lw=1,
+                    label="Average" if n == 0 else "_nolegend_",
+                )
+            if median:
+                ax.plot(
+                    range(len(df) + 1)[i : i + 25],
+                    (df["median"].tolist() + [df["median"].values[0]])[i : i + 24]
+                    + [(df["median"].tolist() + [df["median"].values[0]])[i : i + 24][0]],
+                    c=color,
+                    ls="--",
+                    lw=1,
+                    label="Median" if n == 0 else "_nolegend_",
+                )
 
         # format axes
         ax.set_xlim(0, len(df))