gensbi.diagnostics.sbc#

Functions#

_plot_cdf_region_expected_under_uniformity(...[, ...])

Plot region of empirical cdfs expected under uniformity on the current axis.

_plot_hist_region_expected_under_uniformity(...[, ...])

Plot region of empirical cdfs expected under uniformity.

_plot_ranks_as_cdf(ranks, num_bins, num_repeats[, ...])

Plot ranks as empirical CDFs on the current axis.

_plot_ranks_as_hist(ranks, num_bins, num_posterior_samples)

Plot ranks as histograms on the current axis.

_prepare_reduce_functions(reduce_fns, param_dim)

Prepare reduction functions for SBC analysis.

_run_sbc(thetas, xs, posterior_samples[, reduce_fns, ...])

Calculate ranks for SBC or expected coverage.

_sbc_rank_plot(ranks, num_posterior_samples[, ...])

Plot simulation-based calibration ranks as empirical CDFs or histograms.

_validate_sbc_inputs(thetas, xs, num_sbc_samples, ...)

Validate inputs for the SBC procedure.

check_prior_vs_dap(prior_samples, dap_samples)

Returns the C2ST accuracy between prior and data-averaged posterior samples.

check_sbc(ranks, prior_samples, dap_samples[, ...])

Return uniformity checks and data-averaged posterior checks for SBC.

check_uniformity_c2st(ranks, num_posterior_samples[, ...])

Return C2ST scores for uniformity of the ranks.

check_uniformity_frequentist(ranks, num_posterior_samples)

Return p-values for uniformity of the ranks using Kolmogorov-Smirnov test.

run_sbc(thetas, xs, posterior_samples[, reduce_fns, ...])

Run simulation-based calibration (SBC) or expected coverage.

sbc_rank_plot(ranks, num_posterior_samples[, ...])

Plot simulation-based calibration ranks as empirical CDFs or histograms.

Module Contents#

gensbi.diagnostics.sbc._plot_cdf_region_expected_under_uniformity(num_sbc_runs, num_bins, num_repeats, alpha=0.2, color='gray')[source]#

Plot region of empirical cdfs expected under uniformity on the current axis.

Parameters:

  • num_sbc_runs (int) – Number of SBC runs.

  • num_bins (int) – Number of bins.

  • num_repeats (int) – Number of repeats.

  • alpha (float, optional) – Alpha for the region. Default is 0.2.

  • color (str, optional) – Color for the region. Default is “gray”.

Return type:

None

gensbi.diagnostics.sbc._plot_hist_region_expected_under_uniformity(num_sbc_runs, num_bins, num_posterior_samples, alpha=0.2, color='gray')[source]#

Plot region of empirical cdfs expected under uniformity.

Parameters:

  • num_sbc_runs (int) – Number of SBC runs.

  • num_bins (int) – Number of bins.

  • num_posterior_samples (int) – Number of posterior samples.

  • alpha (float, optional) – Alpha for the region. Default is 0.2.

  • color (str, optional) – Color for the region. Default is “gray”.

Return type:

None

gensbi.diagnostics.sbc._plot_ranks_as_cdf(ranks, num_bins, num_repeats, ranks_label=None, xlabel=None, color=None, alpha=0.8, show_ylabel=True, num_ticks=3)[source]#

Plot ranks as empirical CDFs on the current axis.

Parameters:

  • ranks (np.ndarray) – SBC ranks in shape (num_sbc_runs, ).

  • num_bins (int) – Number of bins for the histogram, recommendation is num_sbc_runs / 20.

  • num_repeats (int) – Number of repeats of each CDF step, i.e., resolution of the eCDF.

  • ranks_label (str, optional) – Label for the ranks, e.g., when comparing ranks of different methods.

  • xlabel (str, optional) – Label for the current parameter.

  • color (str, optional) – Line color for the cdf.

  • alpha (float, optional) – Line transparency. Default is 0.8.

  • show_ylabel (bool, optional) – Whether to show y-label “counts”. Default is True.

  • num_ticks (int, optional) – Number of ticks on the x-axis. Default is 3.

