Construct Custom Models from simulators and JAX callables¶

In [9]:

import os
import pickle

import arviz as az
import jax.numpy as jnp
import lanfactory as lf
import matplotlib.pyplot as plt
import numpy as np
import pymc as pm

import hssm
from hssm.config import ModelConfig

import os import pickle import arviz as az import jax.numpy as jnp import lanfactory as lf import matplotlib.pyplot as plt import numpy as np import pymc as pm import hssm from hssm.config import ModelConfig

Simulate Data¶

As a pre-amble, we will simulate a simple dataset from the DDM model to use through the example.

In [10]:

# simulate some data from the model
obs_ddm = hssm.simulate_data(
    theta=dict(v=1.0, a=1.5, t=0.3, z=0.5), model="ddm", size=1000
)

# simulate some data from the model obs_ddm = hssm.simulate_data( theta=dict(v=1.0, a=1.5, t=0.3, z=0.5), model="ddm", size=1000 )

Construct PyMC distribution from simulator and JAX callable¶

Create JAX Log-likelihood Callable¶

What we need is a callable that takes in a matrix that stacks model parameters and data trialwise, hence with input dimensions: $trials \times (dim-parameters + dim-data)$. The function should return a vector of length $trials$ that contains the log-likelihood for each trial.

If we have a JAX function with this signature, we will be able to proceed to create a valid PyMC distribution form helper functions provided by hssm.

In this particular example, we will reinstantiate a pretrained LAN via utilities from the lanfactory package. This is not necessary, but illustrates how to make use of the braoder ecosystem around HSSM.

In [11]:

# Loaded Net
jax_infer = lf.trainers.MLPJaxFactory(
    network_config=pickle.load(
        open(
            os.path.join(
                "data", "jax_models", "ddm", "ddm_jax_lan_network_config.pickle"
            ),
            "rb",
        )
    ),
    train=False,
)

jax_logp, _ = jax_infer.make_forward_partial(
    seed=42,
    input_dim=4 + 2,  # n-parameters (v,a,z,t) + n-data (rts and choices)
    state=os.path.join("data", "jax_models", "ddm", "ddm_jax_lan_train_state.jax"),
    add_jitted=True,
)

# Loaded Net jax_infer = lf.trainers.MLPJaxFactory( network_config=pickle.load( open( os.path.join( "data", "jax_models", "ddm", "ddm_jax_lan_network_config.pickle" ), "rb", ) ), train=False, ) jax_logp, _ = jax_infer.make_forward_partial( seed=42, input_dim=4 + 2, # n-parameters (v,a,z,t) + n-data (rts and choices) state=os.path.join("data", "jax_models", "ddm", "ddm_jax_lan_train_state.jax"), add_jitted=True, )

passing through identity

Checking the signature of the JAX callable¶

We can test the signature of the JAX callable by passing in a batch of inputs and checking the output shape.

In [12]:

# Testing the signature of the JAX function 1
n_trials = 10
jax_logp(np.tile(np.array([1.0, 1.5, 0.5, 0.3, 1.6, 1.0]), (n_trials, 1)))

# Testing the signature of the JAX function 1 n_trials = 10 jax_logp(np.tile(np.array([1.0, 1.5, 0.5, 0.3, 1.6, 1.0]), (n_trials, 1)))

passing through identity

Out[12]:

Array([[-0.95924026],
       [-0.95924026],
       [-0.95924026],
       [-0.95924026],
       [-0.95924026],
       [-0.95924026],
       [-0.95924026],
       [-0.95924026],
       [-0.95924026],
       [-0.95924026]], dtype=float64)

In [13]:

# Testing the signature of the JAX function 2
n_dim_model_parameters = 4
n_dim_data = 2
in_ = jnp.zeros((n_trials, n_dim_model_parameters + n_dim_data))
out = jax_logp(in_)
print(out.shape)

# Testing the signature of the JAX function 2 n_dim_model_parameters = 4 n_dim_data = 2 in_ = jnp.zeros((n_trials, n_dim_model_parameters + n_dim_data)) out = jax_logp(in_) print(out.shape)

passing through identity
(10, 1)

Decorate and wrap a simulator¶

The simulator-signature expected by hssm is the following:

