Getting Started¶

This tutorial demonstrates how to quickly get started with the HSSM package. We will cover the following steps:

How to create a model
How to create some simple simulated data
How to specify parameters
How to specify parameters with regressions
How to use ArviZ to summarize and visualize the traces.

Colab Instructions¶

If you would like to run this tutorial on Google colab, please click this link.

Once you are in the colab, follow the installation instructions below and then restart your runtime.

Just uncomment the code in the next code cell and run it!

NOTE:

You may want to switch your runtime to have a GPU or TPU. To do so, go to Runtime > Change runtime type and select the desired hardware accelerator.

Note that if you switch your runtime you have to follow the installation instructions again.

In [1]:

Copied!

# !pip install hssm
# !pip install hssm

Import Modules¶

In [2]:

Copied!

import numpy as np

import hssm

%matplotlib inline
%config InlineBackend.figure_format='retina'
import numpy as np

import hssm

%matplotlib inline
%config InlineBackend.figure_format='retina'

Setting the global float type¶

Note: Using the analytical DDM (Drift Diffusion Model) likelihood in PyMC without setting the float type in PyTensor may result in warning messages during sampling, which is a known bug in PyMC v5.6.0 and earlier versions. To avoid these warnings, we provide a convenience function:

In [3]:

Copied!

hssm.set_floatX("float32")
hssm.set_floatX("float32")

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

Simulating a dataset¶

The hssm.simulate_data() function generates data for most SSM types. Here we simulate some data from a Drift Diffusion Model (DDM) with known true parameter values.

In [4]:

Copied!





v_true, a_true, z_true, t_true = [0.5, 1.5, 0.5, 0.5]
dataset = hssm.simulate_data(
    model="ddm",
    theta=[v_true, a_true, z_true, t_true],
    size=1000,
)

dataset
v_true, a_true, z_true, t_true = [0.5, 1.5, 0.5, 0.5]
dataset = hssm.simulate_data(
    model="ddm",
    theta=[v_true, a_true, z_true, t_true],
    size=1000,
)

dataset

Out[4]:

	rt	response
0	4.782969	-1.0
1	3.615409	1.0
2	0.733378	1.0
3	7.904665	1.0
4	2.653816	-1.0
...	...	...
995	1.566862	1.0
996	2.054777	1.0
997	2.213145	-1.0
998	1.103277	1.0
999	4.738029	1.0

1000 rows × 2 columns

Model specification¶

1. DDM using defaults¶

We begin with a simple example. The only information required to create a model in HSSM is a dataset.

A dataset in HSSM is typically a pandas DataFrame with at least rt and response columns, which indicates response time and choices respectively. Right now, response only accepts values of 1 and -1.

If none of the optional parameters is provided, HSSM will assume that we are modeling a classical DDM model with v, a, z, and t as its parameters. HSSM also provides a default analytical likelihood function and some uninformative priors. These can all be overriden by user inputs.

Note

From HSSM v0.1.2 on, lapse distributions will be enabled by default, with `p_outlier` fixed to 0.05. You can set `p_outlier` to 0 or `None` to disable lapse distributions.

In [5]:

Copied!

simple_ddm_model = hssm.HSSM(data=dataset)
simple_ddm_model
simple_ddm_model = hssm.HSSM(data=dataset)
simple_ddm_model

Out[5]:

Hierarchical Sequential Sampling Model
Model: ddm

Response variable: rt,response
Likelihood: analytical
Observations: 1000

Parameters:

v:
    Prior: Normal(mu: 0.0, sigma: 2.0)
    Explicit bounds: (-inf, inf)
a:
    Prior: HalfNormal(sigma: 2.0)
    Explicit bounds: (0.0, inf)
z:
    Prior: Uniform(lower: 0.0, upper: 1.0)
    Explicit bounds: (0.0, 1.0)
t:
    Prior: HalfNormal(sigma: 2.0, initval: 0.10000000149011612)
    Explicit bounds: (0.0, inf)

