pitcp.models.PITCP

class PITCP(estimator, optimizer, *, n_epochs=10, batch_size=None, verbose=True, random_state=None)[source]

PIT conformal predictor using a normalizing flow or mixture density estimator.

This class implements probability integral transform (PIT) conformal prediction. Given potentially black-box nonconformity scores, it fits a conditional density estimator on the score distribution over a training set, then uses the learned conditional CDF to map raw scores to PIT values. Conformal coverage guarantees are obtained by comparing test PIT values against a calibrated threshold.

The estimator must be a zuko subclass, coming from either zuko.flows.Flow (a normalizing flow) or zuko.mixtures.GMM (a mixture density network). The class internally detects which family is used and applies the appropriate CDF computation.

Density estimation settings:

  • estimator: A zuko lazy distribution instance conditioned on features, used to model the score distribution. Must be from zuko.flows or zuko.mixtures.

  • optimizer: PyTorch optimizer bound to estimator.parameters() and used to minimize the negative conditional log-likelihood.

Training settings:

  • n_epochs: Positive number of full passes over the training data. Defaults to 10.

  • batch_size: Positive mini-batch size used during training and inference. None uses the full dataset. Defaults to None.

  • verbose: Boolean or integer controlling the tqdm training progress bar. Defaults to True.

  • random_state: Seed controlling mini-batch shuffling during fit. None uses PyTorch’s current random state. Defaults to None.

Variables:

  • estimator (Flow | GMM) – Conditional density estimator from zuko.flows or zuko.mixtures.

  • optimizer (torch.optim.Optimizer) – Optimizer for training the estimator.

  • n_epochs (int) – Number of training epochs.

  • batch_size (int | None) – Batch size for data loading. None means full-batch training.

  • verbose (bool | int) – Whether to display a progress bar during training.

  • random_state (int | None) – Seed used to shuffle mini-batches during fit.

  • estimator_type (str) – Either flow or mixture, set during fit based on the type of estimator.

  • scores (torch.Tensor | None) – Calibration PIT scores stored after calling conformalize.

Parameters:

  • estimator (Flow | GMM)

  • optimizer (Optimizer)

  • n_epochs (int)

  • batch_size (int | None)

  • verbose (bool | int)

  • random_state (int | None)

Examples

>>> import torch
>>> import zuko
>>> from pitcp import PITCP
>>> X = torch.linspace(-1.0, 1.0, 32).reshape(-1, 1)
>>> s = torch.abs(torch.sin(3.0 * X)) + 0.1
>>> estimator = zuko.flows.NSF(features=1, context=1, bins=4)
>>> optimizer = torch.optim.Adam(estimator.parameters(), lr=1e-2)
>>> model = PITCP(estimator, optimizer, n_epochs=1, verbose=False)
>>> model.fit(X, s)
>>> model.conformalize(X, s)
>>> model.predict(X, confidence_level=0.9)
__init__(estimator, optimizer, *, n_epochs=10, batch_size=None, verbose=True, random_state=None)[source]

Initializes the PITCP instance.

Parameters:

  • estimator (Flow | GMM) – Conditional density estimator.

  • optimizer (torch.optim.Optimizer) – Optimizer for Train.

  • n_epochs (int, optional) – Number of Train epochs. Defaults to 10.

  • batch_size (int | None, optional) – Batch size for data loading. Defaults to None.

  • verbose (bool | int, optional) – Whether to show a Train progress bar. Defaults to True.

  • random_state (int | None, optional) – Seed used to shuffle mini-batches during fit. Defaults to None.

property estimator_type_: str[source]

Detects the type of the estimator.

Raises:

ValueError – If the estimator is not a subclass of zuko.flows.Flow or zuko.mixtures.GMM.

Returns:

Either flow or mixture.

Return type:

str

set_fit_request(*, s='$UNCHANGED$')

Configure whether metadata should be requested to be passed to the fit method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to fit if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to fit.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • s (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for s parameter in fit.

  • self (PITCP)

Returns:

self – The updated object.

Return type:

object

set_predict_request(*, confidence_level='$UNCHANGED$')

Configure whether metadata should be requested to be passed to the predict method.

Note that this method is only relevant when this estimator is used as a sub-estimator within a meta-estimator and metadata routing is enabled with enable_metadata_routing=True (see sklearn.set_config()). Please check the User Guide on how the routing mechanism works.

The options for each parameter are:

  • True: metadata is requested, and passed to predict if provided. The request is ignored if metadata is not provided.

  • False: metadata is not requested and the meta-estimator will not pass it to predict.

  • None: metadata is not requested, and the meta-estimator will raise an error if the user provides it.

  • str: metadata should be passed to the meta-estimator with this given alias instead of the original name.

The default (sklearn.utils.metadata_routing.UNCHANGED) retains the existing request. This allows you to change the request for some parameters and not others.

Added in version 1.3.

Parameters:

  • confidence_level (str, True, False, or None, default=sklearn.utils.metadata_routing.UNCHANGED) – Metadata routing for confidence_level parameter in predict.

  • self (PITCP)

Returns:

self – The updated object.

Return type:

object

fit(X, s)[source]

Fits the conditional density estimator on nonconformity scores.

Parameters:

  • X (np.typing.ArrayLike) – Training features with shape (n_samples, n_features).

  • s (np.typing.ArrayLike) – Training scores with shape (n_samples,).

Returns:

The fitted estimator.

Return type:

Self

conformalize(X, s)[source]

Computes and stores calibration PIT scores from a held-out dataset.

Parameters:

  • X (np.typing.ArrayLike) – Calibration features with shape (n_samples, n_features).

  • s (np.typing.ArrayLike) – Calibration scores with shape (n_samples,).

Returns:

The updated estimator.

Return type:

Self

predict(X, *, confidence_level=0.9)[source]

Predicts conformal regions for test points.

Parameters:

  • X (np.typing.ArrayLike) – Test features with shape (n_samples, n_features).

  • confidence_level (float | Sequence[float], optional) – Target coverage level(s). Defaults to 0.9.

Returns:

Score limits with shape (n_samples,) or ``(n_samples,

n_levels)``.

Return type:

np.ndarray

contains(X, s, *, confidence_level=0.9)[source]

Predicts conformal coverage for test points.

Parameters:

  • X (np.typing.ArrayLike) – Test features with shape (n_samples, n_features).

  • s (np.typing.ArrayLike) – Test scores with shape (n_samples,).

  • confidence_level (float | Sequence[float], optional) – Target coverage level(s). Defaults to 0.9.

Returns:

Coverage indicators with shape (n_samples,) or ``(n_samples,

n_levels)``.

Return type:

np.ndarray