Hierarchical computations#

In this lesson, we extend what we learned about Basic Computation to hierarchical datasets. By the end of the lesson, we will be able to:

  • Apply basic arithmetic and label-aware reductions to xarray DataTree objects

  • Apply arbitrary functions across all nodes across a tree

import numpy as np
import xarray as xr

xr.set_options(keep_attrs=True, display_expand_attrs=False, display_expand_data=False)
<xarray.core.options.set_options at 0x7fb840b35a90>

Example dataset#

First we load the NMC reanalysis air temperature dataset and arrange it to form a hierarchy of temporal resolutions:

ds = xr.tutorial.open_dataset("air_temperature")

ds_daily = ds.resample(time="D").mean("time")
ds_weekly = ds.resample(time="W").mean("time")
ds_monthly = ds.resample(time="ME").mean("time")

tree = xr.DataTree.from_dict(
    {
        "daily": ds_daily,
        "weekly": ds_weekly,
        "monthly": ds_monthly,
        "": xr.Dataset(attrs={"name": "NMC reanalysis temporal pyramid"}),
    }
)
tree
<xarray.DataTree>
Group: /
β”‚   Attributes: (1)
β”œβ”€β”€ Group: /daily
β”‚       Dimensions:  (time: 730, lat: 25, lon: 53)
β”‚       Coordinates:
β”‚         * time     (time) datetime64[ns] 6kB 2013-01-01 2013-01-02 ... 2014-12-31
β”‚         * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
β”‚         * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
β”‚       Data variables:
β”‚           air      (time, lat, lon) float64 8MB 241.9 242.3 242.7 ... 295.9 295.5
β”‚       Attributes: (5)
β”œβ”€β”€ Group: /weekly
β”‚       Dimensions:  (time: 105, lat: 25, lon: 53)
β”‚       Coordinates:
β”‚         * time     (time) datetime64[ns] 840B 2013-01-06 2013-01-13 ... 2015-01-04
β”‚         * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
β”‚         * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
β”‚       Data variables:
β”‚           air      (time, lat, lon) float64 1MB 245.3 245.2 245.0 ... 296.6 296.2
β”‚       Attributes: (5)
└── Group: /monthly
        Dimensions:  (time: 24, lat: 25, lon: 53)
        Coordinates:
          * time     (time) datetime64[ns] 192B 2013-01-31 2013-02-28 ... 2014-12-31
          * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
          * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
        Data variables:
            air      (time, lat, lon) float64 254kB 244.5 244.7 244.7 ... 297.7 297.7
        Attributes: (5)

Arithmetic#

As an extension to Dataset, DataTree objects automatically apply arithmetic to all variables within all nodes:

tree - 273.15
<xarray.DataTree>
Group: /
β”‚   Attributes: (1)
β”œβ”€β”€ Group: /daily
β”‚       Dimensions:  (lat: 25, lon: 53, time: 730)
β”‚       Coordinates:
β”‚         * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
β”‚         * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
β”‚         * time     (time) datetime64[ns] 6kB 2013-01-01 2013-01-02 ... 2014-12-31
β”‚       Data variables:
β”‚           air      (time, lat, lon) float64 8MB -31.28 -30.85 -30.47 ... 22.72 22.39
β”‚       Attributes: (5)
β”œβ”€β”€ Group: /weekly
β”‚       Dimensions:  (lat: 25, lon: 53, time: 105)
β”‚       Coordinates:
β”‚         * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
β”‚         * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
β”‚         * time     (time) datetime64[ns] 840B 2013-01-06 2013-01-13 ... 2015-01-04
β”‚       Data variables:
β”‚           air      (time, lat, lon) float64 1MB -27.87 -27.98 -28.18 ... 23.48 23.03
β”‚       Attributes: (5)
└── Group: /monthly
        Dimensions:  (lat: 25, lon: 53, time: 24)
        Coordinates:
          * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
          * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
          * time     (time) datetime64[ns] 192B 2013-01-31 2013-02-28 ... 2014-12-31
        Data variables:
            air      (time, lat, lon) float64 254kB -28.68 -28.49 -28.48 ... 24.57 24.56
        Attributes: (5)

Indexing#

Just like arithmetic, indexing is simply forwarded to the node datasets. The only difference is that nodes that don’t have a certain coordinate / dimension are skipped instead of raising an error:

