jaxsr.uncertainty

jaxsr.uncertainty#

Uncertainty Quantification for JAXSR.

Provides classical OLS intervals, Pareto front ensemble predictions, Bayesian Model Averaging, conformal prediction, and bootstrap methods.

Since JAXSR models are linear-in-parameters (y = Phi @ beta), classical OLS inference applies directly.

jaxsr.uncertainty.compute_unbiased_variance(Phi: Array, y: Array, coefficients: Array, sample_weight: Array | None = None) → float#

Compute unbiased noise variance estimate: s^2 = SSR / (n - p).

Parameters:

Phi (jnp.ndarray) – Design matrix of shape (n_samples, n_features).
y (jnp.ndarray) – Target values of shape (n_samples,).
coefficients (jnp.ndarray) – Fitted coefficients of shape (n_features,).
sample_weight (jnp.ndarray, optional) – Per-sample weights. The sum of squares becomes sum_i w_i r_i^2, so the estimate is the variance of a unit-weight observation under the model var(y_i) = sigma^2 / w_i.

Returns:

sigma_sq – Unbiased variance estimate.

Return type:

float

Raises:

ValueError – If sample_weight is invalid.

jaxsr.uncertainty.compute_coeff_covariance(Phi: Array, sigma_sq: float, sample_weight: Array | None = None) → Array#

Compute coefficient covariance matrix: Cov(beta) = sigma^2 * (Phi^T W Phi)^{-1}.

Uses SVD for numerical stability.

Parameters:

Phi (jnp.ndarray) – Design matrix of shape (n_samples, n_features).
sigma_sq (float) – Unbiased noise variance estimate.
sample_weight (jnp.ndarray, optional) – Per-sample weights forming the diagonal of W. None gives the ordinary sigma^2 (Phi^T Phi)^{-1}.

Returns:

cov – Covariance matrix of shape (n_features, n_features).

Return type:

jnp.ndarray

Raises:

ValueError – If sample_weight is invalid.

jaxsr.uncertainty.coefficient_intervals(Phi: Array, y: Array, coefficients: Array, names: list[str], alpha: float = 0.05, sample_weight: Array | None = None) → dict[str, tuple[float, float, float, float]]#

Compute t-based confidence intervals for each coefficient.

Parameters:

Phi (jnp.ndarray) – Design matrix of shape (n_samples, n_features).
y (jnp.ndarray) – Target values.
coefficients (jnp.ndarray) – Fitted coefficients.
names (list of str) – Names of basis functions.
alpha (float) – Significance level (default 0.05 for 95% CIs).
sample_weight (jnp.ndarray, optional) – Per-sample weights used for the fit. Standard errors then come from the weighted covariance s^2 (Phi^T W Phi)^{-1}; passing None for a weighted fit understates the uncertainty of coefficients that rest mostly on down-weighted rows.

Returns:

intervals – {name: (estimate, lower, upper, se)} for each coefficient.

Return type:

dict

Raises:

ValueError – If sample_weight is invalid.

jaxsr.uncertainty.prediction_interval(Phi_train: Array, y_train: Array, coefficients: Array, Phi_new: Array, alpha: float = 0.05, sample_weight: Array | None = None) → dict[str, Array]#

Compute prediction and confidence intervals for new observations.

Confidence band: uncertainty in the mean response E[y|x]. Prediction interval: uncertainty in a new observation y.

Parameters:

Phi_train (jnp.ndarray) – Training design matrix of shape (n_train, p).
y_train (jnp.ndarray) – Training target values of shape (n_train,).
coefficients (jnp.ndarray) – Fitted coefficients of shape (p,).
Phi_new (jnp.ndarray) – Design matrix for new points of shape (n_new, p).
alpha (float) – Significance level (default 0.05 for 95% intervals).
sample_weight (jnp.ndarray, optional) – Weights of the training rows, of shape (n_train,). Phi_new is never weighted: the interval it produces is for a new observation of unit weight, i.e. one as precise as an average training point.

