"""Monthly SPEI calculation using Thornthwaite PET and a log-logistic fit. The fitting functions retain parameters by calendar month, grid cell, and timescale, so the same parameters can be used for operational forecasts. """ from __future__ import annotations from typing import Iterable import numpy as np import xarray as xr from scipy.stats import fisk, norm from .precipitation import NEGATIVE_PRECIPITATION_POLICY, clip_negative_precipitation def _require_monthly(da: xr.DataArray, label: str) -> None: if "time" not in da.dims: raise ValueError(f"{label} must have a time dimension") if not np.issubdtype(da.time.dtype, np.datetime64): raise ValueError(f"{label}.time must use datetime values") def fit_thornthwaite_parameters(calibration_t2m: xr.DataArray) -> xr.Dataset: """Fit climatological Thornthwaite heat index on calibration data only.""" _require_monthly(calibration_t2m, "calibration_t2m") if "latitude" not in calibration_t2m.dims and "latitude" not in calibration_t2m.coords: raise ValueError("calibration_t2m must include latitude") months = np.unique(calibration_t2m.time.dt.month.values) if not np.array_equal(months, np.arange(1, 13)): raise ValueError("Thornthwaite calibration must contain all 12 calendar months") climatology = calibration_t2m.clip(min=0).groupby("time.month").mean("time", skipna=True) heat_index = ((climatology / 5) ** 1.514).sum("month", skipna=False).rename("heat_index") exponent = ( 6.75e-7 * heat_index**3 - 7.71e-5 * heat_index**2 + 1.792e-2 * heat_index + 0.49239 ).rename("thornthwaite_exponent") result = xr.Dataset({"heat_index": heat_index, "thornthwaite_exponent": exponent}) result.attrs = { "method": "Thornthwaite heat index from 12 monthly temperature climatologies", "calibration_start": str(calibration_t2m.time.min().values), "calibration_end": str(calibration_t2m.time.max().values), "operational_rule": "reuse saved parameters; do not refit on forecast years", } return result def thornthwaite_pet(t2m: xr.DataArray, parameters: xr.Dataset) -> xr.DataArray: """Return monthly potential evapotranspiration (mm/month). Heat index and exponent must come from ``fit_thornthwaite_parameters`` on the training/calibration period and are reused unchanged operationally. """ _require_monthly(t2m, "t2m") if "latitude" not in t2m.dims and "latitude" not in t2m.coords: raise ValueError("t2m must include a latitude coordinate") required = {"heat_index", "thornthwaite_exponent"} missing = sorted(required - set(parameters.data_vars)) if missing: raise KeyError(f"Thornthwaite parameters are missing: {missing}") positive = t2m.clip(min=0) index = parameters["heat_index"] exponent = parameters["thornthwaite_exponent"] latitude = np.deg2rad(t2m.latitude) # Approximate mid-month solar declination and day length (FAO-56). days = t2m.time.dt.days_in_month doy = t2m.time.dt.dayofyear + (days - 1) / 2 declination = 0.409 * np.sin(2 * np.pi * doy / 365 - 1.39) sunset_angle = np.arccos(np.clip(-np.tan(latitude) * np.tan(declination), -1, 1)) daylight_hours = 24 / np.pi * sunset_angle pet = 16 * (daylight_hours / 12) * (days / 30) * (10 * positive / index) ** exponent result = pet.where((positive > 0) & (index > 0), 0).rename("pet") result.attrs = {"units": "mm month-1", "method": "Thornthwaite", "parameter_calibration_start": parameters.attrs.get("calibration_start", ""), "parameter_calibration_end": parameters.attrs.get("calibration_end", "")} return result def water_balance(pr: xr.DataArray, t2m: xr.DataArray, thornthwaite_parameters: xr.Dataset) -> xr.DataArray: """Return monthly climatic water balance (precipitation minus PET).""" _require_monthly(pr, "pr") nonnegative_pr = clip_negative_precipitation(pr) result = (nonnegative_pr - thornthwaite_pet(t2m, thornthwaite_parameters)).rename("water_balance") result.attrs["precipitation_negative_value_policy"] = NEGATIVE_PRECIPITATION_POLICY return result def accumulated_balance(balance: xr.DataArray, scale: int) -> xr.DataArray: """Return trailing monthly water-balance totals for a positive timescale.""" if scale < 1: raise ValueError("scale must be at least 1") _require_monthly(balance, "balance") return balance.rolling(time=scale, min_periods=scale).sum().rename(f"balance_{scale}") def _fit_1d(values: np.ndarray) -> np.ndarray: values = values[np.isfinite(values)] if values.size < 4 or np.allclose(values, values[0]): return np.array([np.nan, np.nan, np.nan]) shape, loc, scale = fisk.fit(values) return np.array([shape, loc, scale]) def fit_loglogistic(balance: xr.DataArray, scales: Iterable[int] = (1, 3, 6)) -> xr.Dataset: """Fit three-parameter log-logistic distributions by month and grid cell.""" _require_monthly(balance, "balance") fitted = [] for scale in scales: values = accumulated_balance(balance, int(scale)) by_month = [] for month in range(1, 13): subset = values.where(values.time.dt.month == month, drop=True) params = xr.apply_ufunc( _fit_1d, subset, input_core_dims=[["time"]], output_core_dims=[["parameter"]], vectorize=True, dask="forbidden", output_dtypes=[float], dask_gufunc_kwargs={"output_sizes": {"parameter": 3}}, ).assign_coords(parameter=["shape", "loc", "scale"]).expand_dims(month=[month]) by_month.append(params) fitted.append(xr.concat(by_month, dim="month").expand_dims(timescale=[scale])) params = xr.concat(fitted, dim="timescale").rename("loglogistic_parameters") return params.to_dataset() def spei_from_parameters(balance: xr.DataArray, parameters: xr.Dataset, scale: int) -> xr.DataArray: """Convert a balance series to SPEI using saved ``fit_loglogistic`` output.""" if "loglogistic_parameters" not in parameters: raise KeyError("parameters must contain 'loglogistic_parameters'") values = accumulated_balance(balance, scale) # Vectorized selection turns ``month`` into the existing time dimension. params = parameters["loglogistic_parameters"].sel(timescale=scale, month=values.time.dt.month) params = params.reset_coords(["month", "timescale"], drop=True) shape, loc, distribution_scale = (params.sel(parameter=name) for name in ("shape", "loc", "scale")) cdf = xr.apply_ufunc(fisk.cdf, values, shape, loc, distribution_scale, dask="forbidden", output_dtypes=[float]) result = xr.apply_ufunc(norm.ppf, cdf.clip(1e-8, 1 - 1e-8), dask="forbidden", output_dtypes=[float]) return result.rename(f"spei_{scale}")