main.py

"""Run one conformal prediction experiment: one dataset, one CP method.

Pipeline
--------
1. Obtain an ensemble per input.  For most datasets a conditional flow-matching
   model is trained on the 60% training split and sampled on the calibration and
   test splits; results are cached under --model_path so repeated CP runs reuse
   one trained model.  The climate datasets instead arrive as ready-made ensembles
   from the emulator of Wang et al. (2025), so no model is trained.
2. Calibrate the chosen nonconformity score on the calibration split.
3. Invert it on the test split and report coverage, volume and structural
   complexity.

Methods (--CP_type)
-------------------
CP4Gen   cluster the ensemble into K components and score by the dominant one.
         Sweeps K over a coarse grid and saves every K; pick the K minimising
         volume afterwards (see summarise_results.py).
PCP      the K = M corner of CP4Gen: one ball per ensemble member.
HD-PCP   PCP after discarding the least confident ensemble members.  Needs the
         per-sample confidence scores the generative model records, so it is not
         available for the pre-generated climate ensembles.

Example
-------
    python main.py --dataset s_curve --CP_type CP4Gen --n_epochs 20000
    python main.py --dataset precip_2 --CP_type PCP
"""

import os
import argparse
import random

import numpy as np
import torch

import selected_configs
from dataset import CLIMATE_ALIASES, CLIMATE_REDUCTIONS, get_precomputed_ensembles
from cp4gen import CPGen
from generative_models import GenerativeModel
from metrics import get_k_list

# Datasets whose ensembles are supplied directly rather than generated here.
PRECOMPUTED_DATASETS = set(CLIMATE_REDUCTIONS) | set(CLIMATE_ALIASES)

# Fractions of the ensemble HD-PCP retains, swept to pick the best per dataset.
HD_PCP_KEEP_RATES = [1, 0.95, 0.6, 0.5, 0.3]


def load_ensembles(args):
    """Get calibration and test ensembles, observations, and confidence scores.

    Returns (Y_ens_calib, Y_calib, calib_scores, Y_ens_test, Y_test, test_scores).
    The two score arrays are None for the pre-generated climate ensembles, which
    carry no per-sample confidence.
    """
    # Source 1: the climate datasets already carry ensembles from a separate
    # emulator, so there is nothing to train and no confidence scores to return.
    if args.dataset in PRECOMPUTED_DATASETS:
        Y_ens_calib, Y_calib, Y_ens_test, Y_test = get_precomputed_ensembles(
            args.dataset, data_path=args.data_path)
        return Y_ens_calib, Y_calib, None, Y_ens_test, Y_test, None

    # Source 2: a previous run already trained a generator and saved its ensembles.
    # Reusing them matters because training dominates the runtime, while a CP sweep
    # over K is comparatively cheap -- so every K reuses one trained model.
    cached = os.path.join(args.model_path, 'Y_ens_calib.npy')
    if os.path.exists(cached):
        print(f'Reusing cached ensembles from {args.model_path}', flush=True)
        load = lambda name: np.load(os.path.join(args.model_path, f'{name}.npy'))
        return (load('Y_ens_calib'), load('Y_calib'), load('calib_scores'),
                load('Y_ens_test'), load('Y_test'), load('test_scores'))

    # Source 3: nothing cached, so train the generator and sample it now.
    generative_model = GenerativeModel(args)
    generative_model.prep_data()
    generative_model.train()
    (Y_ens_calib, calib_scores, calib_conditions,
     Y_ens_test, test_scores, test_conditions) = generative_model.sample()
    Y_calib, Y_test = generative_model.get_ground_truth()
    generative_model.save()

    # Cache everything the CP methods need, so later runs take the branch above.
    for name, array in [('Y_ens_calib', Y_ens_calib), ('calib_scores', calib_scores),
                        ('Y_calib', Y_calib), ('Y_ens_test', Y_ens_test),
                        ('test_scores', test_scores), ('Y_test', Y_test),
                        ('calib_conditions', calib_conditions),
                        ('test_conditions', test_conditions)]:
        np.save(os.path.join(args.model_path, f'{name}.npy'), array)

    return Y_ens_calib, Y_calib, calib_scores, Y_ens_test, Y_test, test_scores


