From 5212f8de4c852b02ed59a5f66edab5ff522ed277 Mon Sep 17 00:00:00 2001 From: hushen <190065939+918154429@users.noreply.github.com> Date: Mon, 31 Aug 2026 23:40:55 +0800 Subject: [PATCH] ENH support DataTree in apply_ufunc --- xarray/computation/apply_ufunc.py | 35 ++++++++-- xarray/tests/test_computation.py | 103 ++++++++++++++++++++++++++++++ xarray/tests/test_ufuncs.py | 18 ++++++ xarray/ufuncs.py | 1 - 4 files changed, 151 insertions(+), 6 deletions(-) diff --git a/xarray/computation/apply_ufunc.py b/xarray/computation/apply_ufunc.py index 36cdf2bfbf6..d0b2e1d7b79 100644 --- a/xarray/computation/apply_ufunc.py +++ b/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, diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py index d1532df8891..ecdc7eb6e49 100644 --- a/xarray/tests/test_computation.py +++ b/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) diff --git a/xarray/tests/test_ufuncs.py b/xarray/tests/test_ufuncs.py index 6ef3386512e..f4325fa5cea 100644 --- a/xarray/tests/test_ufuncs.py +++ 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) diff --git a/xarray/ufuncs.py b/xarray/ufuncs.py index 83acbde858b..fc5e2e210e3 100644 --- a/xarray/ufuncs.py +++ b/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):