diff --git a/.gitignore b/.gitignore index 80012883..2c1c8d72 100644 --- a/.gitignore +++ b/.gitignore @@ -32,3 +32,4 @@ __pycache__ .ruff-cache .venv +.vscode/settings.json diff --git a/README.md b/README.md index 075f39d1..d1e6b79e 100644 --- a/README.md +++ b/README.md @@ -74,13 +74,13 @@ and produce a hydrograph. In condensed code: ```python -forcing = ewatercycle.forcing.sources['MarrmotForcing'].generate(...) -model = ewatercycle.models.sources['MarrmotM14'](forcing) +forcing = ewatercycle.forcing.sources["MarrmotForcing"].generate(...) +model = ewatercycle.models.sources["MarrmotM14"](forcing) model.setup(...) model.initialize() -while (model.time < model.end_time): +while model.time < model.end_time: model.update() - value = model.get_value_as_xarray('flux_out_Q') + value = model.get_value_as_xarray("flux_out_Q") model.finalize() ewatercycle.analysis.hydrograph(...) ``` @@ -98,14 +98,14 @@ from ewatercycle.testing.fixtures import rhine_shape import shapefile import xarray as xr -forcing = ewatercycle.forcing.sources['MarrmotForcing'].generate( - dataset='ERA5', - start_time='2010-01-01T00:00:00Z', - end_time='2010-12-31T00:00:00Z', - shape=rhine_shape() +forcing = ewatercycle.forcing.sources["MarrmotForcing"].generate( + dataset="ERA5", + start_time="2010-01-01T00:00:00Z", + end_time="2010-12-31T00:00:00Z", + shape=rhine_shape(), ) -model = ewatercycle.models.sources['MarrmotM14'](version='2020.11', forcing=forcing) +model = ewatercycle.models.sources["MarrmotM14"](version="2020.11", forcing=forcing) cfg_file, cfg_dir = model.setup( threshold_flow_generation_evap_change=0.1, @@ -115,31 +115,29 @@ model.initialize(cfg_file) # flux_out_Q unit conversion factor from mm/day to m3/s sf = shapefile.Reader(rhine_shape()) -area = sf.record(0)['SUB_AREA'] * 1e6 # from shapefile in m2 +area = sf.record(0)["SUB_AREA"] * 1e6 # from shapefile in m2 conversion_mmday2m3s = 1 / (1000 * 24 * 60 * 60) conversion = conversion_mmday2m3s * area simulated_discharge = [] -while (model.time < model.end_time): +while model.time < model.end_time: model.update() - simulated_discharge.append( - model.get_value_as_xarray('flux_out_Q') - ) + simulated_discharge.append(model.get_value_as_xarray("flux_out_Q")) observations_ds = ewatercycle.observation.grdc.get_grdc_data( station_id=6335020, # Rees, Germany start_time=model.start_time_as_isostr, end_time=model.end_time_as_isostr, - column='observation', + column="observation", ) # Combine the simulated discharge with the observations -sim_da = xr.concat(simulated_discharge, dim='time') * conversion -sim_da.name = 'simulated' +sim_da = xr.concat(simulated_discharge, dim="time") * conversion +sim_da.name = "simulated" discharge = xr.merge([sim_da, observations_ds["observation"]]).to_dataframe() discharge = discharge[["observation", "simulated"]].dropna() -ewatercycle.analysis.hydrograph(discharge, reference='observation') +ewatercycle.analysis.hydrograph(discharge, reference="observation") model.finalize() ``` diff --git a/pyproject.toml b/pyproject.toml index 433eef23..af748409 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -53,6 +53,7 @@ dependencies = [ "xarray", "fsspec", "cartopy", + "HydroErr==2.0.0", # Pin pyopenssl (conda installs old version) https://github.com/conda/conda/issues/13619 "pyopenssl>=24.0.0", ] diff --git a/src/ewatercycle/_forcings/caravan.py b/src/ewatercycle/_forcings/caravan.py index 3933ee3c..d040e436 100644 --- a/src/ewatercycle/_forcings/caravan.py +++ b/src/ewatercycle/_forcings/caravan.py @@ -199,7 +199,7 @@ def generate( # type: ignore[override] raise ValueError(msg) basin_id = str(kwargs["basin_id"]) - dataset: str = basin_id.split("_")[0] + dataset: str = basin_id.split("_", maxsplit=1)[0] ds = cls.get_dataset(dataset) ds_basin = ds.sel(basin_id=basin_id.encode()) ds_basin_time = crop_ds(ds_basin, start_time, end_time) diff --git a/src/ewatercycle/analysis/__init__.py b/src/ewatercycle/analysis/__init__.py index 473cafa8..2d859d86 100644 --- a/src/ewatercycle/analysis/__init__.py +++ b/src/ewatercycle/analysis/__init__.py @@ -1,159 +1,5 @@ -"""Analysis methods for eWaterCycle.""" +"""ewatercycle analysis module.""" -import os +from ewatercycle.analysis 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"] diff --git a/src/ewatercycle/analysis/hydrograph.py b/src/ewatercycle/analysis/hydrograph.py new file mode 100644 index 00000000..ea804c42 --- /dev/null +++ b/src/ewatercycle/analysis/hydrograph.py @@ -0,0 +1,303 @@ +"""Analysis methods for eWaterCycle.""" + +import os +from typing import Any, Protocol, cast + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import xarray as xr +from HydroErr.HydroErr import function_list +from hydrostats import metrics # noqa: F401 +from matplotlib.axes import Axes +from matplotlib.dates import AutoDateLocator, DateFormatter +from matplotlib.figure import Figure + + +class Metric(Protocol): + """A HydroErr metric function. + + HydroErr attaches the full name and the abbreviation of each metric to its + function object at import time, so they are invisible to type checkers. + """ + + __name__: str + name: str + abbr: str + + def __call__( + self, simulated_array: Any, observed_array: Any, **kwargs: Any + ) -> float: + """Compute the metric for a simulated and an observed timeseries.""" + ... + + +# metrics from https://hydroerr.readthedocs.io/en/stable/list_of_metrics.html callable + +metric_functions = cast("list[Metric]", function_list) + +metric_map = {func.name: func for func in metric_functions} # full names +metric_map.update({func.abbr: func for func in metric_functions}) # abbreviations +metric_map.update({func.__name__: func for func in metric_functions}) # function names +metric_map.update({k.lower(): v for k, v in metric_map.items()}) # lowercase + + +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 _to_pandas(data_in): + """Convert input to pandas DataFrame if it is not already. + + Args: + data_in : supported data types: pd.Dataframe, xr.Dataset + + Returns: + pd.Dataframe + + Raises: + Typerror if the input is not supported + """ + # already a DataFrame + if isinstance(data_in, pd.DataFrame): + return data_in + + # Single series + if isinstance(data_in, pd.Series): + msg = "A panda series contains only a single timeseries, please provide a pandas DataFrame or xr.Dataset." # noqa: E501 + raise TypeError(msg) + + # xarray Dataset + if isinstance(data_in, xr.Dataset): + return data_in.to_pandas() + + # xarray DataArray + if isinstance(data_in, xr.DataArray): + if data_in.ndim == 1: + msg = "A DataArray with a single timeseries is not supported, please provide a DataFrame or xr.Dataset." # noqa: E501 + raise TypeError(msg) + else: # noqa: RET506 + msg = "DataArray with more than one dimension is not supported, please provide a DataFrame or xr.Dataset." # noqa: E501 + raise TypeError(msg) + + # unsupported type + msg = f"Unsupported data type: {type(data_in)}, please provide a DataFrame or xr.Dataset" # noqa: E501 + raise TypeError(msg) + + +def _prepare_discharge(discharge, reference: str, selected_year=None): + """Prepare discharge data for hydrograph.""" + discharge = _to_pandas(discharge) + + # Slice by selected_year if provided + if selected_year is not None: + if not isinstance(discharge.index, pd.DatetimeIndex): + msg = "Discharge index must be a DatetimeIndex to select a year." + raise ValueError(msg) + discharge = discharge[discharge.index.year == selected_year] + y_obs = discharge[reference] + y_sim = discharge.drop(columns=[reference]) + return y_obs, y_sim + + +def _prepare_precipitation(precipitation, nbars=None, selected_year=None): + if precipitation is None: + return None, None + precipitation = _to_pandas(precipitation) + if nbars is not None: + precipitation, barwidth = _downsample(precipitation, nrows=nbars, agg="mean") + else: + barwidth = 0.8 # default value for matplotlib barplot + + if selected_year is not None: + if not isinstance(precipitation.index, pd.DatetimeIndex): + msg = "Precipitation index must be a DatetimeIndex to select a year." + raise ValueError(msg) + precipitation = precipitation[precipitation.index.year == selected_year] + + return precipitation, barwidth + + +def _plot_discharge(ax, y_obs, y_sim, **kwargs): + if hasattr(y_sim, "shape") and y_sim.shape[1] > 1: + y_obs.plot(ax=ax, linewidth=2.5, zorder=10, **kwargs) + y_sim.plot(ax=ax, alpha=0.7, linewidth=1.25, **kwargs) + else: + y_obs.plot(ax=ax, **kwargs, zorder=10) + y_sim.plot(ax=ax, **kwargs) + ax.grid(True) + return ax + + +def _plot_precipitation(ax, precipitation, barwidth, precipitation_units): + 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, + ) + + # adjust ylim + ax_pr.set_ylim(ax_pr.get_ylim()[0] * (7 / 2), 0) + ax.set_ylim(0, ax.get_ylim()[1] * (7 / 5)) + return ax_pr + + +def _calculate_metrics(y_obs, y_sim, metrics_list=None): + if metrics_list is None: + metrics_list = ["NSE", "KGE (2009)", "SA", "ME"] + + metrics_objs = [] + for m in metrics_list: + if isinstance(m, str): + if m in metric_map: + metrics_objs.append(metric_map[m]) + elif m.lower() in metric_map: + metrics_objs.append(metric_map[m.lower()]) + else: + msg = f"Metric '{m}' not found in hydroerr metrics." + raise ValueError(msg) + else: + metrics_objs.append(m) + + def calc_metric(metric) -> float: + return y_sim.apply(metric, observed_array=y_obs) + + df_metrics = pd.DataFrame( + {metric.name: calc_metric(metric) for metric in metrics_objs} + ) + return df_metrics, metrics_objs + + +def _create_metrics_table(ax_tbl, df_metrics, metrics_objs): + col_labels = [f"{metric.name}\n({metric.abbr})" for metric in metrics_objs] + metrs_rounded = df_metrics.round(2) + cell_text = [[f"{item:.2f}" for item in row[1]] for row in metrs_rounded.iterrows()] + table = ax_tbl.table( + cellText=cell_text, + rowLabels=metrs_rounded.index, + colLabels=col_labels, + loc="center", + fontsize=15, + ) + ax_tbl.set_axis_off() + for (i, j), cell in table.get_celld().items(): + if i == 0: # headers + cell.set_fontsize(12) + cell.set_text_props(weight="bold") + cell.set_height(0.15) + elif j == -1: # row labels + cell.set_fontsize(11) + cell.set_text_props(weight="bold") + else: # table values + cell.set_fontsize(10) + table.scale(1, 1.5) + return table + + +def hydrograph( + discharge: pd.DataFrame | pd.Series | xr.DataArray | xr.Dataset, + *, + reference: str, + precipitation: pd.DataFrame | pd.Series | xr.DataArray | xr.Dataset | None = None, + dpi: int | None = None, + title: str | None = None, + 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, + metrics_list: list[object] | None = None, + selected_year: 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` or `xarray.Dataset`. 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: Data 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. + metrics_list: List of metrics to calculate and display in the table below + the hydrograph. If not specified, a default set of metrics is used. + selected_year: Slices a single year from the data + **kwargs: Options to pass to the matplotlib plotting function + + Returns: + First tuple member is a matplotlib figure, the second is a tuple of axes. + """ + y_obs, y_sim = _prepare_discharge(discharge, reference, selected_year) + precipitation, barwidth = _prepare_precipitation( + precipitation, nbars, selected_year + ) + + fig, (ax, ax_tbl) = plt.subplots( + nrows=2, + ncols=1, + dpi=dpi, + figsize=figsize, + gridspec_kw={"height_ratios": [3, 1]}, + ) + fig.subplots_adjust(bottom=0.05, top=0.95, hspace=0.3) + + ax.set_title( + title or f"$\\mathbf{{Hydrograph\\ with\\ Metrics}}$\nReference: {reference}" + ) + + ax.set_ylabel(f"Discharge ({discharge_units})") + + _plot_discharge(ax, y_obs, y_sim, **kwargs) + + if precipitation is not None: + ax_pr = _plot_precipitation(ax, precipitation, barwidth, precipitation_units) + handles, labels = ax_pr.get_legend_handles_labels() + else: + handles, labels = ax.get_legend_handles_labels() + + ax.legend(handles, labels, bbox_to_anchor=(1.10, 1), loc="upper left") + + locator = AutoDateLocator(minticks=5, maxticks=12) # adjust min/max ticks + ax.xaxis.set_major_locator(locator) + ax.xaxis.set_major_formatter(DateFormatter("%Y-%m")) + ax.tick_params(axis="x", rotation=30) + ax.set_xlabel(None) # no xlabel, ticks make it clear + + df_metrics, metrics_objs = _calculate_metrics(y_obs, y_sim, metrics_list) + _create_metrics_table(ax_tbl, df_metrics, metrics_objs) + + if filename: + fig.savefig(filename, bbox_inches="tight", dpi=dpi) + + return fig, (ax, ax_tbl) diff --git a/src/ewatercycle/base/parameter_set.py b/src/ewatercycle/base/parameter_set.py index 41e87ee9..902470c3 100644 --- a/src/ewatercycle/base/parameter_set.py +++ b/src/ewatercycle/base/parameter_set.py @@ -40,7 +40,7 @@ def download_github_repo( """ zip_url = f"https://github.com/{org}/{repo}/archive/refs/heads/{branch}.zip" - https_response = urlopen(zip_url) # noqa: S310 + https_response = urlopen(zip_url) if https_response.status != 200: msg = ( f"HTTP error {https_response.status}\n" diff --git a/src/ewatercycle/esmvaltool/builder.py b/src/ewatercycle/esmvaltool/builder.py index 19942c0a..372fc7e1 100644 --- a/src/ewatercycle/esmvaltool/builder.py +++ b/src/ewatercycle/esmvaltool/builder.py @@ -290,8 +290,8 @@ def add_variable( units: str | None = None, stats: ClimateStatistics | None = None, short_name: str | None = None, - start_year: int | None | Literal[False] = None, - end_year: int | None | Literal[False] = None, + start_year: int | Literal[False] | None = None, + end_year: int | Literal[False] | None = None, ): """Add a variable to the recipe. diff --git a/src/ewatercycle/esmvaltool/datasets.py b/src/ewatercycle/esmvaltool/datasets.py index 2fb380f4..0d89dc3f 100644 --- a/src/ewatercycle/esmvaltool/datasets.py +++ b/src/ewatercycle/esmvaltool/datasets.py @@ -17,6 +17,13 @@ type="reanaly", version=1, ), + "ERA5-Land": Dataset( + dataset="ERA5-Land", + project="OBS6", + tier=3, + type="reanaly", + version=1, + ), } """Dictionary of predefined forcing datasets. @@ -28,7 +35,7 @@ >> from ewatercycle.forcing import DATASETS >> list(DATASETS.keys()) - ['ERA5', 'ERA-Interim'] + ['ERA5', 'ERA-Interim', 'ERA5-Land'] """ diff --git a/src/ewatercycle/observation/caravan.py b/src/ewatercycle/observation/caravan.py index ed432f03..4697dacc 100644 --- a/src/ewatercycle/observation/caravan.py +++ b/src/ewatercycle/observation/caravan.py @@ -60,7 +60,7 @@ def get_caravan_data( NCO: netCDF Operators version 5.0.6 (Homepage = http://nco.sf.... _NCProperties: version=2,netcdf=4.8.1,hdf5=1.10.7 """ # noqa: D214,D410,D411 - dataset: str = basin_id.split("_")[0] + dataset: str = basin_id.split("_", maxsplit=1)[0] ds = CaravanForcing.get_dataset(dataset) ds_basin = ds.sel(basin_id=basin_id.encode()) ds_basin_time = crop_ds(ds_basin, start_time, end_time) diff --git a/src/ewatercycle/observation/grdc.py b/src/ewatercycle/observation/grdc.py index 832055f1..78ca50b0 100644 --- a/src/ewatercycle/observation/grdc.py +++ b/src/ewatercycle/observation/grdc.py @@ -90,7 +90,7 @@ def get_grdc_data( institution: GRDC history: Download from GRDC Database, 21/06/2024 missing_value: -999.000 - """ # noqa: D214,D410,D411 + """ if data_home: data_path = to_absolute_path(data_home) elif CFG.grdc_location: @@ -352,3 +352,248 @@ def _extract_metadata(lines, key, cast=str, default="NA"): return default warnings.warn(f"{key} not found, set to {default}", stacklevel=2) return default + + +def get_grdc_data_monthly( + station_id: str, + start_time: str, + end_time: str, + data_home: str | None = None, + column1: str = "original streamflow", + column2: str = "calculated streamflow", + column3: str = "flag", +) -> xr.Dataset: + """Load monthly GRDC discharge data for a specific station. + + This function is similar to `get_grdc_data` but specifically for **monthly data**. + It returns a Dataset with three time series columns: + - column1: original streamflow (default: "original streamflow") + - column2: calculated streamflow (default: "calculated streamflow") + - column3: flag (default: "flag") + + Currently, this function only reads GRDC `.txt` files. + NetCDF support (.nc) is not implemented. + + Args: + station_id : str + GRDC station identifier. + start_time : str + Start of the period to extract, in ISO format (e.g., "YYYY-MM-DDTHH:MMZ"). + end_time : str + End of the period to extract, in ISO format (e.g., "YYYY-MM-DDTHH:MMZ"). + data_home : str | None, optional + Path to the directory containing the GRDC files. Defaults to None. + If None, falls back to `CFG.grdc_location`. + column1 : str, optional + Name for the first data column (original streamflow). + column2 : str, optional + Name for the second data column (calculated streamflow). + column3 : str, optional + Name for the third data column (flag). + + Returns: + grdc data in a xarray dataset. + Shaped like a filtered version of the GRDC daily NetCDF file. + Data description: + Original - original (provided) data + Calculated - GRDC calculated from daily data + Flag - percentage of valid values used for calculation from daily data + + Raises: + ValueError: If no data for the requested station id + and period could not be found. + NotImplementedError + If a NetCDF (.nc) file is found (support not implemented). + """ + if data_home: + data_path = to_absolute_path(data_home) + elif CFG.grdc_location: + data_path = to_absolute_path(CFG.grdc_location) + else: + msg = ( + "Provide the grdc path using `data_home` argument" + "or using `grdc_location` in ewatercycle configuration file." + ) + raise ValueError(msg) + + if not data_path.exists(): + msg = f"The grdc directory {data_path} does not exist!" + raise ValueError(msg) + + # # Read the NetCDF file + nc_file = data_path / "GRDC-Monthly.nc" + if nc_file.exists(): + msg = ".nc support not implemented at this point" + raise NotImplementedError(msg) + + # Read the text data + raw_file = data_path / f"{station_id}_Q_Month.txt" + if not raw_file.exists(): + if nc_file.exists(): + msg = f"The grdc station {station_id} is not in the {nc_file} file and {raw_file} does not exist!" # noqa: E501 + raise ValueError(msg) + msg = f"The grdc file {raw_file} does not exist!" + raise ValueError(msg) + + # Convert the raw data to an dataframe + metadata, df = _grdc_read_monthly( + raw_file, + start=get_time(start_time).date(), + end=get_time(end_time).date(), + column1=column1, + column2=column2, + column3=column3, + ) + + return xr.Dataset.from_dict( + { + "coords": { + "time": { + "dims": ("time",), + "attrs": {"long_name": "time"}, + "data": df.index.to_numpy(), + }, + "id": { + "dims": (), + "attrs": {"long_name": "grdc number"}, + "data": int(station_id), + }, + }, + "dims": { + "time": len(df.index), + }, + "attrs": { + "title": metadata["dataSetContent"], + "Conventions": "CF-1.7", + "references": "grdc.bafg.de", + "institution": "GRDC", + "history": f"Converted from {raw_file.name} of {metadata['file_generation_date']} to netcdf by eWaterCycle Python package", # noqa: E501 + "missing_value": "-999.000", + }, + "data_vars": { + "Original discharge": { + "dims": ("time",), + "attrs": { + "units": "m3/s", + "long_name": "Mean monthly discharge (MQ)", + }, + "data": df[column1].to_numpy(), + }, + "Calculated discharge": { + "dims": ("time",), + "attrs": { + "units": "m3/s", + "long_name": "Mean monthly discharge (MQ)", + }, + "data": df[column2].to_numpy(), + }, + "Flag": { + "dims": ("time",), + "attrs": { + "units": "%", + "long_name": "percentage of valid daily values used for calculation", # noqa: E501 + }, + "data": df[column3].to_numpy(), + }, + "area": { + "dims": (), + "attrs": {"units": "km2", "long_name": "catchment area"}, + "data": metadata["grdc_catchment_area_in_km2"], + }, + "country": { + "dims": (), + "attrs": { + "long_name": "country name", + "iso2": "ISO 3166-1 alpha-2 - two-letter country code", + }, + "data": metadata["country_code"], + }, + "geo_x": { + "dims": (), + "attrs": { + "units": "degree_east", + "long_name": "station longitude (WGS84)", + }, + "data": metadata["grdc_longitude_in_arc_degree"], + }, + "geo_y": { + "dims": (), + "attrs": { + "units": "degree_north", + "long_name": "station latitude (WGS84)", + }, + "data": metadata["grdc_latitude_in_arc_degree"], + }, + "geo_z": { + "dims": (), + "attrs": { + "units": "m", + "long_name": "station altitude (m above sea level)", + }, + "data": metadata["altitude_masl"], + }, + "owneroforiginaldata": { + "dims": (), + "attrs": {"long_name": "Owner of original data"}, + "data": metadata["Owner of original data"], + }, + "river_name": { + "dims": (), + "attrs": {"long_name": "river name"}, + "data": metadata["river_name"], + }, + "station_name": { + "dims": (), + "attrs": {"long_name": "station name"}, + "data": metadata["station_name"], + }, + "timezone": { + "dims": (), + "attrs": { + "units": "00:00", + "long_name": "utc offset, in relation to the national capital", + }, + "data": nan, + }, + }, + } + ) + + +def _grdc_read_monthly(grdc_station_path, start, end, column1, column2, column3): + """Private helper function for reading monthly grdc data.""" + with grdc_station_path.open("r", encoding="cp1252", errors="ignore") as file: + data = file.read() + + metadata = _grdc_metadata_reader(grdc_station_path, data) + + all_lines = data.split("\n") + header = 0 + for i, line in enumerate(all_lines): + if line.startswith("# DATA"): + header = i + 1 + break + + # Import GRDC data into dataframe and modify dataframe format + grdc_data = pd.read_csv( + grdc_station_path, + encoding="cp1252", + skiprows=header, + delimiter=";", + parse_dates=["YYYY-MM-DD"], + na_values="-999", + ) + grdc_station_df = pd.DataFrame( + { + column1: grdc_data[" Original"].array, + column2: grdc_data[" Calculated"].array, + column3: grdc_data[" Flag"].array, + }, + index=grdc_data["YYYY-MM-DD"].array, + ) + grdc_station_df.index.rename("time", inplace=True) # noqa: PD002 + + # Select GRDC station data that matches the forecast results Date + grdc_station_select = grdc_station_df.loc[start:end] + + return metadata, grdc_station_select diff --git a/src/ewatercycle/util.py b/src/ewatercycle/util.py index 236e8e6e..199d887d 100644 --- a/src/ewatercycle/util.py +++ b/src/ewatercycle/util.py @@ -339,7 +339,7 @@ def extract_package_name(value: str) -> str: E.g. "ewatercycle_HBV.model:HBV" will return "ewatercycle_HBV". """ - source = value.split(":")[0] + source = value.split(":", maxsplit=1)[0] return source.split(".")[0] diff --git a/tests/src/base/test_forcing.py b/tests/src/base/test_forcing.py index dc1ab228..240136a3 100644 --- a/tests/src/base/test_forcing.py +++ b/tests/src/base/test_forcing.py @@ -369,7 +369,7 @@ def test_retrieve_caravan_forcing(tmp_path: Path, mock_retrieve: mock.MagicMock) content = list(ds.data_vars.keys()) expected = ["Q", "evspsblpot", "pr", "tas", "tasmax", "tasmin"] assert content == expected - mock_retrieve.assert_called_once_with(basin_id.split("_")[0]) + mock_retrieve.assert_called_once_with(basin_id.split("_", maxsplit=1)[0]) assert caravan_forcing.to_xarray()["evspsblpot"].attrs["unit"] == "kg m-2 s-1" assert caravan_forcing.to_xarray()["pr"].attrs["unit"] == "kg m-2 s-1" @@ -394,7 +394,7 @@ def test_retrieve_caravan_forcing_empty_vars( content = list(ds.data_vars.keys()) expected = ["Q", "evspsblpot", "pr", "tas", "tasmax", "tasmin"] assert content == expected - mock_retrieve.assert_called_once_with(basin_id.split("_")[0]) + mock_retrieve.assert_called_once_with(basin_id.split("_", maxsplit=1)[0]) def test_retrieve_caravan_forcing_no_basin_id( diff --git a/tests/src/baseline_images/test_analysis/hydrograph_DataFrame.png b/tests/src/baseline_images/test_analysis/hydrograph_DataFrame.png new file mode 100644 index 00000000..72c57a0b Binary files /dev/null and b/tests/src/baseline_images/test_analysis/hydrograph_DataFrame.png differ diff --git a/tests/src/baseline_images/test_analysis/hydrograph_xarray.png b/tests/src/baseline_images/test_analysis/hydrograph_xarray.png new file mode 100644 index 00000000..83744978 Binary files /dev/null and b/tests/src/baseline_images/test_analysis/hydrograph_xarray.png differ diff --git a/tests/src/baseline_images/test_analysis/hydrograph_xarray_single_comparison.png b/tests/src/baseline_images/test_analysis/hydrograph_xarray_single_comparison.png new file mode 100644 index 00000000..657cca67 Binary files /dev/null and b/tests/src/baseline_images/test_analysis/hydrograph_xarray_single_comparison.png differ diff --git a/tests/src/baseline_images/test_analysis/hydrograph_xarray_single_year.png b/tests/src/baseline_images/test_analysis/hydrograph_xarray_single_year.png new file mode 100644 index 00000000..c9d88438 Binary files /dev/null and b/tests/src/baseline_images/test_analysis/hydrograph_xarray_single_year.png differ diff --git a/tests/src/observation/test_grdc.py b/tests/src/observation/test_grdc.py index cf6729fd..10720b99 100644 --- a/tests/src/observation/test_grdc.py +++ b/tests/src/observation/test_grdc.py @@ -7,7 +7,7 @@ from xarray.testing import assert_allclose from ewatercycle import CFG -from ewatercycle.observation.grdc import get_grdc_data +from ewatercycle.observation.grdc import get_grdc_data, get_grdc_data_monthly @pytest.fixture() @@ -321,3 +321,292 @@ def test_get_grdc_data_from_nc_missing_and_no_txtfile(tmp_path, sample_nc_file): "2000-02-01T00:00Z", data_home=str(tmp_path), ) + + +@pytest.fixture() +def sample_grdc_monthly_file(tmp_path): + fn = tmp_path / "30303030_Q_Month.txt" + # Sample with fictive data, but with same structure as real file + body = """# Title: GRDC STATION DATA FILE +# -------------- +# Format: DOS-ASCII +# Field delimiter: ; +# missing values are indicated by -999.000 +# +# file generation date: 2000-02-02 +# +# GRDC-No.: 30303030 +# River: SOME RIVER +# Station: SOME +# Country: NA +# Latitude (DD): 52.356154 +# Longitude (DD): 4.955153 +# Catchment area (km²): 3030.0 +# Altitude (m ASL): 8.0 +# Next downstream station: 42424243 +# Remarks: +# Owner of original data: SOME PROJECT +#************************************************************ +# +# Data Set Content: MEAN MONTHLY DISCHARGE (MQ) +# -------------------- +# Unit of measure: m³/s +# Time series: 2000 - 2000 +# No. of years: 1 +# Last update: 2000-02-01 +# +# Table Header: +# YYYY-MM-DD - Date (DD=00) +# hh:mm - Time +# Original - original (provided) data +# Calculated - GRDC calculated from daily data +# Flag - percentage of valid values used for calculation from daily data +#************************************************************ +# +# Data lines: 3 +# DATA +YYYY-MM-DD;hh:mm; Original; Calculated; Flag +2000-01-01;--:--; 3.000; -999.000; 0 +2000-02-01;--:--; 5.000; -999.000; 0 +2000-03-01;--:--; 8.000; -999.000; 0""" + with fn.open("w", encoding="cp1252") as f: + f.write(body) + return fn + + +@pytest.fixture() +def expected_results_monthly(): + return xr.Dataset.from_dict( + { + "coords": { + "time": { + "dims": ("time",), + "attrs": {"long_name": "time"}, + "data": [ + datetime(2000, 1, 1), + datetime(2000, 2, 1), + datetime(2000, 3, 1), + ], + }, + "id": { + "dims": (), + "attrs": {"long_name": "grdc number"}, + "data": 30303030, + }, + }, + "dims": {"time": 3}, + "attrs": { + "title": "MEAN MONTHLY DISCHARGE (MQ)", + "Conventions": "CF-1.7", + "references": "grdc.bafg.de", + "institution": "GRDC", + "history": ( + "Converted from 30303030_Q_Month.txt " + "of 2000-02-02 to netcdf by eWaterCycle Python package" + ), + "missing_value": "-999.000", + }, + "data_vars": { + "Original discharge": { + "dims": ("time",), + "attrs": { + "units": "m3/s", + "long_name": "Mean monthly discharge (Q)", + }, + "data": [3.0, 5.0, 8.0], + }, + "Calculated discharge": { + "dims": ("time",), + "attrs": { + "units": "m3/s", + "long_name": "Mean monthly discharge (Q)", + }, + "data": [np.nan, np.nan, np.nan], + }, + "Flag": { + "dims": ("time",), + "attrs": { + "units": "%", + "long_name": "percentage of valid daily values used for calculation", + }, + "data": [0, 0, 0], + }, + "area": { + "dims": (), + "attrs": { + "units": "km2", + "long_name": "catchment area", + }, + "data": 3030.0, + }, + "country": { + "dims": (), + "attrs": { + "long_name": "country name", + "iso2": ("ISO 3166-1 alpha-2 - two-letter country code"), + }, + "data": "NA", + }, + "geo_x": { + "dims": (), + "attrs": { + "units": "degree_east", + "long_name": "station longitude (WGS84)", + }, + "data": 4.955153, + }, + "geo_y": { + "dims": (), + "attrs": { + "units": "degree_north", + "long_name": "station latitude (WGS84)", + }, + "data": 52.356154, + }, + "geo_z": { + "dims": (), + "attrs": { + "units": "m", + "long_name": "station altitude (m above sea level)", + }, + "data": 8.0, + }, + "owneroforiginaldata": { + "dims": (), + "attrs": {"long_name": "Owner of original data"}, + "data": "SOME PROJECT", + }, + "river_name": { + "dims": (), + "attrs": {"long_name": "river name"}, + "data": "SOME RIVER", + }, + "station_name": { + "dims": (), + "attrs": {"long_name": "station name"}, + "data": "SOME", + }, + "timezone": { + "dims": (), + "attrs": { + "units": "00:00", + "long_name": ( + "utc offset, in relation to the national capital" + ), + }, + "data": np.nan, + }, + }, + } + ) + + +def test_get_grdc_monthly_data_with_datahome( + tmp_path, expected_results_monthly: xr.Dataset, sample_grdc_monthly_file +): + result_data_monthly = get_grdc_data_monthly( + "30303030", + "2000-01-01T00:00Z", + "2000-03-01T00:00Z", + data_home=str(tmp_path), + ) + + assert_allclose(result_data_monthly, expected_results_monthly) + + +def test_get_grdc_monthly_data_with_cfg( + tmp_path, expected_results_monthly: xr.Dataset, sample_grdc_monthly_file +): + CFG.grdc_location = tmp_path + + result_data_monthly = get_grdc_data_monthly( + "30303030", + "2000-01-01T00:00Z", + "2000-03-01T00:00Z", + ) + + assert_allclose(result_data_monthly, expected_results_monthly) + + +def test_get_grdc_monthly_data_without_datahome_and_cfg(monkeypatch): + class NoLocationConfig: + grdc_location = None + + monkeypatch.setattr("ewatercycle.observation.grdc.CFG", NoLocationConfig()) + + with pytest.raises(ValueError, match=r"Provide the grdc path"): + get_grdc_data_monthly( + "30303030", + "2000-01-01T00:00Z", + "2000-03-01T00:00Z", + ) + + +def test_get_grdc_monthly_data_with_missing_datahome(tmp_path): + missing_dir = tmp_path / "does_not_exist" + + with pytest.raises(ValueError, match=r"The grdc directory .* does not exist!"): + get_grdc_data_monthly( + "30303030", + "2000-01-01T00:00Z", + "2000-03-01T00:00Z", + data_home=str(missing_dir), + ) + + +def test_get_grdc_monthly_data_without_file(tmp_path): + with pytest.raises(ValueError, match=r"The grdc file .* does not exist!"): + get_grdc_data_monthly( + "30303031", + "2000-01-01T00:00Z", + "2000-03-01T00:00Z", + data_home=str(tmp_path), + ) + + +def test_get_grdc_monthly_data_from_nc_not_implemented( + tmp_path, sample_grdc_monthly_file +): + # content is irrelevant, only the presence of the file is checked + (tmp_path / "GRDC-Monthly.nc").touch() + + with pytest.raises(NotImplementedError, match=r".nc support not implemented"): + get_grdc_data_monthly( + "30303030", + "2000-01-01T00:00Z", + "2000-03-01T00:00Z", + data_home=str(tmp_path), + ) + + +def test_get_grdc_monthly_data_custom_column_names( + tmp_path, expected_results_monthly: xr.Dataset, sample_grdc_monthly_file +): + result_data_monthly = get_grdc_data_monthly( + "30303030", + "2000-01-01T00:00Z", + "2000-03-01T00:00Z", + data_home=str(tmp_path), + column1="observation", + column2="derived", + column3="quality", + ) + + # column names only rename the intermediate dataframe columns, + # the resulting dataset variables are unaffected + assert_allclose(result_data_monthly, expected_results_monthly) + + +def test_get_grdc_monthly_data_partial_period( + tmp_path, expected_results_monthly: xr.Dataset, sample_grdc_monthly_file +): + result_data_monthly = get_grdc_data_monthly( + "30303030", + "2000-02-01T00:00Z", + "2000-03-01T00:00Z", + data_home=str(tmp_path), + ) + + expected = expected_results_monthly.isel(time=slice(1, None)) + + assert_allclose(result_data_monthly, expected) diff --git a/tests/src/test_analysis.py b/tests/src/test_analysis.py deleted file mode 100644 index b289d8d1..00000000 --- a/tests/src/test_analysis.py +++ /dev/null @@ -1,35 +0,0 @@ -import numpy as np -import pandas as pd -from matplotlib.testing.decorators import image_comparison - -from ewatercycle.analysis import hydrograph - - -@image_comparison( - baseline_images=["hydrograph"], - extensions=["png"], - savefig_kwarg={"bbox_inches": "tight"}, -) -def test_hydrograph(): - ntime = 3000 - - dti = pd.date_range("2018-01-01", periods=ntime, freq="d") - - np.random.seed(20210416) - - discharge = { - "discharge_a": pd.Series(np.linspace(0, 2, ntime), index=dti), - "discharge_b": pd.Series(3 * np.random.random(ntime) ** 2, index=dti), - "discharge_c": pd.Series(2 * np.random.random(ntime) ** 2, index=dti), - "reference": pd.Series(np.random.random(ntime) ** 2, index=dti), - } - - df_q = pd.DataFrame(discharge) - - precipitation = { - "precipitation_a": pd.Series(np.random.random(ntime) / 20, index=dti), - "precipitation_b": pd.Series(np.random.random(ntime) / 30, index=dti), - } - - df_pr = pd.DataFrame(precipitation) - hydrograph(df_q, reference="reference", precipitation=df_pr, nbars=100) diff --git a/tests/src/test_analysis_hydrograph.py b/tests/src/test_analysis_hydrograph.py new file mode 100644 index 00000000..6d6b1d70 --- /dev/null +++ b/tests/src/test_analysis_hydrograph.py @@ -0,0 +1,246 @@ +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import pytest +import xarray as xr + +from ewatercycle.analysis.hydrograph import hydrograph, metric_map + + +def _create_data(): + """Create sample data for testing.""" + ntime = 3000 + + dti = pd.date_range("2018-01-01", periods=ntime, freq="d") + + np.random.seed(20210416) + t = np.arange(ntime) + + discharge_wave = { + "discharge_a": pd.Series( + 1 + np.sin(2 * np.pi * (t - 60) / 365), # sinus wave 0-2 centered at 1 + index=dti, + ), + "discharge_b": pd.Series( + 1.1 + np.sin(2 * np.pi * t / 365), # sinus wave offset ~2 months + index=dti, + ), + "discharge_c": pd.Series( + 1 + np.cos(2 * np.pi * t / 365), # cosine wave + index=dti, + ), + "reference": pd.Series( + 1 + + np.sin(2 * np.pi * t / 365) + + 0.03 * np.random.randn(ntime), # sinus + noise + index=dti, + ), + } + + df_q_wave = pd.DataFrame(discharge_wave) + + precipitation = { + "precipitation_a": pd.Series(np.random.random(ntime) / 20, index=dti), + "precipitation_b": pd.Series(np.random.random(ntime) / 30, index=dti), + } + + df_pr = pd.DataFrame(precipitation) + return df_q_wave, df_pr + + +def _save_figure(fig, fname): + """Save figure to baseline directory.""" + baseline_dir = "tests/src/baseline_images/test_analysis" + fig_path = Path(baseline_dir) / fname + fig.savefig(fig_path, bbox_inches="tight") + + +def test_hydrograph(): + """Test hydrograph with pandas DataFrame input.""" + df_q, df_pr = _create_data() + fig, (ax, ax_tbl) = hydrograph( + df_q, reference="reference", precipitation=df_pr, nbars=100 + ) + + _save_figure(fig, "hydrograph_DataFrame.png") + + # + assert len(ax.lines) == 4 # 3 discharge + 1 reference + assert ax_tbl.tables + + +def test_hydrograph_xarray(): + """Test hydrograph with xarray Dataset input.""" + df_q = _create_data()[0] + ds_q = xr.Dataset.from_dataframe(df_q) + + fig, (ax, ax_tbl) = hydrograph( + ds_q, reference="reference", metrics_list=["kge_2009", "nse_mod", "male"] + ) + + _save_figure(fig, "hydrograph_xarray.png") + + # + assert len(ax.lines) == 4 # 3 discharge + 1 reference + assert ax_tbl.tables + + +def test_hydrograph_xarray_single_year(): + """Test hydrograph with xarray Dataset input and selecting a single year.""" + df_q = _create_data()[0] + ds_q = xr.Dataset.from_dataframe(df_q) + + fig, (ax, ax_tbl) = hydrograph(ds_q, reference="reference", selected_year=2020) + + _save_figure(fig, "hydrograph_xarray_single_year.png") + + # + assert len(ax.lines) == 4 # 3 discharge + 1 reference + assert ax_tbl.tables + + +def test_hydrograph_xarray_single_hydrograph(): + """Test hydrograph with xarray Dataset input and only one discharge to commpare.""" + df_q = _create_data()[0] + ds_q = xr.Dataset.from_dataframe(df_q) + ds_q = ds_q.drop_vars(["discharge_b", "discharge_c"]) + + fig, (ax, ax_tbl) = hydrograph(ds_q, reference="reference") + + _save_figure(fig, "hydrograph_xarray_single_comparison.png") + + # + assert len(ax.lines) == 2 # 3 discharge + 1 reference + assert ax_tbl.tables + + +def test_hydrograph_series_error(): + """Test hydrograph raises error with pandas Series input.""" + df_q, df_pr = _create_data() + ser_q = df_q["discharge_a"] + + try: + hydrograph(ser_q, reference="discharge_a", precipitation=df_pr, nbars=100) + except TypeError as e: + assert ( + str(e) + == "A panda series contains only a single timeseries, please provide a pandas DataFrame or xr.Dataset." + ) + else: + msg = "TypeError not raised" + raise AssertionError(msg) + + +def test_hydrograph_dataarray_single_timeseries_error(): + """Test hydrograph raises error with a one dimensional xarray DataArray.""" + df_q = _create_data()[0] + da_q = xr.DataArray.from_series(df_q["discharge_a"]) + + with pytest.raises(TypeError, match="A DataArray with a single timeseries"): + hydrograph(da_q, reference="discharge_a") + + +def test_hydrograph_dataarray_multidimensional_error(): + """Test hydrograph raises error with a multi dimensional xarray DataArray.""" + df_q = _create_data()[0] + da_q = xr.DataArray(df_q, dims=("time", "run")) + + with pytest.raises(TypeError, match="DataArray with more than one dimension"): + hydrograph(da_q, reference="reference") + + +def test_hydrograph_unsupported_type_error(): + """Test hydrograph raises error with an unsupported input type.""" + with pytest.raises(TypeError, match="Unsupported data type"): + hydrograph([1.0, 2.0, 3.0], reference="reference") + + +def test_hydrograph_selected_year_without_datetimeindex_error(): + """Test selecting a year requires a DatetimeIndex on the discharge.""" + df_q = _create_data()[0].reset_index(drop=True) + + with pytest.raises(ValueError, match="Discharge index must be a DatetimeIndex"): + hydrograph(df_q, reference="reference", selected_year=2020) + + +def test_hydrograph_precipitation_selected_year_without_datetimeindex_error(): + """Test selecting a year requires a DatetimeIndex on the precipitation.""" + df_q, df_pr = _create_data() + df_pr = df_pr.reset_index(drop=True) + + with pytest.raises(ValueError, match="Precipitation index must be a DatetimeIndex"): + hydrograph(df_q, reference="reference", precipitation=df_pr, selected_year=2020) + + +def test_hydrograph_precipitation_selected_year(): + """Test hydrograph slices both discharge and precipitation to a single year.""" + df_q, df_pr = _create_data() + + fig, (ax, ax_tbl) = hydrograph( + df_q, reference="reference", precipitation=df_pr, selected_year=2020 + ) + + ndays_2020 = 366 + assert len(ax.lines) == 4 # 3 discharge + 1 reference + assert len(ax.lines[0].get_xdata()) == ndays_2020 + assert ax_tbl.tables + plt.close(fig) + + +def test_hydrograph_nbars_larger_than_data(): + """Test precipitation is not downsampled when nbars exceeds the number of rows.""" + df_q, df_pr = _create_data() + + fig, (ax, ax_tbl) = hydrograph( + df_q, reference="reference", precipitation=df_pr, nbars=len(df_pr) + 1 + ) + + # both precipitation series are plotted with all their original values + ax_pr = next(child for child in fig.axes if child not in (ax, ax_tbl)) + assert len(ax_pr.containers) == 2 + assert len(ax_pr.containers[0]) == len(df_pr) + plt.close(fig) + + +def test_hydrograph_metric_name_other_case(): + """Test metric names are looked up case insensitively.""" + df_q = _create_data()[0] + + fig, (_, ax_tbl) = hydrograph(df_q, reference="reference", metrics_list=["Nse"]) + + assert len(ax_tbl.tables[0].get_celld()) > 0 + plt.close(fig) + + +def test_hydrograph_metric_as_object(): + """Test a metric can be passed as a callable instead of a name.""" + df_q = _create_data()[0] + + fig, (_, ax_tbl) = hydrograph( + df_q, reference="reference", metrics_list=[metric_map["nse"]] + ) + + assert len(ax_tbl.tables[0].get_celld()) > 0 + plt.close(fig) + + +def test_hydrograph_unknown_metric_error(): + """Test hydrograph raises error for an unknown metric name.""" + df_q = _create_data()[0] + + with pytest.raises(ValueError, match="Metric 'not_a_metric' not found"): + hydrograph(df_q, reference="reference", metrics_list=["not_a_metric"]) + + +def test_hydrograph_saves_file(tmp_path): + """Test hydrograph writes a copy of the figure to the given filename.""" + df_q = _create_data()[0] + filename = tmp_path / "hydrograph.png" + + fig, _ = hydrograph(df_q, reference="reference", filename=filename) + + assert filename.exists() + assert filename.stat().st_size > 0 + plt.close(fig) diff --git a/tests/src/test_util.py b/tests/src/test_util.py index a4db9c97..4757b873 100644 --- a/tests/src/test_util.py +++ b/tests/src/test_util.py @@ -7,6 +7,7 @@ import ewatercycle from ewatercycle.util import ( + extract_package_name, find_closest_point, fit_extents_to_grid, get_package_versions, @@ -259,3 +260,17 @@ def test_version_getter(): assert versions["ewatercycle"] == ewatercycle.__version__ assert "grpc4bmi" in versions assert "remotebmi" in versions + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("ewatercycle_HBV.model:HBV", "ewatercycle_HBV"), + ("ewatercycle_HBV", "ewatercycle_HBV"), + ("ewatercycle_HBV.model", "ewatercycle_HBV"), + # only the first colon separates the module from the object + ("ewatercycle_HBV.model:HBV:extra", "ewatercycle_HBV"), + ], +) +def test_extract_package_name(value, expected): + assert extract_package_name(value) == expected