tree.isel(lat=slice(None, 10))
<xarray.DataTree>
Group: /
β”‚   Attributes: (1)
β”œβ”€β”€ Group: /daily
β”‚       Dimensions:  (time: 730, lat: 10, lon: 53)
β”‚       Coordinates:
β”‚         * time     (time) datetime64[ns] 6kB 2013-01-01 2013-01-02 ... 2014-12-31
β”‚         * lat      (lat) float32 40B 75.0 72.5 70.0 67.5 65.0 62.5 60.0 57.5 55.0 52.5
β”‚         * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
β”‚       Data variables:
β”‚           air      (time, lat, lon) float64 3MB 241.9 242.3 242.7 ... 276.1 277.5
β”‚       Attributes: (5)
β”œβ”€β”€ Group: /weekly
β”‚       Dimensions:  (time: 105, lat: 10, lon: 53)
β”‚       Coordinates:
β”‚         * time     (time) datetime64[ns] 840B 2013-01-06 2013-01-13 ... 2015-01-04
β”‚         * lat      (lat) float32 40B 75.0 72.5 70.0 67.5 65.0 62.5 60.0 57.5 55.0 52.5
β”‚         * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
β”‚       Data variables:
β”‚           air      (time, lat, lon) float64 445kB 245.3 245.2 245.0 ... 278.1 279.0
β”‚       Attributes: (5)
└── Group: /monthly
        Dimensions:  (time: 24, lat: 10, lon: 53)
        Coordinates:
          * time     (time) datetime64[ns] 192B 2013-01-31 2013-02-28 ... 2014-12-31
          * lat      (lat) float32 40B 75.0 72.5 70.0 67.5 65.0 62.5 60.0 57.5 55.0 52.5
          * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
        Data variables:
            air      (time, lat, lon) float64 102kB 244.5 244.7 244.7 ... 278.5 279.1
        Attributes: (5)
tree.sel(time="2013-11")
<xarray.DataTree>
Group: /
β”‚   Attributes: (1)
β”œβ”€β”€ Group: /daily
β”‚       Dimensions:  (time: 30, lat: 25, lon: 53)
β”‚       Coordinates:
β”‚         * time     (time) datetime64[ns] 240B 2013-11-01 2013-11-02 ... 2013-11-30
β”‚         * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
β”‚         * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
β”‚       Data variables:
β”‚           air      (time, lat, lon) float64 318kB 250.9 251.2 251.2 ... 298.6 298.7
β”‚       Attributes: (5)
β”œβ”€β”€ Group: /weekly
β”‚       Dimensions:  (time: 4, lat: 25, lon: 53)
β”‚       Coordinates:
β”‚         * time     (time) datetime64[ns] 32B 2013-11-03 2013-11-10 ... 2013-11-24
β”‚         * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
β”‚         * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
β”‚       Data variables:
β”‚           air      (time, lat, lon) float64 42kB 256.0 256.4 256.5 ... 298.3 298.3
β”‚       Attributes: (5)
└── Group: /monthly
        Dimensions:  (time: 1, lat: 25, lon: 53)
        Coordinates:
          * time     (time) datetime64[ns] 8B 2013-11-30
          * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
          * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
        Data variables:
            air      (time, lat, lon) float64 11kB 252.3 252.3 252.2 ... 298.8 298.9
        Attributes: (5)

Reductions#

In a similar way, we can reduce all nodes in the datatree at once:

tree.mean(dim=["lat", "lon"])
<xarray.DataTree>
Group: /
β”‚   Attributes: (1)
β”œβ”€β”€ Group: /daily
β”‚       Dimensions:  (time: 730)
β”‚       Coordinates:
β”‚         * time     (time) datetime64[ns] 6kB 2013-01-01 2013-01-02 ... 2014-12-31
β”‚       Data variables:
β”‚           air      (time) float64 6kB 273.6 273.0 273.3 273.5 ... 274.4 274.0 273.3
β”‚       Attributes: (5)
β”œβ”€β”€ Group: /weekly
β”‚       Dimensions:  (time: 105)
β”‚       Coordinates:
β”‚         * time     (time) datetime64[ns] 840B 2013-01-06 2013-01-13 ... 2015-01-04
β”‚       Data variables:
β”‚           air      (time) float64 840B 273.4 273.6 272.7 272.9 ... 275.3 275.2 273.9
β”‚       Attributes: (5)
└── Group: /monthly
        Dimensions:  (time: 24)
        Coordinates:
          * time     (time) datetime64[ns] 192B 2013-01-31 2013-02-28 ... 2014-12-31
        Data variables:
            air      (time) float64 192B 273.0 273.2 275.6 278.4 ... 283.7 278.1 275.2
        Attributes: (5)

