Poisson Race Model Tutorial¶

This short tutorial shows how to (1) simulate synthetic reaction times with the Poisson race simulator from ssms-simulators, (2) fit the analytical Poisson race likelihood provided by HSSM, and (3) inspect the recovered parameters with ArviZ.

In [1]:

import arviz as az
import numpy as np

import hssm

hssm.set_floatX("float32")
az.style.use("seaborn-v0_8-whitegrid")

rng = np.random.default_rng(123)

import arviz as az import numpy as np import hssm hssm.set_floatX("float32") az.style.use("seaborn-v0_8-whitegrid") rng = np.random.default_rng(123)

Setting PyTensor floatX type to float32.

Setting "jax_enable_x64" to False. If this is not intended, please set `jax` to False.

Simulate with ssms-simulators¶

hssm.simulate_data wraps the ssms-simulators Poisson race generator. The simulator names its accumulator-specific parameters with zero-based indices (r0/r1 and k0/k1). We mirror those values into an HSSM-friendly dict so posterior checks line up with the likelihood parameterization (r1, r2, k1, k2, t).

In [2]:

"""
ssms_params = {
    "r0": 2.8,
    "r1": 3.4,
    "k0": 1.6,
    "k1": 1.2,
    "t": 0.25,
}
true_params = {
    "r1": ssms_params["r0"],
    "r2": ssms_params["r1"],
    "k1": ssms_params["k0"],
    "k2": ssms_params["k1"],
    "t": ssms_params["t"],
}
n_trials = 400

data = hssm.simulate_data(
    model="poisson_race",
    theta=true_params,
    size=n_trials,
    random_state=123,
)
data.head()
"""
ssms_params = {
    "r0": 1,
    "r1": 5,
    "k0": 2,
    "k1": 2,
    "t": 0.25,
}
true_params = {
    "r1": ssms_params["r0"],
    "r2": ssms_params["r1"],
    "k1": ssms_params["k0"],
    "k2": ssms_params["k1"],
    "t": ssms_params["t"],
}
n_trials = 500

data = hssm.simulate_data(
    model="poisson_race",
    theta=true_params,
    size=n_trials,
    random_state=123,
)
data.head()

""" ssms_params = { "r0": 2.8, "r1": 3.4, "k0": 1.6, "k1": 1.2, "t": 0.25, } true_params = { "r1": ssms_params["r0"], "r2": ssms_params["r1"], "k1": ssms_params["k0"], "k2": ssms_params["k1"], "t": ssms_params["t"], } n_trials = 400 data = hssm.simulate_data( model="poisson_race", theta=true_params, size=n_trials, random_state=123, ) data.head() """ ssms_params = { "r0": 1, "r1": 5, "k0": 2, "k1": 2, "t": 0.25, } true_params = { "r1": ssms_params["r0"], "r2": ssms_params["r1"], "k1": ssms_params["k0"], "k2": ssms_params["k1"], "t": ssms_params["t"], } n_trials = 500 data = hssm.simulate_data( model="poisson_race", theta=true_params, size=n_trials, random_state=123, ) data.head()

Out[2]:

	rt	response
0	0.938074	-1.0
1	0.444526	1.0
2	0.609057	1.0
3	0.914047	1.0
4	0.509135	1.0

Fit the Poisson race likelihood¶

We pass the simulated data into an HSSM object that uses the analytical Poisson race likelihood. The defaults already enforce positivity for all parameters; you can override them by passing a prior_settings dict.

In [3]:

poisson_model = hssm.HSSM(
    data=data,
    model="poisson_race",
    loglik_kind="analytical",
    prior_settings=None,
)

idata = poisson_model.sample(
    draws=3000,
    tune=3000,
    chains=2,
    cores=2,
    target_accept=0.9,
    random_seed=123,
)

poisson_model = hssm.HSSM( data=data, model="poisson_race", loglik_kind="analytical", prior_settings=None, ) idata = poisson_model.sample( draws=3000, tune=3000, chains=2, cores=2, target_accept=0.9, random_seed=123, )