A simulator is a callable that takes in a matrix that stacks model parameters and data trialwise, hence with input dimensions: $trials \times (dim-parameters)$. The function should return a matrix of shape $trials \times dim-data$.

We will use the decorate_atomic_simulator() utility to annotate the simulator with necessary metadata.

Again, for illustration, as a starting point we will use a simulator from the ssm-simulators package

You can start from any simulator you like.

In [14]:

from functools import partial

from ssms.basic_simulators.simulator import simulator

from hssm.utils import decorate_atomic_simulator


def sim_wrapper(simulator_fun, theta, model, n_samples, random_state, **kwargs):
    """Wrap a simulator function to match HSSM's expected interface.

    Parameters
    ----------
    simulator_fun : callable
        The simulator function to wrap
    theta : array-like
        Model parameters, shape (n_trials, n_parameters)
    model : str
        Name of the model to simulate
    n_samples : int
        Number of samples to generate per trial
    random_state : int or numpy.random.Generator
        Random seed or random number generator
    **kwargs
        Additional keyword arguments passed to simulator_fun

    Returns
    -------
    array-like
        Array of shape (n_trials, 2) containing reaction times and choices
        stacked column-wise
    """
    out = simulator_fun(
        theta=theta,
        model=model,
        n_samples=n_samples,
        random_state=random_state,
        **kwargs,
    )
    return np.column_stack([out["rts"], out["choices"]])


my_wrapped_simulator = partial(
    sim_wrapper, simulator_fun=simulator, model="ddm", n_samples=1
)

decorated_simulator = decorate_atomic_simulator(
    model_name="ddm", choices=[-1, 1], obs_dim=2
)(my_wrapped_simulator)

from functools import partial from ssms.basic_simulators.simulator import simulator from hssm.utils import decorate_atomic_simulator def sim_wrapper(simulator_fun, theta, model, n_samples, random_state, **kwargs): """Wrap a simulator function to match HSSM's expected interface. Parameters ---------- simulator_fun : callable The simulator function to wrap theta : array-like Model parameters, shape (n_trials, n_parameters) model : str Name of the model to simulate n_samples : int Number of samples to generate per trial random_state : int or numpy.random.Generator Random seed or random number generator **kwargs Additional keyword arguments passed to simulator_fun Returns ------- array-like Array of shape (n_trials, 2) containing reaction times and choices stacked column-wise """ out = simulator_fun( theta=theta, model=model, n_samples=n_samples, random_state=random_state, **kwargs, ) return np.column_stack([out["rts"], out["choices"]]) my_wrapped_simulator = partial( sim_wrapper, simulator_fun=simulator, model="ddm", n_samples=1 ) decorated_simulator = decorate_atomic_simulator( model_name="ddm", choices=[-1, 1], obs_dim=2 )(my_wrapped_simulator)

Checking the signature of the decorated simulator¶

We can check the signature of the decorated simulator by passing in a batch of parameters and checking the output shape.

In [15]:

decorated_simulator(
    theta=np.tile(np.array([1.0, 1.5, 0.5, 0.3]), (10, 1)), random_state=42
)

decorated_simulator( theta=np.tile(np.array([1.0, 1.5, 0.5, 0.3]), (10, 1)), random_state=42 )

Out[15]:

array([[ 2.12210941,  1.        ],
       [ 2.6890626 ,  1.        ],
       [ 1.28667963,  1.        ],
       [ 1.16390681,  1.        ],
       [ 1.42591691,  1.        ],
       [ 1.14748347,  1.        ],
       [ 1.47392309,  1.        ],
       [ 0.88049531, -1.        ],
       [ 1.78516912,  1.        ],
       [ 1.20256472,  1.        ]])

Create valid PyMC distribution¶

We have all ingredients to create a valid PyMC distribution, form a few helper functions provided by hssm at this point.

A valid Distribution will need two ingredients:

1. A RandomVariable (RV) that encodes the simulator and parameter names.
2. A likelihood function, which is a valid PyTensor Op that encodes the log-likelihood of the model.

We will use the make_hssm_rv() helper function to create the RandomVariable and the make_likelihood_callable() helper function to create the likelihood Op.

Finally, the make_distribution() helper function will package everything into a valid PyMC distribution.