Returns:

result – Dictionary with keys: - “y_pred”: predicted values - “pred_lower”, “pred_upper”: prediction interval bounds - “conf_lower”, “conf_upper”: confidence band bounds - “pred_se”: prediction standard error - “conf_se”: confidence standard error (mean response)

Return type:

dict

Raises:

ValueError – If sample_weight is invalid.

jaxsr.uncertainty.ensemble_predict(model: SymbolicRegressor, X_new: jnp.ndarray) → dict[str, jnp.ndarray]#

Predictions from all Pareto-front models.

Provides a measure of structural/model uncertainty: how much do predictions vary across plausible model complexities?

Parameters:

model (SymbolicRegressor) – Fitted model with selection path.
X_new (jnp.ndarray) – New input data of shape (n_new, n_features).

Returns:

result – Dictionary with keys: - “y_mean”: mean prediction across Pareto models - “y_std”: std of predictions across Pareto models - “y_min”: min prediction across Pareto models - “y_max”: max prediction across Pareto models - “y_all”: array of shape (n_models, n_new) with all predictions - “models”: list of SelectionResult for the Pareto models

Return type:

dict

class jaxsr.uncertainty.BayesianModelAverage(model: SymbolicRegressor, criterion: str = 'bic', use_pareto: bool = True, top_k: int | None = None)#

Bases: object

IC-weighted model averaging across selection path or Pareto front.

Weights are computed as:: w_k = exp(-0.5 * IC_k) / sum(exp(-0.5 * IC_j))

BMA variance includes within-model and between-model components.

Parameters:

model (SymbolicRegressor) – Fitted model with selection path.
criterion (str) – Information criterion to use for weights: “bic” or “aic”.
use_pareto (bool) – If True, average over Pareto front models only. If False, average over all models in the selection path.
top_k (int, optional) – If specified, only use the top k models by IC value.

property weights: dict[str, float]#: Model weights keyed by expression string.

property expressions: list[str]#: Expressions of models in the average.

predict(X: Array) → tuple[Array, Array]#

BMA prediction with uncertainty.

Parameters:

X (jnp.ndarray) – Input data of shape (n_samples, n_features).

Returns:

y_mean (jnp.ndarray) – Weighted mean prediction.
y_std (jnp.ndarray) – BMA standard deviation (within + between model variance).

predict_interval(X: Array, alpha: float = 0.05) → tuple[Array, Array, Array]#

BMA prediction interval using Gaussian approximation.

Parameters:

X (jnp.ndarray) – Input data.
alpha (float) – Significance level.

Returns:

y_pred (jnp.ndarray) – BMA mean prediction.
lower (jnp.ndarray) – Lower bound.
upper (jnp.ndarray) – Upper bound.

jaxsr.uncertainty.conformal_predict_split(model: SymbolicRegressor, X_cal: jnp.ndarray, y_cal: jnp.ndarray, X_new: jnp.ndarray, alpha: float = 0.05) → dict[str, jnp.ndarray]#

Split conformal prediction intervals.

Uses a held-out calibration set to construct distribution-free prediction intervals with finite-sample coverage guarantee.

Parameters:

model (SymbolicRegressor) – Fitted model (trained on separate training data).
X_cal (jnp.ndarray) – Calibration features of shape (n_cal, n_features).
y_cal (jnp.ndarray) – Calibration targets of shape (n_cal,).
X_new (jnp.ndarray) – New points for prediction.
alpha (float) – Significance level (coverage = 1 - alpha).

Returns:

result – Dictionary with keys: - “y_pred”: point predictions - “lower”: lower interval bound - “upper”: upper interval bound - “quantile”: the conformal quantile used

Return type:

dict

Notes

The calibration set is treated as unweighted, even when model was fitted with sample_weight: coverage here comes from exchangeability of the calibration residuals, not from the fit. If the calibration points have unequal precision, the resulting interval is a valid but uninformatively wide one for the precise points. For a weight-aware interval use method="jackknife+", which reuses the training weights.

