superstats.simulation.generative_model#

Generative-model wrapper for joint priors and simulators.

Classes

GenerativeModel(prior, model[, missing, ...])

A generative model that combines a joint prior with a simulation function.
class superstats.simulation.generative_model.GenerativeModel(prior, model, missing='random', contamination=None)[source]#

Bases: object

A generative model that combines a joint prior with a simulation function.

This class facilitates sampling parameters from a joint prior distribution and generating simulated data using a user-provided model function. It handles parameter broadcasting, flattening, and reshaping to support batched simulations with time-varying parameters. Optionally, a missing-data process can be applied to the simulated data to introduce and record missingness.

Parameters:
priorJointPrior

The joint prior distribution over model parameters, which may include both time-varying transitions and time-invariant priors.

modelCallable

The simulation function that takes parameter values and returns simulated data. The function signature determines the expected parameter names and order.

missingMissingProcess, Callable, “random”, or None, optional, default: “random”

Process applied to simulated data to introduce missingness. - Not provided (default) or “random”: uses RandomMissingProcess(),

the default MCAR missingness process.

  • None: disables missingness augmentation and sample will not include a “missing_mask” entry in its result.

  • MissingProcess instance: used as-is.

  • Plain Callable: must follow the same contract as MissingProcess.__call__, i.e. (data_mapping, rng=None) -> filled_mapping | {“missing_mask”: mask}.

Raises:
TypeError

If model is not callable, or if missing is neither None, “random”, nor callable.

Parameters:
get_fixed_params()[source]#

Return deterministic fixed parameters from the prior for model simulation.

Draws a single pilot sample from self.prior and keeps only the fixed-parameter entries that the model actually consumes.

Returns:
fixed_paramsdict of np.ndarray - mapping from parameter name

to its fixed value, restricted to names in self.param_order

Return type:

Dict[str, ndarray]

plot_push_forward(num_sim=20, num_steps=200, data_dim=0, kind='dist', aggregation=None, uncertainty_fun=None, marginal=True, spaghetti=False, **kwargs)[source]#

Render prior push-forward diagnostics for the generative model.

Parameters:
num_simint, optional, default: 20

Number of simulated datasets to generate.

num_stepsint, optional, default: 200

Number of time steps per simulation.

data_dimint or str, optional, default: 0

Observation variable to plot. Integers index self.data_keys; strings select a variable by name.

kind{“dist”, “trajectory”}, optional, default: “dist”

Plot type.

aggregationcallable() or None, optional, default: None

Aggregation function over the dataset dimension, called as aggregation(x, axis=…) (e.g. np.mean, np.median). If None, individual datasets are shown in separate panels. If specified, all datasets are aggregated into a single panel.

uncertainty_fun{“std”, “95ci”, “mad”, “95hdi”} or callable() or None, optional, default: None

Uncertainty function for aggregate trajectory plots. Forwarded directly to plot_push_forward, so the accepted values must match that function’s own supported set.

marginalbool, optional, default: True

If True, include marginal distributions beside trajectories.

spaghettibool, optional, default: False

If True, include individual trajectories.

**kwargs

Forwarded to plot_push_forward.

Returns:
figplt.Figure - the figure containing the requested plot
Parameters:
Return type:

Figure

sample(batch_size, num_steps, include_fixed=False, tile_to_steps=False, rng=None)[source]#

Sample parameters from the prior and generate simulated data.

This method performs a complete generative process: 1. Samples parameters from the joint prior distribution 2. Prepares parameters for vectorized simulation 3. Runs the simulation model 4. Reshapes outputs back to trajectory format 5. Applies self.missing to the data, if configured

Parameters:
batch_sizeint

Number of independent simulation batches to generate.

num_stepsint

Number of time steps per trajectory.

include_fixedbool, optional, default: False

If True, include fixed_params in the returned dictionary.

tile_to_stepsbool, optional, default: False

If True, tile hyper_params and shared_params from shape (batch_size, 1) to (batch_size, num_steps, 1), aligning them with the time axis of local parameters.

rngnp.random.Generator or None, optional, default: None

Random generator forwarded to self.missing. If None, the missing process falls back to its own default (an unseeded generator).

Returns:
resultdict - flat dictionary with the following entries:

  • one entry per simulated observation variable, each with

shape (batch_size, num_steps), corrupted by self.missing if one is configured. - “time_steps”: shape (batch_size, num_steps), each row equal to 1..num_steps. - “missing_mask”: included only if self.missing is not None; shape matches the mask returned by the process (for RandomMissingProcess, (batch_size, num_steps)). - any additional keys the missing process returns beyond the simulator data keys and “missing_mask” (e.g. RandomMissingProcess also returns “p_missing”, shape (batch_size, 1)); omitted if self.missing is None or returns no extra keys. - one entry per sampled parameter. Local (time-varying) params

have shape (batch_size, num_steps); hyper and shared params have shape (batch_size, 1), or (batch_size, num_steps, 1) when tile_to_steps is True.

  • fixed params are included only when include_fixed is True.

The instance attributes local_keys, hyper_keys, shared_keys, fixed_keys, and data_keys record which keys belong to which group.

Raises:
ValueError

If required parameters are missing from the prior or have invalid shapes.

Parameters:
Return type:

Dict[str, ndarray]

simulate_from_parameters(params, batch_size, num_steps)[source]#

Simulate model outputs for given parameter values.

Parameters:
paramsdict of np.ndarray

Parameter values to simulate from, keyed by model parameter name. See _prepare_flat_params for the accepted shapes.

batch_sizeint

Number of independent simulation batches.

num_stepsint

Number of time steps per trajectory.

Returns:
sim_datadict of np.ndarray

Named simulated variables. Each value has shape (batch_size, num_steps).

Raises:
ValueError

If a required parameter is missing from params and has no default in the model signature, or has an unsupported shape.

Parameters:
Return type:

Dict[str, ndarray]