Return type:

None

gensbi.diagnostics.sbc._plot_ranks_as_hist(ranks, num_bins, num_posterior_samples, ranks_label=None, xlabel=None, color='firebrick', alpha=0.8, show_ylabel=False, num_ticks=3, xlim_offset_factor=0.1)[source]#

Plot ranks as histograms on the current axis.

Parameters:

  • ranks (np.ndarray) – SBC ranks in shape (num_sbc_runs, ).

  • num_bins (int) – Number of bins for the histogram, recommendation is num_sbc_runs / 20.

  • num_posterior_samples (int) – Number of posterior samples used for ranking.

  • ranks_label (str, optional) – Label for the ranks, e.g., when comparing ranks of different methods.

  • xlabel (str, optional) – Label for the current parameter.

  • color (str, optional) – Histogram color, default from Talts et al. Default is “firebrick”.

  • alpha (float, optional) – Histogram transparency. Default is 0.8.

  • show_ylabel (bool, optional) – Whether to show y-label “counts”. Default is False.

  • num_ticks (int, optional) – Number of ticks on the x-axis. Default is 3.

  • xlim_offset_factor (float, optional) – Factor for empty space left and right of the histogram. Default is 0.1.

Return type:

None

gensbi.diagnostics.sbc._prepare_reduce_functions(reduce_fns, param_dim)[source]#

Prepare reduction functions for SBC analysis.

Parameters:

  • reduce_fns (str or Callable or List[Callable]) – Function(s) to reduce parameters to 1D.

  • param_dim (int) – Dimensionality of parameter space.

Returns:

List of callable reduction functions.

Return type:

List[Callable]

gensbi.diagnostics.sbc._run_sbc(thetas, xs, posterior_samples, reduce_fns='marginals', show_progress_bar=True)[source]#

Calculate ranks for SBC or expected coverage.

Parameters:

  • thetas (Array) – Ground-truth parameters.

  • xs (Array) – Observed data corresponding to thetas.

  • posterior_samples (Array) – Samples from posterior distribution.

  • reduce_fns (str or Callable or List[Callable], optional) – Functions to reduce parameter space to 1D. Default is “marginals”.

  • show_progress_bar (bool, optional) – Whether to show progress bar. Default is True.

Returns:

Array of ranks for each parameter and reduction function.

Return type:

Array

gensbi.diagnostics.sbc._sbc_rank_plot(ranks, num_posterior_samples, num_bins=None, plot_type='cdf', parameter_labels=None, ranks_labels=None, colors=None, num_repeats=50, line_alpha=0.8, show_uniform_region=True, uniform_region_alpha=0.3, xlim_offset_factor=0.1, num_cols=4, params_in_subplots=False, show_ylabel=False, sharey=False, fig=None, legend_kwargs=None, ax=None, figsize=None)[source]#

Plot simulation-based calibration ranks as empirical CDFs or histograms.