NOTE:

There are a few helpful settings which we can use to customize our distribution. We will not cover all details here, however it is worth highlighting the params_is_reg argument in make_likelihood_callable(). This argument is used to tell PyMC whether the parameter is a regression parameter or not. Specifically, if a parameter is treated as a regression, the likelihood function is built to assume that the parameter is passed trial-wise, i.e. as a vector of length n_trials.

Note here, we set all parameters to be non-regression parameters, as we expect the parameters to be passed as single values in the simple PyMC model below.

In [16]:

from hssm.distribution_utils.dist import (
    make_distribution,
    make_hssm_rv,
    make_likelihood_callable,
)

# Step 1: Define a pytensor RandomVariable
CustomRV = make_hssm_rv(
    simulator_fun=decorated_simulator, list_params=["v", "a", "z", "t"]
)

# Step 2: Define a likelihood function
logp_jax_op = make_likelihood_callable(
    loglik=jax_logp,
    loglik_kind="approx_differentiable",
    backend="jax",
    params_is_reg=[False, False, False, False],
    params_only=False,
)

# Step 3: Define a distribution
CustomDistribution = make_distribution(
    rv=CustomRV,
    loglik=logp_jax_op,
    list_params=["v", "a", "z", "t"],
    bounds=dict(v=(-3, 3), a=(0.5, 3.0), z=(0.1, 0.9), t=(0, 2.0)),
)

from hssm.distribution_utils.dist import ( make_distribution, make_hssm_rv, make_likelihood_callable, ) # Step 1: Define a pytensor RandomVariable CustomRV = make_hssm_rv( simulator_fun=decorated_simulator, list_params=["v", "a", "z", "t"] ) # Step 2: Define a likelihood function logp_jax_op = make_likelihood_callable( loglik=jax_logp, loglik_kind="approx_differentiable", backend="jax", params_is_reg=[False, False, False, False], params_only=False, ) # Step 3: Define a distribution CustomDistribution = make_distribution( rv=CustomRV, loglik=logp_jax_op, list_params=["v", "a", "z", "t"], bounds=dict(v=(-3, 3), a=(0.5, 3.0), z=(0.1, 0.9), t=(0, 2.0)), )

We can now test the distribution by passing it to a simple PyMC model.

In [17]:

# Test via basic pymc model
with pm.Model() as model:
    v = pm.Normal("v", mu=0, sigma=1)
    a = pm.Uniform("a", lower=0.5, upper=3.0)
    z = pm.Beta("z", alpha=10, beta=10)
    t = pm.Weibull("t", alpha=0.5, beta=1.0)
    CustomDistribution("custom", v=v, a=a, z=z, t=t, observed=obs_ddm.values)
with model:
    idata = pm.sample(draws=1000, tune=200, chains=1, nuts_sampler="numpyro")

# Test via basic pymc model with pm.Model() as model: v = pm.Normal("v", mu=0, sigma=1) a = pm.Uniform("a", lower=0.5, upper=3.0) z = pm.Beta("z", alpha=10, beta=10) t = pm.Weibull("t", alpha=0.5, beta=1.0) CustomDistribution("custom", v=v, a=a, z=z, t=t, observed=obs_ddm.values) with model: idata = pm.sample(draws=1000, tune=200, chains=1, nuts_sampler="numpyro")

passing through identity
passing through identity

  0%|          | 0/1200 [00:00<?, ?it/s]

passing through identity

sample: 100%|██████████| 1200/1200 [02:28<00:00,  8.07it/s, 23 steps of size 1.37e+00. acc. prob=0.92] 
Only one chain was sampled, this makes it impossible to run some convergence checks

In [18]:

az.plot_trace(idata)

az.plot_trace(idata)

Out[18]:

array([[<Axes: title={'center': 'a'}>, <Axes: title={'center': 'a'}>],
       [<Axes: title={'center': 't'}>, <Axes: title={'center': 't'}>],
       [<Axes: title={'center': 'v'}>, <Axes: title={'center': 'v'}>],
       [<Axes: title={'center': 'z'}>, <Axes: title={'center': 'z'}>]],
      dtype=object)

Custom HSSM model from simulator and JAX callable¶

