ENH: Support DataTree in apply_ufunc

xarray/computation/apply_ufunc.py

@@ -929,8 +929,8 @@ def apply_ufunc(
         the style of NumPy universal functions [1]_ (if this is not the case,
         set ``vectorize=True``). If this function returns multiple outputs, you
         must set ``output_core_dims`` as well.
-    *args : Dataset, DataArray, DataArrayGroupBy, DatasetGroupBy, Variable, \
-        numpy.ndarray, dask.array.Array or scalar
+    *args : DataTree, Dataset, DataArray, DataArrayGroupBy, DatasetGroupBy, \
+        Variable, numpy.ndarray, dask.array.Array or scalar
         Mix of labeled and/or unlabeled arrays to which to apply the function.
     input_core_dims : sequence of sequence, optional
         List of the same length as ``args`` giving the list of core dimensions
@@ -1029,8 +1029,10 @@ def apply_ufunc(
 
     Returns
     -------
-    Single value or tuple of Dataset, DataArray, Variable, dask.array.Array or
-    numpy.ndarray, the first type on that list to appear on an input.
+    Single value or tuple of DataTree, Dataset, DataArray, Variable,
+    dask.array.Array or numpy.ndarray, the first type on that list to appear on
+    an input. For DataTree inputs, ``func`` is applied to the datasets at each
+    node and the results are rebuilt into trees with the same structure.
 
     Notes
     -----
@@ -1153,6 +1155,8 @@ def apply_ufunc(
     .. [2] https://numpy.org/doc/stable/reference/c-api/generalized-ufuncs.html
     """
     from xarray.core.dataarray import DataArray
+    from xarray.core.datatree import DataTree
+    from xarray.core.datatree_mapping import map_over_datasets
     from xarray.core.groupby import GroupBy
     from xarray.core.variable import Variable
 
@@ -1231,8 +1235,29 @@ def apply_ufunc(
         dask_gufunc_kwargs=dask_gufunc_kwargs,
     )
 
+    # Apply to the datasets at corresponding nodes before treating DataTree as
+    # a dict-like object. The recursive call then follows the normal Dataset,
+    # DataArray, GroupBy or array dispatch for each node.
+    if any(isinstance(a, DataTree) for a in args):
+        this_apply = functools.partial(
+            apply_ufunc,
+            func,
+            input_core_dims=input_core_dims,
+            output_core_dims=output_core_dims,
+            exclude_dims=exclude_dims,
+            vectorize=vectorize,
+            join=join,
+            dataset_join=dataset_join,
+            dataset_fill_value=dataset_fill_value,
+            keep_attrs=keep_attrs,
+            dask=dask,
+            output_dtypes=output_dtypes,
+            dask_gufunc_kwargs=dask_gufunc_kwargs,
+            on_missing_core_dim=on_missing_core_dim,
+        )
+        return map_over_datasets(this_apply, *args)
     # feed groupby-apply_ufunc through apply_groupby_func
-    if any(isinstance(a, GroupBy) for a in args):
+    elif any(isinstance(a, GroupBy) for a in args):
         this_apply = functools.partial(
             apply_ufunc,
             func,

xarray/tests/test_computation.py

@@ -20,6 +20,7 @@
     ordered_set_union,
     unified_dim_sizes,
 )
+from xarray.core.treenode import TreeIsomorphismError
 from xarray.core.utils import result_name
 from xarray.structure.alignment import broadcast
 from xarray.tests import (
@@ -124,6 +125,108 @@ def test_apply_identity() -> None:
     assert_identical(dataset, apply_identity(dataset.groupby("x")))
 
 
+def test_apply_identity_datatree() -> None:
+    tree = xr.DataTree.from_dict(
+        {
+            "/": xr.Dataset({"x": ("n", [1, 2])}, coords={"n": [10, 20]}),
+            "/child": xr.Dataset({"x": ("n", [3, 4])}),
+        },
+        name="tree",
+    )
+
+    actual = apply_ufunc(identity, tree)
+
+    assert_identical(tree, actual)
+
+
+def test_apply_datatree_with_core_dims_and_scalar() -> None:
+    tree = xr.DataTree.from_dict(
+        {
+            "/": xr.Dataset({"x": ("n", [1, 2])}),
+            "/child": xr.Dataset({"x": ("n", [3, 4])}),
+        }
+    )
+    expected = xr.DataTree.from_dict(
+        {
+            "/": xr.Dataset({"x": 7}),
+            "/child": xr.Dataset({"x": 15}),
+        }
+    )
+
+    actual = apply_ufunc(
+        lambda values, factor, offset: values.sum(axis=-1) * factor + offset,
+        tree,
+        2,
+        input_core_dims=[["n"], []],
+        kwargs={"offset": 1},
+    )
+
+    assert_identical(expected, actual)
+
+
+def test_apply_non_isomorphic_datatrees() -> None:
+    left = xr.DataTree.from_dict({"/child": xr.Dataset({"x": 1})})
+    right = xr.DataTree.from_dict({"/other": xr.Dataset({"x": 2})})
+
+    with pytest.raises(
+        TreeIsomorphismError, match="children at root node do not match"
+    ):
+        apply_ufunc(operator.add, left, right)
+
+
+def test_apply_two_datatrees() -> None:
+    left = xr.DataTree.from_dict(
+        {
+            "/": xr.Dataset({"x": ("n", [1, 2])}),
+            "/child": xr.Dataset({"x": ("n", [3, 4])}),
+        }
+    )
+    right = xr.DataTree.from_dict(
+        {
+            "/": xr.Dataset({"x": ("n", [10, 20])}),
+            "/child": xr.Dataset({"x": ("n", [30, 40])}),
+        }
+    )
+    expected = xr.DataTree.from_dict(
+        {
+            "/": xr.Dataset({"x": ("n", [11, 22])}),
+            "/child": xr.Dataset({"x": ("n", [33, 44])}),
+        }
+    )
+
+    actual = apply_ufunc(operator.add, left, right)
+
+    assert_identical(expected, actual)
+
+
+def test_apply_datatree_two_outputs() -> None:
+    tree = xr.DataTree.from_dict(
+        {
+            "/": xr.Dataset({"x": ("n", [1, 2])}),
+            "/child": xr.Dataset({"x": ("n", [3, 4])}),
+        },
+        name="tree",
+    )
+
+    actual_min, actual_max = apply_ufunc(
+        lambda values: (values.min(axis=-1), values.max(axis=-1)),
+        tree,
+        input_core_dims=[["n"]],
+        output_core_dims=[[], []],
+    )
+
+    expected_min = xr.DataTree.from_dict(
+        {"/": xr.Dataset({"x": 1}), "/child": xr.Dataset({"x": 3})},
+        name="tree",
+    )
+    expected_max = xr.DataTree.from_dict(
+        {"/": xr.Dataset({"x": 2}), "/child": xr.Dataset({"x": 4})},
+        name="tree",
+    )
+    assert_identical(expected_min, actual_min)
+    assert_identical(expected_max, actual_max)
+
+
 def add(a, b):
     return apply_ufunc(operator.add, a, b)
 

xarray/tests/test_ufuncs.py

@@ -235,6 +235,24 @@ def test_ufunc_pickle(self):
         cos_pickled = pickle.loads(pickle.dumps(xu.cos))
         assert_identical(cos_pickled(a), xu.cos(a))
 
+    def test_ufunc_datatree(self):
+        tree = xr.DataTree.from_dict(
+            {
+                "/": xr.Dataset({"x": ("n", [0.0, np.pi])}),
+                "/child": xr.Dataset({"x": ("n", [2 * np.pi, 3 * np.pi])}),
+            }
+        )
+        expected = xr.DataTree.from_dict(
+            {
+                "/": xr.Dataset({"x": ("n", np.cos([0.0, np.pi]))}),
+                "/child": xr.Dataset({"x": ("n", np.cos([2 * np.pi, 3 * np.pi]))}),
+            }
+        )
+
+        actual = xu.cos(tree)
+
+        assert expected.identical(actual)
+
     def test_ufunc_scalar(self):
         actual = xu.sin(1)
         assert isinstance(actual, float)

xarray/ufuncs.py

@@ -11,7 +11,6 @@
 
 def _walk_array_namespaces(obj, namespaces):
     if isinstance(obj, xr.DataTree):
-        # TODO: DataTree doesn't actually support ufuncs yet
         for node in obj.subtree:
             _walk_array_namespaces(node.dataset, namespaces)
     elif isinstance(obj, xr.Dataset):