superstats.simulation

Simulation models and built-in cognitive simulators.

class superstats.simulation.GenerativeModel(prior, model, missing='random', contamination=None)

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:

• prior (JointPrior)

• model (Callable)

• missing (MissingProcess | Callable | Literal['random'] | None)

• contamination (ContaminationProcess | Callable | Literal['random_choice'] | None)

get_fixed_params()

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)

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:

• num_sim (int)

• num_steps (int)

• data_dim (int | str)

• kind (str)

• aggregation (Callable | None)

• uncertainty_fun (str | Callable | None)

• marginal (bool)

• spaghetti (bool)

Return type:

Figure

sample(batch_size, num_steps, include_fixed=False, tile_to_steps=False, rng=None)

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:

• batch_size (int)

• num_steps (int)

• include_fixed (bool)

• tile_to_steps (bool)

• rng (Generator | None)

Return type:

Dict[str, ndarray]

simulate_from_parameters(params, batch_size, num_steps)

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:

• params (Dict[str, ndarray])

• batch_size (int)

• num_steps (int)

Return type:

Dict[str, ndarray]

class superstats.simulation.RandomMissingProcess(p_missing=None, missing_value=-1, shared_across_batch=False)

Bases: MissingProcess

MCAR missingness with a per-dataset missing probability.

Missingness is drawn per (batch, step): whenever a time step is selected as missing, all data dimensions at that step are set to missing_value (an entire observation is dropped, not individual features within it).

Parameters:

p_missingfloat, Prior, or None, default: None

Probability that a time step is missing. - None (default): drawn from DEFAULT_P_MISSING_PRIOR, a Beta(2, 18) prior with mean 0.1. - float: fixed probability, shared across the whole batch. - Prior: sampled to obtain the probability. Sampled once for the whole batch if shared_across_batch=True, or once per dataset (default) otherwise. Prior draws (including the default) are clipped to [0, 1].

missing_valuefloat or np.ndarray, default: -1

Value written into masked entries. A scalar fills every observed variable; a mapping sets a per-variable sentinel; an array of shape (num_variables,) sets sentinels in data-key order. Output dtype is promoted as needed (e.g. np.nan forces float; -1 stays int on int data).

shared_across_batchbool, default: False

If True, one probability and one mask are drawn and applied to every dataset in the batch. If False (default), each dataset gets its own probability draw and its own mask.

Parameters:

• p_missing (float | Prior | None)

• missing_value (float)

• shared_across_batch (bool)

apply(data, rng=None)

Apply the missingness process.

Parameters:

datamapping of np.ndarray

Simulated data to corrupt with missingness.

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

Random generator to use. If None, a fresh, unseeded generator is created via _default_rng, so calling apply directly is safe but not reproducible unless a seeded rng is supplied.

Returns:

resultflat dict with data keys, “missing_mask”, and optional metadata

Parameters:

• data (Mapping[str, ndarray])

• rng (Generator | None)

Return type:

dict

superstats.simulation.sample_cdm(v_angle, v_length, a, tau, sigma=1.0, dt=0.001, max_steps=10000)

Sample from the Circular Diffusion Model (CDM).

Simulates a 2D diffusion process starting from the origin, with a constant drift specified in polar form, evolving until it crosses a circular boundary of radius a. The crossing point determines the response angle and the number of steps determines the response time. On each trial the drift vector has length v_length and points in direction v_angle; the two Cartesian components diffuse independently with noise SD sigma until the squared radius reaches a ** 2.

Parameters:

v_anglenp.ndarray of shape (num_trials,)

Direction of the drift vector (in radians) for each trial.

v_lengthnp.ndarray of shape (num_trials,)

Magnitude of the drift vector for each trial. The Cartesian drift components are v_length * cos(v_angle) and v_length * sin(v_angle).

anp.ndarray of shape (num_trials,)

Radius of the circular decision boundary for each trial.

taunp.ndarray of shape (num_trials,)

Non-decision times for each trial.

sigmafloat, optional, default: 1.0

Diffusion noise standard deviation, shared by both Cartesian components. Fixed (not estimated per trial) for identifiability, since the boundary radius a and drift set the overall scale.

dtfloat, optional, default: 0.001

Time step size.

max_stepsint, optional, default: 10000

Maximum number of diffusion steps per trial before timing out.

Returns:

datadict of np.ndarray

Named decision data. “response_time” contains response times (or -5.0 on timeout) and “choice” contains response angles in radians (or -5.0 on timeout). Each array has shape (num_trials,).

Parameters:

• v_angle (ndarray)

• v_length (ndarray)

• a (ndarray)

• tau (ndarray)

• sigma (float)

• dt (float)