Next, we will create a custom HSSM model from the simulator and JAX callable. After the work we have done above, this is now very straightforward.

The only hssm specific extra step is to define a ModelConfig object, which bundles all information about the model.

Then we pass our ModelConfig object to the HSSM class, along with the data and the log-likelihood function, and hssm will take care of the rest.

Importantly, hssm will automatically understand how to construct the correct likelihood function for the specified model configuration (parameter-wise regression settings, etc.). A step we have to accomplish manually in the code above.

In [19]:

# Define model config
my_custom_model_config = ModelConfig(
    response=["rt", "response"],
    list_params=["v", "a", "z", "t"],
    bounds={
        "v": (-2.5, 2.5),
        "a": (1.0, 3.0),
        "z": (0.0, 0.9),
        "t": (0.001, 2),
    },
    rv=decorated_simulator,
    backend="jax",
    choices=[-1, 1],
)

# Define model config my_custom_model_config = ModelConfig( response=["rt", "response"], list_params=["v", "a", "z", "t"], bounds={ "v": (-2.5, 2.5), "a": (1.0, 3.0), "z": (0.0, 0.9), "t": (0.001, 2), }, rv=decorated_simulator, backend="jax", choices=[-1, 1], )

In [20]:

# Define the HSSM model
model = hssm.HSSM(
    data=obs_ddm,
    model="my_new_model",  # some name for the model
    model_config=my_custom_model_config,
    loglik_kind="approx_differentiable",  # use the blackbox loglik
    loglik=jax_logp,
    p_outlier=0,
)

model.graph()

# Define the HSSM model model = hssm.HSSM( data=obs_ddm, model="my_new_model", # some name for the model model_config=my_custom_model_config, loglik_kind="approx_differentiable", # use the blackbox loglik loglik=jax_logp, p_outlier=0, ) model.graph()

You have specified the `lapse` argument to include a lapse distribution, but `p_outlier` is set to either 0 or None. Your lapse distribution will be ignored.

Model initialized successfully.

Out[20]:

In [21]:

# Test sampling
model.sample(draws=500, tune=200, nuts_sampler="numpyro", discard_tuned_samples=False)

# Test sampling model.sample(draws=500, tune=200, nuts_sampler="numpyro", discard_tuned_samples=False)

Using default initvals. 

