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.0005616998672485352 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: [v2]
>Slice: [t]
>Slice: [a]
>Slice: [v1]
>Slice: [v0]

Output()

Sampling 4 chains for 200 tune and 500 draw iterations (800 + 2_000 draws total) took 28 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
100%|██████████| 2000/2000 [00:00<00:00, 2250.62it/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:
      t        (chain, draw) float64 16kB 0.315 0.3216 0.3237 ... 0.3224 0.319
      v2       (chain, draw) float64 16kB 0.001899 0.03002 ... 0.06186 0.01676
      v0       (chain, draw) float64 16kB 0.7342 0.7491 0.7877 ... 0.7088 0.7291
      a        (chain, draw) float64 16kB 1.23 1.222 1.235 ... 1.225 1.247 1.233
      v1       (chain, draw) float64 16kB 0.3053 0.4118 0.2612 ... 0.3729 0.3705
  Attributes:
      created_at:                  2024-12-25T23:24:24.451015+00:00
      arviz_version:               0.18.0
      inference_library:           pymc
      inference_library_version:   5.16.2
      sampling_time:               28.457155227661133
      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: (5)
    • t
      (chain, draw)
      float64
      0.315 0.3216 ... 0.3224 0.319

      array([[0.31504365, 0.32164031, 0.3237299 , ..., 0.3170411 , 0.3171035 ,
              0.3119193 ],
             [0.31239804, 0.31403069, 0.31505056, ..., 0.32120188, 0.31949881,
              0.31959874],
             [0.31534373, 0.31694127, 0.31629697, ..., 0.31500242, 0.31591025,
              0.31452968],
             [0.30530566, 0.30458979, 0.305918  , ..., 0.31611258, 0.32240671,
              0.31898199]])

    • v2
      (chain, draw)
      float64
      0.001899 0.03002 ... 0.01676

      array([[0.00189915, 0.03002217, 0.01400725, ..., 0.03093113, 0.02856157,
              0.08203611],
             [0.00875186, 0.00151756, 0.02681541, ..., 0.02660003, 0.02366668,
              0.05168057],
             [0.05733873, 0.02544779, 0.11058111, ..., 0.26001511, 0.01913283,
              0.00378113],
             [0.19864728, 0.10696656, 0.24705672, ..., 0.03720765, 0.0618582 ,
              0.01675826]])

    • v0
      (chain, draw)
      float64
      0.7342 0.7491 ... 0.7088 0.7291

      array([[0.73422576, 0.74909365, 0.78766758, ..., 0.80999386, 0.80039266,
              0.80003482],
             [0.84182197, 0.85578976, 0.79164804, ..., 0.71675204, 0.73765147,
              0.76346617],
             [0.76875758, 0.80499241, 0.83578998, ..., 0.70641629, 0.78116282,
              0.77553708],
             [0.89627966, 0.88757486, 0.85798013, ..., 0.76535547, 0.70882056,
              0.72909213]])

    • a
      (chain, draw)
      float64
      1.23 1.222 1.235 ... 1.247 1.233

      array([[1.23041075, 1.2219695 , 1.23471123, ..., 1.2555657 , 1.26521981,
              1.31479592],
             [1.26803611, 1.26448659, 1.25664135, ..., 1.22981427, 1.23272878,
              1.2022575 ],
             [1.2804423 , 1.28639802, 1.25329356, ..., 1.28494268, 1.27962584,
              1.26145475],
             [1.37505953, 1.36013777, 1.3566126 , ..., 1.22459401, 1.24719213,
              1.23317312]])

    • v1
      (chain, draw)
      float64
      0.3053 0.4118 ... 0.3729 0.3705

      array([[0.30534184, 0.41180332, 0.26120931, ..., 0.46591834, 0.27856571,
              0.30864199],
             [0.28168229, 0.43490075, 0.31198052, ..., 0.35126893, 0.31687734,
              0.30444336],
             [0.45146938, 0.49440363, 0.23573725, ..., 0.42571632, 0.37966161,
              0.38032989],
             [0.44002938, 0.54433945, 0.47959925, ..., 0.36574209, 0.3728817 ,
              0.37048313]])

  • 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:24:24.451015+00:00

    arviz_version :

    0.18.0

    inference_library :

    pymc

    inference_library_version :

    5.16.2

    sampling_time :

    28.457155227661133

    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.2108 -2.806 ... -0.8196
  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.2108 -2.806 ... 0.1428 -0.8196

      array([[[ 0.21079947, -2.80578184, -0.90374511, ..., -0.92080194,
                0.1364162 , -0.84094268],
              [ 0.23436444, -2.91908288, -0.92325836, ..., -0.94110435,
                0.15972427, -0.8156417 ],
              [ 0.27326781, -2.75417781, -0.83131272, ..., -0.84859419,
                0.20588839, -0.78791159],
              ...,
              [ 0.2366067 , -2.910007  , -0.89892298, ..., -0.9167577 ,
                0.16517487, -0.84906715],
              [ 0.24560328, -2.72574973, -0.82759064, ..., -0.84457201,
                0.18020169, -0.80621547],
              [ 0.20555201, -2.68535781, -0.81457543, ..., -0.83123332,
                0.14723958, -0.78096628]],
      
             [[ 0.25267854, -2.73311281, -0.826177  , ..., -0.84312379,
                0.18413498, -0.85470456],
              [ 0.25348052, -2.86248875, -0.86735189, ..., -0.88498092,
                0.18233024, -0.88622606],
              [ 0.23357953, -2.77604175, -0.86001939, ..., -0.87712216,
                0.16416328, -0.82674235],
      ...
              [ 0.14841953, -2.95656538, -0.96056664, ..., -0.97830409,
                0.08182101, -0.66450298],
              [ 0.21809348, -2.76563001, -0.85187769, ..., -0.86897624,
                0.15415627, -0.82639349],
              [ 0.2169587 , -2.80270338, -0.8782053 , ..., -0.89538532,
                0.14778927, -0.85161334]],
      
             [[ 0.20051107, -2.77908349, -0.8099004 , ..., -0.8270207 ,
                0.14600161, -0.75816929],
              [ 0.19580024, -2.83917975, -0.84466463, ..., -0.86202657,
                0.13682568, -0.84927565],
              [ 0.17942372, -2.87956023, -0.86972934, ..., -0.88722336,
                0.12026244, -0.72941327],
              ...,
              [ 0.22477996, -2.89916015, -0.92539746, ..., -0.94297636,
                0.14749664, -0.8358168 ],
              [ 0.20782593, -2.83369589, -0.89844519, ..., -0.91587764,
                0.14134422, -0.75960755],
              [ 0.21512723, -2.84971166, -0.91040611, ..., -0.9278093 ,
                0.14275929, -0.81956184]]])

  • 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: 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 1 2 5 1 ... 6 2 6
      nstep_out        (chain, draw, nstep_out_dim_0) int64 80kB 1 0 0 1 ... 0 0 1
  Attributes:
      created_at:                  2024-12-25T23:24:24.456908+00:00
      arviz_version:               0.18.0
      inference_library:           pymc
      inference_library_version:   5.16.2
      sampling_time:               28.457155227661133
      tuning_steps:                200
      modeling_interface:          bambi
      modeling_interface_version:  0.14.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
      1 2 5 1 1 0 1 4 ... 0 3 0 1 6 2 6

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

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

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

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

    2024-12-25T23:24:24.456908+00:00

    arviz_version :

    0.18.0

    inference_library :

    pymc

    inference_library_version :

    5.16.2

    sampling_time :

    28.457155227661133

    tuning_steps :

    200

    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:24:24.458647+00:00
      arviz_version:               0.18.0
      inference_library:           pymc
      inference_library_version:   5.16.2
      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
      0.4902 0.0 1.274 ... 0.0 0.592 2.0

      array([[0.49016428, 0.        ],
             [1.27432394, 0.        ],
             [0.72420627, 0.        ],
             ...,
             [0.72815591, 0.        ],
             [0.50786394, 0.        ],
             [0.59196222, 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 :

    2024-12-25T23:24:24.458647+00:00

    arviz_version :

    0.18.0

    inference_library :

    pymc

    inference_library_version :

    5.16.2

    modeling_interface :

    bambi

    modeling_interface_version :

    0.14.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:
      a        (chain, draw) float64 6kB 3.0 3.0 3.0 3.0 ... 1.283 1.318 1.379
      t        (chain, draw) float64 6kB 0.04004 0.1003 0.1193 ... 0.3096 0.307
      v0       (chain, draw) float64 6kB 2.038 1.805 1.803 ... 0.76 0.9304 0.8776
      v0_mean  (chain, draw, __obs__) float64 6MB 2.038 2.038 ... 0.8776 0.8776
      v1       (chain, draw) float64 6kB 1.305 1.632 1.749 ... 0.3983 0.5868
      v2       (chain, draw) float64 6kB 3.007e-60 0.5101 1.829 ... 0.05378 0.269
  Attributes:
      created_at:                  2024-12-25T23:24:24.455268+00:00
      arviz_version:               0.18.0
      inference_library:           pymc
      inference_library_version:   5.16.2
      sampling_time:               28.457155227661133
      tuning_steps:                200
      modeling_interface:          bambi
      modeling_interface_version:  0.14.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)
    • a
      (chain, draw)
      float64
      3.0 3.0 3.0 ... 1.283 1.318 1.379

      array([[3.        , 3.        , 3.        , 2.99999998, 2.99999988,
              2.99999959, 2.99999957, 2.96071221, 2.91671242, 2.92079302,
              2.84150663, 2.84725663, 2.83014913, 2.82932396, 2.81827468,
              2.54500142, 2.54527578, 2.51264158, 2.48075045, 2.43131042,
              2.37509285, 2.37050433, 2.31849269, 2.20237041, 2.07661864,
              2.08098184, 2.04583156, 2.01049373, 1.93526364, 1.8778792 ,
              1.8600945 , 1.8545043 , 1.80860374, 1.81209312, 1.80235945,
              1.8069167 , 1.81981738, 1.81889154, 1.80735135, 1.81522638,
              1.7834209 , 1.80569539, 1.78279209, 1.7652959 , 1.76231446,
              1.69882144, 1.63610329, 1.64072337, 1.58680526, 1.56510669,
              1.57910097, 1.51807784, 1.52107562, 1.5256099 , 1.46446961,
              1.50553774, 1.46163346, 1.41015396, 1.43502571, 1.38757362,
              1.37388586, 1.38502167, 1.35980554, 1.34078765, 1.34043575,
              1.32675994, 1.32612534, 1.31378052, 1.31149356, 1.32475519,
              1.32549833, 1.2929642 , 1.32392137, 1.33535902, 1.34257737,
              1.34069752, 1.35496971, 1.3917918 , 1.40662588, 1.35416192,
              1.3587426 , 1.35424399, 1.3917709 , 1.39150185, 1.41490014,
              1.42810659, 1.39062014, 1.37761201, 1.37505763, 1.35436773,
              1.34081886, 1.32050135, 1.29664926, 1.2718321 , 1.21378658,
              1.24170064, 1.2773512 , 1.26840261, 1.25780709, 1.26739651,
      ...
              1.40950244, 1.39055854, 1.38927814, 1.36283675, 1.37882383,
              1.37663894, 1.36192068, 1.32898446, 1.34004282, 1.30936447,
              1.33151159, 1.31703494, 1.3044913 , 1.28064548, 1.29153974,
              1.33362307, 1.37559364, 1.32624588, 1.31824242, 1.30571673,
              1.29539321, 1.27622401, 1.28503176, 1.32441183, 1.34646033,
              1.39898532, 1.38727459, 1.34928549, 1.37817881, 1.40396945,
              1.40508546, 1.36697287, 1.3809876 , 1.37236482, 1.36882372,
              1.38969408, 1.33508914, 1.31284369, 1.28548103, 1.27687883,
              1.2782232 , 1.31601884, 1.35637373, 1.33635592, 1.33090272,
              1.32407981, 1.35288312, 1.33927006, 1.35634018, 1.35499621,
              1.31383238, 1.32160385, 1.24086178, 1.24234328, 1.26021152,
              1.26151937, 1.26130254, 1.27173499, 1.31297468, 1.28522534,
              1.26151106, 1.26617048, 1.25857151, 1.29425639, 1.34804343,
              1.3160608 , 1.34392487, 1.33454667, 1.30885181, 1.34230161,
              1.36602341, 1.38313277, 1.35096602, 1.33422978, 1.31153861,
              1.28746576, 1.26452754, 1.26693476, 1.27334368, 1.23868856,
              1.27288812, 1.28159684, 1.28564277, 1.29225178, 1.25751819,
              1.26823933, 1.28024801, 1.25649341, 1.26961469, 1.27205786,
              1.33209051, 1.32359109, 1.31571915, 1.29860212, 1.29787027,
              1.2920736 , 1.2688458 , 1.28264647, 1.3183324 , 1.37895715]])

    • t
      (chain, draw)
      float64
      0.04004 0.1003 ... 0.3096 0.307

      array([[0.04003957, 0.10029952, 0.11929487, 0.13488861, 0.13054776,
              0.10967031, 0.10240416, 0.11159744, 0.11095184, 0.13147894,
              0.13293912, 0.1451427 , 0.14295307, 0.12885496, 0.12011376,
              0.15028032, 0.15023957, 0.17353144, 0.18714659, 0.19004441,
              0.18615328, 0.19528118, 0.20452258, 0.20522445, 0.21315076,
              0.2309544 , 0.22889866, 0.23666743, 0.2454452 , 0.24853254,
              0.25723522, 0.25445735, 0.25216577, 0.25007248, 0.25635454,
              0.25664998, 0.25821666, 0.25486276, 0.25463828, 0.25434709,
              0.2682886 , 0.2593756 , 0.26343175, 0.26813968, 0.26798271,
              0.26603819, 0.27616375, 0.27807392, 0.27784214, 0.2876014 ,
              0.28990419, 0.29361632, 0.29482258, 0.29042274, 0.29362287,
              0.29124615, 0.29656067, 0.29845921, 0.2978607 , 0.29868062,
              0.30065346, 0.30498549, 0.30456662, 0.3110823 , 0.30926474,
              0.30964891, 0.30932706, 0.30969815, 0.31302215, 0.31193384,
              0.30904733, 0.31285934, 0.31333552, 0.30854106, 0.31031522,
              0.31232794, 0.30954681, 0.30796375, 0.30010734, 0.30890663,
              0.30695191, 0.30654713, 0.30675253, 0.30214553, 0.30381397,
              0.30341389, 0.30303059, 0.30333851, 0.30465686, 0.30758263,
              0.30689133, 0.31499804, 0.31509119, 0.31110469, 0.31767806,
              0.31795495, 0.3160385 , 0.31518998, 0.31484341, 0.31403674,
      ...
              0.30199331, 0.30126381, 0.30414334, 0.30392305, 0.30480086,
              0.30625372, 0.30107338, 0.30475157, 0.30646087, 0.31369868,
              0.3105449 , 0.31067825, 0.31087331, 0.3150295 , 0.31402015,
              0.3151685 , 0.30581334, 0.30812498, 0.31021657, 0.31378096,
              0.31358152, 0.31344744, 0.31508121, 0.31446873, 0.30947679,
              0.30555438, 0.30374736, 0.30694133, 0.30403213, 0.30268188,
              0.30201814, 0.30593269, 0.30360134, 0.30488741, 0.3031121 ,
              0.30338251, 0.30506518, 0.30532391, 0.30989938, 0.31018582,
              0.31287459, 0.31164068, 0.30814636, 0.3076451 , 0.30990268,
              0.31218594, 0.30575321, 0.30726788, 0.30818454, 0.30785101,
              0.31111954, 0.30590484, 0.31594746, 0.31652328, 0.31626748,
              0.31598161, 0.31893835, 0.3171408 , 0.31154631, 0.3116241 ,
              0.31624845, 0.3157809 , 0.31839471, 0.3197486 , 0.30813308,
              0.30729899, 0.30826534, 0.3135031 , 0.31132083, 0.31106351,
              0.31090302, 0.30425177, 0.30814743, 0.30888541, 0.30857649,
              0.31658789, 0.31670847, 0.31394998, 0.31446669, 0.31873149,
              0.31829329, 0.31318842, 0.31322924, 0.31902193, 0.31404528,
              0.31353329, 0.3176413 , 0.31819499, 0.31594385, 0.32146203,
              0.31001869, 0.31025486, 0.31227963, 0.31318388, 0.31550415,
              0.31006761, 0.31054354, 0.31375773, 0.30962078, 0.30698131]])

    • v0
      (chain, draw)
      float64
      2.038 1.805 1.803 ... 0.9304 0.8776

      array([[2.03759584, 1.8046471 , 1.80315289, 2.08182891, 2.08106575,
              1.88477045, 1.88839521, 1.83040468, 1.890748  , 1.93281076,
              1.92782818, 1.97363283, 1.9238545 , 1.81682207, 1.82606364,
              1.78990327, 1.65725092, 1.82985886, 1.7626508 , 1.70933734,
              1.67462609, 1.7586272 , 1.65245342, 1.45786911, 1.57351636,
              1.52389352, 1.56637149, 1.5866955 , 1.34644071, 1.29326814,
              1.43946252, 1.3011344 , 1.27831398, 1.33182713, 1.28679138,
              1.23522774, 1.42802596, 1.27917236, 1.32479032, 1.38069365,
              1.29470543, 1.30173119, 1.35673298, 1.30048517, 1.32561221,
              1.13536274, 1.21866389, 1.26868234, 1.22086831, 1.24323009,
              1.11325012, 1.13636369, 1.1237303 , 1.08009045, 1.07727922,
              1.13020664, 0.915067  , 0.94766832, 0.87870746, 0.85194601,
              0.93821465, 0.99679225, 0.88401511, 0.79611672, 0.8490358 ,
              0.8789225 , 0.74051471, 0.73957802, 0.8393775 , 0.86176751,
              0.8722259 , 0.9112642 , 0.82583064, 0.87066003, 0.93067046,
              0.947897  , 0.95484202, 0.93037205, 0.99166397, 0.97602184,
              0.84333577, 0.9124029 , 0.83674635, 1.0615478 , 1.11309343,
              0.90085452, 0.99607943, 0.81638026, 0.86774628, 0.86542022,
              0.91679041, 0.85521235, 0.87221395, 0.77130829, 0.86436241,
              0.67238433, 0.79398342, 0.76698261, 0.81062283, 0.66296398,
      ...
              0.86674119, 0.98453931, 0.913776  , 0.95546498, 0.7906451 ,
              0.8482772 , 0.86868424, 0.82803382, 0.90478302, 0.83822858,
              0.92255336, 0.76932829, 0.79233965, 0.81364844, 0.94590522,
              0.8009403 , 0.75531633, 0.88792194, 0.80420931, 0.8153833 ,
              0.82340312, 0.78220924, 0.79357162, 0.88759979, 0.90254979,
              0.90369498, 0.90074336, 0.83298999, 1.01274503, 1.00893012,
              0.90970302, 0.86026434, 0.80420606, 0.99195721, 0.82907588,
              0.81803974, 0.69452123, 0.75001415, 0.797687  , 0.78189685,
              0.8242912 , 0.92274256, 0.90687105, 0.90946238, 0.85819145,
              0.79347289, 0.89423009, 0.76263588, 0.94134388, 0.87667429,
              0.77371751, 0.83262496, 0.68960009, 0.65200174, 0.77941401,
              0.79729277, 0.93769446, 0.79837339, 0.7688953 , 0.80221816,
              0.85421016, 0.78454814, 0.88426372, 0.89991826, 0.85738375,
              0.81463238, 0.8521204 , 0.91396269, 0.82339666, 0.84480067,
              0.98782971, 0.95197395, 0.95175063, 0.75681615, 0.69734749,
              0.78381085, 0.79411639, 0.78066101, 0.75805404, 0.78282562,
              0.81956801, 0.96663429, 0.94207782, 0.87159701, 0.74334459,
              0.82495031, 0.84290687, 0.7124859 , 0.84001639, 0.7462106 ,
              0.94492534, 0.77648093, 0.89046297, 0.93282297, 0.82626486,
              0.77217046, 0.80281895, 0.76000131, 0.93036272, 0.87762563]])

    • v0_mean
      (chain, draw, __obs__)
      float64
      2.038 2.038 2.038 ... 0.8776 0.8776

      array([[[2.03759584, 2.03759584, 2.03759584, ..., 2.03759584,
               2.03759584, 2.03759584],
              [1.8046471 , 1.8046471 , 1.8046471 , ..., 1.8046471 ,
               1.8046471 , 1.8046471 ],
              [1.80315289, 1.80315289, 1.80315289, ..., 1.80315289,
               1.80315289, 1.80315289],
              ...,
              [0.79615594, 0.79615594, 0.79615594, ..., 0.79615594,
               0.79615594, 0.79615594],
              [0.75114046, 0.75114046, 0.75114046, ..., 0.75114046,
               0.75114046, 0.75114046],
              [0.80178282, 0.80178282, 0.80178282, ..., 0.80178282,
               0.80178282, 0.80178282]],
      
             [[2.15607581, 2.15607581, 2.15607581, ..., 2.15607581,
               2.15607581, 2.15607581],
              [2.1235076 , 2.1235076 , 2.1235076 , ..., 2.1235076 ,
               2.1235076 , 2.1235076 ],
              [2.04161865, 2.04161865, 2.04161865, ..., 2.04161865,
               2.04161865, 2.04161865],
      ...
              [0.74122714, 0.74122714, 0.74122714, ..., 0.74122714,
               0.74122714, 0.74122714],
              [0.86520614, 0.86520614, 0.86520614, ..., 0.86520614,
               0.86520614, 0.86520614],
              [0.88906126, 0.88906126, 0.88906126, ..., 0.88906126,
               0.88906126, 0.88906126]],
      
             [[1.24166652, 1.24166652, 1.24166652, ..., 1.24166652,
               1.24166652, 1.24166652],
              [1.77756772, 1.77756772, 1.77756772, ..., 1.77756772,
               1.77756772, 1.77756772],
              [1.56331798, 1.56331798, 1.56331798, ..., 1.56331798,
               1.56331798, 1.56331798],
              ...,
              [0.76000131, 0.76000131, 0.76000131, ..., 0.76000131,
               0.76000131, 0.76000131],
              [0.93036272, 0.93036272, 0.93036272, ..., 0.93036272,
               0.93036272, 0.93036272],
              [0.87762563, 0.87762563, 0.87762563, ..., 0.87762563,
               0.87762563, 0.87762563]]])

    • v1
      (chain, draw)
      float64
      1.305 1.632 1.749 ... 0.3983 0.5868

      array([[1.30462559, 1.63150928, 1.74907738, 1.68072074, 1.59920136,
              1.50355206, 1.51828514, 1.53726298, 1.52268502, 1.60118978,
              1.58206139, 1.61859764, 1.47749433, 1.39655684, 1.52229236,
              1.12496038, 1.22344678, 1.29768436, 1.44813679, 1.35909869,
              1.28897663, 1.32587645, 1.21331957, 1.10896452, 1.00778881,
              1.13180793, 1.08257818, 1.04150572, 0.92893123, 0.96095676,
              1.08484866, 0.86479228, 0.87677839, 0.86522079, 0.88950499,
              0.89340715, 0.85639533, 0.84066762, 0.90352926, 0.93436432,
              0.91942426, 1.02908862, 0.78738325, 0.96427793, 0.82085663,
              0.82133101, 0.78858725, 0.70266236, 0.73491246, 0.71575261,
              0.73953707, 0.72281424, 0.72317366, 0.68482059, 0.74200534,
              0.55926375, 0.56066972, 0.53944622, 0.50976834, 0.48404047,
              0.47377606, 0.43517445, 0.49387834, 0.47874848, 0.5330483 ,
              0.44361092, 0.44821611, 0.44393218, 0.44163477, 0.4288193 ,
              0.46110692, 0.33216463, 0.41889385, 0.45797705, 0.47933832,
              0.51201202, 0.50256448, 0.46760814, 0.50511322, 0.434934  ,
              0.48341329, 0.55210671, 0.55268882, 0.60887638, 0.60472612,
              0.57932548, 0.59014357, 0.61104809, 0.50975567, 0.49203739,
              0.45783555, 0.45669992, 0.30425778, 0.31431182, 0.25126634,
              0.25338354, 0.2853279 , 0.29624262, 0.3315372 , 0.43841799,
      ...
              0.57616866, 0.43459587, 0.45340873, 0.39200951, 0.62591692,
              0.48363767, 0.5551716 , 0.35094563, 0.55998013, 0.53555919,
              0.38536458, 0.44013415, 0.41223148, 0.41230232, 0.45378091,
              0.50564328, 0.56419768, 0.47793282, 0.40734219, 0.45526322,
              0.39604741, 0.38434553, 0.48203896, 0.51973921, 0.48318641,
              0.57094228, 0.47432672, 0.50248938, 0.52540619, 0.59126091,
              0.48359618, 0.58592022, 0.50428806, 0.45856409, 0.53480994,
              0.53559321, 0.36916857, 0.31611494, 0.29529274, 0.42317752,
              0.3800372 , 0.40481054, 0.57649583, 0.43515512, 0.48165533,
              0.38362112, 0.50275535, 0.42132383, 0.45580724, 0.59273173,
              0.53514877, 0.45437896, 0.26862781, 0.34146971, 0.39099734,
              0.31726122, 0.43561505, 0.36798581, 0.40969439, 0.50695615,
              0.33363291, 0.40554593, 0.42302924, 0.46244499, 0.50714975,
              0.45656166, 0.49763822, 0.35476525, 0.42511115, 0.55230078,
              0.46575948, 0.51649215, 0.4635032 , 0.39722856, 0.35491425,
              0.42190533, 0.51014576, 0.26539481, 0.34853091, 0.37804375,
              0.49320889, 0.28163371, 0.37047251, 0.36644694, 0.34109659,
              0.4600682 , 0.48877218, 0.38682394, 0.39838355, 0.45019114,
              0.49120473, 0.4217414 , 0.40470142, 0.4097634 , 0.44849293,
              0.44291294, 0.3158448 , 0.4081712 , 0.39827618, 0.5868071 ]])

    • v2
      (chain, draw)
      float64
      3.007e-60 0.5101 ... 0.05378 0.269

      array([[3.00671216e-60, 5.10089196e-01, 1.82865485e+00, 1.75419409e+00,
              1.75238293e+00, 1.25015566e+00, 1.33851690e+00, 1.41910399e+00,
              1.38703705e+00, 1.27014187e+00, 1.42949585e+00, 1.34285073e+00,
              1.43780399e+00, 1.22906732e+00, 1.36393016e+00, 1.16298122e+00,
              1.05061915e+00, 1.00285964e+00, 1.25402553e+00, 1.16650635e+00,
              1.13966542e+00, 1.01278686e+00, 1.03905305e+00, 1.04143545e+00,
              9.68012892e-01, 9.04314081e-01, 8.65246983e-01, 7.43353766e-01,
              8.52264600e-01, 6.53031370e-01, 8.19813224e-01, 7.38725785e-01,
              6.27303418e-01, 6.70925353e-01, 6.57409340e-01, 6.15592337e-01,
              5.12512095e-01, 7.11136766e-01, 6.74866155e-01, 6.34949732e-01,
              7.10816772e-01, 6.44906229e-01, 6.31098504e-01, 7.01857918e-01,
              6.89191774e-01, 6.46744034e-01, 6.26554363e-01, 7.06546089e-01,
              4.25941644e-01, 3.73710676e-01, 4.65308563e-01, 4.13473807e-01,
              4.34868667e-01, 3.44710010e-01, 3.81798973e-01, 3.60899683e-01,
              2.83052709e-01, 2.71876990e-01, 3.59886657e-01, 1.94784876e-01,
              3.23098023e-01, 1.07357153e-01, 2.15651583e-01, 2.82073182e-01,
              2.65307362e-01, 1.11452371e-01, 2.14269238e-01, 1.45580359e-01,
              1.96504056e-01, 2.34768887e-01, 8.85930964e-02, 8.13130323e-02,
              9.29735510e-02, 2.10882993e-01, 1.11846912e-01, 1.96091206e-01,
              1.55772517e-01, 2.50520829e-01, 2.37697370e-01, 3.30922346e-01,
      ...
              6.16628555e-02, 8.88838572e-02, 7.64899162e-02, 2.22366015e-01,
              1.37431612e-01, 3.21607449e-01, 2.70074463e-01, 3.12969081e-01,
              2.08707615e-01, 2.12643884e-01, 2.71649978e-01, 2.29995705e-01,
              2.02005564e-01, 2.18586163e-01, 1.29299726e-01, 2.48179294e-01,
              2.78985442e-01, 9.76570496e-02, 8.69448529e-02, 3.07449145e-02,
              3.08709489e-02, 1.67185766e-01, 1.50989120e-01, 2.15896158e-01,
              1.47338275e-01, 1.80619009e-01, 1.82694967e-01, 1.56388317e-01,
              1.45275433e-01, 1.77209968e-01, 1.63861953e-01, 6.52802693e-02,
              1.14545882e-01, 6.33259715e-02, 1.30623066e-01, 9.08593055e-02,
              1.75544450e-01, 1.10646172e-01, 6.09816447e-02, 8.17352667e-02,
              1.00995461e-01, 9.97273526e-02, 1.22132858e-01, 6.01582621e-02,
              3.90658805e-02, 5.85735001e-02, 1.53052028e-01, 2.06373347e-01,
              1.95344481e-01, 2.64640906e-01, 2.83222206e-01, 1.93974378e-01,
              2.83155900e-01, 1.51388807e-01, 2.14884973e-01, 1.56856133e-01,
              9.48029515e-02, 1.20043364e-01, 7.27549006e-02, 8.69149182e-02,
              1.53758138e-01, 1.75142478e-01, 4.74425469e-02, 1.95365739e-01,
              8.50420851e-02, 1.09832364e-01, 8.09276866e-02, 1.79936263e-02,
              1.39411107e-02, 4.20713736e-02, 2.66891061e-02, 1.87723546e-01,
              1.95599684e-01, 1.82847056e-01, 1.01715855e-01, 7.22094366e-02,
              6.04721085e-02, 1.45876475e-01, 5.37794270e-02, 2.69010047e-01]])

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

    2024-12-25T23:24:24.455268+00:00

    arviz_version :

    0.18.0

    inference_library :

    pymc

    inference_library_version :

    5.16.2

    sampling_time :

    28.457155227661133

    tuning_steps :

    200

    modeling_interface :

    bambi

    modeling_interface_version :

    0.14.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 2 ... 8 0 3
      nstep_out        (chain, draw, nstep_out_dim_0) int64 32kB 395 3 786 ... 1 1
  Attributes:
      created_at:                  2024-12-25T23:24:24.457926+00:00
      arviz_version:               0.18.0
      inference_library:           pymc
      inference_library_version:   5.16.2
      sampling_time:               28.457155227661133
      tuning_steps:                200
      modeling_interface:          bambi
      modeling_interface_version:  0.14.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 2 0 2 3 2 ... 1 1 2 1 5 8 0 3

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

    • nstep_out
      (chain, draw, nstep_out_dim_0)
      int64
      395 3 786 0 2 1 1 ... 1 0 0 1 0 1 1

      array([[[395,   3, 786,   0,   2],
              [  1,   1,   1,   1,   0],
              [  0,   0,   0,   0,   0],
              ...,
              [  0,   0,   0,   1,   1],
              [  0,   0,   0,   1,   0],
              [  1,   0,   0,   1,   0]],
      
             [[395,   3, 600,   1,   2],
              [  0,   0,   0,   1,   1],
              [  0,   0,   0,   1,   1],
              ...,
              [  0,   0,   0,   1,   1],
              [  0,   0,   0,   0,   0],
              [  0,   0,   0,   1,   0]],
      
             [[396,   3,   1,   2,   1],
              [  0,   0,   0,   0,   0],
              [  0,   0,   1,   0,   1],
              ...,
              [  0,   0,   0,   0,   0],
              [  0,   0,   0,   1,   1],
              [  1,   0,   0,   1,   0]],
      
             [[395,   3, 834,   0,   1],
              [  0,   0,   1,   0,   1],
              [  0,   0,   0,   1,   0],
              ...,
              [  0,   0,   0,   1,   0],
              [  0,   0,   0,   1,   0],
              [  0,   1,   0,   1,   1]]])

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

    2024-12-25T23:24:24.457926+00:00

    arviz_version :

    0.18.0

    inference_library :

    pymc

    inference_library_version :

    5.16.2

    sampling_time :

    28.457155227661133

    tuning_steps :

    200

    modeling_interface :

    bambi

    modeling_interface_version :

    0.14.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/miniconda3/envs/hssm516/lib/python3.11/site-packages/arviz/utils.py:142: 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/miniconda3/envs/hssm516/lib/python3.11/site-packages/arviz/utils.py:142: UserWarning: Items starting with ~: ['v0_mean'] have not been found and will be ignored
  warnings.warn(