""" Anomaly model for the climate sub-model: predicts (actual - climatology) at a requested horizon, given a recent history of anomalies. Combined with climate_climatology.py's climatology at inference time: forecast(date, horizon) = climatology(day_of_year) + predicted_anomaly(horizon) DESIGN CHOICE, stated directly rather than left implicit: the predicted anomaly is passed through a learnable, per-variable, MONOTONICALLY DECAYING gate as a function of horizon -- built in architecturally, not left for training to discover on its own. This is a direct response to what the real climatology results showed: precipitation has close to zero seasonal predictability even at the baseline level (4/20 real grid cells beat a flat mean, and the margins were negligible), so there is no reason to expect a data-driven model to have genuine multi-month signal for it either. Rather than risk the model outputting a confident-looking but meaningless anomaly at day 180, the gate forces the anomaly's contribution toward zero as horizon grows -- how fast is LEARNED per variable (a temperature-like variable with real persistence should decay slowly; precipitation should decay fast), but the DIRECTION (toward zero, never away from it) is guaranteed by construction, not hoped for. Verification status: prepare_training_windows and decay_gate are pure numpy/pandas, tested directly below. The nn.Module (ClimateAnomalyModel) has been syntax-checked but not run end-to-end -- run scripts/test_climate_anomaly_model.py against real data and confirm behavior before trusting it, the same way the climatology and grid-extraction work was validated. """ from typing import Dict, List, Optional, Tuple import numpy as np import pandas as pd try: import torch import torch.nn as nn _HAS_TORCH = True except ImportError: _HAS_TORCH = False nn = None # only referenced inside the class body below, never at import time otherwise # --------------------------------------------------------------------------- # Data preparation -- pure numpy/pandas, fully testable without torch # --------------------------------------------------------------------------- def prepare_training_windows( anomaly_wide: pd.DataFrame, lookback_days: int, horizons: List[int], ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, List[pd.Timestamp]]: """ Slices a [date x cell] anomaly DataFrame into (history, target) training examples: for every date with a full lookback window behind it AND real target values at every requested horizon ahead of it, extract the lookback window as input and the horizon-ahead anomalies as targets. A window is only kept if the ENTIRE lookback period is real (non-NaN) for that cell -- a history with gaps silently interpolated would teach the model on fabricated input, not real conditions. A target is NaN-masked per (example, horizon) individually rather than requiring every horizon to be real, since discharge/climate series have real, uneven gaps and requiring all horizons real would throw away most of the sparse real data this project already has. Args: anomaly_wide: [date x cell_id], from climate_climatology. compute_anomalies. lookback_days: how much history feeds the encoder. horizons: which day-ahead offsets to produce targets for, e.g. [1, 5, 10, 30, 60, 90, 180]. Returns: X: [n_examples, lookback_days, n_cells] -- history window, same set of cells for every example (aligned columns). Y: [n_examples, n_horizons] -- target anomaly at each horizon, for the SPECIFIC cell/column indicated by `example_cell_idx`. example_cell_idx: [n_examples] -- which column (cell) index in X's last dimension this example's Y targets refer to (one training example = one cell's history predicting that SAME cell's own future, not cross-cell). dates: the anchor date (the "today") for each example, for traceability. """ max_horizon = max(horizons) values = anomaly_wide.values dates = anomaly_wide.index n_dates, n_cells = values.shape X_list, Y_list, cell_idx_list, anchor_dates = [], [], [], [] for t in range(lookback_days, n_dates - max_horizon): window = values[t - lookback_days:t, :] window_complete_mask = ~np.isnan(window).any(axis=0) # per cell if not window_complete_mask.any(): continue for cell in np.where(window_complete_mask)[0]: targets = np.array([values[t + h - 1, cell] for h in horizons]) if np.isnan(targets).all(): continue X_list.append(window[:, cell]) Y_list.append(targets) cell_idx_list.append(cell) anchor_dates.append(dates[t]) if not X_list: return (np.empty((0, lookback_days)), np.empty((0, len(horizons))), np.empty((0,), dtype=int), []) X = np.stack(X_list)[:, :, None] # [n_examples, lookback_days, 1] -- single-variable per call Y = np.stack(Y_list) return X, Y, np.array(cell_idx_list), anchor_dates def decay_gate(horizons: np.ndarray, tau: float) -> np.ndarray: """ exp(-h / tau) -- the monotonic decay curve gating the predicted anomaly toward zero as horizon grows. tau is learned per variable in the real model (see ClimateAnomalyModel); this standalone version is what's actually tested here, since the shape/behavior of the curve is the real design claim, independent of whether it's a fixed constant or a learned nn.Parameter wrapping this same formula. tau is the horizon (in days) at which the gate has decayed to 1/e (~37%) of its starting value -- a small tau means fast decay (appropriate for precipitation, per the real climatology results), a large tau means slow decay (appropriate for temperature). """ return np.exp(-horizons / tau) # --------------------------------------------------------------------------- # Model. Forward-pass shape/behavior verified only via the standalone # decay-gate check above; the full nn.Module has not been independently # run end-to-end. # --------------------------------------------------------------------------- if _HAS_TORCH: class ClimateAnomalyModel(nn.Module): """ Encoder: GRU over the lookback window of anomalies -> final hidden state per example. Decoder: [hidden_state, horizon_embedding] -> MLP -> (mean, log_variance) per horizon. Two complementary mechanisms handle "the model shouldn't be confidently wrong when there's no real signal" -- confirmed as a real, not hypothetical, problem: the first version of this model (mean-only, plain MSE) was measurably WORSE than pure climatology at several horizons for temp_C/solar_Wm2/ wind_speed_ms, while a persistence baseline diagnostic showed there was no real short-horizon signal for those variables to find in the first place -- the model was outputting confident, wrong guesses instead of admitting it didn't know. 1. The MEAN is passed through a learnable, monotonically decaying gate per variable (see decay_gate) -- a HARD architectural constraint: the model is structurally unable to output a large mean anomaly at a horizon its own learned tau says is past the point of real signal. Proven to work correctly on real data already: trained on snow_mm (a genuinely close-to-unpredictable variable), tau collapsed to 2.3 and the model's predictions converged to matching climatology almost exactly at every horizon beyond a few days. 2. The VARIANCE is predicted separately, NOT gated the same way, and trained via Gaussian NLL loss rather than plain MSE. This is what actually targets the diagnosed overconfidence problem: NLL rewards the model for outputting high predicted variance when it genuinely can't predict well, rather than being forced into a single point guess the loss then punishes regardless of whether real signal existed to make a better one. One model per climate variable, not a shared model across all nine -- the real climatology and persistence-diagnostic results showed genuinely different predictability profiles (runoff_mm beat climatology at every horizon; wind_v_ms/precip_mm/snow_mm essentially never did), and forcing one shared decay/variance behavior to serve both would misrepresent at least one of them. """ def __init__(self, hidden_dim: int = 32, horizon_embed_dim: int = 8, init_tau: float = 30.0): super().__init__() self.encoder = nn.GRU(input_size=1, hidden_size=hidden_dim, batch_first=True) self.horizon_embed = nn.Linear(1, horizon_embed_dim) self.decoder = nn.Sequential( nn.Linear(hidden_dim + horizon_embed_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, 2), # [mean, log_variance] per horizon ) # log_tau, not tau directly, so tau = exp(log_tau) is always # positive regardless of what gradient descent does to the # underlying parameter -- a negative or zero tau would make # the decay gate nonsensical (blow up or divide by zero). self.log_tau = nn.Parameter(torch.log(torch.tensor(float(init_tau)))) def forward(self, history: "torch.Tensor", horizons: "torch.Tensor") -> Tuple["torch.Tensor", "torch.Tensor"]: """ Args: history: [batch, lookback_days, 1] -- real anomaly history, already NaN-free (prepare_training_windows only emits complete windows). horizons: [n_horizons] -- day-ahead offsets to predict for, shared across the batch. Returns: (mean, log_variance), each [batch, n_horizons]. The mean is already gated by the learned decay curve; log_variance is the model's own learned, per-example uncertainty, deliberately not put through the same gate -- a model should be able to express high uncertainty at ANY horizon, including short ones, if a specific example genuinely warrants it, not only according to a fixed schedule. """ _, h_n = self.encoder(history) hidden = h_n.squeeze(0) # [batch, hidden_dim] batch_size = hidden.shape[0] n_horizons = horizons.shape[0] hidden_expanded = hidden.unsqueeze(1).expand(batch_size, n_horizons, hidden.shape[-1]) horizon_feat = self.horizon_embed(horizons.float().view(1, n_horizons, 1).expand(batch_size, n_horizons, 1)) decoder_input = torch.cat([hidden_expanded, horizon_feat], dim=-1) raw_output = self.decoder(decoder_input) # [batch, n_horizons, 2] raw_mean = raw_output[..., 0] log_var = raw_output[..., 1] tau = torch.exp(self.log_tau) gate = torch.exp(-horizons.float() / tau) # [n_horizons] mean = raw_mean * gate.unsqueeze(0) return mean, log_var def gaussian_nll_loss(mean: "torch.Tensor", log_var: "torch.Tensor", target: "torch.Tensor", min_log_var: float = -10.0) -> "torch.Tensor": """ Heteroscedastic Gaussian negative log-likelihood, NaN-masked (same masking principle used throughout this project's physics losses -- a target that doesn't exist for a given [example, horizon] pair must not silently poison every other real pair's gradient the same way an un-masked MSE mean would). NLL = 0.5 * log_var + 0.5 * (target - mean)^2 / exp(log_var) min_log_var clamps how confident the model is allowed to claim to be -- without a floor, a model can trivially drive NLL toward -infinity by predicting an extremely small variance on an example it happens to get right by chance, which is a real, well-known failure mode of unconstrained heteroscedastic regression, not a hypothetical one. """ mask = ~torch.isnan(target) if not mask.any(): return mean.sum() * float("nan") log_var_clamped = torch.clamp(log_var[mask], min=min_log_var) residual_sq = (target[mask] - mean[mask]) ** 2 return (0.5 * log_var_clamped + 0.5 * residual_sq / torch.exp(log_var_clamped)).mean() else: ClimateAnomalyModel = None # torch not installed; import guard, not a silent no-op gaussian_nll_loss = None