Skip to content

Add forward race simulation #342

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mariama-design wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Add forward race simulation #342

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5087509
Add forward race simulation utilities
Aug 24, 2026
b885128
Address race simulation review feedback
Aug 31, 2026
f4e153e
Fix race simulator boundary validation
Aug 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
174 changes: 174 additions & 0 deletions notebooks/forward_race_simulator.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
}
152 changes: 152 additions & 0 deletions notebooks/race_npd_numerical_integration.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
}
1 change: 1 addition & 0 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions src/cssm/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -58,6 +59,7 @@
"ddm_flexbound_tradeoff",
# ADDM model
"addm",
"race_multistage",
# Race models
"race_model",
"lca",
Expand Down
Loading