neuralop.models.UQNO

class neuralop.models.UQNO(*args, **kwargs)[source]

Uncertainty Quantification Neural Operator

General N-dim (alpha, delta) Risk-Controlling Neural Operator, as described in [1].

The UQNO is trained to map input functions to a residual function E(a, x) that describes the predicted error between the ground truth and the outputs of a trained model. E(a, x) is then used in combination with a calibrated scaling factor to predict calibrated uncertainty bands around the predictions of the trained model.

Parameters:

base_modelnn.Module

Pre-trained solution operator. Determined by the problem.

residual_modelnn.Module, optional

Architecture to train as the UQNO’s quantile model. If None, a deep copy of base_model is used. Default: None

Methods

forward(*args, **kwargs)

Forward pass returns the solution u(a,x) and the uncertainty ball E(a,x) as a pair for pointwise quantile loss

References

[1]

:

Ma, Z., Pitt, D., Azizzadenesheli, K., and Anandkumar, A. (2024).

“Calibrated Uncertainty Quantification for Operator Learning via Conformal Prediction”. TMLR, https://openreview.net/pdf?id=cGpegxy12T.

forward(*args, **kwargs)[source]

Forward pass returns the solution u(a,x) and the uncertainty ball E(a,x) as a pair for pointwise quantile loss