Use alternative samplers (Bayeux)¶

Sample an HSSM model with an alternative sampler from the bayeux-ml package instead of HSSM's defaults.

Install bayeux-ml with pip¶

We assume that you already have HSSM installed in a virtual environment. If you have not, please follow this tutorial to install HSSM.

In the virtual environment where HSSM is installed, you need to install bayeux-ml exactly as follows:

pip install hssm bayeux-ml flowMC==0.3.4 jax[cuda12]==0.4.38 jaxlib==0.4.38

Bayeux has some conflicts with some dependencies because it has not been updated for a while. For that reason this page is not executed by our continuous integration: the pins above (notably jax[cuda12]==0.4.38) conflict with the JAX version HSSM itself is tested against, so the outputs below are from a manual run rather than a scheduled one. Treat the sampler behaviour as illustrative and verify it in your own environment. You might also need to install other packages that it depends on for samplers, such as blackjax.

If sampling fails to start rather than failing to converge, see Set initial values.

Run this tutorial¶

On Colab, uncomment and run the installation cell below once, then restart the runtime. For local setup, GPU extras, and troubleshooting see the Installation guide.

In [ ]:

# %pip install hssm bayeux-ml flowMC==0.3.4 jax[cuda12]==0.4.38 jaxlib==0.4.38

# %pip install hssm bayeux-ml flowMC==0.3.4 jax[cuda12]==0.4.38 jaxlib==0.4.38

In [ ]:

import arviz as az
import bayeux as bx
import jax

import hssm

import arviz as az import bayeux as bx import jax import hssm

Step 1. Build the model with HSSM¶

You can build a model with HSSM as usual. No special steps required. Below is an example:

In [ ]:

# Load a package-supplied dataset
cav_data = hssm.load_data("cavanagh_theta")

# Define a basic hierarchical model with trial-level covariates
model = hssm.HSSM(
    model="ddm",
    data=cav_data,
    include=[
        {
            "name": "v",
            "prior": {
                "Intercept": {"name": "Normal", "mu": 0.0, "sigma": 0.1},
                "theta": {"name": "Normal", "mu": 0.0, "sigma": 0.1},
            },
            "formula": "v ~ theta + (1|participant_id)",
            "link": "identity",
        },
    ],
)

# Load a package-supplied dataset cav_data = hssm.load_data("cavanagh_theta") # Define a basic hierarchical model with trial-level covariates model = hssm.HSSM( model="ddm", data=cav_data, include=[ { "name": "v", "prior": { "Intercept": {"name": "Normal", "mu": 0.0, "sigma": 0.1}, "theta": {"name": "Normal", "mu": 0.0, "sigma": 0.1}, }, "formula": "v ~ theta + (1|participant_id)", "link": "identity", }, ], )

No common intercept. Bounds for parameter v is not applied due to a current limitation of Bambi. This will change in the future.

WARNING:hssm:No common intercept. Bounds for parameter v is not applied due to a current limitation of Bambi. This will change in the future.

Model initialized successfully.

INFO:hssm:Model initialized successfully.

Step 2. Sample with Bayeux¶

Bayeux is compatible with any PyMC model through bx.Model.from_pymc(), which you can access via model.pymc_model from the HSSM object you have constructed.

The example below samples with the numpyro_nuts sampler that bayeux provides.

In [ ]:

bx_model = bx.Model.from_pymc(model.pymc_model)
idata = bx_model.mcmc.numpyro_nuts(seed=jax.random.key(0))

bx_model = bx.Model.from_pymc(model.pymc_model) idata = bx_model.mcmc.numpyro_nuts(seed=jax.random.key(0))

sample: 100%|██████████| 1500/1500 [05:59<00:00,  4.18it/s]

Step 3. Plot trace with inference data¶

The object returned from the sampling process is an InferenceData object, just as HSSM does. You can use ArviZ to inspect it:

In [ ]:

az.plot_trace(idata)

az.plot_trace(idata)

Out[ ]:

array([[<Axes: title={'center': 'a'}>, <Axes: title={'center': 'a'}>],
       [<Axes: title={'center': 't'}>, <Axes: title={'center': 't'}>],
       [<Axes: title={'center': 'v_1|participant_id_mu'}>,
        <Axes: title={'center': 'v_1|participant_id_mu'}>],
       [<Axes: title={'center': 'v_1|participant_id_offset'}>,
        <Axes: title={'center': 'v_1|participant_id_offset'}>],
       [<Axes: title={'center': 'v_1|participant_id_sigma'}>,
        <Axes: title={'center': 'v_1|participant_id_sigma'}>],
       [<Axes: title={'center': 'v_Intercept'}>,
        <Axes: title={'center': 'v_Intercept'}>],
       [<Axes: title={'center': 'v_theta'}>,
        <Axes: title={'center': 'v_theta'}>],
       [<Axes: title={'center': 'z'}>, <Axes: title={'center': 'z'}>]],
      dtype=object)

You can also use any other sampler provided via bayeux:

In [ ]:

idata = bx_model.mcmc.tfp_nuts(seed=jax.random.key(0))

idata = bx_model.mcmc.tfp_nuts(seed=jax.random.key(0))

In [ ]:

az.plot_trace(idata)

az.plot_trace(idata)

Out[ ]:

array([[<Axes: title={'center': 'a'}>, <Axes: title={'center': 'a'}>],
       [<Axes: title={'center': 't'}>, <Axes: title={'center': 't'}>],
       [<Axes: title={'center': 'v_1|participant_id_mu'}>,
        <Axes: title={'center': 'v_1|participant_id_mu'}>],
       [<Axes: title={'center': 'v_1|participant_id_offset'}>,
        <Axes: title={'center': 'v_1|participant_id_offset'}>],
       [<Axes: title={'center': 'v_1|participant_id_sigma'}>,
        <Axes: title={'center': 'v_1|participant_id_sigma'}>],
       [<Axes: title={'center': 'v_Intercept'}>,
        <Axes: title={'center': 'v_Intercept'}>],
       [<Axes: title={'center': 'v_theta'}>,
        <Axes: title={'center': 'v_theta'}>],
       [<Axes: title={'center': 'z'}>, <Axes: title={'center': 'z'}>]],
      dtype=object)