ablations.py

"""The ablation studies of Appendix E.

Three studies, each reusing one set of cached ensembles so that only the quantity
under study varies:

  --study k          Sensitivity to the number of mixture components K (E.1,
                     Figure 5).  Sweeps K from 1 to the ensemble size.  Coverage
                     stays flat near 1 - alpha while volume traces a V: small K
                     merges distinct modes and swallows the low-density gaps
                     between them, large K approaches PCP's fragmented sets.

  --study ensemble   Effect of the ensemble size M (E.3, Figure 6).  Sweeps M and
                     compares CP4Gen against PCP.  PCP's structural complexity
                     grows linearly in M by construction, whereas CP4Gen's settles
                     once M is large enough to resolve the cluster structure.

  --study em         K-means versus EM as the mixture fitter (E.2, Table 5).  Both
                     use the identical dominant-component score, so the comparison
                     isolates the fitting step.

Ensembles must already be cached by `main.py` (which writes them under
--model_path), or supplied directly as with the climate datasets.

Example
-------
    python main.py --dataset 25-Gaussians --CP_type PCP --n_epochs 50000
    python ablations.py --study k --dataset 25-Gaussians
"""

import argparse
import csv
import os

import numpy as np
import torch

from cp4gen import CPGen
from dataset import CLIMATE_ALIASES, CLIMATE_REDUCTIONS, get_precomputed_ensembles

PRECOMPUTED_DATASETS = set(CLIMATE_REDUCTIONS) | set(CLIMATE_ALIASES)


def load_cached_ensembles(args):
    """Load ensembles produced by a previous `main.py` run, or the supplied ones.

    Returns (Y_ens_calib, Y_calib, Y_ens_test, Y_test).
    """
    if args.dataset in PRECOMPUTED_DATASETS:
        return get_precomputed_ensembles(args.dataset, data_path=args.data_path)

    required = os.path.join(args.model_path, 'Y_ens_calib.npy')
    if not os.path.exists(required):
        raise SystemExit(
            f'No cached ensembles at {args.model_path}.\n'
            f'Run main.py for {args.dataset} first, with the same --model_path '
            'and generative model settings.')

    load = lambda name: np.load(os.path.join(args.model_path, f'{name}.npy'))
    return load('Y_ens_calib'), load('Y_calib'), load('Y_ens_test'), load('Y_test')


def measure(args, k, Y_ens_calib, Y_calib, Y_ens_test, Y_test, fit_mixture='kmeans'):
    """Calibrate at K = k and return the three metrics on the test split."""
    cp_method = CPGen(args, k=k, fit_mixture=fit_mixture)
    cp_method.fit(Y_ens_calib, Y_calib)
    scores, volumes, ks, _ = cp_method.predict(Y_ens_test, Y_test)

    return {
        'coverage': float(np.mean(scores < cp_method.quant_score)),
        'volume': float(np.mean(volumes)),
        'complexity': float(np.mean(ks)),
    }


def study_k(args, ensembles):
    """E.1: sweep K from 1 to the ensemble size at fixed M."""
    Y_ens_calib, Y_calib, Y_ens_test, Y_test = ensembles
    ens_size = Y_ens_calib.shape[1]

    # Dense at small K, where volume changes fastest, thinning out toward M.
    k_values = sorted({1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
                       *range(15, ens_size + 1, 5), ens_size})

    print(f'{"K":>6}  {"coverage":>9}  {"volume":>13}  {"complexity":>11}')
    print('-' * 46)
    results = {}
    for k in k_values:
        metrics = measure(args, k, Y_ens_calib, Y_calib, Y_ens_test, Y_test)
        results[k] = metrics
        print(f'{k:>6}  {metrics["coverage"]:>9.3f}  {metrics["volume"]:>13.6f}  '
              f'{metrics["complexity"]:>11.2f}', flush=True)

    best = min(results, key=lambda k: results[k]['volume'])
    print(f'\nVolume minimised at K = {best} ({results[best]["volume"]:.6f})')
    return results


