The TrainingDataGenerator class¶

This tutorial demonstrates how to use the TrainingDataGenerator class to create training data for sequential sampling models (SSMs). We'll discuss a bit of the overall workflow design here as well.

In [1]:

# Standard library imports
import numpy as np
import matplotlib.pyplot as plt
import time

# SSMS imports
from ssms.dataset_generators.lan_mlp import TrainingDataGenerator
from ssms.config import model_config, get_default_generator_config
from ssms.basic_simulators.simulator_class import Simulator

# Set random seed for reproducibility
np.random.seed(42)

print("✓ Imports successful")
print(f"✓ Available models: {len(model_config)} predefined models")

# Standard library imports import numpy as np import matplotlib.pyplot as plt import time # SSMS imports from ssms.dataset_generators.lan_mlp import TrainingDataGenerator from ssms.config import model_config, get_default_generator_config from ssms.basic_simulators.simulator_class import Simulator # Set random seed for reproducibility np.random.seed(42) print("✓ Imports successful") print(f"✓ Available models: {len(model_config)} predefined models")

✓ Imports successful
✓ Available models: 108 predefined models

What is TrainingDataGenerator?¶

The TrainingDataGenerator orchestrates the creation of training datasets for neural network-based likelihood approximation.

It:

• Samples parameters from valid parameter spaces
• Generates RT/choice data (via simulation or analytical methods)
• Turns the raw data into (feature, label) pairs for various networks we are interested in downstream
• Optionally saves outputs via .pickle files

The TrainingDataGenerator is in fact a light-weight class that puts to work a DataPipeline.

High-Level Flow¶

INPUTS

⚙️

config

Polymorphic parameter

Option 1: Dict (99%)

generator_config dict (samples, bins, estimator_type)

Option 2: Strategy (1%)

Custom strategy object (advanced customization)

📋

model_config

Model specification

Which SSM model? (DDM, angle, LCA...)
Parameter bounds
Boundary/drift functions

⬇

📊

TrainingDataGenerator

Orchestrates the complete data generation workflow

⬇

OUTPUT

📦

Training Data

lan_data
RT/choice pairs for training

lan_labels
Log-likelihoods

theta
Parameter values

cpn_labels
Choice probabilities

... and more

The Single Injection Point: Pipeline¶

Pipeline concept?¶

The generation_pipeline parameter is the single customization point for the TrainingDataGenerator. It encapsulates the entire data generation workflow, from parameter sampling to the composition of the final training data.

Built-in Strategies:¶

1. SimulationPipeline (default for estimator_type='kde')

• Uses the core native simulators of the ssm-simulators package
• Goes through a KDE as the Likelihood Estimator (as in our LAN paper )

2. PyDDMPipeline (used for estimator_type='pyddm')

• Uses analytical Fokker-Planck PDE solver via PyDDM package
• Model choice somewhat more limited (2-choice, Gaussian noise etc.)
• Can be much faster than KDE based strategy where it applies

3. Custom Strategy

• Implement your own DataPipelineProtocol
• Full control over parameter sampling, simulation, and data structuring

How to Control It?¶

Pass a config dict as the first argument

Basic Example¶

In [2]:

model_name = "ddm"
estimator_type = "kde"

model_name = "ddm" estimator_type = "kde"

In [3]:

# 
my_model_config = model_config[model_name]# For simulation-based (default)

pipeline_config = get_default_generator_config()
pipeline_config["estimator"]['type'] = estimator_type  # or just omit, this is the default
pipeline_config["pipeline"]["n_parameter_sets"] = 100
pipeline_config["simulator"]["n_samples"] = 5000
print(pipeline_config)
gen = TrainingDataGenerator(pipeline_config,
                    my_model_config
                    )

# my_model_config = model_config[model_name]# For simulation-based (default) pipeline_config = get_default_generator_config() pipeline_config["estimator"]['type'] = estimator_type # or just omit, this is the default pipeline_config["pipeline"]["n_parameter_sets"] = 100 pipeline_config["simulator"]["n_samples"] = 5000 print(pipeline_config) gen = TrainingDataGenerator(pipeline_config, my_model_config )