def keep_most_confident(Y_ens, scores, n_keep):
    """Retain the `n_keep` highest-confidence members of each ensemble (HD-PCP).

    Confidence is the log-density of the Gaussian noise seed each sample was
    integrated from, recorded at sampling time.  Because the flow is a
    deterministic map out of that noise, samples starting nearer the noise
    distribution's mode are the ones the model places in its high-density region.

    The same filter is applied to calibration and test ensembles, or the
    exchangeability the coverage guarantee rests on would break.
    """
    top_idx = np.argsort(scores.squeeze(-1), axis=1)[:, -n_keep:]
    return Y_ens[np.arange(len(Y_ens))[:, None], top_idx, :]


def evaluate(args, k, Y_ens_calib, Y_calib, Y_ens_test, Y_test, label, tag):
    """Calibrate at K = k, then evaluate on both the calibration and test splits.

    Metrics are reported on both because K has to be chosen somehow.  Choosing it by
    minimising volume on the *test* sweep reports the winner on the same data that
    picked it, which is optimistic.  The calibration split gives a selection signal
    that never touches test, so the honest recipe is: pick K by the calibration
    column, then read that row's test column as the result.

    Coverage is guaranteed either way -- Theorem F.1 only requires K to be fixed
    before calibration -- so this affects the reported volume and complexity, not
    validity.  Calibration metrics are still mildly optimistic in absolute terms,
    since the same points set the quantile, but they rank K without seeing test.
    """
    cp_method = CPGen(args, k=k, fit_mixture=args.fit_mixture)
    cp_method.fit(Y_ens_calib, Y_calib)

    print(f'{label}', flush=True)
    print(f'  {"":<22} {"calibration":>13}  {"test":>13}', flush=True)

    split_metrics = {}
    for split, Y_ens, Y in [('calib', Y_ens_calib, Y_calib), ('test', Y_ens_test, Y_test)]:
        scores, volumes, ks, n_intervals = cp_method.predict(Y_ens, Y)
        split_metrics[split] = {
            'scores': scores, 'volumes': volumes, 'ks': ks, 'n_intervals': n_intervals,
            'coverage': np.mean(scores < cp_method.quant_score),
        }

    for row, key in [('Coverage Rate:', 'coverage'), ('Average Volume:', 'volumes'),
                     ('Structural Complexity:', 'ks'), ('Disjoint Intervals:', 'n_intervals')]:
        if key == 'n_intervals' and Y_test.shape[1] != 1:
            continue
        values = [np.mean(split_metrics[s][key]) for s in ('calib', 'test')]
        print(f'  {row:<22} {values[0]:>13.6f}  {values[1]:>13.6f}', flush=True)
    print(' ', flush=True)

    # Test arrays keep their original filenames; calibration arrays are prefixed, so
    # existing downstream scripts continue to read the test numbers unchanged.
    for split, prefix in [('test', ''), ('calib', 'calib_')]:
        for name in ('scores', 'volumes', 'ks', 'n_intervals'):
            np.save(os.path.join(args.output_saving_path, f'{prefix}{tag}_{name}.npy'),
                    split_metrics[split][name])
    np.save(os.path.join(args.output_saving_path, f'{tag}_quant_score.npy'),
            np.array([cp_method.quant_score]))


