dataset.py

"""Datasets used in the CP4Gen experiments (Section 4.3, Appendix B.1).

Three groups:

  synthetic          25-Gaussians, 8-Gaussians, moon, circle, spiral, s_curve.
                     2-d point clouds in which one coordinate is the covariate and
                     the other the response, so the conditional P(Y | X) is known to
                     be multi-modal.  Generated on the fly, 5000 points each.

  real world         nine public datasets with a 1-d response (meps_19/20/21,
                     facebook_1/2, bio, blog_data, temperature, bike) -- the same
                     ones Wang et al. (2022) used for PCP, so the comparison is like
                     for like -- plus two with a 2-d response, taxi and energy.

  climate emulation  precip_2, precip_2n, precip_3, precip_4.  These arrive as
                     ready-made ensembles rather than (X, Y) pairs, so they are
                     loaded by `get_precomputed_ensembles`, not `get_dataset`.

`get_dataset(name)` returns X of shape (N, dim_x) and Y of shape (N, dim_y);
callers split them 60/20/20 into train, calibration and test.
"""

import torch
import numpy as np
import pandas as pd

from torch.utils.data import Dataset
from sklearn.datasets import make_s_curve
from sklearn.datasets import make_swiss_roll
from sklearn.datasets import make_circles
from sklearn.datasets import make_moons
from sklearn.datasets import make_blobs


class DatasetTensor(Dataset):
    def __init__(self, X, Y):
        """
        Args:
            X: Input data (features), shape (N, dim_x).
            Y: Target data (labels), shape (N, dim_y).
        """
        self.X = torch.tensor(X, dtype=torch.float32)
        self.Y = torch.tensor(Y, dtype=torch.float32)

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

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


# The climate emulation responses are 96 x 192 global precipitation fields, far too
# large to measure a prediction set volume in.  Section 4.3 therefore reduces each
# field to a handful of numbers, four ways, and runs conformal prediction on those.
# The reductions themselves are applied in notebooks/prepare_climate_data.ipynb.
#
#   precip_2    NSM   mean over the northern and southern hemispheres      -> 2-d
#   precip_2n   TP    values at two nearby cities, Key West and Miami      -> 2-d
#   precip_3    NMSM  mean over northern, middle and southern bands        -> 3-d
#   precip_4    QM    mean over the four quadrants of the globe            -> 4-d
#
# precip_2 and precip_2n are both 2-d on purpose: the hemispheres are only weakly
# correlated while the two cities are strongly correlated, which is how the paper
# isolates the effect of response correlation on prediction set volume.
CLIMATE_REDUCTIONS = {
    'precip_2': 'NSM',
    'precip_2n': 'TP',
    'precip_3': 'NMSM',
    'precip_4': 'QM',
}

# Names used while the experiments were being run, kept so older job scripts and
# saved output directories still resolve.
CLIMATE_ALIASES = {
    'Mengze_2': 'precip_2',
    'Mengze_nearest': 'precip_2n',
    'Mengze_3': 'precip_3',
    'Mengze_4': 'precip_4',
}


def get_precomputed_ensembles(name, data_path=None):
    """Load a climate emulation dataset as ready-made ensembles.

    These responses come from the separate flow-matching emulator of Wang et al.
    (2025), so there is no generative model to train here: the ensembles are loaded
    directly and conformal prediction is applied on top.

    Calibration and test come from two different simulation contexts, held in the
    files suffixed 2 and 1 respectively.  The ground truth is taken as member 0 of
    the corresponding true-field array, the remaining members being other draws of
    the same simulated year.

    Returns
    -------
    Y_ens_calib : (n_calib, n_ens, dim_y)
    Y_calib : (n_calib, dim_y)
    Y_ens_test : (n_test, n_ens, dim_y)
    Y_test : (n_test, dim_y)
    """
    name = CLIMATE_ALIASES.get(name, name)
    if name not in CLIMATE_REDUCTIONS:
        raise ValueError(
            f'Unknown climate dataset {name!r}. '
            f'Available: {", ".join(CLIMATE_REDUCTIONS)}'
        )

    reduction = CLIMATE_REDUCTIONS[name]

    Y_ens_calib = np.load(f'{data_path}climate/y_hat2_reduce_{reduction}.npy')
    Y_calib = np.load(f'{data_path}climate/y2_reduce_{reduction}.npy')[:, 0, :]

    Y_ens_test = np.load(f'{data_path}climate/y_hat1_reduce_{reduction}.npy')
    Y_test = np.load(f'{data_path}climate/y1_reduce_{reduction}.npy')[:, 0, :]

    return Y_ens_calib, Y_calib, Y_ens_test, Y_test

