ashvale-station/ashvale/models/nowcast.py

# Copyright 2026 Kemal Yaylali
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Multi-horizon forecasting: one direct head per (target, horizon).

Direct rather than recursive. A recursive one-step model iterated 288
times to reach 24 hours compounds its own bias into a beautifully smooth
lie. Direct heads cost more memory (six horizons x three targets = 18
small models, about 150 kB total) and are worth every byte.

Each head predicts a *delta from now*, then the ensemble blends three
opinions with weights that are themselves learned online:

    persistence  : it will be exactly as it is now
    climatology  : it will be whatever this hour of this day usually is
    learned RLS  : it will be now plus what the regressors imply

Persistence wins at 15 minutes. Climatology wins at 24 hours. The RLS
head wins in the middle, which is exactly the region a physical
forecaster finds hardest. The blend weights are updated by exponentiated
gradient (Hedge), so the ensemble is never worse than its best member by
more than a log factor, and it re-weights itself within a day when the
season turns.
"""

from __future__ import annotations

from typing import Dict, List, Optional, Tuple

import numpy as np

from ..features import N_FEATURES, Standardiser, supervised_pairs
from .rls import AdaptiveConformal, RecursiveLeastSquares

MEMBERS = ("persistence", "climatology", "learned", "setpoint")

# Scored forecasts before the Hedge loss scale switches from this sample's
# worst loss to the running member MAE. Until member_mae has seen anything it
# is zeros, and dividing by that would hand every member the same penalty.
_SCALE_WARMUP = 20


class ForecastHead:
    """One target, one horizon."""

    def __init__(self, target: str, horizon_s: int, n_features: int = N_FEATURES,
                 forgetting: float = 0.9985, delta: float = 100.0,
                 alpha: float = 0.10, conformal_window: int = 400,
                 gamma: float = 0.01, hedge_eta: float = 0.35):
        self.target = target
        self.horizon_s = int(horizon_s)
        self.model = RecursiveLeastSquares(n_features, forgetting, delta)
        self.conformal = AdaptiveConformal(alpha, conformal_window, gamma)
        self.weights = np.ones(len(MEMBERS)) / len(MEMBERS)
        self.eta = float(hedge_eta)
        self.member_mae = np.zeros(len(MEMBERS))
        self.n_scored = 0
        # Validity time of the last outcome the Hedge weights learned from.
        self.last_hedge_ts = -np.inf

    # -------------------------------------------------------- prediction

    def predict(self, x: np.ndarray, anchor: float,
                climatology_delta: float = 0.0,
                setpoint_delta: float = 0.0) -> Dict[str, float]:
        learned_delta = self.model.predict(x)
        deltas = np.array([0.0, float(climatology_delta), float(learned_delta),
                           float(setpoint_delta)])
        blended = float(np.dot(self.weights, deltas))
        mu = float(anchor + blended)
        sigma = self.model.predict_std(x, self.model.noise_var)
        lo, hi = self.conformal.interval(mu, fallback_sigma=sigma)
        return {
            "mu": mu,
            "lo": lo,
            "hi": hi,
            "sigma": sigma,
            "delta": blended,
            "members": {m: float(anchor + d) for m, d in zip(MEMBERS, deltas)},
            "weights": {m: float(w) for m, w in zip(MEMBERS, self.weights)},
        }

    # ---------------------------------------------------------- learning

    def refit_step(self, x: np.ndarray, anchor: float, truth: float) -> None:
        """One regression update from a replayed historical pair.

        This deliberately touches nothing but the RLS. The blend weights and
        the conformal calibrator are statements about how this head's issued
        forecasts actually turned out, and a refit is not an outcome: it is the
        same week of weather being read again.

        Measured on 8.2 days of real station data, the previous arrangement
        (fit() calling a combined learn()) had put 977,078 Hedge updates through
        the 15 minute head from 758 distinct supervised pairs, a factor of 1,289,
        and 296,715 through the 1 day head from 12 pairs, a factor of 24,726.
        Hedge is multiplicative, so an edge far too small to be real compounds
        to certainty: twelve of twelve temperature and humidity heads had
        collapsed onto climatology at a weight of 0.991 or above. The ACI
        integrator, which moves by gamma per observation, had likewise pinned
        against its clips, giving a 6 hour band narrower than the 3 hour one.
        Feeding these two from verify() instead is worth 14.4% of MAE across
        17 of 18 heads, and takes mean absolute coverage error from 0.188
        to 0.059.
        """
        self.model.update(x, truth - anchor)

    def observe_outcome(self, members: np.ndarray, truth: float,
                        covered: Optional[bool] = None,
                        valid_ts: Optional[float] = None) -> float:
        """One matured forecast, scored against what actually happened.

        `members` are the four point predictions this head issued, recovered
        from the forecasts table. The learned one cannot be recomputed here
        because the RLS has moved on since.
        """
        member_pred = np.asarray(members, dtype=float)
        losses = np.abs(member_pred - truth)

        # Score the blend with the weights predict() actually used, before this
        # outcome moves them. Doing it after is look-ahead: the residual handed
        # to the conformal calibrator is then better than anything the
        # forecaster can produce, so the intervals come out about 2% too narrow.
        # Coverage survived it only because ACI notices the extra misses and
        # reopens the band, which is a correction that should not be needed.
        blended = float(np.dot(self.weights, member_pred))
        residual = truth - blended
        self.conformal.observe(residual, covered=covered)

        # The Hedge weights take an outcome only if it does not overlap the last
        # one they took. Forecasts are issued every retrain tick, so at the 1 day
        # horizon roughly two hundred of them per day resolve against what is
        # very nearly the same outcome. The conformal window can absorb that,
        # since a quantile over duplicated scores is merely over-confident about
        # its sample size, but exponentiated gradient cannot: it would apply the
        # same evidence two hundred times and saturate. This is the stride rule
        # from fit(), applied on the scoring side.
        if valid_ts is not None:
            if valid_ts - self.last_hedge_ts < self.horizon_s:
                return residual
            self.last_hedge_ts = float(valid_ts)

        # Hedge / exponentiated gradient on normalised losses.
        #
        # The scale must be a stable quantity, not this sample's worst loss.
        # Dividing by max(losses) means that on a quiet step where every member
        # agrees to within 0.01 C, whichever one happens to be worst still takes
        # the full exp(-eta) penalty, exactly as if it had been wrong by 5 C, so
        # the weights churn on noise. Hedge's regret bound assumes a fixed loss
        # range. Normalising by the running member MAE instead is worth 1.8% of
        # MAE across 106 of 126 heads on real data, with coverage unchanged.
        if self.n_scored > _SCALE_WARMUP:
            scale = max(float(np.mean(self.member_mae)), 1e-6)
        else:
            scale = max(float(np.max(losses)), 1e-6)
        self.weights *= np.exp(-self.eta * losses / scale)
        self.weights = np.clip(self.weights, 1e-4, None)
        self.weights /= self.weights.sum()

        self.member_mae = 0.98 * self.member_mae + 0.02 * losses
        self.n_scored += 1
        return residual

    def to_dict(self) -> Dict:
        return {"target": self.target, "horizon_s": self.horizon_s,
                "model": self.model.to_dict(), "conformal": self.conformal.to_dict(),
                "weights": self.weights.tolist(), "eta": self.eta,
                "member_mae": self.member_mae.tolist(), "n_scored": self.n_scored,
                "last_hedge_ts": (float(self.last_hedge_ts)
                                  if np.isfinite(self.last_hedge_ts) else None)}

    @classmethod
    def from_dict(cls, s: Dict) -> "ForecastHead":
        h = cls(s["target"], s["horizon_s"])
        h.model = RecursiveLeastSquares.from_dict(s["model"])
        h.conformal = AdaptiveConformal.from_dict(s["conformal"])
        w = np.array(s["weights"], dtype=float)
        if w.size != len(MEMBERS):
            # A saved head from before the setpoint member existed. Reinitialise
            # uniformly rather than guessing: the Hedge weights re-converge in
            # about a day, which is far cheaper than silently mismatching a
            # member to the wrong loss and corrupting every blend until someone
            # notices.
            w = np.ones(len(MEMBERS)) / len(MEMBERS)
        h.weights = w
        h.eta = s["eta"]
        # A state file with no last_hedge_ts was written before the Hedge
        # weights and the ACI integrator were cut off from fit()'s replay, so
        # everything they hold is the product of the same week of weather read
        # about a thousand times: weights pinned on one member at 0.99, alpha
        # against a clip, member_mae an EMA over a million duplicated steps.
        # None of that is evidence, and it does not decay on its own, because
        # Hedge needs about twenty independent outcomes to climb back off the
        # 1e-4 floor and the 1 d head sees one a day.
        #
        # The conformal scores are kept. They were also fed by the replay, so
        # they are biased a little narrow, but they are absolute residuals of
        # roughly the right size and the window refreshes within about two days
        # of real outcomes. Clearing them instead would drop the long horizons
        # onto 1.645*sigma for nine days, which is how you get a plus or minus
        # of 115% relative humidity.
        if "last_hedge_ts" not in s:
            h.weights = np.ones(len(MEMBERS)) / len(MEMBERS)
            h.member_mae = np.zeros(len(MEMBERS))
            h.n_scored = 0
            h.conformal.alpha = h.conformal.alpha_target
            h.last_hedge_ts = -np.inf
            return h
        lh = s.get("last_hedge_ts")
        h.last_hedge_ts = -np.inf if lh is None else float(lh)
        mae = np.array(s["member_mae"], dtype=float)
        # Same migration as the weights. Missing this one did not fail on load,
        # it failed later inside the Hedge update on a shape mismatch, which is
        # a worse place to find out.
        if mae.size != len(MEMBERS):
            mae = np.zeros(len(MEMBERS))
        h.member_mae = mae
        h.n_scored = s.get("n_scored", 0)
        return h


class NowcastEnsemble:
    """The full bank of heads plus the shared feature standardiser."""

    def __init__(self, targets: Tuple[str, ...], horizons_s: Tuple[int, ...],
                 cfg_model):
        self.targets = tuple(targets)
        self.horizons = tuple(int(h) for h in horizons_s)
        self.cfg = cfg_model
        self.grid_s = int(cfg_model.grid_s)
        self.scaler = Standardiser(N_FEATURES)
        self.heads: Dict[Tuple[str, int], ForecastHead] = {
            (t, h): ForecastHead(
                t, h, N_FEATURES, cfg_model.rls_forgetting, cfg_model.rls_delta,
                cfg_model.conformal_alpha, cfg_model.conformal_window,
                cfg_model.conformal_gamma,
            )
            for t in self.targets for h in self.horizons
        }
        self.trained_rows = 0
        self.min_pairs = int(getattr(cfg_model, "min_pairs_per_head", 12))
        # Which phase of the stride this refit starts on. Rotated so that over
        # successive retrains every offset is eventually trained on, rather
        # than the model permanently seeing one sample in `steps` forever.
        self.refit_phase = 0

    # ------------------------------------------------------------ train

    def fit(self, X: np.ndarray, valid: np.ndarray, series: Dict[str, np.ndarray],
            passes: int = 1, max_pairs: int = 2500) -> Dict[str, int]:
        """Batch-update every head from history.

        `max_pairs` bounds the work per head to the most recent samples.
        This is not a shortcut: with a forgetting factor of 0.9985 the
        effective memory is about 11 hours, so the 4000th-most-recent
        sample carries a weight of roughly e^-6. Training on it costs
        real seconds on a Cortex-A53 and buys nothing measurable.
        """
        """Batch pass over history. Called on startup and every retrain tick."""
        if X.shape[0] < 10:
            return {"rows": 0}
        # A refit starts from the prior. Without this, every retrain tick replays
        # the same history into a live filter, and RLS with forgetting reads that
        # as new evidence each time: measured on a real station after 1.5 days,
        # 453 grid rows had produced 64,676 updates, cond(P) of 3.1e9 and a
        # weight vector of norm 1680 whose two largest entries were the annual
        # harmonics the record cannot yet resolve. The result was a six hour
        # forecast of 53 C in a 24 C room, with a plus or minus of 0.43.
        #
        # The conformal calibrators and the Hedge weights are left alone here,
        # and refit_step is what enforces that. They are earned from scored
        # forecasts in verify(), not from this regression. The climatology and
        # setpoint members used to be evaluated in this loop purely to feed
        # them, which is why this method no longer needs either.
        for head in self.heads.values():
            head.model.reset()
        self.scaler.partial_fit(X[valid][:: max(1, X.shape[0] // 2000)])
        Xs = self.scaler.transform(X)

        counts = {}
        for target in self.targets:
            y = series[target]
            for h in self.horizons:
                steps = max(int(round(h / self.grid_s)), 1)
                Xa, dy, anchor = supervised_pairs(Xs, valid, y, steps)
                if Xa.shape[0] < 5:
                    counts[f"{target}@{h}"] = 0
                    continue
                if Xa.shape[0] > max_pairs:
                    Xa, dy, anchor = Xa[-max_pairs:], dy[-max_pairs:], anchor[-max_pairs:]
                head = self.heads[(target, h)]
                # One pair per horizon, not one per grid row. Adjacent pairs at
                # the 1 d horizon share 287 of their 288 samples, so training on
                # every row hands the filter the same outcome 288 times and RLS
                # with forgetting reads each as fresh evidence. A 400-score
                # conformal window then holds 1.4 independent outcomes while
                # believing it holds 400.
                #
                # This is not a compute shortcut that costs accuracy. Measured
                # walk-forward on four days of real station data, striding cut
                # MAE at every horizon past an hour (temperature 6h -30%,
                # humidity 6h -48%, pressure 12h -68%) with coverage unchanged,
                # and made the fit 12x faster. The redundancy was not merely
                # wasted work, it was collapsing P onto the repeated direction.
                stride = steps
                if stride > 1 and Xa.shape[0] // stride < self.min_pairs:
                    # A long horizon on a short record would otherwise train
                    # on one or two pairs, which is worse than the redundancy
                    # it avoids. The floor was chosen by sweeping it over five
                    # train splits of real data: 12 was best at every horizon,
                    # and the apparent 1 d regressions at other values were
                    # noise, since a 1 d head on four days of record is fitted
                    # and scored on well under two independent outcomes.
                    stride = max(1, Xa.shape[0] // self.min_pairs)
                idx = np.arange(self.refit_phase % stride, Xa.shape[0], stride)
                if idx.size > max_pairs:
                    idx = idx[-max_pairs:]
                for _ in range(max(int(passes), 1)):
                    for i in idx:
                        head.refit_step(Xa[i], anchor[i], anchor[i] + dy[i])
                counts[f"{target}@{h}"] = int(idx.size)
        self.trained_rows = int(X.shape[0])
        self.refit_phase += 1
        return counts

    # --------------------------------------------------------- inference

    def forecast(self, x_raw: np.ndarray, anchors: Dict[str, float], now: float,
                 climatology=None, setpoint_fn=None) -> Dict[str, Dict[int, Dict[str, float]]]:
        x = self.scaler.transform(np.atleast_2d(x_raw))[0]
        out: Dict[str, Dict[int, Dict[str, float]]] = {}
        for target in self.targets:
            anchor = float(anchors.get(target, 0.0))
            out[target] = {}
            for h in self.horizons:
                clim_delta = 0.0
                if climatology is not None and climatology.ready:
                    clim_delta = float(climatology.predict(target, np.array([now + h]))[0]
                                       - climatology.predict(target, np.array([now]))[0])
                sp = setpoint_fn(target, h, anchor) if setpoint_fn else 0.0
                out[target][h] = self.heads[(target, h)].predict(x, anchor, clim_delta, sp)
        return out

    def diagnostics(self) -> List[Dict]:
        rows = []
        for (target, h), head in sorted(self.heads.items()):
            rows.append({
                "target": target,
                "horizon_s": h,
                "n_updates": head.model.n_updates,
                "n_scored": head.n_scored,
                "weights": {m: round(float(w), 3) for m, w in zip(MEMBERS, head.weights)},
                "member_mae": {m: round(float(v), 3) for m, v in zip(MEMBERS, head.member_mae)},
                "conformal_alpha": round(head.conformal.alpha, 4),
                "conformal_halfwidth": round(float(head.conformal.quantile()), 3)
                if np.isfinite(head.conformal.quantile()) else None,
                "coverage": round(head.conformal.empirical_coverage, 3)
                if np.isfinite(head.conformal.empirical_coverage) else None,
            })
        return rows

    def to_dict(self) -> Dict:
        return {
            "targets": list(self.targets),
            "horizons": list(self.horizons),
            "grid_s": self.grid_s,
            "scaler": self.scaler.to_dict(),
            "heads": [h.to_dict() for h in self.heads.values()],
            "trained_rows": self.trained_rows,
            "refit_phase": self.refit_phase,
        }

    def load_dict(self, s: Dict) -> None:
        self.scaler = Standardiser.from_dict(s["scaler"])
        for hs in s["heads"]:
            head = ForecastHead.from_dict(hs)
            key = (head.target, head.horizon_s)
            if key not in self.heads:
                continue          # a target or horizon this build no longer has
            self.heads[key] = head
        self.trained_rows = s.get("trained_rows", 0)
        self.refit_phase = int(s.get("refit_phase", 0))
        self._apply_config_tuning()

    def _apply_config_tuning(self) -> None:
        """Tuning comes from config; only the estimate comes from the file.

        Every from_dict below this point restores its knobs alongside its data:
        lambda, delta and p_max in the RLS, alpha, gamma and the window in the
        conformal calibrator. So each of those was immutable in the field. Edit
        config.yaml, restart, and the state file quietly puts the old value
        back, which looks exactly like a change that had no effect. The same
        defect cost a Kalman retune here before it was found.
        """
        c = self.cfg
        for head in self.heads.values():
            head.model.lam = float(c.rls_forgetting)
            head.model.delta = float(c.rls_delta)
            head.conformal.retune(c.conformal_alpha, c.conformal_gamma,
                                  c.conformal_window)