"""Module for creating plots to visualize the metrics calculated for NWM evaluation.
Functions:
- get_metric_long_name: Get long names for metrics based on the library used for calculating the metrics.
- filter_by_lead_metric: Filter the metric DataFrame by lead times and metrics based on the configuration.
- gather_all_metrics: Gather metrics from all datasets into a single DataFrame.
- add_tag_to_filename: Add a tag to the filename if it exists in the config.
- save_plot: Save the plot to a file with a name based on the metric, lead time, dataset, and other parameters.
- group_df_by_location: Group DataFrame by location and return a tuple of DataFrames (in_list, out_list).
- get_metric_dataframe: Get the metric DataFrame by gathering all metrics and then filtering by metric subset and lead times.
- create_spatial_map: Create spatial maps for each dataset, metric, and lead time.
- set_up_figure: Set up a matplotlib figure with dynamic sizing based on data.
- add_shared_legend: Add a shared legend to the figure.
- create_boxplot: Create boxplots for each metric in the configuration.
- create_histogram: Create histograms for each metric in the configuration.
"""
import logging
from functools import reduce
from pathlib import Path
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from matplotlib.patches import Patch
from .configuration import PlotsConfig
from .nwm_configs import ForecastConfig
from .settings import (
dict_nwm_eval_metrics,
dict_teehr_metrics,
get_metric_bins,
get_metric_colormap,
)
from .utils import clean_data, read_data
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)
[docs]
def get_metric_long_name(metrics: list, library: str):
"""Get long names for metrics.
Args:
metrics: A list of short names for metrics.
library: The library used for calculating the metrics, which determines the mapping of short names to long names. Valid options: "teehr", "nwm.eval".
Returns:
A list of metric long names corresponding to the input short names. If a short name does not have a mapping, the original short name is returned in the long names list.
"""
if library == "teehr":
dict1 = dict_teehr_metrics
elif library == "nwm.eval":
dict1 = dict_nwm_eval_metrics
else:
raise Exception(f" Metric libray not supported: {library}")
metrics_long = []
for m1 in metrics:
if m1 in dict1.keys():
metrics_long = metrics_long + [dict1.get(m1).title()]
else:
metrics_long = metrics_long + [m1]
return metrics_long
[docs]
def filter_by_lead_metric(
df_metrics: pd.DataFrame, conf: dict, nwm_config: str, fcst_config_file: str
):
"""Filter the metric DataFrame by lead times and metrics.
Args:
df_metrics: A DataFrame containing the metrics to be filtered, with columns including "lead_group" and "metric".
conf: A dictionary containing the configuration for filtering, which may include "lead_times" and "metric_subset".
nwm_config: A string representing the NWM configuration, used for interpreting lead times.
fcst_config_file: A string representing the path to the forecast configuration file, used for interpreting lead times.
Returns:
A filtered DataFrame containing only the specified lead times and metrics, sorted by lead times in a logical order.
"""
# first filter by lead times
lead_times = conf.get("lead_times", [])
if lead_times:
fc = ForecastConfig(fcst_config_file)
leads0, _, _ = fc.interpret_lead_times(lead_times, nwm_config)
leads = df_metrics["lead_group"].unique()
leads1 = [l1 for l1 in leads0 if l1 not in leads]
if len(leads1) > 0:
logger.warning(
f"Lead times {leads1} not found in computed metric results. Skipping them. Available lead times are: {leads}"
)
df_metrics = df_metrics[df_metrics["lead_group"].isin(leads0)]
# then filter by metric (if a metric subset is specified in the config)
metrics = conf.get("metric_subset", [])
if metrics:
mts = df_metrics["metric"].unique()
mts1 = [m1 for m1 in metrics if m1 not in mts]
if len(mts1) > 0:
raise Exception(f"Metrics {mts1} not found in computed metric results")
df_metrics = df_metrics[df_metrics["metric"].isin(metrics)]
# sort the data by lead times in a logical order (e.g., 1, 2, 3, ..., 12, 1-5, 6-10, etc.)
def parse_lead_value(v):
v = str(v).lower()
if v.startswith("m"):
return -float(v[1:])
return float(v)
def lead_key(x):
x = str(x).lower()
if x.startswith("m") and "-" not in x:
return (0, parse_lead_value(x)) # m-values first
elif "-" in x[1:]:
start = x.split("-", 1)[0]
return (2, parse_lead_value(start)) # ranges last
else:
return (1, parse_lead_value(x)) # singles in middle
df_metrics["lead_group"] = pd.Categorical(
df_metrics["lead_group"].astype(str),
categories=sorted(df_metrics["lead_group"].unique(), key=lead_key),
ordered=True,
)
df_metrics = df_metrics.sort_values("lead_group")
return df_metrics
[docs]
def gather_all_metrics(datasets: list, data_paths: dict):
"""Gather metrics from all datasets into a single DataFrame.
Args:
datasets: A list of dataset names to gather metrics for.
data_paths: A dictionary containing the paths to the metric files for each dataset.
Returns:
A DataFrame containing the gathered metrics from all datasets.
"""
df_metrics = pd.DataFrame()
dfs = []
for dataset in datasets:
metric_file = data_paths[dataset]
if not metric_file.exists():
msg = f"Metric file for dataset {dataset} does not exist at {metric_file}; compute metrics first."
logger.error(msg)
raise FileNotFoundError(msg)
if metric_file.suffix.lower() == ".csv":
df = pd.read_csv(metric_file)
else:
df = pd.read_parquet(metric_file)
df = df.melt(
id_vars=["lead_group", "primary_location_id"],
var_name="metric",
value_name=dataset,
)
dfs = dfs + [df]
df_metrics = reduce(
lambda left, right: pd.merge(
left, right, on=["primary_location_id", "lead_group", "metric"], how="inner"
),
dfs,
)
df_metrics = df_metrics.melt(
id_vars=["lead_group", "primary_location_id", "metric"],
value_vars=datasets,
var_name="dataset",
)
return df_metrics
[docs]
def add_tag_to_filename(conf: dict, file_name: str) -> str:
"""Add a tag to the filename if it exists in the config.
Args:
conf: A dictionary containing the configuration, which may include a "tag" key.
file_name: The original filename to which the tag should be added.
Returns:
A string representing the filename with the tag added if it exists, or the original filename if no tag is specified in the config.
"""
path = Path(file_name)
tag = conf.get("tag")
if tag:
path = path.with_name(f"{path.stem}_{tag}{path.suffix}")
return path.as_posix()
[docs]
def save_plot(
plt: plt,
conf: dict,
data_paths: dict,
plt_type: str,
plt_name: str = None,
lead: str = None,
ref_time: str = None,
dataset: str = None,
metric: str = None,
location: str = None,
) -> Path:
"""Save the plot to a file.
Args:
plt: The plot object to be saved.
conf: A dictionary containing the configuration, which may include a "tag" key.
data_paths: A dictionary containing the paths to the data files.
plt_type: The type of the plot.
plt_name: The name of the plot.
lead: The lead time for the plot.
ref_time: The reference time for the plot.
dataset: The dataset name for the plot.
metric: The metric name for the plot.
location: The location name for the plot.
Returns:
A Path object representing the saved plot file.
"""
fig_dir = Path(data_paths["plots"], plt_type)
fig_dir.mkdir(parents=True, exist_ok=True)
if not plt_name:
plt_name = plt_type
fname = f"{plt_name}"
fname += f"_{metric}" if metric else ""
fname += f"_T0_{ref_time}" if ref_time else ""
fname += f"_lead_h{lead}" if lead and lead != "0" else ""
fname += f"_{location}" if location else ""
fname += f"_{dataset}" if dataset else ""
fname += ".png"
file1 = add_tag_to_filename(conf["plots"][plt_type], fname)
fig_file = Path(fig_dir, file1)
plt.savefig(fig_file)
plt.close()
return fig_file
[docs]
def group_df_by_location(
df: pd.DataFrame | gpd.GeoDataFrame, location_list: list, location_col: str
) -> tuple:
"""Group DataFrame by location and return a tuple of DataFrames (in_list, out_list).
Args:
df: A DataFrame or GeoDataFrame containing the data to be grouped.
location_list: A list of location identifiers to be used for grouping.
location_col: The name of the column in the DataFrame that contains the location identifiers.
Returns:
A tuple of two DataFrames: (df_inlist, df_outlist), where df_inlist contains rows with location identifiers in the location_list, and df_outlist contains rows with location identifiers not in the location_list.
"""
df_inlist = df[df[location_col].isin(location_list)]
df_outlist = df[~df[location_col].isin(location_list)]
return df_inlist, df_outlist
[docs]
def get_metric_dataframe(conf: dict, data_paths: dict, plot_type: str) -> pd.DataFrame:
"""Get the metric DataFrame by gathering all metrics and then filtering by metric subset and lead times.
Args:
conf: A dictionary containing the configuration for gathering and filtering the metrics, which may include "lead_times" and "metric_subset".
data_paths: A dictionary containing the paths to the metric files for each dataset.
plot_type: The type of the plot for which the metric DataFrame is being prepared, used to determine which lead times and metrics to filter by.
Returns:
A DataFrame containing the filtered metrics.
"""
# gather all metrics calcualted
datasets = conf["general"]["dataset_name"]
df_metrics = gather_all_metrics(datasets, data_paths["metrics"])
# filter metric dataframe by lead times and metrics
df_metrics = filter_by_lead_metric(
df_metrics,
conf["plots"][plot_type],
conf["general"]["nwm_configuration"],
conf["file_paths"]["fcst_config_file"],
)
# sort by dataset, then metric, then lead times
df_metrics = df_metrics.sort_values(by=["dataset", "metric", "lead_group"])
return df_metrics
[docs]
def create_spatial_map(conf: dict, data_paths: dict):
"""Create spatial maps for each dataset, metric, and lead time.
Args:
conf: A dictionary containing the configuration for creating spatial maps, which may include "lead_times", "metric_subset", and plotting parameters.
data_paths: A dictionary containing the paths to the metric files for each dataset and the crosswalk file.
Returns:
None
"""
# get metric dataframe filtered by metric subset and lead times
df_metrics = get_metric_dataframe(conf, data_paths, "spatial_map")
leads = df_metrics["lead_group"].unique()
metrics = df_metrics["metric"].unique()
metrics_long = get_metric_long_name(metrics, conf["metrics"]["library"])
# add geometry (lat/lon)
df_geo = read_data(data_paths["crosswalk"][list(data_paths["crosswalk"].keys())[0]])
if "geometry" not in df_geo.columns:
logger.warning(
f"Geometry column not found in crosswalk file: {data_paths['crosswalk']}"
f"; spatial maps cannot be created."
)
return
df_geo = df_geo[["primary_location_id", "geometry"]]
gdf_metrics = df_geo.merge(df_metrics, on="primary_location_id", how="inner")
# loop through lead times, metrics, and datasets to create spatial maps
conf1 = conf["plots"]["spatial_map"]
cmap1 = get_metric_colormap(conf1, "map")
for lead1 in leads:
for metric1, metric_long in zip(metrics, metrics_long):
for dataset in gdf_metrics["dataset"].unique():
# filter data based on lead time, metric, and dataset
filtered_gdf = gdf_metrics.query(
f"lead_group == '{lead1}' & metric == '{metric1}' & dataset == '{dataset}'"
)
# clean the data by removing NaN and infinite values
filtered_gdf = clean_data(filtered_gdf, "value")
if filtered_gdf.empty:
logger.warning(
f"No data available for metric {metric1} at lead time {lead1} for dataset {dataset}. "
f"Skipping map creation."
)
continue
logger.debug(
f"Number of locations after filtering nan and inf: {filtered_gdf['primary_location_id'].nunique()}"
)
# clip the data
if not np.isnan(cmap1[metric1]["clim"][0]) and not np.isnan(
cmap1[metric1]["clim"][1]
):
filtered_gdf["value"] = np.clip(
filtered_gdf["value"],
cmap1[metric1]["clim"][0],
cmap1[metric1]["clim"][1],
)
# convert to EPSG:4326 for plotting
filtered_gdf = filtered_gdf.to_crs("EPSG:4326")
# set up the figure and axes with Cartopy projection
fig, ax = plt.subplots(
figsize=(8.5, 6), subplot_kw={"projection": ccrs.PlateCarree()}
)
ax.set_title(
f"{metric1} ({metric_long}), "
f"{'' if lead1 == '0' else f'lead={lead1}h, '}"
f"{dataset}",
fontsize=16,
pad=16,
)
# Add map features
ax.add_feature(cfeature.COASTLINE)
ax.add_feature(cfeature.BORDERS, linestyle=":")
ax.add_feature(cfeature.LAND, facecolor="lightgray")
ax.add_feature(cfeature.OCEAN, facecolor="lightblue")
# Dynamically compute point size
n_points = len(filtered_gdf)
base_size = 10000 # adjust this constant to control density sensitivity
point_size = max(
10, base_size / n_points
) # minimum size to keep the points visible
point_size = min(100, point_size) # don't want them too big either
# Split GeoDataFrame into calibrated and non-calibrated subsets
gdf_calib, gdf_noncalib = group_df_by_location(
filtered_gdf,
conf["general"].get("calib_gages", []),
"primary_location_id",
)
# Plot non-calibrated points (circles)
sc1 = ax.scatter(
gdf_noncalib.geometry.x,
gdf_noncalib.geometry.y,
transform=ccrs.PlateCarree(),
c=gdf_noncalib["value"],
cmap=cmap1[metric1]["cmap"],
marker="o",
edgecolor="k",
linewidth=0.5,
s=point_size,
alpha=0.8,
label=f"Non-calibrated(n={gdf_noncalib['primary_location_id'].nunique()})",
)
# Plot calibrated points (triangle symbols)
if not gdf_calib.empty:
sc2 = ax.scatter(
gdf_calib.geometry.x,
gdf_calib.geometry.y,
transform=ccrs.PlateCarree(),
c=gdf_calib["value"],
cmap=cmap1[metric1]["cmap"],
marker="^",
edgecolor="k",
linewidth=1.0,
s=point_size * 1.2, # make them slightly more visible
alpha=0.8,
label=f"Calibrated(n={gdf_calib['primary_location_id'].nunique()})",
)
# Add legend under the colorbar to avoid being clipped
ax.legend(
loc="upper center",
bbox_to_anchor=(
0.5,
-0.25,
), # centered, below the axes & colorbar
frameon=True,
fontsize=12,
ncol=2,
)
# Colorbar from one of the scatter plots
cbar = plt.colorbar(
sc1,
ax=ax,
label="",
orientation="horizontal",
pad=0.02,
)
cbar.ax.tick_params(labelsize=14)
# set aspect ratio
ax.set_aspect(1.35)
# save plot to png
fig_file = save_plot(
plt,
conf,
data_paths,
"spatial_map",
"map",
lead=str(lead1),
dataset=dataset,
metric=metric1,
)
logger.info(f" Spatial maps created at: {Path(fig_file).parent}")
[docs]
def add_shared_legend(fig, axes, dataset_names, palette):
"""Add a shared legend to the figure.
Args:
fig: The figure object to which the legend will be added.
axes: A list of axes objects from which individual legends will be removed.
dataset_names: A list of dataset names to be included in the legend.
palette: A dictionary mapping dataset names to colors, used for creating legend handles.
Returns:
None
"""
# Create manual legend handles
handles_labels = {
name: Patch(color=palette[name], label=name) for name in dataset_names
}
# Remove axes legends if they exist
for ax in axes:
if ax.get_legend() is not None:
ax.get_legend().remove()
# Add shared legend to the figure
fig.legend(
handles_labels.values(),
handles_labels.keys(),
loc="upper center",
ncol=len(handles_labels),
frameon=True,
fontsize=12,
bbox_to_anchor=(0.5, 0.98),
borderaxespad=0.0,
)
# leave space for the legend
plt.tight_layout(rect=[0, 0, 1, 0.92])
[docs]
def create_boxplot(conf: dict, data_paths: dict):
"""Create boxplots for each metric in the configuration.
Args:
conf: A dictionary containing the configuration for creating boxplots, which may include "lead_times", "metric_subset", and plotting parameters.
data_paths: A dictionary containing the paths to the metric files for each dataset.
Returns:
None
"""
# get metric dataframe filtered by metric subset and lead times
df_metrics = get_metric_dataframe(conf, data_paths, "boxplot")
metrics = df_metrics["metric"].unique()
metrics_long = get_metric_long_name(metrics, conf["metrics"]["library"])
# ensure consistent colors applied to each dataset across metrics
dataset_names = sorted(df_metrics["dataset"].unique())
palette = dict(zip(dataset_names, sns.color_palette("tab10", len(dataset_names))))
# create boxplot for each metric
conf1 = conf["plots"]["boxplot"]
cmap1 = get_metric_colormap(conf1, "boxplot")
for metric1, metric_long in zip(metrics, metrics_long):
# filter the data by metric
df1 = df_metrics[df_metrics["metric"] == metric1]
# clean the data by removing NaN and infinite values
df1 = clean_data(df1, "value")
if df1.empty:
logger.warning(
f"No data available for metric {metric1}. Skipping boxplot creation."
)
continue
logger.debug(
f"Number of locations after filtering nan and inf: {df1['primary_location_id'].nunique()}"
)
# clip the data
if not np.isnan(cmap1[metric1]["clim"][0]) and not np.isnan(
cmap1[metric1]["clim"][1]
):
df1["value"] = np.clip(
df1["value"], cmap1[metric1]["clim"][0], cmap1[metric1]["clim"][1]
)
# break data into calibrated and non-calibrated subsets
df_calib, df_noncalib = group_df_by_location(
df1, conf["general"].get("calib_gages", []), "primary_location_id"
)
# set up figure
fig, axes = set_up_figure(df_calib, df_noncalib, plot_type="boxplot")
# Plot non-calibrated
sns.boxplot(
x="lead_group",
y="value",
data=df_noncalib,
hue="dataset",
hue_order=dataset_names,
showfliers=conf1["show_outliers"],
palette=palette,
ax=axes[0],
)
str1 = df_noncalib["primary_location_id"].nunique()
tit1 = (
f"Non-Calibrated Locations \n(n={str1})"
if len(df_calib) > 0
else f"All Locations \n(n={str1})"
)
axes[0].set_title(tit1)
axes[0].set_ylabel(f"{metric1} ({metric_long})", fontsize=12)
axes[0].set_yticklabels(axes[0].get_yticklabels(), fontsize=10)
# Plot calibrated if it exists
if len(df_calib) > 0:
sns.boxplot(
x="lead_group",
y="value",
data=df_calib,
hue="dataset",
hue_order=dataset_names,
showfliers=conf1["show_outliers"],
palette=palette,
ax=axes[1],
)
str2 = df_calib["primary_location_id"].nunique()
axes[1].set_title(f"Calibrated Locations \n(n={str2})")
# Remove x-ticks and labels from all subplots
if (df1["lead_group"] == "0").all():
for ax in axes:
ax.set_xticks([])
ax.set_xticklabels([])
ax.set_xlabel("")
else:
for ax in axes:
ax.set_xlabel("Lead time (hours)", fontsize=12)
ax.set_xticklabels(ax.get_xticklabels(), rotation=45, fontsize=12)
# set y-axis label font size
for ax in axes:
ax.set_ylabel(f"{metric1} ({metric_long})", fontsize=12)
ax.tick_params(axis="y", labelsize=12)
# Add shared legend
add_shared_legend(fig, axes, dataset_names, palette)
# save plot to png
fig_file = save_plot(
plt,
conf,
data_paths,
"boxplot",
metric=metric1,
)
logger.info(f" Boxplots created at: {Path(fig_file).parent}")
[docs]
def create_histogram(conf: dict, data_paths: dict):
"""Create histograms for each metric in the configuration.
Args:
conf: A dictionary containing the configuration for creating histograms, which may include "lead_times", "metric_subset", and plotting parameters.
data_paths: A dictionary containing the paths to the metric files for each dataset.
Returns:
None
"""
# get metric dataframe filtered by metric subset and lead times
df_metrics = get_metric_dataframe(conf, data_paths, "histogram")
leads = df_metrics["lead_group"].unique()
metrics = df_metrics["metric"].unique()
metrics_long = get_metric_long_name(metrics, conf["metrics"]["library"])
# get custom bin edges for the metrics
conf1 = conf["plots"]["histogram"]
custom_bins = get_metric_bins(conf1)
# ensure consistent colors applied to each dataset across metrics
dataset_names = sorted(df_metrics["dataset"].unique())
palette = dict(zip(dataset_names, sns.color_palette("tab10", len(dataset_names))))
# loop through metrics and lead times to create histograms
for metric1, metric_long in zip(metrics, metrics_long):
# bin edges for the metric
bins1 = custom_bins.get(metric1)
for lead1 in leads:
# filter data by metric and lead time
df = df_metrics[
(df_metrics["metric"] == metric1) & (df_metrics["lead_group"] == lead1)
]
# clean the data by removing NaN and infinite values
df = clean_data(df, "value")
if df.empty:
logger.warning(
f"No data available for metric {metric1} at lead time {lead1}. Skipping histogram creation."
)
continue
logger.debug(
f"Number of locations after filtering nan and inf: {df['primary_location_id'].nunique()}"
)
# bin the data to create customized histograms
if len(bins1) > 0:
df["binned"] = pd.cut(df["value"], bins=bins1, include_lowest=True)
else:
df["binned"] = pd.cut(df["value"], bins=8, include_lowest=True)
df = df.sort_values(by=["binned"])
# check if any nan values after binning
if df["binned"].isnull().any():
logger.warning(
f"Some binned values are NaN for metric {metric1} at lead time {lead1}. Check bin edges."
)
# convert binned column to string for plotting
df["binned"] = df["binned"].astype(str)
# break data into calibrated and non-calibrated subsets
df_calib, df_noncalib = group_df_by_location(
df, conf["general"].get("calib_gages", []), "primary_location_id"
)
# set up figure
fig, axes = set_up_figure(df_calib, df_noncalib, plot_type="histogram")
plt.set_loglevel("WARNING")
# Plot non-calibrated
ax = sns.histplot(
data=df_noncalib,
x="binned",
hue="dataset",
hue_order=dataset_names,
multiple="dodge",
shrink=0.8,
palette=palette,
discrete=True,
ax=axes[0],
)
str1 = df_noncalib["primary_location_id"].nunique()
tit1 = (
f"Non-Calibrated Locations (n={str1})"
if len(df_calib) > 0
else f"All Locations (n={str1})"
)
axes[0].set_title(tit1)
axes[0].set_ylabel(f"{metric1} ({metric_long})", fontsize=12)
axes[0].set_yticklabels(axes[0].get_yticklabels(), fontsize=10)
if len(df_calib) > 0:
ax = sns.histplot(
data=df_calib,
x="binned",
hue="dataset",
hue_order=dataset_names,
multiple="dodge",
shrink=0.8,
palette=palette,
discrete=True,
ax=axes[1],
)
str2 = df_calib["primary_location_id"].nunique()
axes[1].set_title(f"Calibrated Locations (n={str2})")
# set x-tick labels rotation and titles
for ax in axes:
plt.setp(ax.get_xticklabels(), rotation=30, fontsize=10)
ax.set_ylabel("Number of locations", fontsize=12)
ax.tick_params(axis="y", labelsize=12)
ax.set_xlabel(
f"{metric1} ({metric_long})"
f"{'' if lead1 == '0' else f' lead={lead1}h'}",
fontsize=12,
)
plt.subplots_adjust(bottom=0.2)
# add shared legend
add_shared_legend(fig, axes, dataset_names, palette)
# save plot to png
fig_file = save_plot(
plt,
conf,
data_paths,
"histogram",
"hist",
lead=str(lead1),
metric=metric1,
)
logger.info(f" Histograms created at: {Path(fig_file).parent}")
[docs]
def plot_time_series(
conf: dict,
data_paths: dict,
df: pd.DataFrame,
unit: str,
lead: str = None,
ref_time: str = None,
):
"""Plot time series of observed and forecasted streamflow for a single location.
Args:
conf: A dictionary containing the configuration for plotting, which may include plotting parameters.
data_paths: A dictionary containing the paths to the data files.
df: A DataFrame containing the time series data to be plotted, with columns including "value_time", "primary_value", and forecast dataset columns.
unit: A string representing the measurement unit for the streamflow values, used for labeling the y-axis.
lead: A string representing the lead time for the plot, used for labeling the title. Optional if ref_time is provided.
ref_time: A string representing the reference time for the plot, used for labeling the title. Optional if lead is provided.
Returns:
A string representing the file path of the saved plot image.
"""
# lead and ref_time cannot be both not None
if lead is not None and ref_time is not None:
msg = (
"Both lead and reference time are provided for time series plot. "
f"Please provide only one of them. lead={lead}, ref_time={ref_time}"
)
logger.error(msg)
raise ValueError(msg)
lead_ref_str = f"Lead time: {lead}h" if lead else f"Reference time: {ref_time}"
# sort the data by value_time to avoid zigzag lines
df = df.sort_values("value_time")
# Plot
fig, ax = plt.subplots(figsize=(10, 6))
n_points = len(df)
ax.plot(
df["value_time"],
df["primary_value"],
label="Observed",
color="black",
linewidth=1.2,
marker="o" if n_points < 30 else None,
linestyle="-" if n_points >= 2 else "None",
)
for dataset in conf["general"]["dataset_name"]:
ax.plot(
df["value_time"],
df[dataset],
label=dataset,
marker="o" if n_points < 30 else None,
linestyle="-" if n_points >= 2 else "None",
)
ax.set_xlabel("Time", fontsize=14)
ax.set_ylabel(f"Streamflow ({unit})", fontsize=14)
ax.tick_params(axis="both", labelsize=12)
ax.set_title(
f"Simulated vs Observed Streamflow at {conf['general']['location_list'][0]}\n"
f"{conf['general']['nwm_configuration']} {lead_ref_str}",
fontsize=14,
fontweight="bold",
)
ax.legend(fontsize=12)
ax.grid(True)
fig.autofmt_xdate()
# save plot to png
fig_file = save_plot(
plt,
conf,
data_paths,
"time_series",
lead=lead,
ref_time=ref_time,
)
return fig_file
[docs]
def get_time_series(
data_paths: dict, data_type: str, dataset: str = None
) -> tuple[pd.DataFrame, str]:
"""Get forecast or observed time series data (and unit) from forecast or observed data file.
Cannot use paired data file which has forecast data trimmed to the time range of observed data by teehr.
Args:
data_paths: A dictionary containing the paths to the data files.
data_type: A string indicating the type of data to retrieve, either "observed" or "forecast".
dataset: A string representing the dataset name for forecast data. Optional for observed data but required for forecast data.
Returns:
A tuple containing a DataFrame with the time series data and a string representing the measurement unit for the streamflow values.
"""
if data_type == "observed":
path_str = "obs"
elif data_type == "forecast":
path_str = "fcst_link"
if not dataset:
raise ValueError("Dataset name must be provided for forecast data.")
else:
raise ValueError(
f"Invalid data type: {data_type}. Must be 'observed' or 'forecast'."
)
if data_type == "observed":
data_dir = Path(data_paths.get(path_str, {}))
dataset_str = "observed data"
else:
data_dir = Path(data_paths.get(path_str, {}).get(dataset, ""))
dataset_str = f"dataset {dataset}"
if not data_dir.exists():
msg = f"Data directory not found for {dataset_str}: {data_dir}. Cannot create time series plot."
logger.error(msg)
raise FileNotFoundError(msg)
parquet_files = list(data_dir.glob("*.parquet"))
if not parquet_files:
msg = f"No parquet files found in {data_dir} for {dataset_str}. Cannot create time series plot."
logger.error(msg)
raise FileNotFoundError(msg)
df = pd.concat(
[
pd.read_parquet(f)[
["value_time", "value", "reference_time", "measurement_unit"]
]
for f in parquet_files
],
ignore_index=True,
)
# calculate lead time in hours
if data_type == "forecast":
df["lead_time"] = (
(df["value_time"] - df["reference_time"])
.dt.total_seconds()
.div(3600)
.round()
.astype("Int32")
)
df = df.rename(columns={"value": dataset})
else:
df = df.rename(columns={"value": "primary_value"})
unit = (
df["measurement_unit"].iloc[0]
if not df["measurement_unit"].isnull().all()
else None
)
df.drop(columns=["measurement_unit"], inplace=True)
return df, unit
[docs]
def create_time_series(conf: dict, data_paths: dict):
"""Create a time series plot for each dataset in the configuration.
Args:
conf: A dictionary containing the configuration for creating time series plots, which may include "lead_times", "reference_times", and plotting parameters.
data_paths: A dictionary containing the paths to the observed and forecast data files.
Returns:
None
"""
# get observed data (same for all datasets )
merged_df, unit_obs = get_time_series(data_paths, "observed")
merged_df.drop(columns=["reference_time"], inplace=True)
# loop through datasets to get forecast data and merge with observed data
datasets = conf["general"]["dataset_name"]
for dataset in datasets:
df_fcst, unit_fcst = get_time_series(data_paths, "forecast", dataset)
if unit_obs != unit_fcst:
logger.warning(
f"Measurement unit mismatch for dataset {dataset}: observed unit is '{unit_obs}' but forecast unit is '{unit_fcst}'. "
f"Using observed unit '{unit_obs}' for the plot."
)
# subset forecast based on forecast_start_date and forecast_end_date in config (if specified)
df_fcst["reference_time"] = pd.to_datetime(df_fcst["reference_time"])
for time1 in ["forecast_start_date", "forecast_end_date"]:
fcst_time_list = conf["general"].get(time1)
if fcst_time_list:
fcst_time = fcst_time_list[datasets.index(dataset)]
fcst_time = pd.to_datetime(fcst_time, errors="coerce")
if pd.isnull(fcst_time):
logger.warning(
f"Invalid {time1} in config for dataset {dataset}. Ignoring this filter."
)
else:
if time1 == "forecast_start_date":
df_fcst = df_fcst[df_fcst["reference_time"] >= fcst_time]
elif time1 == "forecast_end_date":
df_fcst = df_fcst[df_fcst["reference_time"] <= fcst_time]
# merge observed and forecast data for a given dataset
cols = ["value_time", "reference_time", "lead_time"]
cols_merge = [
c1 for c1 in cols if c1 in df_fcst.columns and c1 in merged_df.columns
]
cols_fcst = cols + [dataset]
cols_fcst = [c1 for c1 in cols_fcst if c1 in df_fcst.columns]
merged_df = pd.merge(merged_df, df_fcst[cols_fcst], on=cols_merge, how="outer")
if merged_df is None or merged_df.empty:
logger.error("No data available to create time series plot.")
return
# remove duplicates if any
merged_df = merged_df.drop_duplicates()
# Ensure datetime
merged_df["value_time"] = pd.to_datetime(merged_df["value_time"])
merged_df["reference_time"] = pd.to_datetime(merged_df["reference_time"])
# if neither lead times nor reference times are specified in config, use the first reference time in the data
lead_times = conf["plots"]["time_series"].get("lead_times", [])
ref_times = conf["plots"]["time_series"].get("reference_times", [])
if not lead_times and not ref_times:
ref_times = merged_df["reference_time"].min()
ref_times = [ref_times]
logger.info(
"No lead times or reference times specified for time series plot; "
f"using the earliest reference time {ref_times} in the data."
)
ref_times = list(pd.to_datetime(ref_times, errors="coerce"))
fig_file = None
# Plot time series for each specified lead time
if lead_times:
logger.info(f"Creating time series plots for lead times: {lead_times}")
merged_df["lead_time"] = merged_df["lead_time"].astype(str)
for lead in lead_times:
df_subset = merged_df[merged_df["lead_time"] == str(lead)]
if df_subset.empty:
logger.warning(f"No data available for lead time {lead}h. Skipping plot.")
continue
fig_file = plot_time_series(
conf, data_paths, df_subset, unit_obs, lead=str(lead)
)
# plot time series for each specified reference time
if ref_times:
logger.info(
f"Creating time series plots for reference times: {[ref.strftime('%Y-%m-%d %H:%M') for ref in ref_times]}"
)
for ref in ref_times:
df_subset = merged_df[merged_df["reference_time"] == ref]
if df_subset.empty:
logger.warning(
f"No data available for reference time {ref}. Skipping plot."
)
continue
fig_file = plot_time_series(
conf, data_paths, df_subset, unit_obs, ref_time=ref.strftime("%Y%m%dT%H")
)
if fig_file:
logger.info(f" Time series plots created at: {Path(fig_file).parent}")
[docs]
def get_metric_groups() -> dict:
"""Get metric groups for the configuration.
Args:
None
Returns:
A dictionary mapping metric group names to lists of metric short names.
"""
# Split metric columns into groups
metric_groups = {
"Standard": [
"CORR",
"KGE",
"NNSE",
"NSE",
"NSElog",
"NSEwt",
"RSR",
"RMSE",
"MAE",
"PBIAS",
],
"Categorical": ["POD", "FAR", "CSI", "FBIAS"],
"Event-based": ["PKBIAS", "PKTE", "EVBIAS"],
"FDC-based": ["HSEG_FDC", "MSEG_FDC", "LSEG_FDC"],
}
return metric_groups
[docs]
def create_metric_table(conf: dict, data_paths: dict):
"""Create a metric table based on the configuration.
Args:
conf: A dictionary containing the configuration for creating metric tables, which may include "lead_times", "metric_subset", and plotting parameters.
data_paths: A dictionary containing the paths to the metric files for each dataset.
Returns:
None
"""
# get metric dataframe filtered by metric subset and lead times
df_metrics_all = get_metric_dataframe(conf, data_paths, "metric_table")
leads = df_metrics_all["lead_group"].unique()
location = df_metrics_all["primary_location_id"].unique()[0]
location = location.split("-")[-1]
# loop through lead times create tables
for lead in leads:
df_metrics = df_metrics_all[df_metrics_all["lead_group"] == lead]
# convert long format to wide format
df_metrics = df_metrics[["dataset", "metric", "value"]].drop_duplicates()
df_metrics.rename(columns={"dataset": "Formulation"}, inplace=True)
df_metrics = df_metrics.pivot(
index="Formulation", columns="metric", values="value"
).reset_index()
metric_groups = get_metric_groups()
n_groups = len(metric_groups)
# Create subplots for each table
max_form_len = df_metrics["Formulation"].str.len().max()
fig, axes = plt.subplots(
n_groups, 1, figsize=(9 + max_form_len * 0.1, 1.2 * n_groups)
)
if n_groups == 1: # if only one group, axes is not iterable
axes = [axes]
for ax, group in zip(axes, metric_groups.keys()):
ax.axis("off") # hide axis
# Round values to 2 digits and format large numbers in scientific notation
def fmt_value(x):
if isinstance(x, (int, float, np.floating)):
if abs(x) >= 1e4:
return f"{x:.2e}" # scientific notation for large numbers
else:
return f"{x:.2f}" # standard float with 2 decimals
return x
cols = ["Formulation"] + metric_groups[group]
cols = [c1 for c1 in cols if c1 in df_metrics.columns]
# if no metrics in this group, display text and skip table
if len(cols) < 2:
ax.text(
0.5,
0.5,
f"No {group} metrics to display",
horizontalalignment="center",
verticalalignment="center",
fontsize=12,
color="dimgrey",
)
continue
display_df = df_metrics[cols].copy().applymap(fmt_value)
cell_values = display_df.values
table = ax.table(
cellText=cell_values,
colLabels=cols,
cellLoc="center",
loc="center",
)
# automatically adjust the column widths
table.auto_set_column_width(col=list(range(len(df_metrics.columns))))
# Header styling
for (i, j), cell in table.get_celld().items():
if i == 0: # first row = header
cell.set_facecolor("teal")
cell.set_text_props(weight="bold", color="white")
table.auto_set_font_size(False)
table.set_fontsize(11)
table.scale(1.2, 2.0)
# Add title above the table
ax.set_title(
f"{group} Metrics", fontsize=14, pad=8, weight="bold", color="dimgrey"
)
# Overall title
conf1 = conf["general"]
title = f"Metrics for {conf1['location_list'][0]} {conf1['nwm_configuration']} "
title += f"(Lead = {lead} hour)"
fig.suptitle(title, fontsize=14, fontweight="bold", y=0.98)
plt.tight_layout()
# save plot to png
fig_file = save_plot(
plt,
conf,
data_paths,
"metric_table",
lead=str(lead),
location=location,
)
logger.info(f" Metric table plots created at: {Path(fig_file).parent}")
[docs]
def create_barchart(conf: dict, data_paths: dict):
"""Create a bar chart comparing datasets for each metric and lead time.
Args:
conf: A dictionary containing the configuration for creating bar charts, which may include "lead_times", "metric_subset", and plotting parameters.
data_paths: A dictionary containing the paths to the metric files for each dataset.
Returns:
None
"""
# get metric dataframe filtered by metric subset and lead times
df_metrics = get_metric_dataframe(conf, data_paths, "barchart")
leads = df_metrics["lead_group"].unique()
if len(leads) > 1:
barchart_by_metric(df_metrics, conf, data_paths)
else:
barchart_all_metrics(df_metrics, conf, data_paths)
[docs]
def barchart_by_metric(df_metrics: pd.DataFrame, conf: dict, data_paths: dict):
"""Create a bar chart comparing datasets for each metric and lead time.
Args:
df_metrics: A DataFrame containing the metric values to be plotted, with columns including "primary_location_id", "metric", "lead_group", "dataset", and "value".
conf: A dictionary containing the configuration for creating bar charts, which may include plotting parameters.
data_paths: A dictionary containing the paths to the metric files for each dataset.
Returns:
None
"""
# get location ID
location = df_metrics["primary_location_id"].unique()[0]
location = location.split("-")[-1]
# plot each metric separately with lead times on x-axis and bars for different datasets
metrics = df_metrics["metric"].unique()
metrics_long = get_metric_long_name(metrics, conf["metrics"]["library"])
for metric, metric_long in zip(metrics, metrics_long):
df_m = df_metrics[df_metrics["metric"] == metric]
# Identify groups
lead_times = df_m["lead_group"].unique()
datasets = df_m["dataset"].unique()
n_leads = len(lead_times)
n_datasets = len(datasets)
fig, ax = set_up_figure(pd.DataFrame(), df_m, plot_type="barchart")
ax = ax[0]
# Bar positioning
x = range(n_leads)
bar_width = 0.8 / n_datasets # spread bars within each lead-time group
for i, dataset in enumerate(datasets):
df_dataset = df_m[df_m["dataset"] == dataset]
# Ensure alignment with lead_times
values = [
df_dataset[df_dataset["lead_group"] == lt]["value"].values[0]
if lt in df_dataset["lead_group"].values
else 0
for lt in lead_times
]
positions = [xi + i * bar_width for xi in x]
ax.bar(positions, values, width=bar_width, label=dataset)
# Formatting
ax.set_title(f"{metric_long} by Lead Time ({location})", fontsize=14)
ax.set_xticks([xi + bar_width * (n_datasets - 1) / 2 for xi in x])
ax.set_xticklabels(lead_times, fontsize=10)
ax.set_xlabel("Lead Time (hours)", fontsize=12)
ax.set_ylabel(f"{metric} ({metric_long})", fontsize=12)
ax.legend(fontsize=12)
plt.tight_layout()
fig_file = save_plot(
plt, conf, data_paths, "barchart", metric=str(metric), location=location
)
logger.info(f" Barchart plots created at: {Path(fig_file).parent}")
[docs]
def barchart_all_metrics(df_metrics: pd.DataFrame, conf: dict, data_paths: dict):
"""Create a bar chart comparing datasets for each metric.
Args:
df_metrics: A DataFrame containing the metric values to be plotted, with columns including "primary_location_id", "metric", "dataset", and "value".
conf: A dictionary containing the configuration for creating bar charts, which may include plotting parameters.
data_paths: A dictionary containing the paths to the metric files for each dataset.
Returns:
None
"""
# get location ID
location = df_metrics["primary_location_id"].unique()[0]
location = location.split("-")[-1]
# Get unique datasets and assign colors
datasets = df_metrics["dataset"].unique()
n_datasets = len(datasets)
colors = plt.cm.tab10.colors # categorical colormap
color_map = {d: colors[i % len(colors)] for i, d in enumerate(datasets)}
# Create subplots
metrics = df_metrics["metric"].unique()
n_metrics = len(metrics)
ncols = 5
nrows = -(-n_metrics // ncols) # ceiling division
fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=(12, 8), squeeze=False)
# bar width
bar_width = 0.25 # fraction of total group width
group_center = 0 # center of the single group
# x positions for each bar in the group, centered at 0
x_positions = np.linspace(
group_center - bar_width * (n_datasets - 1) / 2,
group_center + bar_width * (n_datasets - 1) / 2,
n_datasets,
)
# Plot each metric
for ax, metric in zip(axes.ravel(), metrics):
values = df_metrics[df_metrics["metric"] == metric]["value"].values
for xi, val, f in zip(x_positions, values, datasets):
ax.bar(xi, val, width=bar_width, color=color_map[f])
ax.set_title(metric, fontsize=14)
ax.set_xticks([]) # no x-tick labels
ax.set_xlabel("")
ax.tick_params(axis="y", labelsize=14)
ax.set_xlim(-0.5, 0.5) # expand x-axis limits so bars don’t stretch
# Remove unused subplots
for j in range(len(metrics), len(axes.ravel())):
fig.delaxes(axes.ravel()[j])
# Add overall legend above subplots
handles = [plt.Rectangle((0, 0), 1, 1, color=color_map[f]) for f in datasets]
fig.legend(
handles,
datasets,
loc="upper center",
bbox_to_anchor=(0.5, 0.95),
ncol=len(datasets),
fontsize=16,
frameon=False, # no legend box
)
# Add overall title
conf1 = conf["general"]
title = (
f"Metrics for {conf1['location_list'][0]} {conf1['nwm_configuration']} "
f"(T0 = {conf1['forecast_start_date'][0]})"
)
fig.suptitle(title, fontsize=20)
plt.tight_layout(rect=[0, 0, 1, 0.92]) # leave space for legend and suptitle
# save plot to png
fig_file = save_plot(
plt,
conf,
data_paths,
"barchart",
metric="all_metrics",
location=location,
)
logger.info(f" Barchart plots created at: {Path(fig_file).parent}")
[docs]
def create_all_plots(conf: dict, data_paths: dict):
"""Create all plots based on the configuration.
Args:
conf: A dictionary containing the configuration for creating plots, which may include "lead_times", "metric_subset", and plotting parameters.
data_paths: A dictionary containing the paths to the data files.
Returns:
None
"""
plot_types = list(PlotsConfig.model_fields.keys())
plot_functions = {
pt: globals()[f"create_{pt}"]
for pt in plot_types
if f"create_{pt}" in globals()
}
if conf["general"].get("separate_calibrated"):
logger.info(
" Creating plots distinguishing calibrated and regionalized locations."
)
# retrieve calibrated location IDs and add prefix 'usgs-'
calib_params_file = conf["file_paths"]["calib_param_file"]
calib_gages = (
read_data(calib_params_file)["gage_id"].unique().astype(str).tolist()
)
calib_gages = [f"usgs-{gid}" for gid in calib_gages]
conf["general"]["calib_gages"] = calib_gages
else:
conf["general"]["calib_gages"] = []
for plot_type, func in plot_functions.items():
plot_conf = conf["plots"].get(plot_type) or {}
if plot_conf.get("plot", False):
if plot_type != "time_series":
for [str1, str2] in zip(
["metric_subset", "lead_times"], ["metrics", "lead times"]
):
if str1 not in plot_conf or not plot_conf[str1]:
logger.debug(
f"{str1} not specified for {plot_type} plot. Using all available {str2}."
)
func(conf, data_paths)