CKME-CP: Conditional Kernel Mean Embedding for Conformal Prediction

A comprehensive framework for Distributional Conformal Prediction (DCP) using Conditional Kernel Mean Embedding (CKME) methods.

Table of Contents

• Overview
• Project Structure
• Installation
• Quick Start
• Module Documentation
• Usage Examples
• Experiments
• Dependencies

Overview

This project implements CKME-based conformal prediction methods and compares them with benchmark DCP methods. The framework supports:

• DCP-CKME: Distributional Conformal Prediction using CKME
• CKME-CKME: CKME-native distributional score method
• Benchmark Methods: DCP-QR, DCP-QR*, DCP-DR, HetGP (implemented in R)

Key Features

• Unified data generation for multiple examples (MG1, Exp1)
• Parameter optimization using K-Fold Cross-Validation with MMD loss
• Support for multiple macro-replications
• Automatic result evaluation (coverage, width, interval score)
• Comparison plotting and summary tables
• Consistent data sharing between Python and R code

Project Structure

CKME_CP_111625/
├── core/                    # Core modules
│   ├── models/              # CKME model and weight computation
│   │   └── ckme_weights.py
│   ├── loss/                # Loss functions
│   │   └── mmd_loss.py      # Maximum Mean Discrepancy loss
│   ├── optimizers/          # Parameter optimization
│   │   └── optimizer.py     # K-Fold CV optimizer
│   └── predictors/          # Prediction methods
│       └── dcp_predictor.py # DCP predictor
│
├── data_generators/         # Data generation module
│   ├── generator.py         # Unified data generator
│   ├── example_usage.py    # Usage examples
│   └── README.md           # Data generator documentation
│
├── examples/                # Example definitions
│   ├── mg1.py              # MG1 queue system example
│   └── exp1.py             # Exp1 sinusoidal example
│
├── evaluation/              # Performance evaluation
│   └── evaluator.py        # Coverage, width, interval score
│
├── utils/                   # Utility functions
│   └── kernels.py          # Kernel functions (RBF, Silverman bandwidth)
│
├── experiments/             # Experiment scripts
│   └── exp1/               # Exp1 experiment
│       ├── generate_data.py      # Data generation script
│       ├── runner.py             # Experiment runner
│       ├── run_dcp_ckme.py       # DCP-CKME experiment
│       ├── run_dcp_methods.R     # Benchmark DCP methods (R)
│       └── plot_comparison.py    # Result comparison plots
│
├── benchmark/               # Benchmark methods
│   └── dcp_r/              # R implementation of DCP methods
│       └── functions_final.R
│
├── data/                    # Generated data (timestamped folders)
├── results/                 # Experiment results
│
├── run_examples.py          # Run data generator examples
└── README.md               # This file

Installation

Requirements

• Python 3.7+
• R 4.0+ (for benchmark methods)

Python Dependencies

pip install numpy pandas matplotlib seaborn scikit-learn

R Dependencies

install.packages(c("quantreg", "hetGP", "dplyr"))

Quick Start

1. Generate Data

First, generate data for your experiment:

# Using command line
python experiments/exp1/generate_data.py \
    --example exp1 \
    --n_train_points 100 \
    --n_reps 10 \
    --n_test_points 100 \
    --n_test_reps 1 \
    --n_macrorep 50 \
    --base_seed 42

Or in Python:

from experiments.exp1.generate_data import generate_data

data_dir = generate_data(
    example='exp1',
    n_train_points=100,
    n_reps=10,
    n_test_points=100,
    n_test_reps=1,
    n_macrorep=50
)
# Returns: 'data/data_YYYYMMDD_HHMMSS'

2. Run DCP-CKME Experiment

from experiments.exp1.run_dcp_ckme import run_dcp_ckme_experiment

results_df, macrorep_df, data_dir = run_dcp_ckme_experiment(
    example='exp1',
    data_dir='data/data_20250101_120000'  # Use generated data
)

3. Run Benchmark DCP Methods (R)

Edit experiments/exp1/run_dcp_methods.R and update data_dir:

data_dir <- "data/data_20250101_120000"

Then run:

Rscript experiments/exp1/run_dcp_methods.R