Model initialized successfully.

Using default initvals. 

Initializing NUTS using adapt_diag...

Multiprocess sampling (2 chains in 2 jobs)

NUTS: [k1, r2, t, k2, r1]

/Users/yxu150/.local/share/uv/python/cpython-3.13.13-macos-aarch64-none/lib/python3.13/multiprocessing/popen_fork.py:67: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
  self.pid = os.fork()

/Users/yxu150/.local/share/uv/python/cpython-3.13.13-macos-aarch64-none/lib/python3.13/multiprocessing/popen_fork.py:67: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
  self.pid = os.fork()

Sampling 2 chains for 3_000 tune and 3_000 draw iterations (6_000 + 6_000 draws total) took 15 seconds.

There were 10 divergences after tuning. Increase `target_accept` or reparameterize.

We recommend running at least 4 chains for robust computation of convergence diagnostics

Compare posteriors against the ground truth¶

ArviZ summarises the marginal distributions and allows us to verify that the posterior means/credible intervals overlap the parameters used to simulate the data.

In [4]:

var_names = list(true_params.keys())
summary = az.summary(idata, var_names=var_names)
summary["true_value"] = [true_params[name] for name in summary.index]
summary

var_names = list(true_params.keys()) summary = az.summary(idata, var_names=var_names) summary["true_value"] = [true_params[name] for name in summary.index] summary

Out[4]:

	mean	sd	eti89_lb	eti89_ub	ess_bulk	ess_tail	r_hat	mcse_mean	mcse_sd	true_value
r1	1.43	0.42	0.84	2.2	1637	1834	1.00	0.01	0.0084	1.00
r2	5.31	0.55	4.5	6.3	1755	1867	1.00	0.014	0.011	5.00
k1	2.34	0.39	1.8	3	1573	1838	1.00	0.0099	0.0079	2.00
k2	2.23	0.29	1.9	2.7	1581	1668	1.00	0.0076	0.0074	2.00
t	0.2431	0.0151	0.22	0.26	1691	1647	1.00	0.0004	0.00043	0.25

In [5]:

# ArviZ 1.0 — plot_dist returns a PlotCollection, not axes
pc = az.plot_dist(
    idata,  # xarray.DataTree (was InferenceData)
    var_names=var_names,
    ci_prob=None,  # was hdi_prob=None
    point_estimate=None,
)

for var in var_names:
    ax = pc.get_viz("plot", var)  # was: iterating np.ravel(axes)

    # posterior group is now a DataTree node; .dataset gives the xr.Dataset
    post_vals = idata["posterior"].dataset[var].values.ravel()

    lo, hi = np.quantile(post_vals, [0.005, 0.995])
    post_mean = np.mean(post_vals)

    ax.axvline(
        post_mean, color="C0", linestyle="--", linewidth=2, label="posterior mean"
    )
    ax.axvline(
        true_params[var], color="red", linestyle="-", linewidth=2, label="true value"
    )
    ax.set_xlim(lo, hi)

# ArviZ 1.0 — plot_dist returns a PlotCollection, not axes pc = az.plot_dist( idata, # xarray.DataTree (was InferenceData) var_names=var_names, ci_prob=None, # was hdi_prob=None point_estimate=None, ) for var in var_names: ax = pc.get_viz("plot", var) # was: iterating np.ravel(axes) # posterior group is now a DataTree node; .dataset gives the xr.Dataset post_vals = idata["posterior"].dataset[var].values.ravel() lo, hi = np.quantile(post_vals, [0.005, 0.995]) post_mean = np.mean(post_vals) ax.axvline( post_mean, color="C0", linestyle="--", linewidth=2, label="posterior mean" ) ax.axvline( true_params[var], color="red", linestyle="-", linewidth=2, label="true value" ) ax.set_xlim(lo, hi)

In [ ]: