API reference

Everything below is importable from the top-level kernel_calibration package (e.g. from kernel_calibration import KLCE_test). Plotting helpers live in the kernel_calibration.plots submodule.

The test and statistic

kernel_calibration.kite.KLCE_test

KLCE_test(X, Y, p, prob_kernel_width, iterations, key, x_kernel_width=None, add_one_correction=True)

Perform the KLCE hypothesis test comparing model predictions to true labels.

This function computes the test statistic and p-value by comparing the observed KLCE2 estimator against a null distribution generated by permutations. The null hypothesis is that the model is locally calibrated (KLCE^2 = 0); a small p-value is evidence against it.

Parameters:

Name	Type	Description	Default
X	array_like	Feature matrix of shape (n_samples, n_features). NumPy arrays are accepted.	required
Y	array_like	True label vector of shape (n_samples,).	required
p	array_like	Predicted probability vector of shape (n_samples,).	required
prob_kernel_width	float	Bandwidth for the probability kernel.	required
iterations	int	Number of permutations for null distribution.	required
key	Array or int	PRNG key for random operations. An integer is accepted and converted with jax.random.PRNGKey.	required
x_kernel_width	float	Bandwidth for the feature kernel. If omitted, prob_kernel_width is used for both kernels.	None
add_one_correction	bool	If True (default), use the Monte-Carlo permutation p-value (1 + #{null >= observed}) / (1 + iterations) (Phipson & Smyth, 2010), which is never exactly zero and controls the Type-I error rate. If False, use the uncorrected #{null > observed} / iterations floored at 1 / iterations.	True

Returns:

Type	Description
KLCETestResult	A result object that unpacks as (statistic, pvalue) and also exposes .statistic, .pvalue and .null_distribution (the permutation null samples, useful for plotting or a Type-I error check).

Source code in kernel_calibration/kite.py

def KLCE_test(
    X: jnp.ndarray,
    Y: jnp.ndarray,
    p: jnp.ndarray,
    prob_kernel_width: float,
    iterations: int,
    key,
    x_kernel_width: Optional[float] = None,
    add_one_correction: bool = True,
) -> tuple[float, float]:
    """
    Perform the KLCE hypothesis test comparing model predictions to true labels.

    This function computes the test statistic and p-value by comparing the observed
    KLCE2 estimator against a null distribution generated by permutations. The null
    hypothesis is that the model is locally calibrated (KLCE^2 = 0); a small p-value
    is evidence against it.

    Parameters
    ----------
    X : array_like
        Feature matrix of shape (n_samples, n_features). NumPy arrays are accepted.
    Y : array_like
        True label vector of shape (n_samples,).
    p : array_like
        Predicted probability vector of shape (n_samples,).
    prob_kernel_width : float
        Bandwidth for the probability kernel.
    iterations : int
        Number of permutations for null distribution.
    key : jax.Array or int
        PRNG key for random operations. An integer is accepted and converted with
        ``jax.random.PRNGKey``.
    x_kernel_width : float, optional
        Bandwidth for the feature kernel. If omitted, ``prob_kernel_width``
        is used for both kernels.
    add_one_correction : bool, optional
        If True (default), use the Monte-Carlo permutation p-value
        ``(1 + #{null >= observed}) / (1 + iterations)`` (Phipson & Smyth, 2010),
        which is never exactly zero and controls the Type-I error rate. If False,
        use the uncorrected ``#{null > observed} / iterations`` floored at
        ``1 / iterations``.

    Returns
    -------
    KLCETestResult
        A result object that unpacks as ``(statistic, pvalue)`` and also exposes
        ``.statistic``, ``.pvalue`` and ``.null_distribution`` (the permutation
        null samples, useful for plotting or a Type-I error check).
    """
    if x_kernel_width is None:
        x_kernel_width = prob_kernel_width
    X = jnp.asarray(X)
    Y = jnp.asarray(Y)
    p = jnp.asarray(p)
    if isinstance(key, int):
        key = random.PRNGKey(key)
    K = create_kernel(X, p, prob_kernel_width, x_kernel_width)
    p_err = Y - p
    test_value = KLCE2_estimator(K, p_err)
    test_null = compute_null_distribution(p_err, K, key, iterations)
    if add_one_correction:
        p_value = (1.0 + jnp.sum(test_null >= test_value)) / (1.0 + iterations)
    else:
        resolution = 1.0 / iterations
        p_value = jnp.maximum(resolution, resolution * jnp.sum(test_null > test_value))
    return KLCETestResult(statistic=test_value, pvalue=p_value, null_distribution=test_null)

kernel_calibration.kite.KLCETestResult dataclass

KLCETestResult(statistic, pvalue, null_distribution=None)

Result of :func:KLCE_test.

Unpacks as (statistic, pvalue) for backward compatibility (stat, p = KLCE_test(...)), and also exposes attribute access (.statistic, .pvalue, .null_distribution) in the style of SciPy's hypothesis-test result objects.

kernel_calibration.kite.KLCE2_estimator

KLCE2_estimator(K, err)

Compute the KLCE2 estimator from kernel matrix K and error vector err.

This function sums the off-diagonal elements of the elementwise product between K and the outer product of err with itself, normalized by n*(n-1).

Parameters:

Name	Type	Description	Default
K	ndarray	Kernel matrix of shape (n_samples, n_samples).	required
err	ndarray	Error vector of shape (n_samples,).	required

Returns:

Type	Description
float	KLCE2 estimator value.

Source code in kernel_calibration/kite.py

@jit
def KLCE2_estimator(K: jnp.ndarray, err: jnp.ndarray) -> float:
    """
    Compute the KLCE2 estimator from kernel matrix K and error vector err.

    This function sums the off-diagonal elements of the elementwise product
    between K and the outer product of err with itself, normalized by n*(n-1).

    Parameters
    ----------
    K : jnp.ndarray
        Kernel matrix of shape (n_samples, n_samples).
    err : jnp.ndarray
        Error vector of shape (n_samples,).

    Returns
    -------
    float
        KLCE2 estimator value.
    """
    if K.ndim != 2:
        raise ValueError(f"K must be a 2D array. Got ndim={K.ndim}.")
    if err.ndim != 1:
        raise ValueError(f"err must be a 1D array. Got ndim={err.ndim}.")
    if K.shape[0] != K.shape[1] or K.shape[0] != err.shape[0]:
        raise ValueError(
            f"Shape mismatch: K must be square of size n and err length n. "
            f"Got K.shape={K.shape}, err.shape={err.shape}."
        )

    err_outer = jnp.outer(err, err)
    K_err = K * err_outer

    mask = jnp.ones_like(K_err) - jnp.eye(K_err.shape[0])
    K_err_off_diag = K_err * mask

    n = err.shape[0]
    return jnp.sum(K_err_off_diag) / (n * (n - 1))

kernel_calibration.kite.create_kernel

create_kernel(X, p, prob_kernel_width, x_kernel_width)

Create combined kernel matrix for KLCE.

This function computes the elementwise product of two RBF kernels: one over predicted probabilities p and one over feature matrix X.

Parameters:

Name	Type	Description	Default
X	ndarray	Feature matrix of shape (n_samples, n_features).	required
p	ndarray	Predicted probabilities array of shape (n_samples,).	required
prob_kernel_width	float	Bandwidth parameter for the probability kernel.	required
x_kernel_width	float	Bandwidth parameter for the feature kernel.	required

Returns:

Type	Description
ndarray	Combined kernel matrix of shape (n_samples, n_samples).

Source code in kernel_calibration/kite.py

@jit
def create_kernel(
    X: jnp.ndarray, p: jnp.ndarray, prob_kernel_width: float, x_kernel_width: float
) -> jnp.ndarray:
    """
    Create combined kernel matrix for KLCE.

    This function computes the elementwise product of two RBF kernels:
    one over predicted probabilities p and one over feature matrix X.

    Parameters
    ----------
    X : jnp.ndarray
        Feature matrix of shape (n_samples, n_features).
    p : jnp.ndarray
        Predicted probabilities array of shape (n_samples,).
    prob_kernel_width : float
        Bandwidth parameter for the probability kernel.
    x_kernel_width : float
        Bandwidth parameter for the feature kernel.

    Returns
    -------
    jnp.ndarray
        Combined kernel matrix of shape (n_samples, n_samples).
    """
    if p.ndim != 1:
        raise ValueError(f"p must be a 1D array. Got ndim={p.ndim}.")
    if X.shape[0] != p.shape[0]:
        raise ValueError(
            f"Number of samples in X and p must match. Got {X.shape[0]} and {p.shape[0]}."
        )

    p_reshaped = p.reshape(-1, 1)
    gamma_p = 1.0 / (prob_kernel_width**2)
    gamma_x = 1.0 / (x_kernel_width**2)

    K_pp = rbf_kernel(p_reshaped, p_reshaped, gamma_p)
    K_xx = rbf_kernel(X, X, gamma_x)

    return K_pp * K_xx

kernel_calibration.kite.rbf_kernel

rbf_kernel(X, Y, gamma)

Compute the RBF (Gaussian) kernel matrix between X and Y.

This function computes the Gaussian (radial basis function) kernel elementwise between two datasets X and Y using kernel coefficient gamma.

Parameters:

Name	Type	Description	Default
X	ndarray	First data array of shape (n_samples, n_features) or (n_samples,) for single feature.	required
Y	ndarray	Second data array of shape (m_samples, n_features) or (m_samples,). Must have the same number of features as X, but may have a different number of samples — rectangular kernels are supported (used e.g. by the LCB diagnostic to evaluate at query points).	required
gamma	float	Kernel coefficient, typically defined as 1 / (sigma^2).	required

Returns:

Type	Description
ndarray	Kernel matrix of shape (n_samples, m_samples).

Source code in kernel_calibration/kite.py

def rbf_kernel(X: jnp.ndarray, Y: jnp.ndarray, gamma: float) -> jnp.ndarray:
    """
    Compute the RBF (Gaussian) kernel matrix between X and Y.

    This function computes the Gaussian (radial basis function) kernel
    elementwise between two datasets X and Y using kernel coefficient gamma.

    Parameters
    ----------
    X : jnp.ndarray
        First data array of shape (n_samples, n_features) or (n_samples,) for single feature.
    Y : jnp.ndarray
        Second data array of shape (m_samples, n_features) or (m_samples,). Must
        have the same number of features as X, but may have a different number of
        samples — rectangular kernels are supported (used e.g. by the LCB
        diagnostic to evaluate at query points).
    gamma : float
        Kernel coefficient, typically defined as 1 / (sigma^2).

    Returns
    -------
    jnp.ndarray
        Kernel matrix of shape (n_samples, m_samples).
    """
    if X.ndim == 1:
        X = X[:, None]
    if Y.ndim == 1:
        Y = Y[:, None]
    if X.shape[1] != Y.shape[1]:
        raise ValueError(
            f"X and Y must have the same number of features. Got {X.shape[1]} and {Y.shape[1]}."
        )
    squared_diff = (
        jnp.sum(X**2, axis=1)[:, None] + jnp.sum(Y**2, axis=1)[None, :] - 2 * jnp.dot(X, Y.T)
    )
    return jnp.exp(-gamma * squared_diff)

The diagnostic

kernel_calibration.lcb.local_calibration_bias

local_calibration_bias(X, y, f, prob_kernel_width, x_kernel_width=None, *, X_query=None, f_query=None, leave_one_out=False)

Estimate the local calibration bias at each query point.

For a query point (x', a') the estimator is the Nadaraya–Watson-style kernel-weighted average of residuals (Eq. 7 of the paper)::

LCB(x', a') = sum_i (y_i - f_i) k(f_i, a') l(x_i, x')

                     sum_i k(f_i, a') l(x_i, x')

with k an RBF kernel on predicted probabilities and l an RBF kernel on features. Setting the feature kernel aside (a single feature value) this reduces to the usual reliability-curve bias; with the feature kernel it localises the bias in feature space, revealing e.g. subgroups where the model is miscalibrated.

Parameters:

Name	Type	Description	Default
X	array_like	Reference feature matrix of shape (n_samples, n_features). NumPy arrays are accepted.	required
y	array_like	Reference labels of shape (n_samples,).	required
f	array_like	Reference predicted probabilities of shape (n_samples,).	required
prob_kernel_width	float	Bandwidth for the probability kernel k.	required
x_kernel_width	float	Bandwidth for the feature kernel l. If omitted, prob_kernel_width is used for both kernels.	None
X_query	array_like	Feature points at which to evaluate the bias, shape (m_samples, n_features). Defaults to X (evaluate at every reference point).	None
f_query	array_like	Predicted probabilities at the query points, shape (m_samples,). Required if X_query is given; defaults to f when evaluating at X.	None
leave_one_out	bool	Only used when evaluating at the reference points themselves (no explicit X_query). If True, each point's own residual is excluded from its estimate, giving an honest (leave-one-out) bias. Default is False.	False

Returns:

Type	Description
ndarray	Estimated local calibration bias at each query point, shape (m_samples,).

Examples:

>>> bias = local_calibration_bias(X, y, f, prob_kernel_width=0.1)
>>> # bias[i] > 0  -> model under-confident near sample i
>>> # bias[i] < 0  -> model over-confident near sample i

Source code in kernel_calibration/lcb.py

def local_calibration_bias(
    X: jnp.ndarray,
    y: jnp.ndarray,
    f: jnp.ndarray,
    prob_kernel_width: float,
    x_kernel_width: Optional[float] = None,
    *,
    X_query: Optional[jnp.ndarray] = None,
    f_query: Optional[jnp.ndarray] = None,
    leave_one_out: bool = False,
) -> jnp.ndarray:
    """Estimate the local calibration bias at each query point.

    For a query point ``(x', a')`` the estimator is the Nadaraya–Watson-style
    kernel-weighted average of residuals (Eq. 7 of the paper)::

        LCB(x', a') = sum_i (y_i - f_i) k(f_i, a') l(x_i, x')
                      -------------------------------------------
                             sum_i k(f_i, a') l(x_i, x')

    with ``k`` an RBF kernel on predicted probabilities and ``l`` an RBF kernel on
    features. Setting the feature kernel aside (a single feature value) this reduces
    to the usual reliability-curve bias; with the feature kernel it localises the
    bias in feature space, revealing e.g. subgroups where the model is miscalibrated.

    Parameters
    ----------
    X : array_like
        Reference feature matrix of shape (n_samples, n_features). NumPy arrays are
        accepted.
    y : array_like
        Reference labels of shape (n_samples,).
    f : array_like
        Reference predicted probabilities of shape (n_samples,).
    prob_kernel_width : float
        Bandwidth for the probability kernel ``k``.
    x_kernel_width : float, optional
        Bandwidth for the feature kernel ``l``. If omitted, ``prob_kernel_width``
        is used for both kernels.
    X_query : array_like, optional
        Feature points at which to evaluate the bias, shape (m_samples, n_features).
        Defaults to ``X`` (evaluate at every reference point).
    f_query : array_like, optional
        Predicted probabilities at the query points, shape (m_samples,). Required
        if ``X_query`` is given; defaults to ``f`` when evaluating at ``X``.
    leave_one_out : bool, optional
        Only used when evaluating at the reference points themselves (no explicit
        ``X_query``). If True, each point's own residual is excluded from its
        estimate, giving an honest (leave-one-out) bias. Default is False.

    Returns
    -------
    jnp.ndarray
        Estimated local calibration bias at each query point, shape (m_samples,).

    Examples
    --------
    >>> bias = local_calibration_bias(X, y, f, prob_kernel_width=0.1)
    >>> # bias[i] > 0  -> model under-confident near sample i
    >>> # bias[i] < 0  -> model over-confident near sample i
    """
    if x_kernel_width is None:
        x_kernel_width = prob_kernel_width

    X = jnp.asarray(X)
    y = jnp.asarray(y)
    f = jnp.asarray(f)
    err = y - f

    evaluating_on_self = X_query is None and f_query is None
    if X_query is None:
        X_query = X
    if f_query is None:
        if not evaluating_on_self:
            raise ValueError("f_query must be provided when X_query is given.")
        f_query = f
    X_query = jnp.asarray(X_query)
    f_query = jnp.asarray(f_query)
    if X_query.shape[0] != f_query.shape[0]:
        raise ValueError(
            f"X_query and f_query must have the same number of samples. "
            f"Got {X_query.shape[0]} and {f_query.shape[0]}."
        )

    gamma_p = 1.0 / (prob_kernel_width**2)
    gamma_x = 1.0 / (x_kernel_width**2)

    # Rectangular kernels between the n reference points and m query points.
    K_p = rbf_kernel(f.reshape(-1, 1), f_query.reshape(-1, 1), gamma_p)  # (n, m)
    K_x = rbf_kernel(X, X_query, gamma_x)  # (n, m)
    weights = K_p * K_x  # (n, m)

    if leave_one_out and evaluating_on_self:
        weights = weights * (1.0 - jnp.eye(weights.shape[0]))

    numerator = weights.T @ err  # (m,)
    denominator = jnp.sum(weights, axis=0)  # (m,)
    return numerator / denominator

Recalibration

kernel_calibration.kite.recalibrated_model

recalibrated_model(sigma_k=0.1, sigma_l=1.0, alpha=0.5, beta=0.5, num_steps=1000, learning_rate=0.001, hidden_layer_sizes=(64, 64), seed=121, verbose=False)

Recalibration model combining distillation loss and KLCE penalty.

This class implements a simple MLP-based recalibration of base probabilities trained to minimize a combination of KL divergence and kernel-based calibration error.

Initialize hyperparameters for the recalibration model.

Parameters:

Name	Type	Description	Default
sigma_k	float	Kernel width for probability kernel. Default is 0.1.	0.1
sigma_l	float	Kernel width for feature kernel. Default is 1.0.	1.0
alpha	float	Weight for the distillation loss term. Default is 0.5.	0.5
beta	float	Weight for the KLCE penalty term. Default is 0.5.	0.5
num_steps	int	Number of training steps. Default is 1000.	1000
learning_rate	float	Optimizer learning rate. Default is 0.001.	0.001
hidden_layer_sizes	Tuple[int, ...]	Sizes of hidden MLP layers. Default is (64, 64).	(64, 64)
seed	int	Random seed for initialization. Default is 121.	121
verbose	bool	If True, print the training loss every 100 steps. Default is False.	False

Source code in kernel_calibration/kite.py

def __init__(
    self,
    sigma_k: float = 0.1,
    sigma_l: float = 1.0,
    alpha: float = 0.5,
    beta: float = 0.5,
    num_steps: int = 1000,
    learning_rate: float = 0.001,
    hidden_layer_sizes: tuple[int, ...] = (64, 64),
    seed: int = 121,
    verbose: bool = False,
) -> None:
    """
    Initialize hyperparameters for the recalibration model.

    Parameters
    ----------
    sigma_k : float, optional
        Kernel width for probability kernel. Default is 0.1.
    sigma_l : float, optional
        Kernel width for feature kernel. Default is 1.0.
    alpha : float, optional
        Weight for the distillation loss term. Default is 0.5.
    beta : float, optional
        Weight for the KLCE penalty term. Default is 0.5.
    num_steps : int, optional
        Number of training steps. Default is 1000.
    learning_rate : float, optional
        Optimizer learning rate. Default is 0.001.
    hidden_layer_sizes : Tuple[int, ...], optional
        Sizes of hidden MLP layers. Default is (64, 64).
    seed : int, optional
        Random seed for initialization. Default is 121.
    verbose : bool, optional
        If True, print the training loss every 100 steps. Default is False.
    """
    self.sigma_k = sigma_k
    self.sigma_l = sigma_l
    self.alpha = alpha
    self.beta = beta
    self.num_steps = num_steps
    self.learning_rate = learning_rate
    self.hidden_layer_sizes = hidden_layer_sizes
    self.seed = seed
    self.verbose = verbose
    self.params: Optional[dict[str, jnp.ndarray]] = None
    self.loss_history: Optional[list[float]] = None

accuracy_score

accuracy_score(y_pred, y)

Compute the classification accuracy.

Parameters:

Name	Type	Description	Default
y_pred	Sequence[int]	Predicted labels of shape (n_samples,).	required
y	Sequence[int]	True labels of shape (n_samples,).	required

Returns:

Type	Description
float	Proportion of correctly classified samples.

Source code in kernel_calibration/kite.py

def accuracy_score(self, y_pred: Sequence[int], y: Sequence[int]) -> float:
    """
    Compute the classification accuracy.

    Parameters
    ----------
    y_pred : Sequence[int]
        Predicted labels of shape (n_samples,).
    y : Sequence[int]
        True labels of shape (n_samples,).

    Returns
    -------
    float
        Proportion of correctly classified samples.
    """
    predictions = jnp.array(y_pred)
    actual_labels = jnp.array(y)
    return float(jnp.mean(predictions == actual_labels))

fit

fit(y_proba, x_cal, y)

Train the recalibration model on calibration data.

This method optimizes model parameters to minimize the total loss over specified number of steps using the Adam optimizer.

Parameters:

Name	Type	Description	Default
y_proba	ndarray	Base probability predictions of shape (n_samples,).	required
x_cal	ndarray	Calibration feature matrix of shape (n_samples, n_features).	required
y	ndarray	True labels of shape (n_samples,).	required

Returns:

Type	Description
recalibrated_model	The fitted estimator (self), to allow method chaining.

Source code in kernel_calibration/kite.py

def fit(self, y_proba: jnp.ndarray, x_cal: jnp.ndarray, y: jnp.ndarray) -> "recalibrated_model":
    """
    Train the recalibration model on calibration data.

    This method optimizes model parameters to minimize the total loss
    over specified number of steps using the Adam optimizer.

    Parameters
    ----------
    y_proba : jnp.ndarray
        Base probability predictions of shape (n_samples,).
    x_cal : jnp.ndarray
        Calibration feature matrix of shape (n_samples, n_features).
    y : jnp.ndarray
        True labels of shape (n_samples,).

    Returns
    -------
    recalibrated_model
        The fitted estimator (``self``), to allow method chaining.
    """
    y_proba = jnp.asarray(y_proba)
    x_cal = jnp.asarray(x_cal)
    y = jnp.asarray(y)
    rng = random.PRNGKey(self.seed)
    if x_cal.ndim == 1:
        x_cal = x_cal[:, None]
    input_dim = 2 + x_cal.shape[1]
    layer_sizes = [input_dim] + list(self.hidden_layer_sizes) + [1]
    params = init_recalibrated_model_params(rng, layer_sizes)
    optimizer = optax.adam(self.learning_rate)
    opt_state = optimizer.init(params)

    @jit
    def step(
        params: dict[str, jnp.ndarray], opt_state: optax.OptState
    ) -> tuple[dict[str, jnp.ndarray], optax.OptState, float]:
        loss_val, grads = jax.value_and_grad(self.total_loss)(params, y_proba, x_cal, y)
        updates, opt_state = optimizer.update(grads, opt_state)
        params = optax.apply_updates(params, updates)
        return params, opt_state, loss_val

    loss_history: list[float] = []
    for i in range(self.num_steps):
        params, opt_state, loss_val = step(params, opt_state)
        loss_history.append(float(loss_val))
        if self.verbose and i % 100 == 0:
            print(f"Step {i}: total loss = {loss_val:.6f}")
    self.params = params
    self.loss_history = loss_history
    return self

get_labels

get_labels(y_proba, threshold=0.5)

Convert probability predictions to binary labels.

Parameters:

Name	Type	Description	Default
y_proba	ndarray	Probability vector of shape (n_samples,).	required
threshold	float	Classification threshold. Default is 0.5.	0.5

Returns:

Type	Description
List[int]	Binary labels (0 or 1) for each sample.

Source code in kernel_calibration/kite.py

def get_labels(self, y_proba: jnp.ndarray, threshold: float = 0.5) -> list[int]:
    """
    Convert probability predictions to binary labels.

    Parameters
    ----------
    y_proba : jnp.ndarray
        Probability vector of shape (n_samples,).
    threshold : float, optional
        Classification threshold. Default is 0.5.

    Returns
    -------
    List[int]
        Binary labels (0 or 1) for each sample.
    """
    return [1 if y > threshold else 0 for y in y_proba]

get_params

get_params(deep=True)

Return the model hyperparameters (scikit-learn estimator API).

Source code in kernel_calibration/kite.py

def get_params(self, deep: bool = True) -> dict[str, object]:
    """Return the model hyperparameters (scikit-learn estimator API)."""
    return {
        "sigma_k": self.sigma_k,
        "sigma_l": self.sigma_l,
        "alpha": self.alpha,
        "beta": self.beta,
        "num_steps": self.num_steps,
        "learning_rate": self.learning_rate,
        "hidden_layer_sizes": self.hidden_layer_sizes,
        "seed": self.seed,
        "verbose": self.verbose,
    }

predict_proba

predict_proba(y_proba, x_cal)

Generate recalibrated probability predictions.

This method applies the trained recalibration model to new data, returning corrected probability estimates.

Parameters:

Name	Type	Description	Default
y_proba	ndarray	Base probability predictions of shape (n_samples,).	required
x_cal	ndarray	Calibration features of shape (n_samples, n_features).	required

Returns:

Type	Description
ndarray	Recalibrated probability vector of shape (n_samples,).

Source code in kernel_calibration/kite.py

def predict_proba(self, y_proba: jnp.ndarray, x_cal: jnp.ndarray) -> jnp.ndarray:
    """
    Generate recalibrated probability predictions.

    This method applies the trained recalibration model to new data,
    returning corrected probability estimates.

    Parameters
    ----------
    y_proba : jnp.ndarray
        Base probability predictions of shape (n_samples,).
    x_cal : jnp.ndarray
        Calibration features of shape (n_samples, n_features).

    Returns
    -------
    jnp.ndarray
        Recalibrated probability vector of shape (n_samples,).
    """
    if self.params is None:
        raise RuntimeError("Call fit() before predict_proba().")
    y_proba = jnp.asarray(y_proba)
    x_cal = jnp.asarray(x_cal)
    m = y_proba.shape[0]
    if x_cal.ndim == 1:
        x_cal = x_cal[:, None]
    features_new = jnp.column_stack([jnp.ones(m), y_proba, x_cal])
    correction = recalibrated_model_apply(self.params, features_new).squeeze()
    f_recalibrated_new = y_proba + correction
    return jnp.clip(f_recalibrated_new, _PROB_EPS, 1.0 - _PROB_EPS)

set_params

set_params(**params)

Set model hyperparameters (scikit-learn estimator API). Returns self.

Source code in kernel_calibration/kite.py

def set_params(self, **params) -> "recalibrated_model":
    """Set model hyperparameters (scikit-learn estimator API). Returns self."""
    valid = self.get_params()
    for name, value in params.items():
        if name not in valid:
            raise ValueError(
                f"Invalid parameter {name!r} for recalibrated_model. "
                f"Valid parameters are: {sorted(valid)}."
            )
        setattr(self, name, value)
    return self

total_loss

total_loss(params, base_probs, x, y)

Compute the total loss combining distillation and KLCE penalty.

This function calculates the KL divergence between base_probs and recalibrated predictions, then adds the kernel-based calibration error.

Parameters:

Name	Type	Description	Default
params	Dict[str, ndarray]	Recalibration model parameters.	required
base_probs	ndarray	Original predicted probabilities of shape (n_samples,).	required
x	ndarray	Calibration features of shape (n_samples, n_features).	required
y	ndarray	True labels of shape (n_samples,).	required

Returns:

Type	Description
float	Weighted sum of distillation loss and KLCE penalty.

Source code in kernel_calibration/kite.py

def total_loss(
    self,
    params: dict[str, jnp.ndarray],
    base_probs: jnp.ndarray,
    x: jnp.ndarray,
    y: jnp.ndarray,
) -> float:
    """
    Compute the total loss combining distillation and KLCE penalty.

    This function calculates the KL divergence between base_probs and
    recalibrated predictions, then adds the kernel-based calibration error.

    Parameters
    ----------
    params : Dict[str, jnp.ndarray]
        Recalibration model parameters.
    base_probs : jnp.ndarray
        Original predicted probabilities of shape (n_samples,).
    x : jnp.ndarray
        Calibration features of shape (n_samples, n_features).
    y : jnp.ndarray
        True labels of shape (n_samples,).

    Returns
    -------
    float
        Weighted sum of distillation loss and KLCE penalty.
    """
    n = base_probs.shape[0]
    features = jnp.column_stack([jnp.ones(n), base_probs, x])
    correction = recalibrated_model_apply(params, features).squeeze()
    f_recalibrated = base_probs + correction
    f_recalibrated = jnp.clip(f_recalibrated, _PROB_EPS, 1.0 - _PROB_EPS)
    base_probs_stable = jnp.clip(base_probs, _PROB_EPS, 1.0 - _PROB_EPS)
    log_ratio1 = jnp.log(jnp.maximum(base_probs_stable / f_recalibrated, 1e-10))
    log_ratio2 = jnp.log(jnp.maximum((1 - base_probs_stable) / (1 - f_recalibrated), 1e-10))
    kl_div = base_probs_stable * log_ratio1 + (1 - base_probs_stable) * log_ratio2
    distill_loss = jnp.mean(kl_div)
    x_2d = x[:, None] if x.ndim == 1 else x
    klce_loss = KLCE2_boosting(f_recalibrated, x_2d, y, self.sigma_k, self.sigma_l)
    return self.alpha * distill_loss + self.beta * klce_loss

Bandwidth selection

kernel_calibration.bandwidth.select_bandwidths

select_bandwidths(X, f, max_samples=2000, seed=0)

Median-heuristic bandwidths for the probability and feature kernels.

Convenience wrapper that returns sensible defaults for both kernels used by :func:~kernel_calibration.kite.KLCE_test and :func:~kernel_calibration.lcb.local_calibration_bias.

Parameters:

Name	Type	Description	Default
X	array_like	Feature matrix of shape (n_samples, n_features) or (n_samples,).	required
f	array_like	Predicted probabilities of shape (n_samples,).	required
max_samples	int	Passed through to :func:median_heuristic.	2000
seed	int	Passed through to :func:median_heuristic.	2000

Returns:

Type	Description
Tuple[float, float]	(prob_kernel_width, x_kernel_width).

Examples:

>>> pw, xw = select_bandwidths(X, f)
>>> stat, pval = KLCE_test(X, y, f, prob_kernel_width=pw, iterations=200,
...                        key=0, x_kernel_width=xw)

Source code in kernel_calibration/bandwidth.py

def select_bandwidths(
    X,
    f,
    max_samples: int = 2000,
    seed: int = 0,
) -> tuple[float, float]:
    """Median-heuristic bandwidths for the probability and feature kernels.

    Convenience wrapper that returns sensible defaults for both kernels used by
    :func:`~kernel_calibration.kite.KLCE_test` and
    :func:`~kernel_calibration.lcb.local_calibration_bias`.

    Parameters
    ----------
    X : array_like
        Feature matrix of shape (n_samples, n_features) or (n_samples,).
    f : array_like
        Predicted probabilities of shape (n_samples,).
    max_samples, seed
        Passed through to :func:`median_heuristic`.

    Returns
    -------
    Tuple[float, float]
        ``(prob_kernel_width, x_kernel_width)``.

    Examples
    --------
    >>> pw, xw = select_bandwidths(X, f)
    >>> stat, pval = KLCE_test(X, y, f, prob_kernel_width=pw, iterations=200,
    ...                        key=0, x_kernel_width=xw)
    """
    prob_width = median_heuristic(f, max_samples=max_samples, seed=seed)
    x_width = median_heuristic(X, max_samples=max_samples, seed=seed)
    return prob_width, x_width

kernel_calibration.bandwidth.median_heuristic

median_heuristic(values, max_samples=2000, seed=0)

Median-heuristic bandwidth: the median pairwise Euclidean distance.

Parameters:

Name	Type	Description	Default
values	array_like	Data of shape (n_samples, n_features) or (n_samples,) for a single feature (e.g. predicted probabilities).	required
max_samples	int	If n_samples exceeds this, a random subsample of this size is used to keep the pairwise computation O(max_samples^2). Default is 2000.	2000
seed	int	Seed for the subsampling RNG (only used when subsampling). Default is 0.	0

Returns:

Type	Description
float	The median non-zero pairwise distance, usable directly as a kernel width. Falls back to 1.0 if all points coincide (median distance is zero).

Source code in kernel_calibration/bandwidth.py

def median_heuristic(
    values,
    max_samples: int = 2000,
    seed: int = 0,
) -> float:
    """Median-heuristic bandwidth: the median pairwise Euclidean distance.

    Parameters
    ----------
    values : array_like
        Data of shape (n_samples, n_features) or (n_samples,) for a single feature
        (e.g. predicted probabilities).
    max_samples : int, optional
        If ``n_samples`` exceeds this, a random subsample of this size is used to
        keep the pairwise computation O(max_samples^2). Default is 2000.
    seed : int, optional
        Seed for the subsampling RNG (only used when subsampling). Default is 0.

    Returns
    -------
    float
        The median non-zero pairwise distance, usable directly as a kernel width.
        Falls back to ``1.0`` if all points coincide (median distance is zero).
    """
    values = np.asarray(values, dtype=float)
    if values.ndim == 1:
        values = values[:, None]
    n = values.shape[0]
    if n > max_samples:
        rng = np.random.default_rng(seed)
        idx = rng.choice(n, size=max_samples, replace=False)
        values = values[idx]
        n = max_samples

    # Pairwise squared distances, then the upper triangle (i < j).
    sq = (
        np.sum(values**2, axis=1)[:, None]
        + np.sum(values**2, axis=1)[None, :]
        - 2.0 * values @ values.T
    )
    sq = np.clip(sq, 0.0, None)
    iu = np.triu_indices(n, k=1)
    dists = np.sqrt(sq[iu])
    dists = dists[dists > 0]
    if dists.size == 0:
        return 1.0
    return float(np.median(dists))

Baseline metrics

kernel_calibration.metrics.expected_calibration_error

expected_calibration_error(y, p, n_bins=15)

Expected Calibration Error (ECE) with equal-width bins.

ECE = sum over bins of (bin_weight) * |accuracy - confidence|.

Parameters:

Name	Type	Description	Default
y	array_like	Binary labels of shape (n_samples,).	required
p	array_like	Predicted probabilities of shape (n_samples,).	required
n_bins	int	Number of equal-width bins in [0, 1]. Default is 15.	15

Returns:

Type	Description
float	The ECE. Lower is better; 0 means perfectly calibrated on this binning.

Source code in kernel_calibration/metrics.py

def expected_calibration_error(y, p, n_bins: int = 15) -> float:
    """Expected Calibration Error (ECE) with equal-width bins.

    ECE = sum over bins of ``(bin_weight) * |accuracy - confidence|``.

    Parameters
    ----------
    y : array_like
        Binary labels of shape (n_samples,).
    p : array_like
        Predicted probabilities of shape (n_samples,).
    n_bins : int, optional
        Number of equal-width bins in [0, 1]. Default is 15.

    Returns
    -------
    float
        The ECE. Lower is better; 0 means perfectly calibrated on this binning.
    """
    y = np.asarray(y, dtype=float)
    p = np.asarray(p, dtype=float)
    return float(sum(w * abs(acc - conf) for w, acc, conf in _bin_stats(y, p, n_bins)))

kernel_calibration.metrics.maximum_calibration_error

maximum_calibration_error(y, p, n_bins=15)

Maximum Calibration Error (MCE): the worst per-bin calibration gap.

Parameters:

Name	Type	Description	Default
y	array_like	Labels and predicted probabilities of shape (n_samples,).	required
p	array_like	Labels and predicted probabilities of shape (n_samples,).	required
n_bins	int	Number of equal-width bins in [0, 1]. Default is 15.	15

Returns:

Type	Description
float	max_bin |accuracy - confidence|; 0.0 if there are no samples.

Source code in kernel_calibration/metrics.py

def maximum_calibration_error(y, p, n_bins: int = 15) -> float:
    """Maximum Calibration Error (MCE): the worst per-bin calibration gap.

    Parameters
    ----------
    y, p : array_like
        Labels and predicted probabilities of shape (n_samples,).
    n_bins : int, optional
        Number of equal-width bins in [0, 1]. Default is 15.

    Returns
    -------
    float
        ``max_bin |accuracy - confidence|``; 0.0 if there are no samples.
    """
    y = np.asarray(y, dtype=float)
    p = np.asarray(p, dtype=float)
    gaps = [abs(acc - conf) for _, acc, conf in _bin_stats(y, p, n_bins)]
    return float(max(gaps)) if gaps else 0.0

kernel_calibration.metrics.brier_score

brier_score(y, p)

Brier score: mean squared error between probabilities and outcomes.

Parameters:

Name	Type	Description	Default
y	array_like	Binary labels of shape (n_samples,).	required
p	array_like	Predicted probabilities of shape (n_samples,).	required

Returns:

Type	Description
float	Mean of (p - y) ** 2. Lower is better.

Source code in kernel_calibration/metrics.py

def brier_score(y, p) -> float:
    """Brier score: mean squared error between probabilities and outcomes.

    Parameters
    ----------
    y : array_like
        Binary labels of shape (n_samples,).
    p : array_like
        Predicted probabilities of shape (n_samples,).

    Returns
    -------
    float
        Mean of ``(p - y) ** 2``. Lower is better.
    """
    y = np.asarray(y, dtype=float)
    p = np.asarray(p, dtype=float)
    return float(np.mean((p - y) ** 2))

kernel_calibration.metrics.kernel_calibration_error

kernel_calibration_error(y, p, prob_kernel_width)

Squared Kernel Calibration Error (KCE) of Widmann et al. (2019).

This is the global (non-local) counterpart of KLCE: the same U-statistic on residuals y - p but with only the probability kernel and no feature kernel (equivalently, the feature kernel set to a constant). It measures whether the model is calibrated on average, without localising in feature space.

Parameters:

Name	Type	Description	Default
y	array_like	Binary labels of shape (n_samples,).	required
p	array_like	Predicted probabilities of shape (n_samples,).	required
prob_kernel_width	float	Bandwidth for the RBF probability kernel.	required

Returns:

Type	Description
float	The squared KCE estimate (an unbiased U-statistic, diagonal removed).

Source code in kernel_calibration/metrics.py

def kernel_calibration_error(y, p, prob_kernel_width: float) -> float:
    """Squared Kernel Calibration Error (KCE) of Widmann et al. (2019).

    This is the global (non-local) counterpart of KLCE: the same U-statistic on
    residuals ``y - p`` but with only the probability kernel and no feature kernel
    (equivalently, the feature kernel set to a constant). It measures whether the
    model is calibrated on average, without localising in feature space.

    Parameters
    ----------
    y : array_like
        Binary labels of shape (n_samples,).
    p : array_like
        Predicted probabilities of shape (n_samples,).
    prob_kernel_width : float
        Bandwidth for the RBF probability kernel.

    Returns
    -------
    float
        The squared KCE estimate (an unbiased U-statistic, diagonal removed).
    """
    y = jnp.asarray(y)
    p = jnp.asarray(p)
    err = y - p
    gamma_p = 1.0 / (prob_kernel_width**2)
    K = rbf_kernel(p.reshape(-1, 1), p.reshape(-1, 1), gamma_p)
    return float(KLCE2_estimator(K, err))

Plotting

kernel_calibration.plots.reliability_diagram

reliability_diagram(y, p, n_bins=15, ax=None, label=None)

Plot a reliability diagram (accuracy vs. confidence per bin).

Parameters:

Name	Type	Description	Default
y	array_like	Binary labels of shape (n_samples,).	required
p	array_like	Predicted probabilities of shape (n_samples,).	required
n_bins	int	Number of equal-width bins in [0, 1]. Default is 15.	15
ax	Axes	Axes to draw on. A new figure/axes is created if omitted.	None
label	str	Legend label for the model curve.	None

Returns:

Type	Description
Axes	The axes containing the plot.

Source code in kernel_calibration/plots.py

def reliability_diagram(y, p, n_bins: int = 15, ax=None, label: Optional[str] = None):
    """Plot a reliability diagram (accuracy vs. confidence per bin).

    Parameters
    ----------
    y : array_like
        Binary labels of shape (n_samples,).
    p : array_like
        Predicted probabilities of shape (n_samples,).
    n_bins : int, optional
        Number of equal-width bins in [0, 1]. Default is 15.
    ax : matplotlib.axes.Axes, optional
        Axes to draw on. A new figure/axes is created if omitted.
    label : str, optional
        Legend label for the model curve.

    Returns
    -------
    matplotlib.axes.Axes
        The axes containing the plot.
    """
    plt = _require_matplotlib()
    y = np.asarray(y, dtype=float)
    p = np.asarray(p, dtype=float)
    if ax is None:
        _, ax = plt.subplots(figsize=(4.5, 4.5))

    confidences, accuracies = [], []
    for _, acc, conf in _bin_stats(y, p, n_bins):
        confidences.append(conf)
        accuracies.append(acc)

    ax.plot([0, 1], [0, 1], linestyle="--", color="grey", label="perfectly calibrated")
    ax.plot(confidences, accuracies, marker="o", label=label or "model")
    ax.set_xlabel("Mean predicted probability")
    ax.set_ylabel("Fraction of positives")
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.set_title("Reliability diagram")
    ax.legend(loc="best")
    return ax

kernel_calibration.plots.plot_local_calibration_bias

plot_local_calibration_bias(feature, bias, ax=None, xlabel='feature')

Scatter the local calibration bias against a single feature.

Useful for the "where is the model miscalibrated" view: points far from zero mark regions of feature space where predictions are over- or under-confident.

Parameters:

Name	Type	Description	Default
feature	array_like	A 1-D feature value per sample, shape (n_samples,).	required
bias	array_like	Local calibration bias per sample (e.g. the output of :func:~kernel_calibration.lcb.local_calibration_bias), shape (n_samples,).	required
ax	Axes	Axes to draw on. A new figure/axes is created if omitted.	None
xlabel	str	Label for the feature axis. Default is "feature".	'feature'

Returns:

Type	Description
Axes	The axes containing the plot.

Source code in kernel_calibration/plots.py

def plot_local_calibration_bias(feature, bias, ax=None, xlabel: str = "feature"):
    """Scatter the local calibration bias against a single feature.

    Useful for the "where is the model miscalibrated" view: points far from zero
    mark regions of feature space where predictions are over- or under-confident.

    Parameters
    ----------
    feature : array_like
        A 1-D feature value per sample, shape (n_samples,).
    bias : array_like
        Local calibration bias per sample (e.g. the output of
        :func:`~kernel_calibration.lcb.local_calibration_bias`), shape (n_samples,).
    ax : matplotlib.axes.Axes, optional
        Axes to draw on. A new figure/axes is created if omitted.
    xlabel : str, optional
        Label for the feature axis. Default is "feature".

    Returns
    -------
    matplotlib.axes.Axes
        The axes containing the plot.
    """
    plt = _require_matplotlib()
    feature = np.asarray(feature, dtype=float)
    bias = np.asarray(bias, dtype=float)
    if ax is None:
        _, ax = plt.subplots(figsize=(5.0, 3.5))

    order = np.argsort(feature)
    ax.axhline(0.0, linestyle="--", color="grey")
    ax.scatter(feature, bias, s=14, alpha=0.6)
    ax.plot(feature[order], bias[order], color="tab:blue", alpha=0.4)
    ax.set_xlabel(xlabel)
    ax.set_ylabel("Local calibration bias")
    ax.set_title("Local calibration bias vs. " + xlabel)
    return ax

Example data

kernel_calibration.datasets.make_calibration_data

make_calibration_data(n=1000, miscalibration=0.25, region_threshold=0.0, slope=1.5, seed=0)

Generate synthetic data with a planted local miscalibration.

A single feature x drives the true probability p_true = sigmoid(slope*x). Labels are drawn as y ~ Bernoulli(p_true). The returned model probabilities f equal p_true (locally calibrated) except where x >= region_threshold, in which case the model is made over-confident by adding miscalibration. Setting miscalibration=0 yields a perfectly locally calibrated model — useful for a Type-I error check.

Parameters:

Name	Type	Description	Default
n	int	Number of samples. Default is 1000.	1000
miscalibration	float	Size of the additive bias applied in the affected region. 0 gives a locally calibrated model. Default is 0.25.	0.25
region_threshold	float	The model is miscalibrated where the feature x >= region_threshold. Default is 0.0.	0.0
slope	float	Steepness of the true probability as a function of x. Default is 1.5.	1.5
seed	int	RNG seed. Default is 0.	0

Returns:

Name	Type	Description
X	ndarray	Feature matrix of shape (n, 1).
y	ndarray	Binary labels of shape (n,).
f	ndarray	Model predicted probabilities of shape (n,).

Source code in kernel_calibration/datasets.py

def make_calibration_data(
    n: int = 1000,
    miscalibration: float = 0.25,
    region_threshold: float = 0.0,
    slope: float = 1.5,
    seed: int = 0,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Generate synthetic data with a planted *local* miscalibration.

    A single feature ``x`` drives the true probability ``p_true = sigmoid(slope*x)``.
    Labels are drawn as ``y ~ Bernoulli(p_true)``. The returned model probabilities
    ``f`` equal ``p_true`` (locally calibrated) **except** where ``x >=
    region_threshold``, in which case the model is made over-confident by adding
    ``miscalibration``. Setting ``miscalibration=0`` yields a perfectly locally
    calibrated model — useful for a Type-I error check.

    Parameters
    ----------
    n : int, optional
        Number of samples. Default is 1000.
    miscalibration : float, optional
        Size of the additive bias applied in the affected region. ``0`` gives a
        locally calibrated model. Default is 0.25.
    region_threshold : float, optional
        The model is miscalibrated where the feature ``x >= region_threshold``.
        Default is 0.0.
    slope : float, optional
        Steepness of the true probability as a function of ``x``. Default is 1.5.
    seed : int, optional
        RNG seed. Default is 0.

    Returns
    -------
    X : np.ndarray
        Feature matrix of shape (n, 1).
    y : np.ndarray
        Binary labels of shape (n,).
    f : np.ndarray
        Model predicted probabilities of shape (n,).
    """
    rng = np.random.default_rng(seed)
    x = rng.uniform(-2.0, 2.0, size=n)
    p_true = 1.0 / (1.0 + np.exp(-slope * x))
    y = (rng.uniform(size=n) < p_true).astype(float)

    f = p_true.copy()
    affected = x >= region_threshold
    f[affected] = np.clip(f[affected] + miscalibration, 1e-3, 1.0 - 1e-3)

    return x[:, None], y, f

kernel_calibration.datasets.fetch_compas

fetch_compas(cache_dir=None, filter_propublica=True)

Download the ProPublica COMPAS recidivism dataset (with local caching).

Requires pandas. On the first call the CSV is downloaded and cached; later calls read the cached copy. Set the KERNEL_CALIBRATION_DATA environment variable (or pass cache_dir) to control where it is stored.

Parameters:

Name	Type	Description	Default
cache_dir	str	Directory to cache the raw CSV. Defaults to ~/.cache/kernel_calibration.	None
filter_propublica	bool	Apply ProPublica's standard row filtering (screening-date window, valid recidivism flag, non-ordinary charge degree, non-missing score). Default True.	True

Returns:

Type	Description
DataFrame	The (optionally filtered) COMPAS records.

Source code in kernel_calibration/datasets.py

def fetch_compas(cache_dir: Optional[str] = None, filter_propublica: bool = True):
    """Download the ProPublica COMPAS recidivism dataset (with local caching).

    Requires ``pandas``. On the first call the CSV is downloaded and cached; later
    calls read the cached copy. Set the ``KERNEL_CALIBRATION_DATA`` environment
    variable (or pass ``cache_dir``) to control where it is stored.

    Parameters
    ----------
    cache_dir : str, optional
        Directory to cache the raw CSV. Defaults to ``~/.cache/kernel_calibration``.
    filter_propublica : bool, optional
        Apply ProPublica's standard row filtering (screening-date window, valid
        recidivism flag, non-ordinary charge degree, non-missing score). Default True.

    Returns
    -------
    pandas.DataFrame
        The (optionally filtered) COMPAS records.
    """
    try:
        import pandas as pd
    except ImportError as exc:  # pragma: no cover
        raise ImportError(
            "fetch_compas requires pandas. Install it with `pip install pandas`."
        ) from exc

    cache = Path(cache_dir) if cache_dir else _default_cache_dir()
    cache.mkdir(parents=True, exist_ok=True)
    path = cache / "compas-scores-two-years.csv"
    if not path.exists():
        pd.read_csv(_COMPAS_URL).to_csv(path, index=False)
    df = pd.read_csv(path)

    if filter_propublica:
        df = df[
            (df["days_b_screening_arrest"] <= 30)
            & (df["days_b_screening_arrest"] >= -30)
            & (df["is_recid"] != -1)
            & (df["c_charge_degree"] != "O")
            & (df["score_text"] != "N/A")
        ].reset_index(drop=True)
    return df