Update hydrograph - add xarray support, and improve plotting/metrics

.gitignore

@@ -32,3 +32,4 @@ __pycache__
 
 .ruff-cache
 .venv
+.vscode/settings.json

src/ewatercycle/analysis/__init__.py

@@ -1,159 +1,5 @@
-"""Analysis methods for eWaterCycle."""
+"""ewatercycle analysis module."""
 
-import os
+from . import hydrograph
 
-import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
-from hydrostats import metrics
-from matplotlib.axes import Axes
-from matplotlib.dates import DateFormatter
-from matplotlib.figure import Figure
-
-
-def _downsample(df, nrows=100, agg="mean"):
-    """Resample dataframe with datetimeindex to a fixed number of rows."""
-    if len(df) <= nrows:
-        return df, df.index[1] - df.index[0]
-
-    grouper = np.arange(len(df)) // (len(df) / nrows)
-    new_df = df.groupby(grouper).agg(agg)
-    new_df.index = pd.date_range(df.index[0], df.index[-1], periods=nrows)
-
-    new_period = (df.index[-1] - df.index[0]) / nrows
-
-    return new_df, new_period
-
-
-def hydrograph(
-    discharge: pd.DataFrame,
-    *,
-    reference: str,
-    precipitation: pd.DataFrame | None = None,
-    dpi: int | None = None,
-    title: str = "Hydrograph",
-    discharge_units: str = "m$^3$ s$^{-1}$",
-    precipitation_units: str = "mm day$^{-1}$",
-    figsize: tuple[float, float] = (10, 10),
-    filename: os.PathLike | str | None = None,
-    nbars: int | None = None,
-    **kwargs,
-) -> tuple[Figure, tuple[Axes, Axes]]:
-    """Plot a hydrograph.
-
-    This utility function makes it convenient to create a hydrograph from
-    a set of discharge data from a `pandas.DataFrame`. A column must be marked
-    as the reference, so that the agreement metrics can be calculated.
-
-    Optionally, the corresponding precipitation data can be plotted for
-    comparison.
-
-    Args:
-        discharge: Dataframe containing time series of discharge data to be plotted.
-        reference: Name of the reference data, must correspond to a column
-            in the discharge dataframe. Metrics are calculated between
-            the reference column and each of the other columns.
-        precipitation: Optional dataframe containing time series of precipitation data
-            to be plotted from the top of the hydrograph.
-        dpi: DPI for the plot.
-        title: Title of the hydrograph.
-        discharge_units: Units for the discharge data.
-        precipitation_units: Units for the precipitation data.
-        figsize: With, height of the plot in inches.
-        filename: If specified, a copy of the plot will be saved to this path.
-        nbars: Number of bars to use for downsampling precipitation.
-        **kwargs: Options to pass to the matplotlib plotting function
-
-    Returns:
-        First tuple member is a matplotlib figure, the second is a tuple of axes.
-    """
-    discharge_cols = discharge.columns.drop(reference)
-    y_obs = discharge[reference]
-    y_sim = discharge[discharge_cols]
-
-    fig, (ax, ax_tbl) = plt.subplots(
-        nrows=2,
-        ncols=1,
-        dpi=dpi,
-        figsize=figsize,
-        gridspec_kw={"height_ratios": [3, 1]},
-    )
-
-    ax.set_title(title)
-    ax.set_ylabel(f"Discharge ({discharge_units})")
-
-    y_sim.plot(ax=ax, **kwargs)
-    y_obs.plot(ax=ax, **kwargs)
-
-    handles, labels = ax.get_legend_handles_labels()
-
-    # Add precipitation as bar plot to the top if specified
-    if precipitation is not None:
-        if nbars is not None:
-            precipitation, barwidth = _downsample(
-                precipitation, nrows=nbars, agg="mean"
-            )
-        else:
-            barwidth = 0.8  # default value for matplotlib barplot
-
-        ax_pr = ax.twinx()
-        ax_pr.invert_yaxis()
-        ax_pr.set_ylabel(f"Precipitation ({precipitation_units})")
-
-        for pr_label, pr_timeseries in precipitation.items():
-            ax_pr.bar(
-                pr_timeseries.index.values,
-                pr_timeseries.values,
-                width=barwidth,
-                alpha=0.4,
-                label=pr_label,
-            )
-
-        # tweak ylim to make space at bottom and top
-        ax_pr.set_ylim(ax_pr.get_ylim()[0] * (7 / 2), 0)
-        ax.set_ylim(0, ax.get_ylim()[1] * (7 / 5))
-
-        # prepend handles/labels so they appear at the top
-        handles_pr, labels_pr = ax_pr.get_legend_handles_labels()
-        handles = handles_pr + handles
-        labels = labels_pr + labels
-
-    # Put the legend outside the plot
-    ax.legend(handles, labels, bbox_to_anchor=(1.10, 1), loc="upper left")
-
-    # set formatting for xticks
-    date_fmt = DateFormatter("%Y-%m")
-    ax.xaxis.set_major_formatter(date_fmt)
-    ax.tick_params(axis="x", rotation=30)
-
-    # calculate metrics for data table underneath plot
-    def calc_metric(metric) -> float:
-        return y_sim.apply(metric, observed_array=y_obs)
-
-    metrs = pd.DataFrame(
-        {
-            "nse": calc_metric(metrics.nse),
-            "kge_2009": calc_metric(metrics.kge_2009),
-            "sa": calc_metric(metrics.sa),
-            "me": calc_metric(metrics.me),
-        }
-    )
-
-    # convert data in dataframe to strings
-    cell_text = [[f"{item:.2f}" for item in row[1]] for row in metrs.iterrows()]
-
-    table = ax_tbl.table(
-        cellText=cell_text,
-        rowLabels=metrs.index,
-        colLabels=metrs.columns,
-        loc="center",
-    )
-    ax_tbl.set_axis_off()
-
-    # give more vertical space in cells
-    table.scale(1, 1.5)
-
-    if filename is not None:
-        fig.savefig(filename, bbox_inches="tight", dpi=dpi)
-
-    return fig, (ax, ax_tbl)
+__all__ = ["hydrograph"]