Lapse probability: 0.05
Lapse distribution: Uniform(lower: 0.0, upper: 10.0)

Visualizing the model¶

If you have graphviz installed on your machine, you will also be able to visualize the model. Please uncomment the code and run it if you have graphviz installed.

In [6]:

Copied!

# Uncomment if you have graphviz installed
# simple_ddm_model.graph()
# Uncomment if you have graphviz installed
# simple_ddm_model.graph()

Out[6]:

No description has been provided for this image

Performing MCMC sampling¶

Similar to PyMC, HSSM provides the sample() method once the model is created to perform MCMC sampling. By default, it uses PyMC's NUTS sampler. We can use other samplers, which we will cover soon.

Visualizing the traces with `ArviZ`¶

Like that of pm.sample(), the result of model.sample() is also an az.InferenceData object, which can be used with the ArviZ package. The last sample the model has performed is stored in the model.traces property for eazy access. Here we use the az.summary() and az.plot_trace() functions to generate a summary table and diagnostic plots for the samples.

In [8]:

Copied!

simple_ddm_model.summary()
simple_ddm_model.summary()

Out[8]:

	mean	sd	hdi_3%	hdi_97%	mcse_mean	ess_bulk	ess_tail	r_hat
v	0.567	0.036	0.498	0.632	0.001	2731.0	2657.0	1.0
z	0.482	0.014	0.456	0.509	0.000	2470.0	2508.0	1.0
t	0.509	0.024	0.465	0.555	0.000	2484.0	2613.0	1.0
a	1.497	0.029	1.446	1.555	0.001	2593.0	2738.0	1.0

In [9]:

Copied!

simple_ddm_model.plot_trace();
simple_ddm_model.plot_trace();

Congratulations! You have just created and sampled from your first model in HSSM! Parameter recovery seems to be pretty successful.

2. Specifying different model types¶

The default model in HSSM is the classic DDM, but HSSM supports many other model types. Below is a full list of supported models at the moment. We will add new models, and HSSM users can also contribute models. Please click here for our contribution guidelines.

ddm
ddm_sdv
full_ddm
angle
levy
ornstein
weibull
race_no_bias_angle_4
ddm_seq2_no_bias

For more information about these models, please see here.

To model a different model type, please use the model parameter of the HSSM class constructor.

If model provided is one of the supported models above, HSSM also has default specifications for these models. No other specifications is necessary.

If model provided is not one of the above, then HSSM considers that you are using a custom model. We will cover custom models in more detail in a separate tutorial.

Below is an example of specifying the angle model type.

In [10]:

Copied!





# Simulate data for an angle model
angle_data = hssm.simulate_data(
    model="angle",
    theta=[0.5, 1.5, 0.5, 0.5, 0.3],  # true values
    size=1000,
)

angle_model = hssm.HSSM(data=angle_data, model="angle")
angle_model
# Simulate data for an angle model
angle_data = hssm.simulate_data(
    model="angle",
    theta=[0.5, 1.5, 0.5, 0.5, 0.3],  # true values
    size=1000,
)

angle_model = hssm.HSSM(data=angle_data, model="angle")
angle_model

Out[10]:

Hierarchical Sequential Sampling Model
Model: angle

Response variable: rt,response
Likelihood: approx_differentiable
Observations: 1000

Parameters:

v:
    Prior: Uniform(lower: -3.0, upper: 3.0)
    Explicit bounds: (-3.0, 3.0)
a:
    Prior: Uniform(lower: 0.30000001192092896, upper: 3.0)
    Explicit bounds: (0.3, 3.0)
z:
    Prior: Uniform(lower: 0.10000000149011612, upper: 0.8999999761581421)
    Explicit bounds: (0.1, 0.9)
t:
    Prior: Uniform(lower: 0.0010000000474974513, upper: 2.0)
    Explicit bounds: (0.001, 2.0)
theta:
    Prior: Uniform(lower: -0.10000000149011612, upper: 1.2999999523162842)
    Explicit bounds: (-0.1, 1.3)

