pycsamt.ai.inversion.ensemble

EnsembleInverter — deep ensemble for uncertainty-aware EM inversion.

Trains \(M\) independent copies of a base EMInverter1D estimator with different random seeds. Epistemic (model) uncertainty is estimated from the inter-model variance. Aleatoric uncertainty is not captured here; use MC-Dropout variants for that.

Uncertainty estimation

For a test sample \(\mathbf{x}\), each ensemble member \(f_m\) predicts \(\hat{\mathbf{y}}_m\). The ensemble mean and variance are:

\[ \begin{align}\begin{aligned}\bar{\mathbf{y}} &= \frac{1}{M} \sum_{m=1}^{M} \hat{\mathbf{y}}_m\\\sigma^2 &= \frac{1}{M-1} \sum_{m=1}^{M} (\hat{\mathbf{y}}_m - \bar{\mathbf{y}})^2\end{aligned}\end{align} \]

The ensemble RMSE (expected calibration) is typically lower than any single member; \(\sigma\) provides a calibrated uncertainty estimate for downstream probabilistic interpretation.

References

Lakshminarayanan B. et al. (2017) NeurIPS — Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles.

Classes

EnsembleInverter(base_estimator[, ...])

Deep ensemble of EMInverter1D models.

class pycsamt.ai.inversion.ensemble.EnsembleInverter(base_estimator, n_estimators=5, seeds=None)

Bases: object

Deep ensemble of EMInverter1D models.

Parameters:

• base_estimator (EMInverter1D) – Template estimator (unfitted). Its hyper-parameters are copied for every ensemble member.

• n_estimators (int, default 5) – Number of ensemble members.

• seeds (list of int or None) – Per-member random seeds. When None, seeds [0, 1, …, n_estimators-1] are used.

Examples

>>> from pycsamt.ai.inversion import EMInverter1D, EnsembleInverter
>>> base = EMInverter1D(arch="resnet", n_layers=5)
>>> ens = EnsembleInverter(base, n_estimators=5)
>>> ens.fit(ds, epochs=50)
EnsembleInverter(n_estimators=5, fitted)
>>> mean, std = ens.predict_with_uncertainty(X_test)

fit(X, y=None, *, epochs=100, batch_size=256, lr=0.001, patience=20, val_frac=0.1, grad_clip=1.0, verbose=True)

Train all ensemble members.

Parameters:

• X (ForwardDataset, path, or ndarray) – Training data (forwarded unchanged to each member’s fit).

• y (ndarray or None)

• epochs (int) – Forwarded to each member’s fit call; each member receives a unique seed.

• batch_size (int) – Forwarded to each member’s fit call; each member receives a unique seed.

• lr (float) – Forwarded to each member’s fit call; each member receives a unique seed.

• patience (int) – Forwarded to each member’s fit call; each member receives a unique seed.

• val_frac (float) – Forwarded to each member’s fit call; each member receives a unique seed.

• grad_clip (float | None) – Forwarded to each member’s fit call; each member receives a unique seed.

• verbose (bool) – Forwarded to each member’s fit call; each member receives a unique seed.

Return type:

self

predict(X, **kwargs)

Return the mean ensemble prediction.

Parameters:

• X (ndarray (n_samples, n_features))

• **kwargs (forwarded to each member’s predict)

Returns:

y_mean

Return type:

ndarray (n_samples, n_params)

predict_with_uncertainty(X, _use_calibrated=True, **kwargs)

Return (mean, std) ensemble prediction.

If calibrate() has been called, the returned std is the calibrated standard deviation from the PosteriorCalibrator (i.e. the Gaussianization-flow recalibrated sigma). Pass _use_calibrated=False to retrieve the raw inter-member std.

Parameters:

• X (ndarray (n_samples, n_features))

• _use_calibrated (bool) – When True (default) and a calibrator is attached, return calibrated sigma. Internal flag — not part of the public API.

Returns:

• mean (ndarray (n_samples, n_params))

• std (ndarray (n_samples, n_params)) – Calibrated or raw inter-member standard deviation.

Return type:

tuple[ndarray, ndarray]

predict_quantiles(X, q=(0.05, 0.25, 0.5, 0.75, 0.95), **kwargs)

