Skip to content

ENH: Support DataTree in apply_ufunc - #11552

Open
QinXi-ai wants to merge 1 commit into
pydata:mainfrom
QinXi-ai:feat/datatree-apply-ufunc
Open

ENH: Support DataTree in apply_ufunc#11552
QinXi-ai wants to merge 1 commit into
pydata:mainfrom
QinXi-ai:feat/datatree-apply-ufunc

Conversation

@QinXi-ai

Copy link
Copy Markdown

Closes #9789

Problem

apply_ufunc currently treats DataTree as a dict-like object. Unary calls silently return only the root Dataset, while core-dimension and multi-tree calls can fail when child nodes reach the variable dispatch path.

Solution

  • Dispatch DataTree inputs before the dict-like Dataset path.
  • Reuse map_over_datasets to apply the existing apply_ufunc machinery at corresponding nodes, preserving tree structure, names, inherited coordinates, scalar arguments, kwargs, core dimensions, and tuple outputs.
  • Enable the public xarray.ufuncs wrappers to operate over complete trees.
  • Add regressions for identity, inherited coordinates, core dimensions plus kwargs/scalars, two isomorphic trees, non-isomorphic errors, multiple outputs, and the ufunc wrapper.

Tests

  • python -m pytest -o addopts='' --basetemp D:/CodexCaches/pytest-xarray-9789-full xarray/tests/test_computation.py xarray/tests/test_ufuncs.py xarray/tests/test_datatree_mapping.py -q (323 passed, 131 skipped, 1 xfailed, 2 xpassed)
  • ruff check xarray/computation/apply_ufunc.py xarray/ufuncs.py xarray/tests/test_computation.py xarray/tests/test_ufuncs.py
  • ruff format --check xarray/computation/apply_ufunc.py xarray/ufuncs.py xarray/tests/test_computation.py xarray/tests/test_ufuncs.py
  • git diff --check

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support DataTree in apply_ufunc

1 participant