Parameters:

  • ranks (Array or List[Array]) – Array of ranks to be plotted shape (num_sbc_runs, num_parameters), or list of Arrays when comparing several sets of ranks, e.g., set of ranks obtained from different methods.

  • num_posterior_samples (int) – Number of posterior samples used for ranking.

  • num_bins (int, optional) – Number of bins used for binning the ranks. Default is num_sbc_runs / 20.

  • plot_type (str, optional) – Type of SBC plot, histograms (“hist”) or empirical cdfs (“cdf”). Default is “cdf”.

  • parameter_labels (List[str], optional) – List of labels for each parameter dimension.

  • ranks_labels (List[str], optional) – List of labels for each set of ranks.

  • colors (List[str], optional) – List of colors for each parameter dimension, or each set of ranks.

  • num_repeats (int, optional) – Number of repeats for each empirical CDF step (resolution). Default is 50.

  • line_alpha (float, optional) – Alpha for cdf lines or histograms. Default is 0.8.

  • show_uniform_region (bool, optional) – Whether to plot the region showing the cdfs expected under uniformity. Default is True.

  • uniform_region_alpha (float, optional) – Alpha for region showing the cdfs expected under uniformity. Default is 0.3.

  • xlim_offset_factor (float, optional) – Factor for empty space left and right of the histogram. Default is 0.1.

  • num_cols (int, optional) – Number of subplot columns, e.g., when plotting ranks of many parameters. Default is 4.

  • params_in_subplots (bool, optional) – Whether to show each parameter in a separate subplot, or all in one. Default is False.

  • show_ylabel (bool, optional) – Whether to show ylabels and ticks. Default is False.

  • sharey (bool, optional) – Whether to share the y-labels, ticks, and limits across subplots. Default is False.

  • fig (Figure, optional) – Figure object to plot in.

  • legend_kwargs (Dict, optional) – Kwargs for the legend.

  • ax (Axes, optional) – Axis object, must contain as many sublpots as parameters or len(ranks).

  • figsize (tuple, optional) – Dimensions of figure object, default (8, 5) or (len(ranks) * 4, 5).

Returns:

  • fig (Figure) – Figure object.

  • ax (Axes) – Axis object.

Return type:

Tuple[matplotlib.figure.Figure, matplotlib.axes.Axes]

gensbi.diagnostics.sbc._validate_sbc_inputs(thetas, xs, num_sbc_samples, num_posterior_samples)[source]#

Validate inputs for the SBC procedure.

Parameters:

  • thetas (Array) – Ground-truth parameters.

  • xs (Array) – Observed data.

  • num_sbc_samples (int) – Number of SBC samples.

  • num_posterior_samples (int) – Number of posterior samples.

Return type:

None

gensbi.diagnostics.sbc.check_prior_vs_dap(prior_samples, dap_samples)[source]#

Returns the C2ST accuracy between prior and data-averaged posterior samples.

C2ST is calculated for each dimension separately.

According to simulation-based calibration, the inference method is well-calibrated if the data-averaged posterior samples follow the same distribution as the prior, i.e., if the C2ST score is close to 0.5. If it is not, then this suggests that the inference method is not well-calibrated (see Talts et al, “Simulation-based calibration” for details).

Parameters:

  • prior_samples (Array) – Samples from the prior distribution.

  • dap_samples (Array) – Samples from the data-averaged posterior.

Returns:

Array of C2ST scores for each parameter dimension.

Return type:

Array

gensbi.diagnostics.sbc.check_sbc(ranks, prior_samples, dap_samples, num_posterior_samples=1000, num_c2st_repetitions=1)[source]#

Return uniformity checks and data-averaged posterior checks for SBC.

Parameters:

  • ranks (Array) – Ranks for each SBC run and for each model parameter, shape (N, dim_parameters).

  • prior_samples (Array) – N samples from the prior.

  • dap_samples (Array) – N samples from the data-averaged posterior.

  • num_posterior_samples (int, optional) – Number of posterior samples used for SBC ranking. Default is 1000.

  • num_c2st_repetitions (int, optional) – Number of times C2ST is repeated to estimate robustness. Default is 1.

Returns:

Dictionary containing: - ks_pvals: p-values of the Kolmogorov-Smirnov test of uniformity,

one for each dim_parameters.

  • c2st_ranks: C2ST accuracy between ranks and uniform baseline, one for each dim_parameters.

  • c2st_dap: C2ST accuracy between prior and DAP samples, single value.

Return type:

Dict[str, Array]

gensbi.diagnostics.sbc.check_uniformity_c2st(ranks, num_posterior_samples, num_repetitions=1, seed=1)[source]#

Return C2ST scores for uniformity of the ranks.

Run a C2ST between ranks and uniform samples.