{'pipeline': {'n_parameter_sets': 100, 'n_parameter_sets_rejected': 100, 'n_subruns': 10, 'n_cpus': 'all', 'simulation_filters': {'mode': 20, 'choice_cnt': 0, 'mean_rt': 17, 'std': 0, 'mode_cnt_rel': 0.95}}, 'estimator': {'type': 'kde', 'kde_displace_t': False, 'pdf_interpolation': 'cubic', 'max_undecided_prob': 0.5}, 'training': {'n_samples_per_param': 1000, 'mixture_probabilities': [0.8, 0.1, 0.1], 'separate_response_channels': False, 'negative_rt_log_likelihood': -66.77497}, 'simulator': {'n_samples': 5000, 'delta_t': 0.001, 'max_t': 20.0, 'smooth_unif': True}, 'output': {'folder': 'data/lan_mlp/', 'nbins': 0, 'pickle_protocol': 4, 'bin_pointwise': False}, 'model': 'ddm'}

Generate training data¶

In [4]:

# Generate data
print("Generating training data...")
start_time = time.time()

training_data = gen.generate_data_training()

elapsed = time.time() - start_time
total_trials = pipeline_config['pipeline']['n_parameter_sets'] * pipeline_config['simulator']['n_samples']
print(f"✓ Data generation complete in {elapsed:.2f} seconds")
print(f"  ({total_trials / elapsed:.0f} trials/sec)")

# Generate data print("Generating training data...") start_time = time.time() training_data = gen.generate_data_training() elapsed = time.time() - start_time total_trials = pipeline_config['pipeline']['n_parameter_sets'] * pipeline_config['simulator']['n_samples'] print(f"✓ Data generation complete in {elapsed:.2f} seconds") print(f" ({total_trials / elapsed:.0f} trials/sec)")

Generating training data...
✓ Data generation complete in 5.97 seconds
  (83708 trials/sec)

Inspect output¶

In [5]:

print("Output structure:")
print(f"  Keys: {list(training_data.keys())}")

print("\nData shapes:")
for key, value in training_data.items():
    if value is not None:
        if isinstance(value, np.ndarray):
            print(f"  {key:30s}: {value.shape}")
        elif isinstance(value, dict):
            print(f"  {key:30s}: {value}")
    else:
        print(f"  {key}")
        print(f"  {key:30s}: None")

print("\n--- Understanding the components ---")
print("data: RT/choice pairs [n_parameter_sets, n_samples, 2]")
print("theta: Parameter values [n_parameter_sets, n_params]")
print("choice_p: Choice probabilities for each trial")
print("cpn_*, opn_*, gonogo_*: Additional training labels")

print("Output structure:") print(f" Keys: {list(training_data.keys())}") print("\nData shapes:") for key, value in training_data.items(): if value is not None: if isinstance(value, np.ndarray): print(f" {key:30s}: {value.shape}") elif isinstance(value, dict): print(f" {key:30s}: {value}") else: print(f" {key}") print(f" {key:30s}: None") print("\n--- Understanding the components ---") print("data: RT/choice pairs [n_parameter_sets, n_samples, 2]") print("theta: Parameter values [n_parameter_sets, n_params]") print("choice_p: Choice probabilities for each trial") print("cpn_*, opn_*, gonogo_*: Additional training labels")

Output structure:
  Keys: ['gonogo_data', 'binned_128', 'cpn_data', 'lan_labels', 'lan_data', 'cpn_labels', 'cpn_no_omission_data', 'opn_labels', 'cpn_no_omission_labels', 'binned_256', 'gonogo_labels', 'theta', 'opn_data', 'generator_config', 'model_config']

Data shapes:
  gonogo_data                   : (100, 4)
  binned_128                    : (100, 128, 2)
  cpn_data                      : (100, 4)
  lan_labels                    : (100000,)
  lan_data                      : (100000, 6)
  cpn_labels                    : (100, 1)
  cpn_no_omission_data          : (100, 4)
  opn_labels                    : (100, 1)
  cpn_no_omission_labels        : (100, 1)
  binned_256                    : (100, 256, 2)
  gonogo_labels                 : (100, 1)
  theta                         : (100, 4)
  opn_data                      : (100, 4)
  generator_config              : {'pipeline': {'n_parameter_sets': 100, 'n_parameter_sets_rejected': 100, 'n_subruns': 10, 'n_cpus': 12, 'simulation_filters': {'mode': 20, 'choice_cnt': 0, 'mean_rt': 17, 'std': 0, 'mode_cnt_rel': 0.95}}, 'estimator': {'type': 'kde', 'kde_displace_t': False, 'pdf_interpolation': 'cubic', 'max_undecided_prob': 0.5, 'displace_t': False}, 'training': {'n_samples_per_param': 1000, 'mixture_probabilities': [0.8, 0.1, 0.1], 'separate_response_channels': False, 'negative_rt_log_likelihood': -66.77497}, 'simulator': {'n_samples': 5000, 'delta_t': 0.001, 'max_t': 20.0, 'smooth_unif': True}, 'output': {'folder': 'data/lan_mlp/', 'nbins': 0, 'pickle_protocol': 4, 'bin_pointwise': False}, 'model': 'ddm'}
  model_config                  : {'name': 'ddm', 'params': ['v', 'a', 'z', 't'], 'param_bounds': [[-3.0, 0.3, 0.1, 0.0], [3.0, 2.5, 0.9, 2.0]], 'boundary_name': 'constant', 'boundary': <function constant at 0x11ad24b80>, 'boundary_params': [], 'n_params': 4, 'default_params': [0.0, 1.0, 0.5, 0.001], 'nchoices': 2, 'choices': [-1, 1], 'n_particles': 1, 'simulator': <cyfunction ddm_flexbound at 0x12d85b930>, 'parameter_transforms': {'sampling': [], 'simulation': []}, 'param_bounds_dict': {'v': (-3.0, 3.0), 'a': (0.3, 2.5), 'z': (0.1, 0.9), 't': (0.0, 2.0)}}

