ssms

ssms.Simulator

Simulator(
    model=None,
    boundary=None,
    drift=None,
    simulator_function=None,
    parameter_adapter=None,
    parameter_adaptations=None,
    **config_overrides
)

Class-based interface for Sequential Sampling Model simulations.

This class provides a flexible, extensible way to configure and run SSM simulations. It supports: - Pre-defined models (via string names like "ddm", "angle", etc.) - Custom boundary and drift functions with existing simulators - Fully custom simulator functions - Configuration overrides for any model parameter

Examples:

Basic usage with pre-defined model:

>>> sim = Simulator("ddm")
>>> results = sim.simulate(theta={'v': 0.5, 'a': 1.0, 'z': 0.5, 't': 0.3})

Custom boundary function:

>>> def my_boundary(t, theta, scale):
...     return scale * np.sin(theta * t)
>>> sim = Simulator("ddm", boundary=my_boundary, boundary_params=["theta", "scale"])
>>> results = sim.simulate(theta={'v': 0.5, 'a': 1.0, 'z': 0.5, 't': 0.3,
...                               'theta': 0.5, 'scale': 1.0})

Fully custom simulator:

>>> def my_sim(v, a, z, t, max_t=20, n_samples=1000, **kwargs):
...     rts = np.random.exponential(1/abs(v), n_samples) + t
...     choices = np.where(np.random.random(n_samples) < z, 1, -1)
...     return {'rts': rts, 'choices': choices,
...             'metadata': {'model': 'custom', 'n_samples': n_samples}}
>>> sim = Simulator(simulator_function=my_sim, params=["v", "a", "z", "t"],
...                 nchoices=2)
>>> results = sim.simulate(theta={'v': 0.5, 'a': 1.0, 'z': 0.5, 't': 0.3})

Attributes:

config (dict) –

The full model configuration dictionary

Parameters:

model (str, dict, or None, default: None ) –

Either a model name (e.g., "ddm", "angle"), a full configuration dictionary, or None if providing a custom simulator_function.
boundary (str, Callable, or None, default: None ) –

Boundary function. Can be: - A string name from the boundary registry (e.g., "angle", "weibull_cdf") - A callable function with signature: func(t, **params) -> np.ndarray - None to use the model's default boundary
drift (str, Callable, or None, default: None ) –

Drift function. Can be: - A string name from the drift registry (e.g., "constant", "gamma_drift") - A callable function with signature: func(t, **params) -> np.ndarray - None to use the model's default drift (if applicable)
simulator_function (Callable or None, default: None ) –

Custom simulator function. If provided, this overrides the model's default simulator. Must follow the simulator interface contract.
parameter_adapter (ParameterSimulatorAdapterProtocol or None, default: None ) –

Custom parameter adapter for preparing parameters for simulators. If None, uses ModularParameterSimulatorAdapter with default adaptations for the model.
parameter_adaptations (list[ParameterAdaptation] or None, default: None ) –

Additional parameter adaptations to apply AFTER the model's default adaptations. Useful for adding custom parameter preparation without replacing the entire adapter. Only used if parameter_adapter is None.
**config_overrides –

Additional configuration parameters to override. Common options: - params : list[str] - Parameter names (required if simulator_function provided) - param_bounds : list - [[lower bounds], [upper bounds]] - nchoices : int - Number of choices (required if simulator_function provided) - choices : list - Possible choice values - boundary_params : list[str] - Parameters for custom boundary function - drift_params : list[str] - Parameters for custom drift function - name : str - Model name (for custom simulators)

Raises:

ValueError –

If configuration is invalid or missing required parameters

Examples:

Use default modular adapter:

>>> sim = Simulator("lba2")
>>> results = sim.simulate(theta={'v0': 0.5, 'v1': 0.6, 'A': 0.5, 'b': 1.0})

Add custom adaptations:

>>> from ssms.basic_simulators.parameter_adapters import SetDefaultValue
>>> sim = Simulator("ddm", parameter_adaptations=[SetDefaultValue("custom_param", 42)])

Use a custom parameter adapter:

>>> custom_adapter = ModularParameterSimulatorAdapter()
>>> sim = Simulator("ddm", parameter_adapter=custom_adapter)

config `property`

config

Get the full model configuration.

Returns:

dict –

Configuration dictionary

parameter_adapter `property`

parameter_adapter

Get the configured parameter adapter.

Returns:

ParameterSimulatorAdapterProtocol –

The parameter adapter instance used for parameter preparation

simulate

simulate(
    theta,
    n_samples=1000,
    delta_t=0.001,
    max_t=20,
    no_noise=False,
    sigma_noise=None,
    smooth_unif=True,
    random_state=None,
    return_option="full",
    n_threads=1,
    extra_fields=None,
)

Run simulation with given parameters.

Parameters:

theta (list, np.ndarray, dict, or pd.DataFrame) –

Model parameters. If dict or DataFrame, keys/columns should match parameter names in config. If array, order should match config['params'].
extra_fields (dict or None, default: None ) –