Return per-quantile predictions.

Parameters:

• X (ndarray (n_samples, n_features))

• q (sequence of float) – Quantiles in (0, 1).

Returns:

quantiles

Return type:

dict {q_value: ndarray (n_samples, n_params)}

calibrate(X_cal, y_cal, *, alpha=0.1)

Attach a ConformalPredictor and a PosteriorCalibrator fitted on a held-out calibration set.

After calling this method:

• predict_intervals() returns conformal prediction bands with guaranteed \(\ge 1-\alpha\) marginal coverage (Vovk et al. 2005).

• predict_posterior() returns calibrated posterior samples via the Gaussianization normalising flow (Kuleshov et al. 2018).

• predict_with_uncertainty() returns the calibrated standard deviation instead of the raw inter-member std.

Parameters:

• X_cal (ndarray, shape (n_cal, n_features)) – Calibration inputs (must not overlap with the training set).

• y_cal (ndarray, shape (n_cal, n_params)) – True parameter vectors for the calibration inputs.

• alpha (float) – Default significance level for conformal intervals.

Return type:

self

predict_intervals(X, alpha=None)

Conformal prediction intervals with guaranteed marginal coverage.

Requires a prior call to calibrate().

Parameters:

• X (ndarray, shape (n_samples, n_features))

• alpha (float or None) – Significance level. Defaults to the value used in calibrate().

Returns:

• center (ndarray, shape (n_samples, n_params))

• lower (ndarray, shape (n_samples, n_params))

• upper (ndarray, shape (n_samples, n_params))

Return type:

tuple[ndarray, ndarray, ndarray]

predict_posterior(X, n_samples=500, rng=None)

Draw calibrated posterior samples via the Gaussianization normalising flow (Kuleshov et al. 2018).

Requires a prior call to calibrate().

Parameters:

• X (ndarray, shape (n_pts, n_features))

• n_samples (int) – Number of posterior draws per input point.

• rng (numpy.random.Generator or None)

Returns:

draws

Return type:

ndarray, shape (n_samples, n_pts, n_params)

coverage_diagnostics(X_test, y_test, alphas=None)

Per-alpha empirical coverage table (requires prior calibrate()).

Useful for a reliability diagram: nominal coverage 1-alpha on the x-axis, actual coverage on the y-axis. A perfectly calibrated predictor lies on the diagonal.

Parameters:

• X_test (ndarray)

• y_test (ndarray)

• alphas (sequence of float or None)

Returns:

diag

Return type:

dict {alpha: actual_coverage}

score(X, y, *, metric='rmse')

Score the ensemble mean prediction.

Parameters:

• X (ndarray (n_samples, n_features))

• y (ndarray (n_samples, n_params))

• metric ({'rmse', 'mae', 'r2', 'relative_rmse'})

Returns:

score

Return type:

float

coverage(X, y, *, n_sigma=1.96)

Fraction of true values within ±n_sigma * std of the ensemble mean.

An uncertainty estimate is well-calibrated when coverage(1.96) ≈ 0.95.

Parameters:

• X (ndarray)

• y (ndarray)

• n_sigma (float)

Returns:

coverage

Return type:

float ∈ [0, 1]

save(path)

Save all ensemble members.

Creates a directory <path>/ containing member_00.npz, member_01.npz, ….

Parameters:

path (str or Path)

Return type:

None

classmethod load(path, base_class=None)

Load a saved ensemble.

Parameters:

• path (str or Path) – Directory created by save().

• base_class (type or None) – Estimator class to use for loading. Defaults to EMInverter1D.

Return type:

EnsembleInverter

plot_uncertainty_profile(X, sample_idx=0, *, y_true=None, n_sigma=1.96, **plot_kwargs)

Plot a 1-D resistivity profile with uncertainty bands.

Parameters:

• X (ndarray (n_samples, n_features))

• sample_idx (int) – Which sample to plot.

• y_true (ndarray (n_params,) or None) – True model vector for overlay.

• n_sigma (float) – Band width in standard deviations.

• **plot_kwargs – Forwarded to plot_uncertainty_bands().

Returns:

fig

Return type:

Figure