diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml index 2b6757c0..811d3823 100644 --- a/.github/workflows/run_tests.yml +++ b/.github/workflows/run_tests.yml @@ -49,4 +49,4 @@ jobs: files: /.coverage.xml flags: unittests name: codecov-${{ matrix.python-version }} - fail_ci_if_error: false \ No newline at end of file + fail_ci_if_error: false diff --git a/MANIFEST.in b/MANIFEST.in index a97a4048..5bf8b946 100755 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -27,4 +27,4 @@ global-exclude .git global-exclude *~ global-exclude *.pyc global-exclude .#* -global-exclude __pycache__ \ No newline at end of file +global-exclude __pycache__ diff --git a/docs/api/basic_simulators.md b/docs/api/basic_simulators.md index 9172c388..e8ae1871 100755 --- a/docs/api/basic_simulators.md +++ b/docs/api/basic_simulators.md @@ -1 +1 @@ -:::ssms.basic_simulators \ No newline at end of file +:::ssms.basic_simulators diff --git a/docs/api/config.md b/docs/api/config.md index 080dd388..2e04f0b7 100755 --- a/docs/api/config.md +++ b/docs/api/config.md @@ -1 +1 @@ -:::ssms.config \ No newline at end of file +:::ssms.config diff --git a/docs/api/dataset_generators.md b/docs/api/dataset_generators.md index c1d0f267..6194983b 100755 --- a/docs/api/dataset_generators.md +++ b/docs/api/dataset_generators.md @@ -1 +1 @@ -:::ssms.dataset_generators \ No newline at end of file +:::ssms.dataset_generators diff --git a/docs/api/ssms.md b/docs/api/ssms.md index 10b493dc..66262212 100755 --- a/docs/api/ssms.md +++ b/docs/api/ssms.md @@ -1 +1 @@ -:::ssms \ No newline at end of file +:::ssms diff --git a/docs/api/support_utils.md b/docs/api/support_utils.md index 53d1ca9a..e4da0b1c 100755 --- a/docs/api/support_utils.md +++ b/docs/api/support_utils.md @@ -1 +1 @@ -:::ssms.support_utils \ No newline at end of file +:::ssms.support_utils diff --git a/docs/overrides/main.html b/docs/overrides/main.html index e202da1c..a6675116 100755 --- a/docs/overrides/main.html +++ b/docs/overrides/main.html @@ -15,4 +15,4 @@ {% include ".icons/material/head-question.svg" %} -{% endblock %} \ No newline at end of file +{% endblock %} diff --git a/src/cssm/__init__.py b/src/cssm/__init__.py index ffa66de4..b2f21817 100644 --- a/src/cssm/__init__.py +++ b/src/cssm/__init__.py @@ -23,7 +23,7 @@ ddm_flexbound_tradeoff, ) -from .race_models import race_model, lca +from .race_models import race_model, lca, racing_diffusion_model from .lba_models import lba_vanilla, lba_angle, rlwm_lba_pw_v1, rlwm_lba_race from .sequential_models import ( @@ -54,6 +54,7 @@ # Race models "race_model", "lca", + "racing_diffusion_model", # LBA models "lba_vanilla", "lba_angle", diff --git a/src/cssm/race_models.pyx b/src/cssm/race_models.pyx index 2974bcbb..33dcae8b 100644 --- a/src/cssm/race_models.pyx +++ b/src/cssm/race_models.pyx @@ -446,3 +446,210 @@ def lca(np.ndarray[float, ndim = 2] v, # drift parameters (np.array expect: one else: raise ValueError('return_option must be either "full" or "minimal"') # ----------------------------------------------------------------------------------------------- + +# Simulate (rt, choice) tuples from: Racing Diffusion Model ---------------------------------- +# @cythonboundscheck(False) +# @cythonwraparound(False) + +def racing_diffusion_model(np.ndarray[float, ndim = 2] v, # mean drift rates + np.ndarray[float, ndim = 2] b, # response boundaries (thresholds) + np.ndarray[float, ndim = 2] A, # between-trial variability in starting point (U[0, A]) + np.ndarray[float, ndim = 2] t, # non-decision times + np.ndarray[float, ndim = 2] s, # diffusion coefficients (within-trial noise) + np.ndarray[float, ndim = 1] deadline, + float delta_t = 0.001, # time increment step + float max_t = 20, # maximum rt allowed + int n_samples = 2000, + int n_trials = 1, + random_state = None, + return_option = 'full', + smooth_unif = False, + **kwargs): + """ + Simulate reaction times and choices from the Racing Diffusion Model (RDM) + based on the generative process described in Tillman et al. (2020). + + This model implements a "first-past-the-post" race of N independent + Wiener diffusion processes with no reflecting lower boundary. + + Parameters: + ----------- + v : np.ndarray + Mean drift rates. Shape (n_trials, n_particles). + b : np.ndarray + Response boundaries (thresholds). Shape (n_trials, n_particles). + A : np.ndarray + Upper bound of the uniform starting point distribution (U[0, A]). Shape (n_trials, n_particles). + t : np.ndarray + Non-decision times. Shape (n_trials, 1). + s : np.ndarray + Diffusion coefficients (within-trial noise). Shape (n_trials, n_particles). + deadline : np.ndarray + Maximum reaction time allowed for each trial. Shape (n_trials,). + delta_t : float, optional + Time increment step for simulation (default: 0.001). + max_t : float, optional + Maximum time for simulation (default: 20). + n_samples : int, optional + Number of samples to simulate per trial (default: 2000). + n_trials : int, optional + Number of trials to simulate (default: 1). + random_state : int or None, optional + Seed for random number generator (default: None). + return_option : str, optional + 'full' for complete output, 'minimal' for basic output (default: 'full'). + smooth_unif : bool, optional + Whether to apply uniform smoothing to reaction times (default: False). + **kwargs : dict + Additional keyword arguments. + + Returns: + -------- + dict + A dictionary containing simulated reaction times, choices, and metadata. + The exact contents depend on the 'return_option' parameter. + """ + + set_seed(random_state) + # Param views + cdef float[:, :] v_view = v + cdef float[:, :] b_view = b + cdef float[:, :] A_view = A + cdef float[:, :] t_view = t + cdef float[:, :] s_view = s + cdef float[:] deadline_view = deadline + + cdef float delta_t_sqrt = sqrt(delta_t) + sqrt_st = delta_t_sqrt * s + cdef float[:, :] sqrt_st_view = sqrt_st + + cdef int n_particles = v.shape[1] + rts = np.zeros((n_samples, n_trials, 1), dtype = DTYPE) + cdef float[:, :, :] rts_view = rts + choices = np.zeros((n_samples, n_trials, 1), dtype = np.intc) + cdef int[:, :, :] choices_view = choices + + particles = np.zeros((n_particles), dtype = DTYPE) + cdef float [:] particles_view = particles + + # Trajectory saving (for first trial, first sample) + traj = np.zeros((int(max_t / delta_t) + 1, n_particles), dtype = DTYPE) + traj[:, :] = -999 + cdef float[:, :] traj_view = traj + + # Initialize variables needed for for loop + cdef float t_particle, smooth_u, deadline_tmp + cdef Py_ssize_t n, ix, j, k + cdef Py_ssize_t m = 0 + cdef int winner = -1 + cdef int winner_found = 0 # <-- FIX: Use 'int' (0=False, 1=True) instead of 'bool' + + cdef int num_steps = int((max_t / delta_t) + 1) + cdef int num_draws = num_steps * n_particles + cdef float[:] gaussian_values = draw_gaussian(num_draws) + + for k in range(n_trials): + + deadline_tmp = min(max_t, deadline_view[k] - t_view[k, 0]) + + # Loop over samples + for n in range(n_samples): + + for j in range(n_particles): + particles_view[j] = random_uniform() * A_view[k, 0] + + t_particle = 0.0 # reset time + ix = 0 + winner = -1 # Reset winner for this sample + winner_found = 0 # <-- FIX: Reset to 0 (False) + + # Save initial trajectory + if n == 0: + if k == 0: + for j in range(n_particles): + traj_view[0, j] = particles[j] + + # Random walker + while not winner_found and t_particle <= deadline_tmp: # <-- 'not 0' is True + for j in range(n_particles): + # Standard Wiener diffusion process update + particles_view[j] += (v_view[k, j] * delta_t) + sqrt_st_view[k, 0] * gaussian_values[m] + + # No reflecting boundary for RDM + # The line `particles_view[j] = fmax(0.0, particles_view[j])` is REMOVED. + + m += 1 + if m == num_draws: # Resample random numbers if needed + m = 0 + gaussian_values = draw_gaussian(num_draws) + + # Check for a winner (first-past-the-post) + if particles_view[j] >= b_view[k, 0]: + winner_found = 1 # <-- FIX: Set to 1 (True) + winner = j + break # Stop checking, we have a winner + + if winner_found: # <-- `if 1` is True + break # Stop the while loop, a decision is made + + t_particle += delta_t + ix += 1 + + # Save running trajectory + if n == 0: + if k == 0: + for j in range(n_particles): + traj_view[ix, j] = particles[j] + + # --- End of while loop --- + + # Apply smoothing if specified + if smooth_unif: + if t_particle == 0.0: + smooth_u = random_uniform() * 0.5 * delta_t + elif t_particle < deadline_tmp: # Only smooth if not a deadline response + smooth_u = (0.5 - random_uniform()) * delta_t + else: + smooth_u = 0.0 + else: + smooth_u = 0.0 + + # Store RT and choice + rts_view[n , k, 0] = t_particle + t[k, 0] + smooth_u + choices_view[n, k, 0] = winner + + # Handle non-responses (deadline hit or no decision) + if (rts_view[n, k, 0] >= deadline_view[k]) | (not winner_found): # <-- `not 0` is True + rts_view[n, k, 0] = -999 + choices_view[n, k, 0] = -1 # Ensure choice is also -1 + + + # Create parameter dictionaries for metadata + v_dict = {} + for i in range(n_particles): + v_dict['v' + str(i)] = v[:, i] + + if return_option == 'full': + return {'rts': rts, 'choices': choices, 'metadata': {**v_dict, + 'b': b, + 'A': A, + 't': t, + 'deadline': deadline, + 's': s, + 'delta_t': delta_t, + 'max_t': max_t, + 'n_samples': n_samples, + 'n_trials': n_trials, + 'simulator': 'rdm_simulator', + 'possible_choices': list(np.arange(0, n_particles, 1)), + 'trajectory': traj}} + elif return_option == 'minimal': + return {'rts': rts, 'choices': choices, 'metadata': {'simulator': 'rdm_simulator', + 'possible_choices': list(np.arange(0, n_particles, 1)), + 'n_samples': n_samples, + 'n_trials': n_trials, + }} + + else: + raise ValueError('return_option must be either "full" or "minimal"') +# ----------------------------------------------------------------------------------------------- diff --git a/ssms/config/_modelconfig/__init__.py b/ssms/config/_modelconfig/__init__.py index ca23403c..78c2db75 100644 --- a/ssms/config/_modelconfig/__init__.py +++ b/ssms/config/_modelconfig/__init__.py @@ -122,6 +122,9 @@ get_race_no_z_angle_3_config, get_race_no_z_angle_4_config, ) +from .racing_diffusion import ( + get_racing_diffusion_3_config, +) from .shrink import ( get_shrink_spot_config, get_shrink_spot_extended_config, @@ -260,6 +263,7 @@ def get_model_config(): "race_no_z_4": get_race_no_z_4_config(), "race_no_bias_angle_4": get_race_no_bias_angle_4_config(), "race_no_z_angle_4": get_race_no_z_angle_4_config(), + "racing_diffusion_3": get_racing_diffusion_3_config(), "dev_rlwm_lba_pw_v1": get_dev_rlwm_lba_pw_v1_config(), "dev_rlwm_lba_race_v1": get_dev_rlwm_lba_race_v1_config(), "dev_rlwm_lba_race_v2": get_dev_rlwm_lba_race_v2_config(), diff --git a/ssms/config/_modelconfig/racing_diffusion.py b/ssms/config/_modelconfig/racing_diffusion.py new file mode 100644 index 00000000..9ddbea71 --- /dev/null +++ b/ssms/config/_modelconfig/racing_diffusion.py @@ -0,0 +1,35 @@ +"""Racing Diffusion model configurations.""" + +import cssm +from ssms.basic_simulators import boundary_functions as bf +from ssms.transforms import ( + ColumnStackParameters, + ExpandDimension, +) + + +def get_racing_diffusion_3_config(): + """Get configuration for racing diffusion model with 3 choices.""" + return { + "name": "racing_diffusion_3", + "params": ["v0", "v1", "v2", "A", "b", "t"], + "param_bounds": [ + [0.0, 0.0, 0.0, 0.0, 0.5, 0.0], + [2.5, 2.5, 2.5, 1.0, 3.0, 2.0], + ], + "boundary_name": "constant", + "boundary": bf.constant, + "n_params": 6, + "default_params": [1.0, 1.0, 1.0, 0.5, 1.5, 1e-3], + "nchoices": 3, + "choices": [0, 1, 2], + "n_particles": 3, + "simulator": cssm.racing_diffusion_model, + "parameter_transforms": { + "sampling": [], + "simulation": [ + ColumnStackParameters(["v0", "v1", "v2"], "v", delete_sources=False), + ExpandDimension(["t", "A", "b"]), + ], + }, + }