Applying functions designed for Dataset with map_over_datasets#

What if we wanted to apply a element-wise function, for example to convert the data to log-space? For a DataArray we could just use numpy.log(), but this is not supported for DataTree objects:

np.log(tree)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[7], line 1
----> 1 np.log(tree)

File ~/work/xarray-tutorial/xarray-tutorial/.pixi/envs/default/lib/python3.14/site-packages/xarray/core/datatree.py:876, in DataTree.__array__(self, dtype, copy)
    873 def __array__(
    874     self, dtype: np.typing.DTypeLike | None = None, /, *, copy: bool | None = None
    875 ) -> np.ndarray:
--> 876     raise TypeError(
    877         "cannot directly convert a DataTree into a "
    878         "numpy array. Instead, create an xarray.DataArray "
    879         "first, either with indexing on the DataTree or by "
    880         "invoking the `to_array()` method."
    881     )

TypeError: cannot directly convert a DataTree into a numpy array. Instead, create an xarray.DataArray first, either with indexing on the DataTree or by invoking the `to_array()` method.

To map a function to all nodes, we can use xarray.map_over_datasets() and xarray.DataTree.map_over_datasets():

tree.map_over_datasets(xr.ufuncs.log)
<xarray.DataTree>
Group: /
β”‚   Attributes: (1)
β”œβ”€β”€ Group: /daily
β”‚       Dimensions:  (time: 730, lat: 25, lon: 53)
β”‚       Coordinates:
β”‚         * time     (time) datetime64[ns] 6kB 2013-01-01 2013-01-02 ... 2014-12-31
β”‚         * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
β”‚         * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
β”‚       Data variables:
β”‚           air      (time, lat, lon) float64 8MB 5.488 5.49 5.492 ... 5.691 5.69 5.689
β”‚       Attributes: (5)
β”œβ”€β”€ Group: /weekly
β”‚       Dimensions:  (time: 105, lat: 25, lon: 53)
β”‚       Coordinates:
β”‚         * time     (time) datetime64[ns] 840B 2013-01-06 2013-01-13 ... 2015-01-04
β”‚         * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
β”‚         * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
β”‚       Data variables:
β”‚           air      (time, lat, lon) float64 1MB 5.502 5.502 5.501 ... 5.692 5.691
β”‚       Attributes: (5)
└── Group: /monthly
        Dimensions:  (time: 24, lat: 25, lon: 53)
        Coordinates:
          * time     (time) datetime64[ns] 192B 2013-01-31 2013-02-28 ... 2014-12-31
          * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
          * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
        Data variables:
            air      (time, lat, lon) float64 254kB 5.499 5.5 5.5 ... 5.696 5.696 5.696
        Attributes: (5)

We can also use a custom function to perform more complex operations, like subtracting a group mean:

def demean(ds):
    return ds.groupby("time.day") - ds.groupby("time.day").mean()

Applying that to the dataset raises an error, though:

tree.map_over_datasets(demean)
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
File ~/work/xarray-tutorial/xarray-tutorial/.pixi/envs/default/lib/python3.14/site-packages/xarray/core/dataset.py:1303, in Dataset._construct_dataarray(self, name)
   1301             variable = self._variables[name]
   1302         except KeyError:
-> 1303             _, name, variable = _get_virtual_variable(self._variables, name, self.sizes)
   1304 

KeyError: 'time.day'

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
File ~/work/xarray-tutorial/xarray-tutorial/.pixi/envs/default/lib/python3.14/site-packages/xarray/core/dataset.py:1421, in Dataset.__getitem__(self, key)
   1419                 if isinstance(key, tuple):
   1420                     message += f"\nHint: use a list to select multiple variables, for example `ds[{list(key)}]`"
