[pymc6 migration] Fix all tutorials

[digicosmos86]

HSSM v0.4.0 Notebook Update Report

Breaking Changes Addressed

Based on the v0.4.0 changelog, the following breaking changes required updates:

1. model.summary() and model.plot_trace() methods removed → replaced with az.summary() and az.plot_trace_dist()
2. model.sample() now returns xarray.DataTree instead of arviz.InferenceData
3. ArViZ 1.0+ API changes: az.plot_posterior → az.plot_dist, az.plot_trace → az.plot_trace_dist
4. ArViZ 1.0+ no longer supports kwargs like ref_val, color, hdi_prob, ax, compact, kind, marginals

---

Successfully Updated and Executed Notebooks

Notebook	Changes Made	Status
likelihoods.ipynb	Added import arviz as az, changed model.summary() → az.summary(sample), model.plot_trace() → az.plot_trace_dist(sample)	✅ Success
plotting.ipynb	Updated docs, changed model.summary() → az.summary(cav_data_traces), model.plot_trace() → az.plot_trace_dist(cav_data_traces) + plt.tight_layout()	✅ Success
lapse_prob_and_dist.ipynb	Added import arviz as az, changed model.summary(lapse_trace) → az.summary(lapse_trace), model.plot_trace(lapse_trace) → az.plot_trace_dist(lapse_trace)	✅ Success
tutorial_stim_coding.ipynb	Changed 4 instances of model.summary() → az.summary(model.traces)	✅ Success
initial_values.ipynb	No changes needed	✅ Success
compile_logp.ipynb	No changes needed	✅ Success
pymc.ipynb	No changes needed	✅ Success
rlssm_quickstart.ipynb	No changes needed	✅ Success
save_load_tutorial.ipynb	Use traces.to_netcdf() and vi_traces.to_netcdf() directly; load with az.from_netcdf()	✅ Success
do_operator.ipynb	Changed predictive_idata_to_dataframe → predictive_dt_to_dataframe	✅ Success
choice_only_models.ipynb	az.plot_posterior → az.plot_dist, removed ref_val/ref_val_color; az.plot_trace → az.plot_trace_dist	✅ Success
tutorial_bayesian_t_test.ipynb	az.plot_posterior → az.plot_dist, removed ax, color, hdi_prob params	✅ Success
hmm_ddm_regime_switching.ipynb	az.plot_trace → az.plot_trace_dist, removed compact; az.plot_pair removed kind, marginals, figsize, point_estimate, reference_values, reference_values_kwargs	✅ Success

---

Notebooks with Issues

Issues Fixed but Could Not Execute (Timeout/Kernel Issues)

Notebook	Issue
hssm_tutorial_workshop_1.ipynb	Fixed deprecated calls; kernel hssm519 not found
hssm_tutorial_workshop_2.ipynb	Fixed deprecated calls; execution timed out

Deeper API Changes Required (Not Simple Method Replacements)

Notebook	Error	Root Cause
tutorial_trial_wise_parameters.ipynb	'HSSM' object has no attribute 'add_likelihood_parameters_to_idata'	Method removed in v0.4.0
rlssm_tutorial.ipynb	cannot import name 'decorate_atomic_simulator' from 'hssm.utils'	Utility function removed
rlssm_rlwm_model.ipynb	cannot import name 'decorate_atomic_simulator' from 'hssm.utils'	Utility function removed
add_custom_rlssm_model.ipynb	cannot import name 'decorate_atomic_simulator' from 'hssm.utils'	Utility function removed
main_tutorial.ipynb	'DataTree' object has no attribute 'stack'	DataTree doesn't have InferenceData's stack() method
variational_inference_hierarchical.ipynb	No model on context stack	PyMC/Bambi API change
scientific_workflow_hssm.ipynb	max_subplots exceeds limit (88 vs 40)	Too many variables for pair plot - requires var_names limiting

External Dependency Issues (Not HSSM Code Issues)

Notebook	Error	Root Cause
tutorial_bayeux.ipynb	No module named 'bayeux'	Optional dependency not installed
bayesflow_lre_integration.ipynb	TensorFlow/Keras model loading error	Model format incompatibility
blackbox_contribution_onnx_example.ipynb	cannot pickle 'onnxruntime...InferenceSession'	ONNX runtime issue
jax_callable_contribution_onnx_example.ipynb	max_subplots exceeds limit	Matplotlib config issue
poisson_race.ipynb	Style arviz-whitegrid not found	Matplotlib style removed

Execution Timeouts

Notebook	Note
pymc_to_hssm.ipynb	Timed out
tutorial_p_outlier_regression.ipynb	Timed out
variational_inference.ipynb	Timed out

---

Summary

• 14 notebooks successfully updated and executed
• 6 notebooks require deeper API fixes or have execution issues
• 5 notebooks have external dependency issues (not code issues)
• 3 notebooks timed out

Notes for Fixing Remaining Notebooks

ArViZ 1.0+ API Changes

• az.plot_posterior() → az.plot_dist()
• az.plot_trace() → az.plot_trace_dist()
• Parameters like ref_val, ref_val_color, color, ax, hdi_prob, compact, kind, marginals, figsize, point_estimate, reference_values, reference_values_kwargs are no longer supported in the new ArViZ plots API

DataTree API Compatibility

• xarray.DataTree doesn't have all the methods of arviz.InferenceData (e.g., stack(), to_netcdf() on az module)
• For saving: Use datatree.to_netcdf(path) directly on the DataTree returned by model.sample()
• For loading: Use az.from_netcdf(path) which returns a DataTree

Removed Utility Functions

• hssm.utils.predictive_idata_to_dataframe → hssm.utils.predictive_dt_to_dataframe
• decorate_atomic_simulator has been removed from hssm.utils

External Dependencies

Some notebooks require optional dependencies (e.g., bayeux, bayesflow, onnxruntime) that may not be installed in all environments.

matplotlib max_subplots

Some notebooks generate too many subplots and hit the rcParams['plot.max_subplots']=40 limit. Options:

1. Limit var_names to fewer variables
2. Increase the limit: plt.rcParams['plot.max_subplots'] = 100

[linear-code]

HSSM-631

[digicosmos86]

@AlexanderFengler @krishnbera @cpaniaguam I did some LLM automation and fixed as many notebooks as I could. Please see the above reports on what still need to be updated. Why don't we just pick the ones that we think that we can easily fix, and then randomize the rest? Please feel free to create issues on whatever tutorials that you want to take