jaxsr.uncertainty.conformal_predict_jackknife_plus(model: SymbolicRegressor, X_new: jnp.ndarray, alpha: float = 0.05) → dict[str, jnp.ndarray]#

Jackknife+ conformal prediction intervals.

Uses the LOO infrastructure to compute prediction intervals without a separate calibration set.

Parameters:

model (SymbolicRegressor) – Fitted model (must have stored training data).
X_new (jnp.ndarray) – New points for prediction.
alpha (float) – Significance level (coverage = 1 - alpha).

Returns:

result – Dictionary with keys: - “y_pred”: point predictions - “lower”: lower interval bound - “upper”: upper interval bound

Return type:

dict

Raises:

ValueError – If the model has no stored training data.

Notes

For a weighted fit the leverages are the weighted-least-squares ones and the nonconformity scores are the whitened LOO residuals sqrt(w_i) |r_i| / (1 - h_ii), which are exchangeable under the model var(y_i) = sigma^2 / w_i. The resulting interval is therefore for a new observation of unit weight; scale it by 1 / sqrt(w_new) if the new point has a known precision of its own.

jaxsr.uncertainty.bootstrap_coefficients(model: SymbolicRegressor, n_bootstrap: int = 1000, alpha: float = 0.05, seed: int | None = None) → dict[str, Any]#

Residual bootstrap for coefficient uncertainty.

Resamples residuals, creates y* = y_hat + e*, and refits OLS. No Gaussian assumption needed.

Parameters:

model (SymbolicRegressor) – Fitted model.
n_bootstrap (int) – Number of bootstrap samples.
alpha (float) – Significance level for confidence intervals.
seed (int, optional) – Random seed.

Returns:

result – Dictionary with keys: - “coefficients”: (n_bootstrap, p) array of bootstrap coefficients - “mean”: mean of bootstrap coefficients - “std”: std of bootstrap coefficients - “lower”: lower CI bound for each coefficient - “upper”: upper CI bound for each coefficient - “names”: coefficient names

Return type:

dict

Raises:

ValueError – If the model has no stored training data.

Notes

Weights from fit() are honoured automatically. Raw residuals of a weighted fit are not exchangeable – a down-weighted row is expected to miss by more – so resampling them directly would leak that row’s noise onto precise rows. The resampling therefore happens in whitened space, on sqrt(w_i) * r_i, and each bootstrap replicate is refit by weighted least squares.

jaxsr.uncertainty.bootstrap_predict(model: SymbolicRegressor, X_new: jnp.ndarray, n_bootstrap: int = 1000, alpha: float = 0.05, seed: int | None = None) → dict[str, jnp.ndarray]#

Bootstrap prediction intervals.

Parameters:

model (SymbolicRegressor) – Fitted model.
X_new (jnp.ndarray) – New input data.
n_bootstrap (int) – Number of bootstrap samples.
alpha (float) – Significance level.
seed (int, optional) – Random seed.

Returns:

result – Dictionary with keys: - “y_pred”: point prediction (from original model) - “y_mean”: mean of bootstrap predictions - “y_std”: std of bootstrap predictions - “lower”: lower prediction bound - “upper”: upper prediction bound

Return type:

dict

jaxsr.uncertainty.summarize_selection_replicates(replicates: Sequence[Any], reference: Any | None = None, resampling: str = 'custom') → dict[str, Any]#

Summarise selection stability over replicates the caller produced.

bootstrap_model_selection uses this internally, but it is public so that replicates generated outside JAXSR can be reported the same way. That matters when the rows fed to fit are themselves outputs of an upstream step – a smoother, a derivative estimate, a simulation – because resampling those rows perturbs nothing about the step that produced them. The honest replicate there re-runs the whole pipeline; see bootstrap_model_selection(..., resample_fn=...) or simply collect the fitted models yourself and pass them here.