# The three MEPS panels share one feature list, differing only in which survey
# weight they report.  Kept as one constant rather than three near-identical copies.
MEPS_FEATURE_COLUMNS = [
    'AGE', 'PCS42', 'MCS42', 'K6SUM42', 'PERWT15F', 'REGION=1',
    'REGION=2', 'REGION=3', 'REGION=4', 'SEX=1', 'SEX=2', 'MARRY=1',
    'MARRY=2', 'MARRY=3', 'MARRY=4', 'MARRY=5', 'MARRY=6', 'MARRY=7',
    'MARRY=8', 'MARRY=9', 'MARRY=10', 'FTSTU=-1', 'FTSTU=1', 'FTSTU=2',
    'FTSTU=3', 'ACTDTY=1', 'ACTDTY=2', 'ACTDTY=3', 'ACTDTY=4',
    'HONRDC=1', 'HONRDC=2', 'HONRDC=3', 'HONRDC=4', 'RTHLTH=-1',
    'RTHLTH=1', 'RTHLTH=2', 'RTHLTH=3', 'RTHLTH=4', 'RTHLTH=5',
    'MNHLTH=-1', 'MNHLTH=1', 'MNHLTH=2', 'MNHLTH=3', 'MNHLTH=4',
    'MNHLTH=5', 'HIBPDX=-1', 'HIBPDX=1', 'HIBPDX=2', 'CHDDX=-1',
    'CHDDX=1', 'CHDDX=2', 'ANGIDX=-1', 'ANGIDX=1', 'ANGIDX=2',
    'MIDX=-1', 'MIDX=1', 'MIDX=2', 'OHRTDX=-1', 'OHRTDX=1', 'OHRTDX=2',
    'STRKDX=-1', 'STRKDX=1', 'STRKDX=2', 'EMPHDX=-1', 'EMPHDX=1',
    'EMPHDX=2', 'CHBRON=-1', 'CHBRON=1', 'CHBRON=2', 'CHOLDX=-1',
    'CHOLDX=1', 'CHOLDX=2', 'CANCERDX=-1', 'CANCERDX=1', 'CANCERDX=2',
    'DIABDX=-1', 'DIABDX=1', 'DIABDX=2', 'JTPAIN=-1', 'JTPAIN=1',
    'JTPAIN=2', 'ARTHDX=-1', 'ARTHDX=1', 'ARTHDX=2', 'ARTHTYPE=-1',
    'ARTHTYPE=1', 'ARTHTYPE=2', 'ARTHTYPE=3', 'ASTHDX=1', 'ASTHDX=2',
    'ADHDADDX=-1', 'ADHDADDX=1', 'ADHDADDX=2', 'PREGNT=-1', 'PREGNT=1',
    'PREGNT=2', 'WLKLIM=-1', 'WLKLIM=1', 'WLKLIM=2', 'ACTLIM=-1',
    'ACTLIM=1', 'ACTLIM=2', 'SOCLIM=-1', 'SOCLIM=1', 'SOCLIM=2',
    'COGLIM=-1', 'COGLIM=1', 'COGLIM=2', 'DFHEAR42=-1', 'DFHEAR42=1',
    'DFHEAR42=2', 'DFSEE42=-1', 'DFSEE42=1', 'DFSEE42=2',
    'ADSMOK42=-1', 'ADSMOK42=1', 'ADSMOK42=2', 'PHQ242=-1', 'PHQ242=0',
    'PHQ242=1', 'PHQ242=2', 'PHQ242=3', 'PHQ242=4', 'PHQ242=5',
    'PHQ242=6', 'EMPST=-1', 'EMPST=1', 'EMPST=2', 'EMPST=3', 'EMPST=4',
    'POVCAT=1', 'POVCAT=2', 'POVCAT=3', 'POVCAT=4', 'POVCAT=5',
    'INSCOV=1', 'INSCOV=2', 'INSCOV=3', 'RACE',
]

MEPS_RESPONSE_COLUMN = 'UTILIZATION_reg'