/Users/afengler/miniconda3/envs/hssm516/lib/python3.11/site-packages/pymc/sampling/jax.py:451: UserWarning: There are not enough devices to run parallel chains: expected 4 but got 1. Chains will be drawn sequentially. If you are running MCMC in CPU, consider using `numpyro.set_host_device_count(4)` at the beginning of your program. You can double-check how many devices are available in your system using `jax.local_device_count()`.
  pmap_numpyro = MCMC(

passing through identity

  0%|          | 0/700 [00:00<?, ?it/s]

passing through identity

sample: 100%|██████████| 700/700 [00:11<00:00, 59.57it/s, 15 steps of size 3.33e-01. acc. prob=0.90]

passing through identity

sample: 100%|██████████| 700/700 [00:11<00:00, 58.66it/s, 7 steps of size 2.87e-01. acc. prob=0.92] 

passing through identity

sample: 100%|██████████| 700/700 [00:11<00:00, 63.43it/s, 3 steps of size 3.11e-01. acc. prob=0.94]  

passing through identity

sample: 100%|██████████| 700/700 [00:08<00:00, 87.01it/s, 15 steps of size 3.84e-01. acc. prob=0.93] 
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
 10%|█         | 204/2000 [00:00<00:01, 1053.03it/s]

passing through identity

100%|██████████| 2000/2000 [00:01<00:00, 1166.14it/s]

Out[21]:

arviz.InferenceData

• posterior

  <xarray.Dataset> Size: 68kB
  Dimensions:  (chain: 4, draw: 500)
  Coordinates:
    * chain    (chain) int64 32B 0 1 2 3
    * draw     (draw) int64 4kB 0 1 2 3 4 5 6 7 ... 493 494 495 496 497 498 499
  Data variables:
      z        (chain, draw) float64 16kB 0.398 0.3987 0.3995 ... 0.3791 0.3801
      v        (chain, draw) float64 16kB 1.395 1.433 1.378 ... 1.442 1.422 1.457
      a        (chain, draw) float64 16kB 1.753 1.768 1.755 ... 1.821 1.831 1.774
      t        (chain, draw) float64 16kB 0.1742 0.1857 0.1493 ... 0.1111 0.1386
  Attributes:
      created_at:                  2024-12-25T23:41:36.301669+00:00
      arviz_version:               0.18.0
      inference_library:           numpyro
      inference_library_version:   0.15.2
      sampling_time:               43.397043
      tuning_steps:                200
      modeling_interface:          bambi
      modeling_interface_version:  0.14.0

  xarray.Dataset

  • Dimensions:
    • chain: 4
    • draw: 500
  • Coordinates: (2)
    • chain
      (chain)
      int64
      0 1 2 3

      array([0, 1, 2, 3])

    • draw
      (draw)
      int64
      0 1 2 3 4 5 ... 495 496 497 498 499

      array([  0,   1,   2, ..., 497, 498, 499])

  • Data variables: (4)
    • z
      (chain, draw)
      float64
      0.398 0.3987 ... 0.3791 0.3801

      array([[0.39802242, 0.39870379, 0.3995235 , ..., 0.38775665, 0.38775665,
              0.3916384 ],
             [0.3962099 , 0.3900404 , 0.38827552, ..., 0.38081589, 0.388641  ,
              0.38808292],
             [0.38442812, 0.38190643, 0.38664855, ..., 0.40199733, 0.39126957,
              0.39163304],
             [0.3685333 , 0.36138627, 0.3749158 , ..., 0.37767698, 0.37910233,
              0.38011081]])

    • v
      (chain, draw)
      float64
      1.395 1.433 1.378 ... 1.422 1.457

      array([[1.39534826, 1.43263641, 1.37773822, ..., 1.37386243, 1.37386243,
              1.39960224],
             [1.3806851 , 1.37218431, 1.41756181, ..., 1.41933502, 1.42457808,
              1.44338209],
             [1.41187257, 1.44680853, 1.40547592, ..., 1.39148927, 1.38359145,
              1.3770154 ],
             [1.44466225, 1.42905195, 1.41174314, ..., 1.44174677, 1.42155037,
              1.45702834]])

    • a
      (chain, draw)
      float64
      1.753 1.768 1.755 ... 1.831 1.774

      array([[1.75280387, 1.76775727, 1.75512388, ..., 1.73394364, 1.73394364,
              1.74082372],
             [1.77373812, 1.79549099, 1.66551739, ..., 1.81403889, 1.78765584,
              1.80404876],
             [1.80503584, 1.81160781, 1.76919326, ..., 1.78587591, 1.79011485,
              1.78544739],
             [1.80154644, 1.80474738, 1.806251  , ..., 1.820694  , 1.83051376,
              1.77357071]])

    • t
      (chain, draw)
      float64
      0.1742 0.1857 ... 0.1111 0.1386

      array([[0.17419458, 0.18565053, 0.14928867, ..., 0.16363337, 0.16363337,
              0.14086545],
             [0.13754501, 0.14295353, 0.22690127, ..., 0.1327304 , 0.12744442,
              0.13322138],
             [0.14017347, 0.1231642 , 0.15443876, ..., 0.15114957, 0.13815675,
              0.11718326],
             [0.10579745, 0.10273204, 0.09588006, ..., 0.10341163, 0.11108935,
              0.13856074]])

  • Indexes: (2)
    • chain
      PandasIndex

      PandasIndex(Index([0, 1, 2, 3], dtype='int64', name='chain'))

    • draw
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             490, 491, 492, 493, 494, 495, 496, 497, 498, 499],
            dtype='int64', name='draw', length=500))

  • Attributes: (8)

    created_at :

    2024-12-25T23:41:36.301669+00:00

    arviz_version :

    0.18.0

    inference_library :

    numpyro

    inference_library_version :

    0.15.2

    sampling_time :

    43.397043

    tuning_steps :

    200

    modeling_interface :

    bambi

    modeling_interface_version :

    0.14.0

  
  
