Converting NumPyro objects to DataTree

Contents

Converting NumPyro objects to DataTree#

DataTree is the data format ArviZ relies on.

This page covers multiple ways to generate a DataTree from NumPyro MCMC and SVI objects.

See also

Conversion from Python, numpy or pandas objects
DataTree for Exploratory Analysis of Bayesian Models for an overview of InferenceData and its role within ArviZ.

We will start by importing the required packages and defining the model. The famous 8 school model.

import arviz_base as az
import numpy as np
from numpy.typing import ArrayLike

import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS, SVI, Trace_ELBO, autoguide, Predictive
from jax import random
import jax.numpy as jnp

J = 8
y_obs = np.array([28.0, 8.0, -3.0, 7.0, -1.0, 1.0, 18.0, 12.0])
sigma = np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0])

def eight_schools_model(J, sigma, y=None):
    mu = numpyro.sample("mu", dist.Normal(0, 5))
    tau = numpyro.sample("tau", dist.HalfCauchy(5))
    with numpyro.plate("J", J):
        eta = numpyro.sample("eta", dist.Normal(0, 1))
        theta = numpyro.deterministic("theta", mu + tau * eta)
        return numpyro.sample("obs", dist.Normal(theta, sigma), obs=y)
    

def eight_schools_custom_guide(J, sigma, y=None):

    # Variational parameters for mu
    mu_loc = numpyro.param("mu_loc", 0.0)
    mu_scale = numpyro.param("mu_scale", 1.0, constraint=dist.constraints.positive)
    mu = numpyro.sample("mu", dist.Normal(mu_loc, mu_scale))

    # Variational parameters for tau (positive support)
    tau_loc = numpyro.param("tau_loc", 1.0)
    tau_scale = numpyro.param("tau_scale", 0.5, constraint=dist.constraints.positive)
    tau = numpyro.sample("tau", dist.LogNormal(jnp.log(tau_loc), tau_scale))

    # Variational parameters for eta
    eta_loc = numpyro.param("eta_loc", jnp.zeros(J))
    eta_scale = numpyro.param("eta_scale", jnp.ones(J), constraint=dist.constraints.positive)
    with numpyro.plate("J", J):
        eta = numpyro.sample("eta", dist.Normal(eta_loc, eta_scale))

        # Deterministic transform
        numpyro.deterministic("theta", mu + tau * eta)