Lapse probability: 0.05
Lapse distribution: Uniform(lower: 0.0, upper: 10.0)

angle models, by default, use an approximate differentiable likelihood function that relies on JAX for likelihood computation. At the moment, due to the different ways JAX and Jupyter handle parallelism, parallel sampling is not available for this type of likelihood computations. There are ways to get around this, but for now, let's just perform sequential sampling with just one core.

Note: the sample() method internally calls bambi's fit() fit method, which internally calls pymc's sample() function. This means that HSSM's sample() method will accept most parameters accepted by the two other functions. You could pretty much use it just like the pm.sample() function.

In [12]:

Copied!

angle_model.summary()
angle_model.summary()

Out[12]:

	mean	sd	hdi_3%	hdi_97%	mcse_mean	mcse_sd	ess_bulk	ess_tail	r_hat
t	0.483	0.029	0.432	0.540	0.001	0.001	1222.0	1584.0	1.0
z	0.500	0.014	0.474	0.525	0.000	0.000	1608.0	1710.0	1.0
theta	0.332	0.032	0.272	0.392	0.001	0.001	1277.0	1878.0	1.0
a	1.535	0.066	1.406	1.650	0.002	0.001	1198.0	1702.0	1.0
v	0.548	0.046	0.464	0.634	0.001	0.001	1845.0	1830.0	1.0

In [13]:

Copied!

angle_model.plot_trace();
angle_model.plot_trace();

3. Specifying priors: the non-regression case¶

Next, let's take a look at how to specify priors in the non-regression case. In HSSM, parameter specification can be done in two ways:

through the include parameter, or
through a shortcut

3.1 Specifying priors through the `include` parameter¶

The include parameter accept a list of dictionaries or hssm.Param objects. Both dictionaries and hssm.Param objects are equivalent, since the content of the dictionary will be passed as parameters to hssm.Param class during model creation, so it is more of a matter of preference. We recommend the hssm.Param object because some IDEs will be able to provide prompts for possible options of parameters. In the non-regression case, each dictionary typically looks like this:

{
    "name": "v",
    "prior": {
        "name": "Uniform",
        "lower": -5.0,
        "upper": 5.0,
    },
    "bounds": (-10.0, 10.0)
}

This is the equivalent of writing:

hssm.Param(
    "v",
    prior=dict(name="Uniform", upper=-5.0, lower=5.0),
    bounds=(-10.0, 10.0)
)

The name field corresponds with the name of the parameter being specified.

The prior field specifies the distribution of the prior. There are two ways to achieve this:

A dictionary with the name of the distribution (typically captalized) and the parameters of the distribution that you would typically set if you were specifying a distribution in PyMC. For example, if you would like to specify pm.Normal(mu=0.0, sigma=1.0) as the prior, then in HSSM, this prior dictionary would be:

{
    "name": "Normal", ## Note it is capitalized
    "mu": 0.0,
    "sigma": 1.0,
}

or, using the dict function:

dict(name="Normal", mu=0.0, sigma=1.0)

A hssm.Prior object. This is exactly how you would specify priors using bambi (In fact, hssm.Prior is a subclass of bmb.Prior and for the most part can be used interchangeably with bmb.Prior). Then to specify the same normal prior as above, you would write:

hssm.Prior("Normal", mu=0.0, sigma=1.0)

The bounds field accepts a tuple of floats, indicating the lower and upper bounds for the parameter.

Fixing parameters: sometimes you might want to fix the values of a parameter. You can easily do so by specifying that value to the prior field of the dictionary. In the following example, the paramter v is fixed to 0.5.

{
    "name": "v",
    "prior": 0.5,
}

Now let's make this concrete with an example:

In [14]:

Copied!





# A Normal prior for `v` without explicit bounds
param_v = {
    "name": "v",
    "prior": {
        "name": "Normal",
        "mu": 0.0,
        "sigma": 2.0,
    },
}