def study_ensemble_size(args, ensembles):
    """E.3: sweep the ensemble size M, comparing CP4Gen against PCP."""
    Y_ens_calib, Y_calib, Y_ens_test, Y_test = ensembles
    available = Y_ens_calib.shape[1]

    # Trimming a cached ensemble is the only way to shrink M after the fact, so the
    # sweep is capped by however many members were generated.
    sizes = [m for m in args.ensemble_sizes if m <= available]
    if not sizes:
        raise SystemExit(f'Cached ensembles hold only {available} members; '
                         f'none of {args.ensemble_sizes} fit.')

    print(f'{"M":>5}  {"method":>8}  {"coverage":>9}  {"volume":>13}  {"complexity":>11}')
    print('-' * 55)
    results = {}
    for m in sizes:
        # Take the first m members, so every M in the sweep sees the same samples
        # and the comparison isolates ensemble size from sampling noise.
        calib, test = Y_ens_calib[:, :m], Y_ens_test[:, :m]

        # PCP is the K = M corner; CP4Gen uses the volume-minimising K at this M,
        # since the best K generally shifts as the ensemble grows.
        pcp = measure(args, m, calib, Y_calib, test, Y_test)
        candidates = {k: measure(args, k, calib, Y_calib, test, Y_test)
                      for k in sorted({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, m}) if k <= m}
        best_k = min(candidates, key=lambda k: candidates[k]['volume'])
        cp4gen = candidates[best_k]

        results[m] = {'PCP': pcp, 'CP4Gen': cp4gen, 'best_k': best_k}
        for label, metrics in [('PCP', pcp), ('CP4Gen', cp4gen)]:
            print(f'{m:>5}  {label:>8}  {metrics["coverage"]:>9.3f}  '
                  f'{metrics["volume"]:>13.6f}  {metrics["complexity"]:>11.2f}', flush=True)

    return results


def study_em(args, ensembles):
    """E.2: K-means versus EM as the mixture fitter, at matched K."""
    Y_ens_calib, Y_calib, Y_ens_test, Y_test = ensembles

    print(f'{"K":>4}  {"fitter":>8}  {"coverage":>9}  {"volume":>13}  {"complexity":>11}')
    print('-' * 54)
    results = {}
    # Same K and same ensembles for both fitters, so any difference is the fitting
    # step alone -- the score functional is identical either way.
    for k in args.k_values:
        for fitter in ('kmeans', 'em'):
            metrics = measure(args, k, Y_ens_calib, Y_calib, Y_ens_test, Y_test,
                              fit_mixture=fitter)
            results[(k, fitter)] = metrics
            print(f'{k:>4}  {fitter:>8}  {metrics["coverage"]:>9.3f}  '
                  f'{metrics["volume"]:>13.6f}  {metrics["complexity"]:>11.2f}', flush=True)

    return results



def save_results(results, study, dataset, output_path):
    """Write a study's results to CSV so the figure scripts can plot them.

    One row per configuration, with the varying quantity in the first column.
    """
    os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True)

    with open(output_path, 'w', newline='') as handle:
        writer = csv.writer(handle)

        if study == 'k':
            writer.writerow(['K', 'coverage', 'volume', 'complexity'])
            for k, m in results.items():
                writer.writerow([k, m['coverage'], m['volume'], m['complexity']])

        elif study == 'ensemble':
            writer.writerow(['M', 'method', 'coverage', 'volume', 'complexity', 'best_k'])
            for m_size, entry in results.items():
                for method in ('PCP', 'CP4Gen'):
                    metrics = entry[method]
                    writer.writerow([m_size, method, metrics['coverage'], metrics['volume'],
                                     metrics['complexity'],
                                     entry['best_k'] if method == 'CP4Gen' else m_size])

        elif study == 'em':
            writer.writerow(['K', 'fitter', 'coverage', 'volume', 'complexity'])
            for (k, fitter), m in results.items():
                writer.writerow([k, fitter, m['coverage'], m['volume'], m['complexity']])

    print(f'\nwrote {output_path}')


STUDIES = {'k': study_k, 'ensemble': study_ensemble_size, 'em': study_em}


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('--study', required=True, choices=sorted(STUDIES))
    parser.add_argument('--dataset', default='25-Gaussians', type=str)
    parser.add_argument('--data_path', default='./data/', type=str)
    parser.add_argument('--model_path', default='./generative_models/', type=str,
                        help='directory holding ensembles cached by main.py')
    parser.add_argument('--coverage', type=float, default=0.9, help='1 - alpha')
    parser.add_argument('--volume_method', default='paper', choices=['paper', 'exact'])
    parser.add_argument('--ensemble_sizes', type=int, nargs='+',
                        default=[10, 20, 30, 40, 50, 60, 70, 80, 90, 100],
                        help='M values for --study ensemble')
    parser.add_argument('--k_values', type=int, nargs='+', default=[1, 2, 3, 4, 5],
                        help='K values for --study em')
    parser.add_argument('--save_path', type=str, default=None,
                        help='write results to this CSV, for the figure scripts')
    args = parser.parse_args()
    args.device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

    # Seeded because the volume estimates in 3-d and above are Monte Carlo.
    np.random.seed(100)
    torch.manual_seed(100)

    results = STUDIES[args.study](args, load_cached_ensembles(args))

    save_path = args.save_path or f'output/ablation_{args.study}_{args.dataset}.csv'
    save_results(results, args.study, args.dataset, save_path)