Parameters:

replicates (sequence) – One entry per replicate. Each may be a fitted SymbolicRegressor, a mapping with a "features" key (optionally "expression"), or a sequence of selected feature names.
reference (SymbolicRegressor or sequence of str, optional) – The feature set to measure stability against, typically the model fit to the full data. If omitted, stability_score is the frequency of the most common structure instead.
resampling (str) – Label recorded in the result describing what was resampled, e.g. "rows", "groups", "pipeline".

Returns:

result – Dictionary with keys:

”feature_frequencies”: dict mapping feature name to the fraction of replicates that selected it, highest first. Model replicates are keyed by basis identity, so a parametric basis appears once (as the registered "exp(-a*x)") rather than once per fitted value
”parameter_distributions”: dict mapping basis identity to {param_name: {"mean", "sd", "q05", "q95", "n"}} for parametric basis functions; empty unless the replicates are fitted models with parametric bases
”stability_score”: fraction of replicates whose selected feature set equals reference (or the modal structure’s frequency)
”expressions”: list of expressions, one per replicate, with parametric values rendered as fitted in that replicate
”structures”: list of dicts with “features”, “count”, “frequency”, one per distinct selected feature set, most frequent first
”n_replicates”, “n_distinct_structures”, “n_distinct_expressions”
”reference_features”: the reference feature set, or None
”resampling”: the resampling label

Return type:

dict

Raises:

TypeError – If a replicate is not one of the supported forms.

Examples

>>> models = [fit_one_replicate(i) for i in range(90)]  
>>> summary = summarize_selection_replicates(  
...     models, reference=full_fit, resampling="pipeline"
... )
>>> summary["n_distinct_structures"]  
9

jaxsr.uncertainty.bootstrap_model_selection(model: SymbolicRegressor, X: jnp.ndarray | None, y: jnp.ndarray | None, n_bootstrap: int = 100, seed: int | None = None, groups: Any | None = None, resample_fn: Callable[[np.random.RandomState], tuple[Any, Any]] | None = None, sample_weight: jnp.ndarray | None = None) → dict[str, Any]#

Bootstrap model selection stability at the row, group, or pipeline level.

By default this is a pairs bootstrap: it resamples (X_i, y_i) rows, reruns model.fit(), and tracks which features are selected. That is only the right unit of resampling when rows are independent observations.

Rows are grouped – several rows share one experimental condition, one measured curve, one subject – pass groups. Whole groups are then resampled with replacement, so a replicate never sees part of a group it also trained on. Row resampling here reports a spread far narrower than the real between-group variability.
Rows come out of an upstream fit – spline evaluations, estimated derivatives, simulation output – pass resample_fn. Resampling such rows perturbs nothing about the step that produced them, so the dominant error source is invisible to a row-wise bootstrap. resample_fn gets to regenerate the data (and re-run that upstream step) per replicate.

Parameters:

model (SymbolicRegressor) – Model used as the template for every replicate. If it is fitted, its selected features become the reference for stability_score.
X (jnp.ndarray or None) – Training features. May be None when resample_fn is given.
y (jnp.ndarray or None) – Training targets. May be None when resample_fn is given.
n_bootstrap (int) – Number of bootstrap replicates.
seed (int, optional) – Random seed.
groups (array-like of shape (n_samples,), optional) – Group label per row. When given, groups rather than rows are resampled with replacement. Cannot be combined with resample_fn.
resample_fn (callable, optional) – resample_fn(rng) -> (X_b, y_b), called once per replicate with the bootstrap’s numpy.random.RandomState so the replicates stay reproducible under seed. Use it to regenerate the data and re-run whatever upstream stage produced it. Cannot be combined with groups.
sample_weight (array-like of shape (n_samples,), optional) – Per-sample weights. Each weight follows its row into the resample – under groups too, since a group is drawn as a whole block of rows – so a replicate that happens to draw mostly down-weighted rows is fitted as the thin evidence it is. Defaults to the weights model was fitted with when X/y are that same training set. Cannot be combined with resample_fn: those rows are regenerated, so there is no row for a stored weight to belong to.