Optional per-trial covariates forwarded to the model function (e.g. the aDDM fixations r1, r2, flag, sacc_array, d, sigma), so a covariate-driven model can condition on observed inputs rather than self-sampling them. Empty/None is a no-op for models that take no covariates. Arrays are per-trial (aligned with theta's trial axis).
n_samples (int, default: 1000 ) –

Number of simulation samples per parameter set
delta_t (float, default: 0.001 ) –

Time step size for simulation
max_t (float, default: 20 ) –

Maximum simulation time
no_noise (bool, default: False ) –

If True, disable noise in simulation
sigma_noise (float or None, default: None ) –

Standard deviation of noise. If None, uses model defaults.
smooth_unif (bool, default: True ) –

Whether to add uniform smoothing to RTs
random_state (int or None, default: None ) –

Integer seed for the C-level RNG. Must lie in [-2**31, 2**31 - 1] (fits a 32-bit signed C long, required on Windows). None picks a seed in [0, 2**31 - 1] automatically.
return_option (str, default: "full" ) –

Output format: "full" or "minimal"
n_threads (int, default: 1 ) –

Number of threads for parallel execution. If > 1 and OpenMP is available, uses multi-threaded simulation. Note: trajectory recording is only available with n_threads=1.

Returns:

dict –

Simulation results with keys: - 'rts' : np.ndarray of reaction times - 'choices' : np.ndarray of choices - 'metadata' : dict with simulation metadata - Additional keys depending on model and return_option

Raises:

ValueError –

If parameters are invalid

validate_params

validate_params(theta)

Validate parameter values.

Parameters:

theta (dict) –

Parameter dictionary to validate

Raises:

ValueError –

If parameters are invalid

ssms.basic_simulators

Simulator

Simulator(
    model=None,
    boundary=None,
    drift=None,
    simulator_function=None,
    parameter_adapter=None,
    parameter_adaptations=None,
    **config_overrides
)

Class-based interface for Sequential Sampling Model simulations.

This class provides a flexible, extensible way to configure and run SSM simulations. It supports: - Pre-defined models (via string names like "ddm", "angle", etc.) - Custom boundary and drift functions with existing simulators - Fully custom simulator functions - Configuration overrides for any model parameter

Examples:

Basic usage with pre-defined model:

>>> sim = Simulator("ddm")
>>> results = sim.simulate(theta={'v': 0.5, 'a': 1.0, 'z': 0.5, 't': 0.3})

Custom boundary function:

>>> def my_boundary(t, theta, scale):
...     return scale * np.sin(theta * t)
>>> sim = Simulator("ddm", boundary=my_boundary, boundary_params=["theta", "scale"])
>>> results = sim.simulate(theta={'v': 0.5, 'a': 1.0, 'z': 0.5, 't': 0.3,
...                               'theta': 0.5, 'scale': 1.0})

Fully custom simulator:

>>> def my_sim(v, a, z, t, max_t=20, n_samples=1000, **kwargs):
...     rts = np.random.exponential(1/abs(v), n_samples) + t
...     choices = np.where(np.random.random(n_samples) < z, 1, -1)
...     return {'rts': rts, 'choices': choices,
...             'metadata': {'model': 'custom', 'n_samples': n_samples}}
>>> sim = Simulator(simulator_function=my_sim, params=["v", "a", "z", "t"],
...                 nchoices=2)
>>> results = sim.simulate(theta={'v': 0.5, 'a': 1.0, 'z': 0.5, 't': 0.3})

Attributes:

config (dict) –

The full model configuration dictionary

Parameters:

model (str, dict, or None, default: None ) –

Either a model name (e.g., "ddm", "angle"), a full configuration dictionary, or None if providing a custom simulator_function.
boundary (str, Callable, or None, default: None ) –

Boundary function. Can be: - A string name from the boundary registry (e.g., "angle", "weibull_cdf") - A callable function with signature: func(t, **params) -> np.ndarray - None to use the model's default boundary
drift (str, Callable, or None, default: None ) –

Drift function. Can be: - A string name from the drift registry (e.g., "constant", "gamma_drift") - A callable function with signature: func(t, **params) -> np.ndarray - None to use the model's default drift (if applicable)
simulator_function (Callable or None, default: None ) –

Custom simulator function. If provided, this overrides the model's default simulator. Must follow the simulator interface contract.
parameter_adapter (ParameterSimulatorAdapterProtocol or None, default: None ) –

Custom parameter adapter for preparing parameters for simulators. If None, uses ModularParameterSimulatorAdapter with default adaptations for the model.
parameter_adaptations (list[ParameterAdaptation] or None, default: None ) –

Additional parameter adaptations to apply AFTER the model's default adaptations. Useful for adding custom parameter preparation without replacing the entire adapter. Only used if parameter_adapter is None.
**config_overrides –

Additional configuration parameters to override. Common options: - params : list[str] - Parameter names (required if simulator_function provided) - param_bounds : list - [[lower bounds], [upper bounds]] - nchoices : int - Number of choices (required if simulator_function provided) - choices : list - Possible choice values - boundary_params : list[str] - Parameters for custom boundary function - drift_params : list[str] - Parameters for custom drift function - name : str - Model name (for custom simulators)

Raises:

ValueError –

If configuration is invalid or missing required parameters

Examples:

Use default modular adapter:

>>> sim = Simulator("lba2")
>>> results = sim.simulate(theta={'v0': 0.5, 'v1': 0.6, 'A': 0.5, 'b': 1.0})

Add custom adaptations:

>>> from ssms.basic_simulators.parameter_adapters import SetDefaultValue
>>> sim = Simulator("ddm", parameter_adaptations=[SetDefaultValue("custom_param", 42)])

Use a custom parameter adapter:

>>> custom_adapter = ModularParameterSimulatorAdapter()
>>> sim = Simulator("ddm", parameter_adapter=custom_adapter)

config `property`

config

Get the full model configuration.

Returns:

dict –

Configuration dictionary

parameter_adapter `property`

parameter_adapter

Get the configured parameter adapter.

Returns:

ParameterSimulatorAdapterProtocol –

The parameter adapter instance used for parameter preparation

simulate

simulate(
    theta,
    n_samples=1000,
    delta_t=0.001,
    max_t=20,
    no_noise=False,
    sigma_noise=None,
    smooth_unif=True,
    random_state=None,
    return_option="full",
    n_threads=1,
    extra_fields=None,
)

Run simulation with given parameters.

Parameters:

theta (list, np.ndarray, dict, or pd.DataFrame) –

Model parameters. If dict or DataFrame, keys/columns should match parameter names in config. If array, order should match config['params'].
extra_fields (dict or None, default: None ) –

Optional per-trial covariates forwarded to the model function (e.g. the aDDM fixations r1, r2, flag, sacc_array, d, sigma), so a covariate-driven model can condition on observed inputs rather than self-sampling them. Empty/None is a no-op for models that take no covariates. Arrays are per-trial (aligned with theta's trial axis).
n_samples (int, default: 1000 ) –

Number of simulation samples per parameter set
delta_t (float, default: 0.001 ) –

Time step size for simulation
max_t (float, default: 20 ) –

Maximum simulation time
no_noise (bool, default: False ) –

If True, disable noise in simulation
sigma_noise (float or None, default: None ) –

Standard deviation of noise. If None, uses model defaults.
smooth_unif (bool, default: True ) –

Whether to add uniform smoothing to RTs
random_state (int or None, default: None ) –

Integer seed for the C-level RNG. Must lie in [-2**31, 2**31 - 1] (fits a 32-bit signed C long, required on Windows). None picks a seed in [0, 2**31 - 1] automatically.
return_option (str, default: "full" ) –

Output format: "full" or "minimal"
n_threads (int, default: 1 ) –

Number of threads for parallel execution. If > 1 and OpenMP is available, uses multi-threaded simulation. Note: trajectory recording is only available with n_threads=1.

Returns:

dict –

Simulation results with keys: - 'rts' : np.ndarray of reaction times - 'choices' : np.ndarray of choices - 'metadata' : dict with simulation metadata - Additional keys depending on model and return_option

Raises:

ValueError –

If parameters are invalid

validate_params

validate_params(theta)

Validate parameter values.

Parameters:

theta (dict) –

Parameter dictionary to validate

Raises:

ValueError –

If parameters are invalid

boundary_functions

Define a collection of boundary functions for the simulators in the package.

addm_collapse

addm_collapse(t=0, a=1.0, b=0.5)

Linearly collapsing boundary for the attentional DDM: a - b * t.

The aDDM simulator uses this for the upper boundary and mirrors it (-(a - b * t)) for the lower one. b >= 0 collapses the bounds over time; b = 0 recovers a constant boundary. This documents the boundary the cssm.addm engine computes internally from its own a/b params.

Arguments

t (float or np.ndarray, optional): Time point(s). Defaults to 0.
a (float, optional): Boundary intercept at t=0. Defaults to 1.0.
b (float, optional): Collapse slope (>= 0). Defaults to 0.5.

Returns:

float or np.ndarray: ``a - b * t``. –

angle `module-attribute`

angle = angle

Linear collapsing boundary at angle theta.

Arguments

t (float or np.ndarray, optional): Time point(s). Defaults to 1.
a (float, optional): Threshold parameter (starting height). Defaults to 1.0.
theta (float, optional): Collapse angle in radians. Defaults to 1.0.

Returns:

np.ndarray or float: Boundary value = a + t * tan(theta) –

conflict_gamma `module-attribute`

conflict_gamma = conflict_gamma

Conflict boundary with gamma bump and linear collapse.

Arguments

t: (float, np.ndarray)
    Time points (with arbitrary measure, but in HDDM it is used as seconds),
    at which to evaluate the bound. Defaults to np.arange(0, 20, 0.1).
a: float
    Threshold parameter (starting height). Defaults to 1.0.
theta: float
    Collapse angle. Defaults to 0.5.
scale: float
    Scaling the gamma distribution of the boundary
    (since bound does not have to integrate to one). Defaults to 1.0.
alphaGamma: float
    alpha parameter for a gamma in scale shape parameterization. Defaults to 1.01.
scaleGamma: float
    scale parameter for a gamma in scale shape parameterization. Defaults to 0.3.

Returns:

np.ndarray: Boundary value = a + gamma_bump(t) + t*tan(theta) –

constant `module-attribute`

constant = constant

Constant boundary function.

Arguments

t (float or np.ndarray, optional): Time point(s). Defaults to 0.
a (float, optional): Threshold parameter. Defaults to 1.0.

Returns:

float or np.ndarray: Constant boundary value = a (scalar if t is scalar, array if t is array) –

generalized_logistic `module-attribute`

generalized_logistic = generalized_logistic

Generalized logistic boundary function.

Arguments

t (float or np.ndarray, optional): Time point(s). Defaults to 1.
a (float, optional): Threshold parameter. Defaults to 1.0.
B (float, optional): Growth rate. Defaults to 2.0.
M (float, optional): Time of maximum growth. Defaults to 3.0.
v (float, optional): Affects near which asymptote maximum growth occurs.
Defaults to 0.5.

Returns:

np.ndarray or float: Boundary value = a + logistic_decay(t) –

weibull_cdf `module-attribute`

weibull_cdf = weibull_cdf

Weibull decay boundary function.

Arguments

t (float or np.ndarray, optional): Time point(s). Defaults to 1.
a (float, optional): Threshold parameter (starting height). Defaults to 1.0.
alpha (float, optional): Shape parameter. Defaults to 1.0.
beta (float, optional): Scale parameter. Defaults to 1.0.

Returns:

np.ndarray or float: Boundary value = a * exp(-(t/β)^α) –

drift_functions

Define a collection of drift functions for the simulators in the package.

All drift functions accept v as a parameter and return the FINAL drift value, consistent with boundary functions that accept a and return the final boundary value.

attend_drift `module-attribute`

attend_drift = attend_drift

Shrink spotlight model, which involves a time varying function dependent on a linearly decreasing standard deviation of attention.

The final drift is v + (attention-weighted drift component).

Arguments

t: np.ndarray
    Timepoints at which to evaluate the drift.
    Usually np.arange() of some sort.
v: float
    Base drift rate (typically 0 for shrink_spot models). Defaults to 0.0.
pouter: float
    perceptual input for outer flankers
pinner: float
    perceptual input for inner flankers
ptarget: float
    perceptual input for target flanker
r: float
    rate parameter for sda decrease
sda: float
    width of attentional spotlight

Return

np.ndarray Final drift evaluated at timepoints t

attend_drift_simple

attend_drift_simple(
    t=arange(0, 20, 0.1),
    v=0.0,
    ptarget=-0.3,
    pouter=-0.3,
    r=0.5,
    sda=2,
)

Drift function for shrinking spotlight model, which involves a time varying function dependent on a linearly decreasing standard deviation of attention.

The final drift is v + (attention-weighted drift component).

Arguments

t: np.ndarray
    Timepoints at which to evaluate the drift.
    Usually np.arange() of some sort.
v: float
    Base drift rate (typically 0 for shrink_spot models). Defaults to 0.0.
pouter: float
    perceptual input for outer flankers
ptarget: float
    perceptual input for target flanker
r: float
    rate parameter for sda decrease
sda: float
    width of attentional spotlight

Return

np.ndarray Final drift evaluated at timepoints t

conflict_ds_drift `module-attribute`

conflict_ds_drift = conflict_ds_drift

This drift is inspired by a conflict task which involves a target and a distractor stimuli both presented simultaneously.

Two drift timecourses are linearly combined weighted by the coherence in the respective target and distractor stimuli. Each timecourse follows a dynamical system as described in the ds_support_analytic() function.

The final drift is v + (combined drift timecourse).

Arguments

t: np.ndarray
    Timepoints at which to evaluate the drift.
    Usually np.arange() of some sort.
v: float
    Base drift rate (typically 0 for conflict models). Defaults to 0.0.
tinit: float
    Initial condition of target drift timecourse
dinit: float
    Initial condition of distractor drift timecourse
tslope: float
    Slope parameter for target drift timecourse
dslope: float
    Slope parameter for distractor drift timecourse
tfixedp: float
    Fixed point for target drift timecourse
tcoh: float
    Coefficient for the target drift timecourse
dcoh: float
    Coefficient for the distractor drift timecourse

Return

np.ndarray The full drift timecourse evaluated at the supplied timepoints t.

conflict_dsstimflex_drift `module-attribute`

conflict_dsstimflex_drift = conflict_dsstimflex_drift

Drift function for conflict task with stimuli with potentially variable onset.

The final drift is v + (combined drift timecourse).

Arguments:

t: np.ndarray
    Timepoints at which to evaluate the drift.
    Usually np.arange() of some sort.
v: float
    Base drift rate (typically 0 for conflict models). Defaults to 0.0.
tcoh: float
    Coherence of the target stimulus when 'on'.
dcoh: float
    Coherence of the distractor stimulus when 'on'.
tinit: float
    Initial condition of target drift timecourse.
dinit: float
    Initial condition of distractor drift timecourse.
tslope: float
    Slope parameter for target drift timecourse.
dslope: float
    Slope parameter for distractor drift timecourse.
tfixedp: float
    Fixed point for target drift timecourse.
tonset: float
    Onset time of the target stimulus coherence.
donset: float
    Onset time of the distractor stimulus coherence.
rel_first: bool
    If True, the first stimulus to appear (target or distractor)
    is treated as appearing at time 0, and the other stimulus
    is adjusted accordingly. If False, the onsets are treated
    as absolute times.

conflict_stimflex_drift `module-attribute`

conflict_stimflex_drift = conflict_stimflex_drift

Drift function for conflict task with stimuli with potentially variable onset and duration.

The final drift is v + (combined drift timecourse) when sum_drifts=True. When sum_drifts=False, returns 2D array for dual-drift models (v not added).

Arguments:

t: np.ndarray
    Timepoints at which to evaluate the drift.
    Usually np.arange() of some sort.
v: float
    Base drift rate (typically 0 for conflict models). Defaults to 0.0.
    Only used when sum_drifts=True.
tcoh: float
    Coherence of the target stimulus when 'on'.
dcoh: float
    Coherence of the distractor stimulus when 'on'.
vt: float
    Static drift-rate of target stimulus, when 'on'.
vd: float
    Static drift-rate of distractor stimulus, when 'on'.
tonset: float
    Onset time of the target stimulus coherence.
donset: float
    Onset time of the distractor stimulus coherence.
toffset, doffset: float or None
    Duration of the stimulus coherence pulse. If None, the pulse
    lasts until the end of the trial.
rel_first: bool
    If True, the first stimulus to appear (target or distractor)
    is treated as appearing at time 0, and the other stimulus
    is adjusted accordingly. If False, the onsets are treated
    as absolute times.
sum_drifts: bool
    If True, the drift contributions from target and distractor
    are summed to produce a single drift timecourse. If False,
    a 2D array is returned with separate columns for target
    and distractor drift timecourses (for dual-drift models).

Returns:

np.ndarray: Array of drift values, same length as t. If sum_drifts –

is False, the array has shape (len(t), 2)

constant `module-attribute`

constant = constant

Constant drift function.

Arguments

t: np.ndarray, optional
    Timepoints at which to evaluate the drift. Defaults to
    np.arange(0, 20, 0.1).
v: float, optional
    Drift rate. Defaults to 0.0.

Returns:

np.ndarray: Array of constant drift values = v, same length as t –

ds_support_analytic `module-attribute`

ds_support_analytic = ds_support_analytic

Solve DE.

DE is of the form: x' = slope*(fix_point - x), with initial condition init_p. The solution takes the form: (init_p - fix_point) * exp(-slope * t) + fix_point

Arguments

t: np.ndarray
    Timepoints at which to evaluate the drift. Usually np.arange() of some sort.
init_p: float
    Initial condition of dynamical system
fix_point: float
    Fixed point of dynamical system
slope: float
    Coefficient in exponent of the solution.

Return

np.ndarray The gamma drift evaluated at the supplied timepoints t.

gamma_drift `module-attribute`

gamma_drift = gamma_drift

Drift function that follows a scaled gamma distribution.

The final drift is v + (scaled gamma component).

Arguments

t: np.ndarray
    Timepoints at which to evaluate the drift.
    Usually np.arange() of some sort.
v: float
    Base drift rate. Defaults to 0.0.
shape: float
    Shape parameter of the gamma distribution
scale: float
    Scale parameter of the gamma distribution
c: float
    Scalar parameter that scales the peak of
    the gamma distribution.
    (Note this function follows a gamma distribution
    but does not integrate to 1)

Return

np.ndarray
    The final drift (v + gamma component) evaluated at the supplied timepoints t.

stimflex_support

stimflex_support(t, onset, offset, coh)

Construct a rectangular coherence timecourse, with discrete and potentially variable onsets and offsets of stimulus evidence.

Arguments

t: np.ndarray
    Timepoints within trial.
onset: float
    Onset time of the coherence pulse.
offset: float
    Offset time of the coherence pulse.
coh: float
    Coherence of the stimulus when 'on'.

Returns:

np.ndarray: Array of coherence values, same length as t. –

fixation_continuation

Pluggable fixation-continuation strategies + a positive-distribution factory for the aDDM.

The aDDM posterior-predictive path conditions on observed fixations (Mode 2). When a re-simulated particle has not decided by the last observed fixation, the drift "freezes" at the last gaze. This module makes that tail behaviour pluggable, and shares a scipy.stats positive-distribution factory with the aDDM's Mode-1 self-sampling.

Strategies (name -> callable), each mapping (node_r, d_r, flag_r, T, params, rng) -> (node_r, d_r, flag_r, max_d):

prolong_last_fixation -- no new fixations; the engine holds the last gaze to T (today's behaviour; the default).
sample_continuation -- keep the observed fixations; draw the continuation (tail) durations from a chosen positive distribution and keep alternating past the last observed onset. (Appending an onset after the last observed onset resamples the censored last fixation's duration; further onsets are new alternating fixations.)
resample_all_fixations -- ignore the observed fixation schedule; keep the observed stimulus (r1/r2 untouched upstream), re-sample the first-gaze flag and a fresh schedule. The drift still conditions on the stimulus, so the PPC stays comparable to data.

All durations come from :data:POSITIVE_DISTRIBUTIONS (scipy.stats), so any positive-support distribution is available by name with scipy-native parameters. Mirrors the attention_process registry idiom (name -> callable + a resolve_* helper).

draw_durations

draw_durations(dist_name, dist_params, size, rng)

Draw an array of size positive durations from dist_name (scipy-native params).

generate_schedule

generate_schedule(n, T, dist_name, dist_params, rng)

Self-sample a padded (n, max_d) fixation-onset schedule + per-row stage count d.

Durations are drawn from dist_name and cumulatively summed; column 0 is anchored at 0.0 (trial start). d = number of onsets strictly before T (>=1). The pre-draw budget is sized from the distribution's analytic mean, then over-allocated columns (onsets >= T) are trimmed to the realised max_d.

prolong_last_fixation

prolong_last_fixation(node_r, d_r, flag_r, T, params, rng)

No continuation: the engine holds the last observed gaze to T (today's behaviour).

resample_all_fixations

resample_all_fixations(node_r, d_r, flag_r, T, params, rng)

Ignore observed fixations; re-sample first-gaze + a fresh schedule (keep stimulus r1/r2).

resolve_continuation_mode

resolve_continuation_mode(mode)

Resolve a registry name or a callable to a concrete continuation strategy.

resolve_distribution

resolve_distribution(name)

Return the scipy.stats distribution registered under name.

sample_continuation

sample_continuation(node_r, d_r, flag_r, T, params, rng)

Keep observed fixations; draw the continuation (tail) durations from params['dist'].

inv_temp_softmax

Choice-only inverse-temperature softmax simulators.

inv_temp_softmax

inv_temp_softmax(
    *,
    beta,
    q0,
    q1,
    q2=None,
    q3=None,
    n_samples=1000,
    n_trials=1,
    max_t=20.0,
    random_state=None,
    **kwargs
)

Sample choices from softmax(beta * Q) with placeholder RTs.

modular_parameter_simulator_adapter

Modular parameter simulator adapter using adaptation pipelines.

This module provides the ModularParameterSimulatorAdapter class which applies parameter adaptations defined in model_config['parameter_transforms']['simulation'].

ModularParameterSimulatorAdapter

Modular parameter simulator adapter using config-defined transforms.

This adapter applies a sequence of parameter adaptations defined in the model configuration's parameter_transforms.simulation field.

All built-in models define their simulation transforms directly in their model config, making this adapter simple and transparent.

Examples:

>>> adapter = ModularParameterSimulatorAdapter()
>>> theta = adapter.adapt_parameters(theta, model_config, n_trials)

The model_config should have transforms defined like:

>>> model_config = {
...     "name": "lba_angle_3",
...     "parameter_transforms": {
...         "sampling": [...],
...         "simulation": [
...             ColumnStackParameters(["v0", "v1", "v2"], "v"),
...             ExpandDimension(["a", "z", "theta"]),
...             SetZeroArray("t"),
...         ],
...     },
... }

adapt_parameters

adapt_parameters(theta, model_config, n_trials)

Process theta by applying config-defined transformations.

Parameters:

theta (dict[str, Any]) –

Dictionary of theta parameters
model_config (dict[str, Any]) –

Model configuration dictionary containing parameter_transforms
n_trials (int) –

Number of trials

Returns:

dict[str, Any] –

Processed theta parameters

ParameterSimulatorAdapterProtocol

Bases: Protocol

Protocol defining the interface for parameter adapters.

Any class that implements an adapt_parameters method with this signature can be used as a parameter adapter.

adapt_parameters

adapt_parameters(theta, model_config, n_trials)

Adapt theta parameters for simulator consumption.

Parameters:

theta (dict[str, Any]) –

Dictionary of theta parameters
model_config (dict[str, Any]) –

Model configuration dictionary
n_trials (int) –

Number of trials

Returns:

dict[str, Any] –

Processed theta parameters

parameter_adapters

Parameter adaptation system for modular parameter processing.

This module provides a composable system for adapting parameters before they are passed to simulators. Adaptations are small, focused classes that can be combined to handle complex parameter preparation logic.

Main Components: - ParameterAdaptation: Alias for ParameterTransform (backward compatibility) - Common adaptations: SetDefaultValue, ExpandDimension, etc. - ParameterAdapterRegistry: Maps models to adaptation pipelines - ModularParameterSimulatorAdapter: Applies adaptations in sequence

Example: >>> from ssms.basic_simulators.parameter_adapters import ( ... ExpandDimension, ColumnStackParameters ... ) >>> >>> adaptations = [ ... ColumnStackParameters(["v0", "v1", "v2"], "v"), ... ExpandDimension(["a", "t"]) ... ] >>> >>> for adaptation in adaptations: ... theta = adaptation.apply(theta, model_config, n_trials)

Note: New code should import directly from ssms.transforms.simulation for the common adaptation classes. This module re-exports them for backward compatibility.

ApplyMapping

ApplyMapping(
    source_param,
    target_param,
    mapping_key,
    additional_sources=None,
    delete_source=False,
)

Bases: ParameterTransform

Apply a mapping function to transform a parameter.

Uses a mapping from model_config to transform a parameter value. Commonly used for random variable distributions.

Parameters:

source_param (str) –

Parameter to read from
target_param (str) –

Parameter to write to
mapping_key (str) –

Key in model_config["simulator_param_mappings"]
additional_sources (list[str], default: None ) –

Additional parameters to pass to the mapping function
delete_source (bool, default: False ) –

Whether to delete the source parameter after mapping

apply

apply(theta, model_config=None, n_trials=None)

Apply mapping function to transform parameter.

ColumnStackParameters

ColumnStackParameters(
    source_params, target_param, delete_sources=True
)

Bases: ParameterTransform

Stack multiple parameters into a single multi-column array.

Takes individual parameters (e.g., v0, v1, v2) and stacks them column-wise into a single array (e.g., v with shape (n_trials, 3)).

Parameters:

source_params (list[str]) –

Names of parameters to stack
target_param (str) –

Name of the resulting stacked parameter
delete_sources (bool, default: True ) –

Whether to delete the source parameters after stacking

Examples:

>>> transform = ColumnStackParameters(["v0", "v1", "v2"], "v")
>>> theta = {"v0": [0.5], "v1": [0.6], "v2": [0.7]}
>>> theta = transform.apply(theta)
>>> theta["v"]  # array([[0.5, 0.6, 0.7]])

apply

apply(theta, model_config=None, n_trials=None)

Stack parameters column-wise.

ConditionalAdaptation

ConditionalAdaptation(condition, transformation)

Bases: ParameterTransform

Apply a transform only if a condition is met.

Parameters:

condition (Callable) –

Function that takes (theta, model_config, n_trials) and returns bool
transformation (ParameterTransform) –

Transform to apply if condition is True

Examples:

>>> transform = ConditionalAdaptation(
...     condition=lambda theta, cfg, n: "a" in theta,
...     transformation=ExpandDimension(["a"])
... )

apply

apply(theta, model_config=None, n_trials=None)

Apply transform if condition is met.

DeleteParameters

DeleteParameters(param_names)

Bases: ParameterTransform

Delete specified parameters from theta.

Parameters:

param_names (list[str]) –

Names of parameters to delete

Examples:

>>> transform = DeleteParameters(["A", "b"])
>>> theta = {"v": [0.5], "A": [0.5], "b": [1.0]}
>>> theta = transform.apply(theta)
>>> "A" in theta  # False

apply

apply(theta, model_config=None, n_trials=None)

Delete specified parameters.

ExpandDimension

ExpandDimension(param_names)

Bases: ParameterTransform

Expand dimensions of specified parameters.

Converts 1D arrays of shape (n_trials,) to shape (n_trials, 1) by adding a dimension. This is commonly needed for multi-particle models.

Parameters:

param_names (list[str]) –

Names of parameters to expand

Examples:

>>> transform = ExpandDimension(["a", "t"])
>>> theta = {"a": np.array([1.0, 2.0]), "t": np.array([0.3, 0.3])}
>>> theta = transform.apply(theta)
>>> theta["a"].shape  # (2, 1)

apply

apply(theta, model_config=None, n_trials=None)

Expand dimensions of specified parameters.

LambdaAdaptation

LambdaAdaptation(func, name=None)

Bases: ParameterTransform

Wrap a lambda or callable function as a transform.

Parameters:

func (Callable) –

Function with signature: func(theta, model_config, n_trials) → theta
name (str, default: None ) –

Optional name for debugging

Examples:

>>> transform = LambdaAdaptation(
...     lambda theta, cfg, n: theta.update({"nact": 3}) or theta,
...     name="set_nact_to_3"
... )

apply

apply(theta, model_config=None, n_trials=None)

Apply the wrapped function.

ParameterAdapterRegistry

ParameterAdapterRegistry()

Registry mapping models to parameter adaptation pipelines.

The registry supports two types of lookups: 1. Exact model match: Direct lookup by model name 2. Family match: Pattern-based matching using lambda functions

Exact matches take precedence over family matches, allowing specific models to override family defaults.

Examples:

>>> registry = ParameterAdapterRegistry()
>>>
>>> # Register exact model
>>> registry.register_model("lba2", [
...     SetDefaultValue("nact", 2),
...     ColumnStackParameters(["v0", "v1"], "v")
... ])
>>>
>>> # Register model family
>>> registry.register_family(
...     "race_2",
...     lambda m: m.startswith("race_") and m.endswith("_2"),
...     [ColumnStackParameters(["v0", "v1"], "v")]
... )
>>>
>>> # Lookup
>>> adaptations = registry.get_processor("lba2")  # Exact match
>>> adaptations = registry.get_processor("race_no_bias_2")  # Family match

repr

__repr__()

String representation of registry.

describe

describe(model_name)

Get human-readable description of adaptations for a model.

Parameters:

model_name (str) –

Name of the model

Returns:

str –

Description of the adaptation pipeline

Examples:

>>> print(registry.describe("lba3"))
Model: lba3
Adaptations:
  1. SetDefaultValue(param_name='nact', default_value=3)
  2. ColumnStackParameters(source_params=['v0', 'v1', 'v2'], ...)

get_processor

get_processor(model_name)

Get adaptation pipeline for a model.

Lookup priority: 1. Exact model name match 2. First matching family 3. Empty list (no adaptations)

Parameters:

model_name (str) –

Name of the model to look up

Returns:

list[ParameterAdaptation] –

List of adaptations to apply (empty list if no match)

Examples:

>>> adaptations = registry.get_processor("lba2")
>>> # Returns registered adaptations for "lba2"
>>>
>>> adaptations = registry.get_processor("race_no_bias_2")
>>> # Returns adaptations for "race_2" family (if registered)
>>>
>>> adaptations = registry.get_processor("unknown_model")
>>> # Returns [] (empty list)

has_processor

has_processor(model_name)

Check if a processor is registered for a model.

This checks if the model has a registration (exact or family match), even if the transformation list is empty.

Parameters:

model_name (str) –

Name of the model to check

Returns:

bool –

True if model has a registration (explicit or via family)

list_registered_families

list_registered_families()

Get list of registered family names.

Returns:

list[str] –

List of family names

list_registered_models

list_registered_models()

Get list of explicitly registered model names.

Returns:

list[str] –

List of model names with exact registrations

register_family

register_family(family_name, matcher, adaptations)

Register adaptations for a family of models.

The matcher function determines which models belong to this family. It receives a model name and should return True if the model matches.

Parameters:

family_name (str) –

Name for this model family (used as key)
matcher (Callable[[str], bool]) –

Function that takes a model name and returns True if it matches
adaptations (list[ParameterAdaptation]) –

List of adaptations to apply for matching models

Examples:

>>> registry = ParameterAdapterRegistry()
>>>
>>> # Match all race models with 2 choices
>>> registry.register_family(
...     "race_2",
...     lambda m: m.startswith("race_") and m.endswith("_2"),
...     [ColumnStackParameters(["v0", "v1"], "v")]
... )
>>>
>>> # Match all models with "no_bias" in the name
>>> registry.register_family(
...     "no_bias_models",
...     lambda m: "no_bias" in m,
...     [SetDefaultValue("z", 0.5)]
... )

register_model

register_model(model_name, adaptations)

Register adaptations for a specific model.

Parameters:

model_name (str) –

Exact model name (e.g., "lba2", "ddm", "race_3")
adaptations (list[ParameterAdaptation]) –

List of adaptations to apply for this model

Examples:

>>> registry = ParameterAdapterRegistry()
>>> registry.register_model("lba3", [
...     SetDefaultValue("nact", 3),
...     ColumnStackParameters(["v0", "v1", "v2"], "v")
... ])

RenameParameter

RenameParameter(old_name, new_name, transform_fn=None)

Bases: ParameterTransform

Rename a parameter and optionally transform its value.

Parameters:

old_name (str) –

Current parameter name
new_name (str) –

New parameter name
transform_fn (Callable or None, default: None ) –

Optional function to apply to the value during renaming

Examples:

>>> transform = RenameParameter("A", "z", lambda x: np.expand_dims(x, axis=1))
>>> theta = {"A": np.array([0.5])}
>>> theta = transform.apply(theta)
>>> "z" in theta  # True

apply

apply(theta, model_config=None, n_trials=None)

Rename parameter and optionally transform its value.

SetDefaultValue

SetDefaultValue(param_name, default_value, dtype=float32)

Bases: ParameterTransform

Set a parameter to a default value if not present.

This transform adds a parameter with a default value if it doesn't already exist in theta. The value is automatically tiled to match n_trials.

Parameters:

param_name (str) –

Name of the parameter to set
default_value (float or int or ndarray) –

Default value to use. If scalar, will be tiled to (n_trials,).
dtype (dtype, default: float32 ) –

Data type for the parameter. Defaults to np.float32.

Examples:

>>> transform = SetDefaultValue("nact", 3)
>>> theta = {}
>>> theta = transform.apply(theta, {}, n_trials=5)
>>> theta["nact"]  # array([3., 3., 3., 3., 3.], dtype=float32)

apply

apply(theta, model_config=None, n_trials=None)

Add parameter with default value if not present.

SetZeroArray

SetZeroArray(param_name, shape=None, dtype=float32)

Bases: ParameterTransform

Set a parameter to an array of zeros.

Parameters:

param_name (str) –

Name of the parameter to set
shape (tuple or None, default: None ) –

Shape of the array. Use None for n_trials dimension.
dtype (dtype, default: float32 ) –

Data type for the array. Defaults to np.float32.

Examples:

>>> transform = SetZeroArray("v")
>>> theta = transform.apply({}, {}, n_trials=5)
>>> theta["v"].shape  # (5,)

apply

apply(theta, model_config=None, n_trials=None)

Set parameter to zero array.

TileArray

TileArray(param_name, value, dtype=float32)

Bases: ParameterTransform

Tile a constant value across trials.

Parameters:

param_name (str) –

Name of the parameter to create
value (float or array - like) –

Value(s) to tile
dtype (dtype, default: float32 ) –

Data type for the array. Defaults to np.float32.

Examples:

>>> transform = TileArray("z", 0.5)
>>> theta = transform.apply({}, {}, n_trials=3)
>>> theta["z"]  # array([0.5, 0.5, 0.5], dtype=float32)

apply

apply(theta, model_config=None, n_trials=None)

Tile value across trials.

base

Base classes for parameter adaptations.

This module provides backward compatibility by re-exporting ParameterTransform as ParameterAdaptation.

Note: New code should use ParameterTransform from ssms.transforms.base directly.

ParameterTransform

Bases: ABC

Abstract base class for all parameter transforms.

A parameter transform is a single, focused operation that modifies theta parameters. Transforms can be composed into pipelines and are used at two phases:

Sampling phase: Applied after parameter sampling to enforce constraints (e.g., ensure a > z by swapping values)
Simulation phase: Applied before calling the simulator to prepare parameters for the expected format (e.g., stack v0, v1 → v)

Subclasses must implement the apply method. The method signature includes optional model_config and n_trials arguments to support both simple constraints (which only need theta) and complex adapters (which may need additional context).

Examples:

Create a custom transform:

>>> class ScaleParameter(ParameterTransform):
...     def __init__(self, param_name: str, scale: float):
...         self.param_name = param_name
...         self.scale = scale
...
...     def apply(self, theta, model_config=None, n_trials=None):
...         if self.param_name in theta:
...             theta[self.param_name] = theta[self.param_name] * self.scale
...         return theta

Use in model config:

>>> model_config = {
...     "name": "my_model",
...     "parameter_transforms": {
...         "sampling": [SwapIfLessConstraint("a", "z")],
...         "simulation": [ColumnStackParameters(["v0", "v1"], "v")],
...     }
... }

repr

__repr__()

String representation for debugging.

Returns:

str –

A string describing this transform.

str

__str__()

Human-readable string representation.

apply `abstractmethod`

apply(theta, model_config=None, n_trials=None)

Apply transform to theta parameters.

This method should modify the theta dictionary in place and return it. It may add new parameters, modify existing parameters, or remove parameters as needed.

Parameters:

theta (dict[str, Any]) –

Dictionary of model parameters. Values are typically numpy arrays.
model_config (dict[str, Any] or None, default: None ) –

Model configuration dictionary. Available for transforms that need additional context (e.g., looking up simulator_fixed_params). Sampling-time transforms typically ignore this.
n_trials (int or None, default: None ) –

Number of trials. Available for transforms that need to create arrays of a specific size. Sampling-time transforms typically ignore this.

Returns:

dict[str, Any] –

The modified theta dictionary (usually the same object passed in).

Notes

Transforms should be pure when possible (no side effects)
If creating new arrays, use dtype=np.float32 for consistency
Document any parameters that are added, modified, or removed

get_adapter_registry

get_adapter_registry()

Get the global parameter adapter registry.

Use this to access registry methods like list_registered_models() or get_processor() directly.

Returns:

ParameterAdapterRegistry –

The global parameter adapter registry instance, pre-populated with all built-in model adaptations

Examples:

>>> from ssms.basic_simulators.parameter_adapters import get_adapter_registry
>>>
>>> # List all models with registered adaptations
>>> registry = get_adapter_registry()
>>> print(registry.list_registered_models())
['lba2', 'lba3', 'lca_3', ...]
>>>
>>> # Check if a model has adaptations
>>> if registry.has_processor("my_model"):
...     adaptations = registry.get_processor("my_model")
>>>
>>> # Describe adaptations for a model
>>> print(registry.describe("lba3"))

register_adapter_to_model

register_adapter_to_model(model_name, adaptations)

Register parameter adaptations to a specific model globally.

This associates existing parameter adaptation classes with a model. Once registered, the model will use these adaptations automatically when simulated with ModularParameterSimulatorAdapter.

Note: This registers which adaptations a MODEL uses, not a new adapter type. To create custom adapter types, subclass ParameterAdaptation directly.

Parameters:

model_name (str) –

Unique model name
adaptations (list[ParameterAdaptation]) –

List of parameter adaptation instances to apply for this model

Examples:

Register adaptations for a custom model:

>>> from ssms.basic_simulators.parameter_adapters import register_adapter_to_model
>>> from ssms.basic_simulators.parameter_adapters import (
...     SetDefaultValue, ExpandDimension, ColumnStackParameters
... )
>>>
>>> register_adapter_to_model("my_model", [
...     SetDefaultValue("z", 0.5),
...     ExpandDimension(["a", "t"]),
...     ColumnStackParameters(["v0", "v1"], "v"),
... ])
>>>
>>> # Now when simulating with this model, adaptations are applied automatically
>>> sim = Simulator(model="my_model")

Override built-in model adaptations:

>>> # Register custom adaptations for an existing model
>>> register_adapter_to_model("lba2", [
...     ColumnStackParameters(["v0", "v1"], "v"),
...     # Custom adaptations here
... ])

register_adapter_to_model_family

register_adapter_to_model_family(
    family_name, matcher, adaptations
)

Register parameter adaptations to a family of models globally.

This associates existing parameter adaptation classes with multiple models matching a pattern, without having to register each model individually.

Parameters:

family_name (str) –

Name for this model family (used as identifier)
matcher (Callable[[str], bool]) –

Function that returns True if a model name matches this family
adaptations (list[ParameterAdaptation]) –

List of parameter adaptation instances to apply for matching models

Examples:

Register adaptations for all race models with 2 alternatives:

>>> from ssms.basic_simulators.parameter_adapters import register_adapter_to_model_family
>>> from ssms.basic_simulators.parameter_adapters import ColumnStackParameters
>>>
>>> register_adapter_to_model_family(
...     "race_2",
...     lambda m: m.startswith("race_") and m.endswith("_2"),
...     [ColumnStackParameters(["v0", "v1"], "v")]
... )
>>>
>>> # Now all models matching the pattern use these adaptations
>>> # e.g., "race_no_bias_2", "race_no_z_2", etc.

Register for models with specific naming pattern:

>>> register_adapter_to_model_family(
...     "no_bias_models",
...     lambda m: "no_bias" in m,
...     [SetDefaultValue("z", 0.5)]
... )

registry

Registry system for mapping models to parameter adaptations.

The registry provides a centralized way to associate model names with their required parameter adaptations. It supports both exact model matching and pattern-based family matching.

ParameterAdapterRegistry

ParameterAdapterRegistry()

Registry mapping models to parameter adaptation pipelines.

The registry supports two types of lookups: 1. Exact model match: Direct lookup by model name 2. Family match: Pattern-based matching using lambda functions

Exact matches take precedence over family matches, allowing specific models to override family defaults.

Examples:

>>> registry = ParameterAdapterRegistry()
>>>
>>> # Register exact model
>>> registry.register_model("lba2", [
...     SetDefaultValue("nact", 2),
...     ColumnStackParameters(["v0", "v1"], "v")
... ])
>>>
>>> # Register model family
>>> registry.register_family(
...     "race_2",
...     lambda m: m.startswith("race_") and m.endswith("_2"),
...     [ColumnStackParameters(["v0", "v1"], "v")]
... )
>>>
>>> # Lookup
>>> adaptations = registry.get_processor("lba2")  # Exact match
>>> adaptations = registry.get_processor("race_no_bias_2")  # Family match

repr

__repr__()

String representation of registry.

describe

describe(model_name)

Get human-readable description of adaptations for a model.

Parameters:

model_name (str) –

Name of the model

Returns:

str –

Description of the adaptation pipeline

Examples:

>>> print(registry.describe("lba3"))
Model: lba3
Adaptations:
  1. SetDefaultValue(param_name='nact', default_value=3)
  2. ColumnStackParameters(source_params=['v0', 'v1', 'v2'], ...)

get_processor

get_processor(model_name)

Get adaptation pipeline for a model.

Lookup priority: 1. Exact model name match 2. First matching family 3. Empty list (no adaptations)

Parameters:

model_name (str) –

Name of the model to look up

Returns:

list[ParameterAdaptation] –

List of adaptations to apply (empty list if no match)

Examples:

>>> adaptations = registry.get_processor("lba2")
>>> # Returns registered adaptations for "lba2"
>>>
>>> adaptations = registry.get_processor("race_no_bias_2")
>>> # Returns adaptations for "race_2" family (if registered)
>>>
>>> adaptations = registry.get_processor("unknown_model")
>>> # Returns [] (empty list)

has_processor

has_processor(model_name)

Check if a processor is registered for a model.

This checks if the model has a registration (exact or family match), even if the transformation list is empty.

Parameters:

model_name (str) –

Name of the model to check

Returns:

bool –

True if model has a registration (explicit or via family)

list_registered_families

list_registered_families()

Get list of registered family names.

Returns:

list[str] –

List of family names

list_registered_models

list_registered_models()

Get list of explicitly registered model names.

Returns:

list[str] –

List of model names with exact registrations

register_family

register_family(family_name, matcher, adaptations)

Register adaptations for a family of models.

The matcher function determines which models belong to this family. It receives a model name and should return True if the model matches.

Parameters:

family_name (str) –

Name for this model family (used as key)
matcher (Callable[[str], bool]) –

Function that takes a model name and returns True if it matches
adaptations (list[ParameterAdaptation]) –

List of adaptations to apply for matching models

Examples:

>>> registry = ParameterAdapterRegistry()
>>>
>>> # Match all race models with 2 choices
>>> registry.register_family(
...     "race_2",
...     lambda m: m.startswith("race_") and m.endswith("_2"),
...     [ColumnStackParameters(["v0", "v1"], "v")]
... )
>>>
>>> # Match all models with "no_bias" in the name
>>> registry.register_family(
...     "no_bias_models",
...     lambda m: "no_bias" in m,
...     [SetDefaultValue("z", 0.5)]
... )

register_model

register_model(model_name, adaptations)

Register adaptations for a specific model.

Parameters:

model_name (str) –

Exact model name (e.g., "lba2", "ddm", "race_3")
adaptations (list[ParameterAdaptation]) –

List of adaptations to apply for this model

Examples:

>>> registry = ParameterAdapterRegistry()
>>> registry.register_model("lba3", [
...     SetDefaultValue("nact", 3),
...     ColumnStackParameters(["v0", "v1", "v2"], "v")
... ])

get_adapter_registry

get_adapter_registry()

Get the global parameter adapter registry.

Use this to access registry methods like list_registered_models() or get_processor() directly.

Returns:

ParameterAdapterRegistry –

The global parameter adapter registry instance, pre-populated with all built-in model adaptations

Examples:

>>> from ssms.basic_simulators.parameter_adapters import get_adapter_registry
>>>
>>> # List all models with registered adaptations
>>> registry = get_adapter_registry()
>>> print(registry.list_registered_models())
['lba2', 'lba3', 'lca_3', ...]
>>>
>>> # Check if a model has adaptations
>>> if registry.has_processor("my_model"):
...     adaptations = registry.get_processor("my_model")
>>>
>>> # Describe adaptations for a model
>>> print(registry.describe("lba3"))

register_adapter_to_model

register_adapter_to_model(model_name, adaptations)

Register parameter adaptations to a specific model globally.

This associates existing parameter adaptation classes with a model. Once registered, the model will use these adaptations automatically when simulated with ModularParameterSimulatorAdapter.

Note: This registers which adaptations a MODEL uses, not a new adapter type. To create custom adapter types, subclass ParameterAdaptation directly.

Parameters:

model_name (str) –

Unique model name
adaptations (list[ParameterAdaptation]) –

List of parameter adaptation instances to apply for this model

Examples:

Register adaptations for a custom model:

>>> from ssms.basic_simulators.parameter_adapters import register_adapter_to_model
>>> from ssms.basic_simulators.parameter_adapters import (
...     SetDefaultValue, ExpandDimension, ColumnStackParameters
... )
>>>
>>> register_adapter_to_model("my_model", [
...     SetDefaultValue("z", 0.5),
...     ExpandDimension(["a", "t"]),
...     ColumnStackParameters(["v0", "v1"], "v"),
... ])
>>>
>>> # Now when simulating with this model, adaptations are applied automatically
>>> sim = Simulator(model="my_model")

Override built-in model adaptations:

>>> # Register custom adaptations for an existing model
>>> register_adapter_to_model("lba2", [
...     ColumnStackParameters(["v0", "v1"], "v"),
...     # Custom adaptations here
... ])

register_adapter_to_model_family

register_adapter_to_model_family(
    family_name, matcher, adaptations
)

Register parameter adaptations to a family of models globally.

This associates existing parameter adaptation classes with multiple models matching a pattern, without having to register each model individually.

Parameters:

family_name (str) –

Name for this model family (used as identifier)
matcher (Callable[[str], bool]) –

Function that returns True if a model name matches this family
adaptations (list[ParameterAdaptation]) –

List of parameter adaptation instances to apply for matching models

Examples:

Register adaptations for all race models with 2 alternatives:

>>> from ssms.basic_simulators.parameter_adapters import register_adapter_to_model_family
>>> from ssms.basic_simulators.parameter_adapters import ColumnStackParameters
>>>
>>> register_adapter_to_model_family(
...     "race_2",
...     lambda m: m.startswith("race_") and m.endswith("_2"),
...     [ColumnStackParameters(["v0", "v1"], "v")]
... )
>>>
>>> # Now all models matching the pattern use these adaptations
>>> # e.g., "race_no_bias_2", "race_no_z_2", etc.

Register for models with specific naming pattern:

>>> register_adapter_to_model_family(
...     "no_bias_models",
...     lambda m: "no_bias" in m,
...     [SetDefaultValue("z", 0.5)]
... )

simulator

This module defines the basic simulator function which is the main workshorse of the package. In addition some utility functions are provided that help with preprocessing the output of the simulator function.

bin_simulator_output

bin_simulator_output(
    out=None, bin_dt=0.04, nbins=0, max_t=-1, freq_cnt=False
)

Turns RT part of simulator output into bin-identifier by trial

Arguments

out : dict
    Output of the 'simulator' function
bin_dt : float
    If nbins is 0, this determines the desired
    bin size which in turn automatically
    determines the resulting number of bins.
nbins : int
    Number of bins to bin reaction time data into.
    If supplied as 0, bin_dt instead determines the number of
    bins automatically.
max_t : float <default=-1>
    Override the 'max_t' metadata as part of the simulator output.
    Sometimes useful, but usually default will do the job.
freq_cnt : bool <default=False>
    Decide whether to return proportions (default) or counts in bins.

Returns:

A histogram of counts or proportions. –

bin_simulator_output_pointwise

bin_simulator_output_pointwise(
    out=None, bin_dt=0.04, nbins=0
)

Turns RT part of simulator output into bin-identifier by trial

Arguments

out: dict
    Output of the 'simulator' function
bin_dt: float
    If nbins is 0, this determines the desired
    bin size which in turn automatically
    determines the resulting number of bins.
nbins: int
    Number of bins to bin reaction time data into.
    If supplied as 0, bin_dt instead determines the
    number of bins automatically.

Returns:

2d array. The first columns collects bin-identifiers –

by trial, the second column lists the corresponding choices.

make_boundary_dict

make_boundary_dict(config, theta)

Create a dictionary containing boundary-related parameters and functions.

This function extracts boundary-related parameters from the input theta dictionary, based on the boundary configuration specified in the config. It also retrieves the appropriate boundary function from the boundary registry.

Args: config (dict): A dictionary containing model configuration, including the boundary name. theta (dict): A dictionary of parameter values, potentially including boundary-related parameters.

Returns: dict: A dictionary containing: - boundary_params (dict): Extracted boundary-related parameters (including 'a'). - boundary_fun (callable): The boundary function corresponding to the specified boundary name.

make_drift_dict

make_drift_dict(config, theta)

Create a dictionary containing drift-related parameters and functions.

This function extracts drift-related parameters from the input theta dictionary, based on the drift configuration specified in the config. It also retrieves the appropriate drift function from the drift registry.

Args: config (dict): A dictionary containing model configuration, including the drift name. theta (dict): A dictionary of parameter values, potentially including drift-related parameters.

Returns: dict: A dictionary containing: - drift_fun (callable): The drift function corresponding to the specified drift name. - drift_params (dict): Extracted drift-related parameters. If no drift name is specified in config, returns an empty dictionary.

make_noise_vec

make_noise_vec(sigma_noise, n_trials, n_particles)

Create a noise vector for simulation.

This function generates a noise vector based on the provided sigma_noise parameter, replicating it across trials and particles as needed.

Arguments

sigma_noise (float | np.ndarray): Standard deviation of the noise. Can be a single float value
    or a numpy array.
n_trials (int): Number of trials to simulate.
n_particles (int): Number of particles per trial.

Returns:

np.ndarray: A noise vector with appropriate shape for the simulation, –

containing the replicated sigma_noise values.

simulator

simulator(
    theta,
    model="angle",
    n_samples=1000,
    delta_t=0.001,
    max_t=20,
    no_noise=False,
    sigma_noise=None,
    smooth_unif=True,
    random_state=None,
    return_option="full",
    n_threads=1,
    extra_fields=None,
)

Basic data simulator for the models included in HDDM.

Arguments

theta : list, numpy.array, dict or pd.DataFrame
    Parameters of the simulator. If 2d array, each row is treated as a 'trial'
    and the function runs n_sample * n_trials simulations.
deadline : numpy.array <default=None>
    If supplied, the simulator will run a deadline model. RTs will be returned
model: str <default='angle'>
    Determines the model that will be simulated.
n_samples: int <default=1000>
    Number of simulation runs for each row in the theta argument.
delta_t: float
    Size fo timesteps in simulator (conceptually measured in seconds)
max_t: float
    Maximum reaction the simulator can reach
no_noise: bool <default=False>
    Turn noise of (useful for plotting purposes mostly)
sigma_noise: float | None <default=None>
    Standard deviation of noise in the diffusion process. If None, defaults to 1.0 for most models
    and 0.1 for LBA models. If no_noise is True, sigma_noise will be set to 0.0.
    If 'sd' or 's' is passed via theta dictionary, sigma_noise must be None.
smooth_unif: bool <default=True>
    Whether to add uniform random noise to RTs to smooth the distributions.
random_state: int | None <default=None>
    Integer passed to the C-level RNG seeding.  Must lie in
    ``[-2**31, 2**31 - 1]`` so it fits in a 32-bit signed C ``long`` (required
    on Windows).  ``None`` draws a seed from that range automatically.
    Non-integer RNG objects may be supported on specific code paths.
return_option: str <default='full'>
    Determines what the function returns. Can be either
    'full' or 'minimal'. If 'full' the function returns
    a dictionary with keys 'rts', 'responses' and 'metadata', and
    metadata contains the model parameters and some additional
    information. 'metadata' is a simpler dictionary with less information
    if 'minimal' is chosen.
n_threads: int <default=1>
    Number of threads for parallel execution. If > 1 and OpenMP is available,
    uses multi-threaded simulation. Note: trajectory recording is only
    available with n_threads=1. For models that don't support n_threads yet,
    this parameter is silently ignored.

    **Reproducibility note:** the sequential path (n_threads=1) uses NumPy's
    RNG, while the parallel path (n_threads>1) uses a C-level GSL/Ziggurat
    RNG. Results with a fixed ``random_state`` are reproducible *within* a
    given ``n_threads`` value, but will differ *across* n_threads values.
    Both paths produce statistically equivalent distributions.

Return

dictionary where keys can be (rts, responses, metadata) or (rt-response histogram, metadata) or (rts binned pointwise, responses, metadata)

validate_ssm_parameters

validate_ssm_parameters(model, theta)

Validate the parameters for Sequential Sampling Models (SSM).

This function checks the validity of parameters for different SSM models. It performs specific checks based on the model type.

Args: model (str): The name of the SSM model. theta (dict): A dictionary containing the model parameters.

Raises: ValueError: If any of the parameter validations fail.

simulator_class

Class-based interface for Sequential Sampling Model simulations.

This module provides a modern, object-oriented interface for SSM simulations that supports custom boundary functions, drift functions, and even fully custom simulator implementations.

Simulator

Simulator(
    model=None,
    boundary=None,
    drift=None,
    simulator_function=None,
    parameter_adapter=None,
    parameter_adaptations=None,
    **config_overrides
)

Class-based interface for Sequential Sampling Model simulations.

This class provides a flexible, extensible way to configure and run SSM simulations. It supports: - Pre-defined models (via string names like "ddm", "angle", etc.) - Custom boundary and drift functions with existing simulators - Fully custom simulator functions - Configuration overrides for any model parameter

Examples:

Basic usage with pre-defined model:

>>> sim = Simulator("ddm")
>>> results = sim.simulate(theta={'v': 0.5, 'a': 1.0, 'z': 0.5, 't': 0.3})

Custom boundary function:

>>> def my_boundary(t, theta, scale):
...     return scale * np.sin(theta * t)
>>> sim = Simulator("ddm", boundary=my_boundary, boundary_params=["theta", "scale"])
>>> results = sim.simulate(theta={'v': 0.5, 'a': 1.0, 'z': 0.5, 't': 0.3,
...                               'theta': 0.5, 'scale': 1.0})

Fully custom simulator:

>>> def my_sim(v, a, z, t, max_t=20, n_samples=1000, **kwargs):
...     rts = np.random.exponential(1/abs(v), n_samples) + t
...     choices = np.where(np.random.random(n_samples) < z, 1, -1)
...     return {'rts': rts, 'choices': choices,
...             'metadata': {'model': 'custom', 'n_samples': n_samples}}
>>> sim = Simulator(simulator_function=my_sim, params=["v", "a", "z", "t"],
...                 nchoices=2)
>>> results = sim.simulate(theta={'v': 0.5, 'a': 1.0, 'z': 0.5, 't': 0.3})

Attributes:

config (dict) –

The full model configuration dictionary

Parameters:

model (str, dict, or None, default: None ) –

Either a model name (e.g., "ddm", "angle"), a full configuration dictionary, or None if providing a custom simulator_function.
boundary (str, Callable, or None, default: None ) –

Boundary function. Can be: - A string name from the boundary registry (e.g., "angle", "weibull_cdf") - A callable function with signature: func(t, **params) -> np.ndarray - None to use the model's default boundary
drift (str, Callable, or None, default: None ) –

Drift function. Can be: - A string name from the drift registry (e.g., "constant", "gamma_drift") - A callable function with signature: func(t, **params) -> np.ndarray - None to use the model's default drift (if applicable)
simulator_function (Callable or None, default: None ) –

Custom simulator function. If provided, this overrides the model's default simulator. Must follow the simulator interface contract.
parameter_adapter (ParameterSimulatorAdapterProtocol or None, default: None ) –

Custom parameter adapter for preparing parameters for simulators. If None, uses ModularParameterSimulatorAdapter with default adaptations for the model.
parameter_adaptations (list[ParameterAdaptation] or None, default: None ) –

Additional parameter adaptations to apply AFTER the model's default adaptations. Useful for adding custom parameter preparation without replacing the entire adapter. Only used if parameter_adapter is None.
**config_overrides –

Additional configuration parameters to override. Common options: - params : list[str] - Parameter names (required if simulator_function provided) - param_bounds : list - [[lower bounds], [upper bounds]] - nchoices : int - Number of choices (required if simulator_function provided) - choices : list - Possible choice values - boundary_params : list[str] - Parameters for custom boundary function - drift_params : list[str] - Parameters for custom drift function - name : str - Model name (for custom simulators)

Raises:

ValueError –

If configuration is invalid or missing required parameters

Examples:

Use default modular adapter:

>>> sim = Simulator("lba2")
>>> results = sim.simulate(theta={'v0': 0.5, 'v1': 0.6, 'A': 0.5, 'b': 1.0})

Add custom adaptations:

>>> from ssms.basic_simulators.parameter_adapters import SetDefaultValue
>>> sim = Simulator("ddm", parameter_adaptations=[SetDefaultValue("custom_param", 42)])

Use a custom parameter adapter:

>>> custom_adapter = ModularParameterSimulatorAdapter()
>>> sim = Simulator("ddm", parameter_adapter=custom_adapter)

config `property`

config

Get the full model configuration.

Returns:

dict –

Configuration dictionary

parameter_adapter `property`

parameter_adapter

Get the configured parameter adapter.

Returns:

ParameterSimulatorAdapterProtocol –

The parameter adapter instance used for parameter preparation

simulate

simulate(
    theta,
    n_samples=1000,
    delta_t=0.001,
    max_t=20,
    no_noise=False,
    sigma_noise=None,
    smooth_unif=True,
    random_state=None,
    return_option="full",
    n_threads=1,
    extra_fields=None,
)

Run simulation with given parameters.

Parameters:

theta (list, np.ndarray, dict, or pd.DataFrame) –

Model parameters. If dict or DataFrame, keys/columns should match parameter names in config. If array, order should match config['params'].
extra_fields (dict or None, default: None ) –

Optional per-trial covariates forwarded to the model function (e.g. the aDDM fixations r1, r2, flag, sacc_array, d, sigma), so a covariate-driven model can condition on observed inputs rather than self-sampling them. Empty/None is a no-op for models that take no covariates. Arrays are per-trial (aligned with theta's trial axis).
n_samples (int, default: 1000 ) –

Number of simulation samples per parameter set
delta_t (float, default: 0.001 ) –

Time step size for simulation
max_t (float, default: 20 ) –

Maximum simulation time
no_noise (bool, default: False ) –

If True, disable noise in simulation
sigma_noise (float or None, default: None ) –

Standard deviation of noise. If None, uses model defaults.
smooth_unif (bool, default: True ) –

Whether to add uniform smoothing to RTs
random_state (int or None, default: None ) –

Integer seed for the C-level RNG. Must lie in [-2**31, 2**31 - 1] (fits a 32-bit signed C long, required on Windows). None picks a seed in [0, 2**31 - 1] automatically.
return_option (str, default: "full" ) –

Output format: "full" or "minimal"
n_threads (int, default: 1 ) –

Number of threads for parallel execution. If > 1 and OpenMP is available, uses multi-threaded simulation. Note: trajectory recording is only available with n_threads=1.

Returns:

dict –

Simulation results with keys: - 'rts' : np.ndarray of reaction times - 'choices' : np.ndarray of choices - 'metadata' : dict with simulation metadata - Additional keys depending on model and return_option

Raises:

ValueError –

If parameters are invalid

validate_params

validate_params(theta)

Validate parameter values.

Parameters:

theta (dict) –

Parameter dictionary to validate

Raises:

ValueError –

If parameters are invalid

ssms.config

Configuration module for SSM simulators.

This module provides access to model configurations, boundary and drift function configurations, and various generator configurations used throughout the SSMS package. It centralizes all configuration-related functionality to ensure consistent parameter settings across simulations.

CopyOnAccessDict

Bases: dict

A dict that returns a deep copy of the value on lookup.

ModelConfigBuilder

Helper class for building custom model configurations.

This class provides static methods for creating model configurations in various ways: - Starting from an existing model and overriding specific values - Building a configuration from scratch for fully custom simulators - Creating minimal valid configurations - Validating configurations

Examples:

Start from existing model and override:

>>> config = ModelConfigBuilder.from_model("ddm",
...                                     param_bounds=[[-4, 0.3, 0.1, 0],
...                                                   [4, 3.0, 0.9, 2.0]])

Build from scratch:

>>> config = ModelConfigBuilder.from_scratch(
...     name="my_model",
...     params=["v", "a", "z", "t"],
...     simulator_function=my_sim_fn,
...     nchoices=2
... )

Create minimal configuration:

>>> config = ModelConfigBuilder.minimal_config(
...     params=["v", "a"],
...     simulator_function=my_sim_fn,
...     nchoices=2
... )

add_boundary `staticmethod`

add_boundary(config, boundary, boundary_params=None)

Add or replace boundary function in configuration.

Parameters:

config (dict) –

Configuration to modify
boundary (str or Callable) –

Boundary function name or callable
boundary_params (list[str] or None, default: None ) –

Parameter names for boundary function (required if boundary is callable)
multiplicative (bool, default: True ) –

Whether boundary is multiplicative (True) or additive (False)

Returns:

dict –

Modified configuration (note: modifies in place and returns)

Raises:

ValueError –

If boundary specification is invalid

Examples:

>>> config = ModelConfigBuilder.from_model("ddm")
>>> config = ModelConfigBuilder.add_boundary(config, "angle")

add_drift `staticmethod`

add_drift(config, drift, drift_params=None)

Add or replace drift function in configuration.

Parameters:

config (dict) –

Configuration to modify
drift (str or Callable) –

Drift function name or callable
drift_params (list[str] or None, default: None ) –

Parameter names for drift function (required if drift is callable)

Returns:

dict –

Modified configuration (note: modifies in place and returns)

Raises:

ValueError –

If drift specification is invalid

Examples:

>>> config = ModelConfigBuilder.from_model("ddm")
>>> config = ModelConfigBuilder.add_drift(config, "gamma_drift")

from_model `staticmethod`

from_model(model_name, **overrides)

Create configuration starting from an existing model.

This method automatically parses variant suffixes like "_deadline" from the model name and applies the appropriate transformations.

Parameters:

model_name (str) –

Name of the model, optionally with variant suffixes. Examples: "ddm", "angle", "ddm_deadline", "angle_deadline"
**overrides –

Configuration fields to override. Common options: - params : list[str] - Parameter names - param_bounds : list - [[lower bounds], [upper bounds]] - boundary : Callable - Custom boundary function - boundary_name : str - Boundary name - boundary_params : list[str] - Boundary parameter names - drift : Callable - Custom drift function - drift_name : str - Drift name - drift_params : list[str] - Drift parameter names - simulator : Callable - Custom simulator function - nchoices : int - Number of choices - choices : list - Possible choice values

Returns:

dict –

Configuration dictionary with any variants applied

Raises:

ValueError –

If the base model name is not recognized

Examples:

>>> config = ModelConfigBuilder.from_model("ddm",
...                                     param_bounds=[[-4, 0.3, 0.1, 0],
...                                                   [4, 3.0, 0.9, 2.0]])

>>> # With deadline variant
>>> config = ModelConfigBuilder.from_model("ddm_deadline")
>>> "deadline" in config["params"]
True

from_scratch `staticmethod`

from_scratch(
    name, params, simulator_function, nchoices, **config
)

Build a complete configuration from scratch.

Use this method when creating a fully custom simulator that doesn't build on any existing model.

Parameters:

name (str) –

Model name
params (list[str]) –

List of parameter names
simulator_function (Callable) –

Simulator function
nchoices (int) –

Number of choices
**config –

Additional configuration fields: - param_bounds : list - [[lower bounds], [upper bounds]] - default_params : list - Default parameter values - choices : list - Possible choice values - n_particles : int - Number of particles (default 1) - boundary : Callable - Boundary function - boundary_name : str - Boundary name - boundary_params : list[str] - Boundary parameter names (including 'a') - drift : Callable - Drift function - drift_name : str - Drift name - drift_params : list[str] - Drift parameter names

Returns:

dict –

Complete configuration dictionary

Examples:

>>> def my_sim(v, a, **kwargs):
...     # Custom simulation logic
...     return {'rts': ..., 'choices': ..., 'metadata': ...}
>>>
>>> config = ModelConfigBuilder.from_scratch(
...     name="my_custom_model",
...     params=["v", "a"],
...     simulator_function=my_sim,
...     nchoices=2,
...     param_bounds=[[-2, 0.5], [2, 2.0]],
...     default_params=[0.0, 1.0]
... )

get_sampling_transforms `staticmethod`

get_sampling_transforms(config)

Get parameter sampling transforms from model config.

These transforms are applied during the parameter sampling stage of the training data generation workflow. They enforce parameter relationships (e.g., a > z) when generating synthetic training data for likelihood approximation networks.

Note: These are NOT directly relevant for basic Simulator usage, which uses simulation transforms via ParameterSimulatorAdapters instead.

Parameters:

config (dict) –

Model configuration dictionary

Returns:

list –

List of transform/constraint instances (empty if none defined)

Examples:

>>> config = {
...     "parameter_transforms": {
...         "sampling": [SwapIfLessConstraint("a", "z")],
...     }
... }
>>> transforms = ModelConfigBuilder.get_sampling_transforms(config)

get_simulation_transforms `staticmethod`

get_simulation_transforms(config)

Get simulation transforms from model config.

These transforms are applied via ParameterSimulatorAdapters when running the basic Simulator. They prepare user-provided parameters for the low-level C/Cython simulators (e.g., stacking v0, v1, v2 into a single v array, expanding dimensions).

Parameters:

config (dict) –

Model configuration dictionary

Returns:

list –

List of transform instances (empty if none defined)

Examples:

>>> config = {
...     "parameter_transforms": {
...         "simulation": [
...             ColumnStackParameters(["v0", "v1", "v2"], "v"),
...             ExpandDimension(["a", "z"]),
...         ],
...     }
... }
>>> transforms = ModelConfigBuilder.get_simulation_transforms(config)
>>> len(transforms)
2

get_transforms `staticmethod`

get_transforms(config, phase)

Get transforms for a specific phase from model config.

This method extracts parameter transforms from the unified parameter_transforms field in the model configuration.

Parameters:

config (dict) –

Model configuration dictionary
phase (str) –

Either 'sampling' or 'simulation'

Returns:

list –

List of transform instances (empty if none defined)

Examples:

>>> config = {
...     "name": "lba_angle_3",
...     "parameter_transforms": {
...         "sampling": [SwapIfLessConstraint("a", "z")],
...         "simulation": [ColumnStackParameters(["v0", "v1"], "v")],
...     }
... }
>>> sampling_transforms = ModelConfigBuilder.get_transforms(config, "sampling")
>>> len(sampling_transforms)
1

minimal_config `staticmethod`

minimal_config(
    params, simulator_function, nchoices=2, name="custom"
)

Create a minimal valid configuration.

This is the simplest way to create a configuration for a custom simulator. It includes only the required fields.

Parameters:

params (list[str]) –

List of parameter names
simulator_function (Callable) –

Simulator function
nchoices (int, default: 2 ) –

Number of choices
name (str, default: "custom" ) –

Model name

Returns:

dict –

Minimal configuration dictionary

Examples:

>>> config = ModelConfigBuilder.minimal_config(
...     params=["v", "a", "z", "t"],
...     simulator_function=my_sim_fn,
...     nchoices=2
... )

validate_config `staticmethod`

validate_config(config, strict=False)

Validate a configuration dictionary.

Parameters:

config (dict) –

Configuration to validate
strict (bool, default: False ) –

If True, also check for recommended optional fields

Returns:

is_valid ( bool ) –

Whether the configuration is valid
errors ( list[str] ) –

List of error messages (empty if valid)

Examples:

>>> config = {"params": ["v", "a"], "nchoices": 2}
>>> is_valid, errors = ModelConfigBuilder.validate_config(config)
>>> if not is_valid:
...     print("Errors:", errors)

with_deadline `staticmethod`

with_deadline(config)

Add deadline parameter to a model configuration.

Creates a NEW configuration with the deadline parameter added. This is an immutable operation - the original config is not modified.

The deadline parameter allows models to incorporate response deadlines, where trials are terminated if no response is made within the deadline.

This method is idempotent - calling it on a config that already has the deadline parameter will return an equivalent config.

Parameters:

config (dict) –

Model configuration to extend with deadline support

Returns:

dict –

New configuration with deadline parameter added. Includes: - "deadline" appended to params list - Updated param_bounds (both list and dict formats) - Updated default_params - "_deadline" suffix added to name - Incremented n_params - "deadline" metadata flag set to True

Examples:

>>> base_config = ModelConfigBuilder.from_model("ddm")
>>> deadline_config = ModelConfigBuilder.with_deadline(base_config)
>>> "deadline" in deadline_config["params"]
True
>>> deadline_config["name"]
'ddm_deadline'

boundary_config_to_function_params

boundary_config_to_function_params(config)

Convert boundary configuration to function parameters.

Parameters:

config (dict) –

Dictionary containing the boundary configuration

Returns:

dict –

Dictionary with adjusted key names so that they match function parameters names directly.

boundary_registry

Global registry for boundary functions.

This module provides a centralized registry for boundary functions used in sequential sampling models. It follows the same pattern as the parameter sampling constraint registry for consistency.

Examples:

Register a custom boundary:

>>> from ssms.config import register_boundary
>>>
>>> def my_boundary(t, a=1.0, decay=0.1):
...     return a * np.exp(-decay * t)
>>>
>>> register_boundary("exponential", my_boundary, ["a", "decay"])
>>>
>>> # Use with ModelConfigBuilder
>>> from ssms.config import ModelConfigBuilder
>>> config = ModelConfigBuilder.from_model("ddm")
>>> config = ModelConfigBuilder.add_boundary(config, "exponential")

List available boundaries:

>>> from ssms.config import get_boundary_registry
>>> print(get_boundary_registry().list_boundaries())

BoundaryRegistry

BoundaryRegistry()

Global registry for boundary functions.

This registry maintains a mapping of boundary names to their configuration, including the function and parameters.

All boundary functions accept 'a' as an explicit parameter and return the final boundary value directly.

repr

__repr__()

String representation of registry.

get

get(name)

Get boundary configuration by name.

Parameters:

name (str) –

Name of the registered boundary

Returns:

dict –

Dictionary containing: - 'fun': The boundary function - 'params': List of parameter names (including 'a')

Raises:

KeyError –

If boundary name not registered

Examples:

>>> registry = BoundaryRegistry()
>>> config = registry.get("angle")
>>> print(config["params"])
['theta']

is_registered

is_registered(name)

Check if boundary name is registered.

Parameters:

name (str) –

Boundary name to check

Returns:

bool –

True if boundary is registered, False otherwise

Examples:

>>> registry = BoundaryRegistry()
>>> registry.is_registered("angle")
True
>>> registry.is_registered("my_custom_boundary")
False

list_boundaries

list_boundaries()

List all registered boundary names.

Returns:

list[str] –

Sorted list of all registered boundary names

Examples:

>>> registry = BoundaryRegistry()
>>> boundaries = registry.list_boundaries()
>>> print(boundaries)
['angle', 'constant', 'weibull_cdf', ...]

register

register(name, function, params)

Register a boundary function.

Parameters:

name (str) –

Unique name for the boundary (e.g., "angle", "weibull_cdf")
function (Callable) –

Boundary function with signature (t, **params) -> float or array where t is time and params are boundary-specific parameters (including 'a')
params (list[str]) –

List of parameter names the function expects (e.g., ["a", "theta"] for angle)

Raises:

ValueError –

If name already registered

Examples:

>>> def exponential_decay(t, a=1.0, rate=0.1):
...     return a * np.exp(-rate * t)
>>>
>>> registry = BoundaryRegistry()
>>> registry.register("exp_decay", exponential_decay, ["a", "rate"])

get_boundary_registry

get_boundary_registry()

Get the global boundary registry.

Use this to access registry methods like list_boundaries() or is_registered().

Returns:

BoundaryRegistry –

The global boundary registry instance

Examples:

>>> from ssms.config import get_boundary_registry
>>>
>>> # List all available boundaries
>>> registry = get_boundary_registry()
>>> print(registry.list_boundaries())
['angle', 'constant', 'weibull_cdf', ...]
>>>
>>> # Check if a boundary exists
>>> if registry.is_registered("my_boundary"):
...     config = registry.get("my_boundary")

register_boundary

register_boundary(name, function, params)

Register a boundary function globally.

This is the main entry point for registering custom boundary functions. Once registered, boundaries can be used with ModelConfigBuilder.add_boundary() just like built-in boundaries.

Parameters:

name (str) –

Unique name for the boundary
function (Callable) –

Boundary function with signature (t, **params) -> float or array
params (list[str]) –

List of parameter names the function expects (must include 'a')

Raises:

ValueError –

If name already registered

Examples:

Register a custom exponential decay boundary:

>>> import numpy as np
>>> from ssms.config import register_boundary
>>>
>>> def exponential_decay(t, a=1.0, rate=0.1):
...     return a * np.exp(-rate * t)
>>>
>>> register_boundary(
...     name="exponential_decay",
...     function=exponential_decay,
...     params=["a", "rate"]
... )
>>>
>>> # Now use it with ModelConfigBuilder
>>> from ssms.config import ModelConfigBuilder
>>> config = ModelConfigBuilder.from_model("ddm")
>>> config = ModelConfigBuilder.add_boundary(config, "exponential_decay")

Register a collapsing boundary:

>>> def linear_collapse(t, a=1.0, slope=-0.5):
...     return a + slope * t
>>>
>>> register_boundary("linear_collapse", linear_collapse, ["a", "slope"])

config_utils

Utilities for handling generator config with nested structure.

This module provides utilities for working with the nested generator_config structure.

Nested structure (REQUIRED): { "pipeline": {"n_parameter_sets": 100, "n_subruns": 10, ...}, "estimator": {"type": "kde", "bandwidth": 0.1, ...}, "training": {"mixture_probabilities": [0.8, 0.1, 0.1], ...}, "simulator": {"delta_t": 0.001, "max_t": 20.0, ...}, "output": {"folder": "...", "pickle_protocol": 4, ...}, }

Note: Only nested configs are supported. Flat configs are no longer accepted.

get_nested_config

get_nested_config(config, section, key, default=None)

Get a value from nested config structure.

Args: config: Generator configuration dictionary (must use nested structure) section: Nested section name ("pipeline", "estimator", "training", "simulator", "output") key: Key name within the section default: Default value if key not found

Returns: Value from nested structure if available, else default

Examples: >>> config = {"pipeline": {"n_parameter_sets": 100}} >>> get_nested_config(config, "pipeline", "n_parameter_sets") 100

>>> get_nested_config(config, "pipeline", "missing_key", default=42)
42

has_nested_structure

has_nested_structure(config)

Check if config uses the required nested structure.

Args: config: Generator configuration dictionary

Returns: True if config has nested sections (required format), False otherwise

Note: Only nested configs are supported. This function validates that configs have the correct structure with at least one of the required sections.

drift_registry

Global registry for drift functions.

This module provides a centralized registry for drift functions used in sequential sampling models. It follows the same pattern as the parameter sampling constraint registry and boundary registry for consistency.

Examples:

Register a custom drift:

>>> from ssms.config import register_drift
>>>
>>> def sinusoidal_drift(t, frequency=1.0, amplitude=0.5, baseline=1.0):
...     return baseline + amplitude * np.sin(2 * np.pi * frequency * t)
>>>
>>> register_drift("sinusoidal", sinusoidal_drift, ["frequency", "amplitude", "baseline"])
>>>
>>> # Use with ModelConfigBuilder
>>> from ssms.config import ModelConfigBuilder
>>> config = ModelConfigBuilder.from_model("ddm")
>>> config = ModelConfigBuilder.add_drift(config, "sinusoidal")

List available drifts:

>>> from ssms.config import get_drift_registry
>>> print(get_drift_registry().list_drifts())

DriftRegistry

DriftRegistry()

Global registry for drift functions.

This registry maintains a mapping of drift names to their configuration, including the function and its parameters.

repr

__repr__()

String representation of registry.

get

get(name)

Get drift configuration by name.

Parameters:

name (str) –

Name of the registered drift

Returns:

dict –

Dictionary containing: - 'fun': The drift function - 'params': List of parameter names

Raises:

KeyError –

If drift name not registered

Examples:

>>> registry = DriftRegistry()
>>> config = registry.get("gamma_drift")
>>> print(config["params"])
['shape', 'scale', 'c']

is_registered

is_registered(name)

Check if drift name is registered.

Parameters:

name (str) –

Drift name to check

Returns:

bool –

True if drift is registered, False otherwise

Examples:

>>> registry = DriftRegistry()
>>> registry.is_registered("gamma_drift")
True
>>> registry.is_registered("my_custom_drift")
False

list_drifts

list_drifts()

List all registered drift names.

Returns:

list[str] –

Sorted list of all registered drift names

Examples:

>>> registry = DriftRegistry()
>>> drifts = registry.list_drifts()
>>> print(drifts)
['constant', 'gamma_drift', ...]

register

register(name, function, params)

Register a drift function.

Parameters:

name (str) –

Unique name for the drift (e.g., "gamma_drift", "constant")
function (Callable) –

Drift function with signature (t, **params) -> float or array where t is time and params are drift-specific parameters
params (list[str]) –

List of parameter names the function expects (e.g., ["shape", "scale", "c"])

Raises:

ValueError –

If name already registered

Examples:

>>> def linear_drift(t, slope=0.5, intercept=1.0):
...     return intercept + slope * t
>>>
>>> registry = DriftRegistry()
>>> registry.register("linear", linear_drift, ["slope", "intercept"])

get_drift_registry

get_drift_registry()

Get the global drift registry.

Use this to access registry methods like list_drifts() or is_registered().

Returns:

DriftRegistry –

The global drift registry instance

Examples:

>>> from ssms.config import get_drift_registry
>>>
>>> # List all available drifts
>>> registry = get_drift_registry()
>>> print(registry.list_drifts())
['constant', 'gamma_drift', ...]
>>>
>>> # Check if a drift exists
>>> if registry.is_registered("my_drift"):
...     config = registry.get("my_drift")

register_drift

register_drift(name, function, params)

Register a drift function globally.

This is the main entry point for registering custom drift functions. Once registered, drifts can be used with ModelConfigBuilder.add_drift() just like built-in drifts.

Parameters:

name (str) –

Unique name for the drift
function (Callable) –

Drift function with signature (t, **params) -> float or array
params (list[str]) –

List of parameter names the function expects

Raises:

ValueError –

If name already registered

Examples:

Register a custom sinusoidal drift:

>>> import numpy as np
>>> from ssms.config import register_drift
>>>
>>> def sinusoidal_drift(t, frequency=1.0, amplitude=0.5, baseline=1.0):
...     return baseline + amplitude * np.sin(2 * np.pi * frequency * t)
>>>
>>> register_drift(
...     name="sinusoidal",
...     function=sinusoidal_drift,
...     params=["frequency", "amplitude", "baseline"]
... )
>>>
>>> # Now use it with ModelConfigBuilder
>>> from ssms.config import ModelConfigBuilder
>>> config = ModelConfigBuilder.from_model("ddm")
>>> config = ModelConfigBuilder.add_drift(config, "sinusoidal")

Register a time-varying drift:

>>> def exponential_drift(t, rate=0.1, asymptote=2.0):
...     return asymptote * (1 - np.exp(-rate * t))
>>>
>>> register_drift("exponential", exponential_drift, ["rate", "asymptote"])

get_boundary_registry

get_boundary_registry()

Get the global boundary registry.

Use this to access registry methods like list_boundaries() or is_registered().

Returns:

BoundaryRegistry –

The global boundary registry instance

Examples:

>>> from ssms.config import get_boundary_registry
>>>
>>> # List all available boundaries
>>> registry = get_boundary_registry()
>>> print(registry.list_boundaries())
['angle', 'constant', 'weibull_cdf', ...]
>>>
>>> # Check if a boundary exists
>>> if registry.is_registered("my_boundary"):
...     config = registry.get("my_boundary")

get_drift_registry

get_drift_registry()

Get the global drift registry.

Use this to access registry methods like list_drifts() or is_registered().

Returns:

DriftRegistry –

The global drift registry instance

Examples:

>>> from ssms.config import get_drift_registry
>>>
>>> # List all available drifts
>>> registry = get_drift_registry()
>>> print(registry.list_drifts())
['constant', 'gamma_drift', ...]
>>>
>>> # Check if a drift exists
>>> if registry.is_registered("my_drift"):
...     config = registry.get("my_drift")

get_model_registry

get_model_registry()

Get the global model registry.

Use this to access registry methods like list_models() or has_model().

Returns:

ModelConfigRegistry –

The global model registry instance

Examples:

>>> from ssms.config import get_model_registry
>>>
>>> # List all available models
>>> registry = get_model_registry()
>>> print(registry.list_models())
['ddm', 'angle', 'weibull_cdf', ...]
>>>
>>> # Check if a model exists
>>> if registry.has_model("my_model"):
...     config = registry.get("my_model")

model_config_builder

Utilities for building custom model configurations.

This module provides helper classes and functions for creating valid model configurations for use with the Simulator class.

ModelConfigBuilder

Helper class for building custom model configurations.

This class provides static methods for creating model configurations in various ways: - Starting from an existing model and overriding specific values - Building a configuration from scratch for fully custom simulators - Creating minimal valid configurations - Validating configurations

Examples:

Start from existing model and override:

>>> config = ModelConfigBuilder.from_model("ddm",
...                                     param_bounds=[[-4, 0.3, 0.1, 0],
...                                                   [4, 3.0, 0.9, 2.0]])

Build from scratch:

>>> config = ModelConfigBuilder.from_scratch(
...     name="my_model",
...     params=["v", "a", "z", "t"],
...     simulator_function=my_sim_fn,
...     nchoices=2
... )

Create minimal configuration:

>>> config = ModelConfigBuilder.minimal_config(
...     params=["v", "a"],
...     simulator_function=my_sim_fn,
...     nchoices=2
... )

add_boundary `staticmethod`

add_boundary(config, boundary, boundary_params=None)

Add or replace boundary function in configuration.

Parameters:

config (dict) –

Configuration to modify
boundary (str or Callable) –

Boundary function name or callable
boundary_params (list[str] or None, default: None ) –

Parameter names for boundary function (required if boundary is callable)
multiplicative (bool, default: True ) –

Whether boundary is multiplicative (True) or additive (False)

Returns:

dict –

Modified configuration (note: modifies in place and returns)

Raises:

ValueError –

If boundary specification is invalid

Examples:

>>> config = ModelConfigBuilder.from_model("ddm")
>>> config = ModelConfigBuilder.add_boundary(config, "angle")

add_drift `staticmethod`

add_drift(config, drift, drift_params=None)

Add or replace drift function in configuration.

Parameters:

config (dict) –

Configuration to modify
drift (str or Callable) –

Drift function name or callable
drift_params (list[str] or None, default: None ) –

Parameter names for drift function (required if drift is callable)

Returns:

dict –

Modified configuration (note: modifies in place and returns)

Raises:

ValueError –

If drift specification is invalid

Examples:

>>> config = ModelConfigBuilder.from_model("ddm")
>>> config = ModelConfigBuilder.add_drift(config, "gamma_drift")

from_model `staticmethod`

from_model(model_name, **overrides)

Create configuration starting from an existing model.

This method automatically parses variant suffixes like "_deadline" from the model name and applies the appropriate transformations.

Parameters:

model_name (str) –

Name of the model, optionally with variant suffixes. Examples: "ddm", "angle", "ddm_deadline", "angle_deadline"
**overrides –

Configuration fields to override. Common options: - params : list[str] - Parameter names - param_bounds : list - [[lower bounds], [upper bounds]] - boundary : Callable - Custom boundary function - boundary_name : str - Boundary name - boundary_params : list[str] - Boundary parameter names - drift : Callable - Custom drift function - drift_name : str - Drift name - drift_params : list[str] - Drift parameter names - simulator : Callable - Custom simulator function - nchoices : int - Number of choices - choices : list - Possible choice values

Returns:

dict –

Configuration dictionary with any variants applied

Raises:

ValueError –

If the base model name is not recognized

Examples:

>>> config = ModelConfigBuilder.from_model("ddm",
...                                     param_bounds=[[-4, 0.3, 0.1, 0],
...                                                   [4, 3.0, 0.9, 2.0]])

>>> # With deadline variant
>>> config = ModelConfigBuilder.from_model("ddm_deadline")
>>> "deadline" in config["params"]
True

from_scratch `staticmethod`

from_scratch(
    name, params, simulator_function, nchoices, **config
)

Build a complete configuration from scratch.

Use this method when creating a fully custom simulator that doesn't build on any existing model.

Parameters:

name (str) –

Model name
params (list[str]) –

List of parameter names
simulator_function (Callable) –

Simulator function
nchoices (int) –

Number of choices
**config –

Additional configuration fields: - param_bounds : list - [[lower bounds], [upper bounds]] - default_params : list - Default parameter values - choices : list - Possible choice values - n_particles : int - Number of particles (default 1) - boundary : Callable - Boundary function - boundary_name : str - Boundary name - boundary_params : list[str] - Boundary parameter names (including 'a') - drift : Callable - Drift function - drift_name : str - Drift name - drift_params : list[str] - Drift parameter names

Returns:

dict –

Complete configuration dictionary

Examples:

>>> def my_sim(v, a, **kwargs):
...     # Custom simulation logic
...     return {'rts': ..., 'choices': ..., 'metadata': ...}
>>>
>>> config = ModelConfigBuilder.from_scratch(
...     name="my_custom_model",
...     params=["v", "a"],
...     simulator_function=my_sim,
...     nchoices=2,
...     param_bounds=[[-2, 0.5], [2, 2.0]],
...     default_params=[0.0, 1.0]
... )

get_sampling_transforms `staticmethod`

get_sampling_transforms(config)

Get parameter sampling transforms from model config.

These transforms are applied during the parameter sampling stage of the training data generation workflow. They enforce parameter relationships (e.g., a > z) when generating synthetic training data for likelihood approximation networks.

Note: These are NOT directly relevant for basic Simulator usage, which uses simulation transforms via ParameterSimulatorAdapters instead.

Parameters:

config (dict) –

Model configuration dictionary

Returns:

list –

List of transform/constraint instances (empty if none defined)

Examples:

>>> config = {
...     "parameter_transforms": {
...         "sampling": [SwapIfLessConstraint("a", "z")],
...     }
... }
>>> transforms = ModelConfigBuilder.get_sampling_transforms(config)

get_simulation_transforms `staticmethod`

get_simulation_transforms(config)

Get simulation transforms from model config.

These transforms are applied via ParameterSimulatorAdapters when running the basic Simulator. They prepare user-provided parameters for the low-level C/Cython simulators (e.g., stacking v0, v1, v2 into a single v array, expanding dimensions).

Parameters:

config (dict) –

Model configuration dictionary

Returns:

list –

List of transform instances (empty if none defined)

Examples:

>>> config = {
...     "parameter_transforms": {
...         "simulation": [
...             ColumnStackParameters(["v0", "v1", "v2"], "v"),
...             ExpandDimension(["a", "z"]),
...         ],
...     }
... }
>>> transforms = ModelConfigBuilder.get_simulation_transforms(config)
>>> len(transforms)
2

get_transforms `staticmethod`

get_transforms(config, phase)

Get transforms for a specific phase from model config.

This method extracts parameter transforms from the unified parameter_transforms field in the model configuration.

Parameters:

config (dict) –

Model configuration dictionary
phase (str) –

Either 'sampling' or 'simulation'

Returns:

list –

List of transform instances (empty if none defined)

Examples:

>>> config = {
...     "name": "lba_angle_3",
...     "parameter_transforms": {
...         "sampling": [SwapIfLessConstraint("a", "z")],
...         "simulation": [ColumnStackParameters(["v0", "v1"], "v")],
...     }
... }
>>> sampling_transforms = ModelConfigBuilder.get_transforms(config, "sampling")
>>> len(sampling_transforms)
1

minimal_config `staticmethod`

minimal_config(
    params, simulator_function, nchoices=2, name="custom"
)

Create a minimal valid configuration.

This is the simplest way to create a configuration for a custom simulator. It includes only the required fields.

Parameters:

params (list[str]) –

List of parameter names
simulator_function (Callable) –

Simulator function
nchoices (int, default: 2 ) –

Number of choices
name (str, default: "custom" ) –

Model name

Returns:

dict –

Minimal configuration dictionary

Examples:

>>> config = ModelConfigBuilder.minimal_config(
...     params=["v", "a", "z", "t"],
...     simulator_function=my_sim_fn,
...     nchoices=2
... )

validate_config `staticmethod`

validate_config(config, strict=False)

Validate a configuration dictionary.

Parameters:

config (dict) –

Configuration to validate
strict (bool, default: False ) –

If True, also check for recommended optional fields

Returns:

is_valid ( bool ) –

Whether the configuration is valid
errors ( list[str] ) –

List of error messages (empty if valid)

Examples:

>>> config = {"params": ["v", "a"], "nchoices": 2}
>>> is_valid, errors = ModelConfigBuilder.validate_config(config)
>>> if not is_valid:
...     print("Errors:", errors)

with_deadline `staticmethod`

with_deadline(config)

Add deadline parameter to a model configuration.

Creates a NEW configuration with the deadline parameter added. This is an immutable operation - the original config is not modified.

The deadline parameter allows models to incorporate response deadlines, where trials are terminated if no response is made within the deadline.

This method is idempotent - calling it on a config that already has the deadline parameter will return an equivalent config.

Parameters:

config (dict) –

Model configuration to extend with deadline support

Returns:

dict –

New configuration with deadline parameter added. Includes: - "deadline" appended to params list - Updated param_bounds (both list and dict formats) - Updated default_params - "_deadline" suffix added to name - Incremented n_params - "deadline" metadata flag set to True

Examples:

>>> base_config = ModelConfigBuilder.from_model("ddm")
>>> deadline_config = ModelConfigBuilder.with_deadline(base_config)
>>> "deadline" in deadline_config["params"]
True
>>> deadline_config["name"]
'ddm_deadline'

model_registry

Global registry for model configurations.

This module provides a centralized registry for complete model configurations. It follows the same pattern as boundary and drift registries for consistency, with additional support for factory functions to enable lazy loading.

Examples:

Register a custom model configuration:

>>> from ssms.config import register_model_config
>>>
>>> my_config = {
...     "name": "my_ddm",
...     "params": ["v", "a", "z", "t"],
...     "nchoices": 2,
...     "simulator": my_sim_fn,
... }
>>>
>>> register_model_config("my_ddm", my_config)
>>>
>>> # Now use it with ModelConfigBuilder
>>> from ssms.config import ModelConfigBuilder
>>> config = ModelConfigBuilder.from_model("my_ddm")

Register using a factory function:

>>> def get_my_model_config():
...     return {...}
>>>
>>> register_model_config_factory("my_model", get_my_model_config)

List available models:

>>> from ssms.config import get_model_registry
>>> print(get_model_registry().list_models())

ModelConfigRegistry

ModelConfigRegistry()

Global registry for complete model configurations.

This registry maintains a mapping of model names to their configurations. Supports both direct config registration and factory functions for lazy loading.

repr

__repr__()

String representation of registry.

get

get(name)

Get model configuration by name.

Returns a deep copy of the configuration to prevent accidental mutation of the registered config.

Parameters:

name (str) –

Name of the registered model

Returns:

dict –

Complete model configuration dictionary (deep copy)

Raises:

KeyError –

If model name not registered

Examples:

>>> registry = ModelConfigRegistry()
>>> config = registry.get("ddm")
>>> print(config["params"])
['v', 'a', 'z', 't']

has_model

has_model(name)

Check if model name is registered.

Parameters:

name (str) –

Model name to check

Returns:

bool –

True if model is registered, False otherwise

Examples:

>>> registry = ModelConfigRegistry()
>>> registry.has_model("ddm")
True
>>> registry.has_model("my_custom_model")
False

list_models

list_models()

List all registered model names.

Returns:

list[str] –

Sorted list of all registered model names

Examples:

>>> registry = ModelConfigRegistry()
>>> models = registry.list_models()
>>> print(models[:5])
['angle', 'ddm', 'ddm_par2', 'ddm_sdv', 'ddm_st']

register_config

register_config(name, config)

Register a model configuration directly.

Parameters:

name (str) –

Unique name for the model (e.g., "ddm", "my_custom_model")
config (dict) –

Complete model configuration dictionary containing at minimum: - 'name': Model name - 'params': List of parameter names - 'nchoices': Number of choices - 'simulator': Simulator function

Raises:

ValueError –

If name already registered (either as config or factory)

Examples:

>>> config = {
...     "name": "my_model",
...     "params": ["v", "a", "z", "t"],
...     "param_bounds": [[-3, 0.3, 0.1, 0], [3, 3.0, 0.9, 2]],
...     "nchoices": 2,
...     "simulator": my_sim_fn,
... }
>>>
>>> registry = ModelConfigRegistry()
>>> registry.register_config("my_model", config)

register_factory

register_factory(name, factory)

Register a model config factory function.

Factory functions enable lazy loading - the config is only created when first accessed. This is useful for models with expensive initialization or to reduce memory footprint.

Parameters:

name (str) –

Unique name for the model
factory (Callable[[], dict]) –

Function that returns a complete model configuration dict

Raises:

ValueError –

If name already registered (either as config or factory)

Examples:

>>> def get_my_model_config():
...     # Expensive computation here
...     return {...}
>>>
>>> registry = ModelConfigRegistry()
>>> registry.register_factory("my_model", get_my_model_config)

get_model_registry

get_model_registry()

Get the global model registry.

Use this to access registry methods like list_models() or has_model().

Returns:

ModelConfigRegistry –

The global model registry instance

Examples:

>>> from ssms.config import get_model_registry
>>>
>>> # List all available models
>>> registry = get_model_registry()
>>> print(registry.list_models())
['ddm', 'angle', 'weibull_cdf', ...]
>>>
>>> # Check if a model exists
>>> if registry.has_model("my_model"):
...     config = registry.get("my_model")

register_model_config

register_model_config(name, config)

Register a model configuration globally.

This is the main entry point for registering custom model configurations. Once registered, models can be used with ModelConfigBuilder.from_model() just like built-in models.

Parameters:

name (str) –

Unique name for the model
config (dict) –

Complete model configuration dictionary

Raises:

ValueError –

If name already registered

Examples:

Register a complete custom model:

>>> from ssms.config import register_model_config
>>>
>>> my_model = {
...     "name": "my_custom_ddm",
...     "params": ["v", "a", "z", "t"],
...     "param_bounds": [[-3, 0.3, 0.1, 0], [3, 3.0, 0.9, 2]],
...     "nchoices": 2,
...     "n_params": 4,
...     "default_params": [1.0, 1.5, 0.5, 0.3],
...     "simulator": my_simulator_function,
... }
>>>
>>> register_model_config("my_custom_ddm", my_model)
>>>
>>> # Now use it like any built-in model
>>> from ssms.config import ModelConfigBuilder
>>> config = ModelConfigBuilder.from_model("my_custom_ddm")
>>>
>>> # Or with Simulator
>>> from ssms.basic_simulators import Simulator
>>> sim = Simulator(model="my_custom_ddm")

Register with custom boundary and drift:

>>> advanced_model = {
...     "name": "advanced_ddm",
...     "params": ["v", "a", "z", "t", "theta"],
...     "nchoices": 2,
...     "boundary": my_boundary_fn,
...     "boundary_name": "custom",
...     "boundary_params": ["theta"],
...     "drift": my_drift_fn,
...     "drift_name": "custom",
...     "drift_params": [],
...     "simulator": my_simulator,
... }
>>>
>>> register_model_config("advanced_ddm", advanced_model)

register_model_config_factory

register_model_config_factory(name, factory)

Register a model config factory function globally.

Use this when you want lazy loading of model configurations, or when the config requires computation/processing at access time.

Parameters:

name (str) –

Unique name for the model
factory (Callable[[], dict]) –

Function that returns a complete model configuration dict

Raises:

ValueError –

If name already registered

Examples:

Register with lazy loading:

>>> from ssms.config import register_model_config_factory
>>>
>>> def get_my_model():
...     # This only runs when the model is first accessed
...     return {
...         "name": "my_model",
...         "params": ["v", "a", "z", "t"],
...         "nchoices": 2,
...         "simulator": create_simulator(),  # Expensive operation
...     }
>>>
>>> register_model_config_factory("my_model", get_my_model)
>>>
>>> # Factory is called only when accessing the model
>>> config = ModelConfigBuilder.from_model("my_model")

register_boundary

register_boundary(name, function, params)

Register a boundary function globally.

This is the main entry point for registering custom boundary functions. Once registered, boundaries can be used with ModelConfigBuilder.add_boundary() just like built-in boundaries.

Parameters:

name (str) –

Unique name for the boundary
function (Callable) –

Boundary function with signature (t, **params) -> float or array
params (list[str]) –

List of parameter names the function expects (must include 'a')

Raises:

ValueError –

If name already registered

Examples:

Register a custom exponential decay boundary:

>>> import numpy as np
>>> from ssms.config import register_boundary
>>>
>>> def exponential_decay(t, a=1.0, rate=0.1):
...     return a * np.exp(-rate * t)
>>>
>>> register_boundary(
...     name="exponential_decay",
...     function=exponential_decay,
...     params=["a", "rate"]
... )
>>>
>>> # Now use it with ModelConfigBuilder
>>> from ssms.config import ModelConfigBuilder
>>> config = ModelConfigBuilder.from_model("ddm")
>>> config = ModelConfigBuilder.add_boundary(config, "exponential_decay")

Register a collapsing boundary:

>>> def linear_collapse(t, a=1.0, slope=-0.5):
...     return a + slope * t
>>>
>>> register_boundary("linear_collapse", linear_collapse, ["a", "slope"])

register_drift

register_drift(name, function, params)

Register a drift function globally.

This is the main entry point for registering custom drift functions. Once registered, drifts can be used with ModelConfigBuilder.add_drift() just like built-in drifts.

Parameters:

name (str) –

Unique name for the drift
function (Callable) –

Drift function with signature (t, **params) -> float or array
params (list[str]) –

List of parameter names the function expects

Raises:

ValueError –

If name already registered

Examples:

Register a custom sinusoidal drift:

>>> import numpy as np
>>> from ssms.config import register_drift
>>>
>>> def sinusoidal_drift(t, frequency=1.0, amplitude=0.5, baseline=1.0):
...     return baseline + amplitude * np.sin(2 * np.pi * frequency * t)
>>>
>>> register_drift(
...     name="sinusoidal",
...     function=sinusoidal_drift,
...     params=["frequency", "amplitude", "baseline"]
... )
>>>
>>> # Now use it with ModelConfigBuilder
>>> from ssms.config import ModelConfigBuilder
>>> config = ModelConfigBuilder.from_model("ddm")
>>> config = ModelConfigBuilder.add_drift(config, "sinusoidal")

Register a time-varying drift:

>>> def exponential_drift(t, rate=0.1, asymptote=2.0):
...     return asymptote * (1 - np.exp(-rate * t))
>>>
>>> register_drift("exponential", exponential_drift, ["rate", "asymptote"])

register_model_config

register_model_config(name, config)

Register a model configuration globally.

This is the main entry point for registering custom model configurations. Once registered, models can be used with ModelConfigBuilder.from_model() just like built-in models.

Parameters:

name (str) –

Unique name for the model
config (dict) –

Complete model configuration dictionary

Raises:

ValueError –

If name already registered

Examples:

Register a complete custom model:

>>> from ssms.config import register_model_config
>>>
>>> my_model = {
...     "name": "my_custom_ddm",
...     "params": ["v", "a", "z", "t"],
...     "param_bounds": [[-3, 0.3, 0.1, 0], [3, 3.0, 0.9, 2]],
...     "nchoices": 2,
...     "n_params": 4,
...     "default_params": [1.0, 1.5, 0.5, 0.3],
...     "simulator": my_simulator_function,
... }
>>>
>>> register_model_config("my_custom_ddm", my_model)
>>>
>>> # Now use it like any built-in model
>>> from ssms.config import ModelConfigBuilder
>>> config = ModelConfigBuilder.from_model("my_custom_ddm")
>>>
>>> # Or with Simulator
>>> from ssms.basic_simulators import Simulator
>>> sim = Simulator(model="my_custom_ddm")

Register with custom boundary and drift:

>>> advanced_model = {
...     "name": "advanced_ddm",
...     "params": ["v", "a", "z", "t", "theta"],
...     "nchoices": 2,
...     "boundary": my_boundary_fn,
...     "boundary_name": "custom",
...     "boundary_params": ["theta"],
...     "drift": my_drift_fn,
...     "drift_name": "custom",
...     "drift_params": [],
...     "simulator": my_simulator,
... }
>>>
>>> register_model_config("advanced_ddm", advanced_model)

register_model_config_factory

register_model_config_factory(name, factory)

Register a model config factory function globally.

Use this when you want lazy loading of model configurations, or when the config requires computation/processing at access time.

Parameters:

name (str) –

Unique name for the model
factory (Callable[[], dict]) –

Function that returns a complete model configuration dict

Raises:

ValueError –

If name already registered

Examples:

Register with lazy loading:

>>> from ssms.config import register_model_config_factory
>>>
>>> def get_my_model():
...     # This only runs when the model is first accessed
...     return {
...         "name": "my_model",
...         "params": ["v", "a", "z", "t"],
...         "nchoices": 2,
...         "simulator": create_simulator(),  # Expensive operation
...     }
>>>
>>> register_model_config_factory("my_model", get_my_model)
>>>
>>> # Factory is called only when accessing the model
>>> config = ModelConfigBuilder.from_model("my_model")

ssms.dataset_generators

estimator_builders

Estimator builder implementations.

This module contains builder classes that construct likelihood estimators from configuration dictionaries. Builders follow the EstimatorBuilderProtocol and handle the responsibility of extracting relevant parameters from configs.

Available builders: - KDEEstimatorBuilder: Builds KDE-based likelihood estimators - Works with all models - Requires simulation data

PyDDMEstimatorBuilder: Builds PyDDM analytical PDF estimators
Only for compatible models (single-particle, two-choice, Gaussian noise)
Does not require simulation data (purely analytical)
Requires optional 'pyddm' package: pip install pyddm

Factory: - create_estimator_builder: Factory function for creating builders based on config - Automatically selects appropriate builder based on 'estimator_type' - Handles compatibility checking and error messages

KDEEstimatorBuilder

KDEEstimatorBuilder(generator_config)

Builder for KDE-based likelihood estimators.

This builder is responsible for: 1. Extracting relevant parameters from the generator_config 2. Instantiating KDELikelihoodEstimator with explicit parameters 3. Fitting the estimator with provided simulations

Attributes:

generator_config (dict) –

Configuration dictionary containing KDE settings
displace_t (bool) –

Extracted parameter: whether to displace time by the t parameter

Examples:

>>> builder = KDEEstimatorBuilder(generator_config)
>>> estimator = builder.build(theta, simulations)
>>> log_liks = estimator.evaluate(rts, choices)

Notes

Implements EstimatorBuilderProtocol from protocols.py Extracted from lan_mlp.py as part of Phase 1 refactoring

Arguments

generator_config : dict Configuration dictionary in nested structure containing: - estimator.displace_t (optional): Whether to displace time by t parameter. Defaults to False if not specified.

build

build(theta, simulations=None)

Build and fit a KDE likelihood estimator.

This method: 1. Creates a KDELikelihoodEstimator with extracted parameters 2. Fits it with the provided simulations

Arguments

theta : dict Parameter dictionary (not directly used by KDE, but part of protocol) simulations : dict | None Simulation data containing 'rts', 'choices', and 'metadata' keys. Required for KDE estimators (cannot be None).

Returns:

KDELikelihoodEstimator –

A fitted likelihood estimator ready for evaluate() or sample() calls

Raises:

ValueError –

If simulations is None (KDE requires simulation data)

Examples:

>>> builder = KDEEstimatorBuilder({"kde_displace_t": False})
>>> estimator = builder.build(theta, simulations)
>>> assert estimator._kde is not None  # Verifies it's fitted

Notes

The theta parameter is included in the signature to conform to EstimatorBuilderProtocol. Future estimator types (e.g., PyDDM) will use theta to configure model-specific parameters.

PyDDMEstimatorBuilder

PyDDMEstimatorBuilder(generator_config, model_config)

Builder for PyDDM analytical PDF estimators.

This builder: 1. Validates model compatibility with PyDDM at construction 2. Builds PyDDM models using SSMSToPyDDMMapper for each theta 3. Solves the Fokker-Planck equation analytically 4. Wraps the solution in a PyDDMLikelihoodEstimator

Key advantages: - No simulation required (purely analytical) - Deterministic results - Fast for compatible models

Compatible models: - Single-particle, two-choice models - Gaussian noise only - Examples: ddm, ornstein, angle, weibull, gamma_drift, conflict_*

Incompatible models: - Multi-particle: race_, lca_, lba* - Non-Gaussian noise: levy - Inter-trial variability: full_ddm, ddm_sdv, ddm_st

Validates model compatibility at construction time (fail-fast approach).

Args: generator_config: Generator settings including: - delta_t: Time step for solving - max_t: Maximum time - pdf_interpolation: 'linear' or 'cubic' (optional, defaults to 'cubic') model_config: Model specification from ssms.config

Raises: ValueError: If model is not compatible with PyDDM ImportError: If pyddm package is not installed

build

build(theta, simulations=None)

Build PyDDM likelihood estimator for given parameters.

Unlike KDE builder, this does NOT require simulation data. The estimator is built by solving the Fokker-Planck equation analytically.

Args: theta: Model parameters (v, a, z, t, etc.) simulations: Ignored (PyDDM doesn't need simulations)

Returns: PyDDMLikelihoodEstimator instance fitted to theta

Raises: ValueError: If PyDDM model cannot be built (e.g., invalid parameters)

Note: The simulations parameter is accepted for protocol compliance but is not used by PyDDM estimators.

builder_factory

Builder factory for creating likelihood estimators based on configuration.

This module provides a factory function that selects the appropriate estimator builder based on the generator_config settings. This enables seamless switching between different likelihood estimation methods (KDE, PyDDM, etc.) via configuration alone.

create_estimator_builder

create_estimator_builder(generator_config, model_config)

Create appropriate estimator builder based on configuration.

This factory function examines the generator_config and returns the appropriate builder for likelihood estimation.

Arguments

generator_config : dict Configuration dictionary containing: - 'estimator.type' (str): Estimator type ('kde' or 'pyddm') model_config : dict Model configuration dictionary (used for PyDDM builder)

Returns:

EstimatorBuilderProtocol –

An estimator builder instance

Raises:

ImportError –

If 'pyddm' estimator is requested but pyddm package is not installed
ValueError –

If an unknown estimator_type is specified or if model is incompatible with PyDDM

Examples:

>>> # Default to KDE
>>> builder = create_estimator_builder({}, {})
>>> isinstance(builder, KDEEstimatorBuilder)
True

>>> # Explicit KDE
>>> config = {"estimator": {"type": "kde"}}
>>> builder = create_estimator_builder(config, {})
>>> isinstance(builder, KDEEstimatorBuilder)
True

>>> # PyDDM estimator (requires pyddm package)
>>> config = {"estimator": {"type": "pyddm"}}
>>> builder = create_estimator_builder(config, model_config["ddm"])
>>> isinstance(builder, PyDDMEstimatorBuilder)
True

Notes

PyDDM estimator requires the optional 'pyddm' package. Install with: pip install pyddm or pip install ssms[pyddm]

PyDDM is only compatible with single-particle, two-choice, Gaussian noise models. Incompatible models will raise ValueError at builder construction.

create_estimator_builder

create_estimator_builder(generator_config, model_config)

Create appropriate estimator builder based on configuration.

This factory function examines the generator_config and returns the appropriate builder for likelihood estimation.

Arguments

generator_config : dict Configuration dictionary containing: - 'estimator.type' (str): Estimator type ('kde' or 'pyddm') model_config : dict Model configuration dictionary (used for PyDDM builder)

Returns:

EstimatorBuilderProtocol –

An estimator builder instance

Raises:

ImportError –

If 'pyddm' estimator is requested but pyddm package is not installed
ValueError –

If an unknown estimator_type is specified or if model is incompatible with PyDDM

Examples:

>>> # Default to KDE
>>> builder = create_estimator_builder({}, {})
>>> isinstance(builder, KDEEstimatorBuilder)
True

>>> # Explicit KDE
>>> config = {"estimator": {"type": "kde"}}
>>> builder = create_estimator_builder(config, {})
>>> isinstance(builder, KDEEstimatorBuilder)
True

>>> # PyDDM estimator (requires pyddm package)
>>> config = {"estimator": {"type": "pyddm"}}
>>> builder = create_estimator_builder(config, model_config["ddm"])
>>> isinstance(builder, PyDDMEstimatorBuilder)
True

Notes

PyDDM estimator requires the optional 'pyddm' package. Install with: pip install pyddm or pip install ssms[pyddm]

PyDDM is only compatible with single-particle, two-choice, Gaussian noise models. Incompatible models will raise ValueError at builder construction.

kde_builder

KDE Estimator Builder.

This module implements the builder pattern for KDE-based likelihood estimators. The builder extracts relevant parameters from the generator_config and instantiates a KDELikelihoodEstimator with explicit parameters.

This separation of concerns ensures that: 1. KDELikelihoodEstimator remains focused on likelihood estimation 2. Configuration parsing logic is centralized in the builder 3. New estimator parameters can be added with minimal changes

KDEEstimatorBuilder

KDEEstimatorBuilder(generator_config)

Builder for KDE-based likelihood estimators.

This builder is responsible for: 1. Extracting relevant parameters from the generator_config 2. Instantiating KDELikelihoodEstimator with explicit parameters 3. Fitting the estimator with provided simulations

Attributes:

generator_config (dict) –

Configuration dictionary containing KDE settings
displace_t (bool) –

Extracted parameter: whether to displace time by the t parameter

Examples:

>>> builder = KDEEstimatorBuilder(generator_config)
>>> estimator = builder.build(theta, simulations)
>>> log_liks = estimator.evaluate(rts, choices)

Notes

Implements EstimatorBuilderProtocol from protocols.py Extracted from lan_mlp.py as part of Phase 1 refactoring

Arguments

generator_config : dict Configuration dictionary in nested structure containing: - estimator.displace_t (optional): Whether to displace time by t parameter. Defaults to False if not specified.

build

build(theta, simulations=None)

Build and fit a KDE likelihood estimator.

This method: 1. Creates a KDELikelihoodEstimator with extracted parameters 2. Fits it with the provided simulations

Arguments

theta : dict Parameter dictionary (not directly used by KDE, but part of protocol) simulations : dict | None Simulation data containing 'rts', 'choices', and 'metadata' keys. Required for KDE estimators (cannot be None).

Returns:

KDELikelihoodEstimator –

A fitted likelihood estimator ready for evaluate() or sample() calls

Raises:

ValueError –

If simulations is None (KDE requires simulation data)

Examples:

>>> builder = KDEEstimatorBuilder({"kde_displace_t": False})
>>> estimator = builder.build(theta, simulations)
>>> assert estimator._kde is not None  # Verifies it's fitted

Notes

The theta parameter is included in the signature to conform to EstimatorBuilderProtocol. Future estimator types (e.g., PyDDM) will use theta to configure model-specific parameters.

pyddm_builder

Builder for PyDDM-based likelihood estimators.

PyDDMEstimatorBuilder

PyDDMEstimatorBuilder(generator_config, model_config)

Builder for PyDDM analytical PDF estimators.

This builder: 1. Validates model compatibility with PyDDM at construction 2. Builds PyDDM models using SSMSToPyDDMMapper for each theta 3. Solves the Fokker-Planck equation analytically 4. Wraps the solution in a PyDDMLikelihoodEstimator

Key advantages: - No simulation required (purely analytical) - Deterministic results - Fast for compatible models

Compatible models: - Single-particle, two-choice models - Gaussian noise only - Examples: ddm, ornstein, angle, weibull, gamma_drift, conflict_*

Incompatible models: - Multi-particle: race_, lca_, lba* - Non-Gaussian noise: levy - Inter-trial variability: full_ddm, ddm_sdv, ddm_st

Validates model compatibility at construction time (fail-fast approach).

Args: generator_config: Generator settings including: - delta_t: Time step for solving - max_t: Maximum time - pdf_interpolation: 'linear' or 'cubic' (optional, defaults to 'cubic') model_config: Model specification from ssms.config

Raises: ValueError: If model is not compatible with PyDDM ImportError: If pyddm package is not installed

build

build(theta, simulations=None)

Build PyDDM likelihood estimator for given parameters.

Unlike KDE builder, this does NOT require simulation data. The estimator is built by solving the Fokker-Planck equation analytically.

Args: theta: Model parameters (v, a, z, t, etc.) simulations: Ignored (PyDDM doesn't need simulations)

Returns: PyDDMLikelihoodEstimator instance fitted to theta

Raises: ValueError: If PyDDM model cannot be built (e.g., invalid parameters)

Note: The simulations parameter is accepted for protocol compliance but is not used by PyDDM estimators.

lan_mlp

This module defines a data generator class for use with LANs. The class defined below can be used to generate training data compatible with the expectations of LANs.

Note: Multiprocessing start method is configured in ssms/init.py to use 'spawn' by default for OpenMP safety. Override with SSMS_MP_START_METHOD environment variable if needed.

TrainingDataGenerator

TrainingDataGenerator(config=None, model_config=None)

The TrainingDataGenerator() class is used to generate training data for various likelihood approximators.

Attributes:

generator_config (dict) –
```
Configuration dictionary for the data generator.
(For an example load ssms.get_default_generator_config())
```
model_config: dict Configuration dictionary for the model to be simulated. (For an example load ssms.config.model_config['ddm']) _generation_pipeline: DataGenerationPipelineProtocol Strategy for orchestrating the complete data generation workflow.

Methods:

generate_data_training_uniform –

Generates training data for LANs. _get_ncpus() Helper function for determining the number of cpus to use for parallelization.

Returns:

TrainingDataGenerator object –

Notes

The class supports dependency injection of generation_pipeline, enabling complete customization of the data generation workflow. By default, it auto-creates the appropriate strategy based on generator_config settings (KDE-based simulation or PyDDM analytical methods).

Arguments

config: dict or DataGenerationPipelineProtocol Either: - A dictionary with generator configuration (e.g., from ssms.config.get_default_generator_config()). The TrainingDataGenerator will auto-create the appropriate strategy. - A DataGenerationPipelineProtocol instance for complete custom control over the data generation workflow. model_config: dict, optional Configuration dictionary for the model to be simulated. (For an example load ssms.config.model_config['ddm']) This serves as the single source of truth for model specifications, parameter bounds, transforms, and custom drift/boundary functions.

**Optional** when config is a dict containing a 'model' field - will
auto-derive using ModelConfigBuilder.from_model(config['model']).
**Optional** when config is a pipeline (extracted from pipeline if not provided).

Raises:

ValueError –

If config is None, or if model_config is None and cannot be derived (i.e., config is a dict without 'model' field), or if model_config cannot be extracted from a pipeline.

Returns:

TrainingDataGenerator object –

Notes

Simple Usage (recommended): Pass a generator_config dict with 'model' field set. The TrainingDataGenerator will auto-derive model_config and create the appropriate strategy.

Explicit model_config: Pass both generator_config and model_config for complete control over model configuration (e.g., custom param_bounds).

Advanced Usage: Pass a custom DataGenerationPipelineProtocol instance to completely control the workflow (parameter sampling, simulation, estimation, etc.). In this case, model_config can be omitted (will use pipeline's config).

The pipeline internally manages: - Parameter sampling (uniform, Sobol, custom) - Simulation or analytical PDF computation - Likelihood estimation (KDE, PyDDM analytical) - Training data structuring (ResampleMixture, etc.)

For custom drift/boundary functions, use ModelConfigBuilder to create a custom model_config before passing it to TrainingDataGenerator.

Examples:

Simple (auto-derive model_config from generator_config['model']): >>> from ssms import get_default_generator_config >>> config = get_default_generator_config(model="ddm_deadline") >>> gen = TrainingDataGenerator(config) # model_config auto-derived! >>> data = gen.generate_data_training()

Explicit model_config (for custom bounds): >>> from ssms.config import ModelConfigBuilder >>> model_config = ModelConfigBuilder.from_model("ddm", param_bounds=[...]) >>> gen = TrainingDataGenerator(generator_config, model_config) >>> data = gen.generate_data_training()

Advanced (custom pipeline, model_config from pipeline): >>> from ssms.dataset_generators.pipelines import SimulationPipeline >>> custom_pipeline = SimulationPipeline( ... generator_config, model_config, ... estimator_builder=MyCustomBuilder(...), ... training_strategy=MyCustomStrategy(...), ... ) >>> gen = TrainingDataGenerator(custom_pipeline) # model_config optional! >>> data = gen.generate_data_training()

generate_data_training

generate_data_training(save=False, verbose=False)

Generates training data for LANs.

Arguments

save: bool
    If True, the generated data is saved to disk.
verbose: bool
    If True, progress is printed to the console.

Returns:

data: dict –

Dictionary containing the generated data.

Notes

Phase 4 refactoring: This method now delegates to the
generation_pipeline, which handles the complete workflow from
parameter sampling to training data generation. This eliminates
wasteful simulations for analytical methods (PyDDM) and provides
a clean, modular architecture.

likelihood_estimators

Likelihood estimators for training data generation.

This module provides implementations of LikelihoodEstimatorProtocol: - KDELikelihoodEstimator: KDE-based estimator from simulated samples - Works with all models - Requires simulation data for fitting

PyDDMLikelihoodEstimator: Analytical PDF-based estimator using PyDDM
Only for compatible models (single-particle, two-choice, Gaussian noise)
Does not require simulation data (uses Fokker-Planck solution)
Deterministic results
Requires optional 'pyddm' package: pip install pyddm

KDELikelihoodEstimator

KDELikelihoodEstimator(displace_t=False)

KDE-based likelihood estimator.

This estimator builds a kernel density estimate from simulated (RT, choice) samples and uses it to estimate the likelihood P(RT, choice | theta).

Attributes:

displace_t (bool) –

Whether to displace time by the t parameter (non-decision time)
_kde (LogKDE | None) –

The fitted KDE object (None before fit() is called)
_metadata (dict | None) –

Metadata from simulations (None before fit() is called)

Examples:

>>> estimator = KDELikelihoodEstimator(displace_t=False)
>>> estimator.fit(simulations)
>>> log_liks = estimator.evaluate(rts, choices)
>>> samples = estimator.sample(n_samples=1000)

Arguments

displace_t : bool, optional Whether to displace time by the t parameter (non-decision time). Only works if all simulations have the same t value. Default is False.

evaluate

evaluate(rts, choices)

Evaluate log-likelihood at given (RT, choice) pairs.

Arguments

rts : np.ndarray Reaction times to evaluate, shape (n_samples,) choices : np.ndarray Choices to evaluate, shape (n_samples,)

Returns:

log_likelihoods ( ndarray ) –

Log-likelihood for each (RT, choice) pair, shape (n_samples,)

Raises:

ValueError –

If fit() has not been called yet

Notes

Extracted from lan_mlp.py lines 312, 342-344

fit

fit(simulations)

Build KDE from simulation data.

This method constructs a kernel density estimate from the provided simulations. Must be called before evaluate() or sample().

Arguments

simulations : dict Dictionary containing simulation data with keys: - 'rts': np.ndarray of reaction times - 'choices': np.ndarray of choices - 'metadata': dict with model information

Notes

Extracted from lan_mlp.py lines 305-308

get_metadata

get_metadata()

Return metadata from the simulations.

Returns:

metadata ( dict ) –

Dictionary containing simulation metadata including: - 'max_t': Maximum time value - 'possible_choices': List of possible choice values - Other model-specific metadata

Raises:

ValueError –

If fit() has not been called yet

sample

sample(n_samples, random_state=None)

Sample (RT, choice) pairs from the KDE.

Arguments

n_samples : int Number of samples to generate random_state : int | None, optional Random seed for reproducibility. If None, uses non-reproducible random behavior.

Returns:

samples ( dict ) –

Dictionary with keys: - 'rts': np.ndarray of sampled reaction times, shape (n_samples,) - 'choices': np.ndarray of sampled choices, shape (n_samples,)

Raises:

ValueError –

If fit() has not been called yet

Notes

Extracted from lan_mlp.py line 311

PyDDMLikelihoodEstimator

PyDDMLikelihoodEstimator(
    pyddm_solution, t_domain, interpolation="cubic"
)

Likelihood estimator using PyDDM's analytical Fokker-Planck solutions.

Unlike KDE estimators which resample from simulations, this estimator uses PyDDM's Fokker-Planck solver to compute analytical PDFs and interpolates them for fast likelihood evaluation.

Key advantages: - No simulation required (purely analytical) - Deterministic (no sampling variability) - Fast evaluation via interpolation - Exact solutions for compatible models

Limitations: - Only works for PyDDM-compatible models (single-particle, two-choice, Gaussian noise)

The estimator extracts PDFs from the solution internally and creates interpolators for fast likelihood evaluation.

Args: pyddm_solution: PyDDM Solution object from model.solve() t_domain: Time domain array from model.t_domain() interpolation: Interpolation method ('linear' or 'cubic')

evaluate

evaluate(rts, choices)

Evaluate log-likelihood at (RT, choice) pairs.

Uses interpolation for fast evaluation. RTs outside the time domain are assigned zero probability (log-likelihood = -inf, clipped to -66).

Args: rts: Reaction times, shape (n,) choices: Choices in {-1, 1}, shape (n,)

Returns: Log-likelihoods, shape (n,)

fit

fit(simulations)

Fit is a no-op for analytical estimator (already fitted).

This method exists for protocol compliance with LikelihoodEstimatorProtocol. PyDDM solutions are already "fitted" when the Fokker-Planck equation is solved.

Args: simulations: Ignored (not used by analytical estimator)

get_metadata

get_metadata()

Return metadata about the estimator.

Returns: Dictionary containing: - max_t: Maximum time in domain - dt: Time step - possible_choices: List of possible choice values - n_choices: Number of choices

sample

sample(n_samples, random_state=None)

Sample (RT, choice) pairs from the analytical distribution.

Uses PyDDM's built-in sampling method for efficient sampling.

Args: n_samples: Number of samples to generate random_state: Random seed for reproducibility. If None, uses non-reproducible random behavior.

Returns: Dictionary with keys: - 'rts': Reaction times, shape (n_samples,) - 'choices': Choices in {-1, 1}, shape (n_samples,)

Raises: RuntimeError: If unable to collect enough decided trials after max_attempts (indicates P(undecided) is too high)

Note: - PyDDM may produce "undecided" trials (process didn't reach boundary within T_dur), which appear as NaN. These are filtered out. - PyDDM's Sample object segregates trials by boundary (all upper, then all lower). We always shuffle to provide randomized samples consistent with the rest of ssm-simulators. - If P(undecided) is high, multiple sampling rounds may be needed to collect enough decided trials.

kde_estimator

KDE-based likelihood estimator.

This module implements likelihood estimation via Kernel Density Estimation (KDE). The estimator builds a KDE from simulated samples and uses it to evaluate likelihoods and generate new samples.

KDELikelihoodEstimator

KDELikelihoodEstimator(displace_t=False)

KDE-based likelihood estimator.

This estimator builds a kernel density estimate from simulated (RT, choice) samples and uses it to estimate the likelihood P(RT, choice | theta).

Attributes:

displace_t (bool) –

Whether to displace time by the t parameter (non-decision time)
_kde (LogKDE | None) –

The fitted KDE object (None before fit() is called)
_metadata (dict | None) –

Metadata from simulations (None before fit() is called)

Examples:

>>> estimator = KDELikelihoodEstimator(displace_t=False)
>>> estimator.fit(simulations)
>>> log_liks = estimator.evaluate(rts, choices)
>>> samples = estimator.sample(n_samples=1000)

Arguments

displace_t : bool, optional Whether to displace time by the t parameter (non-decision time). Only works if all simulations have the same t value. Default is False.

evaluate

evaluate(rts, choices)

Evaluate log-likelihood at given (RT, choice) pairs.

Arguments

rts : np.ndarray Reaction times to evaluate, shape (n_samples,) choices : np.ndarray Choices to evaluate, shape (n_samples,)

Returns:

log_likelihoods ( ndarray ) –

Log-likelihood for each (RT, choice) pair, shape (n_samples,)

Raises:

ValueError –

If fit() has not been called yet

Notes

Extracted from lan_mlp.py lines 312, 342-344

fit

fit(simulations)

Build KDE from simulation data.

This method constructs a kernel density estimate from the provided simulations. Must be called before evaluate() or sample().

Arguments

simulations : dict Dictionary containing simulation data with keys: - 'rts': np.ndarray of reaction times - 'choices': np.ndarray of choices - 'metadata': dict with model information

Notes

Extracted from lan_mlp.py lines 305-308

get_metadata

get_metadata()

Return metadata from the simulations.

Returns:

metadata ( dict ) –

Dictionary containing simulation metadata including: - 'max_t': Maximum time value - 'possible_choices': List of possible choice values - Other model-specific metadata

Raises:

ValueError –

If fit() has not been called yet

sample

sample(n_samples, random_state=None)

Sample (RT, choice) pairs from the KDE.

Arguments

n_samples : int Number of samples to generate random_state : int | None, optional Random seed for reproducibility. If None, uses non-reproducible random behavior.

Returns:

samples ( dict ) –

Dictionary with keys: - 'rts': np.ndarray of sampled reaction times, shape (n_samples,) - 'choices': np.ndarray of sampled choices, shape (n_samples,)

Raises:

ValueError –

If fit() has not been called yet

Notes

Extracted from lan_mlp.py line 311

pyddm_estimator

PyDDM-based likelihood estimator using Fokker-Planck solutions.

PyDDMLikelihoodEstimator

PyDDMLikelihoodEstimator(
    pyddm_solution, t_domain, interpolation="cubic"
)

Likelihood estimator using PyDDM's analytical Fokker-Planck solutions.

Unlike KDE estimators which resample from simulations, this estimator uses PyDDM's Fokker-Planck solver to compute analytical PDFs and interpolates them for fast likelihood evaluation.

Key advantages: - No simulation required (purely analytical) - Deterministic (no sampling variability) - Fast evaluation via interpolation - Exact solutions for compatible models

Limitations: - Only works for PyDDM-compatible models (single-particle, two-choice, Gaussian noise)

The estimator extracts PDFs from the solution internally and creates interpolators for fast likelihood evaluation.

Args: pyddm_solution: PyDDM Solution object from model.solve() t_domain: Time domain array from model.t_domain() interpolation: Interpolation method ('linear' or 'cubic')

evaluate

evaluate(rts, choices)

Evaluate log-likelihood at (RT, choice) pairs.

Uses interpolation for fast evaluation. RTs outside the time domain are assigned zero probability (log-likelihood = -inf, clipped to -66).

Args: rts: Reaction times, shape (n,) choices: Choices in {-1, 1}, shape (n,)

Returns: Log-likelihoods, shape (n,)

fit

fit(simulations)

Fit is a no-op for analytical estimator (already fitted).

This method exists for protocol compliance with LikelihoodEstimatorProtocol. PyDDM solutions are already "fitted" when the Fokker-Planck equation is solved.

Args: simulations: Ignored (not used by analytical estimator)

get_metadata

get_metadata()

Return metadata about the estimator.

Returns: Dictionary containing: - max_t: Maximum time in domain - dt: Time step - possible_choices: List of possible choice values - n_choices: Number of choices

sample

sample(n_samples, random_state=None)

Sample (RT, choice) pairs from the analytical distribution.

Uses PyDDM's built-in sampling method for efficient sampling.

Args: n_samples: Number of samples to generate random_state: Random seed for reproducibility. If None, uses non-reproducible random behavior.

Returns: Dictionary with keys: - 'rts': Reaction times, shape (n_samples,) - 'choices': Choices in {-1, 1}, shape (n_samples,)

Raises: RuntimeError: If unable to collect enough decided trials after max_attempts (indicates P(undecided) is too high)

Note: - PyDDM may produce "undecided" trials (process didn't reach boundary within T_dur), which appear as NaN. These are filtered out. - PyDDM's Sample object segregates trials by boundary (all upper, then all lower). We always shuffle to provide randomized samples consistent with the rest of ssm-simulators. - If P(undecided) is high, multiple sampling rounds may be needed to collect enough decided trials.

parameter_samplers

Parameter sampling infrastructure for training data generation.

ParameterSamplerProtocol

Bases: Protocol

Protocol for parameter sampling strategies.

This protocol defines the interface that all parameter samplers must implement. Samplers are responsible for generating parameter sets from a parameter space, potentially with dependencies between parameters and transformations applied.

get_param_space

get_param_space()

Get the parameter space bounds.

Returns: Dictionary mapping parameter names to (lower, upper) bound tuples

sample

sample(n_samples=1)

Sample n_samples parameter sets.

Args: n_samples: Number of parameter sets to sample

Returns: Dictionary mapping parameter names to sampled values (arrays of length n_samples)

ParameterSamplingConstraintProtocol

Bases: Protocol

Protocol for parameter sampling constraints.

Constraints are applied after sampling to enforce relationships or modify parameter values (e.g., ensuring a > z, normalizing drift rates to sum to 1). Constraints can trigger resampling if validation fails.

apply

apply(theta)

Apply constraint to sampled parameters.

Args: theta: Dictionary of parameter values (can be scalars or arrays)

Returns: Constrained parameter dictionary (modified in place or new dict)

UniformParameterSampler

UniformParameterSampler(param_space, constraints=None)

Bases: AbstractParameterSampler

Uniform random sampling of parameters.

This sampler generates parameter values uniformly distributed between lower and upper bounds. It matches the current behavior of the sample_parameters_from_constraints() function.

Example: >>> param_space = {"v": (-1.0, 1.0), "a": (0.5, 2.0)} >>> sampler = UniformParameterSampler(param_space) >>> samples = sampler.sample(n_samples=100) >>> samples["v"].shape (100,)

base_sampler

Base class for parameter samplers with dependency resolution.

AbstractParameterSampler

AbstractParameterSampler(param_space, constraints=None)

Bases: ABC

Base class for parameter samplers with dependency resolution.

This class handles: - Building dependency graphs from parameter bounds - Topological sorting to determine sampling order - Applying transforms after sampling - Validating parameter spaces

Subclasses must implement _sample_parameter() to define the specific sampling strategy (uniform, Sobol, Latin Hypercube, etc.).

Args: param_space: Dictionary mapping parameter names to (lower, upper) bounds. Bounds can be numeric or strings (for dependencies). constraints: List of constraint objects (must have apply() method). Applied after sampling in the order provided.

get_param_space

get_param_space()

Get the parameter space bounds.

Returns: Dictionary mapping parameter names to (lower, upper) bound tuples

sample

sample(n_samples=1, rng=None)

Sample parameters respecting dependencies and applying transforms.

Args: n_samples: Number of parameter sets to sample rng: Random number generator. If None, uses np.random.default_rng()

Returns: Dictionary mapping parameter names to sampled arrays

Raises: ValueError: If parameter dependencies cannot be resolved

constraints

Parameter sampling constraint classes.

NOTE: This module re-exports from ssms.transforms for backwards compatibility. Import directly from ssms.transforms instead:

from ssms.transforms import SwapIfLessConstraint, NormalizeToSumConstraint

NormalizeToSumConstraint

NormalizeToSumConstraint(param_names)

Bases: ParameterTransform

Normalize a set of parameters to sum to 1.

This constraint is commonly used for drift rates in multi-accumulator models (e.g., RLWM models) where the drift rates must sum to 1.

The normalization includes a small epsilon (1e-20) to avoid division by zero and ensure numerical stability.

Parameters:

param_names (list[str]) –

List of parameter names to normalize together

Examples:

>>> constraint = NormalizeToSumConstraint(["v1", "v2", "v3"])
>>> theta = {"v1": 0.2, "v2": 0.3, "v3": 0.5}
>>> result = constraint.apply(theta)
>>> sum([result["v1"], result["v2"], result["v3"]])
1.0

Use in model config:

>>> model_config = {
...     "name": "dev_rlwm_lba_race_v1",
...     "parameter_transforms": {
...         "sampling": [
...             NormalizeToSumConstraint(["vRL0", "vRL1", "vRL2"]),
...             NormalizeToSumConstraint(["vWM0", "vWM1", "vWM2"]),
...         ],
...     }
... }

Parameters:

param_names (list[str]) –

List of parameter names to normalize together

apply

apply(
    theta, model_config=None, n_trials=None, epsilon=1e-20
)

Normalize specified parameters to sum to 1.

Parameters:

theta (dict[str, Any]) –

Dictionary of parameter values (can be scalars or arrays)
model_config (dict[str, Any] or None, default: None ) –

Not used by this transform (included for interface compatibility)
n_trials (int or None, default: None ) –

Not used by this transform (included for interface compatibility)
epsilon (float, default: 1e-20 ) –

Small value to add to the total to avoid division by zero. Default is 1e-20.

Returns:

dict[str, Any] –

Modified theta dictionary with normalized values

SwapIfLessConstraint

SwapIfLessConstraint(param_a, param_b)

Bases: ParameterTransform

Swap two parameters if first is less than or equal to second.

This constraint ensures ordering relationships like a > z by swapping the values when the constraint is violated. Works with both scalar and array inputs.

Parameters:

param_a (str) –

First parameter name (should be > param_b after constraint)
param_b (str) –

Second parameter name (should be < param_a after constraint)

Examples:

>>> constraint = SwapIfLessConstraint("a", "z")
>>> theta = {"a": np.array([0.5, 1.5]), "z": np.array([1.0, 0.5])}
>>> result = constraint.apply(theta)
>>> # First element: a=0.5 <= z=1.0 → swap → a=1.0, z=0.5
>>> # Second element: a=1.5 > z=0.5 → no swap → a=1.5, z=0.5

Use in model config:

>>> model_config = {
...     "name": "lba_angle_3",
...     "parameter_transforms": {
...         "sampling": [SwapIfLessConstraint("a", "z")],
...     }
... }

Parameters:

param_a (str) –

First parameter name (should be > param_b after constraint)
param_b (str) –

Second parameter name (should be < param_a after constraint)

apply

apply(theta, model_config=None, n_trials=None)

Swap param_a and param_b where param_a <= param_b.

Parameters:

theta (dict[str, Any]) –

Dictionary of parameter values (can be scalars or arrays)
model_config (dict[str, Any] or None, default: None ) –

Not used by this transform (included for interface compatibility)
n_trials (int or None, default: None ) –

Not used by this transform (included for interface compatibility)

Returns:

dict[str, Any] –

Modified theta dictionary with swapped values where needed

protocols

Protocols for parameter sampling and constraints.

ParameterSamplerProtocol

Bases: Protocol

Protocol for parameter sampling strategies.

This protocol defines the interface that all parameter samplers must implement. Samplers are responsible for generating parameter sets from a parameter space, potentially with dependencies between parameters and transformations applied.

get_param_space

get_param_space()

Get the parameter space bounds.

Returns: Dictionary mapping parameter names to (lower, upper) bound tuples

sample

sample(n_samples=1)

Sample n_samples parameter sets.

Args: n_samples: Number of parameter sets to sample

Returns: Dictionary mapping parameter names to sampled values (arrays of length n_samples)

ParameterSamplingConstraintProtocol

Bases: Protocol

Protocol for parameter sampling constraints.

Constraints are applied after sampling to enforce relationships or modify parameter values (e.g., ensuring a > z, normalizing drift rates to sum to 1). Constraints can trigger resampling if validation fails.

apply

apply(theta)

Apply constraint to sampled parameters.

Args: theta: Dictionary of parameter values (can be scalars or arrays)

Returns: Constrained parameter dictionary (modified in place or new dict)

uniform_sampler

Uniform random parameter sampler (default behavior).

UniformParameterSampler

UniformParameterSampler(param_space, constraints=None)

Bases: AbstractParameterSampler

Uniform random sampling of parameters.

This sampler generates parameter values uniformly distributed between lower and upper bounds. It matches the current behavior of the sample_parameters_from_constraints() function.

Example: >>> param_space = {"v": (-1.0, 1.0), "a": (0.5, 2.0)} >>> sampler = UniformParameterSampler(param_space) >>> samples = sampler.sample(n_samples=100) >>> samples["v"].shape (100,)

pipelines

Data generation pipelines for orchestrating end-to-end workflows.

PyDDMPipeline

PyDDMPipeline(
    generator_config,
    model_config,
    estimator_builder,
    training_strategy,
)

Generate training data via PyDDM analytical PDF workflow.

This strategy implements an analytical workflow: 1. Sample parameters from parameter space 2. SKIP simulation (use analytical Fokker-Planck solution) 3. Build PyDDM estimator (solves PDE) 4. Compute auxiliary labels analytically 5. Generate training data from analytical PDF

Used for: PyDDM-based likelihood estimation

Benefits: - Much faster (no simulation overhead) - No filtering needed (analytical solution always valid) - Deterministic for given parameters - Memory efficient (no large simulation arrays) - Computes choice probabilities, omission_p, and nogo_p analytically

Limitations: - Only works for compatible models (single-particle, two-choice, Gaussian noise) - Cannot generate binned RT histograms (these require trajectory simulations)

Args: generator_config: Configuration for data generation model_config: Model specification (must contain 'param_bounds_dict') estimator_builder: Builder for PyDDM estimator (class or instance) - If a class is passed, it will be instantiated with generator_config - If an instance is passed, it will be used as-is training_strategy: Strategy for generating training samples (class or instance) - If a class is passed, it will be instantiated with generator_config and model_config - If an instance is passed, it will be used as-is

generate_for_parameter_set

generate_for_parameter_set(
    parameter_sampling_seed, random_seed=None
)

Generate training data for one parameter set (analytical workflow).

Workflow: 1. Sample theta from parameter space 2. Build analytical estimator (no simulations needed) 3. Compute auxiliary labels analytically 4. Generate training data

Args: parameter_sampling_seed: Index for parameter sampling (used as seed) random_seed: Random seed (accepted for API compatibility, may not be used as analytical solutions are deterministic)

Returns: Dictionary with keys: - 'data': Training data dictionary (or None if generation failed) - 'theta': Parameter values used - 'success': Whether generation succeeded - 'error': Error message if success=False (optional)

Note: PyDDM can compute choice probabilities, omission_p, and nogo_p analytically without simulations. Only binned RT histograms require trajectory data and are set to None.

get_param_space

get_param_space()

Get the parameter space dictionary.

SimulationPipeline

SimulationPipeline(
    generator_config,
    model_config,
    estimator_builder,
    training_strategy,
)

Generate training data via simulation + KDE workflow.

This pipeline implements the traditional workflow: 1. Sample parameters from parameter space 2. Run Cython simulator to generate samples 3. Filter simulations for quality (RT variance, choice proportions, etc.) 4. Build KDE estimator from simulations 5. Generate training data from KDE

Used for: KDE-based likelihood estimation

This pipeline is necessary when analytical solutions are not available or when you want to validate analytical methods against simulations.

Args: generator_config: Configuration for data generation model_config: Model specification (must contain 'param_bounds_dict') estimator_builder: Builder for KDE estimator (class or instance) - If a class is passed, it will be instantiated with generator_config - If an instance is passed, it will be used as-is training_strategy: Strategy for generating training samples (class or instance) - If a class is passed, it will be instantiated with generator_config and model_config - If an instance is passed, it will be used as-is

generate_for_parameter_set

generate_for_parameter_set(
    parameter_sampling_seed=None, simulator_seed=None
)

Generate training data for one parameter set (simulation workflow).

Workflow: 1. Sample theta from parameter space 2. Run simulations 3. Filter simulations 4. Build estimator 5. Generate training data

Args: parameter_sampling_seed: Seed for parameter sampling (ensures different workers sample different θ) simulator_seed: Random seed for simulations (controls RT/choice variability)

Returns: Dictionary with keys: - 'data': Training data dictionary (or None if generation failed) - 'theta': Parameter values used - 'success': Whether generation succeeded

get_param_space

get_param_space()

Get the parameter space dictionary.

create_data_generation_pipeline

create_data_generation_pipeline(
    generator_config, model_config
)

Create appropriate data generation pipeline based on configuration.

This factory selects the right pipeline based on estimator_type in generator_config: - 'kde': SimulationPipeline (runs simulations + KDE) - 'pyddm': PyDDMPipeline (analytical Fokker-Planck solution)

The factory internally creates all required components (estimator builder, training strategy, parameter sampler) from the provided configurations. The model_config serves as the single source of truth for model specifications, parameter bounds, and transformations.

Args: generator_config: Configuration dict with 'estimator.type' key model_config: Model specification (must contain 'param_bounds_dict')

Returns: Data generation pipeline instance

Raises: ValueError: If estimator_type is unknown

Note: The same model_config can be used for both simulation-based (KDE) and analytical (PyDDM) pipelines. Custom boundaries and drift functions should be configured through the model_config or using ModelConfigBuilder.

Example: >>> pipeline = create_data_generation_pipeline( ... generator_config=config, ... model_config=model_config, ... ) >>> # Pipeline is ready to use with all components created internally

pipeline_factory

Factory for creating data generation pipelines.

create_data_generation_pipeline

create_data_generation_pipeline(
    generator_config, model_config
)

Create appropriate data generation pipeline based on configuration.

This factory selects the right pipeline based on estimator_type in generator_config: - 'kde': SimulationPipeline (runs simulations + KDE) - 'pyddm': PyDDMPipeline (analytical Fokker-Planck solution)

The factory internally creates all required components (estimator builder, training strategy, parameter sampler) from the provided configurations. The model_config serves as the single source of truth for model specifications, parameter bounds, and transformations.

Args: generator_config: Configuration dict with 'estimator.type' key model_config: Model specification (must contain 'param_bounds_dict')

Returns: Data generation pipeline instance

Raises: ValueError: If estimator_type is unknown

Note: The same model_config can be used for both simulation-based (KDE) and analytical (PyDDM) pipelines. Custom boundaries and drift functions should be configured through the model_config or using ModelConfigBuilder.

Example: >>> pipeline = create_data_generation_pipeline( ... generator_config=config, ... model_config=model_config, ... ) >>> # Pipeline is ready to use with all components created internally

pyddm_pipeline

PyDDM-based data generation pipeline using analytical solutions.

PyDDMPipeline

PyDDMPipeline(
    generator_config,
    model_config,
    estimator_builder,
    training_strategy,
)

Generate training data via PyDDM analytical PDF workflow.

This strategy implements an analytical workflow: 1. Sample parameters from parameter space 2. SKIP simulation (use analytical Fokker-Planck solution) 3. Build PyDDM estimator (solves PDE) 4. Compute auxiliary labels analytically 5. Generate training data from analytical PDF

Used for: PyDDM-based likelihood estimation

Benefits: - Much faster (no simulation overhead) - No filtering needed (analytical solution always valid) - Deterministic for given parameters - Memory efficient (no large simulation arrays) - Computes choice probabilities, omission_p, and nogo_p analytically

Limitations: - Only works for compatible models (single-particle, two-choice, Gaussian noise) - Cannot generate binned RT histograms (these require trajectory simulations)

Args: generator_config: Configuration for data generation model_config: Model specification (must contain 'param_bounds_dict') estimator_builder: Builder for PyDDM estimator (class or instance) - If a class is passed, it will be instantiated with generator_config - If an instance is passed, it will be used as-is training_strategy: Strategy for generating training samples (class or instance) - If a class is passed, it will be instantiated with generator_config and model_config - If an instance is passed, it will be used as-is

generate_for_parameter_set

generate_for_parameter_set(
    parameter_sampling_seed, random_seed=None
)

Generate training data for one parameter set (analytical workflow).

Workflow: 1. Sample theta from parameter space 2. Build analytical estimator (no simulations needed) 3. Compute auxiliary labels analytically 4. Generate training data

Args: parameter_sampling_seed: Index for parameter sampling (used as seed) random_seed: Random seed (accepted for API compatibility, may not be used as analytical solutions are deterministic)

Returns: Dictionary with keys: - 'data': Training data dictionary (or None if generation failed) - 'theta': Parameter values used - 'success': Whether generation succeeded - 'error': Error message if success=False (optional)

Note: PyDDM can compute choice probabilities, omission_p, and nogo_p analytically without simulations. Only binned RT histograms require trajectory data and are set to None.

get_param_space

get_param_space()

Get the parameter space dictionary.

simulation_pipeline

Simulation-based data generation pipeline for KDE estimators.

SimulationPipeline

SimulationPipeline(
    generator_config,
    model_config,
    estimator_builder,
    training_strategy,
)

Generate training data via simulation + KDE workflow.

This pipeline implements the traditional workflow: 1. Sample parameters from parameter space 2. Run Cython simulator to generate samples 3. Filter simulations for quality (RT variance, choice proportions, etc.) 4. Build KDE estimator from simulations 5. Generate training data from KDE

Used for: KDE-based likelihood estimation

This pipeline is necessary when analytical solutions are not available or when you want to validate analytical methods against simulations.

Args: generator_config: Configuration for data generation model_config: Model specification (must contain 'param_bounds_dict') estimator_builder: Builder for KDE estimator (class or instance) - If a class is passed, it will be instantiated with generator_config - If an instance is passed, it will be used as-is training_strategy: Strategy for generating training samples (class or instance) - If a class is passed, it will be instantiated with generator_config and model_config - If an instance is passed, it will be used as-is

generate_for_parameter_set

generate_for_parameter_set(
    parameter_sampling_seed=None, simulator_seed=None
)

Generate training data for one parameter set (simulation workflow).

Workflow: 1. Sample theta from parameter space 2. Run simulations 3. Filter simulations 4. Build estimator 5. Generate training data

Args: parameter_sampling_seed: Seed for parameter sampling (ensures different workers sample different θ) simulator_seed: Random seed for simulations (controls RT/choice variability)

Returns: Dictionary with keys: - 'data': Training data dictionary (or None if generation failed) - 'theta': Parameter values used - 'success': Whether generation succeeded

get_param_space

get_param_space()

Get the parameter space dictionary.

protocols

Protocols for data generation components.

This module defines the interfaces (protocols) for: - Likelihood estimators: Objects that estimate P(RT, choice | theta) - Estimator builders: Objects that create and configure likelihood estimators - Training data strategies: Objects that generate training samples from estimators - Data generation strategies: Objects that orchestrate the end-to-end workflow - Parameter samplers: Objects that sample parameters from parameter spaces

Using protocols allows for dependency injection and clean separation of concerns.

DataGenerationPipelineProtocol

Bases: Protocol

Protocol for end-to-end data generation workflows.

Data generation pipelines orchestrate the complete process from parameter sampling to training data generation. Different pipelines handle different workflows: - Simulation-based: Run simulations, filter, build KDE, generate data - Analytical (PyDDM): Skip simulations, solve Fokker-Planck, generate data - Hybrid: Mix simulation and analytical approaches - Other: Custom workflows

This protocol enables clean separation between simulation-dependent logic (SimulationPipeline) and analytical approaches (PyDDMPipeline), eliminating wasteful computation.

Attributes:

generator_config (dict) –

Configuration for data generation (samples, bins, etc.)
model_config (dict) –

Model specification (parameters, bounds, transforms, etc.)

generate_for_parameter_set

generate_for_parameter_set(
    parameter_sampling_seed, random_seed=None
)

Generate training data for one parameter set.

Arguments

parameter_sampling_seed : int Index of parameter set to generate (used to sample from param space) random_seed : int | None Random seed for reproducibility (may be unused for deterministic methods)

Returns:

result ( dict ) –

Dictionary with keys: - 'data': np.ndarray of training data (or None if generation failed) - 'theta': dict of parameter values used - 'success': bool indicating whether generation succeeded - 'error': str (optional) error message if success=False

get_param_space

get_param_space()

Get the parameter space object used by this strategy.

Returns:

param_space ( Any ) –

Parameter space object for sampling theta values

EstimatorBuilderProtocol

Bases: Protocol

Protocol for building likelihood estimators.

Builders encapsulate all the logic for creating and configuring likelihood estimators. This keeps the TrainingDataGenerator class generic and free of type-specific logic.

build

build(theta, simulations=None)

Build and return a fitted likelihood estimator for given theta.

Arguments

theta : dict Model parameters (e.g., {'v': 1.0, 'a': 2.0, 'z': 0.5, 't': 0.3}) simulations : dict | None Simulation data (required for KDE, may be None for analytical estimators)

Returns:

estimator ( LikelihoodEstimatorProtocol ) –

A fitted likelihood estimator ready for use

Raises:

ValueError –

If required data (e.g., simulations for KDE) is missing

LikelihoodEstimatorProtocol

Bases: Protocol

Protocol for likelihood estimation.

Likelihood estimators compute P(RT, choice | theta) and can sample from this distribution. They may use different methods: - KDE: Build kernel density estimate from simulated samples - Analytical: Use closed-form PDF (e.g., from PyDDM solver) - Other: Custom methods

evaluate

evaluate(rts, choices)

Evaluate log-likelihood at given (RT, choice) pairs.

Arguments

rts : np.ndarray Reaction times to evaluate, shape (n_samples,) choices : np.ndarray Choices to evaluate, shape (n_samples,)

Returns:

log_likelihoods ( ndarray ) –

Log-likelihood for each (RT, choice) pair, shape (n_samples,)

fit

fit(simulations)

Fit estimator to simulation data.

For KDE: Builds the kernel density estimate. For analytical: May be a no-op (PDF already available).

Arguments

simulations : dict Dictionary containing simulation data with keys: - 'rts': np.ndarray of reaction times - 'choices': np.ndarray of choices - 'metadata': dict with model info

get_metadata

get_metadata()

Return metadata about the estimator.

Returns:

metadata ( dict ) –

Dictionary containing: - 'max_t': Maximum time value - 'possible_choices': List of possible choice values - Other model-specific metadata

sample

sample(n_samples, random_state=None)

Sample (RT, choice) pairs from the estimated likelihood.

Arguments

n_samples : int Number of samples to generate random_state : int | None, optional Random seed for reproducibility. If None, uses non-reproducible random behavior.

Returns:

samples ( dict ) –

Dictionary with keys: - 'rts': np.ndarray of sampled reaction times - 'choices': np.ndarray of sampled choices

ParameterSamplerProtocol

Bases: Protocol

Protocol for parameter sampling strategies.

Parameter samplers generate parameter sets from a parameter space, handling dependencies between parameters and applying transforms.

get_param_space

get_param_space()

Get the parameter space bounds.

Returns: Dictionary mapping parameter names to (lower, upper) bound tuples

sample

sample(n_samples=1, rng=None)

Sample n_samples parameter sets.

Args: n_samples: Number of parameter sets to sample

Returns: Dictionary mapping parameter names to sampled values (arrays)

TrainingDataStrategyProtocol

Bases: Protocol

Protocol for training data generation strategies.

Strategies define how to generate training samples from a likelihood estimator. Different strategies might: - Resample from estimator + add uniform samples (current approach) - Evaluate on a uniform grid - Use importance sampling - Other methods

generate

generate(theta, likelihood_estimator, random_state=None)

Generate training data array.

Arguments

theta : dict Model parameters likelihood_estimator : LikelihoodEstimatorProtocol Fitted likelihood estimator random_state : int | None, optional Random seed for reproducibility. If None, uses non-reproducible random behavior. Controls all random operations including sampling and row shuffling.

Returns:

training_data ( ndarray ) –

Array of shape (n_samples, n_params + 3) containing: - Columns 0:n_params: Parameter values (tiled for all samples) - Column -3: Reaction times - Column -2: Choices - Column -1: Log-likelihoods

Note: Rows are shuffled to avoid ordering bias in training batches.

strategies

Training data generation strategies.

This module contains strategy classes for training data generation:

ResampleMixtureStrategy: Generates training data by mixing samples from a likelihood estimator (KDE or PyDDM) with uniform samples in positive and negative RT space. This helps networks learn both the likelihood surface and its boundaries.
MixtureTrainingStrategy: Alias for ResampleMixtureStrategy (for backward compatibility).

Note: For end-to-end data generation workflows, see the pipelines module which contains SimulationPipeline and PyDDMPipeline.

ResampleMixtureStrategy

ResampleMixtureStrategy(generator_config, model_config)

Training data strategy using mixture of resamples and uniform draws.

This strategy generates training data by: 1. Resampling (RT, choice) pairs from the fitted likelihood estimator 2. Generating uniform samples in positive RT space [0, max_t] 3. Generating uniform samples in negative RT space [-1, 0]

For each sample, the strategy evaluates the log-likelihood and combines: - Parameters (theta) - RT and choice - Log-likelihood at (RT, choice)

Attributes:

generator_config (dict) –

Configuration for data generation (sample counts, mixture probabilities, etc.)
model_config (dict) –

Model configuration (params, nchoices, etc.)

Examples:

>>> strategy = ResampleMixtureStrategy(generator_config, model_config)
>>> training_data = strategy.generate(theta, likelihood_estimator)
>>> assert training_data.shape == (n_training_samples, n_features)

Arguments

generator_config : dict Configuration dictionary containing: - ['training']['n_samples_per_param']: Total samples to generate - ['training']['mixture_probabilities']: [p_estimator, p_unif_up, p_unif_down] - ['training']['separate_response_channels']: Whether to one-hot encode choices - ['training']['negative_rt_log_likelihood']: Log-likelihood value for negative RTs model_config : dict Model configuration containing: - 'params': List of parameter names - 'nchoices': Number of possible choices

generate

generate(theta, likelihood_estimator, random_state=None)

Generate training data using mixture strategy.

Arguments

theta : dict Parameter dictionary (e.g., {'v': 1.0, 'a': 2.0, ...}) likelihood_estimator : LikelihoodEstimatorProtocol Fitted likelihood estimator (must support sample() and evaluate()) random_state : int | None, optional Random seed for reproducibility. If None, uses non-reproducible random behavior. Controls all random operations including sampling from estimator and shuffling output.

Returns:

ndarray –

Training data array of shape (n_samples, n_features) where: - First columns: theta parameters - Middle columns: RT and choice (possibly one-hot encoded) - Last column: log-likelihood

Note: Rows are shuffled to avoid ordering bias in training batches.

Raises:

ValueError –

If mixture probabilities don't lead to valid sample counts

Notes

Extracted from lan_mlp.py lines 268-360

resample_mixture_strategy

Mixture training data generation strategy.

This strategy generates training data by mixing three types of samples: 1. Samples drawn from the likelihood estimator (KDE or analytical PDF) 2. Uniform samples in positive RT space (to encourage exploration) 3. Uniform samples in negative RT space (to learn the hard boundary at t=0)

This approach helps the network learn both the likelihood surface and its boundaries.

Note: MixtureTrainingStrategy is provided as an alias for backward compatibility.

ResampleMixtureStrategy

ResampleMixtureStrategy(generator_config, model_config)

Training data strategy using mixture of resamples and uniform draws.

This strategy generates training data by: 1. Resampling (RT, choice) pairs from the fitted likelihood estimator 2. Generating uniform samples in positive RT space [0, max_t] 3. Generating uniform samples in negative RT space [-1, 0]

For each sample, the strategy evaluates the log-likelihood and combines: - Parameters (theta) - RT and choice - Log-likelihood at (RT, choice)

Attributes:

generator_config (dict) –

Configuration for data generation (sample counts, mixture probabilities, etc.)
model_config (dict) –

Model configuration (params, nchoices, etc.)

Examples:

>>> strategy = ResampleMixtureStrategy(generator_config, model_config)
>>> training_data = strategy.generate(theta, likelihood_estimator)
>>> assert training_data.shape == (n_training_samples, n_features)

Arguments

generator_config : dict Configuration dictionary containing: - ['training']['n_samples_per_param']: Total samples to generate - ['training']['mixture_probabilities']: [p_estimator, p_unif_up, p_unif_down] - ['training']['separate_response_channels']: Whether to one-hot encode choices - ['training']['negative_rt_log_likelihood']: Log-likelihood value for negative RTs model_config : dict Model configuration containing: - 'params': List of parameter names - 'nchoices': Number of possible choices

generate

generate(theta, likelihood_estimator, random_state=None)

Generate training data using mixture strategy.

Arguments

theta : dict Parameter dictionary (e.g., {'v': 1.0, 'a': 2.0, ...}) likelihood_estimator : LikelihoodEstimatorProtocol Fitted likelihood estimator (must support sample() and evaluate()) random_state : int | None, optional Random seed for reproducibility. If None, uses non-reproducible random behavior. Controls all random operations including sampling from estimator and shuffling output.

Returns:

ndarray –

Training data array of shape (n_samples, n_features) where: - First columns: theta parameters - Middle columns: RT and choice (possibly one-hot encoded) - Last column: log-likelihood

Note: Rows are shuffled to avoid ordering bias in training batches.

Raises:

ValueError –

If mixture probabilities don't lead to valid sample counts

Notes

Extracted from lan_mlp.py lines 268-360

ssms.external_simulators

External simulator interfaces for ssms.

This module provides wrappers and mappers for integrating external simulation and modeling packages with ssms.

SSMSToPyDDMMapper

Converts ssms model configurations to PyDDM models for analytical Fokker-Planck solutions.

Tested compatibility: - DDM and variants (ddm, ddm_par2, ddm_seq2, ddm_mic2_) - Ornstein models (position-dependent drift) - Collapsing boundary models (angle, weibull) - Custom drift models (gamma_drift, conflict_, attend_, shrink_spot_)

Features: - Handles parameter conflicts ('t' non-decision time vs time variable) - Returns array-compatible drift/boundary functions - Automatic z transformation from [0,1] to [-a, a] - Uses modern PyDDM API (solution.pdf("correct"))

Incompatible models: - Multi-particle: race_, lca_, lba* - Non-Gaussian noise: levy - Inter-trial variability: full_ddm, ddm_sdv, ddm_st

SSMSToPyDDMMapper

Maps ssms model configurations to PyDDM models.

PyDDM can solve the Fokker-Planck equation for: - Single-particle, two-choice models - Gaussian noise only - Arbitrary drift functions (time-dependent, position-dependent, parameter-dependent) - Arbitrary boundary functions (time-dependent, parameter-dependent)

Compatible models include: - ddm and variants (without sv/sz/st) - angle, weibull (collapsing boundaries) - ornstein (leaky/unstable integration) - gamma_drift (time-dependent drift) - conflict_, attend_, shrink_spot_* (custom drift models)

Incompatible models: - race_, lca_, lba* (multi-particle) - levy (non-Gaussian noise) - full_ddm, ddm_sdv, ddm_st (inter-trial variability)

build_pyddm_model `classmethod`

build_pyddm_model(
    model_config, theta, generator_config=None
)

Build PyDDM model for specific parameter values.

Args: model_config: Model structure (name, params, drift_config, boundary_config) theta: Specific parameter values for this instance generator_config: Simulation settings (delta_t, max_t). If None, uses defaults: {'delta_t': 0.001, 'max_t': 20.0}

Returns: pyddm.Model instance ready to solve()

Raises: ValueError: If model is not compatible with PyDDM ImportError: If pyddm package is not installed

create_boundary_function `classmethod`

create_boundary_function(model_config)

Create PyDDM boundary function from ssms model config.

Returns callable with signature: boundary(t, **theta)

Architecture: - Standard model configs specify boundary_name and boundary function - Metadata (params, multiplicative) is looked up from the boundary registry - Direct boundary_config dict is supported for custom runtime configurations

This design avoids duplicating metadata across 100+ model configs and ensures consistent behavior for all instances of a given boundary type (e.g., all "angle" boundaries work the same way).

Args: model_config: Model configuration dictionary from ssms.config

Returns: Callable that accepts (t, **theta) and returns boundary value

create_drift_function `classmethod`

create_drift_function(model_config)

Create PyDDM drift function from ssms model config.

Returns callable with signature: drift(t, x, **theta)

Handles: - Constant drift (v) - Position-dependent drift (Ornstein: v - g*x) - Custom time-dependent drift (gamma_drift, conflict models, etc.)

Architecture: - Standard model configs specify drift_name and drift_fun - Metadata (params) is looked up from the drift registry - Direct drift_config dict is supported for custom runtime configurations

This design avoids duplicating metadata across 100+ model configs.

Args: model_config: Model configuration dictionary from ssms.config

Returns: Callable that accepts (t, x, **theta) and returns drift rate

is_compatible `classmethod`

is_compatible(model_config)

Check if model can use PyDDM analytical solver.

Args: model_config: Model configuration dictionary from ssms.config

Returns: (is_compatible, reason_if_not)

transform_z_to_x0 `classmethod`

transform_z_to_x0(z, bound_at_t0, safety_margin=0.99)

Transform z from [0,1] to PyDDM starting position within valid bounds.

Args: z: Starting position in ssms format (0=lower, 0.5=center, 1=upper) bound_at_t0: Boundary value at t=0 (defines valid range) safety_margin: Scale factor to keep x0 away from exact boundaries (default 0.99)

Returns: Starting position in PyDDM format, within [-bound_at_t0, bound_at_t0]

Note: PyDDM requires starting position to be strictly within the bounds at t=0. The safety margin ensures x0 is not at the exact boundary, avoiding discretization issues with PyDDM's spatial grid. Formula: x0 = (2z - 1) * bound(0) * safety_margin

pyddm_mapper

Mapper for converting ssms model configurations to PyDDM models.

This module provides utilities for mapping ssms model configurations to PyDDM models, enabling the use of PyDDM's analytical Fokker-Planck solver for compatible models.

SSMSToPyDDMMapper

Maps ssms model configurations to PyDDM models.

PyDDM can solve the Fokker-Planck equation for: - Single-particle, two-choice models - Gaussian noise only - Arbitrary drift functions (time-dependent, position-dependent, parameter-dependent) - Arbitrary boundary functions (time-dependent, parameter-dependent)

Compatible models include: - ddm and variants (without sv/sz/st) - angle, weibull (collapsing boundaries) - ornstein (leaky/unstable integration) - gamma_drift (time-dependent drift) - conflict_, attend_, shrink_spot_* (custom drift models)

Incompatible models: - race_, lca_, lba* (multi-particle) - levy (non-Gaussian noise) - full_ddm, ddm_sdv, ddm_st (inter-trial variability)

build_pyddm_model `classmethod`

build_pyddm_model(
    model_config, theta, generator_config=None
)

Build PyDDM model for specific parameter values.

Args: model_config: Model structure (name, params, drift_config, boundary_config) theta: Specific parameter values for this instance generator_config: Simulation settings (delta_t, max_t). If None, uses defaults: {'delta_t': 0.001, 'max_t': 20.0}

Returns: pyddm.Model instance ready to solve()

Raises: ValueError: If model is not compatible with PyDDM ImportError: If pyddm package is not installed

create_boundary_function `classmethod`

create_boundary_function(model_config)

Create PyDDM boundary function from ssms model config.

Returns callable with signature: boundary(t, **theta)

Architecture: - Standard model configs specify boundary_name and boundary function - Metadata (params, multiplicative) is looked up from the boundary registry - Direct boundary_config dict is supported for custom runtime configurations

This design avoids duplicating metadata across 100+ model configs and ensures consistent behavior for all instances of a given boundary type (e.g., all "angle" boundaries work the same way).

Args: model_config: Model configuration dictionary from ssms.config

Returns: Callable that accepts (t, **theta) and returns boundary value

create_drift_function `classmethod`

create_drift_function(model_config)

Create PyDDM drift function from ssms model config.

Returns callable with signature: drift(t, x, **theta)

Handles: - Constant drift (v) - Position-dependent drift (Ornstein: v - g*x) - Custom time-dependent drift (gamma_drift, conflict models, etc.)

Architecture: - Standard model configs specify drift_name and drift_fun - Metadata (params) is looked up from the drift registry - Direct drift_config dict is supported for custom runtime configurations

This design avoids duplicating metadata across 100+ model configs.

Args: model_config: Model configuration dictionary from ssms.config

Returns: Callable that accepts (t, x, **theta) and returns drift rate

is_compatible `classmethod`

is_compatible(model_config)

Check if model can use PyDDM analytical solver.

Args: model_config: Model configuration dictionary from ssms.config

Returns: (is_compatible, reason_if_not)

transform_z_to_x0 `classmethod`

transform_z_to_x0(z, bound_at_t0, safety_margin=0.99)

Transform z from [0,1] to PyDDM starting position within valid bounds.

Args: z: Starting position in ssms format (0=lower, 0.5=center, 1=upper) bound_at_t0: Boundary value at t=0 (defines valid range) safety_margin: Scale factor to keep x0 away from exact boundaries (default 0.99)

Returns: Starting position in PyDDM format, within [-bound_at_t0, bound_at_t0]

Note: PyDDM requires starting position to be strictly within the bounds at t=0. The safety margin ensures x0 is not at the exact boundary, avoiding discretization issues with PyDDM's spatial grid. Formula: x0 = (2z - 1) * bound(0) * safety_margin

ssms.hssm_support

decorate_atomic_simulator

decorate_atomic_simulator(
    model_name, choices=None, obs_dim=2
)

Decorator to add metadata attributes to simulator functions.

This decorator attaches the following attributes to the decorated function as expected of simulators in HSSM: - model_name: Name of the model. - choices: List or array of possible choices/responses. - obs_dim: Number of observation dimensions.

Parameters:

model_name (str) –

Name of the model.
choices (list or ndarray, default: None ) –

List or array of possible choices/responses (default: [-1, 1]).
obs_dim (int, default: 2 ) –

Number of observation dimensions (default: 2).

Returns:

Callable –

Decorator that adds attributes to the simulator function.

get_simulator_fun_internal

get_simulator_fun_internal(simulator_fun)

Get the internal simulator function for a given model.

Parameters:

simulator_fun (Callable or str) –

The simulator function or the name of the model as a string.

Returns:

Callable –

The decorated simulator function.

Raises:

ValueError –

If the simulator argument is not a string or a callable.

hssm_sim_wrapper

hssm_sim_wrapper(
    simulator_fun,
    theta,
    model,
    n_replicas,
    random_state,
    **kwargs
)

Wrap a ssms simulator function to match HSSM's expected interface.

Parameters:

simulator_fun (callable) –

The simulator function to wrap, which should have the following interface: - theta: array-like, shape (n_trials, n_parameters) - model: str, name of the model to simulate - n_samples: int, number of replica datasets to generate - random_state: int, to be used as the random seed internally - **kwargs: additional keyword arguments
theta (array - like) –

Model parameters, shape (n_trials, n_parameters)
model (str) –

Name of the model to simulate
n_replicas (int) –

Number of replica datasets to generate
random_state (int or Generator) –

Random seed or random number generator
**kwargs –

Additional keyword arguments passed to simulator_fun

Returns:

array - like –

Array of shape (n_trials, 2) containing reaction times and choices stacked column-wise

rng_fn

rng_fn(
    arg_arrays,
    size,
    rng,
    simulator_fun,
    obs_dim_int,
    *args,
    **kwargs
)

Generate random variables from this distribution using the provided simulator function.

Parameters:

arg_arrays (list of np.ndarray) –

List of argument arrays corresponding to model parameters.
size (int, tuple, or None) –

The total number of samples to be drawn. If None or 1, only one replica
rng (Generator) –

Random number generator for reproducibility.
simulator_fun (Callable) –

The simulator function to generate samples.
obs_dim_int (int) –

Number of observation dimensions.
*args (tuple, default: () ) –

Model parameters, in the order of _list_params, with the last argument as size.
**kwargs (dict, default: {} ) –

Additional keyword arguments passed to the simulator function.

Returns:

tuple[ndarray, ndarray] –

An array of shape (..., obs_dim_int) containing generated (rt, response) pairs and the p_outlier values if applicable.

validate_simulator_fun

validate_simulator_fun(simulator_fun)

Validate that the simulator function has required attributes.

Parameters:

simulator_fun (Any) –

The simulator function or object to validate.

Returns:

tuple –

A tuple containing model_name, choices, and obs_dim_int.

Raises:

ValueError –

If any required attribute is missing or invalid.

ssms.rl

RLSSM simulation framework for ssm-simulators.

AssembledModel `dataclass`

AssembledModel(
    config,
    learning_backend,
    gradient,
    model_name,
    decision_process,
    list_params,
    bounds,
    params_default,
    response,
    choices,
    context_fields,
    computed_params,
    response_to_choice,
)

Validated executable form of an RLSSM ModelConfig.

The assembled model exposes package-neutral metadata and pure Python/JAX computed-parameter functions that downstream packages can wrap without importing HSSM or PyTensor in ssm-simulators.

assemble_participant_fn

assemble_participant_fn(
    input_fields=None,
    *,
    response_field=DEFAULT_RESPONSE_FIELD,
    output=ARRAY
)

Assemble a participant-wise computed-parameter function.

By default, input_fields are derived from the model config. Pass explicit values only for non-standard layouts.

The returned function accepts a (n_trials, n_fields) array whose columns match input_fields. It computes SSM parameters before each learning update, maps response labels to zero-based action indices, and updates learning state from the response and optional outcome.

from_config `classmethod`

from_config(config, backend=AUTO)

Build an assembled model from a structural model config.

get_participant_input_fields

get_participant_input_fields(
    *, response_field=DEFAULT_RESPONSE_FIELD
)

Return the default participant input columns derived from the config.

participant_input_fields

participant_input_fields(
    *, response_field=DEFAULT_RESPONSE_FIELD
)

Backward-compatible alias for :meth:get_participant_input_fields.

ModelConfig `dataclass`

ModelConfig(
    model_name,
    description,
    decision_process,
    learning_process,
    task_environment,
    list_params=None,
    bounds=None,
    params_default=None,
    choices=None,
    response=(lambda: ["rt", "response"])(),
    response_to_choice="auto",
    learning_backend="auto",
    gradient="auto",
    include_choice=False,
    context_fields=None,
    computed_param_mapping=None,
    ssm_kwargs=(
        lambda: {"delta_t": 0.001, "max_t": 20.0}
    )(),
)

RLSSM model configuration for ssm-simulators.

Describes the structural specification of an RLSSM model: which learning process, which decision process (SSM), and which task environment. Concrete parameter values are NOT stored here — they are passed as theta to Simulator.simulate().

Parameters:

model_name (str) –

Unique identifier for this RLSSM model (e.g., "rlssm_angle_rw").
description (str) –

Human-readable model description.
decision_process (str) –

SSM model name in ssm-simulators registry (e.g., "angle", "ddm"). Must be resolvable via ModelConfigBuilder.from_model().
learning_process (LearningProcess) –

Instance of a class satisfying the LearningProcess protocol.
task_environment (TaskEnvironment | TaskConfig) –

Task environment instance or a TaskConfig to auto-build one. If TaskConfig, build_environment() is called in __post_init__.
list_params (list[str] | None, default: None ) –

All free parameter names (RL + fixed SSM), in order. If None, auto-derived: learning_process.free_params + fixed SSM params.
bounds (dict[str, tuple[float, float]] | None, default: None ) –

Parameter bounds. If None, auto-derived from learning_process.param_bounds + SSM model config param_bounds.
params_default (list[float] | None, default: None ) –

Default values in same order as list_params. If None, auto-derived.
choices (tuple[int, ...] | None, default: None ) –

SSM response labels (e.g., (-1, 1)). If None, taken from task_environment.
response (list[str], default: (lambda: ['rt', 'response'])() ) –

Response column names. Default ["rt", "response"].
response_to_choice (Literal['auto'] | dict[int, int], default: 'auto' ) –

Mapping from SSM response labels to zero-based learning choices. "auto" maps labels by task_environment.response_labels order.
learning_backend (Literal['auto', 'python', 'jax'], default: 'auto' ) –

Learning-process backend used for simulation and exported HSSM metadata. "auto" selects JAX when the process implements it and JAX is installed; otherwise it selects Python.
gradient (Literal['auto', 'available', 'unavailable'], default: 'auto' ) –

Gradient-support policy for HSSM integration metadata.
include_choice (bool, default: False ) –

Whether simulator output includes the derived zero-based choice column. Default False.
context_fields (list[str] | None, default: None ) –

Data/context columns beyond response required by the environment or learning process. Default derives a union from component declarations.
computed_param_mapping (dict[str, str] | None, default: None ) –

Optional override for non-name-matching handshakes. Maps learning process output name -> SSM param name. E.g., {"drift": "v"} if learning process outputs "drift" but SSM expects "v". Default: None (same-name linking).
ssm_kwargs (dict, default: (lambda: {'delta_t': 0.001, 'max_t': 20.0})() ) –

Default kwargs for the underlying SSM simulator call. Default: {"delta_t": 0.001, "max_t": 20.0}.

__post_init__

__post_init__()

Auto-build task environment and derive missing fields.

assemble

assemble(backend='auto')

Return a validated executable assembled model.

participant_contract

participant_contract(
    *, response_field=DEFAULT_RESPONSE_FIELD
)

Return the derived participant input layout for this config.

required_params `property`

required_params

Parameters that simulation requires from theta.

resolved_response_to_choice `property`

resolved_response_to_choice

Concrete response-label -> choice-index map.

__post_init__ normalizes response_to_choice (including the "auto" default) into a plain dict; this accessor exposes that post-init invariant with a narrowed type.

to_hssm_config_dict

to_hssm_config_dict()

Produce a dict compatible with HSSM's RLSSMConfig.from_rlssm_dict().

The output contains all fields from _HSSM_SHARED_FIELDS plus placeholder values for inference-only fields that the user must fill in on the HSSM side.

Returns:

dict[str, Any] –

Dict ready for RLSSMConfig.from_rlssm_dict(result) after user fills in inference-only fields.

validate

validate()

Validate config consistency. Called by Simulator.init().

Checks: 1. decision_process exists in ssm-simulators registry 2. Handshake: computed + fixed params cover all SSM params exactly once 3. No param is both computed and fixed 4. list_params length matches params_default length 5. All list_params have bounds

validate_data

validate_data(data)

Validate trial-level data against this model's RLSSM contract.

Returns a report with readable print() output and raise_for_errors() for fail-fast usage.

Simulator

Simulator(config)

RLSSM simulator composing a learning process with an SSM decision process.

Runs the interleaved trial-by-trial loop: compute SSM params -> simulate SSM -> observe choice -> generate reward -> update learning.

Reuses the existing ssm-simulators simulator() function with n_samples=1 for each trial. No Cython modifications needed — all 40+ SSM models work as decision processes out of the box.

Parameters:

config (ModelConfig) –

Structural model configuration. Validated on construction.

simulate

simulate(
    theta,
    n_trials=200,
    n_participants=None,
    random_state=None,
    mode="generative",
    observed_data=None,
)

Run full RLSSM simulation.

Parameters:

theta (dict[str, Any]) –

Concrete parameter values. Must contain all params required by the learning process and fixed SSM parameters. Each value can be a scalar shared by all participants or a one-dimensional list/array with one value per participant.
n_trials (int, default: 200 ) –

Number of trials per participant. Default 200.
n_participants (int | None, default: None ) –

Number of participants to simulate. If None, inferred from participant-wise theta values when present; otherwise defaults to 20.
random_state (int | None, default: None ) –

Seed for reproducibility. If None, non-deterministic.
mode (('generative', 'ppc'), default: "generative" ) –

Simulation mode. "generative" runs the unconstrained simulator loop. "ppc" runs observed-history-conditioned posterior predictive simulation.
observed_data (DataFrame | None, default: None ) –

Observed participant history required for mode="ppc".

Returns:

DataFrame –

Balanced panel with participant_id, trial_id, configured response columns, configured context fields, and optional derived choice.

assembled

Assembled RLSSM model interfaces.

AssembledFunctionOutput

Bases: StrEnum

Return shape for participant-wise computed-parameter functions.

AssembledModel `dataclass`

AssembledModel(
    config,
    learning_backend,
    gradient,
    model_name,
    decision_process,
    list_params,
    bounds,
    params_default,
    response,
    choices,
    context_fields,
    computed_params,
    response_to_choice,
)

Validated executable form of an RLSSM ModelConfig.

The assembled model exposes package-neutral metadata and pure Python/JAX computed-parameter functions that downstream packages can wrap without importing HSSM or PyTensor in ssm-simulators.

assemble_participant_fn

assemble_participant_fn(
    input_fields=None,
    *,
    response_field=DEFAULT_RESPONSE_FIELD,
    output=ARRAY
)

Assemble a participant-wise computed-parameter function.

By default, input_fields are derived from the model config. Pass explicit values only for non-standard layouts.

The returned function accepts a (n_trials, n_fields) array whose columns match input_fields. It computes SSM parameters before each learning update, maps response labels to zero-based action indices, and updates learning state from the response and optional outcome.

from_config `classmethod`

from_config(config, backend=AUTO)

Build an assembled model from a structural model config.

get_participant_input_fields

get_participant_input_fields(
    *, response_field=DEFAULT_RESPONSE_FIELD
)

Return the default participant input columns derived from the config.

participant_input_fields

participant_input_fields(
    *, response_field=DEFAULT_RESPONSE_FIELD
)

Backward-compatible alias for :meth:get_participant_input_fields.

LearningBackendRequest

Bases: StrEnum

Requested learning backend policy for assembly.

ResolvedLearningBackend

Bases: StrEnum

Concrete learning backend selected after resolving policy.

resolve_model

resolve_model(model)

Resolve a preset name or validate an existing RLSSM model config.

config

ModelConfig — structural model specification for RLSSM simulation.

ModelConfig `dataclass`

ModelConfig(
    model_name,
    description,
    decision_process,
    learning_process,
    task_environment,
    list_params=None,
    bounds=None,
    params_default=None,
    choices=None,
    response=(lambda: ["rt", "response"])(),
    response_to_choice="auto",
    learning_backend="auto",
    gradient="auto",
    include_choice=False,
    context_fields=None,
    computed_param_mapping=None,
    ssm_kwargs=(
        lambda: {"delta_t": 0.001, "max_t": 20.0}
    )(),
)

RLSSM model configuration for ssm-simulators.

Describes the structural specification of an RLSSM model: which learning process, which decision process (SSM), and which task environment. Concrete parameter values are NOT stored here — they are passed as theta to Simulator.simulate().

Parameters:

model_name (str) –

Unique identifier for this RLSSM model (e.g., "rlssm_angle_rw").
description (str) –

Human-readable model description.
decision_process (str) –

SSM model name in ssm-simulators registry (e.g., "angle", "ddm"). Must be resolvable via ModelConfigBuilder.from_model().
learning_process (LearningProcess) –

Instance of a class satisfying the LearningProcess protocol.
task_environment (TaskEnvironment | TaskConfig) –

Task environment instance or a TaskConfig to auto-build one. If TaskConfig, build_environment() is called in __post_init__.
list_params (list[str] | None, default: None ) –

All free parameter names (RL + fixed SSM), in order. If None, auto-derived: learning_process.free_params + fixed SSM params.
bounds (dict[str, tuple[float, float]] | None, default: None ) –

Parameter bounds. If None, auto-derived from learning_process.param_bounds + SSM model config param_bounds.
params_default (list[float] | None, default: None ) –

Default values in same order as list_params. If None, auto-derived.
choices (tuple[int, ...] | None, default: None ) –

SSM response labels (e.g., (-1, 1)). If None, taken from task_environment.
response (list[str], default: (lambda: ['rt', 'response'])() ) –

Response column names. Default ["rt", "response"].
response_to_choice (Literal['auto'] | dict[int, int], default: 'auto' ) –

Mapping from SSM response labels to zero-based learning choices. "auto" maps labels by task_environment.response_labels order.
learning_backend (Literal['auto', 'python', 'jax'], default: 'auto' ) –

Learning-process backend used for simulation and exported HSSM metadata. "auto" selects JAX when the process implements it and JAX is installed; otherwise it selects Python.
gradient (Literal['auto', 'available', 'unavailable'], default: 'auto' ) –

Gradient-support policy for HSSM integration metadata.
include_choice (bool, default: False ) –

Whether simulator output includes the derived zero-based choice column. Default False.
context_fields (list[str] | None, default: None ) –

Data/context columns beyond response required by the environment or learning process. Default derives a union from component declarations.
computed_param_mapping (dict[str, str] | None, default: None ) –

Optional override for non-name-matching handshakes. Maps learning process output name -> SSM param name. E.g., {"drift": "v"} if learning process outputs "drift" but SSM expects "v". Default: None (same-name linking).
ssm_kwargs (dict, default: (lambda: {'delta_t': 0.001, 'max_t': 20.0})() ) –

Default kwargs for the underlying SSM simulator call. Default: {"delta_t": 0.001, "max_t": 20.0}.

__post_init__

__post_init__()

Auto-build task environment and derive missing fields.

assemble

assemble(backend='auto')

Return a validated executable assembled model.

participant_contract

participant_contract(
    *, response_field=DEFAULT_RESPONSE_FIELD
)

Return the derived participant input layout for this config.

required_params `property`

required_params

Parameters that simulation requires from theta.

resolved_response_to_choice `property`

resolved_response_to_choice

Concrete response-label -> choice-index map.

__post_init__ normalizes response_to_choice (including the "auto" default) into a plain dict; this accessor exposes that post-init invariant with a narrowed type.

to_hssm_config_dict

to_hssm_config_dict()

Produce a dict compatible with HSSM's RLSSMConfig.from_rlssm_dict().

The output contains all fields from _HSSM_SHARED_FIELDS plus placeholder values for inference-only fields that the user must fill in on the HSSM side.

Returns:

dict[str, Any] –

Dict ready for RLSSMConfig.from_rlssm_dict(result) after user fills in inference-only fields.

validate

validate()

Validate config consistency. Called by Simulator.init().

Checks: 1. decision_process exists in ssm-simulators registry 2. Handshake: computed + fixed params cover all SSM params exactly once 3. No param is both computed and fixed 4. list_params length matches params_default length 5. All list_params have bounds

validate_data

validate_data(data)

Validate trial-level data against this model's RLSSM contract.

Returns a report with readable print() output and raise_for_errors() for fail-fast usage.

derive_participant_contract

derive_participant_contract(
    config, *, response_field=DEFAULT_RESPONSE_FIELD
)

Derive the participant input layout from a structural RLSSM config.

env

TaskEnvironment protocol and built-in implementations.

Bandit

Bandit(rewards, response_labels=None)

Generic bandit task environment.

Public constructors are Bandit.bernoulli(...) and Bandit.gaussian(...). Rewards are sampled by zero-based choice index; response_labels define the SSM labels mapped onto those choices.

bernoulli `classmethod`

bernoulli(probabilities=None, response_labels=None)

Build a Bernoulli-reward bandit.

gaussian `classmethod`

gaussian(means=None, sds=None, response_labels=None)

Build a Gaussian-reward bandit.

get_extra_data

get_extra_data(trial_idx)

Compatibility wrapper around get_trial_context.

n_arms `property`

n_arms

Alias for :attr:n_choices (bandit terminology).

sample_reward

sample_reward(action, trial_idx)

Compatibility wrapper around sample_context.

DiscreteChoiceEnvironment

Bases: TaskEnvironment, Protocol

Task environment with discrete SSM response labels and learning choices.

n_choices `property`

n_choices

Number of available zero-based learning choices.

response_labels `property`

response_labels

SSM response labels corresponding to choices in order.

TaskConfig

TaskConfig(task='bandit', reward=None, **options)

Convenience configuration for registered task environments.

TaskConfig is a shorthand that delegates task-specific options to a registry builder. Built in support currently includes task="bandit" with reward="bernoulli" or reward="gaussian".

TaskEnvironment

Bases: Protocol

Protocol for RLSSM task environments.

A task environment provides per-trial context and optional post-decision signals. It is stateful and must be reset before each participant.

Discrete response/choice mapping requires :class:DiscreteChoiceEnvironment.

context_fields `property`

context_fields

Names of per-trial context columns this environment provides.

get_trial_context

get_trial_context(trial_idx)

Return pre-decision per-trial context columns.

reset

reset(rng=None)

Reset environment state for a new participant.

sample_context

sample_context(context, trial_idx)

Return post-decision context columns for learning/output.

register_task

register_task(task, builder, *, overwrite=False)

Register a task environment builder for TaskConfig.

registered_tasks

registered_tasks()

List task names available through TaskConfig.

learning

LearningProcess protocol and built-in implementations.

LearningProcess

Bases: Protocol

Protocol for RLSSM learning processes.

A learning process maintains internal state (e.g., Q-values) and computes SSM parameters (e.g., drift rate) from that state on each trial. After each trial's decision and reward, the state is updated.

The computed_params property is the formal handshake between the learning process and the decision process: it declares which SSM parameters the learning process produces. The simulator validates that these, together with fixed SSM params provided by the user, cover all parameters required by the decision process model.

available_backends `property`

available_backends

Learning backends implemented by this process.

compute_python

compute_python(state, params, context)

Compute SSM parameters from explicit Python/NumPy state.

compute_ssm_params

compute_ssm_params(trial_params)

Compute SSM parameters from current learning state.

computed_params `property`

computed_params

SSM parameter names this process computes (e.g., ['v']).

default_params `property`

default_params

Default values for each free param.

free_params `property`

free_params

RL parameter names this process requires from theta.

init_state

init_state()

Return an explicit initial learning state for one participant.

param_bounds `property`

param_bounds

Bounds for each free param.

required_context_fields `property`

required_context_fields

Context keys this process needs for compute/update.

reset

reset(**kwargs)

Reset internal state for a new participant.

supports_gradient `property`

supports_gradient

Whether the differentiable backend supports gradient-based inference.

update

update(action, reward, trial_params)

Update learning state given the choice outcome.

update_python

update_python(state, params, context)

Return the next explicit Python/NumPy state.

RescorlaWagnerDeltaRule

RescorlaWagnerDeltaRule(
    n_actions=2, initial_q=0.5, feedback_field="feedback"
)

Rescorla-Wagner delta learning core.

Updates Q-values via Q[action] += alpha * (reward - Q[action]). This class owns Q-value state and replay/update behavior but emits no SSM parameters by itself. Use RescorlaWagnerDrift for two-action drift models and RescorlaWagnerSoftmax for inverse-temperature softmax models.

compute_ssm_params

compute_ssm_params(trial_params)

Compute pre-update SSM parameters from current learning state.

q_values `property`

q_values

Current Q-values. None if reset() has not been called.

update

update(action, reward, trial_params)

Update Q[action] from the observed outcome.

RescorlaWagnerDrift

RescorlaWagnerDrift(
    n_actions=2, initial_q=0.5, feedback_field="feedback"
)

Bases: RescorlaWagnerDeltaRule

Rescorla-Wagner learner emitting two-action drift v.

Computes drift rate as scaled Q-value difference: v = (Q[1] - Q[0]) * scaler.

RescorlaWagnerDualAlphaDrift

RescorlaWagnerDualAlphaDrift(
    n_actions=2, initial_q=0.5, feedback_field="feedback"
)

Bases: RescorlaWagnerDualAlphaRule

Dual-alpha Rescorla-Wagner learner emitting two-action drift v.

RescorlaWagnerDualAlphaRule

RescorlaWagnerDualAlphaRule(
    n_actions=2, initial_q=0.5, feedback_field="feedback"
)

Bases: RescorlaWagnerDeltaRule

Rescorla-Wagner learning core with separate learning rates.

Positive prediction errors use rl_alpha and negative prediction errors use rl_alpha_neg.

update_python

update_python(state, params, context)

Update Q[action] with sign-dependent learning rates.

RescorlaWagnerDualAlphaSoftmax

RescorlaWagnerDualAlphaSoftmax(
    n_actions=2, initial_q=0.5, feedback_field="feedback"
)

Bases: RescorlaWagnerDualAlphaRule

Dual-alpha Rescorla-Wagner learner emitting Q-values q0..qN.

RescorlaWagnerRaceDrifts

RescorlaWagnerRaceDrifts(
    n_actions=2, initial_q=0.5, feedback_field="feedback"
)

Bases: RescorlaWagnerDeltaRule

Rescorla-Wagner learner emitting scaled race drifts v0..vN.

The scaling contract is explicit: on each trial, before the RW update, v_i = scaler * q_i for every action i.

RescorlaWagnerSoftmax

RescorlaWagnerSoftmax(
    n_actions=2, initial_q=0.5, feedback_field="feedback"
)

Bases: RescorlaWagnerDeltaRule

Rescorla-Wagner learner emitting pre-update Q-values q0..qN.

preset

Registry for RLSSM presets.

PresetInfo

Bases: dict

Dictionary-like preset metadata with a readable string representation.

get

get(name)

Get a named RLSSM preset config. Returns a fresh instance.

info

info(name)

Return readable metadata for a named RLSSM preset.

list

list()

List available RLSSM preset names.

register

register(name, factory, *, metadata=None)

Register a named RLSSM preset.

resolve_model

resolve_model(model)

Resolve a preset name or validate an existing RLSSM model config.

simulator

Simulator — interleaved learning + SSM decision simulation.

Simulator

Simulator(config)

RLSSM simulator composing a learning process with an SSM decision process.

Runs the interleaved trial-by-trial loop: compute SSM params -> simulate SSM -> observe choice -> generate reward -> update learning.

Reuses the existing ssm-simulators simulator() function with n_samples=1 for each trial. No Cython modifications needed — all 40+ SSM models work as decision processes out of the box.

Parameters:

config (ModelConfig) –

Structural model configuration. Validated on construction.

simulate

simulate(
    theta,
    n_trials=200,
    n_participants=None,
    random_state=None,
    mode="generative",
    observed_data=None,
)

Run full RLSSM simulation.

Parameters:

theta (dict[str, Any]) –

Concrete parameter values. Must contain all params required by the learning process and fixed SSM parameters. Each value can be a scalar shared by all participants or a one-dimensional list/array with one value per participant.
n_trials (int, default: 200 ) –

Number of trials per participant. Default 200.
n_participants (int | None, default: None ) –

Number of participants to simulate. If None, inferred from participant-wise theta values when present; otherwise defaults to 20.
random_state (int | None, default: None ) –

Seed for reproducibility. If None, non-deterministic.
mode (('generative', 'ppc'), default: "generative" ) –

Simulation mode. "generative" runs the unconstrained simulator loop. "ppc" runs observed-history-conditioned posterior predictive simulation.
observed_data (DataFrame | None, default: None ) –

Observed participant history required for mode="ppc".

Returns:

DataFrame –

Balanced panel with participant_id, trial_id, configured response columns, configured context fields, and optional derived choice.

validation

Data contract validation for RLSSM panels.

DataValidationIssue `dataclass`

DataValidationIssue(level, code, message, hint=None)

A single validation finding.

DataValidationReport `dataclass`

DataValidationReport(
    issues=list(), n_participants=None, n_trials=None
)

Aggregated validation results for an RLSSM data panel.

ok `property`

ok

True when there are no error-level issues.

print

print()

Print a human-readable summary to stdout.

raise_for_errors

raise_for_errors()

Raise ValueError if any error-level issues were recorded.

validate_rlssm_data

validate_rlssm_data(config, data)

Validate a data panel against the RLSSM model contract.

Parameters:

config (ModelConfig) –

Structural RLSSM configuration. Should already pass config.validate().
data (DataFrame) –

Empirical or simulated trial-level panel.

Returns:

DataValidationReport –

Validation findings. Call raise_for_errors() to fail fast on errors.

ssms.support_utils

kde_class

LogKDE

LogKDE(
    simulator_data,
    bandwidth_type="silverman",
    auto_bandwidth=True,
    displace_t=False,
)

Class for generating kdes from (rt, choice) data. Works for any number of choices.

Attributes:

simulator_data ((dict, default < None)) –
```
Dictionary of the type {'rts':[], 'choices':[], 'metadata':{}}.
Follows the format of simulator returns in this package.
```
bandwidth_type: string type of bandwidth to use, default is 'silverman' auto_bandwidth: boolean whether to compute bandwidths automatically, default is True

Methods:

compute_bandwidths –

Computes bandwidths for each choice from rt data. _generate_base_kdes(auto_bandwidth=True, bandwidth_type='silverman') Generates kdes from rt data. kde_eval(data=([], []), log_eval=True) Evaluates kde log likelihood at chosen points. kde_sample(n_samples=2000, use_empirical_choice_p=True, alternate_choice_p=0) Samples from a given kde. _attach_data_from_simulator(simulator_data={'rts':[0, 2, 4], 'choices':[-1, 1, -1], 'metadata':{}})) Helper function to transform ddm simulator output to dataset suitable for the kde function class.
Returns: –

type: description

Arguments:

simulator_data: Dictionary containing simulation data with keys 'rts', 'choices', and 'metadata'.
    Follows the format returned by simulator functions in this package.
bandwidth_type: Type of bandwidth to use for KDE. Currently only 'silverman' is supported.
    Defaults to 'silverman'.
auto_bandwidth: Whether to automatically compute bandwidths based on the data.
    If False, bandwidths must be set manually. Defaults to True.
displace_t: Whether to shift RTs by the t parameter from metadata.
    Only works if all trials have the same t value. Defaults to False.

Raises:

AssertionError: If displace_t is True but metadata contains multiple t values.

__kde_eval_

__kde_eval_(data, log_eval=True, lb=-66.774, eps=0.0001)

Evaluates kde log likelihood at chosen points.

Arguments:

data: dict Dictionary with keys 'rts', and 'choices' to evaluate the kde at. log_eval: bool Whether to return log likelihood or likelihood, default is True. lb: float Lower bound for log likelihoods, default is -66.774. eps: float Epsilon value to use for lower bounds on rts, default is 10e-5.

Returns:

np.ndarray Array of log likelihoods for each (rt, choice) pair if log_eval is True, otherwise array of likelihoods.

compute_bandwidths

compute_bandwidths(bandwidth_type='silverman')

Computes bandwidths for each choice from rt data.

Arguments:

bandwidth_type: str Type of bandwidth to use, default is 'silverman' which follows silverman rule. return_result: bool Whether to return the result. Defaults to False.

Returns:

bandwidths: list[float] | None List of bandwidths for each choice if return_result is True, otherwise None.

kde_eval

kde_eval(
    data,
    log_eval=True,
    lb=-66.774,
    eps=0.0001,
    filter_rts=OMISSION_SENTINEL,
)

Evaluates kde log likelihood at chosen points.

Arguments:

data: dict Dictionary with keys 'rts', and/or 'log_rts' and 'choices' to evaluate the kde at. If 'rts' is provided, 'log_rts' is ignored. log_eval: boolean Whether to return log likelihood or likelihood, default is True. lb: float Lower bound for log likelihoods, default is -66.774. (This is the log of 1e-29) eps: float Epsilon value to use for lower bounds on rts. filter_rts: float Value to filter rts by, default is OMISSION_SENTINEL (-999). This is the sentinel value returned by simulators when a trial exceeds max_t or deadline (i.e., an omission).

Returns:

log_kde_eval: array Array of log likelihoods for each (rt, choice) pair.

kde_sample

kde_sample(
    n_samples=2000,
    use_empirical_choice_p=True,
    alternate_choice_p=0.0,
    random_state=None,
)

Samples from a given kde.

Arguments:

n_samples: int Number of samples to draw. use_empirical_choice_p: bool Whether to use empirical choice proportions, default is True. (Note 'empirical' here, refers to the originally attached datasets that served as the basis to generate the choice-wise kdes) alternate_choice_p: np.ndarray | float Array of choice proportions to use, default is 0. (Note 'alternate' here refers to 'alternative' to the 'empirical' choice proportions) random_state: int | None Random seed for reproducibility. If None, uses non-reproducible random behavior.

Returns:

dict[str, np.ndarray | dict] Dictionary containing: - 'rts': np.ndarray - Response times - 'log_rts': np.ndarray - Log of response times - 'choices': np.ndarray - Choices made - 'metadata': dict - Simulator information

bandwidth_silverman

bandwidth_silverman(
    sample=(0, 0, 0),
    std_cutoff=0.001,
    std_proc="restrict",
    std_n_1=10.0,
)

Computes silverman bandwidth for an array of samples (rts in our context, but general).

Arguments:

sample: np.ndarray Array of samples to compute bandwidth for. std_cutoff: float Cutoff for std, default is 1e-3. (If sample-std is smaller than this, we either kill it or restrict it to this value) std_proc: str How to deal with small stds, default is 'restrict'. (Options: 'kill', 'restrict') std_n_1: float Value to use if n = 1, default is 10.0. (Not clear if the default is sensible here)

Returns:

float Silverman bandwidth for the given sample. This is applied as the bandwidth parameter when generating gaussian-based kdes in the LogKDE class.

utils

This module provides utility functions for handling parameter dependencies and sampling parameters within specified constraints.

Functions:

parse_bounds –

Parse the bounds of a parameter and extract any dependencies.
build_dependency_graph –

Build a dependency graph based on parameter bounds.
topological_sort_util –

visited: Set[str], stack: List[str], graph: Dict[str, Set[str]], temp_marks: Set[str]) -> None Helper function for performing a depth-first search in the topological sort.
topological_sort –

Perform a topological sort on the dependency graph to determine the sampling order.
sample_parameters_from_constraints –

sample_size: int) -> Dict[str, np.ndarray] Sample parameters uniformly within specified bounds, respecting any dependencies.

build_dependency_graph

build_dependency_graph(param_dict)

Build a dependency graph based on parameter bounds.

The graph is structured as parent -> children, where children depend on parent. For example, if 'st' has upper bound 't', then graph['t'] contains 'st'.

Args: param_dict: Dictionary mapping parameter names to (lower, upper) bounds.

Returns: Dictionary where keys are parameters and values are sets of parameters that depend on them.

Raises: ValueError: If a dependency references an undefined parameter.

Example: >>> param_dict = {"a": (0, 1), "b": ("a", 2)} >>> build_dependency_graph(param_dict) {'a': {'b'}, 'b': set()}

parse_bounds

parse_bounds(bounds)

Parse the bounds of a parameter and extract any dependencies.

Args: bounds: A tuple of (lower, upper) bounds. Can be numeric or strings (for dependencies).

Returns: Set of parameter names that this parameter depends on.

Example: >>> parse_bounds((0.0, 1.0)) set() >>> parse_bounds(("a", 1.0)) {'a'} >>> parse_bounds((0.0, "b")) {'b'} >>> parse_bounds(("c", "d"))

sample_parameters_from_constraints

sample_parameters_from_constraints(
    param_dict, sample_size, random_state=None
)

Sample parameters uniformly within specified bounds, respecting any dependencies.

This is a backward-compatible wrapper around UniformParameterSampler.

Args: param_dict: Dictionary mapping parameter names to (lower, upper) bounds. Bounds can be numeric or strings (for dependencies). sample_size: Number of parameter sets to sample. random_state: Optional random seed for reproducibility.

Returns: Dictionary mapping parameter names to arrays of sampled values.

Raises: ValueError: If bounds are invalid (lower >= upper) or circular dependencies exist.

Example: >>> param_dict = {"v": (-1.0, 1.0), "a": (0.5, 2.0)} >>> samples = sample_parameters_from_constraints(param_dict, sample_size=10) >>> samples['v'].shape (10,)

topological_sort

topological_sort(graph)

Perform a topological sort on the dependency graph to determine the sampling order.

Args: graph: Dependency graph where keys are parameters and values are sets of parameters that depend on them.

Returns: List of parameter names in sampling order (dependencies first).

Raises: ValueError: If circular dependencies are detected.

Example: >>> graph = {"a": {"b"}, "b": set()} >>> topological_sort(graph) ['a', 'b']

topological_sort_util

topological_sort_util(
    node, visited, stack, graph, temp_marks
)

Helper function for performing a depth-first search in the topological sort.

Args: node: Current node to visit visited: Set of permanently visited nodes stack: List to accumulate the topological order (prepended in post-order) graph: Dependency graph temp_marks: Set of temporarily marked nodes (for cycle detection)

Raises: ValueError: If a circular dependency is detected.

ssms.transforms

Unified parameter transforms for the training data generation and simulation workflows.

This module provides a unified system for parameter transformations that can be applied at two different stages:

Sampling transforms (for training data generation): Applied during the parameter sampling stage of the training data generation workflow. These enforce parameter relationships (e.g., a > z) when generating synthetic training data for likelihood approximation networks.
Examples: SwapIfLessConstraint, NormalizeToSumConstraint
Simulation transforms (for basic Simulator): Applied via ParameterSimulatorAdapters when running the basic Simulator. These prepare user-provided parameters for the low-level C/Cython simulators (e.g., stacking v0, v1, v2 into a single v array).
Examples: ColumnStackParameters, ExpandDimension, SetZeroArray

All transforms inherit from ParameterTransform and use the same interface, making it easy to compose them into pipelines.

Example usage in model config:

from ssms.transforms import (
    SwapIfLessConstraint,
    ColumnStackParameters,
    ExpandDimension,
    SetZeroArray,
)

def get_lba_angle_3_config():
    return {
        "name": "lba_angle_3",
        "params": ["v0", "v1", "v2", "a", "z", "theta"],
        # ... other fields ...

        "parameter_transforms": {
            "sampling": [
                SwapIfLessConstraint("a", "z"),
            ],
            "simulation": [
                ColumnStackParameters(["v0", "v1", "v2"], "v"),
                ExpandDimension(["a", "z", "theta"]),
                SetZeroArray("t"),
            ],
        }
    }

ApplyMapping

ApplyMapping(
    source_param,
    target_param,
    mapping_key,
    additional_sources=None,
    delete_source=False,
)

Bases: ParameterTransform

Apply a mapping function to transform a parameter.

Uses a mapping from model_config to transform a parameter value. Commonly used for random variable distributions.

Parameters:

source_param (str) –

Parameter to read from
target_param (str) –

Parameter to write to
mapping_key (str) –

Key in model_config["simulator_param_mappings"]
additional_sources (list[str], default: None ) –

Additional parameters to pass to the mapping function
delete_source (bool, default: False ) –

Whether to delete the source parameter after mapping

apply

apply(theta, model_config=None, n_trials=None)

Apply mapping function to transform parameter.

ColumnStackParameters

ColumnStackParameters(
    source_params, target_param, delete_sources=True
)

Bases: ParameterTransform

Stack multiple parameters into a single multi-column array.

Takes individual parameters (e.g., v0, v1, v2) and stacks them column-wise into a single array (e.g., v with shape (n_trials, 3)).

Parameters:

source_params (list[str]) –

Names of parameters to stack
target_param (str) –

Name of the resulting stacked parameter
delete_sources (bool, default: True ) –

Whether to delete the source parameters after stacking

Examples:

>>> transform = ColumnStackParameters(["v0", "v1", "v2"], "v")
>>> theta = {"v0": [0.5], "v1": [0.6], "v2": [0.7]}
>>> theta = transform.apply(theta)
>>> theta["v"]  # array([[0.5, 0.6, 0.7]])

apply

apply(theta, model_config=None, n_trials=None)

Stack parameters column-wise.

ConditionalAdaptation

ConditionalAdaptation(condition, transformation)

Bases: ParameterTransform

Apply a transform only if a condition is met.

Parameters:

condition (Callable) –

Function that takes (theta, model_config, n_trials) and returns bool
transformation (ParameterTransform) –

Transform to apply if condition is True

Examples:

>>> transform = ConditionalAdaptation(
...     condition=lambda theta, cfg, n: "a" in theta,
...     transformation=ExpandDimension(["a"])
... )

apply

apply(theta, model_config=None, n_trials=None)

Apply transform if condition is met.

DeleteParameters

DeleteParameters(param_names)

Bases: ParameterTransform

Delete specified parameters from theta.

Parameters:

param_names (list[str]) –

Names of parameters to delete

Examples:

>>> transform = DeleteParameters(["A", "b"])
>>> theta = {"v": [0.5], "A": [0.5], "b": [1.0]}
>>> theta = transform.apply(theta)
>>> "A" in theta  # False

apply

apply(theta, model_config=None, n_trials=None)

Delete specified parameters.

ExpandDimension

ExpandDimension(param_names)

Bases: ParameterTransform

Expand dimensions of specified parameters.

Converts 1D arrays of shape (n_trials,) to shape (n_trials, 1) by adding a dimension. This is commonly needed for multi-particle models.

Parameters:

param_names (list[str]) –

Names of parameters to expand

Examples:

>>> transform = ExpandDimension(["a", "t"])
>>> theta = {"a": np.array([1.0, 2.0]), "t": np.array([0.3, 0.3])}
>>> theta = transform.apply(theta)
>>> theta["a"].shape  # (2, 1)

apply

apply(theta, model_config=None, n_trials=None)

Expand dimensions of specified parameters.

LambdaAdaptation

LambdaAdaptation(func, name=None)

Bases: ParameterTransform

Wrap a lambda or callable function as a transform.

Parameters:

func (Callable) –

Function with signature: func(theta, model_config, n_trials) → theta
name (str, default: None ) –

Optional name for debugging

Examples:

>>> transform = LambdaAdaptation(
...     lambda theta, cfg, n: theta.update({"nact": 3}) or theta,
...     name="set_nact_to_3"
... )

apply

apply(theta, model_config=None, n_trials=None)

Apply the wrapped function.

NormalizeToSumConstraint

NormalizeToSumConstraint(param_names)

Bases: ParameterTransform

Normalize a set of parameters to sum to 1.

This constraint is commonly used for drift rates in multi-accumulator models (e.g., RLWM models) where the drift rates must sum to 1.

The normalization includes a small epsilon (1e-20) to avoid division by zero and ensure numerical stability.

Parameters:

param_names (list[str]) –

List of parameter names to normalize together

Examples:

>>> constraint = NormalizeToSumConstraint(["v1", "v2", "v3"])
>>> theta = {"v1": 0.2, "v2": 0.3, "v3": 0.5}
>>> result = constraint.apply(theta)
>>> sum([result["v1"], result["v2"], result["v3"]])
1.0

Use in model config:

>>> model_config = {
...     "name": "dev_rlwm_lba_race_v1",
...     "parameter_transforms": {
...         "sampling": [
...             NormalizeToSumConstraint(["vRL0", "vRL1", "vRL2"]),
...             NormalizeToSumConstraint(["vWM0", "vWM1", "vWM2"]),
...         ],
...     }
... }

Parameters:

param_names (list[str]) –

List of parameter names to normalize together

apply

apply(
    theta, model_config=None, n_trials=None, epsilon=1e-20
)

Normalize specified parameters to sum to 1.

Parameters:

theta (dict[str, Any]) –

Dictionary of parameter values (can be scalars or arrays)
model_config (dict[str, Any] or None, default: None ) –

Not used by this transform (included for interface compatibility)
n_trials (int or None, default: None ) –

Not used by this transform (included for interface compatibility)
epsilon (float, default: 1e-20 ) –

Small value to add to the total to avoid division by zero. Default is 1e-20.

Returns:

dict[str, Any] –

Modified theta dictionary with normalized values

ParameterTransform

Bases: ABC

Abstract base class for all parameter transforms.

A parameter transform is a single, focused operation that modifies theta parameters. Transforms can be composed into pipelines and are used at two phases:

Sampling phase: Applied after parameter sampling to enforce constraints (e.g., ensure a > z by swapping values)
Simulation phase: Applied before calling the simulator to prepare parameters for the expected format (e.g., stack v0, v1 → v)

Subclasses must implement the apply method. The method signature includes optional model_config and n_trials arguments to support both simple constraints (which only need theta) and complex adapters (which may need additional context).

Examples:

Create a custom transform:

>>> class ScaleParameter(ParameterTransform):
...     def __init__(self, param_name: str, scale: float):
...         self.param_name = param_name
...         self.scale = scale
...
...     def apply(self, theta, model_config=None, n_trials=None):
...         if self.param_name in theta:
...             theta[self.param_name] = theta[self.param_name] * self.scale
...         return theta

Use in model config:

>>> model_config = {
...     "name": "my_model",
...     "parameter_transforms": {
...         "sampling": [SwapIfLessConstraint("a", "z")],
...         "simulation": [ColumnStackParameters(["v0", "v1"], "v")],
...     }
... }

repr

__repr__()

String representation for debugging.

Returns:

str –

A string describing this transform.

str

__str__()

Human-readable string representation.

apply `abstractmethod`

apply(theta, model_config=None, n_trials=None)

Apply transform to theta parameters.

This method should modify the theta dictionary in place and return it. It may add new parameters, modify existing parameters, or remove parameters as needed.

Parameters:

theta (dict[str, Any]) –

Dictionary of model parameters. Values are typically numpy arrays.
model_config (dict[str, Any] or None, default: None ) –

Model configuration dictionary. Available for transforms that need additional context (e.g., looking up simulator_fixed_params). Sampling-time transforms typically ignore this.
n_trials (int or None, default: None ) –

Number of trials. Available for transforms that need to create arrays of a specific size. Sampling-time transforms typically ignore this.

Returns:

dict[str, Any] –

The modified theta dictionary (usually the same object passed in).

Notes

Transforms should be pure when possible (no side effects)
If creating new arrays, use dtype=np.float32 for consistency
Document any parameters that are added, modified, or removed

RenameParameter

RenameParameter(old_name, new_name, transform_fn=None)

Bases: ParameterTransform

Rename a parameter and optionally transform its value.

Parameters:

old_name (str) –

Current parameter name
new_name (str) –

New parameter name
transform_fn (Callable or None, default: None ) –

Optional function to apply to the value during renaming

Examples:

>>> transform = RenameParameter("A", "z", lambda x: np.expand_dims(x, axis=1))
>>> theta = {"A": np.array([0.5])}
>>> theta = transform.apply(theta)
>>> "z" in theta  # True

apply

apply(theta, model_config=None, n_trials=None)

Rename parameter and optionally transform its value.

SetDefaultValue

SetDefaultValue(param_name, default_value, dtype=float32)

Bases: ParameterTransform

Set a parameter to a default value if not present.

This transform adds a parameter with a default value if it doesn't already exist in theta. The value is automatically tiled to match n_trials.

Parameters:

param_name (str) –

Name of the parameter to set
default_value (float or int or ndarray) –

Default value to use. If scalar, will be tiled to (n_trials,).
dtype (dtype, default: float32 ) –

Data type for the parameter. Defaults to np.float32.

Examples:

>>> transform = SetDefaultValue("nact", 3)
>>> theta = {}
>>> theta = transform.apply(theta, {}, n_trials=5)
>>> theta["nact"]  # array([3., 3., 3., 3., 3.], dtype=float32)

apply

apply(theta, model_config=None, n_trials=None)

Add parameter with default value if not present.

SetZeroArray

SetZeroArray(param_name, shape=None, dtype=float32)

Bases: ParameterTransform

Set a parameter to an array of zeros.

Parameters:

param_name (str) –

Name of the parameter to set
shape (tuple or None, default: None ) –

Shape of the array. Use None for n_trials dimension.
dtype (dtype, default: float32 ) –

Data type for the array. Defaults to np.float32.

Examples:

>>> transform = SetZeroArray("v")
>>> theta = transform.apply({}, {}, n_trials=5)
>>> theta["v"].shape  # (5,)

apply

apply(theta, model_config=None, n_trials=None)

Set parameter to zero array.

SwapIfLessConstraint

SwapIfLessConstraint(param_a, param_b)

Bases: ParameterTransform

Swap two parameters if first is less than or equal to second.

This constraint ensures ordering relationships like a > z by swapping the values when the constraint is violated. Works with both scalar and array inputs.

Parameters:

param_a (str) –

First parameter name (should be > param_b after constraint)
param_b (str) –

Second parameter name (should be < param_a after constraint)

Examples:

>>> constraint = SwapIfLessConstraint("a", "z")
>>> theta = {"a": np.array([0.5, 1.5]), "z": np.array([1.0, 0.5])}
>>> result = constraint.apply(theta)
>>> # First element: a=0.5 <= z=1.0 → swap → a=1.0, z=0.5
>>> # Second element: a=1.5 > z=0.5 → no swap → a=1.5, z=0.5

Use in model config:

>>> model_config = {
...     "name": "lba_angle_3",
...     "parameter_transforms": {
...         "sampling": [SwapIfLessConstraint("a", "z")],
...     }
... }

Parameters:

param_a (str) –

First parameter name (should be > param_b after constraint)
param_b (str) –

Second parameter name (should be < param_a after constraint)

apply

apply(theta, model_config=None, n_trials=None)

Swap param_a and param_b where param_a <= param_b.

Parameters:

theta (dict[str, Any]) –

Dictionary of parameter values (can be scalars or arrays)
model_config (dict[str, Any] or None, default: None ) –

Not used by this transform (included for interface compatibility)
n_trials (int or None, default: None ) –

Not used by this transform (included for interface compatibility)

Returns:

dict[str, Any] –

Modified theta dictionary with swapped values where needed

TileArray

TileArray(param_name, value, dtype=float32)

Bases: ParameterTransform

Tile a constant value across trials.

Parameters:

param_name (str) –

Name of the parameter to create
value (float or array - like) –

Value(s) to tile
dtype (dtype, default: float32 ) –

Data type for the array. Defaults to np.float32.

Examples:

>>> transform = TileArray("z", 0.5)
>>> theta = transform.apply({}, {}, n_trials=3)
>>> theta["z"]  # array([0.5, 0.5, 0.5], dtype=float32)

apply

apply(theta, model_config=None, n_trials=None)

Tile value across trials.

base

Base class for all parameter transforms.

This module defines the abstract base class that all parameter transforms (both sampling-time and simulation-time) must inherit from.

ParameterTransform

Bases: ABC

Abstract base class for all parameter transforms.

A parameter transform is a single, focused operation that modifies theta parameters. Transforms can be composed into pipelines and are used at two phases:

Sampling phase: Applied after parameter sampling to enforce constraints (e.g., ensure a > z by swapping values)
Simulation phase: Applied before calling the simulator to prepare parameters for the expected format (e.g., stack v0, v1 → v)

Subclasses must implement the apply method. The method signature includes optional model_config and n_trials arguments to support both simple constraints (which only need theta) and complex adapters (which may need additional context).

Examples:

Create a custom transform:

>>> class ScaleParameter(ParameterTransform):
...     def __init__(self, param_name: str, scale: float):
...         self.param_name = param_name
...         self.scale = scale
...
...     def apply(self, theta, model_config=None, n_trials=None):
...         if self.param_name in theta:
...             theta[self.param_name] = theta[self.param_name] * self.scale
...         return theta

Use in model config:

>>> model_config = {
...     "name": "my_model",
...     "parameter_transforms": {
...         "sampling": [SwapIfLessConstraint("a", "z")],
...         "simulation": [ColumnStackParameters(["v0", "v1"], "v")],
...     }
... }

repr

__repr__()

String representation for debugging.

Returns:

str –

A string describing this transform.

str

__str__()

Human-readable string representation.

apply `abstractmethod`

apply(theta, model_config=None, n_trials=None)

Apply transform to theta parameters.

This method should modify the theta dictionary in place and return it. It may add new parameters, modify existing parameters, or remove parameters as needed.

Parameters:

theta (dict[str, Any]) –

Dictionary of model parameters. Values are typically numpy arrays.
model_config (dict[str, Any] or None, default: None ) –

Model configuration dictionary. Available for transforms that need additional context (e.g., looking up simulator_fixed_params). Sampling-time transforms typically ignore this.
n_trials (int or None, default: None ) –

Number of trials. Available for transforms that need to create arrays of a specific size. Sampling-time transforms typically ignore this.

Returns:

dict[str, Any] –

The modified theta dictionary (usually the same object passed in).

Notes

Transforms should be pure when possible (no side effects)
If creating new arrays, use dtype=np.float32 for consistency
Document any parameters that are added, modified, or removed

sampling

Parameter transforms for the training data generation workflow.

These transforms are applied during the parameter sampling stage of the training data generation workflow. They enforce parameter relationships when generating synthetic training data for likelihood approximation networks.

Note: These are NOT directly relevant for basic Simulator usage, which uses simulation transforms via ParameterSimulatorAdapters instead.

Examples: - SwapIfLessConstraint: Ensure a > z by swapping when violated - NormalizeToSumConstraint: Normalize parameters to sum to 1

NormalizeToSumConstraint

NormalizeToSumConstraint(param_names)

Bases: ParameterTransform

Normalize a set of parameters to sum to 1.

This constraint is commonly used for drift rates in multi-accumulator models (e.g., RLWM models) where the drift rates must sum to 1.

The normalization includes a small epsilon (1e-20) to avoid division by zero and ensure numerical stability.

Parameters:

param_names (list[str]) –

List of parameter names to normalize together

Examples:

>>> constraint = NormalizeToSumConstraint(["v1", "v2", "v3"])
>>> theta = {"v1": 0.2, "v2": 0.3, "v3": 0.5}
>>> result = constraint.apply(theta)
>>> sum([result["v1"], result["v2"], result["v3"]])
1.0

Use in model config:

>>> model_config = {
...     "name": "dev_rlwm_lba_race_v1",
...     "parameter_transforms": {
...         "sampling": [
...             NormalizeToSumConstraint(["vRL0", "vRL1", "vRL2"]),
...             NormalizeToSumConstraint(["vWM0", "vWM1", "vWM2"]),
...         ],
...     }
... }

Parameters:

param_names (list[str]) –

List of parameter names to normalize together

apply

apply(
    theta, model_config=None, n_trials=None, epsilon=1e-20
)

Normalize specified parameters to sum to 1.

Parameters:

theta (dict[str, Any]) –

Dictionary of parameter values (can be scalars or arrays)
model_config (dict[str, Any] or None, default: None ) –

Not used by this transform (included for interface compatibility)
n_trials (int or None, default: None ) –

Not used by this transform (included for interface compatibility)
epsilon (float, default: 1e-20 ) –

Small value to add to the total to avoid division by zero. Default is 1e-20.

Returns:

dict[str, Any] –

Modified theta dictionary with normalized values

SwapIfLessConstraint

SwapIfLessConstraint(param_a, param_b)

Bases: ParameterTransform

Swap two parameters if first is less than or equal to second.

This constraint ensures ordering relationships like a > z by swapping the values when the constraint is violated. Works with both scalar and array inputs.

Parameters:

param_a (str) –

First parameter name (should be > param_b after constraint)
param_b (str) –

Second parameter name (should be < param_a after constraint)

Examples:

>>> constraint = SwapIfLessConstraint("a", "z")
>>> theta = {"a": np.array([0.5, 1.5]), "z": np.array([1.0, 0.5])}
>>> result = constraint.apply(theta)
>>> # First element: a=0.5 <= z=1.0 → swap → a=1.0, z=0.5
>>> # Second element: a=1.5 > z=0.5 → no swap → a=1.5, z=0.5

Use in model config:

>>> model_config = {
...     "name": "lba_angle_3",
...     "parameter_transforms": {
...         "sampling": [SwapIfLessConstraint("a", "z")],
...     }
... }

Parameters:

param_a (str) –

First parameter name (should be > param_b after constraint)
param_b (str) –

Second parameter name (should be < param_a after constraint)

apply

apply(theta, model_config=None, n_trials=None)

Swap param_a and param_b where param_a <= param_b.

Parameters:

theta (dict[str, Any]) –

Dictionary of parameter values (can be scalars or arrays)
model_config (dict[str, Any] or None, default: None ) –

Not used by this transform (included for interface compatibility)
n_trials (int or None, default: None ) –

Not used by this transform (included for interface compatibility)

Returns:

dict[str, Any] –

Modified theta dictionary with swapped values where needed

normalize

Normalize transform for drift rate constraints.

NormalizeToSumConstraint

NormalizeToSumConstraint(param_names)

Bases: ParameterTransform

Normalize a set of parameters to sum to 1.

This constraint is commonly used for drift rates in multi-accumulator models (e.g., RLWM models) where the drift rates must sum to 1.

The normalization includes a small epsilon (1e-20) to avoid division by zero and ensure numerical stability.

Parameters:

param_names (list[str]) –

List of parameter names to normalize together

Examples:

>>> constraint = NormalizeToSumConstraint(["v1", "v2", "v3"])
>>> theta = {"v1": 0.2, "v2": 0.3, "v3": 0.5}
>>> result = constraint.apply(theta)
>>> sum([result["v1"], result["v2"], result["v3"]])
1.0

Use in model config:

>>> model_config = {
...     "name": "dev_rlwm_lba_race_v1",
...     "parameter_transforms": {
...         "sampling": [
...             NormalizeToSumConstraint(["vRL0", "vRL1", "vRL2"]),
...             NormalizeToSumConstraint(["vWM0", "vWM1", "vWM2"]),
...         ],
...     }
... }

Parameters:

param_names (list[str]) –

List of parameter names to normalize together

apply

apply(
    theta, model_config=None, n_trials=None, epsilon=1e-20
)

Normalize specified parameters to sum to 1.

Parameters:

theta (dict[str, Any]) –

Dictionary of parameter values (can be scalars or arrays)
model_config (dict[str, Any] or None, default: None ) –

Not used by this transform (included for interface compatibility)
n_trials (int or None, default: None ) –

Not used by this transform (included for interface compatibility)
epsilon (float, default: 1e-20 ) –

Small value to add to the total to avoid division by zero. Default is 1e-20.

Returns:

dict[str, Any] –

Modified theta dictionary with normalized values

swap

Swap transform for ensuring parameter ordering constraints.

SwapIfLessConstraint

SwapIfLessConstraint(param_a, param_b)

Bases: ParameterTransform

Swap two parameters if first is less than or equal to second.

This constraint ensures ordering relationships like a > z by swapping the values when the constraint is violated. Works with both scalar and array inputs.

Parameters:

param_a (str) –

First parameter name (should be > param_b after constraint)
param_b (str) –

Second parameter name (should be < param_a after constraint)

Examples:

>>> constraint = SwapIfLessConstraint("a", "z")
>>> theta = {"a": np.array([0.5, 1.5]), "z": np.array([1.0, 0.5])}
>>> result = constraint.apply(theta)
>>> # First element: a=0.5 <= z=1.0 → swap → a=1.0, z=0.5
>>> # Second element: a=1.5 > z=0.5 → no swap → a=1.5, z=0.5

Use in model config:

>>> model_config = {
...     "name": "lba_angle_3",
...     "parameter_transforms": {
...         "sampling": [SwapIfLessConstraint("a", "z")],
...     }
... }

Parameters:

param_a (str) –

First parameter name (should be > param_b after constraint)
param_b (str) –

Second parameter name (should be < param_a after constraint)

apply

apply(theta, model_config=None, n_trials=None)

Swap param_a and param_b where param_a <= param_b.

Parameters:

theta (dict[str, Any]) –

Dictionary of parameter values (can be scalars or arrays)
model_config (dict[str, Any] or None, default: None ) –

Not used by this transform (included for interface compatibility)
n_trials (int or None, default: None ) –

Not used by this transform (included for interface compatibility)

Returns:

dict[str, Any] –

Modified theta dictionary with swapped values where needed

simulation

Simulation-time parameter transforms.

These transforms are applied right before calling the simulator to prepare parameters for the expected format.

Examples: - ColumnStackParameters: Stack v0, v1, v2 into single array v - ExpandDimension: Add dimension to parameters for multi-particle models - SetZeroArray: Create zero-filled arrays for dynamic drift models

ApplyMapping

ApplyMapping(
    source_param,
    target_param,
    mapping_key,
    additional_sources=None,
    delete_source=False,
)

Bases: ParameterTransform

Apply a mapping function to transform a parameter.

Uses a mapping from model_config to transform a parameter value. Commonly used for random variable distributions.

Parameters:

source_param (str) –

Parameter to read from
target_param (str) –

Parameter to write to
mapping_key (str) –

Key in model_config["simulator_param_mappings"]
additional_sources (list[str], default: None ) –

Additional parameters to pass to the mapping function
delete_source (bool, default: False ) –

Whether to delete the source parameter after mapping

apply

apply(theta, model_config=None, n_trials=None)

Apply mapping function to transform parameter.

ColumnStackParameters

ColumnStackParameters(
    source_params, target_param, delete_sources=True
)

Bases: ParameterTransform

Stack multiple parameters into a single multi-column array.

Takes individual parameters (e.g., v0, v1, v2) and stacks them column-wise into a single array (e.g., v with shape (n_trials, 3)).

Parameters:

source_params (list[str]) –

Names of parameters to stack
target_param (str) –

Name of the resulting stacked parameter
delete_sources (bool, default: True ) –

Whether to delete the source parameters after stacking

Examples:

>>> transform = ColumnStackParameters(["v0", "v1", "v2"], "v")
>>> theta = {"v0": [0.5], "v1": [0.6], "v2": [0.7]}
>>> theta = transform.apply(theta)
>>> theta["v"]  # array([[0.5, 0.6, 0.7]])

apply

apply(theta, model_config=None, n_trials=None)

Stack parameters column-wise.

ConditionalAdaptation

ConditionalAdaptation(condition, transformation)

Bases: ParameterTransform

Apply a transform only if a condition is met.

Parameters:

condition (Callable) –

Function that takes (theta, model_config, n_trials) and returns bool
transformation (ParameterTransform) –

Transform to apply if condition is True

Examples:

>>> transform = ConditionalAdaptation(
...     condition=lambda theta, cfg, n: "a" in theta,
...     transformation=ExpandDimension(["a"])
... )

apply

apply(theta, model_config=None, n_trials=None)

Apply transform if condition is met.

DeleteParameters

DeleteParameters(param_names)

Bases: ParameterTransform

Delete specified parameters from theta.

Parameters:

param_names (list[str]) –

Names of parameters to delete

Examples:

>>> transform = DeleteParameters(["A", "b"])
>>> theta = {"v": [0.5], "A": [0.5], "b": [1.0]}
>>> theta = transform.apply(theta)
>>> "A" in theta  # False

apply

apply(theta, model_config=None, n_trials=None)

Delete specified parameters.

ExpandDimension

ExpandDimension(param_names)

Bases: ParameterTransform

Expand dimensions of specified parameters.

Converts 1D arrays of shape (n_trials,) to shape (n_trials, 1) by adding a dimension. This is commonly needed for multi-particle models.

Parameters:

param_names (list[str]) –

Names of parameters to expand

Examples:

>>> transform = ExpandDimension(["a", "t"])
>>> theta = {"a": np.array([1.0, 2.0]), "t": np.array([0.3, 0.3])}
>>> theta = transform.apply(theta)
>>> theta["a"].shape  # (2, 1)

apply

apply(theta, model_config=None, n_trials=None)

Expand dimensions of specified parameters.

LambdaAdaptation

LambdaAdaptation(func, name=None)

Bases: ParameterTransform

Wrap a lambda or callable function as a transform.

Parameters:

func (Callable) –

Function with signature: func(theta, model_config, n_trials) → theta
name (str, default: None ) –

Optional name for debugging

Examples:

>>> transform = LambdaAdaptation(
...     lambda theta, cfg, n: theta.update({"nact": 3}) or theta,
...     name="set_nact_to_3"
... )

apply

apply(theta, model_config=None, n_trials=None)

Apply the wrapped function.

RenameParameter

RenameParameter(old_name, new_name, transform_fn=None)

Bases: ParameterTransform

Rename a parameter and optionally transform its value.

Parameters:

old_name (str) –

Current parameter name
new_name (str) –

New parameter name
transform_fn (Callable or None, default: None ) –

Optional function to apply to the value during renaming

Examples:

>>> transform = RenameParameter("A", "z", lambda x: np.expand_dims(x, axis=1))
>>> theta = {"A": np.array([0.5])}
>>> theta = transform.apply(theta)
>>> "z" in theta  # True

apply

apply(theta, model_config=None, n_trials=None)

Rename parameter and optionally transform its value.

SetDefaultValue

SetDefaultValue(param_name, default_value, dtype=float32)

Bases: ParameterTransform

Set a parameter to a default value if not present.

This transform adds a parameter with a default value if it doesn't already exist in theta. The value is automatically tiled to match n_trials.

Parameters:

param_name (str) –

Name of the parameter to set
default_value (float or int or ndarray) –

Default value to use. If scalar, will be tiled to (n_trials,).
dtype (dtype, default: float32 ) –

Data type for the parameter. Defaults to np.float32.

Examples:

>>> transform = SetDefaultValue("nact", 3)
>>> theta = {}
>>> theta = transform.apply(theta, {}, n_trials=5)
>>> theta["nact"]  # array([3., 3., 3., 3., 3.], dtype=float32)

apply

apply(theta, model_config=None, n_trials=None)

Add parameter with default value if not present.

SetZeroArray

SetZeroArray(param_name, shape=None, dtype=float32)

Bases: ParameterTransform

Set a parameter to an array of zeros.

Parameters:

param_name (str) –

Name of the parameter to set
shape (tuple or None, default: None ) –

Shape of the array. Use None for n_trials dimension.
dtype (dtype, default: float32 ) –

Data type for the array. Defaults to np.float32.

Examples:

>>> transform = SetZeroArray("v")
>>> theta = transform.apply({}, {}, n_trials=5)
>>> theta["v"].shape  # (5,)

apply

apply(theta, model_config=None, n_trials=None)

Set parameter to zero array.

TileArray

TileArray(param_name, value, dtype=float32)

Bases: ParameterTransform

Tile a constant value across trials.

Parameters:

param_name (str) –

Name of the parameter to create
value (float or array - like) –

Value(s) to tile
dtype (dtype, default: float32 ) –

Data type for the array. Defaults to np.float32.

Examples:

>>> transform = TileArray("z", 0.5)
>>> theta = transform.apply({}, {}, n_trials=3)
>>> theta["z"]  # array([0.5, 0.5, 0.5], dtype=float32)

apply

apply(theta, model_config=None, n_trials=None)

Tile value across trials.

common

Common simulation-time parameter transforms.

This module provides a library of commonly-used parameter transforms that prepare theta parameters for simulator consumption.

ApplyMapping

ApplyMapping(
    source_param,
    target_param,
    mapping_key,
    additional_sources=None,
    delete_source=False,
)

Bases: ParameterTransform

Apply a mapping function to transform a parameter.

Uses a mapping from model_config to transform a parameter value. Commonly used for random variable distributions.

Parameters:

source_param (str) –

Parameter to read from
target_param (str) –

Parameter to write to
mapping_key (str) –

Key in model_config["simulator_param_mappings"]
additional_sources (list[str], default: None ) –

Additional parameters to pass to the mapping function
delete_source (bool, default: False ) –

Whether to delete the source parameter after mapping

apply

apply(theta, model_config=None, n_trials=None)

Apply mapping function to transform parameter.

ColumnStackParameters

ColumnStackParameters(
    source_params, target_param, delete_sources=True
)

Bases: ParameterTransform

Stack multiple parameters into a single multi-column array.

Takes individual parameters (e.g., v0, v1, v2) and stacks them column-wise into a single array (e.g., v with shape (n_trials, 3)).

Parameters:

source_params (list[str]) –

Names of parameters to stack
target_param (str) –

Name of the resulting stacked parameter
delete_sources (bool, default: True ) –

Whether to delete the source parameters after stacking

Examples:

>>> transform = ColumnStackParameters(["v0", "v1", "v2"], "v")
>>> theta = {"v0": [0.5], "v1": [0.6], "v2": [0.7]}
>>> theta = transform.apply(theta)
>>> theta["v"]  # array([[0.5, 0.6, 0.7]])

apply

apply(theta, model_config=None, n_trials=None)

Stack parameters column-wise.

ConditionalAdaptation

ConditionalAdaptation(condition, transformation)

Bases: ParameterTransform

Apply a transform only if a condition is met.

Parameters:

condition (Callable) –

Function that takes (theta, model_config, n_trials) and returns bool
transformation (ParameterTransform) –

Transform to apply if condition is True

Examples:

>>> transform = ConditionalAdaptation(
...     condition=lambda theta, cfg, n: "a" in theta,
...     transformation=ExpandDimension(["a"])
... )

apply

apply(theta, model_config=None, n_trials=None)

Apply transform if condition is met.

DeleteParameters

DeleteParameters(param_names)

Bases: ParameterTransform

Delete specified parameters from theta.

Parameters:

param_names (list[str]) –

Names of parameters to delete

Examples:

>>> transform = DeleteParameters(["A", "b"])
>>> theta = {"v": [0.5], "A": [0.5], "b": [1.0]}
>>> theta = transform.apply(theta)
>>> "A" in theta  # False

apply

apply(theta, model_config=None, n_trials=None)

Delete specified parameters.

ExpandDimension

ExpandDimension(param_names)

Bases: ParameterTransform

Expand dimensions of specified parameters.

Converts 1D arrays of shape (n_trials,) to shape (n_trials, 1) by adding a dimension. This is commonly needed for multi-particle models.

Parameters:

param_names (list[str]) –

Names of parameters to expand

Examples:

>>> transform = ExpandDimension(["a", "t"])
>>> theta = {"a": np.array([1.0, 2.0]), "t": np.array([0.3, 0.3])}
>>> theta = transform.apply(theta)
>>> theta["a"].shape  # (2, 1)

apply

apply(theta, model_config=None, n_trials=None)

Expand dimensions of specified parameters.

LambdaAdaptation

LambdaAdaptation(func, name=None)

Bases: ParameterTransform

Wrap a lambda or callable function as a transform.

Parameters:

func (Callable) –

Function with signature: func(theta, model_config, n_trials) → theta
name (str, default: None ) –

Optional name for debugging

Examples:

>>> transform = LambdaAdaptation(
...     lambda theta, cfg, n: theta.update({"nact": 3}) or theta,
...     name="set_nact_to_3"
... )

apply

apply(theta, model_config=None, n_trials=None)

Apply the wrapped function.

RenameParameter

RenameParameter(old_name, new_name, transform_fn=None)

Bases: ParameterTransform

Rename a parameter and optionally transform its value.

Parameters:

old_name (str) –

Current parameter name
new_name (str) –

New parameter name
transform_fn (Callable or None, default: None ) –

Optional function to apply to the value during renaming

Examples:

>>> transform = RenameParameter("A", "z", lambda x: np.expand_dims(x, axis=1))
>>> theta = {"A": np.array([0.5])}
>>> theta = transform.apply(theta)
>>> "z" in theta  # True

apply

apply(theta, model_config=None, n_trials=None)

Rename parameter and optionally transform its value.

SetDefaultValue

SetDefaultValue(param_name, default_value, dtype=float32)

Bases: ParameterTransform

Set a parameter to a default value if not present.

This transform adds a parameter with a default value if it doesn't already exist in theta. The value is automatically tiled to match n_trials.

Parameters:

param_name (str) –

Name of the parameter to set
default_value (float or int or ndarray) –

Default value to use. If scalar, will be tiled to (n_trials,).
dtype (dtype, default: float32 ) –

Data type for the parameter. Defaults to np.float32.

Examples:

>>> transform = SetDefaultValue("nact", 3)
>>> theta = {}
>>> theta = transform.apply(theta, {}, n_trials=5)
>>> theta["nact"]  # array([3., 3., 3., 3., 3.], dtype=float32)

apply

apply(theta, model_config=None, n_trials=None)

Add parameter with default value if not present.

SetZeroArray

SetZeroArray(param_name, shape=None, dtype=float32)

Bases: ParameterTransform

Set a parameter to an array of zeros.

Parameters:

param_name (str) –

Name of the parameter to set
shape (tuple or None, default: None ) –

Shape of the array. Use None for n_trials dimension.
dtype (dtype, default: float32 ) –

Data type for the array. Defaults to np.float32.

Examples:

>>> transform = SetZeroArray("v")
>>> theta = transform.apply({}, {}, n_trials=5)
>>> theta["v"].shape  # (5,)

apply

apply(theta, model_config=None, n_trials=None)

Set parameter to zero array.

TileArray

TileArray(param_name, value, dtype=float32)

Bases: ParameterTransform

Tile a constant value across trials.

Parameters:

param_name (str) –

Name of the parameter to create
value (float or array - like) –

Value(s) to tile
dtype (dtype, default: float32 ) –

Data type for the array. Defaults to np.float32.

Examples:

>>> transform = TileArray("z", 0.5)
>>> theta = transform.apply({}, {}, n_trials=3)
>>> theta["z"]  # array([0.5, 0.5, 0.5], dtype=float32)

apply

apply(theta, model_config=None, n_trials=None)

Tile value across trials.

ssms

ssms.Simulator

config property

parameter_adapter property

simulate

validate_params

ssms.basic_simulators

Simulator

config property

parameter_adapter property

simulate

validate_params

boundary_functions

addm_collapse

angle module-attribute

conflict_gamma module-attribute

constant module-attribute

generalized_logistic module-attribute

weibull_cdf module-attribute

drift_functions

attend_drift module-attribute

attend_drift_simple

conflict_ds_drift module-attribute

conflict_dsstimflex_drift module-attribute

conflict_stimflex_drift module-attribute

constant module-attribute

ds_support_analytic module-attribute

gamma_drift module-attribute

stimflex_support

fixation_continuation

draw_durations

generate_schedule

prolong_last_fixation

resample_all_fixations

resolve_continuation_mode

resolve_distribution

sample_continuation

inv_temp_softmax

inv_temp_softmax

modular_parameter_simulator_adapter

ModularParameterSimulatorAdapter

adapt_parameters

ParameterSimulatorAdapterProtocol

adapt_parameters

parameter_adapters

ApplyMapping

apply

ColumnStackParameters

apply

ConditionalAdaptation

apply

DeleteParameters

apply

ExpandDimension

apply

LambdaAdaptation

apply

ParameterAdapterRegistry

__repr__

describe

get_processor

has_processor

list_registered_families

list_registered_models

register_family

register_model

RenameParameter

apply

SetDefaultValue

apply

SetZeroArray

apply

TileArray

apply

base

ParameterTransform

__repr__

__str__

apply abstractmethod

get_adapter_registry

config `property`

parameter_adapter `property`

config `property`

parameter_adapter `property`

angle `module-attribute`

conflict_gamma `module-attribute`

constant `module-attribute`

generalized_logistic `module-attribute`

weibull_cdf `module-attribute`

attend_drift `module-attribute`

conflict_ds_drift `module-attribute`

conflict_dsstimflex_drift `module-attribute`

conflict_stimflex_drift `module-attribute`

constant `module-attribute`

ds_support_analytic `module-attribute`

gamma_drift `module-attribute`

repr

repr

str

apply `abstractmethod`

repr

config `property`

parameter_adapter `property`

add_boundary `staticmethod`

add_drift `staticmethod`

from_model `staticmethod`

from_scratch `staticmethod`

get_sampling_transforms `staticmethod`

get_simulation_transforms `staticmethod`

get_transforms `staticmethod`

minimal_config `staticmethod`

validate_config `staticmethod`

with_deadline `staticmethod`

repr

repr

add_boundary `staticmethod`