Attach trial wise parameters to idata post-hoc¶
HSSM automatically cleans up the idata object we return from samplers, to avoid unnecessarily huge objects from being passed around (the decision was taken after having observed many inadvertent out of memory errors).
Instead, if so desired, we can attach the trial wise parameters to the idata object ourselves, using the add_likelihood_parameters_to_idata function.
This quick tutorial shows you how to do this.
Load Modules¶
import pandas as pd
import hssm
Simulate Data¶
# Condition 1
condition_1 = hssm.simulate_data(
model="ddm", theta=dict(v=0.5, a=1.5, z=0.5, t=0.1), size=500
)
# Condition 2
condition_2 = hssm.simulate_data(
model="ddm", theta=dict(v=1.0, a=1.5, z=0.5, t=0.1), size=500
)
condition_1["condition"] = "C1"
condition_2["condition"] = "C2"
data = pd.concat([condition_1, condition_2]).reset_index(drop=True)
Define Model and Sample¶
model = hssm.HSSM(
model="ddm",
data=data,
include=[
{
"name": "v",
"formula": "v ~ 1 + condition",
}
],
)
idata = model.sample(sampler="pymc", tune=500, draws=500)
Model initialized successfully. Using default initvals.
Initializing NUTS using adapt_diag... Multiprocess sampling (4 chains in 4 jobs) NUTS: [a, z, t, v_Intercept, v_condition] /Users/yxu150/.local/share/uv/python/cpython-3.13.13-macos-aarch64-none/lib/python3.13/multiprocessing/popen_fork.py:67: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock. self.pid = os.fork()
Output()
/Users/yxu150/.local/share/uv/python/cpython-3.13.13-macos-aarch64-none/lib/python3.13/multiprocessing/popen_fork.py:67: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock. self.pid = os.fork()
Sampling 4 chains for 500 tune and 500 draw iterations (2_000 + 2_000 draws total) took 8 seconds.
Checking the idata object¶
idata
<xarray.DataTree>
Group: /
├── Group: /posterior
│ Dimensions: (chain: 4, draw: 500, v_condition_dim: 1)
│ 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
│ * v_condition_dim (v_condition_dim) <U2 8B 'C2'
│ Data variables:
│ v_condition (chain, draw, v_condition_dim) float64 16kB 0.5306 ... 0...
│ z (chain, draw) float64 16kB 0.4887 0.4932 ... 0.4967 0.4828
│ v_Intercept (chain, draw) float64 16kB 0.6786 0.6577 ... 0.5924 0.6688
│ t (chain, draw) float64 16kB 0.01325 0.01503 ... 0.03315
│ a (chain, draw) float64 16kB 1.643 1.653 1.611 ... 1.54 1.577
│ Attributes:
│ created_at: 2026-06-30T13:13:24.645662+00:00
│ creation_library: ArviZ
│ creation_library_version: 1.1.0
│ creation_library_language: Python
│ inference_library: pymc
│ inference_library_version: 6.0.1
│ sample_dims: ['chain', 'draw']
│ sampling_time: 8.305341958999634
│ tuning_steps: 500
│ modeling_interface: bambi
│ modeling_interface_version: 0.18.0
├── Group: /sample_stats
│ Dimensions: (chain: 4, draw: 500)
│ Coordinates:
│ * chain (chain) int64 32B 0 1 2 3
│ * draw (draw) int64 4kB 0 1 2 3 4 5 ... 495 496 497 498 499
│ Data variables: (12/18)
│ tree_depth (chain, draw) int64 16kB 4 2 4 3 3 4 ... 4 4 3 3 3 3
│ step_size_bar (chain, draw) float64 16kB 0.3508 0.3508 ... 0.3453
│ max_energy_error (chain, draw) float64 16kB -0.9847 ... -0.08732
│ smallest_eigval (chain, draw) float64 16kB nan nan nan ... nan nan
│ energy_error (chain, draw) float64 16kB -0.1505 ... 0.06108
│ n_steps (chain, draw) float64 16kB 15.0 3.0 15.0 ... 7.0 7.0
│ ... ...
│ step_size (chain, draw) float64 16kB 0.2823 0.2823 ... 0.3354
│ energy (chain, draw) float64 16kB 1.656e+03 ... 1.654e+03
│ divergences (chain, draw) int64 16kB 0 0 0 0 0 0 ... 0 0 0 0 0 0
│ diverging (chain, draw) bool 2kB False False ... False False
│ reached_max_treedepth (chain, draw) bool 2kB False False ... False False
│ perf_counter_diff (chain, draw) float64 16kB 0.007138 ... 0.003213
│ Attributes:
│ created_at: 2026-06-30T13:13:24.650169+00:00
│ creation_library: ArviZ
│ creation_library_version: 1.1.0
│ creation_library_language: Python
│ inference_library: pymc
│ inference_library_version: 6.0.1
│ sample_dims: ['chain', 'draw']
│ sampling_time: 8.305341958999634
│ tuning_steps: 500
│ modeling_interface: bambi
│ modeling_interface_version: 0.18.0
├── Group: /observed_data
│ 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: 2026-06-30T13:13:24.651921+00:00
│ creation_library: ArviZ
│ creation_library_version: 1.1.0
│ creation_library_language: Python
│ inference_library: pymc
│ inference_library_version: 6.0.1
│ sample_dims: []
│ modeling_interface: bambi
│ modeling_interface_version: 0.18.0
└── Group: /log_likelihood
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 -1.07 -0.9731 ... -1.182
Attributes:
modeling_interface: bambi
modeling_interface_version: 0.18.0Note that none of the parameters attached in the posterior have the observation coordinate (__obs__). However, looking at the model graph, we clearly see that for variables that are represented as regression targets, we in fact compute intermediate deterministics which do have dimensionality __obs__.
model.graph()
Computing Trial Wise Parameters¶
We can use the add_likelihood_parameters_to_datatree function to recompute and include in our idata the trial wise deterministics, which are part of our model graph.
### Attach Trial Wise Parameters
idata_trialwise = model.add_likelihood_parameters_to_datatree(idata)
### Checking the `idata` object
idata_trialwise
<xarray.DataTree>
Group: /
├── Group: /posterior
│ Dimensions: (chain: 4, draw: 500, v_condition_dim: 1, __obs__: 1000)
│ 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
│ * v_condition_dim (v_condition_dim) <U2 8B 'C2'
│ * __obs__ (__obs__) int64 8kB 0 1 2 3 4 5 ... 994 995 996 997 998 999
│ Data variables:
│ v_condition (chain, draw, v_condition_dim) float64 16kB 0.5306 ... 0...
│ z (chain, draw) float64 16kB 0.4887 0.4932 ... 0.4967 0.4828
│ v_Intercept (chain, draw) float64 16kB 0.6786 0.6577 ... 0.5924 0.6688
│ t (chain, draw) float64 16kB 0.01325 0.01503 ... 0.03315
│ a (chain, draw) float64 16kB 1.643 1.653 1.611 ... 1.54 1.577
│ v (chain, draw, __obs__) float64 16MB 0.6786 0.6786 ... 1.086
│ Attributes:
│ created_at: 2026-06-30T13:13:24.645662+00:00
│ creation_library: ArviZ
│ creation_library_version: 1.1.0
│ creation_library_language: Python
│ inference_library: pymc
│ inference_library_version: 6.0.1
│ sample_dims: ['chain', 'draw']
│ sampling_time: 8.305341958999634
│ tuning_steps: 500
│ modeling_interface: bambi
│ modeling_interface_version: 0.18.0
├── Group: /sample_stats
│ Dimensions: (chain: 4, draw: 500)
│ Coordinates:
│ * chain (chain) int64 32B 0 1 2 3
│ * draw (draw) int64 4kB 0 1 2 3 4 5 ... 495 496 497 498 499
│ Data variables: (12/18)
│ tree_depth (chain, draw) int64 16kB 4 2 4 3 3 4 ... 4 4 3 3 3 3
│ step_size_bar (chain, draw) float64 16kB 0.3508 0.3508 ... 0.3453
│ max_energy_error (chain, draw) float64 16kB -0.9847 ... -0.08732
│ smallest_eigval (chain, draw) float64 16kB nan nan nan ... nan nan
│ energy_error (chain, draw) float64 16kB -0.1505 ... 0.06108
│ n_steps (chain, draw) float64 16kB 15.0 3.0 15.0 ... 7.0 7.0
│ ... ...
│ step_size (chain, draw) float64 16kB 0.2823 0.2823 ... 0.3354
│ energy (chain, draw) float64 16kB 1.656e+03 ... 1.654e+03
│ divergences (chain, draw) int64 16kB 0 0 0 0 0 0 ... 0 0 0 0 0 0
│ diverging (chain, draw) bool 2kB False False ... False False
│ reached_max_treedepth (chain, draw) bool 2kB False False ... False False
│ perf_counter_diff (chain, draw) float64 16kB 0.007138 ... 0.003213
│ Attributes:
│ created_at: 2026-06-30T13:13:24.650169+00:00
│ creation_library: ArviZ
│ creation_library_version: 1.1.0
│ creation_library_language: Python
│ inference_library: pymc
│ inference_library_version: 6.0.1
│ sample_dims: ['chain', 'draw']
│ sampling_time: 8.305341958999634
│ tuning_steps: 500
│ modeling_interface: bambi
│ modeling_interface_version: 0.18.0
├── Group: /observed_data
│ 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: 2026-06-30T13:13:24.651921+00:00
│ creation_library: ArviZ
│ creation_library_version: 1.1.0
│ creation_library_language: Python
│ inference_library: pymc
│ inference_library_version: 6.0.1
│ sample_dims: []
│ modeling_interface: bambi
│ modeling_interface_version: 0.18.0
└── Group: /log_likelihood
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 -1.07 -0.9731 ... -1.182
Attributes:
modeling_interface: bambi
modeling_interface_version: 0.18.0Note how the idata include the v parameter which has the __obs__ coordinate. Computing these trial-wise parameters, can
serve us well for plotting and other post-hoc analyses, here we would otherwise struggle to be effective without incurring some pre-computation pain (which in addition would be error prone). See e.g. the tutorial on hierarchical Variational Inference for an example.