Returns:

result – The dictionary returned by summarize_selection_replicates, plus:

”n_failed”: number of replicates whose fit raised and was skipped
”resampling”: "rows", "groups", or "pipeline"

Return type:

dict

Raises:

ValueError – If n_bootstrap < 1, if both groups and resample_fn are given, if groups has the wrong length, if fewer than 2 groups are present, if X/y are missing without a resample_fn, or if sample_weight is invalid or combined with resample_fn.
TypeError – If resample_fn does not return a (X_b, y_b) pair.

Examples

>>> bootstrap_model_selection(model, X, y, groups=temperature_id)  
>>> def replicate(rng):  
...     data = simulate(seed=rng.randint(2**31))
...     return features_from_smoother(data), derivative_from_smoother(data)
>>> bootstrap_model_selection(model, None, None, resample_fn=replicate)  

class jaxsr.uncertainty.AnovaRow(source: str, df: int, sum_sq: float, mean_sq: float, f_value: float | None = None, p_value: float | None = None)#

Bases: object

A single row of the ANOVA table.

Parameters:

source (str) – Name of the source of variation (term name, “Model”, “Residual”, or “Total”).
df (int) – Degrees of freedom.
sum_sq (float) – Sum of squares.
mean_sq (float) – Mean sum of squares (sum_sq / df).
f_value (float or None) – F-statistic (None for Residual and Total rows).
p_value (float or None) – p-value from the F-distribution (None for Residual and Total rows).

source: str#

df: int#

sum_sq: float#

mean_sq: float#

f_value: float | None = None#

p_value: float | None = None#

class jaxsr.uncertainty.AnovaResult(rows: list[~jaxsr.uncertainty.AnovaRow] = <factory>, type: str = 'sequential', warnings: list[str] = <factory>)#

Bases: object

Result of an ANOVA decomposition.

Parameters:

rows (list of AnovaRow) – Per-term rows followed by Model, Residual, and Total summary rows.
type (str) – ANOVA type: "sequential" (Type I) or "marginal" (Type III).
warnings (list of str) – Diagnostic messages (e.g. when p-values are approximate).

rows: list[AnovaRow]#

type: str = 'sequential'#

warnings: list[str]#

property term_names: list[str]#: Names of the individual term rows (excludes summary rows).

to_dict() → dict[str, Any]#: Serialise the table to a plain dictionary.

jaxsr.uncertainty.anova(model: SymbolicRegressor, anova_type: str = 'sequential') → AnovaResult#

Perform ANOVA on a fitted SymbolicRegressor.

Decomposes the total sum of squares into contributions from each selected basis-function term plus a residual, and tests each term’s significance with an F-test.

Two decomposition types are supported:

``”sequential”`` (Type I): terms are added in the order they appear in model.selected_features_ and each term’s contribution is the extra sum of squares beyond the previous terms.
``”marginal”`` (Type III): each term’s contribution is the extra sum of squares from adding it last, after all other terms.

Parameters:

model (SymbolicRegressor) – A fitted model.
anova_type (str) – "sequential" (Type I) or "marginal" (Type III).

Returns:

ANOVA table with per-term rows plus Model, Residual, and Total summary rows.

Return type:

AnovaResult

Notes

The F-test p-values assume independent, normally distributed residuals and an unconstrained OLS fit. They are approximate when:

The model contains parametric (nonlinear) basis functions whose parameters were optimised during fitting, because the effective degrees of freedom are larger than reported.
Constraints were applied, since the coefficient estimates are no longer ordinary least-squares.

In these cases the sums-of-squares decomposition is still meaningful, but the p-values should be treated as indicative rather than exact. The returned :pyclass:`AnovaResult` includes diagnostic warnings when these conditions are detected.

