Custom models from onnx files

Build HSSM models starting from ONNX files¶

In this tutorial we build a HSSM model directly from an onnx file. For our purposes, the onnx file-format provides nice translation layer from deep learning frameworks into a common layer from which we can then reconstruct computation graph to use through PyMC.

In [1]:

import os
import matplotlib.pyplot as plt
import numpy as np
import hssm

import os import matplotlib.pyplot as plt import numpy as np import hssm

Loading the network¶

In [2]:

# Networks
network_path = os.path.join("data", "race_3_no_bias_lan_no_batch.onnx")

# Networks network_path = os.path.join("data", "race_3_no_bias_lan_no_batch.onnx")

The network we load here does not have dynamic input dimensions, which prevents us from batching computations.

Instead of fixing things behind the scenes and loading a fixed network, we provide a useful snippet below that shows how to rectify this situation.

In [3]:

import onnx
import onnxruntime as ort

# Load model from path
onnx_model = onnx.load(network_path)

# Change input and output dimensions to be dynamic to allow for batching
# (in case this is not already done)
for input_tensor in onnx_model.graph.input:
    dim_proto = input_tensor.type.tensor_type.shape.dim[0]
    if not dim_proto.dim_param == "None":
        dim_proto.dim_param = "None"

for output_tensor in onnx_model.graph.output:
    dim_proto = output_tensor.type.tensor_type.shape.dim[0]
    if not dim_proto.dim_param == "None":
        dim_proto.dim_param = "None"

input_name = onnx_model.graph.input[0].name

# Please uncomment the below line to save the adjusted model
# onnx.save(onnx_model, "test_files/race_3_no_bias_lan_batch.onnx")

import onnx import onnxruntime as ort # Load model from path onnx_model = onnx.load(network_path) # Change input and output dimensions to be dynamic to allow for batching # (in case this is not already done) for input_tensor in onnx_model.graph.input: dim_proto = input_tensor.type.tensor_type.shape.dim[0] if not dim_proto.dim_param == "None": dim_proto.dim_param = "None" for output_tensor in onnx_model.graph.output: dim_proto = output_tensor.type.tensor_type.shape.dim[0] if not dim_proto.dim_param == "None": dim_proto.dim_param = "None" input_name = onnx_model.graph.input[0].name # Please uncomment the below line to save the adjusted model # onnx.save(onnx_model, "test_files/race_3_no_bias_lan_batch.onnx")

Armed with the corrected network, let's test inference speed on a data-batch of $1000$ trials.

In [4]:

# Load model batch ready model
ort_session = ort.InferenceSession("data/race_3_no_bias_lan_batch.onnx")

# Test inference speed
import time

start = time.time()
for i in range(100):
    ort_session.run(
        None, {input_name: np.random.uniform(size=(1000, 8)).astype(np.float32)}
    )
end = time.time()
print(f"Time taken: {(end - start) / 100} seconds")

# Load model batch ready model ort_session = ort.InferenceSession("data/race_3_no_bias_lan_batch.onnx") # Test inference speed import time start = time.time() for i in range(100): ort_session.run( None, {input_name: np.random.uniform(size=(1000, 8)).astype(np.float32)} ) end = time.time() print(f"Time taken: {(end - start) / 100} seconds")

Time taken: 0.00040030956268310546 seconds

Defining the Likelihood¶

The network we loaded corresponds to a LAN, for a Race model with three choice alternatives. This model has three drift parameters v0, v1, v2, a boundary parameter a, a starting point bias z and a non-decision-time t.

Data from this model has the usual rt, choice format.

We use this to construct a simple blackbox likelihood function below. This likelihood function takes the respective data and model parameters as arguments. The function body shapes these input arguments into a matrix and performs a batched forward pass through the loaded network via the onnx.runtime.

In [5]:

def my_blackbox_race_model(data, v0, v1, v2, a, t, z):
    """Calculate log-likelihood for a 3-choice race model.

    Parameters
    ----------
    data : np.ndarray
        Array of shape (n_trials, 2) containing response times in first column
        and choices (0, 1, or 2) in second column
    v0 : float
        Drift rate for accumulator 0
    v1 : float
        Drift rate for accumulator 1
    v2 : float
        Drift rate for accumulator 2
    a : float
        Decision threshold/boundary
    t : float
        Non-decision time
    z : float
        Starting point bias

    Returns
    -------
    np.ndarray
        Array of log-likelihood values for each trial
    """
    data_nrows = data.shape[0]
    data = np.vstack(
        [np.full(data_nrows, param_) for param_ in [v0, v1, v2, a, t, z]]
        + [data[:, 0], data[:, 1]]
    ).T.astype(np.float32)
    return ort_session.run(None, {input_name: data})[0].squeeze()

def my_blackbox_race_model(data, v0, v1, v2, a, t, z): """Calculate log-likelihood for a 3-choice race model. Parameters ---------- data : np.ndarray Array of shape (n_trials, 2) containing response times in first column and choices (0, 1, or 2) in second column v0 : float Drift rate for accumulator 0 v1 : float Drift rate for accumulator 1 v2 : float Drift rate for accumulator 2 a : float Decision threshold/boundary t : float Non-decision time z : float Starting point bias Returns ------- np.ndarray Array of log-likelihood values for each trial """ data_nrows = data.shape[0] data = np.vstack( [np.full(data_nrows, param_) for param_ in [v0, v1, v2, a, t, z]] + [data[:, 0], data[:, 1]] ).T.astype(np.float32) return ort_session.run(None, {input_name: data})[0].squeeze()

Simulate example data¶

In [6]:

# Set parameters
v0 = 1.0
v1 = 0.5
v2 = 0.25
a = 1.5
t = 0.3
z = 0.5

# simulate some data from the model
obs_race3 = hssm.simulate_data(
    theta=dict(v0=v0, v1=v1, v2=v2, a=a, t=t, z=z), model="race_no_bias_3", size=1000
)

# Set parameters v0 = 1.0 v1 = 0.5 v2 = 0.25 a = 1.5 t = 0.3 z = 0.5 # simulate some data from the model obs_race3 = hssm.simulate_data( theta=dict(v0=v0, v1=v1, v2=v2, a=a, t=t, z=z), model="race_no_bias_3", size=1000 )

Test Likelihood Outputs¶

In [7]:

# Test that outputs are reasonable
for choice in [0, 1, 2]:
    rts = np.linspace(0, 20, 1000)
    choices = np.repeat(choice, 1000)

    data = np.vstack([rts, choices]).T
    out = my_blackbox_race_model(data, v0, v1, v2, a, t, z)

    plt.plot(rts, np.exp(out), label=f"choice: {choice}")
