Uncertainty Module#

Drop-in uncertainty-quantification estimators for fitted MaldiDeepKit classifiers. Three methods share a single predict_with_uncertainty() interface and return a common UncertaintyResult dataclass, so downstream code can swap methods without touching call sites.

  • MCDropoutEstimator - Monte Carlo Dropout (Gal and Ghahramani, 2016) with epistemic / aleatoric decomposition.

  • LaplaceEstimator - last-layer or full-network Laplace approximation via the optional laplace-torch dependency.

  • ConformalPredictor - split conformal prediction with the LAC non-conformity score.

The Laplace estimator requires the uncertainty extra; install it with pip install "maldideepkit[uncertainty]" (see Installation). Monte Carlo Dropout and split conformal prediction need no extras.

UncertaintyResult#

class maldideepkit.uncertainty.UncertaintyResult(predictions, proba_mean, uncertainty, epistemic, aleatoric, method, metadata=<factory>)[source]#

Bases: object

Container for the output of a BaseUncertaintyEstimator.

Variables:

  • predictions (ndarray of shape (n_samples,)) – Hard labels drawn from the classifier’s classes_.

  • proba_mean (ndarray of shape (n_samples, n_classes)) – Mean softmax probabilities. For non-Bayesian methods this is simply the classifier’s predict_proba() output.

  • uncertainty (ndarray of shape (n_samples,)) – Scalar per-sample uncertainty in [0, 1] where possible. The precise semantics are method-specific and documented on the producing class.

  • epistemic (ndarray of shape (n_samples,) or None) – Model-uncertainty component, when the method decomposes total uncertainty into epistemic and aleatoric parts. None for methods that do not decompose.

  • aleatoric (ndarray of shape (n_samples,) or None) – Data-uncertainty component, paired with epistemic. None for methods that do not decompose.

  • method (str) – Name of the method that produced this result (e.g. "mc_dropout", "laplace", "conformal").

  • metadata (dict) – Method-specific extras (per-pass samples for MC Dropout, prediction sets for conformal, predictive variance for Laplace, …).

Parameters:
predictions: ndarray#
proba_mean: ndarray#
uncertainty: ndarray#
epistemic: ndarray | None#
aleatoric: ndarray | None#
method: str#
metadata: dict[str, Any]#
__init__(predictions, proba_mean, uncertainty, epistemic, aleatoric, method, metadata=<factory>)#
Parameters:
Return type:

None

BaseUncertaintyEstimator#

class maldideepkit.uncertainty.BaseUncertaintyEstimator(classifier)[source]#

Bases: object

Abstract base for every estimator in maldideepkit.uncertainty.

Parameters:

classifier (BaseSpectralClassifier) – A fitted MaldiDeepKit classifier. Its preprocessing pipeline (warping, input transform, legacy standardisation) is reused so the uncertainty estimate is produced on inputs identical to those seen at training time.

Raises:

sklearn.exceptions.NotFittedError – If classifier has not been fitted.

__init__(classifier)[source]#
Parameters:

classifier (BaseSpectralClassifier)

Return type:

None

abstractmethod predict_with_uncertainty(X)[source]#

Return predictions and per-sample uncertainty for X.

Parameters:

X (Any) – Spectra to score. Must match classifier.input_dim_.

Returns:

Predictions, probability mean, scalar uncertainty, optional epistemic / aleatoric decomposition, and method-specific metadata.

Return type:

UncertaintyResult

MCDropoutEstimator#

class maldideepkit.uncertainty.MCDropoutEstimator(classifier, n_samples=30, batch_size=None)[source]#

Bases: BaseUncertaintyEstimator

Monte Carlo Dropout estimator (Gal and Ghahramani, 2016).

Runs n_samples stochastic forward passes through classifier.model_ with dropout layers active and aggregates the resulting softmax distribution into a predictive mean plus an epistemic / aleatoric decomposition.

Parameters:

  • classifier (BaseSpectralClassifier) – A fitted classifier whose architecture contains at least one torch.nn.Dropout (or its 1-D / 2-D / 3-D variants). A warning is emitted if it does not, but the estimator still runs - in that degenerate case all n_samples passes are identical and epistemic uncertainty collapses to 0.

  • n_samples (int) – Number of stochastic forward passes per call to predict_with_uncertainty(). Higher values reduce variance of the estimate at a linear cost in compute.

  • batch_size (int | None) – Mini-batch size used to chunk the input through each forward pass. When None, falls back to classifier.batch_size.

Notes

The scalar UncertaintyResult.uncertainty field is the normalised entropy of the predictive mean H(E[p]). The epistemic / aleatoric decomposition follows Depeweg et al. (2018):

  • aleatoric = E_w[H(p)] (mean of per-pass entropies)

  • epistemic = H(E[p]) - E[H(p)] (mutual information)

Both components are normalised to [0, 1] by dividing the raw entropies by log(n_classes).

__init__(classifier, n_samples=30, batch_size=None)[source]#
Parameters:
Return type:

None

predict_with_uncertainty(X, *, store_samples=False)[source]#

Run n_samples MC forward passes and decompose uncertainty.

Parameters:

  • X (Any) – Spectra to score.

  • store_samples (bool) – If True, attach the raw per-pass softmax tensor of shape (n_mc_passes, n_data_samples, n_classes) to metadata["per_pass_proba"]. Off by default to keep memory bounded on large datasets.

Returns:

method="mc_dropout" with both epistemic and aleatoric components populated.

Return type:

UncertaintyResult

LaplaceEstimator#

class maldideepkit.uncertainty.LaplaceEstimator(classifier, subset_of_weights='last_layer', hessian_structure='diag', sigma_noise=1.0)[source]#

Bases: BaseUncertaintyEstimator