When the model was fitted with sample_weight, every sum of squares is the weighted one (sum_i w_i r_i^2) and the total is taken about the weighted mean of y, so the table decomposes the same quantity the fit minimised. Degrees of freedom stay nominal, matching the information criteria.

Examples

>>> from jaxsr.uncertainty import anova
>>> table = anova(model)
>>> print(table)
>>> for row in table.rows:
...     print(row.source, row.f_value, row.p_value)

jaxsr.uncertainty.classification_coefficient_intervals(Phi: Array, y: Array, coefficients: Array, names: list[str], alpha: float = 0.05) → dict[str, tuple[float, float, float, float]]#

Wald confidence intervals for logistic regression coefficients.

Uses the Fisher information matrix I(w) = Phi^T diag(mu*(1-mu)) Phi to compute asymptotic standard errors, then applies Normal quantiles (MLE asymptotics).

Parameters:

Phi (jnp.ndarray) – Design matrix of shape (n, p).
y (jnp.ndarray) – Binary labels of shape (n,).
coefficients (jnp.ndarray) – Fitted logistic regression coefficients.
names (list of str) – Coefficient names.
alpha (float) – Significance level (default 0.05 for 95% CIs).

Returns:

intervals – {name: (estimate, lower, upper, se)} for each coefficient.

Return type:

dict

jaxsr.uncertainty.bootstrap_classification_coefficients(model, n_bootstrap: int = 1000, alpha: float = 0.05, seed: int | None = None) → dict[str, Any]#

Pairs bootstrap for logistic regression coefficient uncertainty.

Resamples (X_i, y_i) pairs, refits IRLS on each bootstrap sample, and collects coefficient distributions.

Parameters:

model (SymbolicClassifier) – Fitted binary classifier.
n_bootstrap (int) – Number of bootstrap resamples.
alpha (float) – Significance level for confidence intervals.
seed (int, optional) – Random seed.

Returns:

result – Dictionary with keys:

"coefficients": (n_bootstrap, p) array
"mean": mean of bootstrap coefficients
"std": std of bootstrap coefficients
"lower": lower CI bound per coefficient
"upper": upper CI bound per coefficient
"names": coefficient names

Return type:

dict

jaxsr.uncertainty.conformal_classification_split(model, X_cal: Array, y_cal: Array, X_new: Array, alpha: float = 0.05) → dict[str, Any]#

Split conformal prediction sets for classification.

Returns prediction sets (sets of labels) with marginal coverage guarantee. Uses the nonconformity score 1 - P(y_true | x).

Parameters:

model (SymbolicClassifier) – Fitted classification model.
X_cal (jnp.ndarray) – Calibration features of shape (n_cal, p).
y_cal (jnp.ndarray) – Calibration labels of shape (n_cal,).
X_new (jnp.ndarray) – New features of shape (n_new, p).
alpha (float) – Significance level (coverage = 1 - alpha).

Returns:

result – Dictionary with keys:

"prediction_sets": list of sets of predicted labels
"quantile": the conformal quantile used
"y_pred": point predictions (most likely class)

Return type:

dict

jaxsr.uncertainty.calibration_curve(y_true: Array, y_prob: Array, n_bins: int = 10) → dict[str, ndarray]#

Compute reliability diagram data for binary classification.

Bins predicted probabilities and computes the observed fraction of positives (fraction_of_positives) and the average predicted probability (mean_predicted_value) in each bin.

Parameters:

y_true (jnp.ndarray) – True binary labels.
y_prob (jnp.ndarray) – Predicted probabilities for the positive class.
n_bins (int) – Number of bins.

Returns:

result – Dictionary with keys:

"fraction_of_positives": observed positive rate per bin
"mean_predicted_value": mean predicted probability per bin
"bin_counts": number of samples per bin

Return type:

dict

jaxsr.uncertainty

Contents

jaxsr.uncertainty#