Linear.py

import numpy as np

from Normalize import fit_X_scaler
from sklearn.preprocessing import PolynomialFeatures
from Dataset import get_station_features, get_station_target


# ── accumulate XtX and Xty ───────────────────────────────────────────

def accumulate_normal_equations(idx, static_ds, time_indicators,
                                era5_ds, madis_ds, variable,
                                poly, mean_, var_):
    """
    Accumulate XtX and Xty over stations in idx.

    Parameters
    ----------
    idx             : np.ndarray of station indices
    static_ds       : xr.Dataset, shape (station_id,)
    time_indicators : dict from compute_time_indicators()
    era5_ds         : xr.Dataset, already time-sliced
    madis_ds        : xr.Dataset, already time-sliced
    variable        : str
    poly            : fitted PolynomialFeatures
    mean_           : np.ndarray, shape (p,)
    var_            : np.ndarray, shape (p,)

    Returns
    -------
    XtX : np.ndarray, shape (p, p)
    Xty : np.ndarray, shape (p,)
    """
    p   = len(mean_)
    XtX = np.zeros((p, p))
    Xty = np.zeros(p)

    for count, i in enumerate(idx):
        X_i      = get_station_features(i, static_ds, time_indicators, variable)
        X_i_poly = poly.transform(X_i)
        X_i_sc   = (X_i_poly - mean_) / np.sqrt(var_)

        y_i      = get_station_target(i, era5_ds, madis_ds, variable)

        XtX     += X_i_sc.T @ X_i_sc
        Xty     += X_i_sc.T @ y_i

        if (count + 1) % 100 == 0:
            print(f'Accumulate {variable}: {count + 1}/{len(idx)} stations', flush=True)

    return XtX, Xty


# ── solve ridge ───────────────────────────────────────────────────────────────

def solve_ridge(XtX, Xty, alpha):
    """
    Solve β = (XtX + α*I)^{-1} Xty.

    Parameters
    ----------
    XtX   : np.ndarray, shape (p, p)
    Xty   : np.ndarray, shape (p,)
    alpha : float, regularization strength

    Returns
    -------
    beta : np.ndarray, shape (p,)
    """
    p    = XtX.shape[0]
    A    = XtX + alpha * np.eye(p)
    beta = np.linalg.solve(A, Xty)
    return beta


# ── evaluate on a set of stations ────────────────────────────────────────────

def eval_stations_linear(idx, static_ds, time_indicators, era5_ds, madis_ds,
                  variable, poly, mean_, var_, beta):
    """
    Compute RMSE over stations in idx.

    Parameters
    ----------
    idx   : np.ndarray of station indices
    beta  : np.ndarray, shape (p,)

    Returns
    -------
    rmse : float
    """
    sq_err = []

    for i in idx:
        X_i      = get_station_features(i, static_ds, time_indicators, variable)
        X_i_poly = poly.transform(X_i)
        X_i_sc   = (X_i_poly - mean_) / np.sqrt(var_)

        y_i      = get_station_target(i, era5_ds, madis_ds, variable)
        e_hat_i  = X_i_sc @ beta

        residual = y_i - e_hat_i
        sq_err.append((residual ** 2).mean())

    return np.sqrt(np.mean(sq_err))


# ── train ─────────────────────────────────────────────────────────────────────

def train_linear(train_idx, val_idx, static_ds, time_indicators,
                 era5_ds, madis_ds, variable):
    """
    Full incremental Ridge training pipeline with alpha selection on val set.

    Parameters
    ----------
    train_idx       : np.ndarray of station indices
    val_idx         : np.ndarray of station indices
    static_ds       : xr.Dataset, shape (station_id,)
    time_indicators : dict from compute_time_indicators()
    era5_ds         : xr.Dataset, already time-sliced
    madis_ds        : xr.Dataset, already time-sliced
    variable        : str

    Returns
    -------
    poly   : fitted PolynomialFeatures
    mean_  : np.ndarray, shape (p,)
    var_   : np.ndarray, shape (p,)
    beta   : np.ndarray, shape (p,) — best solution
    """

    alphas = [0.01, 0.1, 1.0]

    poly = PolynomialFeatures(degree=2, include_bias=True)

    # pass 1: fit poly + accumulate scaler stats on train stations
    print(f'[{variable}] Pass 1: fitting scaler...', flush=True)
    mean_, var_ = fit_X_scaler(
        train_idx, static_ds, time_indicators, variable, poly
    )

    # pass 2: accumulate XtX and Xty on train stations
    print(f'[{variable}] Pass 2: accumulating normal equations...', flush=True)
    XtX, Xty = accumulate_normal_equations(
        train_idx, static_ds, time_indicators,
        era5_ds, madis_ds, variable,
        poly, mean_, var_
    )

    # alpha selection on validation set
    print(f'[{variable}] Selecting alpha on validation set...', flush=True)
    best_alpha, best_rmse, best_beta = None, np.inf, None

    for alpha in alphas:
        beta     = solve_ridge(XtX, Xty, alpha)
        val_rmse = eval_stations_linear(
            val_idx, static_ds, time_indicators,
            era5_ds, madis_ds, variable,
            poly, mean_, var_, beta
        )
        print(f'alpha={alpha}  val RMSE={val_rmse}', flush=True)
        if val_rmse < best_rmse:
            best_rmse  = val_rmse
            best_alpha = alpha
            best_beta  = beta

    print(f'[{variable}] Best alpha: {best_alpha}  val RMSE: {best_rmse}', flush=True)
    return poly, mean_, var_, best_beta


# ── predict ───────────────────────────────────────────────────────────────────

def predict_linear(i, static_ds, time_indicators, era5_ds,
                   variable, poly, mean_, var_, beta):
    """
    Compute bias-corrected ERA5 for a single station i.

    Parameters
    ----------
    i        : int, station index
    era5_ds  : xr.Dataset, already time-sliced

    Returns
    -------
    f_hat : np.ndarray, shape (n_time,), float32
    """
    X_i      = get_station_features(i, static_ds, time_indicators, variable)
    X_i_poly = poly.transform(X_i)
    X_i_sc   = (X_i_poly - mean_) / np.sqrt(var_)

    e_hat    = X_i_sc @ beta
    f_i      = era5_ds[variable].values[i].astype(np.float64)
    f_hat    = (f_i - e_hat).astype(np.float32)

    return f_hat