MLP.py

import torch
import numpy as np
import torch.nn as nn

from Normalize import fit_X_scaler, fit_y_scaler
from torch.utils.data import Dataset, DataLoader
from sklearn.preprocessing import PolynomialFeatures
from Dataset import get_station_features, get_station_target


# ── dataset ───────────────────────────────────────────────────────────────────

class StationDataset(Dataset):
    """
    PyTorch Dataset that pre-loads subsampled features and targets
    for a set of stations into memory as float32 tensors.

    Parameters
    ----------
    station_indices : 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
    mean_           : np.ndarray, shape (n_features,) — feature mean
    var_            : np.ndarray, shape (n_features,) — feature variance
    y_mean_         : float — target mean
    y_var_          : float — target variance
    subsample_frac  : float, fraction of timesteps to sample per station
    seed            : int
    """
    def __init__(self, station_indices, static_ds, time_indicators,
                 era5_ds, madis_ds, variable,
                 mean_, var_, y_mean_, y_var_,
                 subsample_frac=0.1, seed=42):

        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(station_indices):
            X_i = get_station_features(i, static_ds, time_indicators, variable)
            y_i = get_station_target(i, era5_ds, madis_ds, variable)

            # subsample timesteps
            time_idx = rng.choice(n_time, size=n_sub, replace=False)
            time_idx.sort()

            # normalize X
            X_i_sc = (X_i[time_idx] - mean_) / np.sqrt(var_)

            # normalize y
            y_i_sc = (y_i[time_idx] - y_mean_) / np.sqrt(y_var_)

            X_list.append(X_i_sc.astype(np.float32))
            y_list.append(y_i_sc.astype(np.float32))

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

        X_all = np.concatenate(X_list, axis=0)
        y_all = np.concatenate(y_list, axis=0)

        self.X = torch.from_numpy(X_all)
        self.y = torch.from_numpy(y_all)

        print(f'[{variable}] Dataset shape: X={self.X.shape}  y={self.y.shape}', flush=True)

    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]


# ── model architecture ────────────────────────────────────────────────────────

class MLPModel(nn.Module):
    """
    4-layer MLP with Swish activation.
    Architecture from paper: 4 hidden layers x 32 neurons ~ 4000 parameters.

    Parameters
    ----------
    n_features : int, number of input features
    n_hidden   : int, number of neurons per hidden layer (default: 32)
    n_layers   : int, number of hidden layers (default: 4)
    """
    def __init__(self, n_features, n_hidden=32, n_layers=4):
        super().__init__()

        layers = []
        in_dim = n_features
        for _ in range(n_layers):
            layers.append(nn.Linear(in_dim, n_hidden))
            layers.append(nn.SiLU())   # SiLU == Swish
            in_dim = n_hidden
        layers.append(nn.Linear(in_dim, 1))  # output layer, no activation

        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x).squeeze(-1)  # (batch,)


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

def train_mlp(train_idx, val_idx, static_ds, time_indicators,
              era5_ds, madis_ds, variable,
              subsample_frac=0.1, seed=42,
              n_hidden=32, n_layers=4,
              lr=1e-3, batch_size=128, max_epochs=20):
    """
    Train MLP with Adam optimizer and y normalization.
    Matches paper: 4 hidden layers, 32 neurons, Swish, Adam lr=1e-3,
    batch_size=128, 20 epochs.

    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
    n_hidden        : int, neurons per hidden layer
    n_layers        : int, number of hidden layers
    lr              : float, learning rate
    batch_size      : int
    max_epochs      : int

    Returns
    -------
    model   : trained MLPModel
    mean_   : np.ndarray, shape (n_features,) — X mean
    var_    : np.ndarray, shape (n_features,) — X variance
    y_mean_ : float — target mean
    y_var_  : float — target variance
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f'[{variable}] Using device: {device}', flush=True)

    # fit X scaler on train stations
    print(f'[{variable}] Fitting X scaler...', flush=True)
    poly   = PolynomialFeatures(degree=1, include_bias=False)
    mean_, var_ = fit_X_scaler(train_idx, static_ds, time_indicators, variable, poly)

    # fit y scaler on train stations
    print(f'[{variable}] Fitting y scaler...', flush=True)
    y_mean_, y_var_ = fit_y_scaler(train_idx, era5_ds, madis_ds, variable)
    y_mean_ = float(y_mean_)
    y_var_  = float(y_var_)

    # build datasets and loaders
    print(f'[{variable}] Loading train dataset...', flush=True)
    train_ds = StationDataset(
        train_idx, static_ds, time_indicators,
        era5_ds, madis_ds, variable,
        mean_, var_, y_mean_, y_var_,
        subsample_frac=subsample_frac, seed=seed,
    )
    print(f'[{variable}] Loading val dataset...', flush=True)
    val_ds = StationDataset(
        val_idx, static_ds, time_indicators,
        era5_ds, madis_ds, variable,
        mean_, var_, y_mean_, y_var_,
        subsample_frac=subsample_frac, seed=seed,
    )

    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=2)
    val_loader   = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=2)

    # build model
    n_features = train_ds.X.shape[1]
    model      = MLPModel(n_features, n_hidden=n_hidden, n_layers=n_layers).to(device)
    print(f'[{variable}] MLP parameters: '
          f'{sum(p.numel() for p in model.parameters())}', flush=True)

    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    criterion = nn.MSELoss()

    # training loop
    for epoch in range(max_epochs):
        # --- train ---
        model.train()
        train_losses = []
        for X_batch, y_batch in train_loader:
            X_batch, y_batch = X_batch.to(device), y_batch.to(device)
            optimizer.zero_grad()
            loss = criterion(model(X_batch), y_batch)
            loss.backward()
            optimizer.step()
            train_losses.append(loss.item())

        # --- validate ---
        model.eval()
        val_losses = []
        with torch.no_grad():
            for X_batch, y_batch in val_loader:
                X_batch, y_batch = X_batch.to(device), y_batch.to(device)
                val_losses.append(criterion(model(X_batch), y_batch).item())

        # note: losses are in normalized y space
        print(f'[{variable}] Epoch {epoch + 1:02d}/{max_epochs} '
              f'train_loss={np.mean(train_losses)}  '
              f'val_loss={np.mean(val_losses)}  '
              f'val_RMSE={np.sqrt(np.mean(val_losses)) * np.sqrt(y_var_)}',
              flush=True)

    model.eval()
    return model, mean_, var_, y_mean_, y_var_


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

def predict_mlp(i, static_ds, time_indicators, era5_ds,
                variable, model, mean_, var_, y_mean_, y_var_):
    """
    Compute bias-corrected ERA5 for a single station i.

    Parameters
    ----------
    i       : int, station index
    era5_ds : xr.Dataset, already time-sliced
    model   : trained MLPModel
    mean_   : np.ndarray, shape (n_features,)
    var_    : np.ndarray, shape (n_features,)
    y_mean_ : float — target mean
    y_var_  : float — target variance

    Returns
    -------
    f_hat : np.ndarray, shape (n_time,), float32
    """
    device = next(model.parameters()).device

    X_i    = get_station_features(i, static_ds, time_indicators, variable)
    X_i_sc = (X_i - mean_) / np.sqrt(var_)
    X_t    = torch.from_numpy(X_i_sc.astype(np.float32)).to(device)

    model.eval()
    with torch.no_grad():
        e_hat_normalized = model(X_t).cpu().numpy()

    # denormalize predicted error
    e_hat = e_hat_normalized * np.sqrt(y_var_) + y_mean_

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

    return f_hat