# A Uniform prior for `a`. Using the `dict` function
param_a = hssm.Param(
    "a",
    prior=dict(
        name="Uniform",
        lower=0.01,
        upper=5,
    ),
    bounds=(0, np.inf),
)

# A Uniform prior for `z` over (0, 1) set using hssm.Prior.
# bounds are not set, existing default bounds will be used
param_z = {"name": "z", "prior": hssm.Prior("Uniform", lower=0.0, upper=1.0)}

# A fixed value for t
param_t = {"name": "t", "prior": 0.5}

example_ddm_model = hssm.HSSM(
    data=dataset,
    model="ddm",
    include=[
        param_v,
        param_a,
        param_z,
        param_t,
    ],
)

example_ddm_model
# A Normal prior for `v` without explicit bounds
param_v = {
    "name": "v",
    "prior": {
        "name": "Normal",
        "mu": 0.0,
        "sigma": 2.0,
    },
}

# A Uniform prior for `a`. Using the `dict` function
param_a = hssm.Param(
    "a",
    prior=dict(
        name="Uniform",
        lower=0.01,
        upper=5,
    ),
    bounds=(0, np.inf),
)

# A Uniform prior for `z` over (0, 1) set using hssm.Prior.
# bounds are not set, existing default bounds will be used
param_z = {"name": "z", "prior": hssm.Prior("Uniform", lower=0.0, upper=1.0)}

# A fixed value for t
param_t = {"name": "t", "prior": 0.5}

example_ddm_model = hssm.HSSM(
    data=dataset,
    model="ddm",
    include=[
        param_v,
        param_a,
        param_z,
        param_t,
    ],
)

example_ddm_model

Out[14]:

Hierarchical Sequential Sampling Model
Model: ddm

Response variable: rt,response
Likelihood: analytical
Observations: 1000

Parameters:

v:
    Prior: Normal(mu: 0.0, sigma: 2.0)
    Explicit bounds: (-inf, inf)
a:
    Prior: Uniform(lower: 0.009999999776482582, upper: 5.0)
    Explicit bounds: (0, inf)
z:
    Prior: Uniform(lower: 0.0, upper: 1.0)
    Explicit bounds: (0.0, 1.0)
t:
    Prior: 0.5
    Explicit bounds: (0.0, inf)

Lapse probability: 0.05
Lapse distribution: Uniform(lower: 0.0, upper: 10.0)

In [16]:

Copied!

example_ddm_model.summary()
example_ddm_model.summary()

Out[16]:

	mean	sd	hdi_3%	hdi_97%	mcse_mean	ess_bulk	ess_tail	r_hat
v	0.571	0.033	0.508	0.631	0.001	2716.0	2809.0	1.0
z	0.479	0.012	0.457	0.502	0.000	2553.0	2962.0	1.0
a	1.503	0.023	1.461	1.546	0.000	3170.0	2837.0	1.0

3.2 Specifying priors using the shortcut¶

HSSM also supports a syntax very similar to PyMC: You can directly specify priors by passing the prior to the name of the parameter in hssm.HSSM. This is convenient if the prior is simple. Below is an example almost equivalent to the above example:

In [17]:

Copied!





# All ways to specify priors mentioned above are supported in the shortcut syntax
shortcut_ddm_model = hssm.HSSM(
    data=dataset,
    model="ddm",
    v={"name": "Normal", "mu": 0.0, "sigma": 2.0},
    a=dict(name="Uniform", lower=0.01, upper=5),
    z=hssm.Prior("Uniform", lower=0.01, upper=1.0),
    t=0.5,
)

shortcut_ddm_model
# All ways to specify priors mentioned above are supported in the shortcut syntax
shortcut_ddm_model = hssm.HSSM(
    data=dataset,
    model="ddm",
    v={"name": "Normal", "mu": 0.0, "sigma": 2.0},
    a=dict(name="Uniform", lower=0.01, upper=5),
    z=hssm.Prior("Uniform", lower=0.01, upper=1.0),
    t=0.5,
)

shortcut_ddm_model