def run(args):
    """Run the configured CP method on the configured dataset."""
    # Seed every source of randomness: python and numpy drive the data shuffle and
    # the Monte Carlo volume estimates, torch drives training and sampling.
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

    (Y_ens_calib, Y_calib, calib_scores,
     Y_ens_test, Y_test, test_scores) = load_ensembles(args)

    # Trim every ensemble to the requested size; this is what the ensemble-size
    # ablation of Appendix E.3 varies.
    ens_size = min(args.n_ens, Y_ens_calib.shape[1])
    Y_ens_calib = Y_ens_calib[:, :ens_size]
    Y_ens_test = Y_ens_test[:, :ens_size]

    # Confidence scores are per ensemble member, so they must be trimmed in step
    # with the ensembles or HD-PCP would filter on scores for discarded samples.
    if calib_scores is not None:
        calib_scores = calib_scores[:, :ens_size]
        test_scores = test_scores[:, :ens_size]

    print(f'Y_ens_calib: {Y_ens_calib.shape}   Y_calib: {Y_calib.shape}', flush=True)
    print(f'Y_ens_test:  {Y_ens_test.shape}   Y_test:  {Y_test.shape}', flush=True)
    print('-' * 40, flush=True)

    d = Y_ens_calib.shape[-1]

    if args.CP_type == 'PCP':
        # PCP is CP4Gen with one component per ensemble member, so a single
        # evaluation at K = M is the whole method -- there is nothing to tune.
        evaluate(args, ens_size, Y_ens_calib, Y_calib, Y_ens_test, Y_test,
                 label=f'PCP (K = M = {ens_size})', tag='PCP')

    elif args.CP_type == 'CP4Gen':
        # Sweep K over the coarse grid and save each; the choice of a single K is
        # deferred to summarise_results.py so it can be made on the calibration split.
        for k in get_k_list(ens_size, d):
            evaluate(args, k, Y_ens_calib, Y_calib, Y_ens_test, Y_test,
                     label=f'CP4Gen  K = {k}', tag=f'CP4Gen_{k}')

    elif args.CP_type == 'HD-PCP':
        if calib_scores is None:
            raise ValueError(
                'HD-PCP needs per-sample confidence scores, which the pre-generated '
                f'{args.dataset} ensembles do not carry. Use PCP or CP4Gen instead.')

        # Sweep the retained fraction.  selected_configs records which size was
        # chosen per dataset for the comparison in Table 4; it is flagged below so
        # the reported configuration is obvious in the output.
        selected = selected_configs.HD_PCP_KEEP_SIZE.get(args.dataset)

        for keep_rate in HD_PCP_KEEP_RATES:
            n_keep = int(ens_size * keep_rate)
            marker = '  <- selected for this dataset' if n_keep == selected else ''
            evaluate(args, n_keep,
                     keep_most_confident(Y_ens_calib, calib_scores, n_keep), Y_calib,
                     keep_most_confident(Y_ens_test, test_scores, n_keep), Y_test,
                     label=f'HD-PCP  keep {n_keep} of {ens_size}{marker}',
                     tag=f'HD-PCP_{keep_rate}')

    else:
        raise ValueError(f'Unknown CP_type {args.CP_type!r}. '
                         'Choose from CP4Gen, PCP, HD-PCP.')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('--seed', default=100, type=int)
    parser.add_argument('--dataset', default='s_curve', type=str)
    parser.add_argument('--data_path', default='./data/', type=str)
    parser.add_argument('--output_saving_path', default='./output/', type=str)
    parser.add_argument('--model_path', default='./generative_models/', type=str,
                        help='where trained generators and their ensembles are cached')

    # Generative model
    parser.add_argument('--model_type', type=str, default='flow-matching')
    parser.add_argument('--n_epochs', type=int, default=10000)
    parser.add_argument('--batch_size', type=int, default=1000)
    parser.add_argument('--hidden_dim', type=int, default=128)
    parser.add_argument('--timesteps', type=int, default=100)
    parser.add_argument('--lr', type=float, default=1e-3)
    parser.add_argument('--n_samples', type=int, default=1000,
                        help='ensemble members generated per input and cached')

    # Conformal prediction
    parser.add_argument('--CP_type', type=str, default='CP4Gen',
                        choices=['CP4Gen', 'PCP', 'HD-PCP'])
    parser.add_argument('--n_ens', type=int, default=30,
                        help='ensemble size M actually used; the paper uses 30')
    parser.add_argument('--coverage', type=float, default=0.9, help='1 - alpha')
    parser.add_argument('--fit_mixture', type=str, default='kmeans',
                        choices=['kmeans', 'em'],
                        help='mixture fitter; "em" reproduces the Appendix E.2 ablation')
    parser.add_argument('--mc_points', type=int, default=None,
                        help='Monte Carlo points per test point for 3-d and 4-d volumes; '
                             'defaults to the published budget (1e6 in 3-d, 5e5 in 4-d). '
                             '1e5 matches the published mean to ~0.001%% and is ~10x faster')
    parser.add_argument('--volume_method', type=str, default='paper',
                        choices=['paper', 'exact'],
                        help='"paper" reproduces the published tables; "exact" is a '
                             'lower-variance estimator that shifts numbers slightly')

    args = parser.parse_args()
    args.device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

    # Cache trained generators per model configuration and dataset, so that runs
    # differing only in CP settings reuse one trained model.
    model_params = args.model_type + ''.join(
        f'--{key}={getattr(args, key)}'
        for key in ['n_epochs', 'batch_size', 'hidden_dim', 'timesteps', 'lr'])
    args.model_path = os.path.join(args.model_path, model_params, args.dataset)
    os.makedirs(args.model_path, exist_ok=True)
    os.makedirs(args.output_saving_path, exist_ok=True)

    print('Experiment Configuration:', flush=True)
    for key, value in sorted(vars(args).items()):
        print(f'  {key}: {value}', flush=True)
    print('-' * 40, flush=True)

    run(args)