superstats.simulation.augmentation.random_choice_contamination

Wrapper for a contaminated random choice data augmentation process.

Classes

RandomChoiceContamination([p_contaminated, ...])

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

class superstats.simulation.augmentation.random_choice_contamination.RandomChoiceContamination(p_contaminated=None, student_t_df=5, response_time_key='response_time', choice_key='choice')

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)

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