plt.legend()
plt.show()

# Test that outputs are reasonable for choice in [0, 1, 2]: rts = np.linspace(0, 20, 1000) choices = np.repeat(choice, 1000) data = np.vstack([rts, choices]).T out = my_blackbox_race_model(data, v0, v1, v2, a, t, z) plt.plot(rts, np.exp(out), label=f"choice: {choice}") plt.legend() plt.show()

Build HSSM Model¶

We can now build a simple HSSM model that takes in our new blackbox likelihood.

In [8]:

model = hssm.HSSM(
    data=obs_race3,
    model="race_no_bias_3",  # some name for the model
    model_config={
        "list_params": ["v0", "v1", "v2", "a", "z", "t"],
        "bounds": {
            "v0": (0.0, 2.5),
            "v1": (0.0, 2.5),
            "v2": (0.0, 2.5),
            "a": (1.0, 3.0),
            "z": (0.0, 0.9),
            "t": (0.001, 2),
        },
    },  # minimal specification of model parameters and parameter bounds
    loglik_kind="blackbox",  # use the blackbox loglik
    loglik=my_blackbox_race_model,
    choices=[0, 1, 2],  # list the legal choice options
    z=0.5,
    p_outlier=0,
)

model = hssm.HSSM( data=obs_race3, model="race_no_bias_3", # some name for the model model_config={ "list_params": ["v0", "v1", "v2", "a", "z", "t"], "bounds": { "v0": (0.0, 2.5), "v1": (0.0, 2.5), "v2": (0.0, 2.5), "a": (1.0, 3.0), "z": (0.0, 0.9), "t": (0.001, 2), }, }, # minimal specification of model parameters and parameter bounds loglik_kind="blackbox", # use the blackbox loglik loglik=my_blackbox_race_model, choices=[0, 1, 2], # list the legal choice options z=0.5, p_outlier=0, )

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.

In [9]:

model.graph()

model.graph()

Out[9]:

In [10]:

model.sample(draws=500, tune=200, discard_tuned_samples=False)

model.sample(draws=500, tune=200, discard_tuned_samples=False)

Using default initvals. 

Multiprocess sampling (4 chains in 4 jobs)
CompoundStep
>Slice: [v1]
>Slice: [t]
>Slice: [a]
>Slice: [v2]
>Slice: [v0]

Output()

Sampling 4 chains for 200 tune and 500 draw iterations (800 + 2_000 draws total) took 33 seconds.
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
The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details
/Users/afengler/Library/CloudStorage/OneDrive-Personal/proj_hssm/HSSM/.venv/lib/python3.12/site-packages/pymc/pytensorf.py:958: FutureWarning: compile_pymc was renamed to compile. Old name will be removed in a future release of PyMC
  warnings.warn(
100%|██████████| 2000/2000 [00:01<00:00, 1959.13it/s]

Out[10]:

arviz.InferenceData

• posterior

  <xarray.Dataset> Size: 84kB
  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:
      a        (chain, draw) float64 16kB 1.406 1.409 1.413 ... 1.402 1.465 1.494
      v1       (chain, draw) float64 16kB 0.4684 0.4786 0.4823 ... 0.4551 0.5755
      t        (chain, draw) float64 16kB 0.301 0.3007 0.3026 ... 0.2979 0.2961
      v0       (chain, draw) float64 16kB 0.9083 0.9667 0.9723 ... 1.015 1.008
      v2       (chain, draw) float64 16kB 0.209 0.1775 0.1519 ... 0.3455 0.2646
  Attributes:
      created_at:                  2025-07-13T13:38:55.187187+00:00
      arviz_version:               0.21.0
      inference_library:           pymc
      inference_library_version:   5.21.1
      sampling_time:               33.21286725997925
      tuning_steps:                200
      modeling_interface:          bambi
      modeling_interface_version:  0.15.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: (5)
    • a
      (chain, draw)
      float64
      1.406 1.409 1.413 ... 1.465 1.494

      array([[1.40592922, 1.40878083, 1.4126126 , ..., 1.39217642, 1.38674931,
              1.41752162],
             [1.39709415, 1.44401999, 1.45256968, ..., 1.45356977, 1.42172288,
              1.39467821],
             [1.47369161, 1.44055681, 1.46670984, ..., 1.39643025, 1.39297555,
              1.42064556],
             [1.47043019, 1.4720532 , 1.49375593, ..., 1.4022729 , 1.46485852,
              1.49432955]])

    • v1
      (chain, draw)
      float64
      0.4684 0.4786 ... 0.4551 0.5755

      array([[0.46838594, 0.47862572, 0.48233697, ..., 0.33207914, 0.30854439,
              0.3526878 ],
             [0.37124585, 0.47382578, 0.60133901, ..., 0.51397384, 0.50214877,
              0.5468659 ],
             [0.53510531, 0.56360178, 0.54176941, ..., 0.547189  , 0.40957048,
              0.47578033],
             [0.49381614, 0.47806852, 0.55013191, ..., 0.45294303, 0.45514729,
              0.57549072]])

    • t
      (chain, draw)
      float64
      0.301 0.3007 ... 0.2979 0.2961

      array([[0.30100565, 0.30073449, 0.3025787 , ..., 0.30115872, 0.30243843,
              0.30423456],
             [0.30343069, 0.3012835 , 0.29796201, ..., 0.29758497, 0.30059142,
              0.30375849],
             [0.29995877, 0.29980727, 0.29776967, ..., 0.30102802, 0.30211988,
              0.3025669 ],
             [0.29898282, 0.29610292, 0.29361303, ..., 0.3056978 , 0.29786651,
              0.29610822]])

    • v0
      (chain, draw)
      float64
      0.9083 0.9667 ... 1.015 1.008

      array([[0.90825872, 0.96674947, 0.97226234, ..., 0.97220612, 0.88010312,
              1.00846504],
             [0.92722341, 1.04701424, 0.97407736, ..., 0.97041897, 0.99896929,
              0.97225834],
             [1.0762005 , 1.01456578, 0.95441441, ..., 0.87713345, 0.85571168,
              0.9368968 ],
             [1.01039245, 0.98179051, 1.02215291, ..., 0.92771278, 1.01502695,
              1.00828603]])

    • v2
      (chain, draw)
      float64
      0.209 0.1775 ... 0.3455 0.2646

      array([[0.2089603 , 0.17753116, 0.15189631, ..., 0.12758559, 0.23759116,
              0.28398321],
             [0.24977196, 0.29892551, 0.31946899, ..., 0.27305971, 0.15043639,
              0.1502101 ],
             [0.32754931, 0.2487011 , 0.32310058, ..., 0.1902019 , 0.2024323 ,
              0.24206775],
             [0.32348696, 0.35749181, 0.25530057, ..., 0.30722137, 0.34546405,
              0.26463333]])

  • 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 :

    2025-07-13T13:38:55.187187+00:00

    arviz_version :

    0.21.0

    inference_library :

    pymc

    inference_library_version :

    5.21.1

    sampling_time :

    33.21286725997925

    tuning_steps :

    200

    modeling_interface :

    bambi

    modeling_interface_version :

    0.15.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.3937 0.2195 ... -0.7151
  Attributes:
      modeling_interface:          bambi
      modeling_interface_version:  0.15.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.3937 0.2195 ... -0.8749 -0.7151

      array([[[-0.39374053,  0.21951646, -0.53969848, ..., -1.83415735,
               -0.88392496, -0.71343857],
              [-0.42444187,  0.25382274, -0.57551467, ..., -1.85065293,
               -0.92410833, -0.75144583],
              [-0.46547714,  0.25718734, -0.58723229, ..., -1.82965815,
               -0.93328023, -0.76147431],
              ...,
              [-0.42608628,  0.28125912, -0.5907101 , ..., -1.92730331,
               -0.93697762, -0.76597339],
              [-0.33579874,  0.22642225, -0.48783487, ..., -1.92908299,
               -0.82787901, -0.65988976],
              [-0.38695836,  0.27731842, -0.47762096, ..., -1.95824277,
               -0.81656951, -0.64778954]],
      
             [[-0.35861745,  0.24389033, -0.49160874, ..., -1.91298985,
               -0.83486086, -0.66481972],
              [-0.41623402,  0.27065155, -0.50259584, ..., -1.90361488,
               -0.8414681 , -0.67221379],
              [-0.39522213,  0.20771585, -0.50606942, ..., -1.8097527 ,
               -0.84536338, -0.67607421],
      ...
              [-0.38833329,  0.20276956, -0.55826485, ..., -1.78601408,
               -0.91050684, -0.73644817],
              [-0.37464413,  0.2021679 , -0.52175695, ..., -1.84044933,
               -0.86304373, -0.69428593],
              [-0.41310447,  0.22505434, -0.51611173, ..., -1.83577859,
               -0.8554734 , -0.68662417]],
      
             [[-0.43825677,  0.22388588, -0.49159843, ..., -1.85346258,
               -0.81571829, -0.65307504],
              [-0.39255625,  0.20257479, -0.47224408, ..., -1.86702132,
               -0.79419845, -0.63296849],
              [-0.50539142,  0.21126784, -0.57323986, ..., -1.79582405,
               -0.89189482, -0.73189425],
              ...,
              [-0.34518483,  0.23370288, -0.45298988, ..., -1.88167357,
               -0.79963422, -0.627424  ],
              [-0.39973909,  0.23183194, -0.47649533, ..., -1.89705539,
               -0.80251014, -0.63929474],
              [-0.51802039,  0.20058046, -0.55745983, ..., -1.76851726,
               -0.87492377, -0.71512091]]])

  • 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.15.0

  
  
• sample_stats

  <xarray.Dataset> Size: 164kB
  Dimensions:          (chain: 4, draw: 500, nstep_in_dim_0: 5, nstep_out_dim_0: 5)
  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
    * nstep_in_dim_0   (nstep_in_dim_0) int64 40B 0 1 2 3 4
    * nstep_out_dim_0  (nstep_out_dim_0) int64 40B 0 1 2 3 4
  Data variables:
      nstep_in         (chain, draw, nstep_in_dim_0) int64 80kB 5 8 5 0 ... 7 0 3
      nstep_out        (chain, draw, nstep_out_dim_0) int64 80kB 0 1 0 1 ... 0 1 0
  Attributes:
      created_at:                  2025-07-13T13:38:55.199850+00:00
      arviz_version:               0.21.0
      inference_library:           pymc
      inference_library_version:   5.21.1
      sampling_time:               33.21286725997925
      tuning_steps:                200
      modeling_interface:          bambi
      modeling_interface_version:  0.15.0

  xarray.Dataset

  • Dimensions:
    • chain: 4
    • draw: 500
    • nstep_in_dim_0: 5
    • nstep_out_dim_0: 5
  • Coordinates: (4)
    • 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])

    • nstep_in_dim_0
      (nstep_in_dim_0)
      int64
      0 1 2 3 4

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

    • nstep_out_dim_0
      (nstep_out_dim_0)
      int64
      0 1 2 3 4

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

  • Data variables: (2)
    • nstep_in
      (chain, draw, nstep_in_dim_0)
      int64
      5 8 5 0 5 6 1 2 ... 5 2 3 3 2 7 0 3

      array([[[ 5,  8,  5,  0,  5],
              [ 6,  1,  2,  0,  7],
              [10,  2, 10,  1,  6],
              ...,
              [ 6,  2,  3,  0,  2],
              [ 3,  1,  3,  1,  7],
              [ 6,  6,  2,  0,  3]],
      
             [[ 3,  2,  4,  4,  5],
              [ 4,  2,  8,  3,  2],
              [ 1,  1,  3,  2,  0],
              ...,
              [ 7,  6,  4,  3,  0],
              [ 4,  5,  7,  1,  3],
              [ 7,  3,  4,  0,  2]],
      
             [[ 7,  3,  2,  4,  2],
              [ 3,  3,  2,  1,  0],
              [ 1,  2,  3,  0,  0],
              ...,
              [ 2,  4,  0,  3,  0],
              [ 5,  5,  7,  5,  0],
              [ 8,  1,  0,  1,  0]],
      
             [[ 5,  1,  2,  5,  8],
              [ 0,  2,  6,  5,  5],
              [ 3,  2,  3,  7,  8],
              ...,
              [ 7,  1,  4,  4,  9],
              [ 8,  0,  5,  2,  3],
              [ 3,  2,  7,  0,  3]]])

    • nstep_out
      (chain, draw, nstep_out_dim_0)
      int64
      0 1 0 1 1 0 0 0 ... 0 0 0 0 0 0 1 0

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

  • Indexes: (4)
    • 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))

    • nstep_in_dim_0
      PandasIndex

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

    • nstep_out_dim_0
      PandasIndex

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

  • Attributes: (8)

    created_at :

    2025-07-13T13:38:55.199850+00:00

    arviz_version :

    0.21.0

    inference_library :

    pymc

    inference_library_version :

    5.21.1

    sampling_time :

    33.21286725997925

    tuning_steps :

    200

    modeling_interface :

    bambi

    modeling_interface_version :

    0.15.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:                  2025-07-13T13:38:55.201998+00:00
      arviz_version:               0.21.0
      inference_library:           pymc
      inference_library_version:   5.21.1
      modeling_interface:          bambi
      modeling_interface_version:  0.15.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
      0.422 2.0 0.4692 ... 2.0 0.5807 2.0

      array([[0.42201638, 2.        ],
             [0.46916589, 0.        ],
             [0.54321641, 2.        ],
             ...,
             [0.88358653, 1.        ],
             [0.61587334, 2.        ],
             [0.5806542 , 2.        ]])

  • 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: (6)

    created_at :

    2025-07-13T13:38:55.201998+00:00

    arviz_version :

    0.21.0

    inference_library :

    pymc

    inference_library_version :

    5.21.1

    modeling_interface :

    bambi

    modeling_interface_version :

    0.15.0

  
  