• log_likelihood

  <xarray.Dataset> Size: 16MB
  Dimensions:      (chain: 4, draw: 500, __obs__: 1000)
  Coordinates:
    * chain        (chain) int64 32B 0 1 2 3
    * draw         (draw) int64 4kB 0 1 2 3 4 5 6 ... 493 494 495 496 497 498 499
    * __obs__      (__obs__) int64 8kB 0 1 2 3 4 5 6 ... 994 995 996 997 998 999
  Data variables:
      rt,response  (chain, draw, __obs__) float64 16MB -0.4214 -2.193 ... -1.67
  Attributes:
      modeling_interface:          bambi
      modeling_interface_version:  0.14.0

  xarray.Dataset

  • Dimensions:
    • chain: 4
    • draw: 500
    • __obs__: 1000
  • Coordinates: (3)
    • chain
      (chain)
      int64
      0 1 2 3

      array([0, 1, 2, 3])

    • draw
      (draw)
      int64
      0 1 2 3 4 5 ... 495 496 497 498 499

      array([  0,   1,   2, ..., 497, 498, 499])

    • __obs__
      (__obs__)
      int64
      0 1 2 3 4 5 ... 995 996 997 998 999

      array([  0,   1,   2, ..., 997, 998, 999])

  • Data variables: (1)
    • rt,response
      (chain, draw, __obs__)
      float64
      -0.4214 -2.193 ... -0.4149 -1.67

      array([[[-0.42139003, -2.19328821, -0.4247924 , ..., -0.42361219,
               -0.42149198, -1.65929869],
              [-0.40409852, -2.19730165, -0.40625994, ..., -0.40773632,
               -0.40402088, -1.66009839],
              [-0.43349963, -2.20508704, -0.43930199, ..., -0.43320311,
               -0.43393501, -1.67215232],
              ...,
              [-0.44052087, -2.18587642, -0.44313717, ..., -0.44333546,
               -0.44052729, -1.6514956 ],
              [-0.44052087, -2.18587642, -0.44313717, ..., -0.44333546,
               -0.44052729, -1.6514956 ],
              [-0.42286715, -2.23652404, -0.43013812, ..., -0.42122126,
               -0.42349464, -1.69169968]],
      
             [[-0.44376972, -2.18620955, -0.44793877, ..., -0.44493208,
               -0.44399309, -1.65687984],
              [-0.4692077 , -2.12723017, -0.46756308, ..., -0.47584789,
               -0.46866022, -1.60860405],
              [-0.37754136, -2.25922977, -0.38277074, ..., -0.37872102,
               -0.37784538, -1.69968631],
      ...
              [-0.43434228, -2.17965528, -0.43785337, ..., -0.43624553,
               -0.43447141, -1.65238043],
              [-0.45546712, -2.15306973, -0.45619097, ..., -0.45987963,
               -0.45523328, -1.62865005],
              [-0.45742722, -2.17744952, -0.46105205, ..., -0.4589634 ,
               -0.4575864 , -1.65007068]],
      
             [[-0.44774479, -2.17465524, -0.44729758, ..., -0.45326354,
               -0.44735372, -1.63270461],
              [-0.46895676, -2.14213147, -0.46539169, ..., -0.47734032,
               -0.46815661, -1.60525799],
              [-0.46237933, -2.16370863, -0.46291423, ..., -0.46681002,
               -0.46212869, -1.630116  ],
              ...,
              [-0.44645995, -2.17024888, -0.4468295 , ..., -0.45117137,
               -0.44617973, -1.63364187],
              [-0.46310257, -2.13269624, -0.4606867 , ..., -0.47040371,
               -0.46245549, -1.60558783],
              [-0.41489603, -2.21839133, -0.41785058, ..., -0.41748076,
               -0.41494035, -1.66985551]]])

  • Indexes: (3)
    • chain
      PandasIndex

      PandasIndex(Index([0, 1, 2, 3], dtype='int64', name='chain'))

    • draw
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             490, 491, 492, 493, 494, 495, 496, 497, 498, 499],
            dtype='int64', name='draw', length=500))

    • __obs__
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             990, 991, 992, 993, 994, 995, 996, 997, 998, 999],
            dtype='int64', name='__obs__', length=1000))

  • Attributes: (2)

    modeling_interface :

    bambi

    modeling_interface_version :

    0.14.0

  
  
