"""Create the GD1 monthly observed/model climatology NetCDF product.""" from __future__ import annotations import argparse from pathlib import Path import numpy as np import xarray as xr from .config import load_config, project_path from .precipitation import NEGATIVE_PRECIPITATION_POLICY, clip_negative_precipitation def _rename_spatial(da: xr.DataArray) -> xr.DataArray: rename = {} for old, new in (("lon", "longitude"), ("lat", "latitude"), ("time", "time")): if old in da.dims or old in da.coords: rename[old] = new return da.rename(rename) def _require(da: xr.DataArray, names: set[str], label: str) -> None: missing = names - set(da.dims) - set(da.coords) if missing: raise ValueError(f"{label} is missing required coordinate/dimension(s): {sorted(missing)}") def climatology_periods(config: dict) -> tuple[str, str, str, str]: """Return explicit observation and model-initialization periods. The flat ``baseline.start/end`` form remains readable for older configs, but new configs separate observation valid times from GCM initialization times so an operational monthly file cannot enter the baseline by accident. """ baseline = config["baseline"] if "observation" not in baseline and "model_initialization" not in baseline: return baseline["start"], baseline["end"], baseline["start"], baseline["end"] missing = sorted({"observation", "model_initialization"} - set(baseline)) if missing: raise KeyError(f"Baseline configuration is missing period(s): {missing}") observation = baseline["observation"] model = baseline["model_initialization"] return observation["start"], observation["end"], model["start"], model["end"] def observed_climatology( path: Path, variable: str, start: str, end: str, precipitation: bool = False, ) -> xr.DataArray: if not path.exists(): raise FileNotFoundError(f"Observation file not found: {path}") with xr.open_dataset(path) as ds: if variable not in ds: raise KeyError(f"Variable '{variable}' not found in {path}; available: {list(ds.data_vars)}") da = _rename_spatial(ds[variable]).sel(time=slice(start, end)).load() _require(da, {"time", "latitude", "longitude"}, "Observed data") if precipitation: da = clip_negative_precipitation(da) result = da.groupby("time.month").mean("time", skipna=True).transpose("longitude", "latitude", "month") if precipitation: result.attrs["negative_value_policy"] = NEGATIVE_PRECIPITATION_POLICY return result def _model_paths(pattern: str) -> list[Path]: candidate = Path(pattern) if candidate.is_absolute(): return sorted(candidate.parent.glob(candidate.name)) return sorted(project_path(".").glob(pattern)) def model_climatology( pattern: str, variable: str, start: str, end: str, precipitation: bool = False, ) -> xr.DataArray: paths = _model_paths(pattern) if not paths: raise FileNotFoundError(f"No monthly model NetCDF files match: {pattern}") wanted = np.arange(1, 7) sums: np.ndarray | None = None counts: np.ndarray | None = None reference: dict[str, np.ndarray] = {} source_attrs: dict[str, object] = {} seen_initializations: set[str] = set() for path in paths: with xr.open_dataset(path) as ds: if variable not in ds: raise KeyError( f"Variable '{variable}' not found in {path}; available: {list(ds.data_vars)}" ) da = _rename_spatial(ds[variable]).sel(init_time=slice(start, end)).load() _require( da, {"init_time", "member", "lead", "latitude", "longitude"}, "Model data" ) lead_values = np.asarray(da["lead"].values) available = {int(value) for value in lead_values} if not set(wanted).issubset(available): raise ValueError( f"Model lead coordinate must contain 1..6; found {lead_values.tolist()}" ) selected = da.sel(lead=wanted).transpose( "init_time", "member", "lead", "latitude", "longitude" ) if precipitation: selected = clip_negative_precipitation(selected) current = { name: np.asarray(selected[name].values) for name in ("member", "lead", "latitude", "longitude") } if sums is None: reference = {name: values.copy() for name, values in current.items()} shape = ( reference["member"].size, reference["lead"].size, reference["latitude"].size, reference["longitude"].size, 12, ) sums = np.zeros(shape, dtype=np.float64) counts = np.zeros(shape, dtype=np.int32) source_attrs = dict(selected.attrs) else: for name, values in current.items(): if not np.array_equal(values, reference[name]): raise ValueError(f"Inconsistent model {name} coordinate in {path}") assert sums is not None and counts is not None for init_index, init_value in enumerate(selected["init_time"].values): init_month = np.datetime64(init_value, "M") init_key = str(init_month) if init_key in seen_initializations: raise ValueError(f"Duplicate model initialization month: {init_key}") seen_initializations.add(init_key) initial_calendar_month = int(init_key[5:7]) values = np.asarray(selected.isel(init_time=init_index).values) for lead_index, lead in enumerate(wanted): month_index = (initial_calendar_month + int(lead) - 1) % 12 field = values[:, lead_index, :, :] valid = np.isfinite(field) sums[:, lead_index, :, :, month_index] += np.where(valid, field, 0.0) counts[:, lead_index, :, :, month_index] += valid if sums is None or counts is None or not seen_initializations: raise ValueError(f"No model initializations fall inside {start}..{end}") mean = np.full(sums.shape, np.nan, dtype=np.float64) np.divide(sums, counts, out=mean, where=counts > 0) result = xr.DataArray( mean, dims=("member", "lead", "latitude", "longitude", "month"), coords={**reference, "month": np.arange(1, 13)}, attrs=source_attrs, name=variable, ).transpose("longitude", "latitude", "member", "lead", "month") result.attrs["ensemble_processing"] = "members retained; no ensemble mean" result.attrs["initialization_month_count"] = len(seen_initializations) if precipitation: result.attrs["negative_value_policy"] = NEGATIVE_PRECIPITATION_POLICY return result def build(config_path: str = "config/project.json") -> Path: config = load_config(config_path) obs_start, obs_end, model_start, model_end = climatology_periods(config) paths, variables = config["paths"], config["variables"] pr_obs = observed_climatology( project_path(paths["rainfall_observation"]), variables["rainfall_observation"], obs_start, obs_end, precipitation=True, ) t2m_obs = observed_climatology( project_path(paths["temperature_observation"]), variables["temperature_observation"], obs_start, obs_end, ) pr_mod = model_climatology( paths["model_monthly_glob"], variables["rainfall_model"], model_start, model_end, precipitation=True, ) t2m_mod = model_climatology( paths["model_monthly_glob"], variables["temperature_model"], model_start, model_end, ) output = project_path(paths["climatology_output"]) output.parent.mkdir(parents=True, exist_ok=True) xr.Dataset({"pr_mod": pr_mod, "t2m_mod": t2m_mod, "pr_obs": pr_obs, "t2m_obs": t2m_obs}).to_netcdf(output) return output def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--config", default="config/project.json") args = parser.parse_args() print(build(args.config)) if __name__ == "__main__": main()