-> 1421                 raise KeyError(message) from e
   1422 

File ~/work/xarray-tutorial/xarray-tutorial/.pixi/envs/default/lib/python3.14/site-packages/xarray/core/dataset.py:1303, in Dataset._construct_dataarray(self, name)
   1301             variable = self._variables[name]
   1302         except KeyError:
-> 1303             _, name, variable = _get_virtual_variable(self._variables, name, self.sizes)
   1304 

File ~/work/xarray-tutorial/xarray-tutorial/.pixi/envs/default/lib/python3.14/site-packages/xarray/core/dataset_utils.py:82, in _get_virtual_variable(variables, key, dim_sizes)
     81 ref_name, var_name = split_key
---> 82 ref_var = variables[ref_name]
     84 if _contains_datetime_like_objects(ref_var):

KeyError: 'time'

The above exception was the direct cause of the following exception:

KeyError                                  Traceback (most recent call last)
Cell In[10], line 1
----> 1 tree.map_over_datasets(demean)

File ~/work/xarray-tutorial/xarray-tutorial/.pixi/envs/default/lib/python3.14/site-packages/xarray/core/datatree.py:1849, in DataTree.map_over_datasets(self, func, kwargs, *args)
   1818 """
   1819 Apply a function to every dataset in this subtree, returning a new tree which stores the results.
   1820 
   (...)   1846 map_over_datasets
   1847 """
   1848 # TODO this signature means that func has no way to know which node it is being called upon - change?
-> 1849 return map_over_datasets(func, self, *args, kwargs=kwargs)

File ~/work/xarray-tutorial/xarray-tutorial/.pixi/envs/default/lib/python3.14/site-packages/xarray/core/datatree_mapping.py:117, in map_over_datasets(func, kwargs, *args)
    115             node_dataset_args.insert(i, arg)
    116     with add_path_context_to_errors(path):
--> 117         results = func(*node_dataset_args, **kwargs)
    118     out_data_objects[path] = results
    120 num_return_values = _check_all_return_values(out_data_objects)

Cell In[9], line 2, in demean(ds)
      1 def demean(ds):
----> 2     return ds.groupby("time.day") - ds.groupby("time.day").mean()

File ~/work/xarray-tutorial/xarray-tutorial/.pixi/envs/default/lib/python3.14/site-packages/xarray/util/deprecation_helpers.py:119, in _deprecate_positional_args.<locals>._decorator.<locals>.inner(*args, **kwargs)
    115     kwargs.update(zip_args)
    117     return func(*args[:-n_extra_args], **kwargs)
--> 119 return func(*args, **kwargs)

File ~/work/xarray-tutorial/xarray-tutorial/.pixi/envs/default/lib/python3.14/site-packages/xarray/core/dataset.py:10272, in Dataset.groupby(self, group, squeeze, restore_coord_dims, eagerly_compute_group, **groupers)
  10268             _validate_groupby_squeeze,
  10269         )
  10270 
  10271         _validate_groupby_squeeze(squeeze)
> 10272         rgroupers = _parse_group_and_groupers(
  10273             self, group, groupers, eagerly_compute_group=eagerly_compute_group
  10274         )
  10275 

File ~/work/xarray-tutorial/xarray-tutorial/.pixi/envs/default/lib/python3.14/site-packages/xarray/core/groupby.py:439, in _parse_group_and_groupers(obj, group, groupers, eagerly_compute_group)
    436     elif groupers:
    437         grouper_mapping = cast("Mapping[Hashable, Grouper]", groupers)
--> 439     rgroupers = tuple(
    440         ResolvedGrouper(
    441             grouper, group, obj, eagerly_compute_group=eagerly_compute_group
    442         )
    443         for group, grouper in grouper_mapping.items()
    444     )
    445 return rgroupers

File ~/work/xarray-tutorial/xarray-tutorial/.pixi/envs/default/lib/python3.14/site-packages/xarray/core/groupby.py:440, in <genexpr>(.0)
    436     elif groupers:
    437         grouper_mapping = cast("Mapping[Hashable, Grouper]", groupers)
    439     rgroupers = tuple(
--> 440         ResolvedGrouper(
    441             grouper, group, obj, eagerly_compute_group=eagerly_compute_group
    442         )
    443         for group, grouper in grouper_mapping.items()
    444     )
    445 return rgroupers