• sample_stats

  <xarray.Dataset> Size: 102kB
  Dimensions:          (chain: 4, draw: 500)
  Coordinates:
    * chain            (chain) int64 32B 0 1 2 3
    * draw             (draw) int64 4kB 0 1 2 3 4 5 6 ... 494 495 496 497 498 499
  Data variables:
      acceptance_rate  (chain, draw) float64 16kB 0.2345 0.8585 ... 0.9924 0.7746
      diverging        (chain, draw) bool 2kB False False False ... False False
      energy           (chain, draw) float64 16kB 1.229e+03 ... 1.228e+03
      lp               (chain, draw) float64 16kB 1.224e+03 ... 1.225e+03
      n_steps          (chain, draw) int64 16kB 3 5 7 7 11 15 11 ... 7 3 15 7 7 15
      step_size        (chain, draw) float64 16kB 0.3328 0.3328 ... 0.3843 0.3843
      tree_depth       (chain, draw) int64 16kB 2 3 3 3 4 4 4 4 ... 3 3 2 4 3 3 4
  Attributes:
      created_at:                  2024-12-25T23:41:36.307714+00:00
      arviz_version:               0.18.0
      modeling_interface:          bambi
      modeling_interface_version:  0.14.0

  xarray.Dataset

  • Dimensions:
    • chain: 4
    • draw: 500
  • Coordinates: (2)
    • chain
      (chain)
      int64
      0 1 2 3

      array([0, 1, 2, 3])

    • draw
      (draw)
      int64
      0 1 2 3 4 5 ... 495 496 497 498 499

      array([  0,   1,   2, ..., 497, 498, 499])

  • Data variables: (7)
    • acceptance_rate
      (chain, draw)
      float64
      0.2345 0.8585 ... 0.9924 0.7746

      array([[0.23450942, 0.85848199, 0.85276862, ..., 0.98062898, 0.60006106,
              0.59725878],
             [0.92685213, 0.84771319, 0.94042334, ..., 0.84610666, 0.78661087,
              0.93835165],
             [0.92834249, 1.        , 0.97258768, ..., 0.98067234, 0.98697798,
              0.91899995],
             [0.81540984, 0.91433472, 0.87045573, ..., 0.97613743, 0.99242212,
              0.77455976]])

    • diverging
      (chain, draw)
      bool
      False False False ... False False

      array([[False, False, False, ..., False, False, False],
             [False, False, False, ..., False, False, False],
             [False, False, False, ..., False, False, False],
             [False, False, False, ..., False, False, False]])

    • energy
      (chain, draw)
      float64
      1.229e+03 1.226e+03 ... 1.228e+03

      array([[1228.76487832, 1226.44476833, 1227.28480122, ..., 1226.42017166,
              1229.8371812 , 1230.40858816],
             [1226.81911841, 1227.97091985, 1231.68786257, ..., 1225.21937204,
              1227.08832402, 1225.72182285],
             [1224.17329514, 1224.73806921, 1225.09113678, ..., 1225.17180657,
              1225.36448081, 1226.24409451],
             [1229.44661065, 1228.96640263, 1229.30147594, ..., 1227.64919361,
              1224.80466539, 1227.98024561]])

    • lp
      (chain, draw)
      float64
      1.224e+03 1.226e+03 ... 1.225e+03

      array([[1223.84039549, 1225.68008226, 1225.0313331 , ..., 1225.44762651,
              1225.44762651, 1225.84459982],
             [1224.7800351 , 1226.10508851, 1230.25358319, ..., 1224.21913793,
              1224.30849113, 1224.36157665],
             [1223.97390926, 1224.08427381, 1223.57941546, ..., 1224.39179359,
              1224.36420516, 1225.94754887],
             [1225.50591954, 1227.18586854, 1225.34584266, ..., 1224.26800277,
              1224.15963558, 1225.41349293]])

    • n_steps
      (chain, draw)
      int64
      3 5 7 7 11 15 11 ... 7 3 15 7 7 15

      array([[ 3,  5,  7, ...,  7,  3, 15],
             [31,  7,  7, ...,  7,  7,  7],
             [ 3, 15, 15, ..., 15, 15,  3],
             [ 7,  7,  3, ...,  7,  7, 15]])

    • step_size
      (chain, draw)
      float64
      0.3328 0.3328 ... 0.3843 0.3843

      array([[0.33275423, 0.33275423, 0.33275423, ..., 0.33275423, 0.33275423,
              0.33275423],
             [0.28699943, 0.28699943, 0.28699943, ..., 0.28699943, 0.28699943,
              0.28699943],
             [0.31076248, 0.31076248, 0.31076248, ..., 0.31076248, 0.31076248,
              0.31076248],
             [0.38428064, 0.38428064, 0.38428064, ..., 0.38428064, 0.38428064,
              0.38428064]])

    • tree_depth
      (chain, draw)
      int64
      2 3 3 3 4 4 4 4 ... 5 3 3 2 4 3 3 4

      array([[2, 3, 3, ..., 3, 2, 4],
             [5, 3, 3, ..., 3, 3, 3],
             [2, 4, 4, ..., 4, 4, 2],
             [3, 3, 2, ..., 3, 3, 4]])

  • Indexes: (2)
    • chain
      PandasIndex

      PandasIndex(Index([0, 1, 2, 3], dtype='int64', name='chain'))

    • draw
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             490, 491, 492, 493, 494, 495, 496, 497, 498, 499],
            dtype='int64', name='draw', length=500))

  • Attributes: (4)

    created_at :

    2024-12-25T23:41:36.307714+00:00

    arviz_version :

    0.18.0

    modeling_interface :

    bambi

    modeling_interface_version :

    0.14.0

  
  