4. Generate Comparison Plots

from experiments.exp1.plot_comparison import main

main()

Module Documentation

Core Modules

core/models/ckme_weights.py

CKMEModel: Computes conditional kernel mean embedding weights.

Key Methods:

• compute_weights(x_star): Compute weights for a single test point
• compute_weights_batch(X_star): Batch computation for multiple test points
• simplex_weights(x_star): Normalized weights (sum to 1)

Parameters:

• ell_x: Length scale for X-space kernel
• lam: Regularization parameter
• sigma_y: Bandwidth for Y-space kernel

core/loss/mmd_loss.py

MMDLoss: Maximum Mean Discrepancy loss for parameter optimization.

Formula:

MMD² = (1/n) Σᵢ [k_Y(yᵢ, yᵢ) - 2 Σⱼ wⱼ(xᵢ) k_Y(yⱼ, yᵢ) + Σⱼ Σₖ wⱼ(xᵢ) wₖ(xᵢ) k_Y(yⱼ, yₖ)]

core/optimizers/optimizer.py

ParameterOptimizer: K-Fold Cross-Validation parameter optimizer.

Features:

• Uses GroupKFold to split by unique x values (sites)
• Grid search over ell_x, lam, and sigma_y multipliers
• Computes MMD loss on validation sets
• Returns best parameters based on average validation loss

Parameter Grids:

• ell_x: [10⁻², 10⁻¹·⁵, 10⁻¹, 10⁻⁰·⁵, 10⁰]
• lam: [1e-3, 1e-1, 10.0]
• sigma_y_multipliers: [0.5, 1.0, 2.0] × Silverman bandwidth

core/predictors/dcp_predictor.py

DCPPredictor: Distributional Conformal Prediction predictor.

Key Methods:

• compute_ecdf(x_star, y_values): Compute weighted empirical CDF
• compute_calibration_scores(X_cal, Y_cal): Compute DCP calibration scores
• predict_interval(x_star, calibration_scores, alpha): Predict conformal interval

Data Generators

data_generators/generator.py

DataGenerator: Unified data generator for MG1 and Exp1 examples.

Features:

• Generate training data with replicates
• Split training data by replicates (not random)
• Generate independent test data
• Support multiple macro-replications
• Auto-save to timestamped folders with metadata

See data_generators/README.md for detailed documentation.

Examples

examples/mg1.py

MG1 queue system example:

• True function: ζ(x) = 1.5x²/(1-x)
• Noise variance: r(x) = x(20 + 121x - 116x² + 29x³) / (4(1-x)⁴ · 2500)
• Default X range: (0.1, 0.9)

examples/exp1.py

Exp1 sinusoidal example:

• True function: ζ(x) = exp(x/10) · sin(x)
• Noise variance: r(x) = (0.01 + 0.2(x - π)²)²
• Default X range: (0, 2π)

Evaluation

evaluation/evaluator.py

PerformanceEvaluator: Performance metrics for conformal prediction.

Metrics:

• Coverage: Whether true value falls within interval (0 or 1)
• Width: Interval width (upper - lower)
• Interval Score: width + (2/α) · max(0, lower - y) + (2/α) · max(0, y - upper)
• Half Width: 1.96 · std / √n (95% CI half width)

Aggregation:

• Coverage/Width: Average across macro-replications for each test point
• Score: Average across all points within each macro-rep, then across macro-reps

Utilities

utils/kernels.py

Kernel functions:

• rbf_kernel(X, Y, sigma): RBF (Gaussian) kernel
• silverman_bandwidth(Y): Silverman's rule of thumb for bandwidth

Usage Examples

Example 1: Basic Data Generation

from data_generators import DataGenerator
import numpy as np

gen = DataGenerator(
    example='exp1',
    x_range=(0, 2 * np.pi),
    n_train_points=100,
    n_reps=10,
    n_test_points=100,
    n_test_reps=1,
    seed=42
)

# Generate and split data
X_train_full, Y_train_full = gen.generate_training_data()
X_train, Y_train, X_cal, Y_cal = gen.split_training_data(X_train_full, Y_train_full)
X_test, Y_test = gen.generate_test_data()