Out[17]:

Hierarchical Sequential Sampling Model
Model: ddm

Response variable: rt,response
Likelihood: analytical
Observations: 1000

Parameters:

v:
    Prior: Normal(mu: 0.0, sigma: 2.0)
    Explicit bounds: (-inf, inf)
a:
    Prior: Uniform(lower: 0.009999999776482582, upper: 5.0)
    Explicit bounds: (0.0, inf)
z:
    Prior: Uniform(lower: 0.009999999776482582, upper: 1.0)
    Explicit bounds: (0.0, 1.0)
t:
    Prior: 0.5
    Explicit bounds: (0.0, inf)

Lapse probability: 0.05
Lapse distribution: Uniform(lower: 0.0, upper: 10.0)

Note that the shortcut syntax also supports specifying bounds. It will be more polished in a future update. We will skip this step for now.

4. Specifying priors: the regression case¶

Built on top of bambi, HSSM uses an lmer-like syntax that makes it extremely straight-forward to specify regressions.

Parameters that are targets of regressions are also specified using dictionaries in include. Below is an example for such dictionaries.

{
    "name": "v",
    "formula": "v ~ 1 + x + y",
    "prior": {
        # All ways to specify priors in the non-regression case
        # work the same way here.
        "Intercept": {"name": "Uniform", "lower": -10.0, "upper": 10.0},
        "x": dict(name="Normal", mu=0, sigma=1),
        "y": hssm.Prior("HalfNormal", sigma=0.5),
        "z": 1.0
    }
    "link": "identity",
    "bounds": (-10.0, 10.0)
}

We see that in the regression case, name and bounds are specified the exact same way as in the non-regression case. The regression formula is specified in a way that's very similar to the lmer package in R. Users that have experience with R formulas should be very familar with this syntax. In this case, the formula means that the parameter v is regressed on variables x and y, which can be found in the dataframe passed to hssm.HSSM. The 1 explicitly specifies an intercept for the regression.

In addition to the formula, the users typically needs to specify priors for the regression coefficients. This is done in the prior field of the dictionary. Instead of specifying priors for the parameter, the priors are now specified for the corresponding regression coefficients. If not specified, HSSM will use default priors generated in Bambi.

The users might also want to specify a link function for generalized linear models. If left unspecified, the identity link function will be used.

Now let's see an example of a regression:

In [18]:

Copied!

# Generate some fake simulation data

intercept = 1.5
x = np.random.uniform(-5.0, 5.0, size=1000)
y = np.random.uniform(-5.0, 5.0, size=1000)

v = intercept + 0.8 * x + 0.3 * y
# Generate some fake simulation data

intercept = 1.5
x = np.random.uniform(-5.0, 5.0, size=1000)
y = np.random.uniform(-5.0, 5.0, size=1000)

v = intercept + 0.8 * x + 0.3 * y

In [19]:

Copied!

true_values = np.column_stack([v, np.repeat([[1.5, 0.5, 0.5]], axis=0, repeats=1000)])
true_values.shape
true_values = np.column_stack([v, np.repeat([[1.5, 0.5, 0.5]], axis=0, repeats=1000)])
true_values.shape

Out[19]:

(1000, 4)

In [20]:

Copied!





dataset_reg_v = hssm.simulate_data(
    model="ddm",
    theta=true_values,
    size=1,  # Generate one data point for each of the 1000 set of true values
)

dataset_reg_v["x"] = x
dataset_reg_v["y"] = y

dataset_reg_v
dataset_reg_v = hssm.simulate_data(
    model="ddm",
    theta=true_values,
    size=1,  # Generate one data point for each of the 1000 set of true values
)

dataset_reg_v["x"] = x
dataset_reg_v["y"] = y

dataset_reg_v

Out[20]:

	rt	response	x	y