Parameters:

  • ranks (Array) – Ranks for each SBC run and for each model parameter, shape (N, dim_parameters).

  • num_posterior_samples (int) – Number of posterior samples used for SBC ranking.

  • num_repetitions (int, optional) – Repetitions of C2ST tests to estimate classifier variance. Default is 1.

  • seed (int, optional) – Random seed. Default is 1.

Returns:

c2st_ranks: C2ST accuracy between ranks and uniform baseline, one for each dim_parameters.

Return type:

Array

gensbi.diagnostics.sbc.check_uniformity_frequentist(ranks, num_posterior_samples)[source]#

Return p-values for uniformity of the ranks using Kolmogorov-Smirnov test.

Parameters:

  • ranks (Array) – Ranks for each SBC run and for each model parameter, shape (N, dim_parameters).

  • num_posterior_samples (int) – Number of posterior samples used for SBC ranking.

Returns:

ks_pvals: p-values of the Kolmogorov-Smirnov test of uniformity, one for each dim_parameters.

Return type:

Array

gensbi.diagnostics.sbc.run_sbc(thetas, xs, posterior_samples, reduce_fns='marginals', show_progress_bar=True, **kwargs)[source]#

Run simulation-based calibration (SBC) or expected coverage.

Note: This function implements two versions of coverage diagnostics:

Parameters:

  • thetas (Array) – Ground-truth parameters for SBC, simulated from the prior.

  • xs (Array) – Observed data for SBC, simulated from thetas.

  • posterior_samples (Array) – Samples from the posterior. Shape: (num_posterior_samples, num_sbc_samples, dim_theta).

  • reduce_fns (str or Callable or List[Callable], optional) – Function used to reduce the parameter space into 1D. Simulation-based calibration can be recovered by setting this to the string “marginals”. Sample-based expected coverage can be recovered by setting it to posterior.log_prob (as a Callable). Default is “marginals”.

  • show_progress_bar (bool, optional) – Whether to display a progress bar over SBC runs. Default is True.

  • **kwargs – Additional keyword arguments.

Returns:

  • ranks (Array) – Ranks of the ground truth parameters under the inferred posterior.

  • dap_samples (Array) – Samples from the data-averaged posterior.

Return type:

Tuple[jax.Array, jax.Array]

gensbi.diagnostics.sbc.sbc_rank_plot(ranks, num_posterior_samples, num_bins=None, plot_type='cdf', parameter_labels=None, ranks_labels=None, colors=None, fig=None, ax=None, figsize=None, **kwargs)[source]#

Plot simulation-based calibration ranks as empirical CDFs or histograms.

Additional options can be passed via the kwargs argument, see _sbc_rank_plot.

Parameters:

  • ranks (Array or List[Array]) – Array of ranks to be plotted shape (num_sbc_runs, num_parameters), or list of Arrays when comparing several sets of ranks, e.g., set of ranks obtained from different methods.

  • num_posterior_samples (int) – Number of posterior samples used for ranking.

  • num_bins (int, optional) – Number of bins used for binning the ranks. Default is num_sbc_runs / 20.

  • plot_type (str, optional) – Type of SBC plot, histograms (“hist”) or empirical cdfs (“cdf”). Default is “cdf”.

  • parameter_labels (List[str], optional) – List of labels for each parameter dimension.

  • ranks_labels (List[str], optional) – List of labels for each set of ranks.

  • colors (List[str], optional) – List of colors for each parameter dimension, or each set of ranks.

  • fig (Figure, optional) – Figure object to plot in.

  • ax (Axes, optional) – Axis object to plot in.

  • figsize (tuple, optional) – Dimensions of figure object.

  • **kwargs – Additional keyword arguments passed to _sbc_rank_plot.

Returns:

  • fig (Figure) – Figure object.

  • ax (Axes) – Axis object.

Return type:

Tuple[matplotlib.figure.Figure, matplotlib.axes.Axes]