• warmup_posterior

  <xarray.Dataset> Size: 6MB
  Dimensions:  (chain: 4, draw: 200, __obs__: 1000)
  Coordinates:
    * chain    (chain) int64 32B 0 1 2 3
    * draw     (draw) int64 2kB 0 1 2 3 4 5 6 7 ... 193 194 195 196 197 198 199
    * __obs__  (__obs__) int64 8kB 0 1 2 3 4 5 6 7 ... 993 994 995 996 997 998 999
  Data variables:
      v1       (chain, draw) float64 6kB 1.466e-10 2.5 2.5 ... 0.3742 0.621 0.5092
      t        (chain, draw) float64 6kB 0.1457 0.1819 0.1776 ... 0.3038 0.296
      a        (chain, draw) float64 6kB 3.0 2.976 2.977 ... 1.421 1.448 1.441
      v2       (chain, draw) float64 6kB 1.329 1.675 1.623 ... 0.2861 0.2813
      v0       (chain, draw) float64 6kB 2.5 2.496 2.496 ... 1.105 0.9855 1.004
      v0_mean  (chain, draw, __obs__) float64 6MB 2.5 2.5 2.5 ... 1.004 1.004
  Attributes:
      created_at:                  2025-07-13T13:38:55.197877+00:00
      arviz_version:               0.21.0
      inference_library:           pymc
      inference_library_version:   5.21.1
      sampling_time:               33.21286725997925
      tuning_steps:                200
      modeling_interface:          bambi
      modeling_interface_version:  0.15.0

  xarray.Dataset

  • Dimensions:
    • chain: 4
    • draw: 200
    • __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 ... 195 196 197 198 199

      array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,
              14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,
              28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,
              42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,  53,  54,  55,
              56,  57,  58,  59,  60,  61,  62,  63,  64,  65,  66,  67,  68,  69,
              70,  71,  72,  73,  74,  75,  76,  77,  78,  79,  80,  81,  82,  83,
              84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,  95,  96,  97,
              98,  99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
             112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125,
             126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139,
             140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153,
             154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167,
             168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181,
             182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195,
             196, 197, 198, 199])

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

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

  • Data variables: (6)
    • v1
      (chain, draw)
      float64
      1.466e-10 2.5 2.5 ... 0.621 0.5092

      array([[1.46592326e-010, 2.50000000e+000, 2.50000000e+000,
              2.50000000e+000, 2.50000000e+000, 2.50000000e+000,
              2.24900668e+000, 2.09164662e+000, 1.52192846e+000,
              1.43206836e+000, 1.51809942e+000, 1.58278396e+000,
              1.46705975e+000, 1.60313028e+000, 1.62891141e+000,
              1.42289363e+000, 1.35120204e+000, 1.44115721e+000,
              1.52512323e+000, 1.47427728e+000, 1.41245995e+000,
              1.41886438e+000, 1.25712459e+000, 1.13693053e+000,
              1.07030777e+000, 1.06783213e+000, 1.02520733e+000,
              1.06320975e+000, 9.40999672e-001, 1.08632774e+000,
              9.72274473e-001, 1.02882854e+000, 8.46968680e-001,
              9.74959671e-001, 9.10033782e-001, 9.26834385e-001,
              7.97180353e-001, 7.98778765e-001, 8.84149896e-001,
              8.04648207e-001, 7.01882588e-001, 9.02686590e-001,
              7.92392684e-001, 7.61299748e-001, 7.13738415e-001,
              6.53841708e-001, 5.76836869e-001, 6.29574168e-001,
              4.86271297e-001, 4.61191317e-001, 4.10709425e-001,
              4.88127924e-001, 5.00072704e-001, 5.07209262e-001,
              6.50707215e-001, 4.92966710e-001, 5.31507283e-001,
              5.30714785e-001, 5.81164040e-001, 5.62607257e-001,
      ...
              4.99706645e-001, 6.19569011e-001, 6.27552497e-001,
              6.10856547e-001, 5.66780181e-001, 5.57528315e-001,
              5.77869385e-001, 5.54483960e-001, 5.93402646e-001,
              4.05793988e-001, 3.94031575e-001, 3.84307764e-001,
              4.49558055e-001, 2.09866537e-001, 2.39045317e-001,
              2.38509169e-001, 2.29504582e-001, 3.24173235e-001,
              3.60948665e-001, 4.32983327e-001, 4.54080035e-001,
              3.73031380e-001, 4.24424794e-001, 3.78668980e-001,
              4.68606740e-001, 4.36072913e-001, 5.02393675e-001,
              4.19513241e-001, 4.51794177e-001, 4.44429527e-001,
              5.64919166e-001, 3.24749182e-001, 2.47105400e-001,
              5.13889016e-001, 4.39581280e-001, 4.37688194e-001,
              3.52812204e-001, 4.14136990e-001, 3.10042943e-001,
              3.08596719e-001, 4.73501386e-001, 4.87908904e-001,
              2.30044842e-001, 2.96439584e-001, 3.23043737e-001,
              3.63779287e-001, 4.42695972e-001, 4.91769334e-001,
              4.03223917e-001, 3.74969831e-001, 4.40426330e-001,
              3.98148027e-001, 4.66667323e-001, 3.83957801e-001,
              4.75725520e-001, 3.03000864e-001, 3.74169943e-001,
              6.21047416e-001, 5.09198439e-001]])

    • t
      (chain, draw)
      float64
      0.1457 0.1819 ... 0.3038 0.296

      array([[0.14567973, 0.18185052, 0.17762294, 0.18250455, 0.17742804,
              0.16624632, 0.15861905, 0.14934012, 0.10015518, 0.12496131,
              0.12746177, 0.12230423, 0.11970694, 0.12480482, 0.12793813,
              0.12288672, 0.12423736, 0.1250974 , 0.1592832 , 0.16700779,
              0.17452206, 0.18268314, 0.18361008, 0.19436443, 0.20673044,
              0.20956627, 0.21301875, 0.22834882, 0.2289461 , 0.24387931,
              0.2455111 , 0.24506353, 0.24263996, 0.25137075, 0.25452805,
              0.25739486, 0.25211814, 0.251891  , 0.26228786, 0.26319828,
              0.27031585, 0.27135501, 0.2647415 , 0.27309746, 0.2768214 ,
              0.28046885, 0.29077341, 0.29608097, 0.2934577 , 0.29709224,
              0.30145111, 0.30204331, 0.29923445, 0.29302343, 0.2979218 ,
              0.29501142, 0.29459301, 0.29752291, 0.29850647, 0.29780951,
              0.29619386, 0.28481697, 0.28459204, 0.28329845, 0.29063131,
              0.29348508, 0.29769929, 0.29293389, 0.29362501, 0.29631142,
              0.2952969 , 0.29172302, 0.29183735, 0.28951248, 0.29317939,
              0.28814895, 0.28829503, 0.28530729, 0.28790789, 0.28311111,
              0.29086217, 0.28746223, 0.28695876, 0.28705263, 0.28663121,
              0.28696889, 0.28303905, 0.28594928, 0.28153771, 0.28347927,
              0.28357304, 0.29080421, 0.29274023, 0.29056162, 0.29017326,
              0.29289912, 0.28911539, 0.28864402, 0.28606487, 0.28907231,
      ...
              0.29843765, 0.29944538, 0.29053936, 0.28870007, 0.28904386,
              0.28675078, 0.28810143, 0.28783698, 0.29422015, 0.28397951,
              0.29326371, 0.2929658 , 0.29369153, 0.29388065, 0.29459302,
              0.29208706, 0.2874413 , 0.28642243, 0.28721733, 0.29025942,
              0.29237698, 0.29570721, 0.30365935, 0.30195619, 0.29933542,
              0.29903808, 0.303914  , 0.30418176, 0.30385864, 0.30340484,
              0.30897209, 0.31118226, 0.29635893, 0.29805258, 0.29667152,
              0.29969185, 0.29180699, 0.29260137, 0.29850236, 0.28948553,
              0.29000765, 0.29019312, 0.28919058, 0.29425817, 0.29488917,
              0.29544942, 0.30255512, 0.29718746, 0.30064008, 0.29676368,
              0.3007215 , 0.30539339, 0.30507129, 0.30688089, 0.30729585,
              0.31162224, 0.31002343, 0.30929432, 0.30708735, 0.3046742 ,
              0.30706261, 0.30836568, 0.30832625, 0.30027322, 0.3038854 ,
              0.30473839, 0.30354129, 0.30456051, 0.30493058, 0.30433777,
              0.30131255, 0.30065976, 0.30102153, 0.30131775, 0.30168051,
              0.30458177, 0.30398034, 0.3012993 , 0.30372349, 0.30166351,
              0.30309797, 0.30624912, 0.30722361, 0.30900929, 0.30618897,
              0.3070792 , 0.30140133, 0.30076078, 0.30469137, 0.3042356 ,
              0.29916592, 0.29941708, 0.30035261, 0.30220539, 0.30443424,
              0.30510941, 0.30534551, 0.30271836, 0.30381098, 0.29600158]])

    • a
      (chain, draw)
      float64
      3.0 2.976 2.977 ... 1.448 1.441

      array([[2.99999928, 2.97641696, 2.97697299, 2.99746713, 2.9886853 ,
              2.98841693, 2.99434794, 2.99019214, 2.98672895, 2.94603009,
              2.95195859, 2.97208615, 2.97347098, 2.98659871, 2.94234331,
              2.86550301, 2.77972476, 2.8061563 , 2.58740361, 2.55420283,
              2.55691318, 2.49088119, 2.37545764, 2.25174679, 2.2250386 ,
              2.1720471 , 2.11802943, 2.047751  , 2.0572807 , 1.96488367,
              1.99196278, 1.98962343, 1.9605559 , 1.87988572, 1.80953651,
              1.84243273, 1.81826849, 1.81486999, 1.78840648, 1.75650727,
              1.75679629, 1.75095981, 1.73432022, 1.66273405, 1.64963244,
              1.58128422, 1.51051073, 1.47415802, 1.46416199, 1.41094745,
              1.41500893, 1.43161849, 1.45496688, 1.45541351, 1.48376118,
              1.4722175 , 1.45853359, 1.50044044, 1.47437905, 1.44252655,
              1.54898649, 1.63322053, 1.59992671, 1.52080525, 1.52130741,
              1.50854994, 1.46066222, 1.4765174 , 1.48880973, 1.48949878,
              1.50796066, 1.52655718, 1.5460633 , 1.5351502 , 1.54482565,
              1.55774463, 1.5592482 , 1.59637311, 1.58285187, 1.57692787,
              1.56969571, 1.5231698 , 1.56858387, 1.60307305, 1.55728644,
              1.56566866, 1.60881231, 1.6185408 , 1.61551454, 1.6170624 ,
              1.59665339, 1.6026201 , 1.5694914 , 1.53578215, 1.49285749,
              1.52574716, 1.56851579, 1.57518993, 1.5817324 , 1.57867555,
      ...
              1.41916907, 1.48656215, 1.51837595, 1.54441799, 1.58912872,
              1.56775653, 1.55994446, 1.54449416, 1.55019705, 1.55163794,
              1.54357504, 1.49726384, 1.48806439, 1.49877893, 1.51666691,
              1.54466924, 1.57720969, 1.56050386, 1.48828905, 1.50156902,
              1.43191406, 1.40084626, 1.41413126, 1.40540455, 1.40778408,
              1.40858896, 1.38596012, 1.39516406, 1.3893916 , 1.39042343,
              1.37224854, 1.38408927, 1.42575492, 1.44615451, 1.48866832,
              1.48397362, 1.49194974, 1.47997864, 1.49198113, 1.49278935,
              1.53007083, 1.5032683 , 1.50207027, 1.48894354, 1.49524973,
              1.49518487, 1.48038184, 1.4744651 , 1.45098254, 1.4464252 ,
              1.38671369, 1.3834296 , 1.36807613, 1.30665958, 1.30680987,
              1.29206324, 1.29427088, 1.28291465, 1.3443003 , 1.36421896,
              1.34308647, 1.32637139, 1.39681769, 1.39044516, 1.38212275,
              1.39629227, 1.39368587, 1.39898382, 1.3847024 , 1.41795005,
              1.41473731, 1.39511562, 1.38147483, 1.37061751, 1.36892064,
              1.38383594, 1.40100275, 1.41718503, 1.39562343, 1.3880273 ,
              1.36351648, 1.34879123, 1.34055817, 1.3353286 , 1.3214887 ,
              1.38658008, 1.38469638, 1.37601192, 1.36049244, 1.37794238,
              1.44840997, 1.42206726, 1.41898359, 1.40505616, 1.40566666,
              1.36405573, 1.38617731, 1.42087491, 1.4478237 , 1.44094658]])

    • v2
      (chain, draw)
      float64
      1.329 1.675 1.623 ... 0.2861 0.2813

      array([[1.32890849e+00, 1.67480789e+00, 1.62299650e+00, 1.59780413e+00,
              1.59981405e+00, 1.53186363e+00, 1.55469778e+00, 1.25762572e+00,
              1.29194402e+00, 1.42097131e+00, 1.29739062e+00, 1.32770831e+00,
              1.43913411e+00, 1.45033132e+00, 1.33595884e+00, 1.12599021e+00,
              1.10000274e+00, 1.29066390e+00, 1.02940518e+00, 1.11938009e+00,
              1.11675193e+00, 1.00610844e+00, 9.98826748e-01, 9.33068461e-01,
              8.82625510e-01, 8.30029500e-01, 7.69626894e-01, 7.94889394e-01,
              8.46960578e-01, 8.48890174e-01, 8.61567055e-01, 7.96717709e-01,
              8.39581809e-01, 6.35962518e-01, 5.73653541e-01, 5.44106324e-01,
              6.59192036e-01, 6.33163766e-01, 6.39448723e-01, 6.43143554e-01,
              6.05951679e-01, 5.20440161e-01, 6.43847172e-01, 2.90208466e-01,
              3.56499113e-01, 4.28127049e-01, 2.84952778e-01, 3.45642335e-01,
              1.93269897e-01, 1.77963298e-01, 2.49374966e-01, 3.08969247e-01,
              1.72884397e-01, 2.23217906e-01, 1.95206211e-01, 1.42034637e-01,
              3.73771479e-01, 3.68069948e-01, 3.47151315e-01, 2.45018092e-01,
              5.65123867e-01, 4.82484172e-01, 3.13221894e-01, 2.35570747e-01,
              3.45756932e-01, 3.33849634e-01, 2.16936088e-01, 3.68273934e-01,
              2.69883466e-01, 3.58994042e-01, 3.47562420e-01, 3.40250027e-01,
              3.76667626e-01, 3.14356487e-01, 3.48795157e-01, 2.03070553e-01,
              4.65790545e-01, 4.94086442e-01, 3.50868379e-01, 4.64701675e-01,
      ...
              2.46582918e-01, 2.52737502e-01, 2.10860520e-01, 2.04777577e-01,
              1.35309109e-01, 2.50922133e-01, 1.56214067e-01, 2.58324031e-01,
              1.77199570e-01, 1.75851181e-01, 3.08590161e-01, 2.19635089e-01,
              2.41638998e-01, 2.33729516e-01, 2.91495867e-01, 3.74234826e-01,
              2.46785139e-01, 2.55600419e-01, 2.14621979e-01, 3.50901026e-01,
              3.56295608e-01, 2.86189637e-01, 3.13586913e-01, 3.17486348e-01,
              3.65548709e-01, 4.54923639e-01, 2.96737600e-01, 2.69306052e-01,
              1.67902870e-01, 2.15493771e-01, 1.79143272e-01, 1.92908718e-01,
              1.38691346e-01, 1.10922754e-01, 1.00967034e-01, 8.56309425e-02,
              1.31968599e-02, 7.22610307e-03, 1.94648382e-03, 3.23618531e-04,
              3.89612343e-02, 1.49906073e-01, 1.77688318e-01, 1.52061151e-01,
              2.60941144e-01, 2.09322340e-01, 2.43088372e-01, 1.41621371e-01,
              2.23602876e-01, 2.10705301e-01, 1.51173207e-01, 1.31829218e-01,
              9.83780843e-02, 1.11840933e-01, 1.73976633e-01, 1.83966292e-01,
              2.51593380e-01, 1.51205202e-01, 2.07557401e-01, 1.21463399e-01,
              2.07492996e-01, 1.46648008e-01, 3.70167842e-02, 2.75373401e-02,
              1.61534106e-01, 1.55497369e-01, 8.59035696e-02, 1.43736534e-01,
              1.47474066e-01, 1.38128569e-01, 2.89725669e-01, 2.38591270e-01,
              1.88901696e-01, 1.73841933e-01, 1.69686663e-01, 2.02671924e-01,
              2.87498491e-01, 2.39159901e-01, 2.86086510e-01, 2.81288332e-01]])

    • v0
      (chain, draw)
      float64
      2.5 2.496 2.496 ... 0.9855 1.004

      array([[2.49994623e+00, 2.49643190e+00, 2.49621832e+00, 2.48474775e+00,
              2.35169403e+00, 2.34283167e+00, 2.21333691e+00, 2.02800896e+00,
              2.02085304e+00, 2.06737280e+00, 1.93163640e+00, 1.84793161e+00,
              1.91453961e+00, 1.95557589e+00, 1.89680672e+00, 1.84385387e+00,
              1.83311108e+00, 1.79943340e+00, 1.66439461e+00, 1.74983764e+00,
              1.77489389e+00, 1.75682271e+00, 1.51191710e+00, 1.59389496e+00,
              1.62746771e+00, 1.38981834e+00, 1.51941361e+00, 1.31262510e+00,
              1.61380524e+00, 1.61604199e+00, 1.42020181e+00, 1.46064315e+00,
              1.37021882e+00, 1.32749271e+00, 1.32785833e+00, 1.31398957e+00,
              1.25743511e+00, 1.29620175e+00, 1.27775976e+00, 1.31498427e+00,
              1.29493794e+00, 1.28949029e+00, 1.14263610e+00, 1.21556268e+00,
              1.19783246e+00, 1.07859244e+00, 1.05904506e+00, 9.61735639e-01,
              9.20504974e-01, 9.40866364e-01, 9.88978790e-01, 1.00108626e+00,
              9.76765068e-01, 9.47923949e-01, 1.10915123e+00, 1.05350432e+00,
              1.04810765e+00, 1.11498137e+00, 9.08956869e-01, 1.04362904e+00,
              1.15001257e+00, 1.22034198e+00, 1.04697450e+00, 1.03912075e+00,
              1.07967501e+00, 1.06907184e+00, 8.73613454e-01, 9.39843816e-01,
              9.33547983e-01, 1.04599111e+00, 1.12507945e+00, 1.09315600e+00,
              1.10785308e+00, 1.10554958e+00, 1.11065162e+00, 1.08185194e+00,
              1.16768610e+00, 1.19984971e+00, 1.13528471e+00, 1.07712752e+00,
      ...
              8.39920287e-01, 9.18376018e-01, 9.79575480e-01, 9.80956682e-01,
              9.62580645e-01, 9.34735027e-01, 9.26718227e-01, 8.53845538e-01,
              9.55148637e-01, 8.96275421e-01, 9.55960911e-01, 8.37071150e-01,
              8.34055716e-01, 1.06550405e+00, 1.07424550e+00, 1.02240478e+00,
              1.06366048e+00, 9.81274197e-01, 1.14956089e+00, 9.85823325e-01,
              1.10501441e+00, 1.03303118e+00, 1.05014155e+00, 1.01986779e+00,
              9.53019843e-01, 1.12242877e+00, 1.08812316e+00, 1.03052979e+00,
              1.06006528e+00, 8.40247267e-01, 9.46324649e-01, 7.48405610e-01,
              8.52019182e-01, 7.91375314e-01, 7.85417938e-01, 8.12488210e-01,
              8.50874084e-01, 8.42392962e-01, 7.91759598e-01, 9.59168636e-01,
              7.37400377e-01, 9.28532667e-01, 9.43109268e-01, 9.08369453e-01,
              9.24467328e-01, 8.69644745e-01, 9.39666220e-01, 8.26393595e-01,
              8.16711835e-01, 1.01845260e+00, 9.06993389e-01, 8.87651711e-01,
              8.96817581e-01, 8.25869695e-01, 9.52375365e-01, 8.52314711e-01,
              1.00804455e+00, 9.55190746e-01, 9.02392015e-01, 9.25814533e-01,
              8.68093934e-01, 8.43086511e-01, 9.53203531e-01, 9.32838416e-01,
              9.16016333e-01, 8.35586939e-01, 7.99255145e-01, 8.47169593e-01,
              9.21955253e-01, 9.44084955e-01, 9.48042768e-01, 9.13055165e-01,
              9.92536809e-01, 9.01664836e-01, 1.01435874e+00, 7.93175050e-01,
              8.76015395e-01, 1.10475760e+00, 9.85515181e-01, 1.00437554e+00]])

    • v0_mean
      (chain, draw, __obs__)
      float64
      2.5 2.5 2.5 ... 1.004 1.004 1.004

      array([[[2.49994623e+00, 2.49994623e+00, 2.49994623e+00, ...,
               2.49994623e+00, 2.49994623e+00, 2.49994623e+00],
              [2.49643190e+00, 2.49643190e+00, 2.49643190e+00, ...,
               2.49643190e+00, 2.49643190e+00, 2.49643190e+00],
              [2.49621832e+00, 2.49621832e+00, 2.49621832e+00, ...,
               2.49621832e+00, 2.49621832e+00, 2.49621832e+00],
              ...,
              [1.02041450e+00, 1.02041450e+00, 1.02041450e+00, ...,
               1.02041450e+00, 1.02041450e+00, 1.02041450e+00],
              [8.49155634e-01, 8.49155634e-01, 8.49155634e-01, ...,
               8.49155634e-01, 8.49155634e-01, 8.49155634e-01],
              [9.33917203e-01, 9.33917203e-01, 9.33917203e-01, ...,
               9.33917203e-01, 9.33917203e-01, 9.33917203e-01]],
      
             [[1.38249748e+00, 1.38249748e+00, 1.38249748e+00, ...,
               1.38249748e+00, 1.38249748e+00, 1.38249748e+00],
              [1.81100027e+00, 1.81100027e+00, 1.81100027e+00, ...,
               1.81100027e+00, 1.81100027e+00, 1.81100027e+00],
              [1.81048955e+00, 1.81048955e+00, 1.81048955e+00, ...,
               1.81048955e+00, 1.81048955e+00, 1.81048955e+00],
      ...
              [1.06471054e+00, 1.06471054e+00, 1.06471054e+00, ...,
               1.06471054e+00, 1.06471054e+00, 1.06471054e+00],
              [9.93209135e-01, 9.93209135e-01, 9.93209135e-01, ...,
               9.93209135e-01, 9.93209135e-01, 9.93209135e-01],
              [1.04265035e+00, 1.04265035e+00, 1.04265035e+00, ...,
               1.04265035e+00, 1.04265035e+00, 1.04265035e+00]],
      
             [[5.38982981e-16, 5.38982981e-16, 5.38982981e-16, ...,
               5.38982981e-16, 5.38982981e-16, 5.38982981e-16],
              [2.82193738e-13, 2.82193738e-13, 2.82193738e-13, ...,
               2.82193738e-13, 2.82193738e-13, 2.82193738e-13],
              [3.18779366e-12, 3.18779366e-12, 3.18779366e-12, ...,
               3.18779366e-12, 3.18779366e-12, 3.18779366e-12],
              ...,
              [1.10475760e+00, 1.10475760e+00, 1.10475760e+00, ...,
               1.10475760e+00, 1.10475760e+00, 1.10475760e+00],
              [9.85515181e-01, 9.85515181e-01, 9.85515181e-01, ...,
               9.85515181e-01, 9.85515181e-01, 9.85515181e-01],
              [1.00437554e+00, 1.00437554e+00, 1.00437554e+00, ...,
               1.00437554e+00, 1.00437554e+00, 1.00437554e+00]]])

  • 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,
             ...
             190, 191, 192, 193, 194, 195, 196, 197, 198, 199],
            dtype='int64', name='draw', length=200))

    • __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: (8)

    created_at :

    2025-07-13T13:38:55.197877+00:00

    arviz_version :

    0.21.0

    inference_library :

    pymc

    inference_library_version :

    5.21.1

    sampling_time :

    33.21286725997925

    tuning_steps :

    200

    modeling_interface :

    bambi

    modeling_interface_version :

    0.15.0

  
  
