Quick overview#

DataTrees#

DataTree is a tree-like container of xarray.DataArray objects, organised into multiple mutually alignable groups. You can think of it like a (recursive) dict of xarray.Dataset objects.

Let’s first make some example xarray datasets (following on from xarray’s quick overview page):

In [1]: import numpy as np

In [2]: import xarray as xr

In [3]: data = xr.DataArray(np.random.randn(2, 3), dims=("x", "y"), coords={"x": [10, 20]})

In [4]: ds = xr.Dataset(dict(foo=data, bar=("x", [1, 2]), baz=np.pi))

In [5]: ds
Out[5]: 
<xarray.Dataset>
Dimensions:  (x: 2, y: 3)
Coordinates:
  * x        (x) int64 10 20
Dimensions without coordinates: y
Data variables:
    foo      (x, y) float64 0.4691 -0.2829 -1.509 -1.136 1.212 -0.1732
    bar      (x) int64 1 2
    baz      float64 3.142

In [6]: ds2 = ds.interp(coords={"x": [10, 12, 14, 16, 18, 20]})

In [7]: ds2
Out[7]: 
<xarray.Dataset>
Dimensions:  (x: 6, y: 3)
Coordinates:
  * x        (x) int64 10 12 14 16 18 20
Dimensions without coordinates: y
Data variables:
    foo      (x, y) float64 0.4691 -0.2829 -1.509 ... -1.136 1.212 -0.1732
    bar      (x) float64 1.0 1.2 1.4 1.6 1.8 2.0
    baz      float64 3.142

In [8]: ds3 = xr.Dataset(
   ...:     dict(people=["alice", "bob"], heights=("people", [1.57, 1.82])),
   ...:     coords={"species": "human"},
   ...: )
   ...: 

In [9]: ds3
Out[9]: 
<xarray.Dataset>
Dimensions:  (people: 2)
Coordinates:
  * people   (people) <U5 'alice' 'bob'
    species  <U5 'human'
Data variables:
    heights  (people) float64 1.57 1.82

Now we’ll put this data into a multi-group tree:

In [10]: from datatree import DataTree

In [11]: dt = DataTree.from_dict({"simulation/coarse": ds, "simulation/fine": ds2, "/": ds3})

In [12]: dt
Out[12]: 
DataTree('None', parent=None)
│   Dimensions:  (people: 2)
│   Coordinates:
│     * people   (people) <U5 'alice' 'bob'
│       species  <U5 'human'
│   Data variables:
│       heights  (people) float64 1.57 1.82
└── DataTree('simulation')
    ├── DataTree('coarse')
    │       Dimensions:  (x: 2, y: 3)
    │       Coordinates:
    │         * x        (x) int64 10 20
    │       Dimensions without coordinates: y
    │       Data variables:
    │           foo      (x, y) float64 0.4691 -0.2829 -1.509 -1.136 1.212 -0.1732
    │           bar      (x) int64 1 2
    │           baz      float64 3.142
    └── DataTree('fine')
            Dimensions:  (x: 6, y: 3)
            Coordinates:
              * x        (x) int64 10 12 14 16 18 20
            Dimensions without coordinates: y
            Data variables:
                foo      (x, y) float64 0.4691 -0.2829 -1.509 ... -1.136 1.212 -0.1732
                bar      (x) float64 1.0 1.2 1.4 1.6 1.8 2.0
                baz      float64 3.142

This creates a datatree with various groups. We have one root group, containing information about individual people. (This root group can be named, but here is unnamed, so is referred to with "/", same as the root of a unix-like filesystem.) The root group then has one subgroup simulation, which contains no data itself but does contain another two subgroups, named fine and coarse.

The (sub-)sub-groups fine and coarse contain two very similar datasets. They both have an "x" dimension, but the dimension is of different lengths in each group, which makes the data in each group unalignable. In the root group we placed some completely unrelated information, showing how we can use a tree to store heterogenous data.

The constraints on each group are therefore the same as the constraint on dataarrays within a single dataset.

We created the sub-groups using a filesystem-like syntax, and accessing groups works the same way. We can access individual dataarrays in a similar fashion

In [13]: dt["simulation/coarse/foo"]
Out[13]: 
<xarray.DataArray 'foo' (x: 2, y: 3)>
array([[ 0.4691123 , -0.28286334, -1.5090585 ],
       [-1.13563237,  1.21211203, -0.17321465]])
Coordinates:
  * x        (x) int64 10 20
Dimensions without coordinates: y

and we can also pull out the data in a particular group as a Dataset object using .ds:

In [14]: dt["simulation/coarse"].ds
Out[14]: 
<xarray.DatasetView>
Dimensions:  (x: 2, y: 3)
Coordinates:
  * x        (x) int64 10 20
Dimensions without coordinates: y
Data variables:
    foo      (x, y) float64 0.4691 -0.2829 -1.509 -1.136 1.212 -0.1732
    bar      (x) int64 1 2
    baz      float64 3.142

Operations map over subtrees, so we can take a mean over the x dimension of both the fine and coarse groups just by

In [15]: avg = dt["simulation"].mean(dim="x")

In [16]: avg
Out[16]: 
DataTree('simulation', parent=None)
├── DataTree('coarse')
│       Dimensions:  (y: 3)
│       Dimensions without coordinates: y
│       Data variables:
│           foo      (y) float64 -0.3333 0.4646 -0.8411
│           bar      float64 1.5
│           baz      float64 3.142
└── DataTree('fine')
        Dimensions:  (y: 3)
        Dimensions without coordinates: y
        Data variables:
            foo      (y) float64 -0.3333 0.4646 -0.8411
            bar      float64 1.5
            baz      float64 3.142

Here the "x" dimension used is always the one local to that sub-group.

You can do almost everything you can do with Dataset objects with DataTree objects (including indexing and arithmetic), as operations will be mapped over every sub-group in the tree. This allows you to work with multiple groups of non-alignable variables at once.

Note

If all of your variables are mutually alignable (i.e. they live on the same grid, such that every common dimension name maps to the same length), then you probably don’t need DataTree, and should consider just sticking with xarray.Dataset.

Quick overview

Contents

Quick overview#

DataTrees#