SYNTHETIC_DATASETS = ['25-Gaussians', '8-Gaussians', 'moon', 'circle', 'spiral', 's_curve']
REAL_WORLD_DATASETS = ['meps_19', 'meps_20', 'meps_21', 'facebook_1', 'facebook_2',
                       'bio', 'blog_data', 'temperature', 'bike', 'taxi', 'energy']


def meps_columns(panel):
    """MEPS feature list for one panel, with that panel's survey weight substituted.

    Panels 19 and 20 report PERWT15F; panel 21 reports PERWT16F.
    """
    survey_weight = 'PERWT16F' if panel == '21' else 'PERWT15F'
    return [survey_weight if c == 'PERWT15F' else c for c in MEPS_FEATURE_COLUMNS]


def _shuffled(df, feature_columns, response_columns):
    """Shuffle a dataframe's rows, then split it into X and Y arrays.

    Draws from the global RNG that `get_dataset` seeds, so the order -- and hence the
    train/calibration/test split -- is reproducible for a given seed.
    """
    n = len(df)
    idx = np.arange(n)
    np.random.shuffle(idx)

    X = df[feature_columns].values.reshape(n, -1)[idx]
    Y = df[response_columns].values.reshape(n, -1)[idx]
    return X, Y


def get_dataset(name, data_path=None, seed=0):
    """Load a dataset by name.

    Parameters
    ----------
    name : str
        One of `SYNTHETIC_DATASETS` or `REAL_WORLD_DATASETS`.  Climate datasets are
        loaded by `get_precomputed_ensembles` instead, since they arrive as
        ensembles rather than as (X, Y) pairs.
    data_path : str
        Directory holding the data files; see data/README.md.
    seed : int
        Seeds both the synthetic generators and the row shuffle.

    Returns
    -------
    X : (N, dim_x) covariates
    Y : (N, dim_y) responses
    """
    # Seeds the global RNG that `_shuffled` draws from, so a given seed always
    # yields the same row order and therefore the same 60/20/20 split.
    np.random.seed(seed)

    # ---- Synthetic: 2-d point clouds, one coordinate in, the other out ----
    if name == 's_curve':
        # An S-shaped curve, so the response is multi-valued over much of the range.
        points, _ = make_s_curve(n_samples=5000, noise=0, random_state=seed)
        return points[:, 0].reshape(-1, 1), points[:, 2].reshape(-1, 1)

    if name == 'spiral':
        # A 2-d slice of the 3-d swiss roll; the conditional is multi-modal.
        points, _ = make_swiss_roll(n_samples=5000, noise=0, random_state=seed)
        return points[:, 0].reshape(-1, 1), points[:, -1].reshape(-1, 1)

    if name == 'circle':
        # Two concentric circles with radius ratio 0.7 : 1.
        points, _ = make_circles(n_samples=5000, noise=0, factor=0.7, random_state=seed)
        return points[:, 0].reshape(-1, 1), points[:, -1].reshape(-1, 1)

    if name == 'moon':
        # Two interleaved half-moons with light Gaussian noise.
        points, _ = make_moons(n_samples=5000, noise=0.01, random_state=seed)
        return points[:, 0].reshape(-1, 1), points[:, -1].reshape(-1, 1)

    if name == '25-Gaussians':
        # 25 tight blobs on a 5 x 5 grid: conditioning on x leaves up to five
        # well-separated modes in y, the hardest multi-modal case in the paper.
        grid_x, grid_y = np.meshgrid(np.linspace(-1.5, 1.5, 5), np.linspace(-1.5, 1.5, 5))
        centers = np.concatenate([grid_x.reshape(-1, 1), grid_y.reshape(-1, 1)], axis=1)
        points, _ = make_blobs(n_samples=5000, centers=centers, cluster_std=0.01,
                               random_state=seed)
        return points[:, 0].reshape(-1, 1), points[:, -1].reshape(-1, 1)

    if name == '8-Gaussians':
        # 8 tight blobs evenly spaced on a circle of radius 1.5.
        angles = np.linspace(-1, 1, 9)[:-1] * np.pi
        centers = np.concatenate([(np.cos(angles) * 1.5).reshape(-1, 1),
                                  (np.sin(angles) * 1.5).reshape(-1, 1)], axis=1)
        points, _ = make_blobs(n_samples=5000, centers=centers, cluster_std=0.01,
                               random_state=seed)
        return points[:, 0].reshape(-1, 1), points[:, -1].reshape(-1, 1)

    # ---- Real world, 1-d response ----
    if name in ('meps_19', 'meps_20', 'meps_21'):
        # Predict medical expenditure from patient features (Romano et al., 2019).
        panel = name.split('_')[1]
        df = pd.read_csv(data_path + f'meps_{panel}_reg.csv')
        return _shuffled(df, meps_columns(panel), [MEPS_RESPONSE_COLUMN])

    if name in ('facebook_1', 'facebook_2'):
        # Predict a post's comment volume from 53 post features.
        df = pd.read_csv(data_path + f'{name}.csv')
        return _shuffled(df, df.columns[0:53], [df.columns[53]])

    if name == 'bio':
        # Predict protein tertiary-structure RMSD from physicochemical properties.
        df = pd.read_csv(data_path + 'CASP.csv')
        return _shuffled(df, df.columns[1:], [df.columns[0]])

    if name == 'blog_data':
        # Predict a blog post's comment count from 280 features.
        df = pd.read_csv(data_path + 'blogData_train.csv', header=None)
        return _shuffled(df, df.columns[0:280], [df.columns[-1]])

    if name == 'temperature':
        # Predict next-day air temperature from forecasts and in-situ observations.
        df = pd.read_csv(data_path + 'Bias_correction_ucl.csv')
        df = df.drop(columns=['station', 'Date', 'Next_Tmax']).dropna()
        return _shuffled(df, df.columns[:-1], [df.columns[-1]])

    if name == 'bike':
        # Predict hourly bike rentals from weather and calendar features.  Season and
        # weather are categorical so they are one-hot encoded, and the timestamp is
        # expanded into hour / weekday / month / year before being dropped.
        df = pd.read_csv(data_path + 'bike_train.csv')
        df = pd.concat([df, pd.get_dummies(df['season'], prefix='season')], axis=1)
        df = pd.concat([df, pd.get_dummies(df['weather'], prefix='weather')], axis=1)
        df.drop(['season', 'weather'], inplace=True, axis=1)

        df['hour'] = [t.hour for t in pd.DatetimeIndex(df.datetime)]
        df['day'] = [t.dayofweek for t in pd.DatetimeIndex(df.datetime)]
        df['month'] = [t.month for t in pd.DatetimeIndex(df.datetime)]
        df['year'] = [t.year for t in pd.DatetimeIndex(df.datetime)]
        df['year'] = df['year'].map({2011: 0, 2012: 1})

        # 'casual' and 'registered' sum to 'count', so keeping them would leak it.
        df.drop('datetime', axis=1, inplace=True)
        df.drop(['casual', 'registered'], axis=1, inplace=True)
        return _shuffled(df, df.columns.drop('count'), ['count'])

    # ---- Real world, 2-d response ----
    if name == 'taxi':
        # Predict a trip's drop-off latitude and longitude from its pickup location
        # and cyclically encoded pickup time.  The two response dimensions are
        # strongly correlated, which is where CP4Gen's covariance fit pays off most.
        df = pd.read_csv(data_path + 'taxi_data.csv')
        X_cols = ['pickup_time_day_of_week_sin', 'pickup_time_day_of_week_cos',
                  'pickup_time_of_day_sin', 'pickup_time_of_day_cos',
                  'pickup_loc_lat', 'pickup_loc_lon']
        return _shuffled(df, X_cols, ['dropoff_loc_lat', 'dropoff_loc_lon'])

    if name == 'energy':
        # Jointly predict a building's heating load (Y1) and cooling load (Y2) from
        # eight design variables (Tsanas & Xifara, 2012).
        df = pd.read_csv(data_path + 'energy_data.csv')
        return _shuffled(df, [f'X{i}' for i in range(1, 9)], ['Y1', 'Y2'])

    if name in CLIMATE_REDUCTIONS or name in CLIMATE_ALIASES:
        raise ValueError(
            f'{name!r} is a climate dataset supplied as ready-made ensembles; '
            'load it with get_precomputed_ensembles() instead.')

    raise ValueError(f'Unknown dataset {name!r}. Available: '
                     f'{", ".join(SYNTHETIC_DATASETS + REAL_WORLD_DATASETS)}')