add DataTree.subset

doc/api.rst

@@ -789,6 +789,7 @@ Index into all nodes in the subtree simultaneously.
 
    DataTree.isel
    DataTree.sel
+   DataTree.subset
 
 ..    DataTree.drop_sel
 ..    DataTree.drop_isel

doc/whats-new.rst

@@ -24,6 +24,8 @@ v2025.05.0 (unreleased)
 
 New Features
 ~~~~~~~~~~~~
+- Added :py:meth:`~xarray.DataTree.subset` to index variables on all nodes of a datatree (:pull:`10400`)
+  By `Mathias Hauser <https://github.com/mathause>`_.
 - Allow an Xarray index that uses multiple dimensions checking equality with another
   index for only a subset of those dimensions (i.e., ignoring the dimensions
   that are excluded from alignment).

xarray/core/datatree.py

@@ -10,6 +10,7 @@
     Iterable,
     Iterator,
     Mapping,
+    Sequence,
 )
 from html import escape
 from typing import (
@@ -1014,6 +1015,35 @@ def __delitem__(self, key: str) -> None:
         else:
             raise KeyError(key)
 
+    def subset(
+        self, keys: str | Sequence[str], *, errors: ErrorOptions = "raise"
+    ) -> DataTree:
+        """Index DataArrays on each node
+
+        Parameters
+        ----------
+        keys : str | Sequence[str]
+            Name of the data variables to index.
+        errors : "raise", "ignore"
+            Whether to raise a key error if a data variable is missing on a node.
+
+        Returns
+        -------
+        out : DataTree
+        """
+
+        if isinstance(keys, str):
+            keys = [keys]
+
+        def getitem(ds):
+            keys_for_ds = keys
+            if errors == "ignore":
+                keys_for_ds = [key for key in keys if key in ds.data_vars]
+
+            return ds[keys_for_ds]
+
+        return map_over_datasets(getitem, self)
+
     @overload
     def update(self, other: Dataset) -> None: ...
 

xarray/tests/test_datatree.py

@@ -333,6 +333,34 @@ def test_getitem_dict_like_selection_access_to_dataset(self) -> None:
         assert_identical(results[{"temp": 1}], data[{"temp": 1}])  # type: ignore[index]
 
 
+def test_subset():
+    ds1 = xr.Dataset(data_vars={"var1": ("x", [1, 2]), "var2": ("x", [0, 1])})
+    ds2 = xr.Dataset(data_vars={"var1": ("x", [1, 2])})
+    dt = xr.DataTree.from_dict({"ds1": ds1, "ds2": ds2})
+
+    dt_var1 = xr.DataTree.from_dict({"ds1": ds1[["var1"]], "ds2": ds2})
+
+    # errors as map_over_datasets does not skip empty nodes
+    with pytest.raises(KeyError, match="var1"):
+        dt.subset("var1")
+
+    # will still error if map_over_datasets will ever skip empty nodes
+    with pytest.raises(KeyError, match="var2"):
+        dt.subset("var2")
+
+    result = dt.subset("var1", errors="ignore")
+    expected = dt_var1
+    xr.testing.assert_equal(result, expected)
+
+    result = dt.subset(["var1"], errors="ignore")
+    expected = dt_var1
+    xr.testing.assert_equal(result, expected)
+
+    result = dt.subset(["var1", "var2"], errors="ignore")
+    expected = dt
+    xr.testing.assert_equal(result, expected)
+
+
 class TestUpdate:
     def test_update(self) -> None:
         dt = DataTree()