2 Working with DataTree

Published

July 3, 2026

Modified

August 11, 2026

During a modern Bayesian analysis we usually generate many sets of data including posterior samples, prior/posterior predictive samples, statistics generated by the sampling method, etc. To keep all this data tidy and avoid confusion ArviZ relies on the data-structures provided by xarray (Hoyer and Hamman 2017). If you are not familiar with xarray this chapter introduces some basic elements in the context of Bayesian stats. For a deeper understanding we recommend that you check their xarray’s documentation, you may find xarray useful for problems outside Bayesian analysis.

We need to become familiar with 3 Data Structures:

DataArray: A labelled, N-dimensional array. In other words this is like NumPy but you can access the data using meaningful labels instead of numerical indexes. You may also think of this as the N-D generalization of a pandas/polars Series.
Dataset: It is a dict-like container of DataArray objects aligned along any number of shared dimensions. You may also think of this as the N-D generalization of a pandas/polars DataFrame.
DataTree: This is a container of Datasets, each Dataset is associated with a group.

The best way to understand these data-structures is to explore them. In a real scenario DataTree will be directly generated by probabilistic programming languages (ppls), or the output of ppls will be transformed into DataTree, likely using ArviZ. But we don’t need to fit a model to play with DataTree objects, ArviZ comes equipped with a few DataTree objects. Let’s start by loading the centered_eight DataTree.

dt = az.load_arviz_data("centered_eight")

In the context of Bayesian statistics a DataTree has groups like posterior, observed_data, posterior_predictive, log_likelihood, etc.

dt.groups

('/',
 '/posterior',
 '/posterior_predictive',
 '/log_likelihood',
 '/sample_stats',
 '/prior',
 '/prior_predictive',
 '/observed_data',
 '/constant_data')

DataTrees used by ArviZ are expected to follow the InferenceData schema, which defines the groups and structure a DataTree should contain. For example, the posterior group stores posterior samples, while observed_data stores the observed data. Although the schema is not strictly enforced, following it is strongly recommended because many ArviZ functions assume this structure. When a DataTree follows the schema, most functions work out of the box; otherwise, additional preprocessing may be required.

DataTree/Dataset/DataArray objects have a nice HTML representation in supported environments like Jupyter notebooks.

dt

You should be able to interact with this HTML representation of a DataTree. For instance, you can expand each group by clicking on its name (or on folder symbol).

An important concept is that of dimensions and coordinates. The xarray documentation is full of examples related to maps. In that context, the dimensions would be like latitude and longitude, and the coordinates would be the actual values of latitude and longitude. But the idea is very general and applies to any kind of data.

Let’s see the dimensions and coordinates for the posterior in our dt object. We can see 3 dimensions chain, draw, and school.

chain is used to index the different chains of the MCMC sampler. Its coordinates are the integers [0, 1, 2, 3] because we have 4 chains.
draw is used to index the iterations of the MCMC sampler. Its coordinates are the integers [0, 1, 2, ..., 499] because we have 500 draws per chain.
school is problem-specific. The coordinates for school are the names of 8 schools, ['Choate', 'Deerfield', 'Phillips Andover', 'Phillips Exeter', 'Hotchkiss', 'Lawrenceville', "St. Paul's", 'Mt. Hermon']

The dimensions chain and draw are common in MCMC sampling, and hence they are expected to be present in the posterior group of a DataTree generated from MCMC. The school dimension is specific to the model from which the posterior was generated, and it is not expected to be present in all DataTrees.

2.0.1 Get the dataset corresponding to a single group

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

dt["posterior"]

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

dt.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.

Notice that we still get a DataTree. If you want the Dataset you can do, most of the time you probably want to do:

posterior = dt["posterior"].ds
posterior

This will return an immutable Dataset-like view. If you want a new Dataset do dt["posterior"].to_dataset().

2.0.2 Get coordinate values

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

posterior.school

<xarray.DataArray 'school' (school: 8)> Size: 512B
array(['Choate', 'Deerfield', 'Phillips Andover', 'Phillips Exeter',
       'Hotchkiss', 'Lawrenceville', "St. Paul's", 'Mt. Hermon'], dtype='<U16')
Coordinates:
  * school   (school) <U16 512B 'Choate' 'Deerfield' ... 'Mt. Hermon'

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

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(posterior.school)

Notice that we do not need to first obtain the NumPy array and then compute the length. When working with DataTree/sets/Arrays, you may feel tempted to reduce them to NumPy arrays, as you are more familiar with those. But for many problems that is not needed and for many other that is not even recommended as you may lose the benefit of working with labeled array-like structures.

2.0.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:

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:

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 line of code returns the first 200 draws from all chains:

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

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

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

If we try the same operation with the entire DataTree, we will get an error. The reason being that some groups don’t have the draw dimension, for instance the observed_data group has only the school dimension. What we can do instead is to filter those group that have draw and then perform the selection and slicing.

dt.filter(lambda node: "draw" in node.dims).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 now 400 draws compared to the original 500, while the group observed_data remains unaffected.

2.0.4 Combine chains and draws

Our primary goal usually is 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(dt)

You can achieve the same result using dt.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(dt, 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(dt, num_samples=100, random_seed=124)
ll = az.extract(dt, group="log_likelihood", num_samples=100, random_seed=124)

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

2.1 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

/tmp/ipykernel_108120/1785098962.py:1: FutureWarning: Deleting a single level of a MultiIndex is deprecated. Previously, this deleted all levels of a MultiIndex. Please also drop the following variables: {'sample'} to avoid an error in the future.
  posterior["log_tau"] = np.log(posterior["tau"])