Skip to content

data generators

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:

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:

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.

See Also

KDEEstimatorBuilder : Builder for KDE-based likelihood estimators

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:

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.

See Also

KDEEstimatorBuilder : Builder for KDE-based likelihood estimators

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:

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.

ssms.dataset_generators.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.

ssms.dataset_generators.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.

ssms.dataset_generators.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,)

ssms.dataset_generators.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.

ssms.dataset_generators.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:

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.

ssms.dataset_generators.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