RandomForest.py

import numpy as np

from sklearn.ensemble import RandomForestRegressor
from Dataset import TMP_VARS, get_station_features, get_station_target


# ── RF hyperparameters (Table 3, state-independent) ───────────────────────

RF_PARAMS = {
    'TMP':  {'n_estimators': 100, 'max_depth': 13},
    'WIND': {'n_estimators':  50, 'max_depth': 10},
}


# ── subsample training data ───────────────────────────────────────────────────

def subsample_training_data(train_idx, static_ds, time_indicators,
                            era5_ds, madis_ds, variable,
                            subsample_frac=0.01, seed=42):
    """
    Collect a random subsample of training data across all train stations.
    Paper uses 1% of total training data for RF.

    Parameters
    ----------
    train_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
    subsample_frac : float, fraction of timesteps to sample per station
    seed           : int

    Returns
    -------
    X_sub : np.ndarray, shape (n_samples, n_features)
    y_sub : np.ndarray, shape (n_samples,)
    """
    rng    = np.random.default_rng(seed)
    n_time = len(time_indicators['times'])
    n_sub  = max(1, int(n_time * subsample_frac))

    X_list, y_list = [], []

    for count, i in enumerate(train_idx):
        # random time indices for this station
        time_idx = rng.choice(n_time, size=n_sub, replace=False)
        time_idx.sort()

        X_i = get_station_features(i, static_ds, time_indicators, variable)
        y_i = get_station_target(i, era5_ds, madis_ds, variable)

        X_list.append(X_i[time_idx])
        y_list.append(y_i[time_idx])

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

    X_sub = np.concatenate(X_list, axis=0)  # (n_train_stations * n_sub, n_features)
    y_sub = np.concatenate(y_list, axis=0)  # (n_train_stations * n_sub,)

    print(f'[{variable}] Subsample shape: X={X_sub.shape}  y={y_sub.shape}', flush=True)
    return X_sub, y_sub


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

def eval_stations_rf(idx, static_ds, time_indicators, era5_ds, madis_ds,
                     variable, model):
    """
    Compute RMSE over stations in idx.

    Parameters
    ----------
    idx   : np.ndarray of station indices
    model : fitted RandomForestRegressor

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

    for i in idx:
        X_i     = get_station_features(i, static_ds, time_indicators, variable)
        y_i     = get_station_target(i, era5_ds, madis_ds, variable)
        e_hat_i = model.predict(X_i)

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

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


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

def train_rf(train_idx, val_idx, static_ds, time_indicators,
             era5_ds, madis_ds, variable,
             subsample_frac=0.01, seed=42):
    """
    Train RandomForestRegressor on subsampled training data.
    Uses paper hyperparameters from Table 3 (state-independent config).

    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
    subsample_frac  : float, fraction of timesteps per station
    seed            : int

    Returns
    -------
    model : fitted RandomForestRegressor
    """
    params = RF_PARAMS['TMP'] if variable in TMP_VARS else RF_PARAMS['WIND']

    # subsample training data
    print(f'[{variable}] Subsampling training data...', flush=True)
    X_sub, y_sub = subsample_training_data(
        train_idx, static_ds, time_indicators,
        era5_ds, madis_ds, variable,
        subsample_frac=subsample_frac, seed=seed,
    )

    # train RF with paper hyperparameters
    print(f'[{variable}] Training RF: {params}', flush=True)
    model = RandomForestRegressor(
        n_estimators=params['n_estimators'],
        max_depth=params['max_depth'],
        n_jobs=-1,
        random_state=seed,
    )
    model.fit(X_sub, y_sub)

    # evaluate on validation set
    print(f'[{variable}] Evaluating on validation set...', flush=True)
    val_rmse = eval_stations_rf(
        val_idx, static_ds, time_indicators,
        era5_ds, madis_ds, variable, model,
    )
    print(f'[{variable}] Val RMSE={val_rmse}', flush=True)

    return model


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

def predict_rf(i, static_ds, time_indicators, era5_ds, variable, model):
    """
    Compute bias-corrected ERA5 for a single station i.

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

    Returns
    -------
    f_hat : np.ndarray, shape (n_time,), float32
    """
    X_i   = get_station_features(i, static_ds, time_indicators, variable)
    e_hat = model.predict(X_i)
    f_i   = era5_ds[variable].values[i].astype(np.float64)
    f_hat = (f_i - e_hat).astype(np.float32)

    return f_hat