superstats.simulation.augmentation#

Data-augmentation processes for generative models.

class superstats.simulation.augmentation.ContaminationProcess[source]#

Bases: ABC

Introduces contamination into simulated data.

Contract: (data, rng) -> {"data": contaminated}. mask is a boolean array of data.shape (True = contaminated) and contaminated is data with masked entries replaced by draws from the process’s contaminant distribution (e.g. guesses, lapses, outliers). Instances are callable, so a ContaminationProcess, a subclass, or a bare function with this signature are interchangeable.

abstract apply(data, rng=None)[source]#

Apply the contamination process.

Parameters:

datanp.ndarray: Simulated data to corrupt with contamination.
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:

resultdict with keys “data” and “contamination_mask”

Parameters:

data (ndarray)
rng (Generator | None)

Return type:

dict

class superstats.simulation.augmentation.MissingProcess[source]#

Bases: ABC

Introduces missingness into simulated data.

Contract: (data, rng) -> filled | {"missing_mask": mask}, where data is a mapping of named arrays with shape (batch_size, num_steps). mask is a boolean array of shape (batch_size, num_steps) (True = missing), and the returned data keys contain the masked entries set to the process’s missing_value. Instances are callable, so a MissingProcess, a subclass, or a bare function with this signature are interchangeable.

abstract apply(data, rng=None)[source]#

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

class superstats.simulation.augmentation.RandomChoiceContamination(p_contaminated=None, student_t_df=5, response_time_key='response_time', choice_key='choice')[source]#

Bases: ContaminationProcess

Contamination at random for diffusion models with a per-dataset contamination probability.

Contamination is drawn per (batch, step): whenever a time step is selected as contaminated, both the response time and the choice at that step are replaced by draws from a contaminant distribution. Any other keys present in data are passed through unchanged. Non-positive response times are treated as non-finished trials: they are left unchanged and excluded from the contaminant distributions.

Contaminant response times are drawn from a heavy-tailed (Student’s t) distribution centered on each dataset’s own log-RT mean and scaled by its own log-RT standard deviation, so contaminants stay plausible in scale for that dataset while still being outliers relative to it [1].

Contaminant choices are drawn either from the observed unique choice values (if choices are discrete) or from a uniform distribution over the observed choice range (if choices are continuous); this is determined once from the whole batch, not per dataset.

[1] Wu, Y., Radev, S. T., & Tuerlinckx, F. (2026). Testing and improving the: robustness of amortized Bayesian inference for cognitive models. Psychological Methods. https://arxiv.org/abs/2412.20586

Parameters:

p_contaminatedfloat, Prior, or None, default: None: Probability that a time step is contaminated. - None (default): drawn from DEFAULT_P_CONTAMINATED_PRIOR. - float: fixed probability, shared across the whole batch. - Prior: sampled once per dataset to obtain a per-dataset

probability (i.e. _draw_p returns one value per batch element, not one shared value for the whole batch).
student_t_dffloat, default: 5: Degrees of freedom for the Student’s t distribution used to generate contaminant response times. Must be greater than 2.
response_time_keystr, default: “response_time”: Key in data containing response times.
choice_keystr, default: “choice”: Key in data containing choices.

Parameters:

p_contaminated (float | Prior | None)
student_t_df (float)
response_time_key (str)
choice_key (str)

apply(data, rng=None)[source]#

Apply random-choice contamination to response times and choices.

Parameters:

datadict with at least the configured response-time and: choice keys, each an np.ndarray of shape (batch_size, num_steps). Any additional keys are passed through unchanged.
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:

resultdict: A shallow copy of data with the configured response-time and choice keys replaced by their contaminated versions, plus “p_contaminated” (the per-dataset contamination probability used, shape (batch_size,)). All other keys in data are carried over unchanged.

Raises:

KeyError: If data is missing either configured required key.

Parameters:

data (dict[str, ndarray])
rng (Generator | None)

Return type:

dict

class superstats.simulation.augmentation.RandomMissingProcess(p_missing=None, missing_value=-1, shared_across_batch=False)[source]#

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)[source]#

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

Modules

`contamination`	Abstract base class for contamination-data augmentation processes.
`missing`	Abstract base class for missing-data augmentation processes.
`random_choice_contamination`	Wrapper for a contaminated random choice data augmentation process.
`random_missing`	Wrapper for missing at random data augmentation process