Update PyRenew tutorials for ArviZ 1.0

[cdc-mitzimorris]

cf: #737 (comment)

change PyRenew tutorials that use ArviZ to diagnose/visualize NumPyro samples to use az.from_numpyro everywhere so that NumPyro names don't run afoul of xarray.DataTree name restrictions.

[cdc-mitzimorris]

note: there is no way to write tutorials that will run under ArviZ 0.23 and ArviZ 1.0.0.
updating to ArviZ 1.0.0 for the PyRenew package only affects the tutorials - I will create a branch and try things out; this would require merging PR #737 first.

[dylanhmorris]

I agree that depending on ArviZ >= 1.0.0 is the way forward.

[cdc-mitzimorris]

ArviZ 1.0 doesn't seem to have support for NumPyro hierarchical models - cf: arviz-devs/arviz-base#17

When updating the "Building Multisignal Models" tutorial, the call az.from_numpyro(model.mcmc) blows up:

ValueError: Given variables have names containing the '[/](http://localhost:8888/)' character: ['latent_infections[/I0_init_all_subpops_dim_0](http://localhost:8888/I0_init_all_subpops_dim_0)', 'latent_infections[/I0_init_all_subpops_dim_1](http://localhost:8888/I0_init_all_subpops_dim_1)', 'latent_infections[/I0_init_all_subpops](http://localhost:8888/I0_init_all_subpops)', 'latent_infections[/infections_aggregate](http://localhost:8888/infections_aggregate)', 'latent_infections[/log_rt_baseline](http://localhost:8888/log_rt_baseline)', 'latent_infections[/rt_baseline](http://localhost:8888/rt_baseline)', 'latent_infections[/rt_subpop](http://localhost:8888/rt_subpop)', 'latent_infections[/subpop_deviations](http://localhost:8888/subpop_deviations)']. Variables stored in DataTree objects cannot have names containing '[/](http://localhost:8888/)' characters, as this would make path-like access to variables ambiguous.

[dylanhmorris]

@cdc-mitzimorris I think the problem is here:

PyRenew/pyrenew/latent/hierarchical_infections.py

Line 265 in 9134324

numpyro.deterministic("rt_baseline", jnp.exp(log_rt_baseline))

By default, numpyro.handlers.scope uses a / as a separator. In #727, I'm planning to enforce _ instead.

Quick fix: use _ as a separator in that scope call and see if this resolves.

Longer term, from_numpyro() probably should support numpyro sites with / in their names given the default behavior of numpyro.handlers.scope(). Could you make a minimal reprex and open an issue in arviz?

[cdc-mitzimorris]

Quick fix: use _ as a separator in that scope call and see if this resolves

that seems reasonable. the DAG structure of the model shouldn't be anything that ArviZ knows or cares about.

[cdc-mitzimorris]

@cdc-mitzimorris I think the problem is here:

By default, numpyro.handlers.scope uses a / as a separator. In #727, I'm planning to enforce _ instead.

Quick fix: use _ as a separator in that scope call and see if this resolves.

I don't think that using _ is a good way to go here - since names often contain underscore, this flattens the structure - is that what you want?

Per discussion with the ArviZ devs, they suggested using :: instead.

The fix goes into file pyrenew/latent/hierarchical_infections.py, line 262:

with numpyro.handlers.scope(prefix=self.name, divider="::"):

With this, az.from_numpyro just works. I returns an xarray.DataTree object. Here are the changes necessary to get this working for PyRenew:

main...mem_msm_tutorial_edits

[dylanhmorris]

I like the idea of ::. Let's go with that. I will adjust #727 to use it.

UPDATE: ea699ff

[dylanhmorris]

@cdc-mitzimorris open a PR when ready?

[cdc-mitzimorris]

opened #754