Adding racing diffusion model

.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
+          fail_ci_if_error: false

MANIFEST.in

@@ -27,4 +27,4 @@ global-exclude .git
 global-exclude *~
 global-exclude *.pyc
 global-exclude .#*
-global-exclude __pycache__
+global-exclude __pycache__

docs/api/basic_simulators.md

@@ -1 +1 @@
-:::ssms.basic_simulators
+:::ssms.basic_simulators

docs/api/config.md

@@ -1 +1 @@
-:::ssms.config
+:::ssms.config

docs/api/dataset_generators.md

@@ -1 +1 @@
-:::ssms.dataset_generators
+:::ssms.dataset_generators

docs/api/ssms.md

@@ -1 +1 @@
-:::ssms
+:::ssms

docs/api/support_utils.md

@@ -1 +1 @@
-:::ssms.support_utils
+:::ssms.support_utils

docs/overrides/main.html

@@ -15,4 +15,4 @@
         {% include ".icons/material/head-question.svg" %}
     </span>
 </span>
-{% endblock %}
+{% endblock %}

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",

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"')
+# -----------------------------------------------------------------------------------------------

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(),

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"]),
+            ],
+        },
+    }