Example 2: Run Experiment with Custom Parameters

from experiments.exp1.run_dcp_ckme import run_dcp_ckme_experiment

results_df, macrorep_df, data_dir = run_dcp_ckme_experiment(
    example='exp1',
    n_train_points=200,
    n_reps=10,
    n_test_points=100,
    n_test_reps=1,
    n_macrorep=50,
    alpha=0.1,
    optimize_params=True,
    base_seed=42,
    data_dir=None  # Will generate new data
)

Example 3: Use Existing Data

# Use previously generated data
results_df, macrorep_df, _ = run_dcp_ckme_experiment(
    example='exp1',
    data_dir='data/data_20250101_120000'  # Use existing data
)

Example 4: Parameter Optimization

from core import CKMEModel, ParameterOptimizer
from data_generators import DataGenerator
import numpy as np

# Generate data
gen = DataGenerator(example='exp1', x_range=(0, 2*np.pi), 
                    n_train_points=100, n_reps=10, seed=42)
X_train, Y_train, _, _ = gen.split_training_data(*gen.generate_training_data())

# Optimize parameters
class Config:
    optimize_params = True

optimizer = ParameterOptimizer(Config(), k=5)

class Params:
    def __init__(self, ell_x, lam, sigma_y):
        self.ell_x = ell_x
        self.lam = lam
        self.sigma_y = sigma_y

best_params = optimizer.optimize(X_train, Y_train, Params)
print(f"Best ell_x: {best_params.ell_x}, lam: {best_params.lam}, sigma_y: {best_params.sigma_y}")

Experiments

Exp1 Experiment

The Exp1 experiment compares DCP-CKME with benchmark DCP methods on the Exp1 example.

Workflow:

1. Generate Data:

   python experiments/exp1/generate_data.py --example exp1 --n_macrorep 50

2. Run DCP-CKME (Python):

   from experiments.exp1.run_dcp_ckme import run_dcp_ckme_experiment
   run_dcp_ckme_experiment(example='exp1', data_dir='data/data_YYYYMMDD_HHMMSS')

3. Run Benchmark Methods (R):

   # Edit run_dcp_methods.R: data_dir <- "data/data_YYYYMMDD_HHMMSS"
   Rscript experiments/exp1/run_dcp_methods.R

4. Generate Comparison Plots:

   from experiments.exp1.plot_comparison import main
   main()

Output:

• Coverage comparison plot
• Width comparison plot
• Summary table with mean ± 95% CI half width

MG1 Experiment

Similar workflow, use example='mg1' in all scripts.

Data Consistency

To ensure Python and R code use the same data:

1. Always generate data first using generate_data.py
2. Use the same data_dir in both Python and R scripts
3. Check metadata.txt to verify data configuration

The data generation module automatically:

• Creates timestamped folders: data/data_YYYYMMDD_HHMMSS/
• Saves metadata.txt with configuration information
• Uses fixed random seeds for reproducibility

Dependencies

Python

• numpy: Numerical computations
• pandas: Data processing and CSV operations
• matplotlib: Plotting
• seaborn: Statistical visualization
• scikit-learn: Machine learning utilities (for GroupKFold)

R

• quantreg: Quantile regression (for DCP-QR methods)
• hetGP: Heteroscedastic Gaussian Process (for HetGP)
• dplyr: Data manipulation

File Naming Conventions

• Data files: train_{macrorep_id}.csv, cal_{macrorep_id}.csv, test_{macrorep_id}.csv
• Result files: {example}_dcp_ckme_{config}_{timestamp}.csv
• Data directories: data/data_{timestamp}/
• Result directories: results/{example}_dcp_ckme_{config}_{timestamp}/

Important Notes

1. Data Splitting: Always splits by replicates, ensuring each location appears in both train and calibration sets
2. Parameter Optimization: Uses K-Fold CV with GroupKFold to prevent information leakage
3. Score Aggregation: Interval score is averaged across all points within each macro-rep, then across macro-reps
4. Coverage/Width Aggregation: Averaged across macro-replications for each test point
5. Random Seeds: Fixed seeds ensure reproducibility; each macro-rep uses base_seed + macrorep_id