from __future__ import annotations from dataclasses import dataclass from pathlib import Path from typing import Dict, Iterable, List, Tuple import numpy as np import pandas as pd import xarray as xr from src.safe_ops import safe_mean MODEL_VARIABLES = [ "pr", "t2", "tx", "tm", "h500", "h850", "olr", "q2", "rh2", "slp", "u200", "u850", "w700" ] OBS_VARIABLES = ["pr_o", "t2_o", "tx_o", "tm_o"] ALL_DATA_VARIABLES = OBS_VARIABLES + MODEL_VARIABLES TARGET_TO_MODEL = {"pr_o": "pr", "t2_o": "t2", "tx_o": "tx", "tm_o": "tm"} TARGET_SHORT = {"pr_o": "pr", "t2_o": "t2", "tx_o": "tx", "tm_o": "tm"} @dataclass class GridData: ds: xr.Dataset mask: xr.Dataset years: np.ndarray months: np.ndarray time: pd.DatetimeIndex valid_mask: np.ndarray # lat, lon bool reg: np.ndarray # lat, lon float/int with nan outside regions lon2d: np.ndarray lat2d: np.ndarray def _replace_missing(ds: xr.Dataset, missing_value: float) -> xr.Dataset: out = ds.copy() for name in list(out.data_vars): arr = out[name] fill_values = [missing_value] for att in ["missing_value", "_FillValue"]: if att in arr.attrs: try: fill_values.append(float(arr.attrs[att])) except Exception: pass mask = xr.zeros_like(arr, dtype=bool) for mv in set(fill_values): mask = mask | np.isclose(arr, mv) out[name] = arr.where(~mask) return out def open_lead_dataset(data_dir: str | Path, file_pattern: str, source_model: str, lead: int, missing_value: float) -> xr.Dataset: fname = file_pattern.format(lead=lead, source_model=source_model) path = Path(data_dir) / fname if not path.exists(): raise FileNotFoundError(f"Input file not found: {path}") ds = xr.open_dataset(path, decode_times=False) # Ensure standard dimension order for data variables. for name in list(ds.data_vars): if set(["mon", "lat", "lon"]).issubset(ds[name].dims): ds[name] = ds[name].transpose("mon", "lat", "lon") return _replace_missing(ds, missing_value) def open_mask(mask_file: str | Path, missing_value: float) -> xr.Dataset: path = Path(mask_file) if not path.exists(): raise FileNotFoundError(f"Mask file not found: {path}") mask = xr.open_dataset(path, decode_times=False) for name in list(mask.data_vars): if set(["lat", "lon"]).issubset(mask[name].dims): mask[name] = mask[name].transpose("lat", "lon") return _replace_missing(mask, missing_value) def make_time_axis(year_start: int, ntime: int) -> Tuple[pd.DatetimeIndex, np.ndarray, np.ndarray]: # User convention: mon=1 means YYYY-01, not one month after YYYY-01. time = pd.date_range(f"{year_start}-01-01", periods=ntime, freq="MS") years = time.year.to_numpy() months = time.month.to_numpy() return time, years, months def load_grid_data( data_dir: str | Path, mask_file: str | Path, file_pattern: str, source_model: str, lead: int, year_start: int, missing_value: float, ) -> GridData: ds = open_lead_dataset(data_dir, file_pattern, source_model, lead, missing_value) mask = open_mask(mask_file, missing_value) # Basic grid compatibility checks. if ds.sizes["lat"] != mask.sizes["lat"] or ds.sizes["lon"] != mask.sizes["lon"]: raise ValueError("Lead file and mask file have different lat/lon sizes.") if not np.allclose(ds["lat"].values, mask["lat"].values) or not np.allclose(ds["lon"].values, mask["lon"].values): raise ValueError("Lead file and mask file have different lat/lon coordinates.") ntime = ds.sizes["mon"] time, years, months = make_time_axis(year_start, ntime) reg = mask["reg"].values.astype(float) valid_mask = np.isfinite(reg) & (reg > 0) lon2d, lat2d = np.meshgrid(ds["lon"].values, ds["lat"].values) return GridData(ds=ds, mask=mask, years=years, months=months, time=time, valid_mask=valid_mask, reg=reg, lon2d=lon2d, lat2d=lat2d) def year_range_mask(years: np.ndarray, year_pair: Iterable[int]) -> np.ndarray: y1, y2 = list(year_pair) return (years >= y1) & (years <= y2) def calendar_month_mask(months: np.ndarray, month: int) -> np.ndarray: return months == month def compute_train_climatology( ds: xr.Dataset, years: np.ndarray, train_years: Iterable[int], variables: List[str], ) -> Dict[str, np.ndarray]: """Return climatology[var][month_index, lat, lon] computed from training years only.""" train_mask = year_range_mask(years, train_years) clim: Dict[str, np.ndarray] = {} for var in variables: arr = ds[var].values # mon, lat, lon out = np.full((12, arr.shape[1], arr.shape[2]), np.nan, dtype=np.float32) for m in range(1, 13): tmask = train_mask & (((np.arange(arr.shape[0]) % 12) + 1) == m) if np.any(tmask): out[m - 1] = safe_mean(arr[tmask, :, :], axis=0, min_count=1, dtype=np.float32) clim[var] = out return clim def anomaly_for_var(ds: xr.Dataset, clim: Dict[str, np.ndarray], var: str, time_indices: np.ndarray, month: int) -> np.ndarray: arr = ds[var].values[time_indices, :, :] # time, lat, lon return arr - clim[var][month - 1][None, :, :] def ratio_for_var(ds: xr.Dataset, clim: Dict[str, np.ndarray], var: str, time_indices: np.ndarray, month: int, eps: float = 1.0e-6) -> np.ndarray: """Return var / monthly climatology for positive variables such as precipitation.""" arr = ds[var].values[time_indices, :, :].astype(float) den = clim[var][month - 1][None, :, :].astype(float) out = np.full_like(arr, np.nan, dtype=float) ok = np.isfinite(arr) & np.isfinite(den) & (np.abs(den) > float(eps)) np.divide(arr, den, out=out, where=ok) return out def transformed_predictor_for_var(ds: xr.Dataset, clim: Dict[str, np.ndarray], var: str, time_indices: np.ndarray, month: int, eps: float = 1.0e-6) -> np.ndarray: """Predictor transform used in GĐ1. Precipitation is represented by a monthly factor (pr/pr_cli); other variables are represented by monthly anomalies. This keeps the pr correction consistent with the factor method while preserving anomaly features for temperature and circulation variables. """ if var == "pr": return ratio_for_var(ds, clim, var, time_indices, month, eps=eps) return anomaly_for_var(ds, clim, var, time_indices, month) def target_training_transform(ds: xr.Dataset, clim: Dict[str, np.ndarray], target: str, time_indices: np.ndarray, month: int, eps: float = 1.0e-6) -> np.ndarray: """Return the AI target transform for a predictand. * t2_o: observed anomaly = t2_o - t2_o_cli * pr_o: observed factor = pr_o / pr_o_cli """ if target == "pr_o": return ratio_for_var(ds, clim, target, time_indices, month, eps=eps) return anomaly_for_var(ds, clim, target, time_indices, month) def inverse_target_transform(pred_transformed: np.ndarray, clim: Dict[str, np.ndarray], target: str, month: int, clip_precipitation_to_zero: bool = True) -> np.ndarray: """Convert model output from transformed target space back to physical units.""" if target == "pr_o": out = pred_transformed * clim[target][month - 1][None, :, :] if clip_precipitation_to_zero: out = np.where(np.isfinite(out), np.maximum(out, 0.0), np.nan) return out return pred_transformed + clim[target][month - 1][None, :, :] def raw_for_var(ds: xr.Dataset, var: str, time_indices: np.ndarray) -> np.ndarray: return ds[var].values[time_indices, :, :] def target_obs_climatology(clim: Dict[str, np.ndarray], target: str, month: int) -> np.ndarray: return clim[target][month - 1]