src/graph/climate_anomaly_model.py

"""
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