from __future__ import annotations import json import warnings from pathlib import Path from typing import Any, Dict, Iterable, List, Tuple import numpy as np import pandas as pd import xarray as xr from sklearn.exceptions import ConvergenceWarning from .data_utils import ( ALL_DATA_VARIABLES, TARGET_TO_MODEL, compute_train_climatology, load_grid_data, raw_for_var, year_range_mask, target_training_transform, inverse_target_transform, ) from .features import build_feature_cube, predict_to_grid, select_samples from .metrics import grid_metrics, metrics_to_dataset, summarize_by_region from .models import make_model from .plotting import plot_all_metric_maps warnings.filterwarnings("ignore", category=ConvergenceWarning) def task_label(task: Dict[str, Any]) -> str: return ( f"{task['source_model']}_lead{int(task['lead'])}_" f"{task['target']}_{task['spatial_mode']}_{task['predictor_set']}_{task['method']}" ) def _target_units(target: str) -> str: if target == "pr_o": return "mm/month" if target == "t2_o": return "C" return "" def _spatial_masks(valid_mask: np.ndarray, reg: np.ndarray, spatial_mode: str, regions: Iterable[int]) -> List[Tuple[str, np.ndarray]]: if spatial_mode == "national": return [("ALL", valid_mask)] if spatial_mode == "regional": return [(f"R{int(r)}", valid_mask & (reg.astype(float) == float(r))) for r in regions] raise ValueError(f"Unknown spatial_mode: {spatial_mode}") def _fit_predict_month( gd, clim: Dict[str, np.ndarray], cfg: Dict[str, Any], task: Dict[str, Any], month: int, train_time_idx: np.ndarray, test_time_idx: np.ndarray, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """Return pred, obs, raw grids for one calendar month over test_time_idx. Shapes are (len(test_time_idx), lat, lon). """ target = task["target"] method = task["method"].lower() predictor_set = task["predictor_set"] spatial_mode = task["spatial_mode"] direct_model_var = TARGET_TO_MODEL[target] obs_raw = raw_for_var(gd.ds, target, test_time_idx) raw_pred = raw_for_var(gd.ds, direct_model_var, test_time_idx) pred_out = np.full_like(obs_raw, np.nan, dtype=np.float32) if method == "raw": pred_out[:] = raw_pred.astype(np.float32) if target == "pr_o" and cfg.get("clip_precipitation_to_zero", True): pred_out = np.where(np.isfinite(pred_out), np.maximum(pred_out, 0.0), np.nan).astype(np.float32) return pred_out, obs_raw.astype(np.float32), raw_pred.astype(np.float32) # Target transform: t2 anomaly; pr factor. target_train_anom = target_training_transform(gd.ds, clim, target, train_time_idx, month) feature_train, _ = build_feature_cube( gd.ds, clim, target, predictor_set, train_time_idx, month, gd.lat2d, gd.lon2d, gd.reg, spatial_mode, include_lat_lon=cfg.get("include_lat_lon", True), include_region_feature=cfg.get("include_region_feature", True), ) feature_test, _ = build_feature_cube( gd.ds, clim, target, predictor_set, test_time_idx, month, gd.lat2d, gd.lon2d, gd.reg, spatial_mode, include_lat_lon=cfg.get("include_lat_lon", True), include_region_feature=cfg.get("include_region_feature", True), ) for region_label, smask in _spatial_masks(gd.valid_mask, gd.reg, spatial_mode, cfg.get("valid_regions", [1, 2, 3, 4, 5, 6, 7])): if not np.any(smask): continue X_train, y_train, *_ = select_samples(feature_train, target_train_anom, smask) if len(y_train) < 20: continue model = make_model(method, cfg, random_seed=int(cfg.get("random_seed", 42))) model.fit(X_train, y_train) pred_anom = predict_to_grid(model, feature_test, smask) pred_raw = inverse_target_transform(pred_anom, clim, target, month, clip_precipitation_to_zero=cfg.get("clip_precipitation_to_zero", True)) pred_out[:, smask] = pred_raw[:, smask] if target == "pr_o" and cfg.get("clip_precipitation_to_zero", True): pred_out = np.where(np.isfinite(pred_out), np.maximum(pred_out, 0.0), np.nan).astype(np.float32) return pred_out.astype(np.float32), obs_raw.astype(np.float32), raw_pred.astype(np.float32) def run_task(cfg: Dict[str, Any], task: Dict[str, Any], dirs: Dict[str, Path]) -> Dict[str, Path]: label = task_label(task) done_file = dirs["done"] / f"{label}.done" metrics_file = dirs["metrics"] / f"grid_metrics_{label}.nc" summary_file = dirs["metrics"] / f"summary_{label}.csv" pred_file = dirs["corrected"] / f"pred_{label}.nc" if done_file.exists() and not cfg.get("overwrite", False): return {"status": "skipped", "label": label, "metrics": metrics_file, "summary": summary_file, "pred": pred_file} gd = load_grid_data( data_dir=cfg["data_dir"], mask_file=cfg["mask_file"], file_pattern=cfg.get("file_pattern", "lead{lead}_{source_model}_01_BL_Test_VN.nc"), source_model=task["source_model"], lead=int(task["lead"]), year_start=int(cfg.get("year_start", 1982)), missing_value=float(cfg.get("missing_value", -99.0)), ) variables_needed = list(dict.fromkeys(ALL_DATA_VARIABLES)) clim = compute_train_climatology(gd.ds, gd.years, cfg.get("climatology_years", cfg["train_years"]), variables_needed) train_year_mask = year_range_mask(gd.years, cfg["train_years"]) test_year_mask = year_range_mask(gd.years, cfg["test_years"]) ntest = int(test_year_mask.sum()) ntrain = int(train_year_mask.sum()) nlat = gd.ds.sizes["lat"] nlon = gd.ds.sizes["lon"] pred = np.full((ntest, nlat, nlon), np.nan, dtype=np.float32) obs = np.full_like(pred, np.nan) raw = np.full_like(pred, np.nan) pred_train = np.full((ntrain, nlat, nlon), np.nan, dtype=np.float32) obs_train = np.full_like(pred_train, np.nan) raw_train = np.full_like(pred_train, np.nan) test_global_indices = np.where(test_year_mask)[0] train_global_indices = np.where(train_year_mask)[0] for month in range(1, 13): train_idx = np.where(train_year_mask & (gd.months == month))[0] test_idx = np.where(test_year_mask & (gd.months == month))[0] if len(test_idx) == 0: continue pred_m, obs_m, raw_m = _fit_predict_month(gd, clim, cfg, task, month, train_idx, test_idx) # Map global test indices to compact test-period index. compact_positions = np.searchsorted(test_global_indices, test_idx) pred[compact_positions, :, :] = pred_m obs[compact_positions, :, :] = obs_m raw[compact_positions, :, :] = raw_m # Training diagnostics requested in updated Dr_SFC.docx. These are # fitted and predicted on the training period, used only for web # comparison against validation and not for independent skill claims. if cfg.get("save_train_predictions", True) and len(train_idx) > 0: pred_tr_m, obs_tr_m, raw_tr_m = _fit_predict_month(gd, clim, cfg, task, month, train_idx, train_idx) tr_positions = np.searchsorted(train_global_indices, train_idx) pred_train[tr_positions, :, :] = pred_tr_m obs_train[tr_positions, :, :] = obs_tr_m raw_train[tr_positions, :, :] = raw_tr_m # Restrict metrics to valid climate-region cells only. invalid = ~gd.valid_mask pred[:, invalid] = np.nan obs[:, invalid] = np.nan raw[:, invalid] = np.nan pred_train[:, invalid] = np.nan obs_train[:, invalid] = np.nan raw_train[:, invalid] = np.nan metric_values = grid_metrics(pred, obs, rmae_percent=bool(cfg.get("rmae_percent", True))) metric_ds = metrics_to_dataset(metric_values, gd.ds["lat"].values, gd.ds["lon"].values) metric_ds.attrs.update({k: str(v) for k, v in task.items()}) metric_ds.attrs["train_years"] = f"{cfg['train_years'][0]}-{cfg['train_years'][1]}" metric_ds.attrs["test_years"] = f"{cfg['test_years'][0]}-{cfg['test_years'][1]}" metric_ds.to_netcdf(metrics_file) meta = { "source_model": task["source_model"], "lead": int(task["lead"]), "target": task["target"], "spatial_mode": task["spatial_mode"], "predictor_set": task["predictor_set"], "method": task["method"], "member": task.get("member", f"{task['method']}_{task['predictor_set']}_{task['spatial_mode']}"), } rows = summarize_by_region(metric_values, gd.reg, cfg.get("valid_regions", [1, 2, 3, 4, 5, 6, 7]), meta) pd.DataFrame(rows).to_csv(summary_file, index=False) if cfg.get("save_predictions", True): time_test = gd.time[test_global_indices] ds_pred = xr.Dataset( coords={"time": time_test, "lat": gd.ds["lat"].values, "lon": gd.ds["lon"].values}, data_vars={ "obs": (("time", "lat", "lon"), obs.astype(np.float32)), "raw": (("time", "lat", "lon"), raw.astype(np.float32)), "pred": (("time", "lat", "lon"), pred.astype(np.float32)), }, attrs={**{k: str(v) for k, v in task.items()}, "units": _target_units(task["target"])}, ) ds_pred["obs"].attrs["long_name"] = f"Observed {task['target']}" ds_pred["raw"].attrs["long_name"] = f"Raw model forecast for {task['target']}" ds_pred["pred"].attrs["long_name"] = f"Bias-corrected prediction for {task['target']}" ds_pred.to_netcdf(pred_file) if cfg.get("save_train_predictions", True): train_file = dirs.get("corrected_train", dirs["corrected"].parent / "corrected_train") / f"pred_train_{label}.nc" time_train = gd.time[train_global_indices] ds_tr = xr.Dataset( coords={"time": time_train, "lat": gd.ds["lat"].values, "lon": gd.ds["lon"].values}, data_vars={ "obs": (("time", "lat", "lon"), obs_train.astype(np.float32)), "raw": (("time", "lat", "lon"), raw_train.astype(np.float32)), "pred": (("time", "lat", "lon"), pred_train.astype(np.float32)), }, attrs={**{k: str(v) for k, v in task.items()}, "units": _target_units(task["target"]), "period": "training_diagnostic"}, ) ds_tr["obs"].attrs["long_name"] = f"Observed {task['target']}" ds_tr["raw"].attrs["long_name"] = f"Raw model forecast for {task['target']}" ds_tr["pred"].attrs["long_name"] = f"Bias-corrected training-period prediction for {task['target']}" ds_tr.to_netcdf(train_file) plot_all_metric_maps( gd.ds["lon"].values, gd.ds["lat"].values, metric_values, label, dirs["maps"], ) done_file.write_text(json.dumps({"label": label, "status": "done"}, indent=2), encoding="utf-8") return {"status": "done", "label": label, "metrics": metrics_file, "summary": summary_file, "pred": pred_file}