[2]{.chapter-number}  [Working with InferenceData]{.chapter-title}

2 Working with InferenceData

WARNING (pytensor.tensor.blas): Using NumPy C-API based implementation for BLAS functions.

2.1 InferenceData

During a modern Bayesian analysis we usually generate many sets of data including posterior samples and posterior predictive samples. But we also have observed data, and statistics generated by the sampling method, samples from the prior and/or prior predictive distribution, etc. All this data can be stored in an InferenceData object, to keep them tidy and avoid confusion each one of these sets is assigned to a different group. For instance, the posterior samples are stored in the posterior group. The observed data is stored in the observed_data group and so on. InferenceData is based on xarray . If you are not familiar with xarray imagine NumPy multidimensional arrays but with labels! This makes many operations easier as you don’t have to remember the order of the dimensions. Each group in an InferenceData is a DataSet and Datasets are dictionary-like containers of DataArrays. You can read more about the terminology used in xarray here

The best way to understand InferenceData is to explore it. ArviZ comes equipped with a few InferenceData objects so we can start playing with them even without the need to fit a model. Let’s start by loading the centered_eight InferenceData. For the moment we do not need to know about the details of the model.

idata = az.load_arviz_data("centered_eight")
idata

We can immediately see that we have many groups. This is an HTML representation of the InferenceData, so if you are reading this from a browser you should be able to interact with it. If you click on the posterior group you will see that we have three dimensions, with the names chain, draw, and school, you can think of dimensions as the axes of a plot. This means that the posterior samples were generated by running an MCMC sampler with 4 chains, each one of 500 draws. At least for one of the parameters in the posterior we have and additional dimension called school. If you click on coordinates you will be able to see the actual values that each dimension can take, like the integers [0, 1, 2, 3] for chain and the strings ['Choate', 'Deerfield', 'Phillips Andover', 'Phillips Exeter', 'Hotchkiss', 'Lawrenceville', "St. Paul's", 'Mt. Hermon'] for school, notice that we have an array of dtype=object. Furthermore, if you click on the symbol by the school coordinate, you will be able to see the names of each school.

2.1.1 Get the dataset corresponding to a single group

We can access each group using a dictionary-like notation:

idata["posterior"];

Alternatively, we can use the dot notation, as groups are attributes of the InferenceData object. For instance, to access the posterior group we can write:

idata.posterior

The dot notation works at the group level and for DataSets and DataArrays as long as there is no conflict with a method or attribute of these objects. If there is a conflict, you can always use the dictionary-like notation.

2.1.2 Get coordinate values

As we have seen, we have 8 schools in this InferenceData with their names. If we want to programmatically access the names we can do

idata.posterior.school.values

array(['Choate', 'Deerfield', 'Phillips Andover', 'Phillips Exeter',
       'Hotchkiss', 'Lawrenceville', "St. Paul's", 'Mt. Hermon'],
      dtype='<U16')

Which returns a DataSet with the names of the schools. To obtain a NumPy array we can do

idata.posterior.school.values

array(['Choate', 'Deerfield', 'Phillips Andover', 'Phillips Exeter',
       'Hotchkiss', 'Lawrenceville', "St. Paul's", 'Mt. Hermon'],
      dtype='<U16')

If we want to get the number of schools we can write:

len(idata.observed_data.school)

Notice that we do not need to first obtain the NumPy array and then compute the length.

2.1.3 Get a subset of chains

Because we have labels for the names of the schools we can use them to access their associated information. Labels are usually much easier to remember than numerical indices. For instance, to access the posterior samples of the school Choate we can write:

idata.posterior.sel(school="Choate")

The draw and chain coordinates are indexed using numbers, the following code will return the last draw from chain 1 and chain 2:

idata.posterior.sel(draw=499, chain=[1, 2])

Usually, in Bayesian statistics, we don’t need to access individual draws or chains, a more common operation is to select a range. For that purpose, we can use Python’s slice function. For example, the following returns the first 200 draws from all chains:

