Saving and loading models
Saving and loading models¶
In this short how-to, tutorial, we show how to save a HSSM model instance and its inference results to disk and then re-instantiate the model from the saved files.
Load data and instantiate HSSM model¶
In [1]:
Copied!
import hssm
cav_data = hssm.load_data("cavanagh_theta")
basic_hssm_model = hssm.HSSM(
data=cav_data,
process_initvals=True,
link_settings="log_logit",
model="angle",
include=[
{
"name": "v",
"formula": "v ~ 1 + C(stim)",
}
],
)
import hssm
cav_data = hssm.load_data("cavanagh_theta")
basic_hssm_model = hssm.HSSM(
data=cav_data,
process_initvals=True,
link_settings="log_logit",
model="angle",
include=[
{
"name": "v",
"formula": "v ~ 1 + C(stim)",
}
],
)
Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.
Model initialized successfully.
In [2]:
Copied!
traces = basic_hssm_model.sample(sampler="numpyro", tune=100, draws=100, chains=2)
traces = basic_hssm_model.sample(sampler="numpyro", tune=100, draws=100, chains=2)
Using default initvals.
/Users/yxu150/HSSM/.venv/lib/python3.13/site-packages/pytensor/gradient.py:1327: FutureWarning: LANLogpOp should implement `pullback` instead of `L_op`/`grad`. Direct `L_op`/`grad` implementations are deprecated and will stop being called in a future version. input_grads = node.op.pullback(inputs, node.outputs, new_output_grads) /Users/yxu150/HSSM/.venv/lib/python3.13/site-packages/bambi/backend/pymc.py:224: UserWarning: `init='adapt_diag'` is ignored by `nuts_sampler='numpyro'`; the external sampler uses its own initialization. idata = pm.sample( NUTS[numpyro]: [t, z, a, theta, v_Intercept, v_C(stim)]
/Users/yxu150/HSSM/.venv/lib/python3.13/site-packages/pymc/sampling/jax.py:463: UserWarning: There are not enough devices to run parallel chains: expected 2 but got 1. Chains will be drawn sequentially. If you are running MCMC in CPU, consider using `numpyro.set_host_device_count(2)` at the beginning of your program. You can double-check how many devices are available in your system using `jax.local_device_count()`. pmap_numpyro = MCMC(
0%| | 0/200 [00:00<?, ?it/s]
warmup: 5%|▌ | 10/200 [00:00<00:15, 12.57it/s, 15 steps of size 1.79e-02. acc. prob=0.60]
warmup: 10%|█ | 20/200 [00:01<00:14, 12.22it/s, 63 steps of size 3.76e-03. acc. prob=0.67]
warmup: 15%|█▌ | 30/200 [00:02<00:10, 16.25it/s, 7 steps of size 1.38e-02. acc. prob=0.72]
warmup: 20%|██ | 40/200 [00:02<00:08, 19.67it/s, 7 steps of size 6.04e-03. acc. prob=0.73]
warmup: 25%|██▌ | 50/200 [00:02<00:07, 19.61it/s, 7 steps of size 1.17e-02. acc. prob=0.74]
warmup: 30%|███ | 60/200 [00:03<00:07, 18.72it/s, 47 steps of size 6.16e-03. acc. prob=0.74]
warmup: 35%|███▌ | 70/200 [00:03<00:06, 20.18it/s, 7 steps of size 1.36e-02. acc. prob=0.75]
warmup: 40%|████ | 80/200 [00:04<00:05, 23.96it/s, 7 steps of size 1.82e-02. acc. prob=0.76]
warmup: 45%|████▌ | 90/200 [00:04<00:04, 25.93it/s, 7 steps of size 1.22e-02. acc. prob=0.76]
warmup: 50%|█████ | 100/200 [00:09<00:18, 5.45it/s, 35 steps of size 2.61e-02. acc. prob=0.76]
sample: 55%|█████▌ | 110/200 [00:11<00:16, 5.39it/s, 63 steps of size 2.61e-02. acc. prob=0.94]
sample: 60%|██████ | 120/200 [00:13<00:16, 4.88it/s, 127 steps of size 2.61e-02. acc. prob=0.94]
sample: 65%|██████▌ | 130/200 [00:16<00:15, 4.61it/s, 127 steps of size 2.61e-02. acc. prob=0.94]
sample: 70%|███████ | 140/200 [00:19<00:14, 4.06it/s, 95 steps of size 2.61e-02. acc. prob=0.94]
sample: 75%|███████▌ | 150/200 [00:22<00:12, 3.93it/s, 127 steps of size 2.61e-02. acc. prob=0.95]
sample: 80%|████████ | 160/200 [00:24<00:10, 3.84it/s, 127 steps of size 2.61e-02. acc. prob=0.95]
sample: 85%|████████▌ | 170/200 [00:27<00:07, 4.05it/s, 11 steps of size 2.61e-02. acc. prob=0.95]
sample: 90%|█████████ | 180/200 [00:30<00:05, 3.74it/s, 127 steps of size 2.61e-02. acc. prob=0.95]
sample: 95%|█████████▌| 190/200 [00:32<00:02, 3.74it/s, 31 steps of size 2.61e-02. acc. prob=0.96]
sample: 100%|██████████| 200/200 [00:35<00:00, 3.80it/s, 63 steps of size 2.61e-02. acc. prob=0.96]
sample: 100%|██████████| 200/200 [00:35<00:00, 5.64it/s, 63 steps of size 2.61e-02. acc. prob=0.96]
0%| | 0/200 [00:00<?, ?it/s]
warmup: 5%|▌ | 10/200 [00:00<00:02, 66.25it/s, 15 steps of size 1.32e-02. acc. prob=0.59]
warmup: 10%|█ | 20/200 [00:00<00:08, 22.09it/s, 63 steps of size 3.74e-03. acc. prob=0.67]
warmup: 15%|█▌ | 30/200 [00:01<00:07, 23.34it/s, 15 steps of size 8.93e-03. acc. prob=0.71]
warmup: 20%|██ | 40/200 [00:01<00:06, 26.01it/s, 3 steps of size 2.55e-03. acc. prob=0.72]
warmup: 25%|██▌ | 50/200 [00:02<00:06, 23.77it/s, 15 steps of size 5.50e-03. acc. prob=0.74]
warmup: 30%|███ | 60/200 [00:02<00:05, 26.38it/s, 15 steps of size 3.42e-03. acc. prob=0.74]
warmup: 35%|███▌ | 70/200 [00:02<00:05, 23.02it/s, 39 steps of size 9.97e-03. acc. prob=0.75]
warmup: 40%|████ | 80/200 [00:03<00:04, 25.99it/s, 7 steps of size 1.92e-02. acc. prob=0.76]
warmup: 45%|████▌ | 90/200 [00:03<00:04, 26.09it/s, 7 steps of size 2.36e-02. acc. prob=0.76]
warmup: 50%|█████ | 100/200 [00:04<00:05, 17.11it/s, 3 steps of size 1.87e-01. acc. prob=0.76]
sample: 55%|█████▌ | 110/200 [00:04<00:04, 18.95it/s, 7 steps of size 1.87e-01. acc. prob=0.92]
sample: 60%|██████ | 120/200 [00:05<00:04, 18.75it/s, 15 steps of size 1.87e-01. acc. prob=0.93]
sample: 65%|██████▌ | 130/200 [00:06<00:03, 17.51it/s, 31 steps of size 1.87e-01. acc. prob=0.94]
sample: 70%|███████ | 140/200 [00:06<00:03, 18.25it/s, 23 steps of size 1.87e-01. acc. prob=0.93]
sample: 75%|███████▌ | 150/200 [00:07<00:02, 19.92it/s, 23 steps of size 1.87e-01. acc. prob=0.92]
sample: 80%|████████ | 160/200 [00:07<00:01, 20.77it/s, 15 steps of size 1.87e-01. acc. prob=0.92]
sample: 85%|████████▌ | 170/200 [00:08<00:01, 20.03it/s, 15 steps of size 1.87e-01. acc. prob=0.93]
sample: 90%|█████████ | 180/200 [00:08<00:01, 19.35it/s, 15 steps of size 1.87e-01. acc. prob=0.92]
sample: 95%|█████████▌| 190/200 [00:09<00:00, 19.36it/s, 31 steps of size 1.87e-01. acc. prob=0.92]
sample: 100%|██████████| 200/200 [00:09<00:00, 21.60it/s, 15 steps of size 1.87e-01. acc. prob=0.92]
sample: 100%|██████████| 200/200 [00:09<00:00, 21.19it/s, 15 steps of size 1.87e-01. acc. prob=0.92]
We recommend running at least 4 chains for robust computation of convergence diagnostics
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/yxu150/HSSM/.venv/lib/python3.13/site-packages/pytensor/link/numba/dispatch/basic.py:214: UserWarning: Numba will use object mode to run LANLogpOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details. warnings.warn(
VI¶
In [3]:
Copied!
vi_traces = basic_hssm_model.vi(method="advi", niter=5000)
vi_traces = basic_hssm_model.vi(method="advi", niter=5000)
Using MCMC starting point defaults.
/Users/yxu150/HSSM/.venv/lib/python3.13/site-packages/pytensor/gradient.py:1327: FutureWarning: LANLogpOp should implement `pullback` instead of `L_op`/`grad`. Direct `L_op`/`grad` implementations are deprecated and will stop being called in a future version. input_grads = node.op.pullback(inputs, node.outputs, new_output_grads)
/Users/yxu150/HSSM/.venv/lib/python3.13/site-packages/pytensor/link/numba/dispatch/basic.py:214: UserWarning: Numba will use object mode to run LANLogpOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details. warnings.warn( /Users/yxu150/HSSM/.venv/lib/python3.13/site-packages/pytensor/link/numba/dispatch/basic.py:214: UserWarning: Numba will use object mode to run LANLogpVJPOp's perform method. Set `pytensor.config.compiler_verbose = True` to see more details. warnings.warn(
Finished [100%]: Average Loss = 7,721.6
Saving and Loading the model¶
In [4]:
Copied!
import os
# Create directory for model files
model_dir = "hssm_models/test_model"
os.makedirs(model_dir, exist_ok=True)
# Save traces using xarray's to_netcdf (v0.4.0 returns DataTree)
traces.to_netcdf(f"{model_dir}/traces.nc")
vi_traces.to_netcdf(f"{model_dir}/vi_traces.nc")
import os
# Create directory for model files
model_dir = "hssm_models/test_model"
os.makedirs(model_dir, exist_ok=True)
# Save traces using xarray's to_netcdf (v0.4.0 returns DataTree)
traces.to_netcdf(f"{model_dir}/traces.nc")
vi_traces.to_netcdf(f"{model_dir}/vi_traces.nc")
We are using the defaults here, which saves the model traces to the hssm_models/test_model/ directory inside your curerent working directory.
Two files are saved in the model directory:
traces.nc: The MCMC traces.vi_traces.nc: The VI traces.
Note: In v0.4.0, model.sample() returns an xarray.DataTree instead of arviz.InferenceData, so we save traces directly using xarray's to_netcdf() method.
We can now load the model traces from the directory we just created.
In [5]:
Copied!
import arviz as az
# Load traces using arviz (v0.4.0 compatible)
loaded_traces = az.from_netcdf(f"{model_dir}/traces.nc")
loaded_vi_traces = az.from_netcdf(f"{model_dir}/vi_traces.nc")
loaded_traces
import arviz as az
# Load traces using arviz (v0.4.0 compatible)
loaded_traces = az.from_netcdf(f"{model_dir}/traces.nc")
loaded_vi_traces = az.from_netcdf(f"{model_dir}/vi_traces.nc")
loaded_traces
Out[5]:
<xarray.DataTree>
Group: /
├── Group: /posterior
│ Dimensions: (chain: 2, draw: 100, v_C(stim)_dim: 2)
│ Coordinates:
│ * chain (chain) int64 16B 0 1
│ * draw (draw) int64 800B 0 1 2 3 4 5 6 7 ... 92 93 94 95 96 97 98 99
│ * v_C(stim)_dim (v_C(stim)_dim) <U2 16B 'WL' 'WW'
│ Data variables:
│ v_Intercept (chain, draw) float64 2kB ...
│ t (chain, draw) float64 2kB ...
│ v_C(stim) (chain, draw, v_C(stim)_dim) float64 3kB ...
│ z (chain, draw) float64 2kB ...
│ a (chain, draw) float64 2kB ...
│ theta (chain, draw) float64 2kB ...
│ Attributes:
│ created_at: 2026-06-29T19:03:03.541744+00:00
│ creation_library: ArviZ
│ creation_library_version: 1.1.0
│ creation_library_language: Python
│ sample_dims: ['chain', 'draw']
│ inference_library: numpyro
│ inference_library_version: 0.21.0
│ sampling_time: 46.577575
│ tuning_steps: 100
│ modeling_interface: bambi
│ modeling_interface_version: 0.18.0
├── Group: /sample_stats
│ Dimensions: (chain: 2, draw: 100)
│ Coordinates:
│ * chain (chain) int64 16B 0 1
│ * draw (draw) int64 800B 0 1 2 3 4 5 6 7 ... 93 94 95 96 97 98 99
│ Data variables:
│ acceptance_rate (chain, draw) float64 2kB ...
│ step_size (chain, draw) float64 2kB ...
│ diverging (chain, draw) bool 200B ...
│ energy (chain, draw) float64 2kB ...
│ n_steps (chain, draw) int64 2kB ...
│ tree_depth (chain, draw) int64 2kB ...
│ lp (chain, draw) float64 2kB ...
│ Attributes:
│ created_at: 2026-06-29T19:03:03.545554+00:00
│ creation_library: ArviZ
│ creation_library_version: 1.1.0
│ creation_library_language: Python
│ sample_dims: ['chain', 'draw']
│ modeling_interface: bambi
│ modeling_interface_version: 0.18.0
├── Group: /observed_data
│ Dimensions: (__obs__: 3988, rt,response_extra_dim_0: 2)
│ Coordinates:
│ * __obs__ (__obs__) int64 32kB 0 1 2 3 ... 3985 3986 3987
│ * 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 64kB ...
│ Attributes:
│ created_at: 2026-06-29T19:03:03.546074+00:00
│ creation_library: ArviZ
│ creation_library_version: 1.1.0
│ creation_library_language: Python
│ sample_dims: []
│ modeling_interface: bambi
│ modeling_interface_version: 0.18.0
├── Group: /constant_data
│ Attributes:
│ created_at: 2026-06-29T19:03:03.546133+00:00
│ creation_library: ArviZ
│ creation_library_version: 1.1.0
│ creation_library_language: Python
│ sample_dims: []
│ modeling_interface: bambi
│ modeling_interface_version: 0.18.0
└── Group: /log_likelihood
Dimensions: (chain: 2, draw: 100, __obs__: 3988)
Coordinates:
* chain (chain) int64 16B 0 1
* draw (draw) int64 800B 0 1 2 3 4 5 6 7 8 ... 92 93 94 95 96 97 98 99
* __obs__ (__obs__) int64 32kB 0 1 2 3 4 5 ... 3983 3984 3985 3986 3987
Data variables:
rt,response (chain, draw, __obs__) float64 6MB ...
Attributes:
modeling_interface: bambi
modeling_interface_version: 0.18.0With this simple workflow your models are portable across sessions and machines.