--- Understanding the components ---
data: RT/choice pairs [n_parameter_sets, n_samples, 2]
theta: Parameter values [n_parameter_sets, n_params]
choice_p: Choice probabilities for each trial
cpn_*, opn_*, gonogo_*: Additional training labels

Advanced Example¶

In [6]:

from ssms.dataset_generators.estimator_builders.kde_builder import KDEEstimatorBuilder

from ssms.dataset_generators.pipelines import SimulationPipeline
from ssms.dataset_generators.strategies import ResampleMixtureStrategy

model_config_advanced = model_config[model_name]# For simulation-based (default)

pipeline_config_advanced = get_default_generator_config()
pipeline_config_advanced["estimator"]["type"] = "kde"  # or just omit, this is the default
pipeline_config_advanced["pipeline"]["n_parameter_sets"] = 100
pipeline_config_advanced["simulator"]["n_samples"] = 5000

# Create custom pipeline with specialized components
custom_pipeline = SimulationPipeline(
    generator_config=pipeline_config_advanced,
    model_config=model_config_advanced,
    estimator_builder=KDEEstimatorBuilder,
    training_strategy=ResampleMixtureStrategy,
)

# Pass the strategy directly as the first positional argument
gen_advanced = TrainingDataGenerator(
    config = custom_pipeline
    ) # no need for `model_config` it's an attribute of the pipeline

from ssms.dataset_generators.estimator_builders.kde_builder import KDEEstimatorBuilder from ssms.dataset_generators.pipelines import SimulationPipeline from ssms.dataset_generators.strategies import ResampleMixtureStrategy model_config_advanced = model_config[model_name]# For simulation-based (default) pipeline_config_advanced = get_default_generator_config() pipeline_config_advanced["estimator"]["type"] = "kde" # or just omit, this is the default pipeline_config_advanced["pipeline"]["n_parameter_sets"] = 100 pipeline_config_advanced["simulator"]["n_samples"] = 5000 # Create custom pipeline with specialized components custom_pipeline = SimulationPipeline( generator_config=pipeline_config_advanced, model_config=model_config_advanced, estimator_builder=KDEEstimatorBuilder, training_strategy=ResampleMixtureStrategy, ) # Pass the strategy directly as the first positional argument gen_advanced = TrainingDataGenerator( config = custom_pipeline ) # no need for `model_config` it's an attribute of the pipeline

Advanced Example: Generate Training Data¶

In [7]:

# Generate data
print("Generating training data...")
start_time = time.time()

training_data = gen.generate_data_training()

elapsed = time.time() - start_time
total_trials = pipeline_config['pipeline']['n_parameter_sets'] * pipeline_config['simulator']['n_samples']
print(f"✓ Data generation complete in {elapsed:.2f} seconds")
print(f"  ({total_trials / elapsed:.0f} trials/sec)")

# Generate data print("Generating training data...") start_time = time.time() training_data = gen.generate_data_training() elapsed = time.time() - start_time total_trials = pipeline_config['pipeline']['n_parameter_sets'] * pipeline_config['simulator']['n_samples'] print(f"✓ Data generation complete in {elapsed:.2f} seconds") print(f" ({total_trials / elapsed:.0f} trials/sec)")

Generating training data...
✓ Data generation complete in 2.29 seconds
  (218395 trials/sec)

Default Behavior¶

When you pass a config dict as the first argument, TrainingDataGenerator auto-creates the appropriate strategy based on estimator_type:

• estimator_type='kde' → SimulationPipeline
• estimator_type='pyddm' → PyDDMPipeline

This should cover most basic use cases!