Laplace-approximation uncertainty estimator.

A thin wrapper around the laplace-torch package (aleximmer/Laplace) that fits a Gaussian posterior over the classifier’s weights and turns the predictive variance into a per-sample uncertainty estimate.

Parameters:

  • classifier (BaseSpectralClassifier) – A fitted classifier; its model_ is reused as the network whose posterior is approximated.

  • subset_of_weights (str) – Which subset of weights to model. "last_layer" is the standard, cheap default and works for any of the MaldiDeepKit backbones whose final layer is a torch.nn.Linear.

  • hessian_structure (str) – Approximation structure for the Hessian. "diag" is the cheapest; "kron" (Kronecker-factored) is more accurate at a moderate compute cost.

  • sigma_noise (float) – Forwarded to laplace-torch. For classification it has no effect (it controls the regression noise scale) but is exposed for interface symmetry.

Raises:

ImportError – If laplace-torch is not installed.

__init__(classifier, subset_of_weights='last_layer', hessian_structure='diag', sigma_noise=1.0)[source]#
Parameters:
Return type:

None

calibrate(X_cal, y_cal, *, batch_size=None)[source]#

Fit the Laplace approximation on (X_cal, y_cal).

Applies the classifier’s preprocessing to X_cal, builds an internal DataLoader, fits the Laplace approximation, and runs marginal-likelihood prior precision optimisation.

Parameters:

  • X_cal (Any) – Calibration spectra.

  • y_cal (Any) – Calibration labels using the original label space stored in classifier.classes_. Re-encoded to 0..n_classes-1 internally.

  • batch_size (int | None) – DataLoader batch size. When None, falls back to classifier.batch_size.

Returns:

self, with the fitted Laplace approximation stored on la_.

Return type:

LaplaceEstimator

predict_with_uncertainty(X)[source]#

Return predictions and Laplace-derived uncertainty for X.

Uses pred_type="glm" and link_approx="probit" to map the weight-space posterior into a softmax-domain predictive distribution. Per-sample predictive variance is summarised as the mean diagonal entry, then squashed into [0, 1] via 1 - exp(-v) and stored as epistemic. The scalar UncertaintyResult.uncertainty field is the normalised entropy of the predictive mean; aleatoric is the non-negative residual uncertainty - epistemic.

Raises:

RuntimeError – If calibrate() has not been called.

Parameters:

X (Any)

Return type:

UncertaintyResult

ConformalPredictor#

class maldideepkit.uncertainty.ConformalPredictor(classifier, alpha=0.1, score='lac')[source]#

Bases: BaseUncertaintyEstimator

Split conformal predictor with the LAC non-conformity score.

Parameters:

  • classifier (BaseSpectralClassifier) – A fitted classifier whose predict_proba()-style softmax scores act as the underlying probability estimate.

  • alpha (float) – Miscoverage level in (0, 1). The target marginal coverage is 1 - alpha.

  • score (str) – Non-conformity score. Currently only "lac" is supported: s(x, y) = 1 - p_hat(y | x).

Notes

The UncertaintyResult.uncertainty field is the prediction set size normalised by n_classes: 1 / n_classes for a singleton set, 1.0 when every class is included. Empty sets (which can occur with very small calibration sets) yield 0 and are flagged in metadata via the empty-set count.

__init__(classifier, alpha=0.1, score='lac')[source]#
Parameters:
Return type:

None

calibrate(X_cal, y_cal)[source]#

Compute the conformal quantile from calibration data.

Parameters:

  • X_cal (Any) – Calibration spectra.

  • y_cal (Any) – Calibration labels using the original label space stored in classifier.classes_.

Returns:

self, with quantile_, calibration_coverage_, and n_calibration_ populated.

Return type:

ConformalPredictor

predict_with_uncertainty(X)[source]#

Return conformal predictions and prediction sets for X.

Returns:

method="conformal" with epistemic and aleatoric set to None. Boolean prediction sets are stored in metadata["prediction_sets"] with shape (n_samples, n_classes); the empirical calibration coverage is stored in metadata["calibration_coverage"].

Return type:

UncertaintyResult

Raises:

RuntimeError – If calibrate() has not been called.

Parameters:

X (Any)

Examples#

Monte Carlo Dropout on a fitted attention-MLP:

import numpy as np
from maldideepkit import MaldiMLPClassifier
from maldideepkit.uncertainty import MCDropoutEstimator

rng = np.random.default_rng(0)
X = rng.standard_normal((200, 6000)).astype("float32")
y = rng.integers(0, 2, size=200)

clf = MaldiMLPClassifier(random_state=0).fit(X, y)
est = MCDropoutEstimator(clf, n_samples=30)
result = est.predict_with_uncertainty(X[:10])
# result.predictions, result.proba_mean, result.uncertainty,
# result.epistemic, result.aleatoric

Split conformal prediction with calibration:

from sklearn.model_selection import train_test_split
from maldideepkit.uncertainty import ConformalPredictor

X_tr, X_cal, y_tr, y_cal = train_test_split(
    X, y, test_size=0.3, stratify=y, random_state=0,
)
clf = MaldiMLPClassifier(random_state=0).fit(X_tr, y_tr)

cp = ConformalPredictor(clf, alpha=0.1).calibrate(X_cal, y_cal)
result = cp.predict_with_uncertainty(X[:10])
sets = result.metadata["prediction_sets"]   # (10, n_classes) bool

Laplace approximation (requires the optional laplace-torch dependency):

from maldideepkit.uncertainty import LaplaceEstimator

la = LaplaceEstimator(clf, subset_of_weights="last_layer",
                      hessian_structure="diag")
la.calibrate(X_cal, y_cal)
result = la.predict_with_uncertainty(X[:10])