• warmup_sample_stats

  <xarray.Dataset> Size: 66kB
  Dimensions:          (chain: 4, draw: 200, nstep_in_dim_0: 5, nstep_out_dim_0: 5)
  Coordinates:
    * chain            (chain) int64 32B 0 1 2 3
    * draw             (draw) int64 2kB 0 1 2 3 4 5 6 ... 194 195 196 197 198 199
    * nstep_in_dim_0   (nstep_in_dim_0) int64 40B 0 1 2 3 4
    * nstep_out_dim_0  (nstep_out_dim_0) int64 40B 0 1 2 3 4
  Data variables:
      nstep_in         (chain, draw, nstep_in_dim_0) int64 32kB 0 0 0 0 ... 3 3 2
      nstep_out        (chain, draw, nstep_out_dim_0) int64 32kB 344 3 131 ... 0 0
  Attributes:
      created_at:                  2025-07-13T13:38:55.200992+00:00
      arviz_version:               0.21.0
      inference_library:           pymc
      inference_library_version:   5.21.1
      sampling_time:               33.21286725997925
      tuning_steps:                200
      modeling_interface:          bambi
      modeling_interface_version:  0.15.0

  xarray.Dataset

  • Dimensions:
    • chain: 4
    • draw: 200
    • nstep_in_dim_0: 5
    • nstep_out_dim_0: 5
  • Coordinates: (4)
    • chain
      (chain)
      int64
      0 1 2 3

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

    • draw
      (draw)
      int64
      0 1 2 3 4 5 ... 195 196 197 198 199

      array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,  13,
              14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,
              28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,
              42,  43,  44,  45,  46,  47,  48,  49,  50,  51,  52,  53,  54,  55,
              56,  57,  58,  59,  60,  61,  62,  63,  64,  65,  66,  67,  68,  69,
              70,  71,  72,  73,  74,  75,  76,  77,  78,  79,  80,  81,  82,  83,
              84,  85,  86,  87,  88,  89,  90,  91,  92,  93,  94,  95,  96,  97,
              98,  99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
             112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125,
             126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139,
             140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153,
             154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167,
             168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181,
             182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195,
             196, 197, 198, 199])

    • nstep_in_dim_0
      (nstep_in_dim_0)
      int64
      0 1 2 3 4

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

    • nstep_out_dim_0
      (nstep_out_dim_0)
      int64
      0 1 2 3 4

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

  • Data variables: (2)
    • nstep_in
      (chain, draw, nstep_in_dim_0)
      int64
      0 0 0 0 0 1 5 2 ... 7 1 3 4 0 3 3 2

      array([[[ 0,  0,  0,  0,  0],
              [ 1,  5,  2,  0,  0],
              [ 3,  5,  8,  2,  2],
              ...,
              [ 5,  2,  6,  6,  4],
              [ 4,  0,  2,  2,  2],
              [ 4,  1,  3,  0,  0]],
      
             [[ 0,  1,  0,  0,  0],
              [ 5,  2,  0,  1,  1],
              [ 1,  0,  3,  0,  2],
              ...,
              [ 0,  1, 10,  2,  4],
              [ 5,  2, 10,  1,  0],
              [ 3,  2,  9,  1,  4]],
      
             [[ 0,  2,  2,  0,  3],
              [ 4,  0,  0,  0,  0],
              [ 5,  4,  4,  2,  0],
              ...,
              [ 3,  2,  5,  3,  1],
              [ 4,  0,  0,  2,  5],
              [ 4,  4,  0,  1,  0]],
      
             [[ 0,  2,  0,  0,  0],
              [ 2,  3,  0,  0,  1],
              [ 0,  2,  0,  4,  0],
              ...,
              [ 4,  3,  6,  5,  1],
              [ 0,  5,  7,  1,  3],
              [ 4,  0,  3,  3,  2]]])

    • nstep_out
      (chain, draw, nstep_out_dim_0)
      int64
      344 3 131 0 82 1 0 ... 0 0 0 0 0 0

      array([[[344,   3, 131,   0,  82],
              [  1,   0,   0,   1,   0],
              [  0,   0,   0,   0,   1],
              ...,
              [  0,   0,   0,   1,   0],
              [  0,   0,   0,   0,   1],
              [  0,   0,   0,   1,   0]],
      
             [[343,   3, 812,   0,   1],
              [  1,   1,   0,   0,   1],
              [  0,   1,   0,   0,   0],
              ...,
              [  0,   0,   0,   0,   1],
              [  0,   0,   0,   0,   0],
              [  0,   0,   0,   1,   1]],
      
             [[344,   3,   2,  58,   1],
              [  1,   0,   0,   0,   1],
              [  0,   0,   1,   0,   3],
              ...,
              [  0,   0,   0,   1,   1],
              [  0,   1,   0,   0,   1],
              [  0,   1,   1,   0,   0]],
      
             [[346,   3, 128, 226,  66],
              [  1,   0,   0,   1,   0],
              [  0,   0,   0,   1,   0],
              ...,
              [  0,   0,   0,   1,   0],
              [  0,   0,   0,   0,   0],
              [  0,   0,   0,   0,   0]]])

  • Indexes: (4)
    • 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,
             ...
             190, 191, 192, 193, 194, 195, 196, 197, 198, 199],
            dtype='int64', name='draw', length=200))

    • nstep_in_dim_0
      PandasIndex

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

    • nstep_out_dim_0
      PandasIndex

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

  • Attributes: (8)

    created_at :

    2025-07-13T13:38:55.200992+00:00

    arviz_version :

    0.21.0

    inference_library :

    pymc

    inference_library_version :

    5.21.1

    sampling_time :

    33.21286725997925

    tuning_steps :

    200

    modeling_interface :

    bambi

    modeling_interface_version :

    0.15.0

  
  

In [11]:

import arviz as az

az.plot_trace(model.traces, var_names=["~v0_mean"])
plt.tight_layout()

import arviz as az az.plot_trace(model.traces, var_names=["~v0_mean"]) plt.tight_layout()

/Users/afengler/Library/CloudStorage/OneDrive-Personal/proj_hssm/HSSM/.venv/lib/python3.12/site-packages/arviz/utils.py:146: UserWarning: Items starting with ~: ['v0_mean'] have not been found and will be ignored
  warnings.warn(

In [12]:

az.plot_pair(model.traces, var_names=["~v0_mean"])
plt.tight_layout()

az.plot_pair(model.traces, var_names=["~v0_mean"]) plt.tight_layout()

/Users/afengler/Library/CloudStorage/OneDrive-Personal/proj_hssm/HSSM/.venv/lib/python3.12/site-packages/arviz/utils.py:146: UserWarning: Items starting with ~: ['v0_mean'] have not been found and will be ignored
  warnings.warn(