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, ) 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) # Anomaly target, training and test feature cubes. target_train_anom = gd.ds[target].values[train_time_idx, :, :] - clim[target][month - 1][None, :, :] 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 = pred_anom + clim[target][month - 1][None, :, :] 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["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()) 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) test_global_indices = np.where(test_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 # Restrict metrics to valid climate-region cells only. invalid = ~gd.valid_mask pred[:, invalid] = np.nan obs[:, invalid] = np.nan raw[:, 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"], } 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) 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}