Dataset.py

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

from pathlib import Path


# ── constants ────────────────────────────────────────────────────────────────

STATIC_TMP = ['sdor', 'lsm', 'cvl', 'cvh', 'tvl', 'tvh', 'station_elevation', 'elev_diff', 'log_elev_diff', 'latitude', 'cos_longitude']
STATIC_WIND = ['sdor', 'lsm', 'cvl', 'cvh', 'latitude', 'longitude', 'cos_longitude', 'elev_diff', 'log_elev_diff']
TMP_VARS  = ['t2m', 'd2m']
WIND_VARS = ['u10', 'v10']
VARIABLES = TMP_VARS + WIND_VARS

# ── station split ─────────────────────────────────────────────────────

def load_station_ids(split_yaml_path):
    with open(split_yaml_path, 'r') as f:
        split = yaml.safe_load(f)

    station_ids = split['train'] + split['val'] + split['test']
    return np.array(station_ids)

def split_stations(split_yaml_path):
    """Load train/val/test station split from yaml file."""
    with open(split_yaml_path, 'r') as f:
        split = yaml.safe_load(f)

    station_ids = split['train'] + split['val'] + split['test']

    id_to_idx = {sid: i for i, sid in enumerate(station_ids)}

    train_idx = np.array([id_to_idx[s] for s in split['train']])
    val_idx   = np.array([id_to_idx[s] for s in split['val']  ])
    test_idx  = np.array([id_to_idx[s] for s in split['test'] ])

    print(f'Train: {len(train_idx)}  Val: {len(val_idx)}  Test: {len(test_idx)} stations', flush=True)

    return train_idx, val_idx, test_idx


# ── time indicators ───────────────────────────────────────────────────

def compute_time_indicators(lats, lons, start_year=2020, end_year=2023):
    """
    Compute time indicators for state-independent config.

    Parameters
    ----------
    lats : np.ndarray, shape (n_stations,)
    lons : np.ndarray, shape (n_stations,)
    start_year : int, default=2020
    end_year : int, default=2023

    Returns
    -------
    dict of arrays:
        time-only indicators  → (n_time,)
        time-location indicators → (n_stations, n_time)
    """
    times      = np.arange(
        np.datetime64(f'{start_year}-01-01T00:00'),
        np.datetime64(f'{end_year+1}-01-01T00:00'),
        np.timedelta64(1, 'h'),
        dtype='datetime64[h]'
    )
    timestamps = pd.DatetimeIndex(times, tz='UTC')

    doy        = timestamps.day_of_year.values # (n_time,)
    utc_tod    = timestamps.hour.values + timestamps.minute.values / 60.0 # (n_time, )
    utc_offset = lons / 15.0  # (n_stations,)

    cos_doy = np.cos(2 * np.pi * doy / 365.25)
    sin_doy = np.sin(2 * np.pi * doy / 365.25)

    local_tod = (utc_tod[None, :] + utc_offset[:, None]) % 24 # (n_stations, n_time)
    cos_tod = np.cos(2 * np.pi * local_tod / 24.0) # (n_stations, n_time)

    return {
        'times':   times,     # (n_time,)
        'cos_doy': cos_doy,   # (n_time,)
        'sin_doy': sin_doy,   # (n_time,)
        'cos_tod': cos_tod,   # (n_stations, n_time)
        'tod':     local_tod, # (n_stations, n_time)
    }


# ── per-station feature vector ──────────────────────────────────────────

def get_station_features(i, static_ds, time_indicators, variable):
    """
    Build raw feature matrix for a single station i.
    Returns X_i of shape (n_time, n_raw_features).
    """  
    static_vars = STATIC_TMP if variable in TMP_VARS else STATIC_WIND
    n_time     = len(time_indicators['times'])

    # --- static block ---
    static_block = np.stack(
        [np.full(n_time, static_ds[v].values[i]) for v in static_vars],
        axis=-1
    )  # (n_time, n_static)

    # --- time indicator block ---
    if variable in ['t2m', 'd2m']:
        time_block = np.stack([
            time_indicators['cos_doy'],
            time_indicators['sin_doy'],
            time_indicators['cos_tod'][i],
            time_indicators['tod'][i],
        ], axis=-1)
    else:  # u10 / v10
        time_block = np.stack([
            time_indicators['cos_doy'],
            time_indicators['cos_tod'][i],
        ], axis=-1)
    # (n_time, n_time_features)

    # --- concatenate and flatten ---
    X = np.concatenate([static_block, time_block], axis=-1).astype(np.float64)
    return X


# ── build target ──────────────────────────────────────────────────────

def slice_datasets(era5_ds, madis_ds, start_year=2020, end_year=2023):
    data_start = np.datetime64('2020-01-01T00:00', 'h')
    start_dt   = np.datetime64(f'{start_year}-01-01T00:00', 'h')
    end_dt     = np.datetime64(f'{end_year}-12-31T23:00',   'h')

    start_idx  = int((start_dt - data_start) / np.timedelta64(1, 'h'))
    end_idx    = int((end_dt   - data_start) / np.timedelta64(1, 'h'))

    return (
        era5_ds.isel(time=slice(start_idx, end_idx + 1)),
        madis_ds.isel(time=slice(start_idx, end_idx + 1)),
    )

def get_station_target(i, era5_ds, madis_ds, variable):
    """
    Compute forecast error e = ERA5 - MADIS = f - o
    """
    f = era5_ds[variable].values[i].astype(np.float64)
    o = madis_ds[variable].values[i].astype(np.float64)
    return f - o