File <string>:7, in __create_fn__.<locals>.__init__(self, grouper, group, obj, eagerly_compute_group)

File ~/work/xarray-tutorial/xarray-tutorial/.pixi/envs/default/lib/python3.14/site-packages/xarray/core/groupby.py:333, in ResolvedGrouper.__post_init__(self)
    329 from xarray.groupers import BinGrouper, UniqueGrouper
    331 self.grouper = copy.deepcopy(self.grouper)
--> 333 self.group = _resolve_group(self.obj, self.group)
    335 if self.eagerly_compute_group:
    336     raise ValueError(
    337         f""""Eagerly computing the DataArray you're grouping by ({self.group.name!r}) "
    338         has been removed.
   (...)    343         `.groupby({self.group.name}=BinGrouper(bins=...))`; as appropriate."""
    344     )

File ~/work/xarray-tutorial/xarray-tutorial/.pixi/envs/default/lib/python3.14/site-packages/xarray/core/groupby.py:499, in _resolve_group(obj, group)
    493 if not hashable(group):
    494     raise TypeError(
    495         "`group` must be an xarray.DataArray or the "
    496         "name of an xarray variable or dimension. "
    497         f"Received {group!r} instead."
    498     )
--> 499 group_da: DataArray = obj[group]
    500 if group_da.name not in obj._indexes and group_da.name in obj.dims:
    501     # DummyGroups should not appear on groupby results
    502     newgroup = _DummyGroup(obj, group_da.name, group_da.coords)

File ~/work/xarray-tutorial/xarray-tutorial/.pixi/envs/default/lib/python3.14/site-packages/xarray/core/datatree.py:323, in DatasetView.__getitem__(self, key)
    320 def __getitem__(self, key) -> DataArray | Dataset:
    321     # TODO call the `_get_item` method of DataTree to allow path-like access to contents of other nodes
    322     # For now just call Dataset.__getitem__
--> 323     return Dataset.__getitem__(self, key)

File ~/work/xarray-tutorial/xarray-tutorial/.pixi/envs/default/lib/python3.14/site-packages/xarray/core/dataset.py:1421, in Dataset.__getitem__(self, key)
   1417 
   1418                 # If someone attempts `ds['foo' , 'bar']` instead of `ds[['foo', 'bar']]`
   1419                 if isinstance(key, tuple):
   1420                     message += f"\nHint: use a list to select multiple variables, for example `ds[{list(key)}]`"
-> 1421                 raise KeyError(message) from e
   1422 
   1423         if utils.iterable_of_hashable(key):
   1424             return self._copy_listed(key)

KeyError: "No variable named 'time.day'. Variables on the dataset include []"
Raised whilst mapping function over node(s) with path '.'

The reason for this error is that the root node does not have any variables, and thus in particular no "time" coordinate. To avoid the error, we have to skip computing the function for that node:

def demean(ds):
    if "time" not in ds.coords:
        return ds
    return ds.groupby("time.day") - ds.groupby("time.day").mean()


tree.map_over_datasets(demean)
<xarray.DataTree>
Group: /
β”‚   Attributes: (1)
β”œβ”€β”€ Group: /daily
β”‚       Dimensions:  (lat: 25, lon: 53, time: 730)
β”‚       Coordinates:
β”‚         * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
β”‚         * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
β”‚         * time     (time) datetime64[ns] 6kB 2013-01-01 2013-01-02 ... 2014-12-31
β”‚           day      (time) int64 6kB 1 2 3 4 5 6 7 8 9 ... 23 24 25 26 27 28 29 30 31
β”‚       Data variables:
β”‚           air      (time, lat, lon) float64 8MB -19.41 -18.72 -17.92 ... -1.263 -1.549
β”‚       Attributes: (5)
β”œβ”€β”€ Group: /weekly
β”‚       Dimensions:  (lat: 25, lon: 53, time: 105)
β”‚       Coordinates:
β”‚         * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
β”‚         * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
β”‚         * time     (time) datetime64[ns] 840B 2013-01-06 2013-01-13 ... 2015-01-04
β”‚           day      (time) int64 840B 6 13 20 27 3 10 17 24 3 ... 16 23 30 7 14 21 28 4
β”‚       Data variables:
β”‚           air      (time, lat, lon) float64 1MB -14.98 -14.71 ... -0.2929 -0.5685
β”‚       Attributes: (5)
└── Group: /monthly
        Dimensions:  (lat: 25, lon: 53, time: 24)
        Coordinates:
          * lat      (lat) float32 100B 75.0 72.5 70.0 67.5 65.0 ... 22.5 20.0 17.5 15.0
          * lon      (lon) float32 212B 200.0 202.5 205.0 207.5 ... 325.0 327.5 330.0
          * time     (time) datetime64[ns] 192B 2013-01-31 2013-02-28 ... 2014-12-31
            day      (time) int64 192B 31 28 31 30 31 30 31 31 ... 30 31 31 30 31 30 31
        Data variables:
            air      (time, lat, lon) float64 254kB -15.81 -15.45 ... 0.399 0.3784
        Attributes: (5)

