Normalize.py

import yaml
import numpy as np
import xarray as xr
import pandas as pd

from pathlib import Path
from Dataset import get_station_features, get_station_target


# ── fit poly and accumulate mean/var for input features ──────────────

def fit_X_scaler(train_idx, static_ds, time_indicators, variable, poly):
    """
    Pass over training stations.
    Fits PolynomialFeatures on first station, then accumulates 
    mean and variance of poly-transformed features incrementally.
    Returns mean_, var_ of shape (p,).
    """
    EPS = 1e-8
    mean_ = None
    var_  = None
    n_    = 0

    for count, i in enumerate(train_idx):
        X_i     = get_station_features(i, static_ds, time_indicators, variable)
        X_i_poly = poly.fit_transform(X_i) if count == 0 else poly.transform(X_i)
        n_i     = X_i_poly.shape[0]

        if mean_ is None:
            p      = X_i_poly.shape[1]
            mean_  = np.zeros(p)
            var_   = np.zeros(p)

        # Welford-style incremental mean and sum of squares
        mean_ += X_i_poly.sum(axis=0)
        var_  += (X_i_poly ** 2).sum(axis=0)
        n_    += n_i

        if (count + 1) % 100 == 0:
            print(f'Fit X scaler: {count + 1}/{len(train_idx)} stations', flush=True)

    mean_ /= n_
    var_   = var_ / n_ - mean_ ** 2  # E[x^2] - E[x]^2
    var_   = np.maximum(var_, EPS)    # numerical safety

    return mean_, var_


# ── fit accumulate mean/var for target ──────────────────────────────

def fit_y_scaler(train_idx, era5_ds, madis_ds, variable):
    """
    Pass over training stations.
    Accumulates mean and variance of target features incrementally.
    Returns mean_, var_ of shape (1,).
    """
    EPS = 1e-8
    mean_ = None
    var_  = None
    n_    = 0

    for count, i in enumerate(train_idx):
        y_i     = get_station_target(i, era5_ds, madis_ds, variable) # (n_time,)
        n_i     = y_i.shape[0]

        if mean_ is None:
            mean_ = np.zeros(1)
            var_  = np.zeros(1)

        mean_ += y_i.sum()
        var_  += (y_i ** 2).sum()
        n_    += n_i

        if (count + 1) % 100 == 0:
            print(f'Fit y scaler: {count + 1}/{len(train_idx)} stations', flush=True)
    
    mean_ /= n_
    var_   = var_ / n_ - mean_ ** 2  # E[x^2] - E[x]^2
    var_   = np.maximum(var_, EPS)    # numerical safety

    return mean_, var_