diff --git a/README.md b/README.md index a5f96f5..f2fb91e 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ The diagram above illustrates the complete workflow from data ingestion to opera ### Key Features -- **Multiple Model Families**: Linear regression baselines and advanced tree-based models (XGBoost, LightGBM, CatBoost) +- **Multiple Model Families**: Linear regression baselines and gradient boosted tree-based models (XGBoost, LightGBM, CatBoost) - **Advanced Feature Engineering**: Time-series features, elevation band aggregation, and glacier-related features - **Ensemble Methods**: Naive averaging, temporal meta-models, and uncertainty quantification - **Comprehensive Evaluation**: Interactive dashboards and extensive metrics diff --git a/dev_tools/visualization/gla_eval.py b/dev_tools/visualization/gla_eval.py index 1c07eb9..da396e6 100644 --- a/dev_tools/visualization/gla_eval.py +++ b/dev_tools/visualization/gla_eval.py @@ -39,12 +39,12 @@ } model_colors = { - "Linear Regression Base": "#787878", # gray - "Base Case Ensemble": "#787878", # orange + "Linear Regression Base": "#505050", # gray + "Base Case Ensemble": "#A15B5B", # orange "SnowMapper Ensemble": "#1F77B4", # orange "MC ALD": "#2CA02C", # green, "Glacier Mapper Ensemble": "#0ECFFF", # red - "Gla MC ALD": "#7BDBAB", # red + "Gla MC ALD": "#7741E3", # red } metric_renamer = { @@ -101,6 +101,15 @@ def plot_monthly_overall( ) plot_df = plot_df.sort_values("month") + # only plot months that exist in the data + existing_months = [ + m for m in plot_df["month"].cat.categories if m in plot_df["month"].values + ] + plot_df = plot_df[plot_df["month"].isin(existing_months)] + + # Remove unused categories from the categorical to prevent empty bars in plot + plot_df["month"] = plot_df["month"].cat.remove_unused_categories() + # order of models in legend plot_df["model"] = pd.Categorical( plot_df["model"], categories=models, ordered=True @@ -158,8 +167,19 @@ def create_monthly_and_overall_performance_plots( models_to_plot: List[str], rename_dict: Dict[str, str], save_dir: str, + agricultural: bool = False, ): """Create and save monthly and overall performance plots for selected models and metrics.""" + if agricultural: + # filter to only include april to september + df_metrics = df_metrics[ + df_metrics["month"].isin( + ["April", "May", "June", "July", "August", "September", "Overall"] + ) + ].copy() + print("Filtered metrics to agricultural months (April to September)") + print("unique months in data:", df_metrics["month"].unique()) + fig, ax = plt.subplots() ax = plot_monthly_overall( df_metrics, @@ -1122,7 +1142,7 @@ def draw_feature_importance(folders_to_check): # run with uv run python dev_tools/visualization/gla_eval.py def main(): - save_dir = "../monthly_forecasting_results/figures/KGZ_Glacier_Eval" + save_dir = "../monthly_forecasting_results/figures/KGZ_Glacier_Eval_2" os.makedirs(save_dir, exist_ok=True) config_plotting() @@ -1133,20 +1153,22 @@ def main(): "../monthly_forecasting_models/SnowMapper_Based/Snow_GBT", "../monthly_forecasting_models/SnowMapper_Based/Snow_GBT_Norm", "../monthly_forecasting_models/SnowMapper_Based/Snow_GBT_LR", - "../monthly_forecasting_models/BaseCase/GBT", + # "../monthly_forecasting_models/BaseCase/GBT", + "../monthly_forecasting_models/BaseCase/LR_Base", ] # Analyze gl_fr feature importance across models print("\n" + "=" * 60) print("FEATURE IMPORTANCE ANALYSIS") print("=" * 60) - draw_feature_importance(folders_to_check=folders_to_check) + # + # draw_feature_importance(folders_to_check=folders_to_check) # Plot top 20 features for all GBT models print("\n" + "=" * 60) print("TOP 20 FEATURES ACROSS GBT MODELS") print("=" * 60) - top_features_save_path = Path(save_dir) / "top_20_features_all_models.png" + """top_features_save_path = Path(save_dir) / "top_20_features_all_models.png" df_top_features, fig_top_features = collect_and_plot_top_features( folders=folders_to_check, top_n=20, save_path=str(top_features_save_path) ) @@ -1163,10 +1185,10 @@ def main(): ) print(f"\nMost common features across all models:") feature_counts = df_top_features["feature"].value_counts().head(15) - print(feature_counts) + print(feature_counts)""" - plt.show() - plt.close("all") + # plt.show() + # plt.close("all") # Continue with existing analysis print("\n" + "=" * 60) @@ -1182,16 +1204,8 @@ def main(): # to int static_df["code"] = static_df["code"].astype(int) - models_to_plot = [ - # "BaseCase_Ensemble", - "SnowMapper_Based_Ensemble", - "Uncertainty_MC_ALD", - "GlacierMapper_Based_Ensemble", - "Uncertainty_Gla_MC_ALD", - ] - rename_dict = { - # "BaseCase_LR_Base": "Linear Regression Base", + "BaseCase_LR_Base": "Linear Regression Base", "BaseCase_Ensemble": "Base Case Ensemble", "SnowMapper_Based_Ensemble": "SnowMapper Ensemble", "Uncertainty_MC_ALD": "MC ALD", @@ -1209,16 +1223,27 @@ def main(): metric_to_plot = "r2" + models_to_plot = [ + "BaseCase_Ensemble", + "BaseCase_LR_BaseSnowMapper_Based_Ensemble", + "Uncertainty_MC_ALD", + "GlacierMapper_Based_Ensemble", + "Uncertainty_Gla_MC_ALD", + ] + create_monthly_and_overall_performance_plots( df_metrics=df_metrics, metric_to_plot=metric_to_plot, models_to_plot=models_to_plot, rename_dict=rename_dict, save_dir=save_dir, + agricultural=True, ) plt.close("all") + sys.exit(0) + combinations = [ ("Glacier Mapper Ensemble", "SnowMapper Ensemble", "Gla_Snow"), ("Glacier Mapper Ensemble", "Base Case Ensemble", "Gla_Base"), diff --git a/docs/Monthly_Workflow_updated.drawio b/docs/Monthly_Workflow_updated.drawio index 60e501f..ca36863 100644 --- a/docs/Monthly_Workflow_updated.drawio +++ b/docs/Monthly_Workflow_updated.drawio @@ -1,74 +1,71 @@ - + - + - - + + - - + + - - - - + - - + + - - - + + + - + - - + + - + - - + + - + - - + + - - + + - + - - + + - + - + - - + + - + @@ -77,94 +74,103 @@ - - + + - - + + - + - - + + - - + + - - + + - - + + - - + + - - + + - - + + + + + - - + + - - + + - - + + - - + + - + - - + + - + - - + + - + - + - - + + - - + + + + + + + + - - + + - + - - + + diff --git a/lt_forecasting/forecast_models/LINEAR_REGRESSION.py b/lt_forecasting/forecast_models/LINEAR_REGRESSION.py index fdfa0c7..08e2caa 100644 --- a/lt_forecasting/forecast_models/LINEAR_REGRESSION.py +++ b/lt_forecasting/forecast_models/LINEAR_REGRESSION.py @@ -163,14 +163,14 @@ def get_highest_corr_features( # Select features based on correlation threshold if len(best_features_above_threshold) == 0: # No features above threshold - use top features regardless - logger.warning( + logger.debug( f"No features with a correlation greater than {threshold} for {this_period}. " f"Selecting top {num_low_corr_features} features regardless of correlation." ) features_week = best_features[:num_low_corr_features] elif len(best_features_above_threshold) < num_low_corr_features: # Some features above threshold but fewer than minimum - backfill with top features - logger.info( + logger.debug( f"Only {len(best_features_above_threshold)} features above threshold {threshold} for {this_period}. " f"Backfilling to {num_low_corr_features} features with top-correlated features." ) @@ -213,16 +213,13 @@ def predict_on_sample( prediction_data = prediction_data.dropna(subset=features, how="any").copy() if len(prediction_data) == 0: - logger.warning("No prediction data available after dropping NaNs.") + logger.debug("No prediction data available after dropping NaNs.") # get the columns with nan in features nan_columns = prediction_data[features].isna().any() missing_features = nan_columns[nan_columns].index.tolist() - logger.warning( - f"Features with NaN values in prediction data: {missing_features}" - ) return [np.nan], [np.nan], np.nan, None if len(calibration_data) < self.general_config["num_features"] * 2: - logger.warning("Not enough calibration data available after dropping NaNs.") + logger.debug("Not enough calibration data available after dropping NaNs.") return [np.nan], [np.nan], np.nan, None X_calibration = calibration_data[features] @@ -234,7 +231,18 @@ def predict_on_sample( if lr_type == "linear": model = LinearRegression() elif lr_type == "ridge": - model = Pipeline([("scaler", StandardScaler()), ("ridge", RidgeCV())]) + model = Pipeline( + [ + ("scaler", StandardScaler()), + ( + "ridge", + RidgeCV( + cv=None, # Uses Efficient Leave-One-Out CV + alphas=[1e-6, 1e-3, 1e-2, 0.1, 1.0, 10.0, 50.0], + ), + ), + ] + ) elif lr_type == "lasso": model = Pipeline( [ @@ -243,7 +251,7 @@ def predict_on_sample( "lasso", LassoCV( alphas=np.logspace(-4, 0, 30), - cv=len(X_calibration), + cv=None, # Uses Efficient Leave-One-Out CV max_iter=100, ), ), # Leave-One-Out CV diff --git a/lt_forecasting/forecast_models/deep_models/uncertainty_mixture.py b/lt_forecasting/forecast_models/deep_models/uncertainty_mixture.py index be3c3f9..a76617f 100644 --- a/lt_forecasting/forecast_models/deep_models/uncertainty_mixture.py +++ b/lt_forecasting/forecast_models/deep_models/uncertainty_mixture.py @@ -62,6 +62,8 @@ def __init__( model_config: Dict[str, Any], feature_config: Dict[str, Any], path_config: Dict[str, Any], + base_predictors: Optional[pd.DataFrame] = None, + base_model_names: Optional[List[str]] = None, ) -> None: """ Initialize the DeepMetaLearner model. @@ -73,6 +75,8 @@ def __init__( model_config: Deep learning model hyperparameters feature_config: Feature engineering configuration path_config: Path configuration for saving/loading + base_predictors: Optional[pd.DataFrame] = None, + base_model_names: Optional[List[str]] = None, """ super().__init__( data=data, @@ -83,6 +87,9 @@ def __init__( path_config=path_config, ) + self.base_predictors = base_predictors + self.base_model_names = base_model_names + # this will be overwritten in calibration / fit self.is_fitted = False @@ -165,11 +172,22 @@ def __preprocess_data__(self): logger.info("Ground Truth created") # 2. Load base predictors using inherited method - logger.info("Loading base predictors") - base_predictors, model_names = self.__load_base_predictors__( - use_mean_pred=self.general_config.get("use_ens_mean", False) - ) - self.model_names = model_names + if self.base_predictors is not None and self.base_model_names is not None: + logger.info("Using provided base predictors and model names") + base_predictors = self.base_predictors + model_names = self.base_model_names + else: + logger.info("Loading base predictors - NOT DATABASE COMPATIBLE YET") + base_predictors, model_names = self.__load_base_predictors__( + use_mean_pred=self.general_config.get("use_ens_mean", False) + ) + self.model_names = model_names + + # sort by date so the newsest date is at the top + base_predictors = base_predictors.sort_values(by=["date"], ascending=[False]) + today = pd.to_datetime(datetime.datetime.now().date()) + base_predictors = base_predictors[base_predictors["date"] <= today] + logger.info(f"Head of Base Predictors data:\n{base_predictors.head(10)}") logger.info("converting base predictors to mm/day") for model in model_names: @@ -288,8 +306,10 @@ def _add_temporal_features(self, data: pd.DataFrame) -> pd.DataFrame: def _m3_to_mm(self, data: pd.DataFrame, col: str = "discharge") -> pd.DataFrame: for code in data.code.unique(): + if code in self.rivers_to_remove: + continue if code not in self.static_data["code"].values: - logger.warning( + logger.debug( f"Code {code} not found in static data. Skipping this code." ) self.rivers_to_remove.append(code) @@ -306,7 +326,7 @@ def _m3_to_mm(self, data: pd.DataFrame, col: str = "discharge") -> pd.DataFrame: def _mm_to_m3(self, data: pd.DataFrame, col: Union[List[str], str]) -> pd.DataFrame: for code in data.code.unique(): if code not in self.static_data["code"].values: - logger.warning( + logger.debug( f"Code {code} not found in static data. Skipping this code." ) self.rivers_to_remove.append(code) @@ -867,8 +887,8 @@ def predict_operational(self, today=None): ) logger.info(f"Total number of unique codes: {len(all_codes)}") - logger.info("Data Tail before filtering:") - logger.info(self.data.tail()) + # logger.info("Data Tail before filtering:") + # logger.info(self.data.tail()) ## Cutoff between cutoff day and today operational_data = self.data[ @@ -889,6 +909,12 @@ def predict_operational(self, today=None): code_data = code_data.sort_values(by="date", ascending=False) # take the most recent date where all features are available + features_with_na_values = code_data[self.features].isna().any(axis=1) + # logg the missing features + if features_with_na_values.any(): + logger.warning( + f"The following features are missing for code {code}: {code_data[self.features].columns[code_data[self.features].isna().any()].tolist()}" + ) code_data = code_data.dropna(subset=self.features) if code_data.empty: logger.warning(f"No complete feature data available for code {code}.") @@ -1069,6 +1095,24 @@ def calibrate_model_and_hindcast(self) -> pd.DataFrame: hindcast = hindcast[pred_cols] hindcast = hindcast.merge(ground_truth, on=["date", "code"], how="left") + # Calculate valid period + if not self.general_config.get("offset"): + self.general_config["offset"] = self.general_config["prediction_horizon"] + shift = ( + self.general_config["offset"] - self.general_config["prediction_horizon"] + ) + valid_from = ( + hindcast["date"] + + datetime.timedelta(days=1) + + datetime.timedelta(days=shift) + ) + valid_to = valid_from + datetime.timedelta( + days=self.general_config["prediction_horizon"] + ) + + hindcast["valid_from"] = valid_from + hindcast["valid_to"] = valid_to + # save model self.save_model() @@ -1076,51 +1120,6 @@ def calibrate_model_and_hindcast(self) -> pd.DataFrame: logger.info("Head of hindcast DataFrame:") logger.info(hindcast.head()) - # quick evaluation of coverage - coverage_90 = np.mean( - (hindcast["Q_obs"] >= hindcast["Q5"]) - & (hindcast["Q_obs"] <= hindcast["Q95"]) - ) - logger.info(f"Approximate 90% coverage: {coverage_90:.2f}") - - coverage_50 = np.mean( - (hindcast["Q_obs"] >= hindcast["Q25"]) - & (hindcast["Q_obs"] <= hindcast["Q75"]) - ) - - logger.info(f"Approximate 50% coverage: {coverage_50:.2f}") - - hindcast["in_90"] = (hindcast["Q_obs"] >= hindcast["Q5"]) & ( - hindcast["Q_obs"] <= hindcast["Q95"] - ) - hindcast["in_50"] = (hindcast["Q_obs"] >= hindcast["Q25"]) & ( - hindcast["Q_obs"] <= hindcast["Q75"] - ) - - # group by code and calculate coverage per code - coverage_by_code = ( - hindcast.groupby("code") - .agg( - coverage_90=("in_90", "mean"), - coverage_50=("in_50", "mean"), - n_samples=("Q_obs", "count"), - ) - .reset_index() - ) - - import matplotlib.pyplot as plt - import seaborn as sns - - # boxplot of coverage_90 and coverage_50 - plt.figure(figsize=(10, 6)) - sns.boxplot(data=coverage_by_code[["coverage_90", "coverage_50"]]) - plt.title("Coverage by Code") - plt.ylabel("Coverage") - plt.ylim(0, 1) - plt.grid(axis="y") - plt.tight_layout() - plt.show() - return hindcast def objective( diff --git a/lt_forecasting/scr/FeatureExtractor.py b/lt_forecasting/scr/FeatureExtractor.py index 8a13c03..16af8ce 100644 --- a/lt_forecasting/scr/FeatureExtractor.py +++ b/lt_forecasting/scr/FeatureExtractor.py @@ -203,9 +203,13 @@ def __init__(self, feature_configs, prediction_horizon=30, offset=None): else: self.offset = offset - assert self.offset >= self.prediction_horizon, ( - "Offset must be greater or equal than prediction horizon" - ) + if self.offset <= self.prediction_horizon: + logger.warning( + "Offset is smaller than or equal to prediction horizon. " + "This may lead to data leakage. But it can be a valid use case." + "Targets are accessed over the period [t+k-H + 1, t+k] where k is the offset and H is the prediction horizon." + "Please ensure this is intended behavior." + ) # Define feature configurations self.feature_configs = feature_configs diff --git a/lt_forecasting/scr/FeatureProcessingArtifacts.py b/lt_forecasting/scr/FeatureProcessingArtifacts.py index 6eecef5..9e6d177 100644 --- a/lt_forecasting/scr/FeatureProcessingArtifacts.py +++ b/lt_forecasting/scr/FeatureProcessingArtifacts.py @@ -547,9 +547,9 @@ def process_training_data( artifacts.static_features = static_features if static_features else [] artifacts.target_col = target - logger.info(f"Numeric features: {len(artifacts.num_features)}") - logger.info(f"Categorical features: {len(artifacts.cat_features)}") - logger.info(f"Static features: {len(artifacts.static_features)}") + logger.debug(f"Numeric features: {len(artifacts.num_features)}") + logger.debug(f"Categorical features: {len(artifacts.cat_features)}") + logger.debug(f"Static features: {len(artifacts.static_features)}") # We scale the static features globally static_scaler = {} @@ -560,7 +560,7 @@ def process_training_data( mean = df_processed[feature].mean() std = df_processed[feature].std() static_scaler[feature] = (mean, std) - logger.info(f"Static feature {feature} - mean: {mean}, std: {std}") + logger.debug(f"Static feature {feature} - mean: {mean}, std: {std}") else: logger.warning(f"Static feature {feature} not found in training data") @@ -765,7 +765,7 @@ def _handle_missing_values_training( df, long_term_mean=artifacts.long_term_means, features=numeric_features ) logger.info("Created long-term mean artifacts") - logger.info( + logger.debug( f"Long Term Stats shape: {artifacts.long_term_stats.shape if hasattr(artifacts.long_term_stats, 'shape') else 'N/A'}" ) @@ -885,7 +885,7 @@ def _normalization_training( features_for_stats.append(target) artifacts.relative_features.append(target) # Track target as relative - logger.info( + logger.debug( f"Using target '{target}' as a relative feature for normalization" ) @@ -894,8 +894,8 @@ def _normalization_training( df, features=features_for_stats ) - logger.info(f"Long-term stats shape: {artifacts.long_term_stats.shape}") - logger.info(f"Long term stats head: {artifacts.long_term_stats.head()}") + logger.debug(f"Long-term stats shape: {artifacts.long_term_stats.shape}") + logger.debug(f"Long term stats head: {artifacts.long_term_stats.head()}") # Apply long-term mean scaling to relative features df = du.apply_long_term_mean_scaling( @@ -905,10 +905,10 @@ def _normalization_training( features_to_scale=features_for_stats, ) - logger.info( + logger.debug( f"Applied long-term mean scaling to features: {features_for_stats}" ) - logger.info( + logger.debug( f"Description of the long-term stats: {df[features_for_stats].describe()}" ) @@ -961,7 +961,7 @@ def _normalization_training( "Use 'global', 'per_basin'." ) - logger.info("Created normalization scaler") + logger.debug("Created normalization scaler") return df, artifacts @@ -1114,7 +1114,7 @@ def _apply_normalization( features_to_scale=relative_features_in_df, ) - logger.info( + logger.debug( f"Applied long-term mean scaling to relative features: {relative_features_in_df}" ) @@ -1219,12 +1219,12 @@ def post_process_predictions( prediction_col=prediction_column, target_col=target, ) - logger.info(f"Applied period-based denormalization to {prediction_column}") + logger.debug(f"Applied period-based denormalization to {prediction_column}") else: # Target uses global/per_basin normalization normalization_process = experiment_config.get("normalization_type", "global") - logger.info(f"Denormalization process: {normalization_process}") + logger.debug(f"Denormalization process: {normalization_process}") if normalization_process == "per_basin": if artifacts.scaler is None: @@ -1245,7 +1245,7 @@ def post_process_predictions( var_used_for_scaling=target, ) - logger.info(f"Applied per-basin denormalization to {prediction_column}") + logger.debug(f"Applied per-basin denormalization to {prediction_column}") elif normalization_process == "global": if artifacts.scaler is None: @@ -1264,7 +1264,7 @@ def post_process_predictions( var_used_for_scaling=target, ) - logger.info(f"Applied global denormalization to {prediction_column}") + logger.debug(f"Applied global denormalization to {prediction_column}") else: raise ValueError( @@ -1272,6 +1272,6 @@ def post_process_predictions( "Use 'global', 'per_basin'." ) - logger.info(f"Applied denormalization to {prediction_column}") + logger.debug(f"Applied denormalization to {prediction_column}") return df_predictions diff --git a/scripts/evaluate_operational.py b/scripts/evaluate_operational.py new file mode 100644 index 0000000..93a202b --- /dev/null +++ b/scripts/evaluate_operational.py @@ -0,0 +1,1921 @@ +import os +import re +import sys +import logging +from pathlib import Path + +logger = logging.getLogger(__name__) +logger.setLevel(logging.DEBUG) + +# Add a stream handler to output logs to the terminal +if not logger.handlers: + handler = logging.StreamHandler(sys.stdout) + handler.setLevel(logging.DEBUG) + formatter = logging.Formatter( + "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + ) + handler.setFormatter(formatter) + logger.addHandler(handler) + + +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import seaborn as sns +from sklearn.metrics import r2_score, mean_squared_error + +# load the .env file +from dotenv import load_dotenv + +load_dotenv(dotenv_path=Path(__file__).parent.parent / ".env") + + +day_of_forecast = { + # "month_0": 10, + "month_1": 1, + "month_2": 1, + "month_3": 1, + "month_4": 1, + "month_5": 1, + "month_6": 1, + # "month_7": 25, + # "month_8": 25, + # "month_9": 25, +} + +kgz_path_config = { + "pred_dir": os.getenv("kgz_path_discharge"), + "obs_file": os.getenv("kgz_path_base_pred"), +} + +taj_path_config = { + "pred_dir": os.getenv("taj_path_base_pred"), + "obs_file": os.getenv("taj_path_discharge"), +} + +output_dir = os.getenv("out_dir_op_lt") + +horizons = list(day_of_forecast.keys()) + +# Models to exclude from ensemble (includes Q_obs variants as safety measure) +models_not_to_ensemble = [ + "MC_ALD", + "SM_GBT", + "SM_GBT_Norm", + "SM_GBT_LR", + "MC_ALD_loc", + "obs", + "Obs", + "OBS", # Exclude any observation-based "models" +] + +models_plot = ["LR_Base", "LR_SM", "Ensemble", "MC_ALD"] + +month_renaming = { + 1: "January", + 2: "February", + 3: "March", + 4: "April", + 5: "May", + 6: "June", + 7: "July", + 8: "August", + 9: "September", + 10: "October", + 11: "November", + 12: "December", +} + +# Add configuration for which issue months to evaluate +# Set to None to evaluate all months, or specify a list like [3, 4, 5] for Mar-May only +issue_months_to_evaluate: list[int] | None = None # Evaluate all months + + +def metric_pipeline( + y_true: np.ndarray, + y_pred: np.ndarray, + quantiles_pred: dict[str, np.ndarray] | None = None, +) -> pd.DataFrame: + """ + Pipeline to evaluate operational metrics for forecasting models. + + Metrics computed: + R2 + RMSE + MSE + MAE + Accuracy + Efficiency + Coverage (if quantiles provided) + """ + + nan_mask = ~np.isnan(y_true) & ~np.isnan(y_pred) + y_true_clean = y_true[nan_mask] + y_pred_clean = y_pred[nan_mask] + + if len(y_true_clean) == 0: + logger.warning( + "No valid data points after removing NaNs for metric computation." + ) + return {} + + # Need at least 2 samples for meaningful metrics + if len(y_true_clean) < 2: + logger.warning( + f"Only {len(y_true_clean)} data point(s) - insufficient for metric computation." + ) + return {} + + for quantile in quantiles_pred.keys() if quantiles_pred is not None else []: + quantiles_pred[quantile] = quantiles_pred[quantile][nan_mask] + + metrics = {} + + # R2 + metrics["R2"] = r2_score(y_true_clean, y_pred_clean) + # RMSE + metrics["RMSE"] = np.sqrt(mean_squared_error(y_true_clean, y_pred_clean)) + # MSE + metrics["MSE"] = mean_squared_error(y_true_clean, y_pred_clean) + # MAE + metrics["MAE"] = np.mean(np.abs(y_true_clean - y_pred_clean)) + + obs_mean = np.mean(y_true_clean) + metrics["nRMSE"] = metrics["RMSE"] / obs_mean if obs_mean != 0 else np.nan + metrics["nMAE"] = metrics["MAE"] / obs_mean if obs_mean != 0 else np.nan + + # Accuracy + sigma_obs = np.std(y_true_clean) + abs_errors = np.abs(y_true_clean - y_pred_clean) + # if abs error smaller than 0.674 * sigma_obs, count as accurate -> 1 else 0 + if sigma_obs > 0: + accurate_preds = abs_errors < (0.674 * sigma_obs) + metrics["Accuracy"] = np.mean(accurate_preds) + else: + metrics["Accuracy"] = np.nan + + std_abs_errors = np.std(abs_errors) + metrics["Efficiency"] = std_abs_errors / sigma_obs if sigma_obs > 0 else np.nan + + # Coverage for quantiles, 90% interval and 50% interval + if quantiles_pred is not None: + + def coverage(y_true, lower_bound, upper_bound): + return np.mean((y_true >= lower_bound) & (y_true <= upper_bound)) + + if "Q5" in quantiles_pred and "Q95" in quantiles_pred: + metrics["Coverage_90"] = coverage( + y_true_clean, quantiles_pred["Q5"], quantiles_pred["Q95"] + ) + if "Q25" in quantiles_pred and "Q75" in quantiles_pred: + metrics["Coverage_50"] = coverage( + y_true_clean, quantiles_pred["Q25"], quantiles_pred["Q75"] + ) + + return pd.DataFrame([metrics]) + + +def compute_long_term_means( + observations_df: pd.DataFrame, + daily_obs_path: str, +) -> tuple[pd.DataFrame, pd.DataFrame]: + """ + Compute long-term means for: + 1. Each calendar month (for target month normalization) + 2. Daily observations (for period-based normalization) + + Args: + observations_df: Monthly aggregated observations with columns ["code", "date", "discharge"] + daily_obs_path: Path to daily observation file + + Returns: + Tuple of: + - monthly_means: DataFrame with ["code", "month", "discharge_ltm"] + - daily_obs: DataFrame with daily observations ["code", "date", "discharge", "day_of_year"] + """ + # Compute monthly long-term means from aggregated monthly data + obs_copy = observations_df.copy() + obs_copy["month"] = obs_copy["date"].dt.month + + monthly_means = ( + obs_copy.groupby(["code", "month"]) + .agg(discharge_ltm=("discharge", "mean")) + .reset_index() + ) + + logger.info( + f"Computed monthly long-term means for {monthly_means['code'].nunique()} stations" + ) + + # Load daily observations for period-based calculations + daily_obs_path = Path(daily_obs_path) + daily_df = pd.read_csv(daily_obs_path) + + # Find date column + date_col = None + for candidate in ["date", "Date", "DATE", "time", "Time"]: + if candidate in daily_df.columns: + date_col = candidate + break + if date_col is None: + date_col = daily_df.columns[0] + + daily_df["date"] = pd.to_datetime(daily_df[date_col]) + + # Find discharge column + discharge_col = None + for candidate in ["discharge", "Discharge", "Q", "q", "runoff", "Runoff"]: + if candidate in daily_df.columns: + discharge_col = candidate + break + if discharge_col is None: + non_date_cols = [ + c for c in daily_df.columns if c not in [date_col, "code", "Code"] + ] + if non_date_cols: + discharge_col = non_date_cols[0] + + # Find code column + code_col = None + for candidate in ["code", "Code", "CODE", "station_id", "basin_id"]: + if candidate in daily_df.columns: + code_col = candidate + break + + daily_df["code"] = daily_df[code_col].astype(int) + daily_df["discharge"] = pd.to_numeric(daily_df[discharge_col], errors="coerce") + daily_df["day_of_year"] = daily_df["date"].dt.dayofyear + daily_df["month"] = daily_df["date"].dt.month + daily_df["day"] = daily_df["date"].dt.day + + logger.info( + f"Loaded {len(daily_df)} daily observations for long-term mean calculation" + ) + + return monthly_means, daily_df + + +def compute_period_long_term_mean( + daily_obs: pd.DataFrame, + code: int, + start_month: int, + start_day: int, + end_month: int, + end_day: int, +) -> float: + """ + Compute the long-term mean for a specific period pattern (e.g., Jul 25 - Aug 25) across all years. + + Args: + daily_obs: DataFrame with daily observations + code: Station code + start_month: Start month of the period (1-12) + start_day: Start day of the period + end_month: End month of the period (1-12) + end_day: End day of the period + + Returns: + Long-term mean discharge for the specified period pattern + """ + station_data = daily_obs[daily_obs["code"] == code] + + if station_data.empty: + return np.nan + + # Handle period within same month or crossing month boundary + if start_month == end_month: + # Same month: filter by month and day range + mask = ( + (station_data["month"] == start_month) + & (station_data["day"] >= start_day) + & (station_data["day"] <= end_day) + ) + elif end_month == start_month + 1 or (start_month == 12 and end_month == 1): + # Crosses one month boundary + mask = ( + (station_data["month"] == start_month) & (station_data["day"] >= start_day) + ) | ((station_data["month"] == end_month) & (station_data["day"] <= end_day)) + else: + # Spans multiple months + if start_month < end_month: + mask = ( + ( + (station_data["month"] == start_month) + & (station_data["day"] >= start_day) + ) + | ( + (station_data["month"] > start_month) + & (station_data["month"] < end_month) + ) + | ( + (station_data["month"] == end_month) + & (station_data["day"] <= end_day) + ) + ) + else: + # Crosses year boundary + mask = ( + ( + (station_data["month"] == start_month) + & (station_data["day"] >= start_day) + ) + | (station_data["month"] > start_month) + | (station_data["month"] < end_month) + | ( + (station_data["month"] == end_month) + & (station_data["day"] <= end_day) + ) + ) + + period_data = station_data[mask] + + if period_data.empty: + return np.nan + + return period_data["discharge"].mean() + + +def _get_period_mask( + station_data: pd.DataFrame, + start_month: int, + start_day: int, + end_month: int, + end_day: int, +) -> pd.Series: + """ + Create a boolean mask for observations falling within a specific period pattern. + + Handles periods within the same month, crossing month boundaries, and crossing year boundaries. + """ + if start_month == end_month: + # Same month: filter by month and day range + mask = ( + (station_data["month"] == start_month) + & (station_data["day"] >= start_day) + & (station_data["day"] <= end_day) + ) + elif end_month == start_month + 1 or (start_month == 12 and end_month == 1): + # Crosses one month boundary + mask = ( + (station_data["month"] == start_month) & (station_data["day"] >= start_day) + ) | ((station_data["month"] == end_month) & (station_data["day"] <= end_day)) + else: + # Spans multiple months + if start_month < end_month: + mask = ( + ( + (station_data["month"] == start_month) + & (station_data["day"] >= start_day) + ) + | ( + (station_data["month"] > start_month) + & (station_data["month"] < end_month) + ) + | ( + (station_data["month"] == end_month) + & (station_data["day"] <= end_day) + ) + ) + else: + # Crosses year boundary + mask = ( + ( + (station_data["month"] == start_month) + & (station_data["day"] >= start_day) + ) + | (station_data["month"] > start_month) + | (station_data["month"] < end_month) + | ( + (station_data["month"] == end_month) + & (station_data["day"] <= end_day) + ) + ) + return mask + + +def precompute_period_ltms_loo( + predictions_raw: pd.DataFrame, + daily_obs: pd.DataFrame, +) -> dict[tuple[int, int, int, int, int, int], float]: + """ + Precompute leave-one-out (LOO) long-term means for period patterns. + + For each unique (code, start_month, start_day, end_month, end_day, exclude_year) combination, + computes the mean discharge excluding the specified year. This prevents data leakage + when evaluating hindcasts. + + Efficiency: Uses vectorized per-year sum/count aggregation, then computes LOO means + in O(1) per lookup by subtracting the excluded year's contribution. + + Args: + predictions_raw: DataFrame with raw predictions including valid_from, valid_to + daily_obs: DataFrame with daily observations + + Returns: + Dictionary mapping (code, start_month, start_day, end_month, end_day, exclude_year) to LOO period LTM + """ + predictions_raw = predictions_raw.copy() + predictions_raw["start_month"] = predictions_raw["valid_from"].dt.month + predictions_raw["start_day"] = predictions_raw["valid_from"].dt.day + predictions_raw["end_month"] = predictions_raw["valid_to"].dt.month + predictions_raw["end_day"] = predictions_raw["valid_to"].dt.day + predictions_raw["target_year"] = predictions_raw["valid_to"].dt.year + + # Get unique period patterns (without year) + unique_periods = predictions_raw[ + ["code", "start_month", "start_day", "end_month", "end_day"] + ].drop_duplicates() + + # Get all years we need to exclude + all_years = predictions_raw["target_year"].unique() + + logger.info( + f"Precomputing LOO period LTMs for {len(unique_periods)} patterns × {len(all_years)} years" + ) + + # Add year column to daily_obs if not present + if "year" not in daily_obs.columns: + daily_obs = daily_obs.copy() + daily_obs["year"] = daily_obs["date"].dt.year + + period_ltm_loo_cache: dict[tuple[int, int, int, int, int, int], float] = {} + + for _, period_row in unique_periods.iterrows(): + code = period_row["code"] + start_month = period_row["start_month"] + start_day = period_row["start_day"] + end_month = period_row["end_month"] + end_day = period_row["end_day"] + + # Filter station data once + station_data = daily_obs[daily_obs["code"] == code] + if station_data.empty: + continue + + # Apply period mask + mask = _get_period_mask( + station_data, start_month, start_day, end_month, end_day + ) + period_data = station_data[mask] + + if period_data.empty: + continue + + # Compute per-year sum and count (vectorized) + yearly_stats = period_data.groupby("year")["discharge"].agg(["sum", "count"]) + total_sum = yearly_stats["sum"].sum() + total_count = yearly_stats["count"].sum() + + # For each year that needs to be excluded, compute LOO mean + for exclude_year in all_years: + cache_key = ( + code, + start_month, + start_day, + end_month, + end_day, + int(exclude_year), + ) + + if exclude_year in yearly_stats.index: + year_sum = yearly_stats.loc[exclude_year, "sum"] + year_count = yearly_stats.loc[exclude_year, "count"] + loo_sum = total_sum - year_sum + loo_count = total_count - year_count + else: + # Year not in data, use full stats + loo_sum = total_sum + loo_count = total_count + + if loo_count > 0: + period_ltm_loo_cache[cache_key] = loo_sum / loo_count + else: + period_ltm_loo_cache[cache_key] = np.nan + + logger.info(f"Precomputed {len(period_ltm_loo_cache)} LOO period LTMs") + + return period_ltm_loo_cache + + +def precompute_monthly_ltms_loo( + observations_df: pd.DataFrame, +) -> dict[tuple[int, int, int], float]: + """ + Precompute leave-one-out monthly long-term means. + + For each (code, month, exclude_year) combination, computes the mean monthly discharge + excluding the specified year. + + Args: + observations_df: Monthly aggregated observations with columns ["code", "date", "discharge"] + + Returns: + Dictionary mapping (code, month, exclude_year) to LOO monthly LTM + """ + obs_copy = observations_df.copy() + obs_copy["month"] = obs_copy["date"].dt.month + obs_copy["year"] = obs_copy["date"].dt.year + + all_years = obs_copy["year"].unique() + + # Compute per-year, per-month sum and count (vectorized) + yearly_monthly_stats = ( + obs_copy.groupby(["code", "month", "year"])["discharge"] + .agg(["sum", "count"]) + .reset_index() + ) + + # Compute total sum and count per (code, month) + totals = ( + yearly_monthly_stats.groupby(["code", "month"]) + .agg(total_sum=("sum", "sum"), total_count=("count", "sum")) + .reset_index() + ) + + # Create lookup for yearly stats + yearly_lookup: dict[tuple[int, int, int], tuple[float, int]] = {} + for _, row in yearly_monthly_stats.iterrows(): + key = (int(row["code"]), int(row["month"]), int(row["year"])) + yearly_lookup[key] = (row["sum"], row["count"]) + + # Create lookup for totals + totals_lookup: dict[tuple[int, int], tuple[float, int]] = {} + for _, row in totals.iterrows(): + key = (int(row["code"]), int(row["month"])) + totals_lookup[key] = (row["total_sum"], row["total_count"]) + + monthly_ltm_loo_cache: dict[tuple[int, int, int], float] = {} + + unique_code_months = totals[["code", "month"]].values + + for code, month in unique_code_months: + code = int(code) + month = int(month) + total_sum, total_count = totals_lookup[(code, month)] + + for exclude_year in all_years: + exclude_year = int(exclude_year) + cache_key = (code, month, exclude_year) + + yearly_key = (code, month, exclude_year) + if yearly_key in yearly_lookup: + year_sum, year_count = yearly_lookup[yearly_key] + loo_sum = total_sum - year_sum + loo_count = total_count - year_count + else: + loo_sum = total_sum + loo_count = total_count + + if loo_count > 0: + monthly_ltm_loo_cache[cache_key] = loo_sum / loo_count + else: + monthly_ltm_loo_cache[cache_key] = np.nan + + logger.info( + f"Precomputed {len(monthly_ltm_loo_cache)} LOO monthly LTMs for " + f"{len(unique_code_months)} (code, month) combinations" + ) + + return monthly_ltm_loo_cache + + +def transform_predictions_with_ratio( + predictions_raw: pd.DataFrame, + observations_df: pd.DataFrame, + daily_obs: pd.DataFrame, +) -> pd.DataFrame: + """ + Transform predictions using the ratio approach with leave-one-out (LOO) normalization. + + Uses LOO long-term means to prevent data leakage during hindcast evaluation: + 1. Calculate ratio = raw_prediction / LOO_long_term_mean(valid_period, exclude_target_year) + 2. Determine target month from issue_date + horizon + 3. Final prediction = LOO_target_month_long_term_mean(exclude_target_year) * ratio + + Args: + predictions_raw: DataFrame with raw predictions including valid_from, valid_to + observations_df: Monthly aggregated observations for computing LOO monthly LTMs + daily_obs: DataFrame with daily observations for period-based LOO LTM calculation + + Returns: + DataFrame with transformed predictions using unbiased LOO normalization + """ + # Precompute all LOO period LTMs (efficient: O(patterns × years) precomputation) + period_ltm_loo_cache = precompute_period_ltms_loo(predictions_raw, daily_obs) + + # Precompute all LOO monthly LTMs + monthly_ltm_loo_cache = precompute_monthly_ltms_loo(observations_df) + + # Vectorized extraction of date components for fast lookup + predictions_raw = predictions_raw.copy() + predictions_raw["start_month"] = predictions_raw["valid_from"].dt.month + predictions_raw["start_day"] = predictions_raw["valid_from"].dt.day + predictions_raw["end_month"] = predictions_raw["valid_to"].dt.month + predictions_raw["end_day"] = predictions_raw["valid_to"].dt.day + predictions_raw["target_year"] = predictions_raw["valid_to"].dt.year + + # Pre-extract arrays for faster iteration + codes = predictions_raw["code"].values + issue_dates = predictions_raw["issue_date"].values + horizons = predictions_raw["horizon"].values + q_preds_raw = predictions_raw["Q_pred"].values + models = predictions_raw["model"].values + start_months = predictions_raw["start_month"].values + start_days = predictions_raw["start_day"].values + end_months = predictions_raw["end_month"].values + end_days = predictions_raw["end_day"].values + target_years = predictions_raw["target_year"].values + + # Pre-extract quantile arrays if they exist + quantile_arrays = {} + for q_col in ["Q5", "Q25", "Q75", "Q95"]: + if q_col in predictions_raw.columns: + quantile_arrays[q_col] = predictions_raw[q_col].values + + result_rows = [] + skipped_period = 0 + skipped_monthly = 0 + + for i in range(len(predictions_raw)): + code = int(codes[i]) + issue_date = pd.Timestamp(issue_dates[i]) + horizon = int(horizons[i]) + q_pred_raw = q_preds_raw[i] + model = models[i] + target_year = int(target_years[i]) + + # Build cache key for LOO period LTM lookup + period_cache_key = ( + code, + int(start_months[i]), + int(start_days[i]), + int(end_months[i]), + int(end_days[i]), + target_year, + ) + period_ltm = period_ltm_loo_cache.get(period_cache_key, np.nan) + + if pd.isna(period_ltm) or period_ltm == 0: + skipped_period += 1 + continue + + # Calculate ratio using LOO period LTM + ratio = q_pred_raw / period_ltm + + # Determine target month from issue_date + horizon + issue_period = issue_date.to_period("M") + target_period = issue_period + horizon + target_month = target_period.month + target_date = target_period.to_timestamp() + + # Get LOO target month long-term mean (exclude target year) + monthly_cache_key = (code, target_month, target_year) + target_ltm = monthly_ltm_loo_cache.get(monthly_cache_key, np.nan) + + if pd.isna(target_ltm) or target_ltm == 0: + skipped_monthly += 1 + continue + + # Calculate final prediction using LOO LTMs + q_pred_transformed = target_ltm * ratio + + result_row = { + "code": code, + "issue_date": issue_date, + "target_date": target_date, + "Q_pred": q_pred_transformed, + "Q_pred_raw": q_pred_raw, + "ratio": ratio, + "period_ltm": period_ltm, + "target_ltm": target_ltm, + "horizon": horizon, + "model": model, + } + + # Transform quantile columns using the same LOO LTMs + for q_col in ["Q5", "Q25", "Q75", "Q95"]: + if q_col in quantile_arrays and pd.notna(quantile_arrays[q_col][i]): + q_ratio = quantile_arrays[q_col][i] / period_ltm + result_row[q_col] = target_ltm * q_ratio + else: + result_row[q_col] = np.nan + + result_rows.append(result_row) + + if skipped_period > 0 or skipped_monthly > 0: + logger.debug( + f"Skipped {skipped_period} rows due to missing period LTM, " + f"{skipped_monthly} rows due to missing monthly LTM" + ) + + if not result_rows: + logger.error("No predictions transformed successfully") + return pd.DataFrame() + + result_df = pd.DataFrame(result_rows) + logger.info( + f"Transformed {len(result_df)} predictions using LOO ratio approach " + f"(unbiased evaluation)" + ) + + return result_df + + +def load_predictions( + base_path: str, + horizons: list[str], + issue_months: list[int] | None = None, +) -> pd.DataFrame: + """ + This function loads all the model predictions from the specified directory. + + Args: + base_path: Base directory containing horizon subdirectories + horizons: List of horizon identifiers (e.g., ["month_0", "month_1", ...]) + issue_months: Optional list of issue months to filter (1-12). If None, all months are kept. + + Returns: + DataFrame with concatenated predictions filtered by issue_months if specified. + """ + base_path = Path(base_path) + all_predictions = [] + + set_codes = set() + for horizon in horizons: + horizon_path = base_path / horizon + if not horizon_path.exists(): + logger.warning(f"Horizon directory not found: {horizon_path}") + continue + + # Extract horizon number from 'month_X' format + horizon_num = int(horizon.split("_")[1]) + forecast_day = day_of_forecast[horizon] + + # Iterate through model subdirectories + for model_dir in horizon_path.iterdir(): + if not model_dir.is_dir(): + continue + + model_name = model_dir.name + hindcast_file = model_dir / f"{model_name}_hindcast.csv" + + if not hindcast_file.exists(): + logger.warning(f"Hindcast file not found: {hindcast_file}") + continue + + try: + df = pd.read_csv(hindcast_file) + except Exception as e: + logger.error(f"Failed to read {hindcast_file}: {e}") + continue + + # Convert dates to datetime + df["date"] = pd.to_datetime(df["date"]) + df["valid_from"] = pd.to_datetime(df["valid_from"]) + df["valid_to"] = pd.to_datetime(df["valid_to"]) + + # Convert code to int + df["code"] = df["code"].astype(int) + + # update the set code to only have codes which are present in all models + set_codes = ( + set_codes.intersection(set(df["code"].unique().tolist())) + if set_codes + else set(df["code"].unique().tolist()) + ) + + # Keep track of unique codes + set_codes.update(df["code"].unique().tolist()) + + # Sort by code and date + df = df.sort_values(["code", "date"]).reset_index(drop=True) + + # Filter to keep only the day of the month matching forecast day + df = df[df["date"].dt.day == forecast_day].copy() + + # Filter by issue months early if specified + if issue_months is not None: + df = df[df["date"].dt.month.isin(issue_months)].copy() + if df.empty: + continue + + # Find all prediction columns (Q_* except quantiles) + # Quantile columns match pattern Q followed by digits only (Q5, Q25, Q50, Q75, Q95, etc.) + quantile_cols = [col for col in df.columns if re.fullmatch(r"Q\d+", col)] + # Exclude Q_obs (observed discharge) - should not be treated as a prediction + excluded_cols = ["Q_obs", "Q_Obs", "Q_OBS"] + q_cols = [ + c + for c in df.columns + if c.startswith("Q_") + and c not in quantile_cols + and c not in excluded_cols + ] + + if not q_cols: + logger.warning(f"No prediction column found in {hindcast_file}") + continue + + # Create a result DataFrame for each Q column (submodel) + for q_col in q_cols: + # Extract submodel name from column (e.g., Q_xgb -> xgb, Q_SM_GBT -> SM_GBT) + submodel_name = q_col[2:] # Remove 'Q_' prefix + + # Create full model name: model_dir/submodel or just submodel if it matches model_dir + if submodel_name == model_name: + full_model_name = model_name + else: + full_model_name = f"{model_name}_{submodel_name}" + + # Restructure the DataFrame - keep valid_from and valid_to for ratio calculation + result_df = pd.DataFrame( + { + "code": df["code"], + "issue_date": df["date"], + "valid_from": df["valid_from"], + "valid_to": df["valid_to"], + "Q_pred": df[q_col], + "horizon": horizon_num, + "model": full_model_name, + } + ) + + # Add quantile columns if they exist (only for the main model, not submodels) + for quantile_col in quantile_cols: + if quantile_col in df.columns and submodel_name == model_name: + result_df[quantile_col] = df[quantile_col].values + else: + result_df[quantile_col] = np.nan + + all_predictions.append(result_df) + + if not all_predictions: + logger.error("No predictions loaded from any horizon/model combination.") + return pd.DataFrame() + + combined_df = pd.concat(all_predictions, ignore_index=True) + logger.info( + f"Loaded {len(combined_df)} prediction records from {len(all_predictions)} files." + ) + + # Filter combined_df to only include codes present in all models + combined_df = combined_df[combined_df["code"].isin(set_codes)].copy() + + return combined_df + + +def load_ground_truth(path_obs: str) -> pd.DataFrame: + """ + This functions loads the daily observed discharge data from the specified path. + 1. convert dates to datetime + 2. convert "code" to int + 3. sort by code and date + 4. aggregate to monthly data (mean, min_obs = 20 days per month) + Returns a pandas DataFrame with columns: ["code", "date", "discharge"] + """ + path_obs = Path(path_obs) + + if not path_obs.exists(): + raise FileNotFoundError(f"Observation file not found: {path_obs}") + + try: + df = pd.read_csv(path_obs) + except Exception as e: + logger.error(f"Failed to read observation file {path_obs}: {e}") + raise + + # Convert date column to datetime + # Try common date column names + date_col = None + for candidate in ["date", "Date", "DATE", "time", "Time"]: + if candidate in df.columns: + date_col = candidate + break + + if date_col is None: + # Assume first column is date + date_col = df.columns[0] + + df["date"] = pd.to_datetime(df[date_col]) + + # filter to only include dates from 2000 onwards + df = df[df["date"].dt.year >= 2000].copy() + + # Find discharge column + discharge_col = None + for candidate in ["discharge", "Discharge", "Q", "q", "runoff", "Runoff"]: + if candidate in df.columns: + discharge_col = candidate + break + + if discharge_col is None: + # Assume second column after date and code is discharge + non_date_cols = [c for c in df.columns if c not in [date_col, "code", "Code"]] + if non_date_cols: + discharge_col = non_date_cols[0] + else: + raise ValueError("Could not identify discharge column in observation file.") + + # Find code column + code_col = None + for candidate in ["code", "Code", "CODE", "station_id", "basin_id"]: + if candidate in df.columns: + code_col = candidate + break + + if code_col is None: + raise ValueError("Could not identify code/station column in observation file.") + + # Convert code to int + df["code"] = df[code_col].astype(int) + df["discharge"] = pd.to_numeric(df[discharge_col], errors="coerce") + + # Sort by code and date + df = df.sort_values(["code", "date"]).reset_index(drop=True) + + # Create year-month column for aggregation + df["year_month"] = df["date"].dt.to_period("M") + + # Aggregate to monthly data with minimum 20 days requirement + monthly_agg = ( + df.groupby(["code", "year_month"]) + .agg(discharge_mean=("discharge", "mean"), day_count=("discharge", "count")) + .reset_index() + ) + + # Filter for months with at least 20 valid days + monthly_agg = monthly_agg[monthly_agg["day_count"] >= 20].copy() + + # Convert period back to timestamp (first day of month) + monthly_agg["date"] = monthly_agg["year_month"].dt.to_timestamp() + + # Select and rename columns + result_df = monthly_agg[["code", "date", "discharge_mean"]].copy() + result_df = result_df.rename(columns={"discharge_mean": "discharge"}) + + logger.info(f"Loaded {len(result_df)} monthly observation records.") + + return result_df + + +def create_ensemble( + predictions_df: pd.DataFrame, + models_to_exclude: list[str] | None = None, + ensemble_name: str = "Ensemble", +) -> pd.DataFrame: + """ + Create ensemble predictions by averaging all models for each code, horizon, and issue_date. + + For each unique combination of (code, horizon, issue_date, target_date), this function + computes the mean of Q_pred and quantile columns across all models not in the exclusion list, + then adds these ensemble predictions as new rows with model="Ensemble". + + Args: + predictions_df: DataFrame with predictions from multiple models. + Expected columns: ["code", "issue_date", "target_date", "Q_pred", + "Q5", "Q25", "Q75", "Q95", "horizon", "model"] + models_to_exclude: List of model names to exclude from ensemble averaging. + If None, all models are included. + ensemble_name: Name to assign to the ensemble model. Defaults to "Ensemble". + + Returns: + DataFrame with original predictions plus ensemble predictions appended. + + Raises: + ValueError: If predictions_df is empty or missing required columns. + """ + if predictions_df.empty: + logger.warning("Empty predictions DataFrame provided for ensemble creation.") + return predictions_df + + required_cols = ["code", "issue_date", "target_date", "Q_pred", "horizon", "model"] + missing_cols = [col for col in required_cols if col not in predictions_df.columns] + if missing_cols: + raise ValueError( + f"Missing required columns for ensemble creation: {missing_cols}" + ) + + if models_to_exclude is None: + models_to_exclude = [] + + # Filter out excluded models for ensemble calculation + ensemble_eligible_df = predictions_df[ + ~predictions_df["model"].isin(models_to_exclude) + ].copy() + + if ensemble_eligible_df.empty: + logger.warning("No models remaining after exclusion for ensemble creation.") + return predictions_df + + # Define columns to average + numeric_cols = ["Q_pred"] + quantile_cols = ["Q5", "Q25", "Q75", "Q95"] + for q_col in quantile_cols: + if q_col in ensemble_eligible_df.columns: + numeric_cols.append(q_col) + + # Group by code, horizon, issue_date, and target_date + grouping_cols = ["code", "horizon", "issue_date", "target_date"] + + # Compute ensemble averages + ensemble_df = ensemble_eligible_df.groupby(grouping_cols, as_index=False).agg( + {col: "mean" for col in numeric_cols} + ) + + # Add model name + ensemble_df["model"] = ensemble_name + + # Ensure all quantile columns exist (fill with NaN if not present in original) + for q_col in quantile_cols: + if q_col not in ensemble_df.columns: + ensemble_df[q_col] = np.nan + + # Reorder columns to match original DataFrame + column_order = [ + "code", + "issue_date", + "target_date", + "Q_pred", + "Q5", + "Q25", + "Q75", + "Q95", + "horizon", + "model", + ] + existing_cols = [c for c in column_order if c in ensemble_df.columns] + ensemble_df = ensemble_df[existing_cols] + + # Concatenate with original predictions + combined_df = pd.concat([predictions_df, ensemble_df], ignore_index=True) + + logger.info( + f"Created {len(ensemble_df)} ensemble predictions from " + f"{ensemble_eligible_df['model'].nunique()} models." + ) + + return combined_df + + +def evaluate( + predictions_df: pd.DataFrame, + observations_df: pd.DataFrame, +) -> pd.DataFrame: + """ + 1. merges predictions with observations based on the code and the target_date + 2. computes the metrics for each model and each basin and each forecast horizon + 3. returns a pandas DataFrame with the metrics with the format + ["code", "horizon", "issue_month", "model", "R2", "RMSE", "MSE", "MAE", "nRMSE", "nMAE", "Accuracy", "Efficiency", "Coverage_90", "Coverage_50"] + """ + if predictions_df.empty or observations_df.empty: + logger.warning("Empty predictions or observations DataFrame provided.") + return pd.DataFrame() + + # Merge predictions with observations + # Match on code and target_date (predictions) with date (observations) + merged_df = predictions_df.merge( + observations_df, + left_on=["code", "target_date"], + right_on=["code", "date"], + how="inner", + suffixes=("_pred", "_obs"), + ) + + if merged_df.empty: + logger.warning( + "No matching records found between predictions and observations." + ) + return pd.DataFrame() + + logger.info(f"Merged {len(merged_df)} prediction-observation pairs.") + + # Extract issue_month from issue_date + merged_df["issue_month"] = merged_df["issue_date"].dt.month + + # Group by code, horizon, issue_month, and model + grouping_cols = ["code", "horizon", "issue_month", "model"] + all_metrics = [] + + for group_keys, group_df in merged_df.groupby(grouping_cols): + code, horizon, issue_month, model = group_keys + + y_true = group_df["discharge"].values + y_pred = group_df["Q_pred"].values + + # Prepare quantiles if available + quantiles_pred = {} + for q_col in ["Q5", "Q25", "Q75", "Q95"]: + if q_col in group_df.columns and group_df[q_col].notna().any(): + quantiles_pred[q_col] = group_df[q_col].values + + if not quantiles_pred: + quantiles_pred = None + + # Compute metrics + metrics_df = metric_pipeline(y_true, y_pred, quantiles_pred) + + if isinstance(metrics_df, pd.DataFrame) and not metrics_df.empty: + metrics_dict = metrics_df.iloc[0].to_dict() + elif isinstance(metrics_df, dict): + metrics_dict = metrics_df + else: + continue + + # Add grouping information + metrics_dict["code"] = code + metrics_dict["horizon"] = horizon + metrics_dict["issue_month"] = issue_month + metrics_dict["model"] = model + + all_metrics.append(metrics_dict) + + if not all_metrics: + logger.warning("No metrics computed for any group.") + return pd.DataFrame() + + result_df = pd.DataFrame(all_metrics) + + # Reorder columns + column_order = [ + "code", + "horizon", + "issue_month", + "model", + "R2", + "RMSE", + "MSE", + "MAE", + "nRMSE", + "nMAE", + "Accuracy", + "Efficiency", + "Coverage_90", + "Coverage_50", + ] + + # Only include columns that exist + existing_cols = [c for c in column_order if c in result_df.columns] + result_df = result_df[existing_cols] + + # rename the issue_month to month name + result_df["issue_month"] = result_df["issue_month"].map(month_renaming) + + logger.info(f"Computed metrics for {len(result_df)} groups.") + + return result_df + + +def draw_overall_plot( + metrics_df: pd.DataFrame, + models: list[str], + metric_name: str, + start_month: str, + output_path: str, +) -> None: + """ + Create a boxplot showing metric performance across forecast horizons for multiple models. + + Filters the metrics_df for the specified models and start_month (issue_month), + then plots the specified metric_name across forecast horizons for each model. + The plot shows transparent boxplots with individual data points in the background. + + Args: + metrics_df: DataFrame containing evaluation metrics with columns + ["code", "horizon", "issue_month", "model", metric_name, ...]. + models: List of model names to include in the plot. + metric_name: Name of the metric column to plot (e.g., "R2", "nRMSE"). + start_month: Issue month name to filter on (e.g., "March"). + output_path: File path where the plot will be saved. + + Raises: + ValueError: If metric_name is not found in metrics_df columns. + ValueError: If no data remains after filtering. + """ + if metric_name not in metrics_df.columns: + raise ValueError( + f"Metric '{metric_name}' not found in DataFrame columns: " + f"{metrics_df.columns.tolist()}" + ) + + # Filter for specified models and start_month + filtered_df = metrics_df[ + (metrics_df["model"].isin(models)) & (metrics_df["issue_month"] == start_month) + ].copy() + + if filtered_df.empty: + raise ValueError( + f"No data found for models {models} and issue_month {start_month}. " + f"Available models: {metrics_df['model'].unique().tolist()}, " + f"Available issue_months: {metrics_df['issue_month'].unique().tolist()}" + ) + + # Set up the plot + fig, ax = plt.subplots(figsize=(12, 6)) + + # Plot individual points in the background using stripplot + sns.stripplot( + data=filtered_df, + x="horizon", + y=metric_name, + hue="model", + hue_order=models, + dodge=True, + alpha=0.4, + size=4, + jitter=0.15, + ax=ax, + legend=False, + ) + + # Plot transparent boxplots on top + sns.boxplot( + data=filtered_df, + x="horizon", + y=metric_name, + hue="model", + hue_order=models, + boxprops={"alpha": 0.6}, + whiskerprops={"alpha": 0.6}, + capprops={"alpha": 0.6}, + medianprops={"alpha": 0.9, "color": "black"}, + flierprops={"alpha": 0}, # Hide outlier points (already shown by stripplot) + ax=ax, + ) + + # Set y-axis limits based on metric + if metric_name == "R2": + ax.set_ylim(-0.5, 1.0) + # set y=0 line + ax.axhline(0, color="black", linestyle="-", linewidth=1) + elif metric_name == "Accuracy": + ax.set_ylim(0.0, 1.0) + + elif metric_name == "Efficiency": + ax.set_ylim(0.0, 2.0) + ax.axhline(0.6, color="black", linestyle="-", linewidth=1) + + # Customize the plot + ax.set_xlabel("Forecast Horizon (months)", fontsize=12) + ax.set_ylabel(metric_name, fontsize=12) + ax.set_title( + f"{metric_name} by Forecast Horizon (Issue Month: {start_month})", + fontsize=14, + ) + + # Adjust legend + ax.legend(title="Model", loc="best", framealpha=0.9) + + # Add grid for better readability + ax.yaxis.grid(True, linestyle="--", alpha=0.3) + ax.set_axisbelow(True) + + # Adjust layout + plt.tight_layout() + + # Save the plot + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + plt.savefig(output_path, dpi=150, bbox_inches="tight") + plt.close(fig) + + logger.info(f"Saved overall plot to {output_path}") + + +def draw_single_basin_plot( + metrics_df: pd.DataFrame, + basin_codes: list[int], + models: list[str], + metric_name: str, + output_path: str, +) -> None: + """ + Create a 4x3 grid figure showing metric vs horizon for specified basins across all months. + + Each subplot represents one issue month, with lines for each basin showing the metric + value across forecast horizons. Different models are shown as separate lines. + + Args: + metrics_df: DataFrame containing evaluation metrics with columns + ["code", "horizon", "issue_month", "model", metric_name, ...]. + basin_codes: List of basin codes to include in the plot (e.g., [16936, 15194, 16100]). + models: List of model names to include in the plot. + metric_name: Name of the metric column to plot (e.g., "R2", "nRMSE"). + output_path: File path where the plot will be saved. + + Raises: + ValueError: If metric_name is not found in metrics_df columns. + ValueError: If no data remains after filtering. + """ + if metric_name not in metrics_df.columns: + raise ValueError( + f"Metric '{metric_name}' not found in DataFrame columns: " + f"{metrics_df.columns.tolist()}" + ) + + # Filter for specified basins and models + filtered_df = metrics_df[ + (metrics_df["code"].isin(basin_codes)) & (metrics_df["model"].isin(models)) + ].copy() + + if filtered_df.empty: + raise ValueError( + f"No data found for basins {basin_codes} and models {models}. " + f"Available codes: {metrics_df['code'].unique().tolist()[:10]}..., " + f"Available models: {metrics_df['model'].unique().tolist()}" + ) + + # Get ordered list of months for consistent subplot arrangement + month_order = [ + "January", + "February", + "March", + "April", + "May", + "June", + "July", + "August", + "September", + "October", + "November", + "December", + ] + available_months = [ + m for m in month_order if m in filtered_df["issue_month"].unique() + ] + + # Set up 4x3 grid figure + fig, axes = plt.subplots(4, 3, figsize=(16, 14), sharey=True) + axes_flat = axes.flatten() + + # Define colors and markers for basins and line styles for models + basin_colors = plt.cm.tab10.colors[: len(basin_codes)] + model_linestyles = ["-", "--", "-.", ":"] + model_markers = ["o", "s", "^", "D"] + + # Create subplot for each month + for idx, month_name in enumerate(month_order): + ax = axes_flat[idx] + + month_data = filtered_df[filtered_df["issue_month"] == month_name] + + if month_data.empty: + ax.set_title(month_name, fontsize=11, fontweight="bold") + ax.text( + 0.5, + 0.5, + "No data", + ha="center", + va="center", + transform=ax.transAxes, + fontsize=10, + color="gray", + ) + ax.set_xlabel("Horizon") + ax.set_ylabel(metric_name if idx % 3 == 0 else "") + continue + + # Plot each combination of basin and model + for basin_idx, basin_code in enumerate(basin_codes): + for model_idx, model_name in enumerate(models): + basin_model_data = month_data[ + (month_data["code"] == basin_code) + & (month_data["model"] == model_name) + ].sort_values("horizon") + + if basin_model_data.empty: + continue + + # Use basin color, model linestyle/marker + color = basin_colors[basin_idx] + linestyle = model_linestyles[model_idx % len(model_linestyles)] + marker = model_markers[model_idx % len(model_markers)] + + label = f"{basin_code} - {model_name}" + ax.plot( + basin_model_data["horizon"], + basin_model_data[metric_name], + color=color, + linestyle=linestyle, + marker=marker, + markersize=5, + linewidth=1.5, + alpha=0.8, + label=label, + ) + + ax.set_title(month_name, fontsize=11, fontweight="bold") + ax.set_xlabel("Horizon") + ax.set_ylabel(metric_name if idx % 3 == 0 else "") + ax.grid(True, linestyle="--", alpha=0.3) + + # Set y-axis limits based on metric + if metric_name == "R2": + ax.set_ylim(-0.5, 1.0) + ax.axhline(0, color="black", linestyle="-", linewidth=0.8) + elif metric_name == "Accuracy": + ax.set_ylim(0.0, 1.0) + elif metric_name == "Efficiency": + ax.set_ylim(0.0, 2.0) + ax.axhline(0.6, color="black", linestyle="-", linewidth=0.8) + + # Create a single legend for all subplots + # Get handles and labels from the last non-empty subplot + handles, labels = [], [] + for ax in axes_flat: + h, l = ax.get_legend_handles_labels() + if h: + handles, labels = h, l + break + + # Remove duplicates while preserving order + unique_labels = [] + unique_handles = [] + for handle, label in zip(handles, labels): + if label not in unique_labels: + unique_labels.append(label) + unique_handles.append(handle) + + # Add legend at the bottom of the figure + fig.legend( + unique_handles, + unique_labels, + loc="lower center", + ncol=min(len(unique_labels), 4), + fontsize=9, + bbox_to_anchor=(0.5, -0.02), + ) + + # Add overall title + fig.suptitle( + f"{metric_name} vs Forecast Horizon for Basins {basin_codes}", + fontsize=14, + fontweight="bold", + y=0.98, + ) + + # Adjust layout + plt.tight_layout(rect=[0, 0.05, 1, 0.96]) + + # Save the plot + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + plt.savefig(output_path, dpi=150, bbox_inches="tight") + plt.close(fig) + + logger.info(f"Saved single basin plot to {output_path}") + + +def compute_seasonal_aggregates( + predictions_df: pd.DataFrame, + observations_df: pd.DataFrame, + target_months: list[int], + issue_months: list[int], +) -> pd.DataFrame: + """ + Compute aggregated (mean) predictions and observations for a seasonal period. + + For each unique combination of (code, model, issue_month, year), this function + computes the mean of predictions and observations over the specified target months + (e.g., April-September). + + Args: + predictions_df: DataFrame with predictions including columns + ["code", "issue_date", "target_date", "Q_pred", "model"]. + observations_df: DataFrame with observations including columns + ["code", "date", "discharge"]. + target_months: List of target months to aggregate (e.g., [4, 5, 6, 7, 8, 9] for Apr-Sep). + issue_months: List of issue months to filter on (e.g., [1, 2, 3] for Jan-Mar). + + Returns: + DataFrame with aggregated seasonal values per year: + ["code", "model", "issue_month", "year", "Q_pred_seasonal", "Q_obs_seasonal"] + """ + if predictions_df.empty or observations_df.empty: + logger.warning("Empty predictions or observations DataFrame provided.") + return pd.DataFrame() + + # Merge predictions with observations + merged_df = predictions_df.merge( + observations_df, + left_on=["code", "target_date"], + right_on=["code", "date"], + how="inner", + suffixes=("_pred", "_obs"), + ) + + if merged_df.empty: + logger.warning( + "No matching records found between predictions and observations." + ) + return pd.DataFrame() + + # Extract issue_month and target_month + merged_df["issue_month"] = merged_df["issue_date"].dt.month + merged_df["target_month"] = merged_df["target_date"].dt.month + merged_df["year"] = merged_df["issue_date"].dt.year + + # Filter for specified issue months and target months + filtered_df = merged_df[ + (merged_df["issue_month"].isin(issue_months)) + & (merged_df["target_month"].isin(target_months)) + ].copy() + + if filtered_df.empty: + logger.warning( + f"No data found for issue months {issue_months} and target months {target_months}." + ) + return pd.DataFrame() + + logger.info( + f"Filtered {len(filtered_df)} records for seasonal aggregation " + f"(issue months: {issue_months}, target months: {target_months})." + ) + + # Group by code, model, issue_month, year and compute mean + grouping_cols = ["code", "model", "issue_month", "year"] + aggregated_df = filtered_df.groupby(grouping_cols, as_index=False).agg( + Q_pred_seasonal=("Q_pred", "mean"), + Q_obs_seasonal=("discharge", "mean"), + n_months=("target_month", "nunique"), + ) + + # Only keep aggregates with all target months present + expected_months = len(target_months) + complete_aggregates = aggregated_df[ + aggregated_df["n_months"] == expected_months + ].copy() + + if complete_aggregates.empty: + logger.warning( + f"No complete seasonal aggregates found with all {expected_months} months." + ) + # Fall back to partial aggregates with at least half the months + min_months = expected_months // 2 + complete_aggregates = aggregated_df[ + aggregated_df["n_months"] >= min_months + ].copy() + + # Map issue_month to month name + complete_aggregates["issue_month_name"] = complete_aggregates["issue_month"].map( + month_renaming + ) + + logger.info(f"Computed {len(complete_aggregates)} seasonal aggregates.") + + return complete_aggregates + + +def draw_seasonal_obs_vs_pred_plot( + seasonal_df: pd.DataFrame, + basin_codes: list[int], + models: list[str], + issue_months: list[str], + output_path: str, +) -> None: + """ + Create a 3x3 grid of obs vs pred scatter plots for seasonal aggregates. + + Rows represent issue months (e.g., January, February, March). + Columns represent basins. + Each subplot shows observed vs predicted seasonal mean with different models as colors. + + Args: + seasonal_df: DataFrame with seasonal aggregates from compute_seasonal_aggregates. + Expected columns: ["code", "model", "issue_month_name", "year", + "Q_pred_seasonal", "Q_obs_seasonal"]. + basin_codes: List of basin codes to plot (one per column). + models: List of model names to include in the plot. + issue_months: List of issue month names (e.g., ["January", "February", "March"]). + output_path: File path where the plot will be saved. + + Raises: + ValueError: If no data is available for plotting. + """ + # Filter for specified basins and models + filtered_df = seasonal_df[ + (seasonal_df["code"].isin(basin_codes)) + & (seasonal_df["model"].isin(models)) + & (seasonal_df["issue_month_name"].isin(issue_months)) + ].copy() + + if filtered_df.empty: + raise ValueError( + f"No data found for basins {basin_codes}, models {models}, " + f"and issue months {issue_months}." + ) + + # Set up 3x3 grid figure (3 issue months x 3 basins) + n_rows = len(issue_months) + n_cols = len(basin_codes) + fig, axes = plt.subplots(n_rows, n_cols, figsize=(4 * n_cols, 4 * n_rows)) + + # Ensure axes is 2D array + if n_rows == 1 and n_cols == 1: + axes = np.array([[axes]]) + elif n_rows == 1: + axes = axes.reshape(1, -1) + elif n_cols == 1: + axes = axes.reshape(-1, 1) + + # Define colors for models + model_colors = dict(zip(models, plt.cm.tab10.colors[: len(models)])) + + # Create subplot for each combination + for row_idx, issue_month in enumerate(issue_months): + for col_idx, basin_code in enumerate(basin_codes): + ax = axes[row_idx, col_idx] + + subplot_data = filtered_df[ + (filtered_df["issue_month_name"] == issue_month) + & (filtered_df["code"] == basin_code) + ] + + if subplot_data.empty: + ax.text( + 0.5, + 0.5, + "No data", + ha="center", + va="center", + transform=ax.transAxes, + fontsize=10, + color="gray", + ) + ax.set_title(f"Basin {basin_code}\n(Issue: {issue_month})", fontsize=10) + continue + + # Plot each model + for model_name in models: + model_data = subplot_data[subplot_data["model"] == model_name] + if model_data.empty: + continue + + ax.scatter( + model_data["Q_obs_seasonal"], + model_data["Q_pred_seasonal"], + c=[model_colors[model_name]], + label=model_name, + alpha=0.7, + s=50, + edgecolors="white", + linewidth=0.5, + ) + + # Add 1:1 line + all_values = pd.concat( + [ + subplot_data["Q_obs_seasonal"], + subplot_data["Q_pred_seasonal"], + ] + ) + min_val = all_values.min() * 0.9 + max_val = all_values.max() * 1.1 + ax.plot( + [min_val, max_val], [min_val, max_val], "k--", linewidth=1, alpha=0.5 + ) + + # Compute R2 for each model and display + r2_texts = [] + for model_name in models: + model_data = subplot_data[subplot_data["model"] == model_name] + if len(model_data) >= 2: # Need at least 2 points for R2 + r2_model = r2_score( + model_data["Q_obs_seasonal"], + model_data["Q_pred_seasonal"], + ) + r2_texts.append(f"{model_name}: {r2_model:.2f}") + + if r2_texts: + r2_display = "R²\n" + "\n".join(r2_texts) + ax.text( + 0.05, + 0.95, + r2_display, + transform=ax.transAxes, + fontsize=8, + verticalalignment="top", + bbox={"boxstyle": "round", "facecolor": "white", "alpha": 0.8}, + ) + + ax.set_xlim(min_val, max_val) + ax.set_ylim(min_val, max_val) + ax.set_aspect("equal", adjustable="box") + ax.grid(True, linestyle="--", alpha=0.3) + + # Set titles and labels + if row_idx == 0: + ax.set_title(f"Basin {basin_code}", fontsize=11, fontweight="bold") + if col_idx == 0: + ax.set_ylabel(f"{issue_month}\nPredicted (m³/s)", fontsize=10) + else: + ax.set_ylabel("") + if row_idx == n_rows - 1: + ax.set_xlabel("Observed (m³/s)", fontsize=10) + else: + ax.set_xlabel("") + + # Create a single legend for all subplots + handles, labels = [], [] + for ax_row in axes: + for ax in ax_row: + h, l = ax.get_legend_handles_labels() + if h: + handles, labels = h, l + break + if handles: + break + + # Remove duplicates while preserving order + unique_labels = [] + unique_handles = [] + for handle, label in zip(handles, labels): + if label not in unique_labels: + unique_labels.append(label) + unique_handles.append(handle) + + # Add legend at the bottom of the figure + fig.legend( + unique_handles, + unique_labels, + loc="lower center", + ncol=len(unique_labels), + fontsize=10, + bbox_to_anchor=(0.5, -0.02), + ) + + # Add overall title + fig.suptitle( + "Seasonal Mean (Apr-Sep) Observed vs Predicted Discharge", + fontsize=14, + fontweight="bold", + y=0.98, + ) + + # Adjust layout + plt.tight_layout(rect=[0, 0.05, 1, 0.96]) + + # Save the plot + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + plt.savefig(output_path, dpi=150, bbox_inches="tight") + plt.close(fig) + + logger.info(f"Saved seasonal obs vs pred plot to {output_path}") + + +def main(): + region = "taj" # or "taj" + if region == "kgz": + path_config = kgz_path_config + region_output_dir = os.path.join(output_dir, "kgz") + elif region == "taj": + path_config = taj_path_config + region_output_dir = os.path.join(output_dir, "taj") + else: + raise ValueError("Invalid region specified. Choose 'kgz' or 'taj'.") + + # Load raw predictions (with valid_from, valid_to) - filtered by issue months + predictions_raw = load_predictions( + base_path=path_config["pred_dir"], + horizons=horizons, + issue_months=issue_months_to_evaluate, + ) + + # Debug: Check which months are in the raw predictions + logger.info( + f"Issue months in raw predictions: " + f"{sorted(predictions_raw['issue_date'].dt.month.unique().tolist())}" + ) + + # Load observations + observations_df = load_ground_truth(path_obs=path_config["obs_file"]) + + # Load daily observations for LOO LTM calculation + _, daily_obs = compute_long_term_means( + observations_df=observations_df, + daily_obs_path=path_config["obs_file"], + ) + + # Transform predictions using LOO ratio approach (unbiased evaluation) + predictions_df = transform_predictions_with_ratio( + predictions_raw=predictions_raw, + observations_df=observations_df, + daily_obs=daily_obs, + ) + + # Create ensemble predictions + predictions_df = create_ensemble( + predictions_df=predictions_df, + models_to_exclude=models_not_to_ensemble, + ensemble_name="Ensemble", + ) + + # Create snow mapper ensemble + models_no_snow = [m for m in predictions_df["model"].unique() if "SM" not in m] + predictions_df = create_ensemble( + predictions_df=predictions_df, + models_to_exclude=models_no_snow, + ensemble_name="SM_Ensemble", + ) + + # Evaluate + metrics_df = evaluate( + predictions_df=predictions_df, + observations_df=observations_df, + ) + print(metrics_df.head()) + + # Save metrics to CSV + metrics_output_path = os.path.join(region_output_dir, "operational_metrics.csv") + os.makedirs(region_output_dir, exist_ok=True) + metrics_df.to_csv(metrics_output_path, index=False) + logger.info(f"Saved operational metrics to {metrics_output_path}") + + # Generate plots for different metrics and start months + metrics_to_plot = ["R2", "Accuracy", "Efficiency"] + + # Use all available issue months from the metrics, or specify a subset + available_months = metrics_df["issue_month"].unique().tolist() + start_months_to_plot = available_months # Plot all available months + + # Or if you want specific months, filter to only those that exist: + # desired_months = ["January", "March", "April", "May", "June", "July"] + # start_months_to_plot = [m for m in desired_months if m in available_months] + + for metric in metrics_to_plot: + if metric not in metrics_df.columns: + logger.warning(f"Metric {metric} not found in results, skipping.") + continue + for start_month in start_months_to_plot: + try: + draw_overall_plot( + metrics_df=metrics_df, + models=models_plot, + metric_name=metric, + start_month=start_month, + output_path=os.path.join( + region_output_dir, f"overall_{metric}_month{start_month}.png" + ), + ) + except ValueError as e: + logger.warning( + f"Could not generate plot for {metric}, month {start_month}: {e}" + ) + + # Generate single basin plots for specified basins + basins_to_plot = [16936, 15194, 16100] + for metric in metrics_to_plot: + if metric not in metrics_df.columns: + logger.warning( + f"Metric {metric} not found in results, skipping single basin plot." + ) + continue + try: + draw_single_basin_plot( + metrics_df=metrics_df, + basin_codes=basins_to_plot, + models=models_plot, + metric_name=metric, + output_path=os.path.join( + region_output_dir, f"single_basin_{metric}_all_months.png" + ), + ) + except ValueError as e: + logger.warning(f"Could not generate single basin plot for {metric}: {e}") + + # Generate seasonal aggregated obs vs pred plot + # Aggregate April-September predictions/observations issued in January, February, March + seasonal_target_months = [4, 5, 6, 7, 8, 9] # April to September + seasonal_issue_months = [1, 2, 3] # January, February, March + + try: + seasonal_aggregates = compute_seasonal_aggregates( + predictions_df=predictions_df, + observations_df=observations_df, + target_months=seasonal_target_months, + issue_months=seasonal_issue_months, + ) + + if not seasonal_aggregates.empty: + draw_seasonal_obs_vs_pred_plot( + seasonal_df=seasonal_aggregates, + basin_codes=basins_to_plot, + models=models_plot, + issue_months=["January", "February", "March"], + output_path=os.path.join( + region_output_dir, "seasonal_obs_vs_pred_apr_sep.png" + ), + ) + + # Save seasonal aggregates to CSV + seasonal_output_path = os.path.join( + region_output_dir, "seasonal_aggregates_apr_sep.csv" + ) + seasonal_aggregates.to_csv(seasonal_output_path, index=False) + logger.info(f"Saved seasonal aggregates to {seasonal_output_path}") + except ValueError as e: + logger.warning(f"Could not generate seasonal obs vs pred plot: {e}") + + +if __name__ == "__main__": + main() diff --git a/scripts/examine_operational_fc.py b/scripts/examine_operational_fc.py new file mode 100644 index 0000000..df99a53 --- /dev/null +++ b/scripts/examine_operational_fc.py @@ -0,0 +1,1859 @@ +import os +import re +import sys +import logging +from pathlib import Path + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + +# Add a stream handler to output logs to the terminal +if not logger.handlers: + handler = logging.StreamHandler(sys.stdout) + handler.setLevel(logging.INFO) + formatter = logging.Formatter( + "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + ) + handler.setFormatter(formatter) + logger.addHandler(handler) + + +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import seaborn as sns +from sklearn.metrics import r2_score, mean_squared_error + +# load the .env file +from dotenv import load_dotenv + +load_dotenv(dotenv_path=Path(__file__).parent.parent / ".env") + + +day_of_forecast = { + "month_0": 15, + "month_1": 25, + "month_2": 25, + "month_3": 25, + "month_4": 25, + "month_5": 25, + # "month_6": 25, + # "month_7": 25, + # "month_8": 25, + # "month_9": 25, +} + +kgz_path_config = { + "pred_dir": os.getenv("kgz_path_discharge"), + "obs_file": os.getenv("kgz_path_base_pred"), +} + +taj_path_config = { + "pred_dir": os.getenv("taj_path_base_pred"), + "obs_file": os.getenv("taj_path_discharge"), +} + +output_dir = os.getenv("out_dir_op_lt") + +horizons = list(day_of_forecast.keys()) + +# Models to exclude from ensemble (includes Q_obs variants as safety measure) +models_not_to_ensemble = [ + "MC_ALD", + "SM_GBT", + "SM_GBT_Norm", + "SM_GBT_LR", + "LR_SM_ROF", + "LR_SM", + "LR_Base", + "MC_ALD_loc", + "obs", + "Obs", + "OBS", # Exclude any observation-based "models" +] + +models_plot = ["LR_Base", "LR_SM", "MC_ALD"] + +month_renaming = { + 1: "January", + 2: "February", + 3: "March", + 4: "April", + 5: "May", + 6: "June", + 7: "July", + 8: "August", + 9: "September", + 10: "October", + 11: "November", + 12: "December", +} + +# Add configuration for which issue months to evaluate +# Set to None to evaluate all months, or specify a list like [3, 4, 5] for Mar-May only +issue_months_to_evaluate: list[int] | None = None # Evaluate all months + +# Flag to control whether to use Q_obs from predictions or load from observations file +# True: Use Q_obs directly from prediction files (if available) +# False: Load observations from file and aggregate to monthly means +use_Q_obs: bool = False + +# Flag to control whether to apply climatology-based ratio correction +# True: Calculate ratio = Q_pred / Q_ltm_period, then Q_pred_corrected = ratio * Q_ltm_monthly +# False: Use Q_pred directly without correction +apply_climatology_correction: bool = True + + +def calculate_metrics(y_true: pd.Series, y_pred: pd.Series) -> dict: + """ + Calculate various performance metrics. + + 1. R2 + 2. nRMSE = RMSE / mean(observed) + 3. MAE + 4. nMAE = MAE / mean(observed) + 5. Accuracy : |y_true - y_pred| <= 0.675 * std of y_true -> 1 else 0 + 6. Efficiency: std (|y_true - y_pred|) / std(y_true) + + Args: + y_true: Series of observed values + y_pred: Series of predicted values + + Returns: + Dictionary with calculated metrics + """ + # Drop NaN values + mask = ~(y_true.isna() | y_pred.isna()) + y_true_clean = y_true[mask] + y_pred_clean = y_pred[mask] + + if len(y_true_clean) < 2: + return { + "r2": np.nan, + "rmse": np.nan, + "nrmse": np.nan, + "mae": np.nan, + "nmae": np.nan, + "accuracy": np.nan, + "efficiency": np.nan, + "n_samples": len(y_true_clean), + } + + # Calculate metrics + r2 = r2_score(y_true_clean, y_pred_clean) + rmse = np.sqrt(mean_squared_error(y_true_clean, y_pred_clean)) + mae = np.mean(np.abs(y_true_clean - y_pred_clean)) + + mean_obs = y_true_clean.mean() + std_obs = y_true_clean.std() + + nrmse = rmse / mean_obs if mean_obs != 0 else np.nan + nmae = mae / mean_obs if mean_obs != 0 else np.nan + + # Accuracy: fraction of predictions within 0.675 * std of observed + threshold = 0.675 * std_obs + accuracy = np.mean(np.abs(y_true_clean - y_pred_clean) <= threshold) + + # Efficiency: std(errors) / std(observed) + errors = np.abs(y_true_clean - y_pred_clean) + efficiency = errors.std() / std_obs if std_obs != 0 else np.nan + + return { + "r2": r2, + "rmse": rmse, + "nrmse": nrmse, + "mae": mae, + "nmae": nmae, + "accuracy": accuracy, + "efficiency": efficiency, + "n_samples": len(y_true_clean), + } + + +def load_observations(obs_file: str) -> pd.DataFrame: + """ + Load observed discharge data from a CSV file. + + Args: + obs_file: Path to the CSV file containing observed discharge data. + + Returns: + DataFrame with columns: date, code, discharge (daily observations) + """ + obs_df = pd.read_csv(obs_file) + obs_df["date"] = pd.to_datetime(obs_df["date"]) + obs_df["code"] = obs_df["code"].astype(int) + + return obs_df + + +def calculate_target(obs: pd.DataFrame) -> pd.DataFrame: + """ + Aggregates daily observations to monthly means for each code. + + Args: + obs: DataFrame with columns: date, code, discharge (daily observations) + + Returns: + DataFrame with columns: code, year, month, Q_obs_monthly (monthly mean discharge) + """ + # Extract year and month from date + obs = obs.copy() + obs["year"] = obs["date"].dt.year + obs["month"] = obs["date"].dt.month + + # Group by code, year, month and calculate mean discharge + monthly_obs = ( + obs.groupby(["code", "year", "month"])["discharge"] + .mean() + .reset_index() + .rename(columns={"discharge": "Q_obs_monthly"}) + ) + + return monthly_obs + + +def calculate_leave_one_out_monthly_mean_fast( + monthly_obs: pd.DataFrame, +) -> pd.DataFrame: + """ + Calculate leave-one-out long-term mean and std for each code and month (FAST vectorized version). + + For each (code, year, month) combination, calculates the mean and std of all other + years' values for that month (excluding the current year). + + Uses vectorized operations: LOO_mean = (total_sum - current_value) / (n - 1) + + Args: + monthly_obs: DataFrame with columns: code, year, month, Q_obs_monthly + + Returns: + DataFrame with columns: code, year, month, Q_obs_monthly, + Q_ltm_monthly (leave-one-out long-term mean), + Q_std_monthly (leave-one-out standard deviation) + """ + df = monthly_obs.copy() + + # Calculate sum, sum of squares, and count for each (code, month) group + agg = ( + df.groupby(["code", "month"])["Q_obs_monthly"] + .agg(["sum", "count", lambda x: (x**2).sum()]) + .reset_index() + ) + agg.columns = ["code", "month", "total_sum", "n_years", "total_sum_sq"] + + # Merge back to get total_sum, total_sum_sq and n_years for each row + df = df.merge(agg, on=["code", "month"], how="left") + + # Calculate leave-one-out mean: (total_sum - current_value) / (n_years - 1) + df["Q_ltm_monthly"] = (df["total_sum"] - df["Q_obs_monthly"]) / (df["n_years"] - 1) + + # Calculate leave-one-out std using the formula: + # LOO_var = (sum_sq - x^2 - (n-1)*LOO_mean^2) / (n-2) for sample variance + # Or more directly: recalculate from remaining values + # Using: var = E[X^2] - E[X]^2, adjusted for leave-one-out + loo_sum = df["total_sum"] - df["Q_obs_monthly"] + loo_sum_sq = df["total_sum_sq"] - df["Q_obs_monthly"] ** 2 + loo_n = df["n_years"] - 1 + + # LOO variance = (sum_sq / n) - mean^2, then multiply by n/(n-1) for sample variance + # = (loo_sum_sq - loo_sum^2/loo_n) / (loo_n - 1) + df["Q_std_monthly"] = np.sqrt((loo_sum_sq - (loo_sum**2 / loo_n)) / (loo_n - 1)) + + # Handle edge case where n_years == 1 (use the value itself, std = 0) + df.loc[df["n_years"] == 1, "Q_ltm_monthly"] = df.loc[ + df["n_years"] == 1, "Q_obs_monthly" + ] + df.loc[df["n_years"] <= 2, "Q_std_monthly"] = ( + np.nan + ) # Can't compute std with < 3 samples for LOO + + # Drop helper columns + df = df.drop(columns=["total_sum", "n_years", "total_sum_sq"]) + + logger.info( + f"Calculated leave-one-out monthly means and stds for {df['code'].nunique()} codes" + ) + + return df + + +def precompute_daily_climatology( + daily_obs: pd.DataFrame, +) -> pd.DataFrame: + """ + Pre-compute daily climatology statistics for fast leave-one-out calculations. + + Creates a lookup table with sum and count of discharge for each (code, month, day) + combination across all years, enabling O(1) leave-one-out calculations. + + Args: + daily_obs: DataFrame with columns: date, code, discharge + + Returns: + DataFrame with columns: code, month, day, total_sum, n_years, yearly data + """ + df = daily_obs.copy() + df["year"] = df["date"].dt.year + df["month"] = df["date"].dt.month + df["day"] = df["date"].dt.day + + # Aggregate by code, month, day, year first (in case of duplicates) + daily_by_year = ( + df.groupby(["code", "year", "month", "day"])["discharge"].mean().reset_index() + ) + + # Create summary stats for each (code, month, day) + daily_stats = ( + daily_by_year.groupby(["code", "month", "day"]) + .agg( + total_sum=("discharge", "sum"), + n_years=("discharge", "count"), + ) + .reset_index() + ) + + # Also create a pivot of year -> discharge for each (code, month, day) + # This allows us to quickly subtract specific years + yearly_pivot = daily_by_year.pivot_table( + index=["code", "month", "day"], + columns="year", + values="discharge", + aggfunc="mean", + ).reset_index() + + # Merge stats with yearly data + result = daily_stats.merge(yearly_pivot, on=["code", "month", "day"], how="left") + + # Ensure month and day are integers + result["month"] = result["month"].astype(int) + result["day"] = result["day"].astype(int) + + logger.info(f"Pre-computed daily climatology for {result['code'].nunique()} codes") + + return result + + +def calculate_period_ltm_fast( + daily_climatology: pd.DataFrame, + predictions_df: pd.DataFrame, + daily_obs: pd.DataFrame, +) -> pd.DataFrame: + """ + Calculate leave-one-out period means and stds for all predictions using vectorized operations. + + For each prediction, calculates the long-term mean and std for the valid_from to valid_to period, + excluding all years that overlap with the prediction (valid_from.year, valid_to.year, target_year). + + Args: + daily_climatology: Pre-computed daily climatology from precompute_daily_climatology + predictions_df: DataFrame with predictions (must have valid_from, valid_to, code, target_year) + daily_obs: Original daily observations (for year lookup) + + Returns: + predictions_df with Q_ltm_period and Q_std_period columns added + """ + df = predictions_df.copy() + + # Extract period bounds + df["from_month"] = df["valid_from"].dt.month + df["from_day"] = df["valid_from"].dt.day + df["to_month"] = df["valid_to"].dt.month + df["to_day"] = df["valid_to"].dt.day + df["from_year"] = df["valid_from"].dt.year + df["to_year"] = df["valid_to"].dt.year + + # Get unique years in the daily observations + year_cols = [ + col for col in daily_climatology.columns if isinstance(col, (int, np.integer)) + ] + + # Helper function to convert month/day to day-of-year (for period filtering) + def to_doy(month: int, day: int) -> int: + days_in_months = [0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31] + return sum(days_in_months[: int(month)]) + min( + int(day), days_in_months[int(month)] + ) + + # Add day-of-year to climatology (vectorized) + clim = daily_climatology.copy() + days_in_months_cumsum = np.array( + [0, 0, 31, 59, 90, 120, 151, 181, 212, 243, 273, 304, 334] + ) + clim["doy"] = days_in_months_cumsum[clim["month"].values] + clim["day"].values + + # Process in batches by unique (code, valid_from, valid_to, target_year) combinations + # This is much faster than row-by-row + unique_periods = df[ + [ + "code", + "valid_from", + "valid_to", + "target_year", + "from_month", + "from_day", + "to_month", + "to_day", + "from_year", + "to_year", + ] + ].drop_duplicates() + + logger.info(f"Calculating period LTM for {len(unique_periods)} unique periods...") + + period_ltm_results = [] + + for _, period in unique_periods.iterrows(): + code = period["code"] + from_doy = to_doy(period["from_month"], period["from_day"]) + to_doy_val = to_doy(period["to_month"], period["to_day"]) + + # Years to exclude (leave-one-out) + exclude_years = set( + [period["from_year"], period["to_year"], period["target_year"]] + ) + + # Filter climatology for this code and period + code_clim = clim[clim["code"] == code].copy() + + # Filter by day-of-year range + if from_doy <= to_doy_val: + period_clim = code_clim[ + (code_clim["doy"] >= from_doy) & (code_clim["doy"] <= to_doy_val) + ] + else: + # Period spans year boundary + period_clim = code_clim[ + (code_clim["doy"] >= from_doy) | (code_clim["doy"] <= to_doy_val) + ] + + if period_clim.empty: + period_ltm_results.append( + { + "code": code, + "valid_from": period["valid_from"], + "valid_to": period["valid_to"], + "target_year": period["target_year"], + "Q_ltm_period": np.nan, + "Q_std_period": np.nan, + } + ) + continue + + # Calculate leave-one-out mean and std + # Sum all year columns except excluded years + valid_year_cols = [y for y in year_cols if y not in exclude_years] + + if not valid_year_cols: + period_ltm_results.append( + { + "code": code, + "valid_from": period["valid_from"], + "valid_to": period["valid_to"], + "target_year": period["target_year"], + "Q_ltm_period": np.nan, + "Q_std_period": np.nan, + } + ) + continue + + # Calculate mean across valid years for each day, then mean across days + # This properly weights each day equally + daily_means = period_clim[valid_year_cols].mean(axis=1, skipna=True) + period_mean = daily_means.mean() + + # Calculate std: for each year, get the period mean, then std across years + # This gives the inter-annual variability of period means + year_period_means = [] + for y in valid_year_cols: + year_data = period_clim[y].dropna() + if len(year_data) > 0: + year_period_means.append(year_data.mean()) + + if len(year_period_means) >= 2: + period_std = np.std(year_period_means, ddof=1) # sample std + else: + period_std = np.nan + + period_ltm_results.append( + { + "code": code, + "valid_from": period["valid_from"], + "valid_to": period["valid_to"], + "target_year": period["target_year"], + "Q_ltm_period": period_mean, + "Q_std_period": period_std, + } + ) + + period_ltm_df = pd.DataFrame(period_ltm_results) + + # Merge back to predictions + df = df.merge( + period_ltm_df, on=["code", "valid_from", "valid_to", "target_year"], how="left" + ) + + # Clean up temporary columns + df = df.drop( + columns=["from_month", "from_day", "to_month", "to_day", "from_year", "to_year"] + ) + + logger.info(f"Calculated period LTM for {len(unique_periods)} unique periods") + + return df + + +def apply_climatology_ratio_correction_fast( + predictions_df: pd.DataFrame, + daily_obs: pd.DataFrame, + monthly_obs_with_ltm: pd.DataFrame, + region: str, +) -> pd.DataFrame: + """ + Apply climatology-based z-score correction to predictions (FAST vectorized version). + + For each prediction: + 1. Calculate leave-one-out long-term mean and std for the valid_from to valid_to period + (excluding ALL years that overlap with valid_from/valid_to to prevent data leakage) + 2. Calculate z-score: z = (Q_pred - Q_ltm_period) / Q_std_period + 3. Clip z-score at 2% and 98% probability levels (z ≈ -2.054 to +2.054) + 4. Final prediction = Q_ltm_monthly + z_clipped * Q_std_monthly + + IMPORTANT: Leave-one-out is applied on a yearly basis to prevent data leakage. + + Args: + predictions_df: DataFrame with predictions (must have valid_from, valid_to, Q_pred) + daily_obs: DataFrame with daily observations (date, code, discharge) + monthly_obs_with_ltm: DataFrame with monthly obs and leave-one-out means/stds + region: Region name ("Kyrgyzstan" or "Tajikistan") + + Returns: + DataFrame with added columns: + - Q_ltm_period: Leave-one-out long-term mean for the valid period + - Q_std_period: Leave-one-out standard deviation for the valid period + - Q_ltm_monthly: Leave-one-out long-term mean for the target month + - Q_std_monthly: Leave-one-out standard deviation for the target month + - z_score: (Q_pred - Q_ltm_period) / Q_std_period + - z_score_clipped: z_score clipped at 2% and 98% levels + - Q_pred_corrected: Q_ltm_monthly + z_score_clipped * Q_std_monthly + """ + from scipy import stats + + df = predictions_df.copy() + + # Calculate target date based on region + if region == "Kyrgyzstan": + # Vectorized: add horizon months + df["target_date"] = df.apply( + lambda row: row["issue_date"] + pd.DateOffset(months=int(row["horizon"])), + axis=1, + ) + elif region == "Tajikistan": + df["target_date"] = df.apply( + lambda row: row["issue_date"] + + pd.DateOffset(months=int(row["horizon"]) - 1), + axis=1, + ) + else: + raise ValueError(f"Unknown region: {region}") + + df["target_year"] = df["target_date"].dt.year + df["target_month"] = df["target_date"].dt.month + df["issue_month"] = df["issue_date"].dt.month + + logger.info( + "Calculating climatology-based z-score corrections (fast vectorized)..." + ) + logger.info(" NOTE: Using leave-one-out on yearly basis to prevent data leakage") + logger.info(" NOTE: Z-scores clipped at 2% and 98% probability levels") + + # Pre-compute daily climatology + logger.info(" Pre-computing daily climatology...") + daily_climatology = precompute_daily_climatology(daily_obs) + + # Calculate period LTM and STD for all predictions + logger.info(" Calculating period leave-one-out means and stds...") + df = calculate_period_ltm_fast(daily_climatology, df, daily_obs) + + # Merge with monthly leave-one-out means and stds + logger.info(" Merging with monthly leave-one-out means and stds...") + df = df.merge( + monthly_obs_with_ltm[ + ["code", "year", "month", "Q_ltm_monthly", "Q_std_monthly"] + ], + left_on=["code", "target_year", "target_month"], + right_on=["code", "year", "month"], + how="left", + ) + df = df.drop(columns=["year", "month"], errors="ignore") + + # Calculate z-score: (Q_pred - Q_ltm_period) / Q_std_period + logger.info(" Calculating z-scores and corrected predictions...") + df["z_score"] = (df["Q_pred"] - df["Q_ltm_period"]) / df["Q_std_period"] + + # Clip z-score at 1% and 99% probability levels + # For a standard normal distribution: + z_lower = stats.norm.ppf(0.01) + z_upper = stats.norm.ppf(0.99) + df["z_score_clipped"] = df["z_score"].clip(lower=z_lower, upper=z_upper) + + # Transform z-score to monthly prediction: Q_ltm_monthly + z_clipped * Q_std_monthly + df["Q_pred_corrected"] = ( + df["Q_ltm_monthly"] + df["z_score_clipped"] * df["Q_std_monthly"] + ) + + # Handle edge cases (missing std, etc.) + df.loc[df["Q_std_period"] <= 0, "z_score"] = np.nan + df.loc[df["Q_std_period"] <= 0, "z_score_clipped"] = np.nan + df.loc[df["Q_std_period"].isna(), "Q_pred_corrected"] = np.nan + df.loc[df["Q_std_monthly"].isna(), "Q_pred_corrected"] = np.nan + + # Ensure Q_pred_corrected is non-negative (discharge can't be negative) + df["Q_pred_corrected"] = df["Q_pred_corrected"].clip(lower=0) + + logger.info(f"Finished processing {len(df)} rows") + logger.info(f" Rows with Q_ltm_period: {df['Q_ltm_period'].notna().sum()}") + logger.info(f" Rows with Q_std_period: {df['Q_std_period'].notna().sum()}") + logger.info(f" Rows with Q_pred_corrected: {df['Q_pred_corrected'].notna().sum()}") + logger.info( + f" Z-score stats: min={df['z_score'].min():.3f}, median={df['z_score'].median():.3f}, max={df['z_score'].max():.3f}" + ) + logger.info( + f" Z-score clipped stats: min={df['z_score_clipped'].min():.3f}, median={df['z_score_clipped'].median():.3f}, max={df['z_score_clipped'].max():.3f}" + ) + logger.info(f" Z-score clip bounds: [{z_lower:.3f}, {z_upper:.3f}] (1% - 99%)") + + return df + + +def create_ensemble(predictions_df: pd.DataFrame) -> pd.DataFrame: + """ + Creates ensemble mean across all models except those in models_not_to_ensemble. + + Args: + predictions_df: DataFrame with predictions containing columns: + - code, issue_date, valid_from, valid_to, Q_pred, Q_obs, horizon, model + + Returns: + DataFrame with ensemble predictions added as a new model "Ensemble" + """ + # Filter models to ensemble + ensemble_models = predictions_df[ + ~predictions_df["model"].isin(models_not_to_ensemble) + ].copy() + + logger.info(f"Creating ensemble from {ensemble_models['model'].nunique()} models") + logger.info( + f" Names of models included: {ensemble_models['model'].unique().tolist()}" + ) + + if ensemble_models.empty: + logger.warning("No models available for ensemble creation") + return predictions_df + + # Group by code, issue_date, horizon and calculate mean prediction + ensemble = ( + ensemble_models.groupby( + ["code", "issue_date", "horizon", "valid_from", "valid_to"] + ) + .agg( + { + "Q_pred": "mean", + "Q_obs": "first", # Q_obs should be the same for all models + } + ) + .reset_index() + ) + + # Add model name + ensemble["model"] = "Ensemble" + + # Add quantile columns as NaN (ensemble doesn't have quantiles) + quantile_cols = [ + col for col in predictions_df.columns if re.fullmatch(r"Q\d+", col) + ] + for col in quantile_cols: + ensemble[col] = np.nan + + # Concatenate with original predictions + combined = pd.concat([predictions_df, ensemble], ignore_index=True) + + logger.info(f"Created ensemble from {ensemble_models['model'].nunique()} models") + + return combined + + +def aggregate( + predictions_df: pd.DataFrame, monthly_obs: pd.DataFrame, region: str +) -> pd.DataFrame: + """ + Merge predictions with monthly aggregated observations based on region-specific logic. + + For Kyrgyzstan: target month = issue_date + horizon (months) + e.g., 15.4 + month_0 = April, month_1 = May, etc. + For Tajikistan: target month = issue_date + horizon - 1 (months) + e.g., 15.4 + month_1 = April, etc. + + Args: + predictions_df: DataFrame with predictions (code, issue_date, horizon, Q_pred, model, etc.) + monthly_obs: DataFrame with monthly observations (code, year, month, Q_obs_monthly) + region: Region name ("Kyrgyzstan" or "Tajikistan") + + Returns: + DataFrame with merged predictions and monthly observations, including target_month + """ + df = predictions_df.copy() + + # Calculate target date based on region + if region == "Kyrgyzstan": + # Target month = issue_date + horizon + df["target_date"] = df.apply( + lambda row: row["issue_date"] + pd.DateOffset(months=int(row["horizon"])), + axis=1, + ) + elif region == "Tajikistan": + # Target month = issue_date + horizon - 1 + df["target_date"] = df.apply( + lambda row: row["issue_date"] + + pd.DateOffset(months=int(row["horizon"]) - 1), + axis=1, + ) + else: + raise ValueError( + f"Unknown region: {region}. Must be 'Kyrgyzstan' or 'Tajikistan'" + ) + + # Extract year and month from target_date + df["target_year"] = df["target_date"].dt.year + df["target_month"] = df["target_date"].dt.month + df["issue_month"] = df["issue_date"].dt.month + + # Merge with monthly observations + merged = df.merge( + monthly_obs, + left_on=["code", "target_year", "target_month"], + right_on=["code", "year", "month"], + how="left", + ) + + # Drop redundant columns + merged = merged.drop(columns=["year", "month"], errors="ignore") + + # Replace Q_obs with Q_obs_monthly if available + merged["Q_obs"] = merged["Q_obs_monthly"].combine_first(merged["Q_obs"]) + merged = merged.drop(columns=["Q_obs_monthly"], errors="ignore") + + logger.info(f"Aggregated predictions with monthly observations for {region}") + logger.info(f" Total records: {len(merged)}") + logger.info(f" Records with Q_obs: {merged['Q_obs'].notna().sum()}") + + return merged + + +def compute_metrics_dataframe(aggregated_df: pd.DataFrame) -> pd.DataFrame: + """ + Compute metrics for all combinations of issue_month, horizon, target_month, model, and code. + + Args: + aggregated_df: DataFrame with aggregated predictions and observations + + Returns: + DataFrame with columns: code, model, issue_month, horizon, target_month, + r2, rmse, nrmse, mae, nmae, accuracy, efficiency, n_samples + """ + metrics_list = [] + + # Group by code, model, issue_month, horizon, target_month + for ( + code, + model, + issue_month, + horizon, + target_month, + ), group in aggregated_df.groupby( + ["code", "model", "issue_month", "horizon", "target_month"] + ): + # Calculate metrics for this group + metrics = calculate_metrics(group["Q_obs"], group["Q_pred"]) + + metrics_list.append( + { + "code": code, + "model": model, + "issue_month": issue_month, + "horizon": horizon, + "target_month": target_month, + **metrics, + } + ) + + metrics_df = pd.DataFrame(metrics_list) + + logger.info(f"Computed metrics for {len(metrics_df)} combinations") + + return metrics_df + + +def load_predictions( + base_path: str, + horizons: list[str], + issue_months: list[int] | None = None, +) -> pd.DataFrame: + """ + This function loads all the model predictions from the specified directory. + + Args: + base_path: Base directory containing horizon subdirectories + horizons: List of horizon identifiers (e.g., ["month_0", "month_1", ...]) + issue_months: Optional list of issue months to filter (1-12). If None, all months are kept. + + Returns: + DataFrame with concatenated predictions filtered by issue_months if specified. + """ + base_path = Path(base_path) + all_predictions = [] + + set_codes = set() + for horizon in horizons: + horizon_path = base_path / horizon + if not horizon_path.exists(): + logger.warning(f"Horizon directory not found: {horizon_path}") + continue + + # Extract horizon number from 'month_X' format + horizon_num = int(horizon.split("_")[1]) + forecast_day = day_of_forecast[horizon] + + # Iterate through model subdirectories + for model_dir in horizon_path.iterdir(): + if not model_dir.is_dir(): + continue + + model_name = model_dir.name + hindcast_file = model_dir / f"{model_name}_hindcast.csv" + + if not hindcast_file.exists(): + logger.warning(f"Hindcast file not found: {hindcast_file}") + continue + + try: + df = pd.read_csv(hindcast_file) + except Exception as e: + logger.error(f"Failed to read {hindcast_file}: {e}") + continue + + if df.empty: + logger.debug(f"Empty dataframe in {hindcast_file}") + continue + + # Convert dates to datetime + df["date"] = pd.to_datetime(df["date"]) + df["valid_from"] = pd.to_datetime(df["valid_from"]) + df["valid_to"] = pd.to_datetime(df["valid_to"]) + + # Convert code to int + df["code"] = df["code"].astype(int) + + logger.debug( + f"Loaded {len(df)} rows from {hindcast_file}, date range: {df['date'].min()} to {df['date'].max()}, days: {df['date'].dt.day.unique()}" + ) + + # Keep track of unique codes across all models + if not set_codes: + set_codes = set(df["code"].unique().tolist()) + else: + set_codes.update(df["code"].unique().tolist()) + + # Sort by code and date + df = df.sort_values(["code", "date"]).reset_index(drop=True) + + # Filter to keep only the day of the month matching forecast day + df = df[df["date"].dt.day == forecast_day].copy() + + if df.empty: + logger.debug( + f"No data after filtering by forecast day {forecast_day} in {hindcast_file}" + ) + continue + + # Filter by issue months early if specified + if issue_months is not None: + df = df[df["date"].dt.month.isin(issue_months)].copy() + if df.empty: + logger.debug( + f"No data after filtering by issue months in {hindcast_file}" + ) + continue + + # Find all prediction columns (Q_* except quantiles) + # Quantile columns match pattern Q followed by digits only (Q5, Q25, Q50, Q75, Q95, etc.) + quantile_cols = [col for col in df.columns if re.fullmatch(r"Q\d+", col)] + # Exclude Q_obs (observed discharge) - should not be treated as a prediction + excluded_cols = ["Q_obs", "Q_Obs", "Q_OBS"] + q_cols = [ + c + for c in df.columns + if c.startswith("Q_") + and c not in quantile_cols + and c not in excluded_cols + ] + + # Find the Q_obs column if it exists + q_obs_col = next( + (col for col in df.columns if col.lower() == "q_obs"), None + ) + + if not q_cols: + logger.warning(f"No prediction column found in {hindcast_file}") + continue + + # Create a result DataFrame for each Q column (submodel) + for q_col in q_cols: + # Extract submodel name from column (e.g., Q_xgb -> xgb, Q_SM_GBT -> SM_GBT) + submodel_name = q_col[2:] # Remove 'Q_' prefix + + # Create full model name: model_dir/submodel or just submodel if it matches model_dir + if submodel_name == model_name: + full_model_name = model_name + else: + full_model_name = f"{model_name}_{submodel_name}" + + # Restructure the DataFrame - keep valid_from and valid_to for ratio calculation + result_df = pd.DataFrame( + { + "code": df["code"], + "issue_date": df["date"], + "valid_from": df["valid_from"], + "valid_to": df["valid_to"], + "Q_pred": df[q_col], + "Q_obs": df[q_obs_col] if q_obs_col else np.nan, + "horizon": horizon_num, + "model": full_model_name, + } + ) + + # Add quantile columns if they exist (only for the main model, not submodels) + for quantile_col in quantile_cols: + if quantile_col in df.columns and submodel_name == model_name: + result_df[quantile_col] = df[quantile_col].values + else: + result_df[quantile_col] = np.nan + + all_predictions.append(result_df) + + if not all_predictions: + logger.error("No predictions loaded from any horizon/model combination.") + return pd.DataFrame() + + combined_df = pd.concat(all_predictions, ignore_index=True) + logger.info( + f"Loaded {len(combined_df)} prediction records from {len(all_predictions)} files." + ) + logger.info(f"Unique codes found: {sorted(set_codes)}") + logger.info( + f"Unique models found: {sorted(combined_df['model'].unique().tolist())}" + ) + + return combined_df + + +def plot_metric_by_horizon( + metrics_df: pd.DataFrame, + horizon: int, + models: list[str], + metric: str = "r2", + y_limits: tuple[float, float] | None = None, + output_path: str | Path | None = None, +) -> plt.Figure: + """ + Plot metric distribution by target month for a given forecast horizon. + + Creates a boxplot showing metric distribution for all models across all target months + for a specific forecast horizon. X-axis = target month (1-12), Y-axis = metric. + If all points for a model/month fall below the y-axis limit, a visual marker + showing the percentage of samples below threshold is displayed. + + Args: + metrics_df: DataFrame with metrics containing columns: + - code, model, issue_month, horizon, target_month, r2, rmse, nrmse, etc. + horizon: The forecast horizon to filter by (e.g., 1, 2, 3) + models: List of model names to include in the plot + metric: Metric to plot (default: "r2"). Options: "r2", "rmse", "nrmse", "mae", "nmae", "accuracy", "efficiency" + y_limits: Optional tuple (y_min, y_max) for y-axis limits. If None, uses metric defaults. + output_path: Optional path to save the figure + + Returns: + matplotlib Figure object + """ + # Filter for the specified horizon + df = metrics_df[metrics_df["horizon"] == horizon].copy() + + if df.empty: + logger.warning(f"No data found for horizon {horizon}") + return plt.figure() + + # Filter for specified models + df = df[df["model"].isin(models)].copy() + + if df.empty: + logger.warning(f"No data found for models {models} at horizon {horizon}") + return plt.figure() + + # Set default y-axis limits based on metric + if y_limits is None: + if metric == "r2": + y_limits = (-1.0, 1.0) + elif metric == "accuracy": + y_limits = (0.0, 1.0) + elif metric in ["nrmse", "nmae", "efficiency"]: + y_limits = (0.0, 2.0) + else: + # For rmse, mae - use data range + y_limits = (df[metric].min() * 0.9, df[metric].max() * 1.1) + + y_min, y_max = y_limits + + # Create a single plot + fig, ax = plt.subplots(1, 1, figsize=(16, 8)) + + # Get unique target months and sort them + target_months = sorted(df["target_month"].unique()) + + # Define color palette for models + model_colors = sns.color_palette("husl", len(models)) + color_map = dict(zip(models, model_colors)) + + # Prepare data for grouped boxplot + positions = [] + data_to_plot = [] + colors_to_use = [] + below_threshold_markers = [] # Store info for markers + + box_width = 0.8 / len(models) + + # Debug: track what we find for each model/month + missing_combinations = [] + + for m_idx, month in enumerate(target_months): + for model_idx, model in enumerate(models): + model_month_data = df[ + (df["target_month"] == month) & (df["model"] == model) + ][metric] + + position = m_idx + (model_idx - len(models) / 2 + 0.5) * box_width + + if len(model_month_data) == 0: + missing_combinations.append(f"{model} @ month {month}") + continue + + if len(model_month_data) > 0: + # Filter out NaN values - they can occur when metrics calculation had insufficient data + values = model_month_data.dropna().values + + if len(values) == 0: + missing_combinations.append(f"{model} @ month {month} (all NaN)") + continue + + n_total = len(values) + n_below = np.sum(values < y_min) + n_above = np.sum(values > y_max) + + # Clip values to y_limits for plotting, but track outliers + values_clipped = np.clip(values, y_min, y_max) + + # If ALL values are below threshold, don't plot box, just marker + if n_below == n_total: + below_threshold_markers.append( + { + "position": position, + "pct_below": 100.0, + "median": np.median(values), + "color": color_map[model], + "model": model, + "month": month, + "direction": "below", + } + ) + elif n_above == n_total: + below_threshold_markers.append( + { + "position": position, + "pct_above": 100.0, + "median": np.median(values), + "color": color_map[model], + "model": model, + "month": month, + "direction": "above", + } + ) + else: + positions.append(position) + data_to_plot.append(values) + colors_to_use.append(color_map[model]) + + # Track partial below/above threshold + if n_below > 0: + below_threshold_markers.append( + { + "position": position, + "pct_below": (n_below / n_total) * 100, + "color": color_map[model], + "model": model, + "month": month, + "direction": "below", + "partial": True, + } + ) + if n_above > 0: + below_threshold_markers.append( + { + "position": position, + "pct_above": (n_above / n_total) * 100, + "color": color_map[model], + "model": model, + "month": month, + "direction": "above", + "partial": True, + } + ) + + # Create boxplots if there's data + if data_to_plot: + bp = ax.boxplot( + data_to_plot, + positions=positions, + widths=box_width * 0.85, + patch_artist=True, + showfliers=True, + flierprops=dict(marker="o", markersize=3, alpha=0.5), + ) + + # Color the boxes + for patch, color in zip(bp["boxes"], colors_to_use): + patch.set_facecolor(color) + patch.set_alpha(0.7) + else: + logger.warning(f"Horizon {horizon}: No data to plot in boxplot!") + + # Add markers for values below/above threshold + for marker_info in below_threshold_markers: + pos = marker_info["position"] + color = marker_info["color"] + + if marker_info.get("partial", False): + # Partial: some values outside range - add small annotation + if marker_info["direction"] == "below": + pct = marker_info["pct_below"] + ax.annotate( + f"↓{pct:.0f}%", + xy=(pos, y_min), + xytext=(pos, y_min + 0.02 * (y_max - y_min)), + fontsize=7, + color=color, + ha="center", + va="bottom", + fontweight="bold", + ) + else: + pct = marker_info["pct_above"] + ax.annotate( + f"↑{pct:.0f}%", + xy=(pos, y_max), + xytext=(pos, y_max - 0.02 * (y_max - y_min)), + fontsize=7, + color=color, + ha="center", + va="top", + fontweight="bold", + ) + else: + # All values outside range - prominent marker + if marker_info["direction"] == "below": + median = marker_info["median"] + ax.scatter( + [pos], + [y_min + 0.03 * (y_max - y_min)], + marker="v", + s=80, + color=color, + edgecolors="black", + linewidths=0.5, + zorder=10, + ) + ax.annotate( + f"100%↓\n(med={median:.2f})", + xy=(pos, y_min + 0.03 * (y_max - y_min)), + xytext=(pos, y_min + 0.12 * (y_max - y_min)), + fontsize=8, + color=color, + ha="center", + va="bottom", + fontweight="bold", + bbox=dict( + boxstyle="round,pad=0.2", + facecolor="white", + alpha=0.8, + edgecolor=color, + ), + ) + else: + median = marker_info["median"] + ax.scatter( + [pos], + [y_max - 0.03 * (y_max - y_min)], + marker="^", + s=80, + color=color, + edgecolors="black", + linewidths=0.5, + zorder=10, + ) + ax.annotate( + f"100%↑\n(med={median:.2f})", + xy=(pos, y_max - 0.03 * (y_max - y_min)), + xytext=(pos, y_max - 0.12 * (y_max - y_min)), + fontsize=8, + color=color, + ha="center", + va="top", + fontweight="bold", + bbox=dict( + boxstyle="round,pad=0.2", + facecolor="white", + alpha=0.8, + edgecolor=color, + ), + ) + + # Customize the plot + ax.set_xlabel("Target Month", fontsize=12, fontweight="bold") + ax.set_ylabel( + metric.upper() if len(metric) <= 4 else metric.capitalize(), + fontsize=12, + fontweight="bold", + ) + ax.set_title( + f"Forecast Horizon {horizon} - {metric.upper()} by Target Month", + fontsize=14, + fontweight="bold", + ) + + # Set x-axis ticks and labels + ax.set_xticks(range(len(target_months))) + ax.set_xticklabels( + [month_renaming[m][:3] for m in target_months], rotation=45, ha="right" + ) + + # Set y-axis limits + ax.set_ylim(y_min, y_max) + + # Add horizontal reference line + if metric == "r2": + ax.axhline(y=0, color="gray", linestyle="--", linewidth=0.8, alpha=0.7) + elif metric in ["nrmse", "nmae", "efficiency"]: + ax.axhline(y=1.0, color="gray", linestyle="--", linewidth=0.8, alpha=0.7) + elif metric == "accuracy": + ax.axhline(y=0.5, color="gray", linestyle="--", linewidth=0.8, alpha=0.7) + + ax.grid(axis="y", alpha=0.3) + + # Create legend + legend_handles = [ + plt.Rectangle((0, 0), 1, 1, fc=color_map[model], alpha=0.7) for model in models + ] + ax.legend(legend_handles, models, loc="upper right", framealpha=0.9, fontsize=10) + + # Log missing combinations + if missing_combinations: + logger.warning( + f"Horizon {horizon}: Missing data for {len(missing_combinations)} combinations: {missing_combinations[:10]}..." + ) + + # Log below/above threshold statistics + n_full_below = len( + [ + m + for m in below_threshold_markers + if not m.get("partial", False) and m["direction"] == "below" + ] + ) + n_full_above = len( + [ + m + for m in below_threshold_markers + if not m.get("partial", False) and m["direction"] == "above" + ] + ) + if n_full_below > 0 or n_full_above > 0: + logger.info( + f"Horizon {horizon}: {n_full_below} model/month combos have 100% values below y_min={y_min}, {n_full_above} have 100% above y_max={y_max}" + ) + # Print details of which ones + for m in below_threshold_markers: + if not m.get("partial", False): + logger.info( + f" -> {m['model']} @ month {m['month']}: direction={m['direction']}, median={m.get('median', 'N/A')}" + ) + + plt.tight_layout() + + # Save if output path provided + if output_path: + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(output_path, dpi=150, bbox_inches="tight") + logger.info( + f"Saved {metric} by target month plot for horizon {horizon} to {output_path}" + ) + + return fig + + +def plot_metric_by_lead_time( + metrics_df: pd.DataFrame, + models: list[str], + metric: str = "r2", + y_limits: tuple[float, float] | None = None, + output_path: str | Path | None = None, + aggregate_by: str = "target_month", +) -> plt.Figure: + """ + Plot metric vs forecasting lead time (horizon) for all months. + + Creates a line plot showing metric values for each model across different + forecast horizons, with separate lines for each month (target or issue). + X-axis = forecast lead time (horizon), Y-axis = metric. + + Args: + metrics_df: DataFrame with metrics containing columns: + - code, model, issue_month, horizon, target_month, r2, rmse, nrmse, etc. + models: List of model names to include in the plot + metric: Metric to plot (default: "r2"). Options: "r2", "rmse", "nrmse", "mae", "nmae", "accuracy", "efficiency" + y_limits: Optional tuple (y_min, y_max) for y-axis limits. If None, uses metric defaults. + output_path: Optional path to save the figure + aggregate_by: Which month to use for grouping - "target_month" or "issue_month" + + Returns: + matplotlib Figure object + """ + # Filter for specified models + df = metrics_df[metrics_df["model"].isin(models)].copy() + + if df.empty: + logger.warning(f"No data found for models {models}") + return plt.figure() + + # Set default y-axis limits based on metric + if y_limits is None: + if metric == "r2": + y_limits = (-0.5, 1.0) + elif metric == "accuracy": + y_limits = (0.0, 1.0) + elif metric in ["nrmse", "nmae", "efficiency"]: + y_limits = (0.0, 2.0) + else: + y_limits = (df[metric].min() * 0.9, df[metric].max() * 1.1) + + y_min, y_max = y_limits + + # Get unique horizons and months + horizons = sorted(df["horizon"].unique()) + months = sorted(df[aggregate_by].unique()) + + # Create subplots - one per model + n_models = len(models) + n_cols = min(3, n_models) + n_rows = (n_models + n_cols - 1) // n_cols + + fig, axes = plt.subplots( + n_rows, n_cols, figsize=(7 * n_cols, 5 * n_rows), squeeze=False + ) + axes_flat = axes.flatten() + + # Define color palette for months (12 distinct colors) + month_colors = sns.color_palette("husl", 12) + color_map = {m: month_colors[m - 1] for m in range(1, 13)} + + for model_idx, model in enumerate(models): + ax = axes_flat[model_idx] + model_data = df[df["model"] == model] + + if model_data.empty: + ax.set_title(f"{model}\n(No data)", fontsize=12, fontweight="bold") + ax.set_visible(True) + continue + + # Aggregate metrics by horizon and month (mean across codes) + agg_data = ( + model_data.groupby(["horizon", aggregate_by])[metric] + .agg(["mean", "std", "count"]) + .reset_index() + ) + agg_data.columns = ["horizon", "month", "mean", "std", "count"] + + # Plot line for each month + for month in months: + month_data = agg_data[agg_data["month"] == month].sort_values("horizon") + + if month_data.empty: + continue + + color = color_map.get(month, "gray") + month_name = month_renaming.get(month, str(month))[:3] + + # Plot mean line with error bands (± 1 std / sqrt(n) for SEM) + ax.plot( + month_data["horizon"], + month_data["mean"], + marker="o", + markersize=6, + linewidth=2, + color=color, + label=month_name, + alpha=0.8, + ) + + # Add shaded error band (standard error of the mean) + sem = month_data["std"] / np.sqrt(month_data["count"]) + ax.fill_between( + month_data["horizon"], + month_data["mean"] - sem, + month_data["mean"] + sem, + color=color, + alpha=0.15, + ) + + # Customize subplot + ax.set_xlabel("Forecast Lead Time (months)", fontsize=11, fontweight="bold") + ax.set_ylabel( + metric.upper() if len(metric) <= 4 else metric.capitalize(), + fontsize=11, + fontweight="bold", + ) + ax.set_title(f"{model}", fontsize=12, fontweight="bold") + ax.set_xticks(horizons) + ax.set_xticklabels([str(h) for h in horizons]) + ax.set_ylim(y_min, y_max) + + # Add horizontal reference line + if metric == "r2": + ax.axhline(y=0, color="gray", linestyle="--", linewidth=0.8, alpha=0.7) + elif metric in ["nrmse", "nmae", "efficiency"]: + ax.axhline(y=1.0, color="gray", linestyle="--", linewidth=0.8, alpha=0.7) + elif metric == "accuracy": + ax.axhline(y=0.5, color="gray", linestyle="--", linewidth=0.8, alpha=0.7) + + ax.grid(axis="both", alpha=0.3) + ax.legend( + loc="best", + fontsize=8, + ncol=3, + framealpha=0.9, + title=aggregate_by.replace("_", " ").title(), + ) + + # Hide unused subplots + for idx in range(n_models, len(axes_flat)): + axes_flat[idx].set_visible(False) + + # Add overall title + fig.suptitle( + f"{metric.upper()} vs Forecast Lead Time (by {aggregate_by.replace('_', ' ').title()})", + fontsize=14, + fontweight="bold", + y=1.02, + ) + + plt.tight_layout() + + # Save if output path provided + if output_path: + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(output_path, dpi=150, bbox_inches="tight") + logger.info(f"Saved {metric} by lead time plot to {output_path}") + + return fig + + +def plot_metric_by_lead_time_combined( + metrics_df: pd.DataFrame, + models: list[str], + metric: str = "r2", + y_limits: tuple[float, float] | None = None, + output_path: str | Path | None = None, + aggregate_by: str = "target_month", +) -> plt.Figure: + """ + Plot metric vs forecasting lead time (horizon) for all months in a single combined plot. + + Creates a single plot showing metric values across different forecast horizons, + with separate grouped bars/lines for each model and month combination. + X-axis = forecast lead time (horizon), Y-axis = metric. + + Args: + metrics_df: DataFrame with metrics containing columns: + - code, model, issue_month, horizon, target_month, r2, rmse, nrmse, etc. + models: List of model names to include in the plot + metric: Metric to plot (default: "r2") + y_limits: Optional tuple (y_min, y_max) for y-axis limits + output_path: Optional path to save the figure + aggregate_by: Which month to use for grouping - "target_month" or "issue_month" + + Returns: + matplotlib Figure object + """ + # Filter for specified models + df = metrics_df[metrics_df["model"].isin(models)].copy() + + if df.empty: + logger.warning(f"No data found for models {models}") + return plt.figure() + + # Set default y-axis limits based on metric + if y_limits is None: + if metric == "r2": + y_limits = (-0.5, 1.0) + elif metric == "accuracy": + y_limits = (0.0, 1.0) + elif metric in ["nrmse", "nmae", "efficiency"]: + y_limits = (0.0, 2.0) + else: + y_limits = (df[metric].min() * 0.9, df[metric].max() * 1.1) + + y_min, y_max = y_limits + + # Aggregate metrics by model, horizon, and month (mean across codes) + agg_data = ( + df.groupby(["model", "horizon", aggregate_by])[metric] + .agg(["mean", "std", "count"]) + .reset_index() + ) + agg_data.columns = ["model", "horizon", "month", "mean", "std", "count"] + + # Calculate overall mean across all months for each model and horizon + overall_mean = ( + agg_data.groupby(["model", "horizon"]) + .agg( + mean=("mean", "mean"), + std=("mean", "std"), # std of monthly means + count=("count", "sum"), + ) + .reset_index() + ) + + # Get unique horizons + horizons = sorted(df["horizon"].unique()) + + # Create a single plot + fig, ax = plt.subplots(figsize=(12, 7)) + + # Define color palette for models + model_colors = sns.color_palette("husl", len(models)) + color_map = dict(zip(models, model_colors)) + + # Define line styles for different visualization + line_styles = ["-", "--", "-.", ":"] + + # Plot line for each model (showing mean across all months with individual month points) + for model_idx, model in enumerate(models): + model_overall = overall_mean[overall_mean["model"] == model].sort_values( + "horizon" + ) + model_monthly = agg_data[agg_data["model"] == model] + + if model_overall.empty: + continue + + color = color_map[model] + + # Plot mean line across all months + ax.plot( + model_overall["horizon"], + model_overall["mean"], + marker="o", + markersize=10, + linewidth=3, + color=color, + label=model, + alpha=0.9, + linestyle=line_styles[model_idx % len(line_styles)], + ) + + # Add shaded error band (std across months) + ax.fill_between( + model_overall["horizon"], + model_overall["mean"] - model_overall["std"], + model_overall["mean"] + model_overall["std"], + color=color, + alpha=0.15, + ) + + # Plot individual month points as scatter (smaller, semi-transparent) + for horizon in horizons: + month_points = model_monthly[model_monthly["horizon"] == horizon] + if not month_points.empty: + # Jitter x positions slightly for visibility + jitter = (model_idx - len(models) / 2) * 0.05 + ax.scatter( + [horizon + jitter] * len(month_points), + month_points["mean"], + color=color, + alpha=0.4, + s=30, + marker="o", + ) + + # Customize plot + ax.set_xlabel("Forecast Lead Time (months)", fontsize=12, fontweight="bold") + ax.set_ylabel( + metric.upper() if len(metric) <= 4 else metric.capitalize(), + fontsize=12, + fontweight="bold", + ) + ax.set_title( + f"{metric.upper()} vs Forecast Lead Time\n(Lines = mean across all {aggregate_by.replace('_', ' ')}s, points = individual months)", + fontsize=14, + fontweight="bold", + ) + ax.set_xticks(horizons) + ax.set_xticklabels([str(h) for h in horizons]) + ax.set_ylim(y_min, y_max) + + # Add horizontal reference line + if metric == "r2": + ax.axhline(y=0, color="gray", linestyle="--", linewidth=0.8, alpha=0.7) + elif metric in ["nrmse", "nmae", "efficiency"]: + ax.axhline(y=1.0, color="gray", linestyle="--", linewidth=0.8, alpha=0.7) + elif metric == "accuracy": + ax.axhline(y=0.5, color="gray", linestyle="--", linewidth=0.8, alpha=0.7) + + ax.grid(axis="both", alpha=0.3) + ax.legend(loc="best", fontsize=11, framealpha=0.9, title="Model") + + plt.tight_layout() + + # Save if output path provided + if output_path: + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(output_path, dpi=150, bbox_inches="tight") + logger.info(f"Saved {metric} by lead time combined plot to {output_path}") + + return fig + + +def main(): + region = "Kyrgyzstan" # Options: "Kyrgyzstan" or "Tajikistan" + save_dir = Path(output_dir) / f"{region.lower()}" + # Configuration for plotting + metric_to_plot = ( + "r2" # Options: "r2", "rmse", "nrmse", "mae", "nmae", "accuracy", "efficiency" + ) + + # Load predictions + if region == "Tajikistan": + pred_config = taj_path_config + else: + pred_config = kgz_path_config + + logger.info(f"Loading predictions for {region}...") + predictions_df = load_predictions( + base_path=pred_config["pred_dir"], + horizons=horizons, + issue_months=issue_months_to_evaluate, + ) + + if predictions_df.empty: + logger.error("No predictions loaded. Exiting.") + return + + logger.info(f"Loaded {len(predictions_df)} prediction records") + logger.info(f"Available models: {predictions_df['model'].unique().tolist()}") + logger.info(f"Available horizons: {predictions_df['horizon'].unique().tolist()}") + + # Create ensemble + logger.info("Creating ensemble...") + predictions_df = create_ensemble(predictions_df) + + # Handle observations based on use_Q_obs flag + if use_Q_obs: + logger.info("Using Q_obs from prediction files...") + # Check if Q_obs is available + if predictions_df["Q_obs"].isna().all(): + logger.error( + "Q_obs not available in prediction files. Set use_Q_obs=False to load from observations file." + ) + return + + # Add target_month, issue_month based on region logic + if region == "Kyrgyzstan": + predictions_df["target_date"] = predictions_df.apply( + lambda row: row["issue_date"] + + pd.DateOffset(months=int(row["horizon"])), + axis=1, + ) + elif region == "Tajikistan": + predictions_df["target_date"] = predictions_df.apply( + lambda row: row["issue_date"] + + pd.DateOffset(months=int(row["horizon"]) - 1), + axis=1, + ) + + predictions_df["target_month"] = predictions_df["target_date"].dt.month + predictions_df["issue_month"] = predictions_df["issue_date"].dt.month + aggregated_df = predictions_df + + else: + logger.info(f"Loading observations from {pred_config['obs_file']}...") + obs_df = load_observations(pred_config["obs_file"]) + + # Calculate monthly targets + logger.info("Calculating monthly observation targets...") + monthly_obs = calculate_target(obs_df) + + # Apply climatology-based ratio correction if enabled + if apply_climatology_correction: + logger.info("Applying climatology-based ratio correction...") + + # Calculate leave-one-out monthly means (fast vectorized version) + monthly_obs_with_ltm = calculate_leave_one_out_monthly_mean_fast( + monthly_obs + ) + + # Apply ratio correction to predictions (fast vectorized version) + predictions_df = apply_climatology_ratio_correction_fast( + predictions_df=predictions_df, + daily_obs=obs_df, + monthly_obs_with_ltm=monthly_obs_with_ltm, + region=region, + ) + + # Use corrected predictions as Q_pred for evaluation + # Keep original Q_pred as Q_pred_original + predictions_df["Q_pred_original"] = predictions_df["Q_pred"] + predictions_df["Q_pred"] = predictions_df["Q_pred_corrected"] + + logger.info( + "Ratio correction applied. Q_pred now contains corrected values." + ) + + # Aggregate predictions with observations + logger.info("Aggregating predictions with observations...") + aggregated_df = aggregate(predictions_df, monthly_obs, region) + + # Compute metrics dataframe + logger.info("Computing metrics...") + metrics_df = compute_metrics_dataframe(aggregated_df) + + logger.info(f"Metrics computed for {len(metrics_df)} combinations") + logger.info(f"\nMetrics DataFrame preview:") + print(metrics_df.head(20)) + + # Diagnostic: Print coverage per model and horizon + print("\n" + "=" * 80) + print( + "DIAGNOSTIC: Data coverage per model and horizon (number of target months with data)" + ) + print("=" * 80) + for model in models_plot: + print(f"\n{model}:") + model_data = metrics_df[metrics_df["model"] == model] + if model_data.empty: + print(" NO DATA FOUND!") + # Check if it's in the predictions_df + pred_model_data = aggregated_df[aggregated_df["model"] == model] + if not pred_model_data.empty: + print(f" But found {len(pred_model_data)} rows in aggregated_df") + print( + f" Horizons in aggregated_df: {sorted(pred_model_data['horizon'].unique())}" + ) + print( + f" Target months in aggregated_df: {sorted(pred_model_data['target_month'].unique())}" + ) + else: + for horizon in sorted(model_data["horizon"].unique()): + horizon_data = model_data[model_data["horizon"] == horizon] + target_months = sorted(horizon_data["target_month"].unique()) + n_codes = horizon_data["code"].nunique() + print( + f" Horizon {horizon}: {len(target_months)} target months {target_months}, {n_codes} codes, {len(horizon_data)} records" + ) + + # Also print raw predictions coverage for comparison + print("\n" + "=" * 80) + print("DIAGNOSTIC: Raw predictions coverage (before metrics calculation)") + print("=" * 80) + for model in models_plot: + print(f"\n{model}:") + model_data = aggregated_df[aggregated_df["model"] == model] + if model_data.empty: + print(" NO DATA FOUND in aggregated_df!") + else: + for horizon in sorted(model_data["horizon"].unique()): + horizon_data = model_data[model_data["horizon"] == horizon] + target_months = sorted(horizon_data["target_month"].unique()) + n_records = len(horizon_data) + n_with_obs = horizon_data["Q_obs"].notna().sum() + n_with_pred = horizon_data["Q_pred"].notna().sum() + print( + f" Horizon {horizon}: {n_records} records, {n_with_obs} with Q_obs, {n_with_pred} with Q_pred, target_months: {target_months}" + ) + print("=" * 80 + "\n") + + # Add Ensemble to models_plot if not already there + models_to_plot = ( + models_plot + ["Ensemble"] if "Ensemble" not in models_plot else models_plot + ) + + # Plot metric by target month for each forecast horizon + available_horizons = sorted(metrics_df["horizon"].unique()) + logger.info(f"\nPlotting {metric_to_plot} by target month for each horizon...") + + for horizon in available_horizons: + logger.info(f" Plotting horizon {horizon}...") + fig = plot_metric_by_horizon( + metrics_df=metrics_df, + horizon=horizon, + models=models_to_plot, + metric=metric_to_plot, + output_path=Path(save_dir) + / f"{metric_to_plot}_horizon_{horizon}_by_target_month.png" + if save_dir + else None, + ) + # plt.show() + plt.close(fig) # Close to avoid memory issues + + # Plot metric by lead time (horizon) for all months + logger.info(f"\nPlotting {metric_to_plot} by forecast lead time...") + + # Plot with subplots per model (showing each month as a separate line) + fig = plot_metric_by_lead_time( + metrics_df=metrics_df, + models=models_to_plot, + metric=metric_to_plot, + output_path=Path(save_dir) / f"{metric_to_plot}_by_lead_time_per_model.png" + if save_dir + else None, + aggregate_by="target_month", + ) + plt.close(fig) + + # Plot combined view (all models on one plot, mean across months) + fig = plot_metric_by_lead_time_combined( + metrics_df=metrics_df, + models=models_to_plot, + metric=metric_to_plot, + output_path=Path(save_dir) / f"{metric_to_plot}_by_lead_time_combined.png" + if save_dir + else None, + aggregate_by="target_month", + ) + plt.close(fig) + + logger.info("Processing complete!") + + +if __name__ == "__main__": + main() diff --git a/scripts/investigate_mcald.py b/scripts/investigate_mcald.py new file mode 100644 index 0000000..3cb8228 --- /dev/null +++ b/scripts/investigate_mcald.py @@ -0,0 +1,1027 @@ +""" +Script to investigate MC_ALD model predictions and test correction strategies. + +This script analyzes the MC_ALD model performance across different regions and +forecast horizons, training linear regression correction models using leave-one-out +cross-validation. +""" + +import os +import sys +import logging +from pathlib import Path + +logger = logging.getLogger(__name__) +logger.setLevel(logging.DEBUG) + +# Add a stream handler to output logs to the terminal +if not logger.handlers: + handler = logging.StreamHandler(sys.stdout) + handler.setLevel(logging.DEBUG) + formatter = logging.Formatter( + "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + ) + handler.setFormatter(formatter) + logger.addHandler(handler) + + +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import seaborn as sns +from sklearn.metrics import r2_score +from sklearn.linear_model import LinearRegression +from sklearn.model_selection import LeaveOneOut + +# load the .env file +from dotenv import load_dotenv + +load_dotenv(dotenv_path=Path(__file__).parent.parent / ".env") + + +# ============================================================================= +# CONFIGURATION +# ============================================================================= + +# Region path configurations from environment variables +kgz_path_config = { + "pred_dir": os.getenv("kgz_path_discharge"), + "obs_file": os.getenv("kgz_path_base_pred"), +} + +taj_path_config = { + "pred_dir": os.getenv("taj_path_base_pred"), + "obs_file": os.getenv("taj_path_discharge"), +} + +output_dir = os.getenv("out_dir_op_lt") + +# Available forecast horizons +ALL_HORIZONS = [ + "month_0", + "month_1", + "month_2", + "month_3", + "month_4", + "month_5", + "month_6", + "month_7", + "month_8", + "month_9", +] + +# Day of forecast for each horizon (when forecasts are issued) +day_of_forecast = { + "month_0": 15, + "month_1": 25, + "month_2": 25, + "month_3": 25, + "month_4": 25, + "month_5": 25, + "month_6": 25, + "month_7": 25, + "month_8": 25, + "month_9": 25, +} + +month_renaming = { + 1: "January", + 2: "February", + 3: "March", + 4: "April", + 5: "May", + 6: "June", + 7: "July", + 8: "August", + 9: "September", + 10: "October", + 11: "November", + 12: "December", +} + +# ============================================================================= +# USER CONFIGURATION - MODIFY THESE +# ============================================================================= + +# Regions to analyze: ["Kyrgyzstan", "Tajikistan"] or just one +REGIONS_TO_ANALYZE: list[str] = ["Kyrgyzstan"] + +# Horizons to analyze: list of horizon names or "all" +# Examples: ["month_1", "month_2", "month_3"] or ["month_0"] or "all" +HORIZONS_TO_ANALYZE: list[str] | str = [ + "month_0", + "month_1", + "month_2", + "month_3", + "month_4", + "month_5", +] + +# Whether to save plots to files (True) or display interactively (False) +SAVE_PLOTS: bool = True + +# ============================================================================= +# FUNCTIONS +# ============================================================================= + + +def get_path_config(region: str) -> dict[str, str]: + """Get path configuration for a specific region.""" + if region == "Kyrgyzstan": + return kgz_path_config + elif region == "Tajikistan": + return taj_path_config + else: + raise ValueError( + f"Unknown region: {region}. Must be 'Kyrgyzstan' or 'Tajikistan'" + ) + + +def get_mc_ald_path(pred_dir: str, horizon: str) -> Path: + """ + Construct the path to MC_ALD hindcast file for a given horizon. + + Args: + pred_dir: Base prediction directory + horizon: Horizon name (e.g., "month_1") + + Returns: + Path to the MC_ALD hindcast CSV file + """ + return Path(pred_dir) / horizon / "MC_ALD" / "MC_ALD_hindcast.csv" + + +def load_mc_ald_data( + pred_dir: str, + horizon: str, +) -> pd.DataFrame | None: + """ + Load MC_ALD hindcast data for a specific horizon. + + Args: + pred_dir: Base prediction directory + horizon: Horizon name (e.g., "month_1") + + Returns: + DataFrame with MC_ALD data or None if file not found + """ + file_path = get_mc_ald_path(pred_dir, horizon) + + if not file_path.exists(): + logger.warning(f"MC_ALD hindcast file not found: {file_path}") + return None + + try: + data = pd.read_csv(file_path, parse_dates=["date"]) + + # Filter by day of forecast if specified + forecast_day = day_of_forecast.get(horizon) + if forecast_day is not None: + data = data[data["date"].dt.day == forecast_day].copy() + + # Select relevant columns + required_cols = ["date", "code", "Q_loc", "Q_MC_ALD", "Q_obs"] + available_cols = [col for col in required_cols if col in data.columns] + + if len(available_cols) < len(required_cols): + missing = set(required_cols) - set(available_cols) + logger.warning(f"Missing columns in {file_path}: {missing}") + return None + + data = data[available_cols].dropna() + data["month"] = data["date"].dt.month + data["horizon"] = horizon + data["horizon_num"] = int(horizon.split("_")[1]) + + logger.info( + f"Loaded {len(data)} records from {horizon} " + f"({data['code'].nunique()} codes, months: {sorted(data['month'].unique())})" + ) + + return data + + except Exception as e: + logger.error(f"Failed to load {file_path}: {e}") + return None + + +def train_correction_model_loocv( + data: pd.DataFrame, +) -> pd.DataFrame: + """ + Train a linear regression per code and month to predict the correction + (Q_obs - Q_loc) based on Q_loc using leave-one-out cross-validation. + + Args: + data: DataFrame with columns 'code', 'month', 'Q_loc', 'Q_obs', 'Q_MC_ALD' + + Returns: + DataFrame with additional columns for predictions and corrected values + """ + # Calculate the target correction + data = data.copy() + data["target_correction"] = data["Q_obs"] - data["Q_loc"] + + # Initialize columns for predictions + data["Q_loc_corrected"] = np.nan + + # Get unique codes + codes = data["code"].unique() + + for code in codes: + code_mask = data["code"] == code + code_data = data[code_mask] + + # Get unique months for this code + months = code_data["month"].unique() + + for month in months: + month_mask = (data["code"] == code) & (data["month"] == month) + month_data = data[month_mask] + + if len(month_data) < 2: + # Not enough data for LOOCV, use mean correction + logger.warning( + f"Code {code}, month {month}: only {len(month_data)} sample(s), " + "using Q_loc as prediction" + ) + data.loc[month_mask, "Q_loc_corrected"] = month_data["Q_loc"] + continue + + # Prepare features and target + X = month_data[["Q_loc"]].values + y = month_data["target_correction"].values + indices = month_data.index.tolist() + + # Leave-one-out cross-validation + loo = LeaveOneOut() + predictions = np.zeros(len(X)) + + for train_idx, test_idx in loo.split(X): + X_train, X_test = X[train_idx], X[test_idx] + y_train = y[train_idx] + + model = LinearRegression() + model.fit(X_train, y_train) + predictions[test_idx] = model.predict(X_test) + + # Calculate corrected Q_loc + corrected_values = month_data["Q_loc"].values + predictions + + # Update the dataframe + for i, idx in enumerate(indices): + data.loc[idx, "Q_loc_corrected"] = corrected_values[i] + + return data + + +def calculate_r2_by_month(data: pd.DataFrame) -> pd.DataFrame: + """ + Calculate R² scores per month for different predictions vs Q_obs. + + Args: + data: DataFrame with Q_obs, Q_loc, Q_MC_ALD, and Q_loc_corrected + + Returns: + DataFrame with R² scores per month + """ + results = [] + + for month in sorted(data["month"].unique()): + month_data = data[data["month"] == month].dropna( + subset=["Q_obs", "Q_loc", "Q_MC_ALD", "Q_loc_corrected"] + ) + + if len(month_data) < 2: + continue + + r2_q_loc = r2_score(month_data["Q_obs"], month_data["Q_loc"]) + r2_mc_ald = r2_score(month_data["Q_obs"], month_data["Q_MC_ALD"]) + r2_corrected = r2_score(month_data["Q_obs"], month_data["Q_loc_corrected"]) + + results.append( + { + "month": month, + "R² Q_loc (original)": r2_q_loc, + "R² Q_MC_ALD": r2_mc_ald, + "R² Q_loc (corrected)": r2_corrected, + "n_samples": len(month_data), + } + ) + + return pd.DataFrame(results) + + +def calculate_r2_by_code_month(data: pd.DataFrame) -> pd.DataFrame: + """ + Calculate R² scores per code and month for different predictions vs Q_obs. + + Args: + data: DataFrame with Q_obs, Q_loc, Q_MC_ALD, and Q_loc_corrected + + Returns: + DataFrame with R² scores per code and month + """ + results = [] + + for code in data["code"].unique(): + for month in sorted(data["month"].unique()): + mask = (data["code"] == code) & (data["month"] == month) + subset = data[mask].dropna( + subset=["Q_obs", "Q_loc", "Q_MC_ALD", "Q_loc_corrected"] + ) + + if len(subset) < 2: + continue + + r2_q_loc = r2_score(subset["Q_obs"], subset["Q_loc"]) + r2_mc_ald = r2_score(subset["Q_obs"], subset["Q_MC_ALD"]) + r2_corrected = r2_score(subset["Q_obs"], subset["Q_loc_corrected"]) + + results.append( + { + "code": code, + "month": month, + "R² Q_loc (original)": r2_q_loc, + "R² Q_MC_ALD": r2_mc_ald, + "R² Q_loc (corrected)": r2_corrected, + "n_samples": len(subset), + } + ) + + return pd.DataFrame(results) + + +def plot_r2_distribution_by_month( + r2_df: pd.DataFrame, + region: str, + horizon: str, + save_path: Path | None = None, +) -> plt.Figure: + """ + Plot the distribution of R² scores per month for all three methods. + + Args: + r2_df: DataFrame with R² scores per code and month + region: Region name for the title + horizon: Horizon name for the title + save_path: Optional path to save the figure + + Returns: + matplotlib Figure object + """ + # Reshape data for plotting + plot_data = [] + for _, row in r2_df.iterrows(): + plot_data.append( + { + "month": row["month"], + "Method": "Q_loc (original)", + "R²": row["R² Q_loc (original)"], + } + ) + plot_data.append( + {"month": row["month"], "Method": "Q_MC_ALD", "R²": row["R² Q_MC_ALD"]} + ) + plot_data.append( + { + "month": row["month"], + "Method": "Q_loc (corrected)", + "R²": row["R² Q_loc (corrected)"], + } + ) + + plot_df = pd.DataFrame(plot_data) + + # Create the boxplot and violin plot side by side + fig, axes = plt.subplots(1, 2, figsize=(18, 7)) + + # Define colors for each method + palette = { + "Q_loc (original)": "#3498db", # Blue + "Q_MC_ALD": "#e74c3c", # Red + "Q_loc (corrected)": "#2ecc71", # Green + } + + # Boxplot + sns.boxplot( + x="month", y="R²", hue="Method", data=plot_df, palette=palette, ax=axes[0] + ) + axes[0].axhline(0, color="gray", linestyle="--", alpha=0.5) + axes[0].set_title( + f"{region} - {horizon}\nR² Distribution per Month (Leave-One-Out CV)", + fontsize=12, + fontweight="bold", + ) + axes[0].set_xlabel("Month", fontsize=11) + axes[0].set_ylabel("R² Score", fontsize=11) + axes[0].legend(title="Method", loc="lower right", fontsize=9) + axes[0].grid(True, alpha=0.3) + + # Violin plot + sns.violinplot( + x="month", + y="R²", + hue="Method", + data=plot_df, + palette=palette, + inner="box", + ax=axes[1], + ) + axes[1].axhline(0, color="gray", linestyle="--", alpha=0.5) + axes[1].set_title( + f"{region} - {horizon}\nR² Violin Plot (Leave-One-Out CV)", + fontsize=12, + fontweight="bold", + ) + axes[1].set_xlabel("Month", fontsize=11) + axes[1].set_ylabel("R² Score", fontsize=11) + axes[1].legend(title="Method", loc="lower right", fontsize=9) + axes[1].grid(True, alpha=0.3) + + plt.tight_layout() + + if save_path: + save_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(save_path, dpi=150, bbox_inches="tight") + logger.info(f"Saved R² distribution plot to {save_path}") + + return fig + + +def plot_summary_statistics( + r2_df: pd.DataFrame, + region: str, + horizon: str, + save_path: Path | None = None, +) -> plt.Figure: + """ + Plot summary statistics comparing the three methods. + + Args: + r2_df: DataFrame with R² scores per code and month + region: Region name for the title + horizon: Horizon name for the title + save_path: Optional path to save the figure + + Returns: + matplotlib Figure object + """ + # Calculate median R² per month for each method + summary = ( + r2_df.groupby("month") + .agg( + { + "R² Q_loc (original)": ["median", "mean", "std", "count"], + "R² Q_MC_ALD": ["median", "mean", "std"], + "R² Q_loc (corrected)": ["median", "mean", "std"], + } + ) + .round(3) + ) + + logger.info(f"\n{'=' * 80}") + logger.info(f"{region} - {horizon}: Summary Statistics per Month") + logger.info("=" * 80) + print(summary.to_string()) + + # Calculate overall statistics + overall_stats = pd.DataFrame( + { + "Method": ["Q_loc (original)", "Q_MC_ALD", "Q_loc (corrected)"], + "Median R²": [ + r2_df["R² Q_loc (original)"].median(), + r2_df["R² Q_MC_ALD"].median(), + r2_df["R² Q_loc (corrected)"].median(), + ], + "Mean R²": [ + r2_df["R² Q_loc (original)"].mean(), + r2_df["R² Q_MC_ALD"].mean(), + r2_df["R² Q_loc (corrected)"].mean(), + ], + "Std R²": [ + r2_df["R² Q_loc (original)"].std(), + r2_df["R² Q_MC_ALD"].std(), + r2_df["R² Q_loc (corrected)"].std(), + ], + } + ).round(3) + + logger.info(f"\n{'=' * 80}") + logger.info(f"{region} - {horizon}: Overall Statistics") + logger.info("=" * 80) + print(overall_stats.to_string(index=False)) + + # Plot bar chart of median R² per month + fig, axes = plt.subplots(1, 2, figsize=(16, 6)) + + months = sorted(r2_df["month"].unique()) + x = np.arange(len(months)) + width = 0.25 + + medians_orig = r2_df.groupby("month")["R² Q_loc (original)"].median() + medians_mcald = r2_df.groupby("month")["R² Q_MC_ALD"].median() + medians_corr = r2_df.groupby("month")["R² Q_loc (corrected)"].median() + + # Median R² bar chart + axes[0].bar( + x - width, medians_orig, width, label="Q_loc (original)", color="#3498db" + ) + axes[0].bar(x, medians_mcald, width, label="Q_MC_ALD", color="#e74c3c") + axes[0].bar( + x + width, medians_corr, width, label="Q_loc (corrected)", color="#2ecc71" + ) + + axes[0].set_xlabel("Month", fontsize=11) + axes[0].set_ylabel("Median R²", fontsize=11) + axes[0].set_title( + f"{region} - {horizon}\nMedian R² per Month by Method", + fontsize=12, + fontweight="bold", + ) + axes[0].set_xticks(x) + axes[0].set_xticklabels([month_renaming.get(m, m)[:3] for m in months]) + axes[0].legend(fontsize=9) + axes[0].axhline(0, color="gray", linestyle="--", alpha=0.5) + axes[0].grid(True, alpha=0.3) + + # Improvement comparison + improvement_vs_orig = medians_corr - medians_orig + improvement_vs_mcald = medians_corr - medians_mcald + + axes[1].bar( + x - width / 2, + improvement_vs_orig, + width, + label="vs Q_loc (original)", + color="#3498db", + ) + axes[1].bar( + x + width / 2, improvement_vs_mcald, width, label="vs Q_MC_ALD", color="#e74c3c" + ) + + axes[1].set_xlabel("Month", fontsize=11) + axes[1].set_ylabel("R² Improvement", fontsize=11) + axes[1].set_title( + f"{region} - {horizon}\nR² Improvement of Corrected Q_loc", + fontsize=12, + fontweight="bold", + ) + axes[1].set_xticks(x) + axes[1].set_xticklabels([month_renaming.get(m, m)[:3] for m in months]) + axes[1].legend(fontsize=9) + axes[1].axhline(0, color="gray", linestyle="--", alpha=0.5) + axes[1].grid(True, alpha=0.3) + + plt.tight_layout() + + if save_path: + save_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(save_path, dpi=150, bbox_inches="tight") + logger.info(f"Saved summary statistics plot to {save_path}") + + return fig + + +def plot_r2_by_horizon( + all_r2_data: pd.DataFrame, + region: str, + save_path: Path | None = None, +) -> plt.Figure: + """ + Plot R² vs forecast horizon for all methods. + + Args: + all_r2_data: DataFrame with R² scores including horizon_num column + region: Region name for the title + save_path: Optional path to save the figure + + Returns: + matplotlib Figure object + """ + # Aggregate by horizon + agg_data = ( + all_r2_data.groupby("horizon_num") + .agg( + { + "R² Q_loc (original)": ["mean", "std"], + "R² Q_MC_ALD": ["mean", "std"], + "R² Q_loc (corrected)": ["mean", "std"], + } + ) + .reset_index() + ) + + # Flatten column names + agg_data.columns = [ + "horizon", + "Q_loc_mean", + "Q_loc_std", + "Q_MC_ALD_mean", + "Q_MC_ALD_std", + "Q_loc_corr_mean", + "Q_loc_corr_std", + ] + + fig, ax = plt.subplots(figsize=(12, 7)) + + horizons = agg_data["horizon"].values + + # Plot lines with error bands + ax.plot( + horizons, + agg_data["Q_loc_mean"], + "o-", + color="#3498db", + linewidth=2, + markersize=8, + label="Q_loc (original)", + ) + ax.fill_between( + horizons, + agg_data["Q_loc_mean"] - agg_data["Q_loc_std"], + agg_data["Q_loc_mean"] + agg_data["Q_loc_std"], + color="#3498db", + alpha=0.2, + ) + + ax.plot( + horizons, + agg_data["Q_MC_ALD_mean"], + "s-", + color="#e74c3c", + linewidth=2, + markersize=8, + label="Q_MC_ALD", + ) + ax.fill_between( + horizons, + agg_data["Q_MC_ALD_mean"] - agg_data["Q_MC_ALD_std"], + agg_data["Q_MC_ALD_mean"] + agg_data["Q_MC_ALD_std"], + color="#e74c3c", + alpha=0.2, + ) + + ax.plot( + horizons, + agg_data["Q_loc_corr_mean"], + "^-", + color="#2ecc71", + linewidth=2, + markersize=8, + label="Q_loc (corrected)", + ) + ax.fill_between( + horizons, + agg_data["Q_loc_corr_mean"] - agg_data["Q_loc_corr_std"], + agg_data["Q_loc_corr_mean"] + agg_data["Q_loc_corr_std"], + color="#2ecc71", + alpha=0.2, + ) + + ax.set_xlabel("Forecast Lead Time (months)", fontsize=12, fontweight="bold") + ax.set_ylabel("Mean R² Score", fontsize=12, fontweight="bold") + ax.set_title( + f"{region}: R² vs Forecast Lead Time\n(Mean ± Std across all codes and months)", + fontsize=14, + fontweight="bold", + ) + ax.set_xticks(horizons) + ax.axhline(0, color="gray", linestyle="--", alpha=0.5) + ax.legend(loc="best", fontsize=11) + ax.grid(True, alpha=0.3) + + plt.tight_layout() + + if save_path: + save_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(save_path, dpi=150, bbox_inches="tight") + logger.info(f"Saved R² by horizon plot to {save_path}") + + return fig + + +def plot_bias_by_month( + data: pd.DataFrame, + region: str, + horizon: str, + ax: plt.Axes | None = None, +) -> plt.Axes: + """ + Plot bias boxplots by month as percentage of observed, averaged per code: + mean((Q_obs - Q_loc) / Q_obs * 100) and mean((Q_MC_ALD - Q_loc) / Q_obs * 100) per code. + + Args: + data: DataFrame with Q_obs, Q_loc, Q_MC_ALD, code, month columns + region: Region name for the title + horizon: Horizon name for the title + ax: Optional matplotlib Axes to plot on + + Returns: + matplotlib Axes object + """ + # Calculate biases as percentage of observed + data = data.copy() + + # Skip rows where Q_obs is zero or very small + data = data[data["Q_obs"] > 0.01] + + # Calculate bias percentages + data["true_bias_pct"] = (data["Q_obs"] - data["Q_loc"]) / data["Q_obs"] * 100 + data["model_corr_pct"] = (data["Q_MC_ALD"] - data["Q_loc"]) / data["Q_obs"] * 100 + + # Average per code and month + avg_bias = ( + data.groupby(["code", "month"]) + .agg( + { + "true_bias_pct": "mean", + "model_corr_pct": "mean", + } + ) + .reset_index() + ) + + # Reshape for plotting + plot_data = [] + for _, row in avg_bias.iterrows(): + plot_data.append( + { + "month": row["month"], + "Bias Type": "(Q_obs - Q_loc) / Q_obs (true bias)", + "Bias (%)": row["true_bias_pct"], + } + ) + plot_data.append( + { + "month": row["month"], + "Bias Type": "(Q_MC_ALD - Q_loc) / Q_obs (model correction)", + "Bias (%)": row["model_corr_pct"], + } + ) + + plot_df = pd.DataFrame(plot_data) + + if ax is None: + fig, ax = plt.subplots(figsize=(14, 6)) + + # Define colors + palette = { + "(Q_obs - Q_loc) / Q_obs (true bias)": "#2ecc71", # Green + "(Q_MC_ALD - Q_loc) / Q_obs (model correction)": "#e74c3c", # Red + } + + sns.boxplot( + x="month", y="Bias (%)", hue="Bias Type", data=plot_df, palette=palette, ax=ax + ) + + ax.set_ylim(-100, 100) + ax.axhline(0, color="gray", linestyle="--", linewidth=1.5, alpha=0.7) + ax.set_title(f"{region} - {horizon}", fontsize=11, fontweight="bold") + ax.set_xlabel("Month", fontsize=10) + ax.set_ylabel("Bias (% of Q_obs, avg per code)", fontsize=10) + ax.legend(title="", loc="upper right", fontsize=8) + ax.grid(True, alpha=0.3, axis="y") + + return ax + + +def plot_bias_all_horizons( + all_data: list[pd.DataFrame], + horizons: list[str], + region: str, + save_path: Path | None = None, +) -> plt.Figure: + """ + Plot bias boxplots for all horizons in a single figure. + + Args: + all_data: List of DataFrames, one per horizon + horizons: List of horizon names + region: Region name for the title + save_path: Optional path to save the figure + + Returns: + matplotlib Figure object + """ + n_horizons = len(horizons) + n_cols = min(3, n_horizons) + n_rows = (n_horizons + n_cols - 1) // n_cols + + fig, axes = plt.subplots( + n_rows, n_cols, figsize=(7 * n_cols, 5 * n_rows), squeeze=False + ) + axes_flat = axes.flatten() + + for idx, (data, horizon) in enumerate(zip(all_data, horizons)): + ax = axes_flat[idx] + plot_bias_by_month(data, region, horizon, ax=ax) + + # Only show legend on first subplot + if idx > 0: + ax.get_legend().remove() + + # Hide unused subplots + for idx in range(n_horizons, len(axes_flat)): + axes_flat[idx].set_visible(False) + + fig.suptitle( + f"{region}: Bias Comparison by Month (% of Q_obs)\n((Q_obs - Q_loc) / Q_obs vs (Q_MC_ALD - Q_loc) / Q_obs)", + fontsize=14, + fontweight="bold", + y=1.02, + ) + + plt.tight_layout() + + if save_path: + save_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(save_path, dpi=150, bbox_inches="tight") + logger.info(f"Saved bias comparison plot to {save_path}") + + return fig + + +def analyze_region_horizon( + region: str, + horizon: str, + save_dir: Path | None = None, +) -> tuple[pd.DataFrame | None, pd.DataFrame | None]: + """ + Analyze MC_ALD for a specific region and horizon. + + Args: + region: Region name + horizon: Horizon name (e.g., "month_1") + save_dir: Optional directory to save plots + + Returns: + Tuple of (data_with_corrections, r2_by_code_month) or (None, None) if no data + """ + logger.info(f"\n{'=' * 80}") + logger.info(f"Analyzing {region} - {horizon}") + logger.info("=" * 80) + + # Get path configuration + path_config = get_path_config(region) + + # Load data + data = load_mc_ald_data(path_config["pred_dir"], horizon) + + if data is None or data.empty: + logger.warning(f"No data available for {region} - {horizon}") + return None, None + + logger.info( + f"Loaded {len(data)} records with {data['code'].nunique()} unique codes" + ) + + # Train correction model using LOOCV + logger.info("Training linear regression correction model with Leave-One-Out CV...") + data_with_corrections = train_correction_model_loocv(data) + + # Calculate R² scores per code and month + logger.info("Calculating R² scores per code and month...") + r2_by_code_month = calculate_r2_by_code_month(data_with_corrections) + + if r2_by_code_month.empty: + logger.error("No R² scores could be calculated. Check your data.") + return data_with_corrections, None + + logger.info(f"Calculated R² for {len(r2_by_code_month)} code-month combinations") + + # Add horizon info to R² data + r2_by_code_month["horizon"] = horizon + r2_by_code_month["horizon_num"] = int(horizon.split("_")[1]) + + # Create plots + if save_dir: + dist_path = save_dir / f"{region.lower()}_{horizon}_r2_distribution.png" + summary_path = save_dir / f"{region.lower()}_{horizon}_summary_stats.png" + else: + dist_path = None + summary_path = None + + fig1 = plot_r2_distribution_by_month(r2_by_code_month, region, horizon, dist_path) + fig2 = plot_summary_statistics(r2_by_code_month, region, horizon, summary_path) + + if not SAVE_PLOTS: + plt.show() + else: + plt.close(fig1) + plt.close(fig2) + + # Show aggregated R² per month + logger.info(f"\n{'=' * 80}") + logger.info(f"{region} - {horizon}: Aggregated R² per Month (all samples pooled)") + logger.info("=" * 80) + r2_aggregated = calculate_r2_by_month(data_with_corrections) + print(r2_aggregated.to_string(index=False)) + + return data_with_corrections, r2_by_code_month + + +def main() -> None: + """Main function to run MC_ALD analysis.""" + + # Determine horizons to analyze + if HORIZONS_TO_ANALYZE == "all": + horizons_to_use = ALL_HORIZONS + else: + horizons_to_use = HORIZONS_TO_ANALYZE + + logger.info("=" * 80) + logger.info("MC_ALD Investigation Script") + logger.info("=" * 80) + logger.info(f"Regions to analyze: {REGIONS_TO_ANALYZE}") + logger.info(f"Horizons to analyze: {horizons_to_use}") + logger.info(f"Save plots: {SAVE_PLOTS}") + logger.info("=" * 80) + + # Analyze each region + for region in REGIONS_TO_ANALYZE: + logger.info(f"\n{'#' * 80}") + logger.info(f"REGION: {region}") + logger.info("#" * 80) + + # Determine save directory for this region + if SAVE_PLOTS and output_dir: + save_dir = Path(output_dir) / region.lower() / "mc_ald_investigation" + save_dir.mkdir(parents=True, exist_ok=True) + else: + save_dir = None + + all_r2_data = [] + all_raw_data = [] # Store raw data for bias plots + processed_horizons = [] # Track which horizons have data + + # Analyze each horizon + for horizon in horizons_to_use: + data, r2_data = analyze_region_horizon( + region=region, + horizon=horizon, + save_dir=save_dir, + ) + + if r2_data is not None: + all_r2_data.append(r2_data) + + if data is not None: + all_raw_data.append(data) + processed_horizons.append(horizon) + + # Create combined bias plot if we have data from multiple horizons + if len(all_raw_data) > 0: + bias_plot_path = None + if save_dir: + bias_plot_path = save_dir / f"{region.lower()}_bias_all_horizons.png" + + fig = plot_bias_all_horizons( + all_raw_data, processed_horizons, region, bias_plot_path + ) + + if not SAVE_PLOTS: + plt.show() + else: + plt.close(fig) + + # Create combined horizon plot if multiple horizons analyzed + if len(all_r2_data) > 1: + combined_r2 = pd.concat(all_r2_data, ignore_index=True) + + horizon_plot_path = None + if save_dir: + horizon_plot_path = save_dir / f"{region.lower()}_r2_by_horizon.png" + + fig = plot_r2_by_horizon(combined_r2, region, horizon_plot_path) + + if not SAVE_PLOTS: + plt.show() + else: + plt.close(fig) + + # Print overall summary + logger.info(f"\n{'=' * 80}") + logger.info(f"{region}: Overall Summary Across All Horizons") + logger.info("=" * 80) + + overall_summary = ( + combined_r2.groupby("horizon_num") + .agg( + { + "R² Q_loc (original)": ["mean", "median", "std"], + "R² Q_MC_ALD": ["mean", "median", "std"], + "R² Q_loc (corrected)": ["mean", "median", "std"], + } + ) + .round(3) + ) + print(overall_summary.to_string()) + + logger.info("\n" + "=" * 80) + logger.info("MC_ALD Investigation Complete!") + logger.info("=" * 80) + + +if __name__ == "__main__": + main()