Escape hatches#

For some more complex operations, it might make sense to work on xarray.Dataset or xarray.DataArray objects and reassemble the tree afterwards.

Let’s look at a new dataset:

precipitation = xr.tutorial.open_datatree("precipitation.nc4").load()
precipitation
<xarray.DataTree>
Group: /
β”‚   Dimensions:  (time: 10)
β”‚   Coordinates:
β”‚     * time     (time) datetime64[ns] 80B 2021-08-29T07:30:00 ... 2021-08-29T16:...
β”œβ”€β”€ Group: /observed
β”‚       Dimensions:        (time: 10, lon: 320, lat: 150)
β”‚       Coordinates:
β”‚         * lon            (lon) float32 1kB -109.9 -109.8 -109.8 ... -78.15 -78.05
β”‚         * lat            (lat) float32 600B 20.05 20.15 20.25 ... 34.75 34.85 34.95
β”‚       Data variables:
β”‚           precipitation  (time, lon, lat) float32 2MB 0.0 0.0 0.0 0.0 ... 0.0 0.0 0.0
└── Group: /reanalysis
        Dimensions:        (time: 10, lat: 31, lon: 52)
        Coordinates:
          * lat            (lat) float64 248B 20.0 20.5 21.0 21.5 ... 34.0 34.5 35.0
          * lon            (lon) float64 416B -110.0 -109.4 -108.8 ... -78.75 -78.12
        Data variables:
            precipitation  (time, lat, lon) float32 64kB 0.3941 0.6236 ... 0.0 6.129e-05

Suppose we wanted to interpolate the observed precipitation to the modelled precipitation. We could use map_over_datasets for this, but we can also have a bit more control:

interpolated = xr.DataTree.from_dict(
    {
        "/": precipitation.ds,
        "/observed": precipitation["/observed"].ds.interp(
            lat=precipitation["/reanalysis/lat"],
            lon=precipitation["/reanalysis/lon"],
        ),
        "/reanalysis": precipitation["/reanalysis"],
    }
)
interpolated
<xarray.DataTree>
Group: /
β”‚   Dimensions:  (time: 10)
β”‚   Coordinates:
β”‚     * time     (time) datetime64[ns] 80B 2021-08-29T07:30:00 ... 2021-08-29T16:...
β”œβ”€β”€ Group: /observed
β”‚       Dimensions:        (time: 10, lon: 52, lat: 31)
β”‚       Coordinates:
β”‚         * lon            (lon) float64 416B -110.0 -109.4 -108.8 ... -78.75 -78.12
β”‚         * lat            (lat) float64 248B 20.0 20.5 21.0 21.5 ... 34.0 34.5 35.0
β”‚       Data variables:
β”‚           precipitation  (time, lon, lat) float64 129kB nan nan nan ... 0.0 0.0 nan
└── Group: /reanalysis
        Dimensions:        (time: 10, lat: 31, lon: 52)
        Coordinates:
          * lat            (lat) float64 248B 20.0 20.5 21.0 21.5 ... 34.0 34.5 35.0
          * lon            (lon) float64 416B -110.0 -109.4 -108.8 ... -78.75 -78.12
        Data variables:
            precipitation  (time, lat, lon) float32 64kB 0.3941 0.6236 ... 0.0 6.129e-05

Exercise

Compute the difference between total observed and modelled precipitation, and plot the result.