0	0.694455	1.0	4.293505	1.462886
1	0.685728	1.0	4.523834	4.172541
2	0.817345	1.0	4.907615	-3.951474
3	1.122517	-1.0	-4.882512	1.316306
4	0.917183	1.0	0.386110	-0.028240
...	...	...	...	...
995	1.438341	1.0	1.074268	-1.578958
996	1.276988	1.0	-1.546509	-0.707855
997	0.748494	1.0	2.516730	3.707871
998	0.672862	1.0	4.995486	-1.544320
999	1.252353	-1.0	-4.960653	-4.543217

1000 rows × 4 columns

In [21]:

Copied!





model_reg_v = hssm.HSSM(
    data=dataset_reg_v,
    include=[
        {
            "name": "v",
            "formula": "v ~ 1 + x + y",
            "prior": {
                "Intercept": {"name": "Uniform", "lower": 0.0, "upper": 0.5},
                "x": dict(name="Uniform", lower=0.0, upper=1.0),
                "y": hssm.Prior("Uniform", lower=0.0, upper=1.0),
            },
            "link": "identity",
        }
    ],
)
model_reg_v
model_reg_v = hssm.HSSM(
    data=dataset_reg_v,
    include=[
        {
            "name": "v",
            "formula": "v ~ 1 + x + y",
            "prior": {
                "Intercept": {"name": "Uniform", "lower": 0.0, "upper": 0.5},
                "x": dict(name="Uniform", lower=0.0, upper=1.0),
                "y": hssm.Prior("Uniform", lower=0.0, upper=1.0),
            },
            "link": "identity",
        }
    ],
)
model_reg_v

Out[21]:

Hierarchical Sequential Sampling Model
Model: ddm

Response variable: rt,response
Likelihood: analytical
Observations: 1000

Parameters:

v:
    Formula: v ~ 1 + x + y
    Priors:
        v_Intercept ~ Uniform(lower: 0.0, upper: 0.5)
        v_x ~ Uniform(lower: 0.0, upper: 1.0)
        v_y ~ Uniform(lower: 0.0, upper: 1.0)
    Link: identity
    Explicit bounds: (-inf, inf)
a:
    Prior: HalfNormal(sigma: 2.0)
    Explicit bounds: (0.0, inf)
z:
    Prior: Uniform(lower: 0.0, upper: 1.0)
    Explicit bounds: (0.0, 1.0)
t:
    Prior: HalfNormal(sigma: 2.0, initval: 0.10000000149011612)
    Explicit bounds: (0.0, inf)

Lapse probability: 0.05
Lapse distribution: Uniform(lower: 0.0, upper: 10.0)

In [22]:

Copied!

# Uncomment to see model graph if you have graphviz installed
# model_reg_v.graph()
# Uncomment to see model graph if you have graphviz installed
# model_reg_v.graph()

Out[22]:

In [23]:

Copied!

trace_reg_v = model_reg_v.sample()
trace_reg_v = model_reg_v.sample()

Auto-assigning NUTS sampler...
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [z, t, a, v_Intercept, v_x, v_y]

100.00% [8000/8000 00:21<00:00 Sampling 4 chains, 0 divergences]

Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 21 seconds.

In [24]:

Copied!

# Looks like parameter recovery was successful
model_reg_v.summary()
# Looks like parameter recovery was successful
model_reg_v.summary()

Out[24]:

	mean	sd	hdi_3%	hdi_97%	mcse_mean	ess_bulk	ess_tail	r_hat
z	0.633	0.009	0.616	0.649	0.000	3879.0	3277.0	1.0
t	0.543	0.009	0.526	0.559	0.000	3704.0	3180.0	1.0
a	1.358	0.033	1.293	1.418	0.001	3601.0	3037.0	1.0
v_Intercept	0.471	0.004	0.464	0.475	0.000	3400.0	2531.0	1.0
v_x	0.631	0.019	0.594	0.665	0.000	3474.0	3027.0	1.0
v_y	0.240	0.014	0.214	0.267	0.000	4755.0	3082.0	1.0

In [25]:

Copied!

model_reg_v.plot_trace();
model_reg_v.plot_trace();

In [ ]: