Migrate to ArviZ >= 1.0 - #346
Conversation
ArviZ 1.x replaces InferenceData with xarray.DataTree, renames stats arguments (hdi_prob -> prob/ci_prob), and moves plotting to the new arviz-plots PlotCollection API. The default (non-science) test suite, lint, and docs build all pass against the refreshed resolution (arviz 1.3.0, jax 0.11.1, numpyro 0.21.0) on Python 3.12.
tests/arviz_utils.py concentrates all knowledge of the new ArviZ return schemas in one place: hdi_bounds() wraps arviz_stats.hdi (raw 1-D input returns [lower, upper]) and save_posterior_plot() replaces the 0.x az.plot_posterior + plt.savefig idiom via az.plot_dist/PlotCollection. tests/test_arviz_utils.py runs in the default CI job and acts as a runtime guard on those schemas.
Mechanical rewrite of all 24 arviz-using files in tests/test_science/: - az.hdi(draws, hdi_prob=...) + ["x"].sel(hdi="lower"/"higher") (15 sites) -> hdi_bounds(draws[, prob=...]) from tests.arviz_utils. - az.plot_posterior + plt.savefig + plt.close (34 sites) -> save_posterior_plot(...), preserving output paths, probs and ref values exactly. Assertion semantics are unchanged (arviz 1.3 hdi returns bit-identical bounds to 0.x on the same draws). Verified on a representative subset (test_lti_gaussian, test_hmm, test_hierarchical_svi, test_discrete_time_l63_svi_pf): all arviz-related paths pass and the figures render. The remaining failures in that subset (test_hmm's numpy 2.5 int()-on-array in dynestyx plotting_utils, and test_discrete_time_l63_svi_pf's rho-recovery assertion) reproduce identically on unmodified main with the same freshly-resolved environment and are unrelated to this migration.
- az.plot_posterior -> az.plot_dist with explicit ci_kind="hdi" /
ci_prob and true-value reference lines; multi-panel ax=axes[i]
figures become native per-variable facets (11, 11c).
- az.style.use("arviz-white") -> "arviz-variat" (the 0.x style no
longer exists in ArviZ 1.x).
- quick_example.ipynb additionally needed updating to the current
dynestyx API (predict_times kwarg and (sample, 1, T, dim) output
indexing); it predated both and could not execute on main.
- All notebooks re-executed end-to-end on the arviz 1.3 environment
(jupytext round-trip, outputs regenerated; stale ArviZ
FutureWarnings and other users' local paths purged from stored
outputs).
11b_missing_observations_latent_path_mcmc.ipynb follows in a separate
commit: its migrated source is ready but its ~5 latent-path NUTS runs
take hours to re-execute.
- az.from_dict(posterior={...}) keyword form (TypeError in 1.x) ->
az.from_dict({"posterior": {...}}).
- az.ess(..., method="bulk") usage is unchanged in ArviZ 1.x and kept.
- 01_numpyro_bayesian_workflow: the list-valued ref_val plot becomes
az.plot_dist with per-variable reference lines, and
plot_forest(ess=True) (parameter removed in 1.x) becomes a nested
50%/94% HDI forest plot plus az.summary(kind="diagnostics") for
ESS/R-hat; adjacent prose and the ArviZ docs link updated.
- All three re-executed end-to-end on the arviz 1.3 environment.
LLMs like to do an annoying thing with formatting progress bars, so I reran each notebook.
DanWaxman
left a comment
DanWaxman
left a comment
There was a problem hiding this comment.
Thanks so much Juan!! Just a few small comments below.
git is also horrifically bad at notebooks, and LLMs format progress bars in notebooks extra weirdly. I reran the changed notebooks manually and opened juanitorduz#2 to fix these.
I think things are ready barring these.
|
Thanks @DanWaxman ! I merged your PR and addressed the feedback :) |
|
Great!! As per the PR description, the |
@DanWaxman I have tried running this notebook, and it takes a really long time (more than 2h). Do you wanna run it yourself? |
Reruns the 11b_missing_observations_latent_path_mcmc notebook. Also changes the layout of the final reconstructed missing observations plot to be more legible.
Hmm, that's weird... I don't remember this taking a long time. I think @mattlevine22 wrote that notebook and also didn't have such issues. Indeed, I just reran it on my laptop and it took ~1 minute end-to-end. I guess this leaves two things:
|
|
Hey @juanitorduz just wanted to check in and see if you were able to look at the PR (but no rush :)). |
Hey @DanWaxman ....I had to disconnect during the weekend :)
|
Sure no worries! Thanks so much for taking the time on the PR, it looks good to me now! |
|
Here is the issue #353 |
2 participants
Closes #345
Migrates dynestyx from
arviz>=0.18,<1.0to ArviZ >= 1.3 (the new arviz-base / arviz-stats / arviz-plots stack). The core library has zero arviz usage, so the migration surface is the science tests and the docs notebooks. One commit per milestone:tests/arviz_utils.pyconcentrates the new ArviZ schemas in one module (hdi_boundswrapsarviz_stats.hdi;save_posterior_plotreplaces theplot_posterior+savefigidiom viaaz.plot_dist/PlotCollection).tests/test_arviz_utils.pyruns in default CI as a runtime schema guard.plot_posterior→plot_distwith explicitci_kind="hdi", multi-panelax=figures → native facets,arviz-white→arviz-variat; all re-executed end-to-end.from_dictsignature;plot_forest(ess=True)→ nested-HDI forest +az.summary(kind="diagnostics"); re-executed (stale FutureWarnings and leaked local paths purged from stored outputs).Verification
scripts/lint.sh(ty + ruff + format): clean.mkdocs build: clean.scripts/test.sh): 696 passed, 2 xfailed.test_lti_gaussian,test_hmm,test_hierarchical_svi,test_discrete_time_l63_svi_pf, serial): all arviz code paths pass and figures render. The remaining 20 files use the identical helper pattern.Known issues (pre-existing, reproduced on unmodified
mainwith the same freshly-resolved environment)test_hmmfails in coreplotting_utils.py(int()on a 1-element array, removed in numpy 2.5).test_discrete_time_l63_svi_pffails its rho-recovery assertion under the refreshed jax stack.test_hierarchical_sviis load-flaky (passes in isolation).These are environment/core issues outside this PR's scope (no core changes here). The full serial science suite (
scripts/test_full.sh, multi-hour) was deliberately not run for this PR.Follow-up on this branch
11b_missing_observations_latent_path_mcmc.ipynbis migrated and currently re-executing (~5 latent-path NUTS runs, multi-hour); it will be pushed to this branch as a follow-up commit when execution completes. Also of note:quick_example.ipynbpredated the currentpredict_timesAPI and could not execute on main; it was updated as part of re-execution.