• max_steps (int)

Return type:

dict[str, ndarray]

superstats.simulation.sample_ddm(v, a, tau, bias, sigma=1.0, dt=0.001, max_steps=10000)

Sample from the Diffusion Decision Model (DDM) for decision making.

This function simulates decision processes using the DDM, where evidence accumulates over time with drift rate v, boundary separation a, and noise. The simulation stops when a boundary is reached or max_steps is exceeded.

Parameters:

vnp.ndarray of shape (num_steps,)

Drift rates for each trial.

anp.ndarray of shape (num_steps,)

Boundary separation for each trial; decision boundaries are at 0 (lower) and a (upper).

taunp.ndarray of shape (num_steps,)

Non-decision times for each trial.

biasnp.ndarray of shape (num_steps,)

Starting point, as a fraction of a (i.e. the initial evidence is bias * a). 0.5 starts at the midpoint between the two boundaries; values > 0.5 start closer to the upper boundary, values < 0.5 closer to the lower one. Must lie in (0, 1).

sigmafloat, optional, default: 1.0

Diffusion noise standard deviation.

dtfloat, optional, default: 0.001

Time step size.

max_stepsint, optional, default: 10000

Maximum number of diffusion steps per trial before timing out.

Returns:

datadict of np.ndarray

Named decision data. “response_time” contains response times (or -1.0 on timeout) and “choice” contains choices (1.0 for the upper boundary, 0.0 for the lower boundary, -1.0 on timeout). Each array has shape (num_steps,).

Parameters:

• v (ndarray)

• a (ndarray)

• tau (ndarray)

• bias (ndarray)

• sigma (float)

• dt (float)

• max_steps (int)

Return type:

dict[str, ndarray]

superstats.simulation.sample_rdm(v_base, v_diff, a_base, tau, bias, sigma_diff, num_accumulators=2, correct_idx=None, sigma_base=1.0, dt=0.001, max_steps=10000)

Sample from the Racing Diffusion Model (RDM).

Simulates num_accumulators independent diffusion accumulators racing from a starting point of 0 toward their own threshold; the first to cross wins and determines the response and response time. On each trial, the accumulator at index correct_idx[i] is treated as the correct/target accumulator: it receives a drift advantage of v_diff, a bias-scaled threshold, and noise scaled by sigma_diff. All other accumulators on that trial share the disadvantaged drift, an unscaled threshold, and noise fixed at sigma_base.

Parameters:

v_basenp.ndarray of shape (num_trials,)

Base drift rate shared by all accumulators before the correct/incorrect adjustment.

v_diffnp.ndarray of shape (num_trials,)

Drift rate difference between the correct and incorrect accumulators. The correct accumulator gets v_base + v_diff / 2; all other accumulators get v_base - v_diff / 2.

a_basenp.ndarray of shape (num_trials,)

Base threshold distance from the origin for each trial.

taunp.ndarray of shape (num_trials,)

Non-decision times for each trial.

biasnp.ndarray of shape (num_trials,)

Threshold scaling factor in [0, 1] for the correct/target accumulator: its threshold is a_base * bias. All other accumulators use the unscaled threshold a_base.

sigma_diffnp.ndarray of shape (num_trials,)

Noise scaling factor in [0, +inf) for the correct/target accumulator: its noise SD is sigma_base * sigma_diff. All other accumulators always use sigma_base directly.

num_accumulatorsint

Number of racing accumulators per trial (fixed across trials).

correct_idxnp.ndarray of shape (num_trials,), optional

Index (into 0 .. num_accumulators - 1) of the correct/target accumulator for each trial. If left empty, accumulator 0 is treated as correct on every trial.

sigma_basefloat, optional, default: 1.0

Diffusion noise standard deviation of the non-correct accumulators. Fixed (not estimated per trial) for identifiability.

dtfloat, optional, default: 0.001

Time step size.

max_stepsint, optional, default: 10000

Maximum number of diffusion steps per trial before timing out.

Returns:

datadict of np.ndarray

Named decision data. “response_time” contains response times (or -1.0 on timeout) and “choice” contains the index of the winning accumulator (or -1.0 on timeout). Each array has shape (num_trials,).

Parameters:

• v_base (ndarray)

• v_diff (ndarray)

• a_base (ndarray)

• tau (ndarray)

• bias (ndarray)

• sigma_diff (ndarray)

• num_accumulators (int)

• correct_idx (ndarray | None)

• sigma_base (float)

• dt (float)

• max_steps (int)

Return type:

dict[str, ndarray]

Modules

augmentation	Data-augmentation processes for generative models.
cognitive	Cognitive-model simulators.
generative_model	Generative-model wrapper for joint priors and simulators.