From 5087509a0ec445ac256ea9e8ff9f5fe8a70e69f0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9Cmariama=E2=80=9D?= <“maria_ma@brown.edu”> Date: Mon, 24 Aug 2026 14:40:52 +0800 Subject: [PATCH 1/3] Add forward race simulation utilities --- notebooks/forward_race_simulator.ipynb | 174 +++++++++ .../race_npd_numerical_integration.ipynb | 152 ++++++++ setup.py | 1 + src/cssm/__init__.py | 2 + src/cssm/race_multistage_models.pyx | 363 ++++++++++++++++++ ssms/basic_simulators/race_math.py | 105 +++++ 6 files changed, 797 insertions(+) create mode 100644 notebooks/forward_race_simulator.ipynb create mode 100644 notebooks/race_npd_numerical_integration.ipynb create mode 100644 src/cssm/race_multistage_models.pyx create mode 100644 ssms/basic_simulators/race_math.py diff --git a/notebooks/forward_race_simulator.ipynb b/notebooks/forward_race_simulator.ipynb new file mode 100644 index 00000000..18aef962 --- /dev/null +++ b/notebooks/forward_race_simulator.ipynb @@ -0,0 +1,174 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "title", + "metadata": {}, + "source": [ + "### Example: one accumulator in a race model\n", + "#### 0.1 Simulate and visualize the model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "imports", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "\n", + "from cssm import race_multistage\n", + "from ssms.basic_simulators.race_math import big_F, q, small_f\n", + "\n", + "RANDOM_STATE = 20260820\n", + "OMISSION = -999.0" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "parameters-and-trajectories", + "metadata": {}, + "outputs": [], + "source": [ + "mu = 0.75\n", + "sigma = 1.0\n", + "x0 = 0.0\n", + "a = 1.0\n", + "b = 0.0\n", + "T = 1.5\n", + "dt = 1e-3\n", + "num_trajs = 2_000\n", + "\n", + "def simulate_trajs(mu, sigma, x0, T, dt, num, rng):\n", + " \"\"\"Euler-Maruyama paths for the trajectory plot.\"\"\"\n", + " t_grid = np.arange(0.0, T + dt, dt)\n", + " X = np.empty((num, t_grid.size))\n", + " X[:, 0] = x0\n", + " for step in range(1, t_grid.size):\n", + " X[:, step] = X[:, step - 1] + mu * dt + sigma * np.sqrt(dt) * rng.normal(size=num)\n", + " return t_grid, X\n", + "\n", + "t_grid, X_grids = simulate_trajs(mu, sigma, x0, T, dt, num_trajs, np.random.default_rng(RANDOM_STATE))\n", + "expected_mean = x0 + mu * t_grid\n", + "empirical_mean = X_grids.mean(axis=0)\n", + "expected_std = sigma * np.sqrt(t_grid)\n", + "empirical_std = X_grids.std(axis=0)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "trajectory-plot", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(figsize=(8, 4))\n", + "ax.plot(t_grid, empirical_mean, label='empirical mean', color='tab:blue')\n", + "ax.plot(t_grid, expected_mean, label='expected mean', linewidth=2.5, linestyle='--', color='tab:blue')\n", + "ax.plot(t_grid, empirical_mean + 2 * empirical_std, label=r'empirical mean $\\pm$ 2 std', color='tab:green')\n", + "ax.plot(t_grid, empirical_mean - 2 * empirical_std, color='tab:green')\n", + "ax.plot(t_grid, expected_mean + 2 * expected_std, label=r'expected mean $\\pm$ 2 std', linewidth=2.5, linestyle='--', color='tab:green')\n", + "ax.plot(t_grid, expected_mean - 2 * expected_std, linewidth=2.5, linestyle='--', color='tab:green')\n", + "ax.plot(t_grid, X_grids[:10, :].T, alpha=0.5)\n", + "ax.plot(t_grid, a + b * t_grid, color='black', linestyle='--', label='upper boundary')\n", + "ax.autoscale(axis='x', tight=True)\n", + "ax.set(xlabel='time', ylabel='evidence')\n", + "ax.legend(fontsize=9)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "fptd-npd-heading", + "metadata": {}, + "source": [ + "#### 0.2 Simulate first-passage times, compute FPTD, CDF, and NPD" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "one-accumulator-cssm", + "metadata": {}, + "outputs": [], + "source": [ + "race_arrays = {\n", + " 'mu_array': np.array([[[mu]]]),\n", + " 'sigma_array': np.array([[[sigma]]]),\n", + " 'node_array': np.zeros((1, 1, 1)),\n", + " 'd_array': np.ones((1, 1), dtype=np.int32),\n", + " 'upper_intercept_array': np.array([[[a]]]),\n", + " 'upper_slope_array': np.array([[[b]]]),\n", + " 'x0_array': np.array([[x0]]),\n", + "}\n", + "\n", + "num_fpt = 50_000\n", + "out = race_multistage(\n", + " **race_arrays, n_samples=num_fpt, delta_t=dt, max_t=T, random_state=RANDOM_STATE,\n", + ")\n", + "rt = out['rts'].reshape(-1)\n", + "x_final = out['metadata']['x_final'].reshape(-1)\n", + "exited = rt != OMISSION\n", + "fp_times = rt[exited]\n", + "np_positions = x_final[~exited]\n", + "\n", + "print(f'empirical F(T): {exited.mean():.4f}')\n", + "print(f'analytic F(T): {float(big_F(T, mu, sigma, a, b, T, x0)):.4f}')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fptd-and-npd-plot", + "metadata": {}, + "outputs": [], + "source": [ + "ts = np.linspace(1e-3, T, 1_000)\n", + "xs = np.linspace(-4.0, a + b * T, 1_000)\n", + "\n", + "fig, ax = plt.subplots(1, 2, figsize=(12, 4))\n", + "fpt_bins = np.linspace(0.0, T, 101)\n", + "fpt_width = fpt_bins[1] - fpt_bins[0]\n", + "fpt_counts, _ = np.histogram(fp_times, bins=fpt_bins)\n", + "ax[0].stairs(fpt_counts / (num_fpt * fpt_width), fpt_bins, color='black', label='Monte Carlo')\n", + "ax[0].plot(ts, small_f(ts, mu, sigma, a, b, T, x0), color='tab:blue', label=r'$f(t)$')\n", + "ax[0].set(xlabel=r'$t$', ylabel='FPTD')\n", + "ax[0].legend()\n", + "\n", + "np_bins = np.linspace(-4.0, a + b * T, 101)\n", + "np_width = np_bins[1] - np_bins[0]\n", + "np_counts, _ = np.histogram(np_positions, bins=np_bins)\n", + "ax[1].stairs(np_counts / (num_fpt * np_width), np_bins, color='gray', label='Monte Carlo')\n", + "ax[1].plot(xs, q(xs, mu, sigma, a, b, T, x0), color='tab:red', label=r'$q(x, T)$')\n", + "ax[1].set(xlabel=r'$x$', ylabel='NPD')\n", + "ax[1].legend()\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cdf-plot", + "metadata": {}, + "outputs": [], + "source": [ + "empirical_cdf = np.array([(fp_times <= t).sum() / num_fpt for t in ts])\n", + "fig, ax = plt.subplots(figsize=(8, 4))\n", + "ax.plot(ts, empirical_cdf, color='black', label='Monte Carlo')\n", + "ax.plot(ts, big_F(ts, mu, sigma, a, b, T, x0), color='tab:blue', label=r'$F(t)$')\n", + "ax.set(xlabel=r'$t$', ylabel='CDF', ylim=(0, 1))\n", + "ax.legend()\n", + "plt.show()" + ] + } + ], + "metadata": { + "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, + "language_info": {"name": "python", "version": "3.12"} + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/race_npd_numerical_integration.ipynb b/notebooks/race_npd_numerical_integration.ipynb new file mode 100644 index 00000000..2f6a4a12 --- /dev/null +++ b/notebooks/race_npd_numerical_integration.ipynb @@ -0,0 +1,152 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Numerical integration of the non-passage density" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "imports-and-parameters", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "\n", + "from ssms.basic_simulators.race_math import big_F, q, small_f\n", + "\n", + "mu = 0.75\n", + "sigma = 1.0\n", + "x0 = 0.0\n", + "a = 1.0\n", + "b = 0.0\n", + "T = 1.5\n", + "\n", + "# Finite replacements for the mathematical lower bound -infinity.\n", + "# For this example, -100 is the conservative default; the plots below check smaller cutoffs.\n", + "lower_bounds = np.array([-4.0, -10.0, -25.0, -50.0, -100.0])\n", + "n_trapezoid_points = 20_000\n", + "n_gauss_nodes = 128" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "integration-methods", + "metadata": {}, + "outputs": [], + "source": [ + "gauss_nodes, gauss_weights = np.polynomial.legendre.leggauss(n_gauss_nodes)\n", + "\n", + "def integrate_trapezoid(t, lower_x):\n", + " x_grid = np.linspace(lower_x, a + b * t, n_trapezoid_points)\n", + " return np.trapezoid(q(x_grid, mu, sigma, a, b, t, x0), x_grid)\n", + "\n", + "def integrate_gauss_legendre(t, lower_x):\n", + " upper_x = a + b * t\n", + " x_grid = 0.5 * (upper_x - lower_x) * gauss_nodes + 0.5 * (upper_x + lower_x)\n", + " weights = 0.5 * (upper_x - lower_x) * gauss_weights\n", + " return np.sum(weights * q(x_grid, mu, sigma, a, b, t, x0))\n", + "\n", + "ts = np.linspace(1e-3, T, 300)\n", + "survival = 1.0 - big_F(ts, mu, sigma, a, b, T, x0)\n", + "trapezoid_mass = np.array([[integrate_trapezoid(t, lower_x) for t in ts] for lower_x in lower_bounds])\n", + "gauss_mass = np.array([[integrate_gauss_legendre(t, lower_x) for t in ts] for lower_x in lower_bounds])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "integration-method-plots", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(1, 2, figsize=(13, 4), sharey=True)\n", + "for i, lower_x in enumerate(lower_bounds):\n", + " ax[0].plot(ts, trapezoid_mass[i], label=fr'$L={lower_x:g}$')\n", + " ax[1].plot(ts, gauss_mass[i], label=fr'$L={lower_x:g}$')\n", + "for axis, title in zip(ax, ['trapezoid', 'Gauss-Legendre']):\n", + " axis.plot(ts, survival, color='black', linestyle='--', linewidth=2, label=r'$1-F(t)$')\n", + " axis.set(title=title, xlabel=r'$t$', ylim=(0, 1))\n", + " axis.legend(fontsize=8)\n", + "ax[0].set_ylabel('non-passage probability')\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fptd-to-cdf-methods", + "metadata": {}, + "outputs": [], + "source": [ + "def integrate_f_trapezoid(t):\n", + " t_grid = np.linspace(1e-8, t, n_trapezoid_points)\n", + " return np.trapezoid(small_f(t_grid, mu, sigma, a, b, T, x0), t_grid)\n", + "\n", + "def integrate_f_gauss_legendre(t):\n", + " t_grid = 0.5 * t * (gauss_nodes + 1.0)\n", + " weights = 0.5 * t * gauss_weights\n", + " return np.sum(weights * small_f(t_grid, mu, sigma, a, b, T, x0))\n", + "\n", + "f_cdf_trapezoid = np.array([integrate_f_trapezoid(t) for t in ts])\n", + "f_cdf_gauss = np.array([integrate_f_gauss_legendre(t) for t in ts])\n", + "analytic_cdf = big_F(ts, mu, sigma, a, b, T, x0)\n", + "\n", + "print(f'F(T) from trapezoid: {f_cdf_trapezoid[-1]:.6f}')\n", + "print(f'F(T) from Gauss-Legendre: {f_cdf_gauss[-1]:.6f}')\n", + "print(f'analytic F(T): {analytic_cdf[-1]:.6f}')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fptd-cdf-plot", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(figsize=(8, 4))\n", + "ax.plot(ts, f_cdf_trapezoid, color='tab:orange', label=r'$\\int_0^t f(s)\\,ds$: trapezoid')\n", + "ax.plot(ts, f_cdf_gauss, color='tab:green', label=r'$\\int_0^t f(s)\\,ds$: Gauss-Legendre')\n", + "ax.plot(ts, analytic_cdf, color='black', linestyle='--', label=r'$F(t)$')\n", + "ax.set(xlabel=r'$t$', ylabel='CDF', ylim=(0, 1))\n", + "ax.legend(fontsize=8)\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "final-time-errors", + "metadata": {}, + "outputs": [], + "source": [ + "target_T = survival[-1]\n", + "trapezoid_error = np.abs(trapezoid_mass[:, -1] - target_T)\n", + "gauss_error = np.abs(gauss_mass[:, -1] - target_T)\n", + "\n", + "print('lower bound trapezoid Gauss-Legendre')\n", + "for lower_x, trap, gauss in zip(lower_bounds, trapezoid_mass[:, -1], gauss_mass[:, -1]):\n", + " print(f'{lower_x:10.0f} {trap:.8f} {gauss:.8f}')\n", + "print(f'1 - F(T) = {target_T:.8f}')\n", + "\n", + "fig, ax = plt.subplots(figsize=(7, 4))\n", + "ax.semilogy(-lower_bounds, trapezoid_error, 'o-', label='trapezoid')\n", + "ax.semilogy(-lower_bounds, gauss_error, 's-', label='Gauss-Legendre')\n", + "ax.set(xlabel=r'magnitude of lower cutoff $-L$', ylabel=r'absolute error at $T$')\n", + "ax.legend()\n", + "plt.show()" + ] + } + ], + "metadata": { + "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, + "language_info": {"name": "python", "version": "3.12"} + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/setup.py b/setup.py index 22d9af0c..0cf959bc 100644 --- a/setup.py +++ b/setup.py @@ -38,6 +38,7 @@ OPENMP_MODULES = [ "_openmp_status", # Runtime OpenMP/GSL detection "addm_models", # aDDM simulator (ported efpt engine, prange + inline xoshiro) + "race_multistage_models", # covariate-conditioned multi-stage race model "ddm_models", # DDM simulators with n_threads support "levy_models", # Levy simulators with n_threads support "ornstein_models", # Ornstein-Uhlenbeck with n_threads support diff --git a/src/cssm/__init__.py b/src/cssm/__init__.py index 0b79eb90..044735cb 100644 --- a/src/cssm/__init__.py +++ b/src/cssm/__init__.py @@ -26,6 +26,7 @@ ) from .addm_models import addm +from .race_multistage_models import race_multistage from .race_models import race_model, lca, racing_diffusion_model from .poisson_race_models import poisson_race @@ -58,6 +59,7 @@ "ddm_flexbound_tradeoff", # ADDM model "addm", + "race_multistage", # Race models "race_model", "lca", diff --git a/src/cssm/race_multistage_models.pyx b/src/cssm/race_multistage_models.pyx new file mode 100644 index 00000000..92edeca7 --- /dev/null +++ b/src/cssm/race_multistage_models.pyx @@ -0,0 +1,363 @@ +# cython: language_level=3 +# cython: boundscheck=False +# cython: wraparound=False +# cython: cdivision=True +# cython: initializedcheck=False + +"""Forward simulator for independent, multi-stage race models. + +This module implements the Monte Carlo model in ``race_model.pdf``: every +accumulator has an independent Brownian motion, a native stage partition, +piecewise-constant drift/diffusion, and a piecewise-linear *upper* boundary. +The race ends at the first upper-boundary crossing. It intentionally has no +lower decision boundary and does not clip paths at zero. + +The random-number implementation follows ``addm_models.pyx``. One xoshiro +state is seeded for every (sample, trial) pair before the OpenMP loop, making +results reproducible for a fixed ``random_state`` independently of scheduling +or ``n_threads``. +""" + +import numpy as np +cimport numpy as np +from libc.math cimport sqrt, log, cos, sin, M_PI +from libc.stdint cimport uint64_t +from cython.parallel cimport prange + +from cssm._utils import ( + setup_simulation, + build_minimal_metadata, + build_full_metadata, + build_return_dict, +) + +cdef double OMISSION = -999.0 +DEF MAX_ACCUMULATORS = 32 + + +# Inline xoshiro256++ / Box-Muller RNG, matching the Efficient-FPT-derived +# aDDM engine. Keep this local for now; a future shared RNG module can remove +# the duplication once the race simulator interface has settled. +cdef struct Xoshiro256State: + uint64_t s0 + uint64_t s1 + uint64_t s2 + uint64_t s3 + + +cdef inline uint64_t _rotl(uint64_t x, int k) noexcept nogil: + return (x << k) | (x >> (64 - k)) + + +cdef inline uint64_t _next(Xoshiro256State *state) noexcept nogil: + cdef uint64_t result = _rotl(state.s0 + state.s3, 23) + state.s0 + cdef uint64_t t = state.s1 << 17 + state.s2 ^= state.s0 + state.s3 ^= state.s1 + state.s1 ^= state.s2 + state.s0 ^= state.s3 + state.s2 ^= t + state.s3 = _rotl(state.s3, 45) + return result + + +cdef inline uint64_t _splitmix64(uint64_t *state) noexcept nogil: + state[0] += 0x9e3779b97f4a7c15 + cdef uint64_t z = state[0] + z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9 + z = (z ^ (z >> 27)) * 0x94d049bb133111eb + return z ^ (z >> 31) + + +cdef inline void _seed(Xoshiro256State *state, uint64_t seed) noexcept nogil: + cdef uint64_t sm_state = seed + state.s0 = _splitmix64(&sm_state) + state.s1 = _splitmix64(&sm_state) + state.s2 = _splitmix64(&sm_state) + state.s3 = _splitmix64(&sm_state) + + +cdef struct BoxMullerState: + double spare + int has_spare + + +cdef inline double _normal(Xoshiro256State *rng, BoxMullerState *bm) noexcept nogil: + cdef double u1, u2, magnitude + if bm.has_spare: + bm.has_spare = 0 + return bm.spare + u1 = (_next(rng) >> 11) * (1.0 / 9007199254740992.0) + u2 = (_next(rng) >> 11) * (1.0 / 9007199254740992.0) + if u1 < 1e-300: + u1 = 1e-300 + magnitude = sqrt(-2.0 * log(u1)) + bm.spare = magnitude * sin(2.0 * M_PI * u2) + bm.has_spare = 1 + return magnitude * cos(2.0 * M_PI * u2) + + +cdef void _run_race_trial( + double[:, :, ::1] mu, + double[:, :, ::1] sigma, + double[:, :, ::1] nodes, + int[:, ::1] d, + double[:, :, ::1] upper_intercept, + double[:, :, ::1] upper_slope, + int row, + double[:, ::1] x0, + double dt, + int max_steps, + double horizon, + uint64_t seed, + double *rt_out, + int *choice_out, + double *x_final_out, + int n_accumulators, +) noexcept nogil: + """Simulate one race. A same-grid-step tie uses the lowest index. + + Continuous-time ties have probability zero. The deterministic grid-tie rule + therefore only resolves a discretisation artefact and keeps seeded runs + deterministic. + """ + cdef: + Xoshiro256State rng + BoxMullerState bm + double particle[MAX_ACCUMULATORS] + double t_particle, dt_current, sqrt_dt, boundary + int stage[MAX_ACCUMULATORS] + int i, step, winner + + _seed(&rng, seed) + bm.has_spare = 0 + t_particle = 0.0 + rt_out[0] = -1.0 + choice_out[0] = 0 + for i in range(n_accumulators): + particle[i] = x0[row, i] + stage[i] = 0 + + for step in range(max_steps): + dt_current = horizon - t_particle + if dt_current <= 0.0: + break + if dt_current > dt: + dt_current = dt + sqrt_dt = sqrt(dt_current) + + for i in range(n_accumulators): + particle[i] += ( + mu[row, i, stage[i]] * dt_current + + sigma[row, i, stage[i]] * sqrt_dt * _normal(&rng, &bm) + ) + t_particle += dt_current + + winner = -1 + for i in range(n_accumulators): + boundary = ( + upper_intercept[row, i, stage[i]] + + upper_slope[row, i, stage[i]] + * (t_particle - nodes[row, i, stage[i]]) + ) + if particle[i] >= boundary and winner < 0: + winner = i + + if winner >= 0: + # The aDDM Efficient-FPT-compatible simulator reports the midpoint + # of the Euler step; use the same first-order convention here. + rt_out[0] = t_particle - 0.5 * dt_current + choice_out[0] = winner + break + + for i in range(n_accumulators): + while stage[i] + 1 < d[row, i] and t_particle >= nodes[row, i, stage[i] + 1]: + stage[i] += 1 + + for i in range(n_accumulators): + x_final_out[i] = particle[i] + + +def _simulate_race_multistage( + double[:, :, ::1] mu, + double[:, :, ::1] sigma, + double[:, :, ::1] nodes, + int[:, ::1] d, + double[:, :, ::1] upper_intercept, + double[:, :, ::1] upper_slope, + double[:, ::1] x0, + double dt, + double horizon, + uint64_t[::1] seeds, + int n_threads=1, +): + """Low-level batch kernel used by :func:`race_multistage` and tests.""" + cdef: + int n_rows = mu.shape[0] + int n_accumulators = mu.shape[1] + int max_steps = int(np.ceil(horizon / dt)) if horizon > 0.0 else 0 + int row + + if n_accumulators > MAX_ACCUMULATORS: + raise ValueError( + f"race_multistage supports at most {MAX_ACCUMULATORS} accumulators; " + f"got {n_accumulators}" + ) + if dt <= 0.0: + raise ValueError("dt must be positive") + if ( + sigma.shape[0] != n_rows or sigma.shape[1] != n_accumulators or sigma.shape[2] != mu.shape[2] + or nodes.shape[0] != n_rows or nodes.shape[1] != n_accumulators or nodes.shape[2] != mu.shape[2] + or upper_intercept.shape[0] != n_rows or upper_intercept.shape[1] != n_accumulators or upper_intercept.shape[2] != mu.shape[2] + or upper_slope.shape[0] != n_rows or upper_slope.shape[1] != n_accumulators or upper_slope.shape[2] != mu.shape[2] + ): + raise ValueError("stage arrays must have the same (rows, accumulators, stages) shape") + if d.shape[0] != n_rows or d.shape[1] != n_accumulators: + raise ValueError("d must have shape (rows, accumulators)") + if x0.shape[0] != n_rows or x0.shape[1] != n_accumulators: + raise ValueError("x0 must have shape (rows, accumulators)") + if seeds.shape[0] != n_rows: + raise ValueError("seeds must contain one seed per row") + if np.any(np.asarray(d) < 1) or np.any(np.asarray(d) > mu.shape[2]): + raise ValueError("each d entry must lie between 1 and the padded stage count") + + rt = np.empty(n_rows, dtype=np.float64) + choice = np.empty(n_rows, dtype=np.int32) + x_final = np.empty((n_rows, n_accumulators), dtype=np.float64) + cdef double[::1] rt_view = rt + cdef int[::1] choice_view = choice + cdef double[:, ::1] final_view = x_final + + # Match the Efficient-FPT-derived aDDM batch engine: there is no path to + # evolve at a zero/negative horizon, so avoid starting the parallel loop. + if horizon <= 0.0: + rt.fill(-1.0) + choice.fill(0) + x_final[:] = np.asarray(x0) + return rt, choice, x_final + + for row in prange(n_rows, nogil=True, num_threads=n_threads, schedule='dynamic'): + _run_race_trial( + mu, sigma, nodes, d, upper_intercept, upper_slope, row, x0, + dt, max_steps, horizon, seeds[row], &rt_view[row], + &choice_view[row], &final_view[row, 0], n_accumulators, + ) + return rt, choice, x_final + + +def race_multistage( + mu_array, + sigma_array, + node_array, + d_array, + upper_intercept_array, + upper_slope_array, + x0_array, + nondecision_time=None, + deadline=None, + float delta_t=0.001, + float max_t=20.0, + int n_samples=1000, + int n_trials=0, + return_option='full', + random_state=None, + int n_threads=1, + **kwargs, +): + """Simulate independent multi-stage race-model trajectories. + + Array inputs describe one row per *experimental trial*, with shape + ``(n_trials, K, max_stages)`` except ``d_array`` and ``x0_array``, whose + shapes are ``(n_trials, K)``. ``node_array`` contains stage start times; + stage parameters are used from a start time until the next node. The + boundary in stage ``k`` is ``upper_intercept + upper_slope * elapsed``. + + ``n_samples`` repeats each experimental trial. Outputs follow the SSMS + contract: arrays are ``(n_samples, n_trials, 1)``, choices are zero based, + and an unobserved response has RT ``-999.0``. + """ + if n_samples < 1: + raise ValueError("n_samples must be positive") + + mu = np.ascontiguousarray(mu_array, dtype=np.float64) + sigma = np.ascontiguousarray(sigma_array, dtype=np.float64) + nodes = np.ascontiguousarray(node_array, dtype=np.float64) + d = np.ascontiguousarray(d_array, dtype=np.int32) + intercept = np.ascontiguousarray(upper_intercept_array, dtype=np.float64) + slope = np.ascontiguousarray(upper_slope_array, dtype=np.float64) + x0 = np.ascontiguousarray(x0_array, dtype=np.float64) + if mu.ndim != 3: + raise ValueError("mu_array must have shape (n_trials, K, max_stages)") + if n_trials == 0: + n_trials = mu.shape[0] + if n_trials < 1: + raise ValueError("n_trials must be positive") + if mu.shape[0] != n_trials: + raise ValueError("n_trials must equal the first dimension of mu_array") + + setup = setup_simulation(n_samples, n_trials, max_t, delta_t, random_state) + N = n_samples * n_trials + seed = random_state if random_state is not None else np.random.randint(0, 2**31) + rng = np.random.default_rng(seed) + seeds = np.ascontiguousarray(rng.integers(0, 2**64, size=N, dtype=np.uint64)) + + # sample-major tiling matches SSMS' (sample, trial) output layout. + mu_rows = np.ascontiguousarray(np.tile(mu, (n_samples, 1, 1))) + sigma_rows = np.ascontiguousarray(np.tile(sigma, (n_samples, 1, 1))) + nodes_rows = np.ascontiguousarray(np.tile(nodes, (n_samples, 1, 1))) + d_rows = np.ascontiguousarray(np.tile(d, (n_samples, 1))) + intercept_rows = np.ascontiguousarray(np.tile(intercept, (n_samples, 1, 1))) + slope_rows = np.ascontiguousarray(np.tile(slope, (n_samples, 1, 1))) + x0_rows = np.ascontiguousarray(np.tile(x0, (n_samples, 1))) + + rt, choice, x_final = _simulate_race_multistage( + mu_rows, sigma_rows, nodes_rows, d_rows, intercept_rows, slope_rows, + x0_rows, delta_t, max_t, seeds, n_threads, + ) + + ndt = np.zeros(n_trials, dtype=np.float64) if nondecision_time is None else np.asarray(nondecision_time, dtype=np.float64).reshape(-1) + ddl = np.full(n_trials, max_t, dtype=np.float64) if deadline is None else np.asarray(deadline, dtype=np.float64).reshape(-1) + if ndt.size == 1: + ndt = np.full(n_trials, ndt[0]) + if ddl.size == 1: + ddl = np.full(n_trials, ddl[0]) + if ndt.size != n_trials or ddl.size != n_trials: + raise ValueError("nondecision_time and deadline must be scalars or length n_trials") + ndt_rows = np.tile(ndt, n_samples) + ddl_rows = np.tile(ddl, n_samples) + shifted_rt = rt + ndt_rows + omitted = (rt < 0.0) | (shifted_rt > ddl_rows) + final_rt = np.where(omitted, OMISSION, shifted_rt) + final_choice = np.where(omitted, 0, choice) + + rts = setup['rts'] + choices = setup['choices'] + rts[:] = final_rt.reshape(n_samples, n_trials, 1).astype(np.float32) + choices[:] = final_choice.reshape(n_samples, n_trials, 1).astype(np.int32) + + possible_choices = list(range(mu.shape[1])) + minimal_meta = build_minimal_metadata( + simulator_name='race_multistage', + possible_choices=possible_choices, + n_samples=n_samples, + n_trials=n_trials, + boundary_fun_name='piecewise_linear_upper', + ) + if return_option == 'minimal': + metadata = minimal_meta + elif return_option == 'full': + metadata = build_full_metadata( + minimal_metadata=minimal_meta, + params={ + 'mu_array': mu, 'sigma_array': sigma, 'node_array': nodes, + 'd_array': d, 'upper_intercept_array': intercept, + 'upper_slope_array': slope, 'x0_array': x0, + }, + sim_config={'delta_t': delta_t, 'max_t': max_t, 'n_threads': n_threads}, + traj=setup['traj'], + boundary=np.array([], dtype=np.float32), + ) + metadata['x_final'] = x_final.reshape(n_samples, n_trials, mu.shape[1]) + else: + raise ValueError("return_option must be either 'full' or 'minimal'") + return build_return_dict(rts, choices, metadata) diff --git a/ssms/basic_simulators/race_math.py b/ssms/basic_simulators/race_math.py new file mode 100644 index 00000000..24495615 --- /dev/null +++ b/ssms/basic_simulators/race_math.py @@ -0,0 +1,105 @@ +"""Analytical one-sided race-model quantities from equation (14). + +These functions describe one independent accumulator with a piecewise-linear +upper boundary. They are the mathematical reference for the CSSM forward +simulator and the later multi-stage race likelihood implementation. +""" + +from __future__ import annotations + +from math import erf, sqrt + +import numpy as np + + +_SQRT_2 = sqrt(2.0) +_SQRT_2PI = sqrt(2.0 * np.pi) + + +def _normal_cdf(x: np.ndarray | float) -> np.ndarray: + """Standard-normal CDF without requiring SciPy.""" + x = np.asarray(x, dtype=float) + return 0.5 * (1.0 + np.vectorize(erf, otypes=[float])(x / _SQRT_2)) + + +def small_f( + t: np.ndarray | float, + mu: float, + sigma: float, + a: float, + b: float, + T: float, + x0: float, +) -> np.ndarray: + """One-sided FPT density ``f_tau(t)`` in race-model equation (14).""" + if sigma <= 0.0 or T <= 0.0: + raise ValueError("sigma and T must be positive") + t = np.asarray(t, dtype=float) + out = np.zeros_like(t) + valid = (t > 0.0) & (t <= T) + distance = a - x0 + relative_drift = mu - b + t_valid = t[valid] + out[valid] = ( + distance + / (_SQRT_2PI * sigma * t_valid**1.5) + * np.exp( + -((distance - relative_drift * t_valid) ** 2) / (2.0 * sigma**2 * t_valid) + ) + ) + return out + + +def big_F( + t: np.ndarray | float, + mu: float, + sigma: float, + a: float, + b: float, + T: float, + x0: float, +) -> np.ndarray: + """One-sided FPT CDF ``F_tau(t)`` in race-model equation (14).""" + if sigma <= 0.0 or T <= 0.0: + raise ValueError("sigma and T must be positive") + t = np.asarray(t, dtype=float) + out = np.zeros_like(t) + valid = t > 0.0 + elapsed = np.minimum(t[valid], T) + distance = a - x0 + relative_drift = mu - b + root_elapsed = np.sqrt(elapsed) + out[valid] = _normal_cdf( + (relative_drift * elapsed - distance) / (sigma * root_elapsed) + ) + np.exp(2.0 * relative_drift * distance / sigma**2) * _normal_cdf( + (-distance - relative_drift * elapsed) / (sigma * root_elapsed) + ) + return out + + +def q( + x: np.ndarray | float, + mu: float, + sigma: float, + a: float, + b: float, + T: float, + x0: float, +) -> np.ndarray: + """Killed/non-passage density ``q(x; ..., T, x0)`` in equation (14).""" + if sigma <= 0.0 or T <= 0.0: + raise ValueError("sigma and T must be positive") + x = np.asarray(x, dtype=float) + boundary = a + b * T + out = np.zeros_like(x) + inside = x < boundary + x_inside = x[inside] + gaussian = np.exp(-((x_inside - x0 - mu * T) ** 2) / (2.0 * sigma**2 * T)) + killed_factor = 1.0 - np.exp( + 2.0 * (a - x0) * (x_inside - boundary) / (sigma**2 * T) + ) + out[inside] = gaussian * killed_factor / (_SQRT_2PI * sigma * sqrt(T)) + return out + + +__all__ = ["small_f", "big_F", "q"] From b88512856b36afe62888848c90bb0cfb62086218 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9Cmariama=E2=80=9D?= <“maria_ma@brown.edu”> Date: Mon, 31 Aug 2026 23:49:36 +0800 Subject: [PATCH 2/3] Address race simulation review feedback --- src/cssm/race_multistage_models.pyx | 211 +++++++++++++++++----- ssms/basic_simulators/race_math.py | 66 +++++-- tests/test_race_math_validation.py | 28 +++ tests/test_race_multistage_stage_nodes.py | 27 +++ 4 files changed, 268 insertions(+), 64 deletions(-) create mode 100644 tests/test_race_math_validation.py create mode 100644 tests/test_race_multistage_stage_nodes.py diff --git a/src/cssm/race_multistage_models.pyx b/src/cssm/race_multistage_models.pyx index 92edeca7..e2f39384 100644 --- a/src/cssm/race_multistage_models.pyx +++ b/src/cssm/race_multistage_models.pyx @@ -21,6 +21,7 @@ or ``n_threads``. import numpy as np cimport numpy as np from libc.math cimport sqrt, log, cos, sin, M_PI +from libc.stdlib cimport free, malloc from libc.stdint cimport uint64_t from cython.parallel cimport prange @@ -33,6 +34,21 @@ from cssm._utils import ( cdef double OMISSION = -999.0 DEF MAX_ACCUMULATORS = 32 +cdef int UNIFORM_BITS = 11 +cdef double UINT64_TO_DOUBLE = 1.0 / 9007199254740992.0 +cdef double MIN_UNIFORM = 1e-300 +cdef double TWO_PI = 2.0 * M_PI + + +cdef struct SimulationConfig: + double dt + double horizon + int max_steps + + +cdef struct TrialResult: + double rt + int choice # Inline xoshiro256++ / Box-Muller RNG, matching the Efficient-FPT-derived @@ -78,23 +94,64 @@ cdef inline void _seed(Xoshiro256State *state, uint64_t seed) noexcept nogil: cdef struct BoxMullerState: + # Box-Muller produces two normals from a pair of uniforms. ``spare`` + # stores the second value, and ``has_spare`` marks it for the next call. double spare int has_spare cdef inline double _normal(Xoshiro256State *rng, BoxMullerState *bm) noexcept nogil: + """Return one standard normal and cache its Box-Muller companion. + + A call first consumes ``bm.spare`` when available. Otherwise it generates + a pair of normals, returns the cosine component, and caches the sine + component for the following call. + """ cdef double u1, u2, magnitude + # Consume the cached companion before drawing new uniforms. if bm.has_spare: bm.has_spare = 0 return bm.spare - u1 = (_next(rng) >> 11) * (1.0 / 9007199254740992.0) - u2 = (_next(rng) >> 11) * (1.0 / 9007199254740992.0) - if u1 < 1e-300: - u1 = 1e-300 + u1 = (_next(rng) >> UNIFORM_BITS) * UINT64_TO_DOUBLE + u2 = (_next(rng) >> UNIFORM_BITS) * UINT64_TO_DOUBLE + if u1 < MIN_UNIFORM: + u1 = MIN_UNIFORM magnitude = sqrt(-2.0 * log(u1)) - bm.spare = magnitude * sin(2.0 * M_PI * u2) + bm.spare = magnitude * sin(TWO_PI * u2) bm.has_spare = 1 - return magnitude * cos(2.0 * M_PI * u2) + return magnitude * cos(TWO_PI * u2) + + +cdef inline double _upper_boundary_at( + double[:, :, ::1] upper_intercept, + double[:, :, ::1] upper_slope, + double[:, :, ::1] nodes, + int row, + int accumulator, + int stage, + double time, +) noexcept nogil: + """Return an accumulator's upper boundary at ``time`` in ``stage``.""" + return ( + upper_intercept[row, accumulator, stage] + + upper_slope[row, accumulator, stage] + * (time - nodes[row, accumulator, stage]) + ) + + +cdef inline bint _has_reached_next_stage( + double[:, :, ::1] nodes, + int[:, ::1] d, + int row, + int accumulator, + int current_stage, + double time, +) noexcept nogil: + """Return whether an accumulator should advance to its next stage.""" + return ( + current_stage + 1 < d[row, accumulator] + and time >= nodes[row, accumulator, current_stage + 1] + ) cdef void _run_race_trial( @@ -106,14 +163,10 @@ cdef void _run_race_trial( double[:, :, ::1] upper_slope, int row, double[:, ::1] x0, - double dt, - int max_steps, - double horizon, + SimulationConfig config, uint64_t seed, - double *rt_out, - int *choice_out, + TrialResult *result, double *x_final_out, - int n_accumulators, ) noexcept nogil: """Simulate one race. A same-grid-step tie uses the lowest index. @@ -125,40 +178,75 @@ cdef void _run_race_trial( Xoshiro256State rng BoxMullerState bm double particle[MAX_ACCUMULATORS] - double t_particle, dt_current, sqrt_dt, boundary + double t_particle, dt_current, sqrt_dt, boundary, next_node + double drift_increment, diffusion_increment int stage[MAX_ACCUMULATORS] - int i, step, winner + int i, step, winner, stage_changed, n_accumulators _seed(&rng, seed) bm.has_spare = 0 t_particle = 0.0 - rt_out[0] = -1.0 - choice_out[0] = 0 + result.rt = -1.0 + result.choice = 0 + n_accumulators = mu.shape[1] for i in range(n_accumulators): particle[i] = x0[row, i] stage[i] = 0 - for step in range(max_steps): - dt_current = horizon - t_particle + for step in range(config.max_steps): + # A node reached by the preceding propagation begins its new stage + # before this iteration can draw additional noise. + stage_changed = 0 + for i in range(n_accumulators): + while _has_reached_next_stage( + nodes, d, row, i, stage[i], t_particle + ): + stage[i] += 1 + stage_changed = 1 + + # A discontinuity in the boundary can itself end the race. Use the + # previous Euler-step midpoint convention (or zero at the start). + if stage_changed: + winner = -1 + for i in range(n_accumulators): + boundary = _upper_boundary_at( + upper_intercept, upper_slope, nodes, + row, i, stage[i], t_particle, + ) + if particle[i] >= boundary and winner < 0: + winner = i + if winner >= 0: + result.rt = t_particle - 0.5 * dt_current if step > 0 else 0.0 + result.choice = winner + break + + dt_current = config.horizon - t_particle if dt_current <= 0.0: break - if dt_current > dt: - dt_current = dt + dt_current = min(dt_current, config.dt) + + # Do not propagate through a stage node with the preceding stage's + # dynamics. The earliest pending node controls this Euler step. + for i in range(n_accumulators): + if stage[i] + 1 < d[row, i]: + next_node = nodes[row, i, stage[i] + 1] + if next_node < t_particle + dt_current: + dt_current = next_node - t_particle sqrt_dt = sqrt(dt_current) for i in range(n_accumulators): - particle[i] += ( - mu[row, i, stage[i]] * dt_current - + sigma[row, i, stage[i]] * sqrt_dt * _normal(&rng, &bm) + drift_increment = mu[row, i, stage[i]] * dt_current + diffusion_increment = ( + sigma[row, i, stage[i]] * sqrt_dt * _normal(&rng, &bm) ) + particle[i] += drift_increment + diffusion_increment t_particle += dt_current winner = -1 for i in range(n_accumulators): - boundary = ( - upper_intercept[row, i, stage[i]] - + upper_slope[row, i, stage[i]] - * (t_particle - nodes[row, i, stage[i]]) + boundary = _upper_boundary_at( + upper_intercept, upper_slope, nodes, + row, i, stage[i], t_particle, ) if particle[i] >= boundary and winner < 0: winner = i @@ -166,19 +254,18 @@ cdef void _run_race_trial( if winner >= 0: # The aDDM Efficient-FPT-compatible simulator reports the midpoint # of the Euler step; use the same first-order convention here. - rt_out[0] = t_particle - 0.5 * dt_current - choice_out[0] = winner + result.rt = t_particle - 0.5 * dt_current + result.choice = winner break - for i in range(n_accumulators): - while stage[i] + 1 < d[row, i] and t_particle >= nodes[row, i, stage[i] + 1]: - stage[i] += 1 + # Stage updates happen at the beginning of the next iteration, where + # the new boundary is checked before another noise draw. for i in range(n_accumulators): x_final_out[i] = particle[i] -def _simulate_race_multistage( +cdef void _validate_race_inputs( double[:, :, ::1] mu, double[:, :, ::1] sigma, double[:, :, ::1] nodes, @@ -187,16 +274,12 @@ def _simulate_race_multistage( double[:, :, ::1] upper_slope, double[:, ::1] x0, double dt, - double horizon, uint64_t[::1] seeds, - int n_threads=1, -): - """Low-level batch kernel used by :func:`race_multistage` and tests.""" +) except *: + """Validate the low-level multi-stage race simulator input contract.""" cdef: int n_rows = mu.shape[0] int n_accumulators = mu.shape[1] - int max_steps = int(np.ceil(horizon / dt)) if horizon > 0.0 else 0 - int row if n_accumulators > MAX_ACCUMULATORS: raise ValueError( @@ -221,6 +304,33 @@ def _simulate_race_multistage( if np.any(np.asarray(d) < 1) or np.any(np.asarray(d) > mu.shape[2]): raise ValueError("each d entry must lie between 1 and the padded stage count") + +def _simulate_race_multistage( + double[:, :, ::1] mu, + double[:, :, ::1] sigma, + double[:, :, ::1] nodes, + int[:, ::1] d, + double[:, :, ::1] upper_intercept, + double[:, :, ::1] upper_slope, + double[:, ::1] x0, + double dt, + double horizon, + uint64_t[::1] seeds, + int n_threads=1, +): + """Low-level batch kernel used by :func:`race_multistage` and tests.""" + cdef: + int n_rows = mu.shape[0] + int n_accumulators = mu.shape[1] + int max_steps = int(np.ceil(horizon / dt)) if horizon > 0.0 else 0 + int row + SimulationConfig config + TrialResult *trial_results + + _validate_race_inputs( + mu, sigma, nodes, d, upper_intercept, upper_slope, x0, dt, seeds + ) + rt = np.empty(n_rows, dtype=np.float64) choice = np.empty(n_rows, dtype=np.int32) x_final = np.empty((n_rows, n_accumulators), dtype=np.float64) @@ -236,12 +346,23 @@ def _simulate_race_multistage( x_final[:] = np.asarray(x0) return rt, choice, x_final - for row in prange(n_rows, nogil=True, num_threads=n_threads, schedule='dynamic'): - _run_race_trial( - mu, sigma, nodes, d, upper_intercept, upper_slope, row, x0, - dt, max_steps, horizon, seeds[row], &rt_view[row], - &choice_view[row], &final_view[row, 0], n_accumulators, - ) + config.dt = dt + config.horizon = horizon + config.max_steps = max_steps + trial_results = malloc(n_rows * sizeof(TrialResult)) + if trial_results == NULL: + raise MemoryError("could not allocate race trial results") + try: + for row in prange(n_rows, nogil=True, num_threads=n_threads, schedule='dynamic'): + _run_race_trial( + mu, sigma, nodes, d, upper_intercept, upper_slope, row, x0, + config, seeds[row], &trial_results[row], &final_view[row, 0], + ) + for row in range(n_rows): + rt_view[row] = trial_results[row].rt + choice_view[row] = trial_results[row].choice + finally: + free(trial_results) return rt, choice, x_final diff --git a/ssms/basic_simulators/race_math.py b/ssms/basic_simulators/race_math.py index 24495615..a0efb729 100644 --- a/ssms/basic_simulators/race_math.py +++ b/ssms/basic_simulators/race_math.py @@ -7,19 +7,48 @@ from __future__ import annotations -from math import erf, sqrt +from math import sqrt import numpy as np +from scipy.special import ndtr -_SQRT_2 = sqrt(2.0) _SQRT_2PI = sqrt(2.0 * np.pi) def _normal_cdf(x: np.ndarray | float) -> np.ndarray: - """Standard-normal CDF without requiring SciPy.""" - x = np.asarray(x, dtype=float) - return 0.5 * (1.0 + np.vectorize(erf, otypes=[float])(x / _SQRT_2)) + """Standard-normal CDF.""" + return ndtr(np.asarray(x, dtype=float)) + + +def _validate_race_parameters(sigma: float, T: float, a: float, x0: float) -> None: + """Validate the shared scalar parameters of the one-sided race model.""" + if sigma <= 0.0: + raise ValueError("sigma must be positive") + if T <= 0.0: + raise ValueError("T must be positive") + if x0 >= a: + raise ValueError("x0 must be less than a") + + +def _nonpassage_density( + x: np.ndarray, + mu: float, + sigma: float, + boundary: float, + T: float, + a: float, + x0: float, +) -> np.ndarray: + """Return the Gaussian density corrected for absorption at the boundary.""" + distance_to_boundary = a - x0 + terminal_mean = x0 + mu * T + variance = sigma**2 * T + gaussian = np.exp(-((x - terminal_mean) ** 2) / (2.0 * variance)) + killed_factor = 1.0 - np.exp( + 2.0 * distance_to_boundary * (x - boundary) / variance + ) + return gaussian * killed_factor def small_f( @@ -32,8 +61,7 @@ def small_f( x0: float, ) -> np.ndarray: """One-sided FPT density ``f_tau(t)`` in race-model equation (14).""" - if sigma <= 0.0 or T <= 0.0: - raise ValueError("sigma and T must be positive") + _validate_race_parameters(sigma, T, a, x0) t = np.asarray(t, dtype=float) out = np.zeros_like(t) valid = (t > 0.0) & (t <= T) @@ -60,8 +88,7 @@ def big_F( x0: float, ) -> np.ndarray: """One-sided FPT CDF ``F_tau(t)`` in race-model equation (14).""" - if sigma <= 0.0 or T <= 0.0: - raise ValueError("sigma and T must be positive") + _validate_race_parameters(sigma, T, a, x0) t = np.asarray(t, dtype=float) out = np.zeros_like(t) valid = t > 0.0 @@ -69,10 +96,13 @@ def big_F( distance = a - x0 relative_drift = mu - b root_elapsed = np.sqrt(elapsed) - out[valid] = _normal_cdf( - (relative_drift * elapsed - distance) / (sigma * root_elapsed) - ) + np.exp(2.0 * relative_drift * distance / sigma**2) * _normal_cdf( - (-distance - relative_drift * elapsed) / (sigma * root_elapsed) + standard_error = sigma * root_elapsed + passage_z_score = (relative_drift * elapsed - distance) / standard_error + survival_z_score = (-distance - relative_drift * elapsed) / standard_error + reflection_factor = np.exp(2.0 * relative_drift * distance / sigma**2) + out[valid] = ( + _normal_cdf(passage_z_score) + + reflection_factor * _normal_cdf(survival_z_score) ) return out @@ -87,18 +117,16 @@ def q( x0: float, ) -> np.ndarray: """Killed/non-passage density ``q(x; ..., T, x0)`` in equation (14).""" - if sigma <= 0.0 or T <= 0.0: - raise ValueError("sigma and T must be positive") + _validate_race_parameters(sigma, T, a, x0) x = np.asarray(x, dtype=float) boundary = a + b * T out = np.zeros_like(x) inside = x < boundary x_inside = x[inside] - gaussian = np.exp(-((x_inside - x0 - mu * T) ** 2) / (2.0 * sigma**2 * T)) - killed_factor = 1.0 - np.exp( - 2.0 * (a - x0) * (x_inside - boundary) / (sigma**2 * T) + out[inside] = _nonpassage_density( + x_inside, mu, sigma, boundary, T, a, x0 ) - out[inside] = gaussian * killed_factor / (_SQRT_2PI * sigma * sqrt(T)) + out[inside] /= _SQRT_2PI * sigma * sqrt(T) return out diff --git a/tests/test_race_math_validation.py b/tests/test_race_math_validation.py new file mode 100644 index 00000000..7849376f --- /dev/null +++ b/tests/test_race_math_validation.py @@ -0,0 +1,28 @@ +"""Validation tests for analytical one-sided race-model quantities.""" + +import pytest + +from ssms.basic_simulators.race_math import big_F, q, small_f + + +PARAMS = dict(mu=0.75, sigma=1.0, a=1.0, b=0.0, T=1.5, x0=0.0) + + +@pytest.mark.parametrize( + ("function", "argument"), + [(small_f, 0.5), (big_F, 0.5), (q, 0.0)], +) +@pytest.mark.parametrize( + ("parameter", "value", "message"), + [ + ("sigma", 0.0, "sigma must be positive"), + ("T", 0.0, "T must be positive"), + ("x0", 1.0, "x0 must be less than a"), + ], +) +def test_analytical_race_functions_validate_shared_parameters( + function, argument, parameter, value, message +): + params = {**PARAMS, parameter: value} + with pytest.raises(ValueError, match=message): + function(argument, **params) diff --git a/tests/test_race_multistage_stage_nodes.py b/tests/test_race_multistage_stage_nodes.py new file mode 100644 index 00000000..4e56bdb4 --- /dev/null +++ b/tests/test_race_multistage_stage_nodes.py @@ -0,0 +1,27 @@ +"""Regression tests for multi-stage race simulator node transitions.""" + +import numpy as np + +import cssm + + +def test_euler_step_stops_at_the_next_stage_node(): + """Dynamics after a node must not be applied retroactively to its step.""" + inputs = dict( + mu_array=np.array([[[0.0, 2.0]]]), + sigma_array=np.zeros((1, 1, 2)), + node_array=np.array([[[0.0, 0.5]]]), + d_array=np.array([[2]], dtype=np.int32), + upper_intercept_array=np.ones((1, 1, 2)), + upper_slope_array=np.zeros((1, 1, 2)), + x0_array=np.zeros((1, 1)), + ) + + out = cssm.race_multistage( + **inputs, n_samples=1, delta_t=0.6, max_t=2.0, random_state=3 + ) + + assert out["choices"][0, 0, 0] == 0 + # The 0.5 node splits the first 0.6 step. The second stage then crosses + # during [0.5, 1.1], whose midpoint is 0.8. + np.testing.assert_allclose(out["rts"][0, 0, 0], 0.8, atol=1e-12) From f4e153eb96b0260f9ad6a10ce179c1cf3e26f7f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9Cmariama=E2=80=9D?= <“maria_ma@brown.edu”> Date: Tue, 1 Sep 2026 00:09:36 +0800 Subject: [PATCH 3/3] Fix race simulator boundary validation --- src/cssm/race_multistage_models.pyx | 23 +++++++++++++ ssms/basic_simulators/race_math.py | 13 +++----- tests/test_race_multistage_stage_nodes.py | 40 +++++++++++++++++++++++ 3 files changed, 67 insertions(+), 9 deletions(-) diff --git a/src/cssm/race_multistage_models.pyx b/src/cssm/race_multistage_models.pyx index e2f39384..9930b607 100644 --- a/src/cssm/race_multistage_models.pyx +++ b/src/cssm/race_multistage_models.pyx @@ -193,6 +193,23 @@ cdef void _run_race_trial( particle[i] = x0[row, i] stage[i] = 0 + # An initial value can already cross its stage-0 boundary. Report this + # before drawing noise; iterating in index order preserves tie-breaking. + winner = -1 + for i in range(n_accumulators): + boundary = _upper_boundary_at( + upper_intercept, upper_slope, nodes, + row, i, stage[i], t_particle, + ) + if particle[i] >= boundary and winner < 0: + winner = i + if winner >= 0: + result.rt = 0.0 + result.choice = winner + for i in range(n_accumulators): + x_final_out[i] = particle[i] + return + for step in range(config.max_steps): # A node reached by the preceding propagation begins its new stage # before this iteration can draw additional noise. @@ -280,6 +297,7 @@ cdef void _validate_race_inputs( cdef: int n_rows = mu.shape[0] int n_accumulators = mu.shape[1] + int row, accumulator, stage if n_accumulators > MAX_ACCUMULATORS: raise ValueError( @@ -303,6 +321,11 @@ cdef void _validate_race_inputs( raise ValueError("seeds must contain one seed per row") if np.any(np.asarray(d) < 1) or np.any(np.asarray(d) > mu.shape[2]): raise ValueError("each d entry must lie between 1 and the padded stage count") + for row in range(n_rows): + for accumulator in range(n_accumulators): + for stage in range(d[row, accumulator] - 1): + if nodes[row, accumulator, stage + 1] < nodes[row, accumulator, stage]: + raise ValueError("active stage nodes must be nondecreasing") def _simulate_race_multistage( diff --git a/ssms/basic_simulators/race_math.py b/ssms/basic_simulators/race_math.py index a0efb729..707d1228 100644 --- a/ssms/basic_simulators/race_math.py +++ b/ssms/basic_simulators/race_math.py @@ -45,9 +45,7 @@ def _nonpassage_density( terminal_mean = x0 + mu * T variance = sigma**2 * T gaussian = np.exp(-((x - terminal_mean) ** 2) / (2.0 * variance)) - killed_factor = 1.0 - np.exp( - 2.0 * distance_to_boundary * (x - boundary) / variance - ) + killed_factor = 1.0 - np.exp(2.0 * distance_to_boundary * (x - boundary) / variance) return gaussian * killed_factor @@ -100,9 +98,8 @@ def big_F( passage_z_score = (relative_drift * elapsed - distance) / standard_error survival_z_score = (-distance - relative_drift * elapsed) / standard_error reflection_factor = np.exp(2.0 * relative_drift * distance / sigma**2) - out[valid] = ( - _normal_cdf(passage_z_score) - + reflection_factor * _normal_cdf(survival_z_score) + out[valid] = _normal_cdf(passage_z_score) + reflection_factor * _normal_cdf( + survival_z_score ) return out @@ -123,9 +120,7 @@ def q( out = np.zeros_like(x) inside = x < boundary x_inside = x[inside] - out[inside] = _nonpassage_density( - x_inside, mu, sigma, boundary, T, a, x0 - ) + out[inside] = _nonpassage_density(x_inside, mu, sigma, boundary, T, a, x0) out[inside] /= _SQRT_2PI * sigma * sqrt(T) return out diff --git a/tests/test_race_multistage_stage_nodes.py b/tests/test_race_multistage_stage_nodes.py index 4e56bdb4..7616e278 100644 --- a/tests/test_race_multistage_stage_nodes.py +++ b/tests/test_race_multistage_stage_nodes.py @@ -1,6 +1,7 @@ """Regression tests for multi-stage race simulator node transitions.""" import numpy as np +import pytest import cssm @@ -25,3 +26,42 @@ def test_euler_step_stops_at_the_next_stage_node(): # The 0.5 node splits the first 0.6 step. The second stage then crosses # during [0.5, 1.1], whose midpoint is 0.8. np.testing.assert_allclose(out["rts"][0, 0, 0], 0.8, atol=1e-12) + + +def test_initial_boundary_crossing_has_zero_response_time(): + """An initial crossing must win before the first Euler propagation.""" + out = cssm.race_multistage( + mu_array=np.zeros((1, 2, 1)), + sigma_array=np.zeros((1, 2, 1)), + node_array=np.zeros((1, 2, 1)), + d_array=np.ones((1, 2), dtype=np.int32), + upper_intercept_array=np.array([[[1.0], [0.5]]]), + upper_slope_array=np.zeros((1, 2, 1)), + x0_array=np.array([[1.0, 0.5]]), + n_samples=1, + delta_t=0.1, + max_t=1.0, + random_state=3, + ) + + # Both particles cross at time zero; the lowest accumulator index wins. + assert out["choices"][0, 0, 0] == 0 + np.testing.assert_allclose(out["rts"][0, 0, 0], 0.0, atol=1e-12) + + +def test_decreasing_active_nodes_are_rejected(): + """The simulator must not take a negative Euler step at a stale node.""" + with pytest.raises(ValueError, match="active stage nodes must be nondecreasing"): + cssm.race_multistage( + mu_array=np.zeros((1, 1, 3)), + sigma_array=np.ones((1, 1, 3)), + node_array=np.array([[[0.0, 0.5, 0.2]]]), + d_array=np.array([[3]], dtype=np.int32), + upper_intercept_array=np.ones((1, 1, 3)), + upper_slope_array=np.zeros((1, 1, 3)), + x0_array=np.zeros((1, 1)), + n_samples=1, + delta_t=0.1, + max_t=1.0, + random_state=3, + )