import logging
import os
import warnings

# This tutorial does not need a GPU. Keep Molab and headless documentation
# builds away from host CUDA plugins before importing HSSM/JAX.
os.environ["JAX_PLATFORMS"] = "cpu"
os.environ["JAX_SKIP_CUDA_CONSTRAINTS_CHECK"] = "1"
warnings.filterwarnings("ignore")
logging.getLogger("jax._src.xla_bridge").setLevel(logging.CRITICAL)

import arviz as az
import marimo as mo
import numpy as np

import hssm

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

FULL_RUN = os.environ.get("FULL_RUN") == "1"

import logging import os import warnings # This tutorial does not need a GPU. Keep Molab and headless documentation # builds away from host CUDA plugins before importing HSSM/JAX. os.environ["JAX_PLATFORMS"] = "cpu" os.environ["JAX_SKIP_CUDA_CONSTRAINTS_CHECK"] = "1" warnings.filterwarnings("ignore") logging.getLogger("jax._src.xla_bridge").setLevel(logging.CRITICAL) import arviz as az import marimo as mo import numpy as np import hssm hssm.set_floatX("float32") az.style.use("seaborn-v0_8-whitegrid") FULL_RUN = os.environ.get("FULL_RUN") == "1"

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

Poisson race models¶

This short tutorial shows how to:

1. simulate synthetic reaction times with the Poisson race simulator from ssm-simulators;
2. fit HSSM's analytical Poisson race likelihood; and
3. compare the recovered posterior with the generating parameters.

Simulate with ssm-simulators¶

hssm.simulate_data wraps the ssm-simulators Poisson race generator. The simulator and HSSM likelihood share the same accumulator-specific parameter names: r1, r2, k1, k2, and the non-decision time t.

true_params = {
    "r1": 1.0,
    "r2": 5.0,
    "k1": 2.0,
    "k2": 2.0,
    "t": 0.25,
}

data = hssm.simulate_data(
    model="poisson_race",
    theta=true_params,
    size=500,
    random_state=123,
)
mo.md(data.head().to_markdown(index=False))

true_params = { "r1": 1.0, "r2": 5.0, "k1": 2.0, "k2": 2.0, "t": 0.25, } data = hssm.simulate_data( model="poisson_race", theta=true_params, size=500, random_state=123, ) mo.md(data.head().to_markdown(index=False))

rt	response
0.938074	-1
0.444526	1
0.609057	1
0.914047	1
0.509135	1

Fit the Poisson race likelihood¶

This model has no regression terms, so its positive simple-parameter priors come from HSSM's model configuration. The model-level prior_settings selector is not a prior dictionary: it accepts only "safe" or None and controls generated regression-term priors. To replace a prior here, provide it through a parameter specification such as include=[...] or a parameter keyword.

Because this is a synthetic recovery check, we start NUTS adaptation at the known generating point. Initial values choose where adaptation begins; they do not change the posterior target. With real data, use scientifically plausible starts and diagnose multiple chains.

_draws = 3_000 if FULL_RUN else 250
_tune = 3_000 if FULL_RUN else 250

_poisson_model = hssm.HSSM(
    data=data,
    model="poisson_race",
    loglik_kind="analytical",
)

idata = _poisson_model.sample(
    draws=_draws,
    tune=_tune,
    chains=2,
    cores=1,
    target_accept=0.95,
    initvals=true_params,
    random_seed=123,
    progressbar=False,
)

_draws = 3_000 if FULL_RUN else 250 _tune = 3_000 if FULL_RUN else 250 _poisson_model = hssm.HSSM( data=data, model="poisson_race", loglik_kind="analytical", ) idata = _poisson_model.sample( draws=_draws, tune=_tune, chains=2, cores=1, target_accept=0.95, initvals=true_params, random_seed=123, progressbar=False, )

Model initialized successfully.

Initializing NUTS using adapt_diag...
Sequential sampling (2 chains in 1 job)
NUTS: [k1, r2, t, k2, r1]
Sampling 2 chains for 3_000 tune and 3_000 draw iterations (6_000 + 6_000 draws total) took 47 seconds.
We recommend running at least 4 chains for robust computation of convergence diagnostics

Compare posteriors against the ground truth¶

ArviZ summarizes the marginal posterior distributions. Adding the generating values to the table and plot makes it easy to check whether the fitted uncertainty covers the parameters used to simulate the data. In the plot, solid red lines mark the generating values and dashed blue lines mark posterior means.

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

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

	mean	sd	eti89_lb	eti89_ub	ess_bulk	ess_tail	r_hat	mcse_mean	mcse_sd	true_value
r1	1.41982	0.411784	0.834143	2.12525	1927.56	2252.93	1.00205	0.00930975	0.00769077	1
r2	5.2836	0.531954	4.50744	6.19299	1414.92	1904.43	1.00105	0.0142903	0.0109482	5
k1	2.32912	0.390063	1.76538	2.98161	1800.45	2194.27	1.00188	0.00914875	0.00749123	2
k2	2.21355	0.268641	1.84602	2.68826	1265.31	1574.02	1.00131	0.00772825	0.00656554	2
t	0.244213	0.014064	0.218543	0.261764	1393.78	1734.46	1.00139	0.000398223	0.000369556	0.25

_plot_collection = az.plot_dist(
    idata,
    var_names=var_names,
    ci_prob=None,
    point_estimate=None,
)

for _var in var_names:
    _axis = _plot_collection.get_viz("plot", _var)
    _posterior_values = idata["posterior"].dataset[_var].values.ravel()
    _lower, _upper = np.quantile(_posterior_values, [0.005, 0.995])

    _axis.axvline(
        np.mean(_posterior_values),
        color="C0",
        linestyle="--",
        linewidth=2,
        label="posterior mean",
    )
    _axis.axvline(
        true_params[_var],
        color="red",
        linestyle="-",
        linewidth=2,
        label="true value",
    )
    _axis.set_xlim(_lower, _upper)

_plot_collection

_plot_collection = az.plot_dist( idata, var_names=var_names, ci_prob=None, point_estimate=None, ) for _var in var_names: _axis = _plot_collection.get_viz("plot", _var) _posterior_values = idata["posterior"].dataset[_var].values.ravel() _lower, _upper = np.quantile(_posterior_values, [0.005, 0.995]) _axis.axvline( np.mean(_posterior_values), color="C0", linestyle="--", linewidth=2, label="posterior mean", ) _axis.axvline( true_params[_var], color="red", linestyle="-", linewidth=2, label="true value", ) _axis.set_xlim(_lower, _upper) _plot_collection

Takeaway¶

HSSM's analytical Poisson race likelihood can recover the generating accumulator rates, thresholds, and non-decision time in this synthetic example. Model-level prior presets are relevant when those parameters are regression targets; explicit parameter priors remain the appropriate tool for changing this intercept-only model.