idata.posterior.sel(draw=slice(0, 200))

Using the slice function we can also remove the first 100 samples.

idata.posterior.sel(draw=slice(100, None))

We can apply the same operations to the entire InferenceData object. Can you anticipate the result?

idata.sel(draw=slice(100, None))

If you check the object you will see that the groups posterior, posterior_predictive, log_likelihood, sample_stats, prior, and prior_predictive have 400 draws compared to the original 500. The group observed_data has not been affected because it does not have the draw dimension. Alternatively, you can specify which group or groups you want to change.

idata.sel(draw=slice(100, None), groups="posterior")

2.1.4 Compute posterior mean

We can perform operations on the InferenceData object. For instance, to compute the mean of the first 200 draws we can write:

idata.posterior.sel(draw=slice(0, 200)).mean()

In NumPy, it is common to perform operations like this along a given axis. With InferecenData/xarray we can do the same by specifying the dimension along which we want to operate. For instance, to compute the mean along the draw dimension we can write:

idata.posterior.mean("draw")

This returns the mean for each chain and school. Can you anticipate how different this would be if the dimension was chain instead of draw? And what about if we use school?

We can also specify multiple dimensions. For instance, to compute the mean along the draw and chain dimensions we can write:

idata.posterior.mean(["chain", "draw"])

2.1.5 Combine chains and draws

Our primary goal is usually to obtain posterior samples and thus we aren’t concerned with chains and draws. In those cases, we can use the az.extract function. This combines the chain and draw into a sample coordinate which can make further operations easier. By default, az.extract works on the posterior, but you can specify other groups using the group argument.

az.extract(idata)

You can achieve the same result using idata.posterior.stack(sample=("chain", "draw")). But extract can be more flexible because it takes care of the most common subsetting operations with MCMC samples. It can:

Combine chains and draws
Return a subset of variables (with optional filtering with regular expressions or string matching)
Return a subset of samples. Moreover, by default, it returns a random subset to prevent getting non-representative samples due to bad mixing.
Access any group

To get a subsample we can specify the number of samples we want with the num_samples argument. For instance, to get 100 samples we can write:

az.extract(idata, num_samples=100);

If you need to extract subsets from multiple groups, you should use a random seed. This will ensure that subsamples match. For example, if you do

posterior = az.extract(idata, num_samples=100, rng=124)
ll = az.extract(idata, group="log_likelihood", num_samples=100, rng=124)

You can inspect the samples in the posterior and ll variables and see that they match.

2.2 Ploting

Xarray has some plotting capabilities, for instance, we can do:

idata.posterior["mu"].plot.hist(figsize=(9, 3));

But in most scenarios calling a plotting function from ArviZ and passing the InfereceData as an argument will be a much better idea.

2.3 Add a new variable

We can add variables to existing groups. For instance, we may want to transform a parameter from the posterior. Like computing and adding the \(\log\) of the parameter \(\tau\) to the posterior group.

posterior["log_tau"] = np.log(posterior["tau"])
posterior

2.4 Advance operations with InferenceData

Now we delve into more advanced operations with InferenceData. While these operations are not essential to use ArviZ, they can be useful in some cases. Exploring these advanced functionalities will help you become more familiar with InferenceData and provide additional insights that may enhance your overall experience with ArviZ.

2.4.1 Compute and store posterior pushforward quantities

We use “posterior pushforward quantities” to refer to quantities that are not variables in the posterior but deterministic computations using posterior variables.

You can use xarray for these pushforward operations and store them as a new variable in the posterior group. You’ll then be able to plot them with ArviZ functions, calculate stats and diagnostics on them (like mcse), or save and share the InferenceData object with the push forward quantities included.

The first thing we are going to do is to store the posterior group in a variable called post to make the code more readable. And to compute the log of \(\tau\).

post = idata.posterior
post["log_tau"] = np.log(post["tau"])

Compute the rolling mean of \(\log(\tau)\) with xarray.DataArray.rolling, storing the result in the posterior:

post["mlogtau"] = post["log_tau"].rolling({"draw": 50}).mean()

Using xarray for pushforward calculations has all the advantages of working with xarray. It also inherits the disadvantages of working with xarray, but we believe those to be outweighed by the advantages, and we have already shown how to extract the data as NumPy arrays. Working with InferenceData is working mainly with xarray objects and this is what is shown in this guide.

Some examples of these advantages are specifying operations with named dimensions instead of positional ones (as seen in some previous sections), automatic alignment and broadcasting of arrays (as we’ll see now), or integration with Dask (as shown in the dask_for_arviz guide).

In this cell, you will compute pairwise differences between schools on their mean effects (variable theta). To do so, subtract the variable theta after renaming the school dimension to the original variable. Xarray then aligns and broadcasts the two variables because they have different dimensions, and the result is a 4D variable with all the pointwise differences.

Eventually, store the result in the theta_school_diff variable. Notice that the theta_shool_diff variable in the posterior has kept the named dimensions and coordinates:

post["theta_school_diff"] = post.theta - post.theta.rename(school="school_bis")
post

Note

This same operation using NumPy would require manual alignment of the two arrays to make sure they broadcast correctly. The code could be something like:

theta_school_diff = theta[:, :, :, None] - theta[:, :, None, :]

2.4.2 Advanced subsetting

To select the value corresponding to the difference between the Choate and Deerfield schools do:

post["theta_school_diff"].sel(school="Choate", school_bis="Deerfield")

For more advanced subsetting (the equivalent to what is sometimes called “fancy indexing” in NumPy) you need to provide the indices as DataArray objects:

school_idx = xr.DataArray(["Choate", "Hotchkiss", "Mt. Hermon"], dims=["pairwise_school_diff"])
school_bis_idx = xr.DataArray(
    ["Deerfield", "Choate", "Lawrenceville"], dims=["pairwise_school_diff"]
)
post["theta_school_diff"].sel(school=school_idx, school_bis=school_bis_idx)

Using lists or NumPy arrays instead of DataArrays does column/row-based indexing. As you can see, the result has 9 values of theta_shool_diff instead of the 3 pairs of difference we selected in the previous cell:

post["theta_school_diff"].sel(
    school=["Choate", "Hotchkiss", "Mt. Hermon"],
    school_bis=["Deerfield", "Choate", "Lawrenceville"],
)

2.4.3 Add new chains using concat

After checking the mcse and realizing you need more samples, you rerun the model with two chains and obtain an idata_rerun object.

idata_rerun = (
    idata.sel(chain=[0, 1])
    .copy()
    .assign_coords(coords={"chain": [4, 5]}, groups="posterior_groups")
)

You can combine the two into a single InferenceData object using the concat function from ArviZ:

idata_complete = az.concat(idata, idata_rerun, dim="chain")
idata_complete.posterior.dims["chain"]

/tmp/ipykernel_2263/149534111.py:2: FutureWarning: The return type of `Dataset.dims` will be changed to return a set of dimension names in future, in order to be more consistent with `DataArray.dims`. To access a mapping from dimension names to lengths, please use `Dataset.sizes`.
  idata_complete.posterior.dims["chain"]

2.4.4 Add groups to InferenceData objects

To add new groups to InferenceData objects you can use the extend method if the new groups are already in an InferenceData object or the add_groups method if the new groups are dictionaries or xarray.Dataset objects.

rng = np.random.default_rng(3)
idata.add_groups(
    {"predictions": {"obs": rng.normal(size=(4, 500, 2))}},
    dims={"obs": ["new_school"]},
    coords={"new_school": ["Essex College", "Moordale"]},
)
idata

2.5 Final remarks

We have discussed a few of the most common operations with InferenceData objects. If you want to learn more about InferenceData, you can check the InferenceData API documentation.

If you have doubts about how to use InferenceData with ArviZ functions, you can ask questions at PyMC’s discourse