🎯 PIT-CP

pitcp is a Python package for conformal prediction using probability integral transform (PIT) pivotal scores. Given any black-box nonconformity score, it fits a conditional density estimator on the score distribution and maps raw scores to PIT values, yielding valid marginal coverage at any user-specified level.

Our contribution is PITCP. The package also reimplements the state-of-the-art SCP, CQR, HPD, and CONTRA baselines behind a consistent scikit-learn API.

✨ Features

• PIT conformal prediction: PITCP maps base nonconformity scores through a learned conditional CDF.
• Split conformal prediction: SCP calibrates arbitrary scalar nonconformity scores without a learned correction.
• Conformalized quantile regression: CQR accepts multiple outputs and provides a scikit-learn gradient-boosting implementation of state-of-the-art conformalized quantile regression.
• Highest-density regions: HPD calibrates conditional highest-predictive-density sets.
• Latent-space regions: CONTRA maps targets through a conditional normalizing flow and calibrates a Euclidean norm-based score in latent space.
• Model-agnostic: Works with any callable nonconformity score s(x, y), including distance, residual, or likelihood-based scores.
• Flexible Density Estimation: Supports normalizing flows and mixture density networks from the zuko library.
• Marginal Coverage Guarantee: Provably valid conformal coverage at any target level via finite-sample calibration.
• scikit-learn integration: Native BaseEstimator integration with familiar fit, conformalize, predict, and contains methods.

🚀 Installation

python -m pip install pitcp

🔧 Usage

Example

import torch
import zuko
from pitcp import PITCP


def std(x):
    return torch.where((x > -0.9) & (x < 0.9), torch.cos(torch.pi * x / 2), 1.0)


def gen_data(n):
    x = torch.rand(n, 1) * 2 - 1
    return x, torch.randn(n, 1) * std(x)


torch.manual_seed(42)

(X_train, y_train), (X_cal, y_cal), (X_test, y_test) = [
    gen_data(5000) for _ in range(3)
]


# Define a nonconformity score
def score(x, y):
    return y.abs()


# Build a normalizing flow density estimator
model = zuko.flows.NSF(features=1, context=1, bins=4, hidden_features=(32, 32, 32))
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)

# Compute nonconformity scores
s_train = score(X_train, y_train)
s_cal = score(X_cal, y_cal)
s_test = score(X_test, y_test)

# Fit and conformalize
pitcp = PITCP(model, optimizer, n_epochs=10, batch_size=128)
pitcp.fit(X_train, s_train)
pitcp.conformalize(X_cal, s_cal)

# Predict conformal regions (max score thresholds) at multiple quantiles
limits = pitcp.predict(X_test, confidence_level=[0.7, 0.8, 0.9])

# Predict conformal coverage
covered = pitcp.contains(X_test, s_test, confidence_level=[0.7, 0.8, 0.9])
print(f"Empirical coverages: {covered.mean(axis=0)}")

📖 Learn More

For tutorials, API reference, visit the official site:
👉 pitcp's documentation

The method is described in A Post-Processing Conformal Prediction Approach for Conditional Coverage via Pivotal Scores.

📊 Reproducing Results

Clone the repository, create and activate a virtual environment, and install the exact package versions used for the paper:

python -m pip install -r scripts/requirements.txt

This installs the frozen pitcp release from PyPI together with the experimental dependencies recorded in scripts/requirements.txt. Do not subsequently install the repository in editable mode when reproducing the paper, because that would replace the frozen PyPI release with the local source checkout.

For local development rather than exact paper reproduction, install the current checkout in editable mode from the repository root:

python -m pip install -e .

The editable installation makes local source changes immediately available without reinstalling the package.

Preparing the Real Data

The repository contains the raw SARCOS (data/sarcos_inv.mat) and Naval Propulsion Plants (data/naval.txt) datasets. To regenerate the training, validation-test, and prediction CSV files, install the TabPFN client, set the TABPFN_ACCESS_TOKEN environment variable to a valid access token, and run both dataset modes:

python -m pip install tabpfn-client

On PowerShell:

$env:TABPFN_ACCESS_TOKEN="your-token"
python data/predict.py --sarcos
python data/predict.py --naval

On Linux or macOS:

export TABPFN_ACCESS_TOKEN="your-token"
python data/predict.py --sarcos
python data/predict.py --naval

The token remains in the current terminal session and is not written to the repository.

These commands write {dataset}-train.csv, {dataset}-valtest.csv, and {dataset}-pred.csv to data/. The generated CSV files are already included, so this step can be skipped unless the predictions must be regenerated.

The paper experiments use Python 3.13 and the dependencies in scripts/requirements.txt. Verify downloaded and generated data against data/SHA256SUMS before running the experiments. The committed prediction files are the canonical reproduction artifacts because results returned by the remote TabPFN service may change independently of this repository.

On systems with sha256sum, verify the artifacts with:

cd data
sha256sum --check SHA256SUMS

Dataset Attribution and Licensing

The SARCOS inverse-dynamics dataset was provided by Sethu Vijayakumar and is distributed through the Gaussian Processes for Machine Learning dataset repository. Users should cite Vijayakumar and Schaal (2000), LWPR: An O(n) Algorithm for Incremental Real Time Learning in High Dimensional Space. No explicit redistribution license has been identified for SARCOS. Its inclusion in this repository should not be interpreted as granting reuse or redistribution rights.

The Naval Propulsion Plants dataset is licensed under the Creative Commons Attribution 4.0 International license. Users should cite Coraddu et al. (2014), Condition Based Maintenance of Naval Propulsion Plants, UCI Machine Learning Repository.

Running the Experiments

Run the experiment scripts from the repository root:

python scripts/convergence-plots.py
python scripts/synthetic-plots.py
python scripts/real-data-diagnostics.py --sarcos
python scripts/real-data-diagnostics.py --naval

The scripts resolve data and output paths from their file locations, so they do not depend on the current working directory. Figures and diagnostic tables are written to figures/.

Script Descriptions