• observed_data

  <xarray.Dataset> Size: 24kB
  Dimensions:                  (__obs__: 1000, rt,response_extra_dim_0: 2)
  Coordinates:
    * __obs__                  (__obs__) int64 8kB 0 1 2 3 4 ... 996 997 998 999
    * rt,response_extra_dim_0  (rt,response_extra_dim_0) int64 16B 0 1
  Data variables:
      rt,response              (__obs__, rt,response_extra_dim_0) float64 16kB ...
  Attributes:
      created_at:                  2024-12-25T23:41:36.308543+00:00
      arviz_version:               0.18.0
      inference_library:           numpyro
      inference_library_version:   0.15.2
      sampling_time:               43.397043
      tuning_steps:                200
      modeling_interface:          bambi
      modeling_interface_version:  0.14.0

  xarray.Dataset

  • Dimensions:
    • __obs__: 1000
    • rt,response_extra_dim_0: 2
  • Coordinates: (2)
    • __obs__
      (__obs__)
      int64
      0 1 2 3 4 5 ... 995 996 997 998 999

      array([  0,   1,   2, ..., 997, 998, 999])

    • rt,response_extra_dim_0
      (rt,response_extra_dim_0)
      int64
      0 1

      array([0, 1])

  • Data variables: (1)
    • rt,response
      (__obs__, rt,response_extra_dim_0)
      float64
      1.099 1.0 2.952 ... 1.0 2.426 1.0

      array([[1.09931123, 1.        ],
             [2.95180678, 1.        ],
             [1.14392102, 1.        ],
             ...,
             [1.06000042, 1.        ],
             [1.10505605, 1.        ],
             [2.42620206, 1.        ]])

  • Indexes: (2)
    • __obs__
      PandasIndex

      PandasIndex(Index([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,
             ...
             990, 991, 992, 993, 994, 995, 996, 997, 998, 999],
            dtype='int64', name='__obs__', length=1000))

    • rt,response_extra_dim_0
      PandasIndex

      PandasIndex(Index([0, 1], dtype='int64', name='rt,response_extra_dim_0'))

  • Attributes: (8)

    created_at :

    2024-12-25T23:41:36.308543+00:00

    arviz_version :

    0.18.0

    inference_library :

    numpyro

    inference_library_version :

    0.15.2

    sampling_time :

    43.397043

    tuning_steps :

    200

    modeling_interface :

    bambi

    modeling_interface_version :

    0.14.0

  
  

In [22]:

az.plot_trace(model.traces)
plt.tight_layout()

az.plot_pair(model.traces)
plt.tight_layout()

az.plot_trace(model.traces) plt.tight_layout() az.plot_pair(model.traces) plt.tight_layout()

We hope you find it easy to use the above example to leverage hssm to fit your own custom models.