superstats.workflow#

Workflow orchestration for amortized Bayesian inference.

class superstats.workflow.Workflow(simulator=None, adapter=None, summary_network='recurrent', inference_network='coupling', checkpoint_filepath=None, restore_approximator=True, restore_history=True, **kwargs)[source]#

Bases: object

Lightweight amortized Bayesian inference workflow wrapper.

Wraps bf.BasicWorkflow with sensible defaults for the summary and inference networks, an auto-built adapter when one isn’t supplied, and optional checkpoint/history restoration.

Parameters:

simulatorGenerativeModel or None, optional, default: None: The simulator used for training and, when adapter is not provided, for building a default adapter. Required in that case.
adapterAdapter or None, optional, default: None: Data adapter for the workflow. If None, a default adapter is built from the stochastic simulator.local_keys, simulator.hyper_keys, and simulator.shared_keys (which requires simulator to be set).
summary_network{“recurrent”, “transformer”} or keras.Layer, optional, default: “recurrent”.: String names build a default summary network; otherwise, an already-created Keras layer is used directly.
inference_network{“coupling”, “coupling_flow”} or keras.Layer, optional, default: “coupling”.: String names build a default inference network; otherwise, an already-created Keras layer is used directly.
checkpoint_filepathstr or None, optional, default: None: Directory for saving/restoring the approximator and training history.
restore_approximatorbool, optional, default: True: If True and a checkpoint directory exists at checkpoint_filepath, restore the approximator from it. Otherwise, a warning is issued and training starts from scratch.
restore_historybool, optional, default: True: If True and a history.pkl file exists at checkpoint_filepath, restore self.history from it. Otherwise, a warning is issued.
**kwargs: Forwarded to bf.BasicWorkflow.

Parameters:

simulator (GenerativeModel | None)
adapter (Adapter | None)
summary_network (Literal['recurrent', 'transformer'] | ~keras.src.layers.layer.Layer)
inference_network (Literal['coupling', 'coupling_flow'] | ~keras.src.layers.layer.Layer)
checkpoint_filepath (str | None)
restore_approximator (bool)
restore_history (bool)

property approximator#

The underlying trained BayesFlow approximator object.

Reads through to self.workflow.approximator by default (kept in sync automatically by bf.BasicWorkflow during training), but can be explicitly assigned - e.g. when restoring a checkpoint from disk in __init__.

df_to_dict(df, id_col, data_mapping, missing_value=None, time_col=None)[source]#

Convert a long-format DataFrame into the model’s dict-of-arrays format.

Groups df by id_col and reshapes the columns named in data_mapping into arrays of shape (batch_size, num_steps).

If time_col is given, it must contain discrete integer-like values. The actual labels may be negative or non-contiguous; they are normalized to positions 1..num_steps in sorted time order. Otherwise, rows are placed by order of appearance within each id_col group. Missing or padded positions are flagged in “missing_mask” and filled with the simulator’s missing-value convention when one exists.

Parameters:

dfpd.DataFrame: Long-format data with one row per (dataset, step). Must contain id_col, time_col (if given), and every key in data_mapping.
id_colstr: Name of the column in df identifying which dataset/sequence each row belongs to. Rows are grouped by this column, in order of first appearance, to form the batch dimension.
data_mappingMapping[str, str]: Maps a column name in df to the corresponding key expected by the generative model, e.g. {“rt”: “response_time”, “correct”: “choice”}. The set of values (not keys) must exactly match self.simulator.data_keys.
missing_valueint or float: Sentinel value marking a missing observation, and used to initialize/pad positions with no corresponding row.
time_colstr or None, optional, default: None: Name of the column in df giving each row’s discrete time label. If None, rows are placed by order of appearance within their id_col group instead.

Returns:

datadict of np.ndarray: One entry per generative-model data key, each of shape (batch_size, num_steps), plus “missing_mask” (1 where any mapped column equals missing_value at that step, 0 otherwise) and “time_steps” (each row equal to 1..num_steps).

Parameters:

df (DataFrame)
id_col (str)
data_mapping (Mapping[str, str])
missing_value (int | float | None)
time_col (str | None)

Return type:

dict[str, ndarray]

fit_offline(data, validation_data, epochs=100, batch_size=32, save_history=True, **kwargs)[source]#

Train the approximator on a fixed, pre-simulated dataset.

Parameters:

dataAny: Training data, in the format expected by bf.BasicWorkflow.fit_offline.
validation_dataAny: Validation data, in the same format as data.
epochsint, optional, default: 100: Number of training epochs.
batch_sizeint, optional, default: 32: Training batch size.
save_historybool, optional, default: True: If True, merge this run’s history into self.history and persist it to checkpoint_filepath (if set).
**kwargs: Forwarded to bf.BasicWorkflow.fit_offline.

Returns:

historykeras.callbacks.History - the training history for: this run

Parameters:

epochs (int)
batch_size (int)
save_history (bool)

Return type:

History

fit_online(num_steps, epochs=100, num_batches_per_epoch=100, batch_size=32, save_history=True, **kwargs)[source]#

Train the approximator by simulating data on the fly.

Temporarily binds self.simulator.sample to always draw trajectories of length num_steps with tile_to_steps=True, then restores the original method afterward (even if training raises).

Parameters:

num_stepsint: Number of time steps per simulated trajectory during training.
epochsint, optional, default: 100: Number of training epochs.
num_batches_per_epochint, optional, default: 100: Number of simulated batches per epoch.
batch_sizeint, optional, default: 32: Training batch size.
save_historybool, optional, default: True: If True, merge this run’s history into self.history and persist it to checkpoint_filepath (if set).
**kwargs: Forwarded to bf.BasicWorkflow.fit_online.

Returns:

historykeras.callbacks.History - the training history for: this run

Parameters:

num_steps (int)
epochs (int)
num_batches_per_epoch (int)
batch_size (int)
save_history (bool)

Return type:

History

property history#: keras.callbacks.History or None - the workflow’s training history.

plot_history(history)[source]#

Plot training loss curves.

Parameters:

historykeras.callbacks.History: Training history, e.g. from fit_offline, fit_online, or self.history.

Returns:

figplt.Figure - the loss curve figure

plot_time_invariant_posterior(estimates, targets=None, variable_keys=None, variable_names=None, aggregation=None, mixture_names=None, **kwargs)[source]#

Plot time-invariant posterior diagnostics.

Parameters:

estimatesMapping[str, np.ndarray] or np.ndarray: Posterior samples. If a dict, values of shape (num_datasets, num_post_samples, num_steps, num_components), keyed by variable. If an array, shape (num_datasets, num_post_samples, num_steps, num_params) directly.
targetsMapping[str, np.ndarray], np.ndarray, or None, optional, default: None: Ground-truth values, matching the input type of estimates. If given, drawn as black dashed vertical lines - per dataset when aggregation is None, or collapsed with aggregation into a single line per panel otherwise.
variable_keyssequence of str or None, optional, default: None: Which variables to select and plot, and in what order, when estimates is a dict. Defaults to self.simulator.hyper_keys + self.simulator.shared_keys when not supplied. Ignored for array input.
variable_namessequence of str or None, optional, default: None: Display names for the plotted panels. Defaults to variable_keys (dict input) or param_0, param_1, … (array input).
aggregationcallable() or None, optional, default: None: Controls both the posterior layout and the target summary. If None: one panel per (dataset, parameter) pair. If a callable (e.g. np.mean, np.median): posterior samples (and targets, if given) are pooled/aggregated across datasets into one panel per parameter.
mixture_namesdict or None, optional, default: None: Mapping from parameter name to a list of component names. Defaults to self.simulator.prior._mixture_names() when not supplied.
**kwargs: Forwarded to plot_time_invariant_posterior (e.g. num_cols, color, title_fontsize, label_fontsize, tick_fontsize, figsize).

Returns:

figplt.Figure - the figure instance for optional saving

Parameters:

estimates (Mapping[str, ndarray] | ndarray)
targets (Mapping[str, ndarray] | ndarray | None)
variable_keys (Sequence[str] | None)
variable_names (Sequence[str] | None)
aggregation (Callable | None)
mixture_names (dict | None)

plot_time_varying_posterior(estimates, targets=None, variable_keys=None, variable_names=None, aggregation=None, aggregate_strategy='full_uncertainty', uncertainty_fun='95ci', smoothing=None, smoothing_window=5, marginal=True, **kwargs)[source]#

Plot time-varying posterior diagnostics.

Parameters:

estimatesMapping[str, np.ndarray] or np.ndarray: Posterior samples. If a dict, values of shape (num_datasets, num_post_samples, num_steps, 1), keyed by variable. If an array, shape (num_datasets, num_post_samples, num_steps, num_params) directly.
targetsMapping[str, np.ndarray], np.ndarray, or None, optional, default: None: Ground-truth trajectories, matching the input type of estimates. If a dict, values of shape (num_datasets, num_steps, 1). If an array, shape (num_datasets, num_steps, num_params) directly. If given, drawn as a black dashed line on top of each panel: the raw per-dataset trajectory when aggregation is None, or aggregated across datasets (using aggregation) when aggregation is not None.
variable_keyssequence of str or None, optional, default: None: Which variables to select and plot, and in what order, when estimates/targets are dicts. Defaults to self.simulator.local_keys when not supplied. Ignored for array input.
variable_namessequence of str or None, optional, default: None: Display names (used for panel labels/titles), in the same order as variable_keys (or the array’s last axis). Defaults to variable_keys for dict input, or param_0, param_1, … for array input.
aggregationcallable() or None, optional, default: None: None: one panel per (param, dataset). callable: one panel per param, aggregated across datasets. Called as aggregation(trajectories, axis=0) and must return a (T,) center. The same function aggregates targets across datasets when both targets and aggregation are given.
aggregate_strategy{“full_uncertainty”, “no_epistemic”}, optional, default: “full_uncertainty”: Only used when aggregation is not None. “full_uncertainty”: flatten datasets and posterior samples, then summarize. “no_epistemic”: median across posterior samples per dataset first, then aggregate.
uncertainty_fun{“std”, “95ci”, “mad”, “95hdi”} or callable() or None, optional, default: “95ci”: Band drawn around the center line. A callable receives (N, T) trajectories and must return (lo, hi), each of shape (T,).
smoothing{“sma”, “ema”} or None, optional, default: None: Applied to each trajectory before computing the center, uncertainty, and marginal.
smoothing_windowint, optional, default: 5: Window size for sma, or span parameter for ema.
marginalbool, optional, default: True: Attach a marginal KDE panel to the right of each trajectory axis. The KDE is computed on the same array used for the uncertainty band.
**kwargs: Forwarded to plot_time_varying_posterior (e.g. num_cols, color, alpha, title_fontsize, label_fontsize, tick_fontsize, figsize).

Returns:

figplt.Figure - the figure instance for optional saving

Parameters:

estimates (Mapping[str, ndarray] | ndarray)
targets (Mapping[str, ndarray] | ndarray | None)
variable_keys (Sequence[str] | None)
variable_names (Sequence[str] | None)
aggregation (Callable | None)
aggregate_strategy (Literal['full_uncertainty', 'no_epistemic'])
uncertainty_fun (Literal['std', '95ci', 'mad', '95hdi'] | ~typing.Callable | None)
smoothing (Literal['sma', 'ema'] | None)
smoothing_window (int)
marginal (bool)

resimulate_posterior(posterior_samples, num_sims=10, rng=None)[source]#

Generate posterior predictive simulations from posterior parameter draws.

Parameters:

posterior_samplesdict of np.ndarray: Posterior samples returned by self.sample. Each array should have shape (batch_size, num_samples, num_steps, dim) or (batch_size, num_samples, num_steps).
num_simsint, optional, default: 10: Number of posterior predictive trajectories to simulate per dataset.
rngint or np.random.Generator or None, optional, default: None: Random seed or generator for sampling posterior indices.

Returns:

sim_datadict of np.ndarray: Named simulated variables. Each value has shape (batch_size, num_sims, num_steps).

Raises:

ValueError: If posterior_samples is empty, if a posterior array has fewer than 3 dimensions, if a parameter’s batch size doesn’t match the others, if a parameter has an unsupported number of dimensions, or if a parameter’s shape can’t be reshaped to collapse the sample axis into the batch axis.

Parameters:

posterior_samples (Mapping[str, ndarray])
num_sims (int)

Return type:

dict[str, ndarray]

sample(data, num_samples=500, batch_size=4, **kwargs)[source]#

Run inference on observed data.

Parameters:

datadict of np.ndarray: Observed data to condition on, keyed by the simulator’s named observation variables. Each value should have shape (num_datasets, num_steps). If time_steps is omitted, it is generated automatically. If the simulator was configured with missingness and missing_mask is omitted, an all-observed mask is generated automatically.
num_samplesint, optional, default: 500: Number of posterior samples per dataset.
batch_sizeint, optional, default: 4: Datasets per GPU batch, to avoid out-of-memory errors.
**kwargs: Forwarded to self.approximator.sample.

Returns:

samplesdict of {param_name: np.ndarray} - posterior samples: per parameter

Parameters:

data (dict[str, ndarray])
num_samples (int)
batch_size (int)

Return type:

dict[str, ndarray]

verify_time_invariant(targets, estimates, variable_keys=None, variable_names=None, **kwargs)[source]#

Plot time-invariant parameter recovery and calibration.

Parameters:

targetsMapping[str, np.ndarray] or np.ndarray: If a dict, mapping from parameter name to an np.ndarray of shape (num_sims, dim). If an array, the fully-prepared target array of shape (num_sims, num_params) directly (e.g. mixture components already expanded).
estimatesMapping[str, np.ndarray] or np.ndarray: If a dict, mapping from parameter name to an np.ndarray of shape (num_sims, num_samples, steps, dim). If an array, the fully-prepared estimate array of shape (num_sims, num_pooled_samples, num_params) directly. Must use the same input type (dict or array) as targets.
variable_keyssequence of str or None, optional, default: None: Which time-invariant parameters to include, and in what order, when targets/estimates are dicts. Defaults to self.simulator.hyper_keys + self.simulator.shared_keys when not supplied. Mixture parameters (dim > 1) are expanded into one column per component regardless of this selection. Ignored for array input.
variable_namessequence of str or None, optional, default: None: Display names for the final, expanded columns. For dict input, must match the number of expanded columns (not len(variable_keys)) and defaults to the auto-derived per-component names. For array input, defaults to param_0, param_1, …
**kwargs: Forwarded to both plot_recovery and plot_calibration (e.g. label_fontsize, title_fontsize, tick_fontsize). Note plot_recovery takes color while plot_calibration takes rank_ecdf_color - pass whichever applies, or both, via **kwargs.

Returns:

figstuple - (fig_recovery, fig_calibration), the recovery: and calibration diagnostic figures

Raises:

ValueError: If no time-invariant parameters are found for dict input.

Parameters:

targets (Mapping[str, ndarray] | ndarray)
estimates (Mapping[str, ndarray] | ndarray)
variable_keys (Sequence[str] | None)
variable_names (Sequence[str] | None)

verify_time_varying(targets, estimates, variable_keys=None, variable_names=None, aggregation=<function median>, **kwargs)[source]#

Plot recovery diagnostics over steps for time-varying parameters.

Parameters:

targetsdict: Ground-truth local parameter trajectories, keyed by parameter name; each value has shape (batch_size, num_steps, 1).
estimatesdict: Posterior estimates for the same parameters, keyed by name; each value has shape (batch_size, num_post_samples, num_steps, 1).
variable_keyslist of str or None, optional, default: None: Which parameters to select and plot, and in what order. Defaults to self.simulator.local_keys when not supplied.
variable_nameslist of str or None, optional, default: None: Display names for the plotted columns, in the same order as variable_keys. Defaults to variable_keys when not supplied.
aggregationcallable(), optional, default: np.median: Aggregation function forwarded to plot_time_varying_verification, used to collapse each metric across simulations. Typically np.mean or np.median.
**kwargs: Additional keyword arguments forwarded to plot_time_varying_verification (e.g. colors, title_fontsize).

Returns:

figplt.Figure - the figure instance for optional saving

Parameters:

targets (dict)
estimates (dict)
variable_keys (list | None)
variable_names (list | None)
aggregation (Callable)

Modules

workflow

High-level workflow wrapper around BayesFlow.