from __future__ import annotations from pathlib import Path from typing import Dict, Optional import numpy as np import pandas as pd import xarray as xr from .dates import days_in_month # Canonical operational variable names mapped to GĐ1/hindcast feature names. CANONICAL_TO_HINDCAST_VAR = { 'pr': 'pr', 'tas': 't2', 'tasmax': 'tx', 'tasmin': 'tm', 'h500': 'h500', 'h850': 'h850', 'olr': 'olr', 'q2': 'q2', 'rh2': 'rh2', 'slp': 'slp', 'u200': 'u200', 'u850': 'u850', 'w700': 'w700', } ALIASES = { 'tas': ['tas', 't2'], 'tasmax': ['tasmax', 'tx'], 'tasmin': ['tasmin', 'tm'], 'pr': ['pr'], 'h500': ['h500'], 'h850': ['h850'], 'olr': ['olr', 'Olr', 'OLR'], 'q2': ['q2'], 'rh2': ['rh2'], 'slp': ['slp', 'Slp', 'SLP'], 'u200': ['u200'], 'u850': ['u850'], 'w700': ['w700'], } def _find_var(ds: xr.Dataset, canonical: str) -> Optional[str]: direct = ALIASES.get(canonical, [canonical]) lower_map = {str(v).lower(): str(v) for v in ds.data_vars} for name in direct: if name in ds.data_vars: return name if name.lower() in lower_map: return lower_map[name.lower()] return None def _replace_missing(ds: xr.Dataset, missing_value: float = -32767.0) -> xr.Dataset: out = ds.copy() for v in list(out.data_vars): arr = out[v] fill_values = [missing_value] for att in ['_FillValue', 'missing_value']: 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[v] = arr.where(~mask) return out def _standardize_time_dim(ds: xr.Dataset) -> xr.Dataset: if 'time' in ds.dims: return ds for cand in ['lead', 'month', 'mon', 'LT']: if cand in ds.dims and int(ds.sizes[cand]) == 6: return ds.rename({cand: 'time'}) raise ValueError('Operational file must contain a 6-step time/lead dimension.') def _standardize_horizontal_dims(arr: xr.DataArray) -> xr.DataArray: """Return an array with dimensions exactly time, lat, lon. Operational NetCDF files are not fully uniform. Some variables are stored as time,lat,lon while others include singleton auxiliary dimensions such as m2 (for 2-m variables). For a deterministic operational issue we collapse non-spatial auxiliary dimensions safely before prediction. Singleton auxiliary dimensions are squeezed; if a non-singleton auxiliary dimension is encountered, it is averaged with skipna=True and recorded by xarray. """ rename = {} for d in arr.dims: dl = str(d).lower() if dl in {'latitude', 'y'}: rename[d] = 'lat' elif dl in {'longitude', 'x'}: rename[d] = 'lon' if rename: arr = arr.rename(rename) required = {'time', 'lat', 'lon'} missing = required.difference(arr.dims) if missing: raise ValueError(f'missing required dimensions {sorted(missing)}; actual dims={arr.dims}') for d in list(arr.dims): if d in required: continue if int(arr.sizes[d]) == 1: arr = arr.isel({d: 0}, drop=True) else: arr = arr.mean(dim=d, skipna=True) return arr.transpose('time', 'lat', 'lon') def read_operational_rcm_file(path: str | Path, init_year: int, init_month: int, missing_value: float = -32767.0) -> xr.Dataset: """Read operational RCM file and standardize to GĐ1 feature names. The input 6-step time dimension is interpreted as LT=1..6. Output variables are named exactly as in the GĐ1 hindcast files: pr, t2, tx, tm, h500, ... . Operational tas/tasmax/tasmin are converted to t2/tx/tm. Temperature in K is converted to degC when values/units indicate Kelvin. pr is treated as monthly total unless units clearly indicate flux. """ path = Path(path) ds0 = xr.open_dataset(path, decode_times=False) ds0 = _standardize_time_dim(ds0) ds0 = _replace_missing(ds0, missing_value) if int(ds0.sizes['time']) != 6: raise ValueError(f'Operational file must have 6 lead steps: {path}') for coord in ['lat', 'lon']: if coord not in ds0.coords and coord not in ds0.variables: raise ValueError(f'Missing {coord} in operational file: {path}') leads = np.arange(1, 7, dtype=np.int16) target_dates = [] target_year = [] target_month = [] for lead in leads: t = pd.Timestamp(int(init_year), int(init_month), 1) + pd.DateOffset(months=int(lead)) target_dates.append(t) target_year.append(int(t.year)) target_month.append(int(t.month)) data_vars = {} missing_canonical = [] for canonical, out_name in CANONICAL_TO_HINDCAST_VAR.items(): in_name = _find_var(ds0, canonical) if in_name is None: missing_canonical.append(canonical) continue arr = ds0[in_name] if 'time' not in arr.dims: continue # Prefer common lat/lon dimension names and collapse auxiliary singleton dims # such as m2 that appear in some operational files. try: arr_std = _standardize_horizontal_dims(arr) vals = arr_std.values.astype(np.float32) except Exception as exc: raise ValueError( f'Variable {in_name} in {path} must contain time,lat,lon dimensions; ' f'auxiliary dimensions are allowed only if collapsible. actual dims={arr.dims}' ) from exc units_in = str(ds0[in_name].attrs.get('units', '')).strip() units_norm = units_in.lower().replace(' ', '') if canonical in {'tas', 'tasmax', 'tasmin'}: # Convert Kelvin only when units or values indicate K. if ('k' == units_norm) or ('kelvin' in units_norm) or np.nanmedian(vals) > 100.0: vals = vals - 273.15 units = 'degC' elif canonical == 'pr': if ('s-1' in units_norm) or ('/s' in units_norm) or ('sec' in units_norm): for k, (yy, mm) in enumerate(zip(target_year, target_month)): vals[k] = vals[k] * 86400.0 * days_in_month(yy, mm) vals = np.where(np.isfinite(vals), np.maximum(vals, 0.0), np.nan).astype(np.float32) units = 'mm/month' else: units = units_in data_vars[out_name] = (('lead', 'lat', 'lon'), vals, {'units': units, 'source_variable': in_name}) if 'pr' not in data_vars or 't2' not in data_vars: raise ValueError(f'Operational file {path} must contain pr and tas/t2 after standardization. Missing canonical variables: {missing_canonical}') out = xr.Dataset( coords={ 'lead': leads, 'target_time': ('lead', np.array(target_dates, dtype='datetime64[ns]')), 'lat': ds0['lat'].values, 'lon': ds0['lon'].values, }, data_vars=data_vars, attrs={'init_year': int(init_year), 'init_month': int(init_month), 'source_file': str(path)}, ) ds0.close() return out