Reimplement Datatree typed ops #9619

TomNicholas · 2024-10-13T21:13:36Z

Follows on from #9589, but for the ops defined in _typed_ops.py.

Does this mean we can remove most/all of datatree_ops.py?

EDIT: Answering my own question - none of datatree_ops.py is currently used since we disabled it. It does contain lists of methods we should implement, many of which are now implemented in this PR and #9589. As the only purpose it serves right now is tracking remaining un-implemented methods, we should probably remove it in favour of creating a new github issue.

Closes datatree ops.py migration cleanup #8963, and also datatree: Bug with arithmetic between datasets and datatrees (e.g. ds * dt) #9365 (comment). However does not (yet) implement the restriction described in datatree: Bug with arithmetic between datasets and datatrees (e.g. ds * dt) #9365 (comment).
Tests added
New functions/methods are listed in api.rst
Type hints passing

cc @shoyer, @flamingbear

TomNicholas · 2024-10-13T21:34:15Z

xarray/core/datatree.py

+                "a grouped object are not permitted"
+            )
+
+        # TODO requires an implementation of map_over_subtree_inplace


I think I will just do in-place ops in a follow-up PR

xarray/tests/test_datatree.py

…as Mapping

TomNicholas · 2024-10-14T15:16:36Z

Pretty sure the mypy failure in the last commit is unrelated to this PR (which passed mypy before today) - see #9618 (comment).

kmuehlbauer · 2024-10-14T15:36:17Z

@TomNicholas #9621 for the mypy issue

shoyer

Thanks Tom!

I think we should also revise DataTree arithmetic to support arithmetic between trees with child nodes defined in different orders, but that can come later. I will try to work on the implementation of the core machinery on my long flight today....

shoyer · 2024-10-15T05:15:24Z

xarray/core/datatree.py

+        if isinstance(other, GroupBy):
+            # TODO should we be trying to make this work?
+            raise NotImplementedError


Suggested change

if isinstance(other, GroupBy):

# TODO should we be trying to make this work?

raise NotImplementedError

return NotImplemented

NotImplemented is a sentinel value that tells Python that an arithmetic operator is not implemented, and allows the other argument to try implementing it. If all special methods return NotImplemented, then Python raises an informative TypeError.

Should we also explicitly exclude Dataset here, or are the "mapped over all nodes" semantics of DataTree + Dataset arithmetic obvious enough? #9365

If all special methods return NotImplemented, then Python raises an informative TypeError.

Ah I did not realise that part. So NotImplemented is fine here for now.

Should we also explicitly exclude Dataset here, or are the "mapped over all nodes" semantics of DataTree + Dataset arithmetic obvious enough?

I think it's fine to allow Dataset - I did that deliberately. It's also tested now.

done in 9e2dfad

shoyer · 2024-10-15T05:19:00Z

xarray/core/datatree.py

+                "a grouped object are not permitted"
+            )
+
+        # TODO requires an implementation of map_over_subtree_inplace


I suspect this actually needs a different implementation, in order to handle error recovery properly.

I would suggest handling this like Dataset._inplace_binary_op, which builds result Dataset and then replaces the contents of the current dataset with those of the new dataset.

Makes sense. I've raised #9629 to track that, so I'll remove this pseudocode from this PR.

done in 304eb19

…into datatree-ops

xarray/util/generate_ops.py

xarray/tests/test_datatree.py

headtr1ck · 2024-10-15T17:01:33Z

@TomNicholas this should definitely go into whats-new!

TomNicholas · 2024-10-16T00:54:15Z

this should definitely go into whats-new!

Generally with the datatree stuff we have been taking the approach that since the previous versions of xarray never had DataTree publicly available, that "what's new" isn't really appropriate for all the internal work going on. Instead we're just going to have one what's new entry announcing the new class, and summarise the important changes relative to the old version of datatree in the other repository in a dedicated migration guide.

* main: Fix multiple grouping with missing groups (pydata#9650) flox: Properly propagate multiindex (pydata#9649) Update Datatree html repr to indicate inheritance (pydata#9633) Re-implement map_over_datasets using group_subtrees (pydata#9636) fix zarr intersphinx (pydata#9652) Replace black and blackdoc with ruff-format (pydata#9506) Fix error and missing code cell in io.rst (pydata#9641) Support alternative names for the root node in DataTree.from_dict (pydata#9638) Updates to DataTree.equals and DataTree.identical (pydata#9627) DOC: Clarify error message in open_dataarray (pydata#9637) Add zip_subtrees for paired iteration over DataTrees (pydata#9623) Type check datatree tests (pydata#9632) Add missing `memo` argument to DataTree.__deepcopy__ (pydata#9631) Bug fixes for DataTree indexing and aggregation (pydata#9626) Add inherit=False option to DataTree.copy() (pydata#9628) docs(groupby): mention deprecation of `squeeze` kwarg (pydata#9625) Migration guide for users of old datatree repo (pydata#9598) Reimplement Datatree typed ops (pydata#9619)

* main: (63 commits) Add close() method to DataTree and use it to clean-up open files in tests (pydata#9651) Change URL for pydap test (pydata#9655) Fix multiple grouping with missing groups (pydata#9650) flox: Properly propagate multiindex (pydata#9649) Update Datatree html repr to indicate inheritance (pydata#9633) Re-implement map_over_datasets using group_subtrees (pydata#9636) fix zarr intersphinx (pydata#9652) Replace black and blackdoc with ruff-format (pydata#9506) Fix error and missing code cell in io.rst (pydata#9641) Support alternative names for the root node in DataTree.from_dict (pydata#9638) Updates to DataTree.equals and DataTree.identical (pydata#9627) DOC: Clarify error message in open_dataarray (pydata#9637) Add zip_subtrees for paired iteration over DataTrees (pydata#9623) Type check datatree tests (pydata#9632) Add missing `memo` argument to DataTree.__deepcopy__ (pydata#9631) Bug fixes for DataTree indexing and aggregation (pydata#9626) Add inherit=False option to DataTree.copy() (pydata#9628) docs(groupby): mention deprecation of `squeeze` kwarg (pydata#9625) Migration guide for users of old datatree repo (pydata#9598) Reimplement Datatree typed ops (pydata#9619) ...

TomNicholas added 9 commits October 13, 2024 15:56

test unary op

812d207

implement and generate unary ops

1b9c089

test for unary op with inherited coordinates

2d9cef4

re-enable arithmetic tests

4bb8902

implementation for binary ops

8e7c1da

test ds * dt commutativity

909ae0e

ensure other types defer to DataTree, thus fixing pydata#9365

03ce2c5

test for inplace binary op

12112f1

pseudocode implementation of inplace binary op, and xfail test

b823290

TomNicholas added the topic-DataTree Related to the implementation of a DataTree class label Oct 13, 2024

TomNicholas requested a review from shoyer October 13, 2024 21:13

TomNicholas commented Oct 13, 2024

View reviewed changes

TomNicholas added 8 commits October 13, 2024 23:13

remove some unneeded type: ignore comments

7147bb3

return type should be DataTree

368d456

type datatree ops as accepting dataset-compatible types too

2c740b4

use same type hinting hack as Dataset does for __eq__ not being same …

1eae418

…as Mapping

ignore return type

bb019da

add some methods to api docs

20199cb

don't try to import DataTree.astype in API docs

7634b6c

test to check that single-node trees aren't broadcast

2814801

TomNicholas changed the title ~~Reimplement Datatree ops~~ Reimplement Datatree typed ops Oct 14, 2024

Merge branch 'main' into datatree-ops

1fcbe02

shoyer approved these changes Oct 15, 2024

View reviewed changes

TomNicholas mentioned this pull request Oct 15, 2024

datatree: Implement in-place binary ops #9629

Open

TomNicholas added 4 commits October 15, 2024 09:50

return NotImplemented

9e2dfad

Merge branch 'datatree-ops' of https://github.com/TomNicholas/xarray …

399c6e1

…into datatree-ops

remove pseudocode for inplace binary ops

304eb19

Merge branch 'main' into datatree-ops

7dbc817

map_over_subtree -> map_over_datasets

c7d8060

TomNicholas enabled auto-merge (squash) October 15, 2024 13:57

TomNicholas merged commit 97ec434 into pydata:main Oct 15, 2024
28 checks passed

TomNicholas deleted the datatree-ops branch October 15, 2024 14:16

headtr1ck approved these changes Oct 15, 2024

View reviewed changes

xarray/util/generate_ops.py Show resolved Hide resolved

xarray/tests/test_datatree.py Show resolved Hide resolved

This was referenced Oct 16, 2024

Type check datatree tests #9632

Merged

Enable datatree * dataset commutativity #7497

Closed

datatree: Bug with arithmetic between datasets and datatrees (e.g. ds * dt) #9365

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Reimplement Datatree typed ops #9619

Reimplement Datatree typed ops #9619

TomNicholas commented Oct 13, 2024 •

edited

Loading

TomNicholas Oct 13, 2024

TomNicholas commented Oct 14, 2024

kmuehlbauer commented Oct 14, 2024

shoyer left a comment

shoyer Oct 15, 2024

TomNicholas Oct 15, 2024

TomNicholas Oct 15, 2024

shoyer Oct 15, 2024

TomNicholas Oct 15, 2024

TomNicholas Oct 15, 2024

headtr1ck commented Oct 15, 2024

TomNicholas commented Oct 16, 2024

Reimplement Datatree typed ops #9619

Reimplement Datatree typed ops #9619

Conversation

TomNicholas commented Oct 13, 2024 • edited Loading

TomNicholas Oct 13, 2024

Choose a reason for hiding this comment

TomNicholas commented Oct 14, 2024

kmuehlbauer commented Oct 14, 2024

shoyer left a comment

Choose a reason for hiding this comment

shoyer Oct 15, 2024

Choose a reason for hiding this comment

TomNicholas Oct 15, 2024

Choose a reason for hiding this comment

TomNicholas Oct 15, 2024

Choose a reason for hiding this comment

shoyer Oct 15, 2024

Choose a reason for hiding this comment

TomNicholas Oct 15, 2024

Choose a reason for hiding this comment

TomNicholas Oct 15, 2024

Choose a reason for hiding this comment

headtr1ck commented Oct 15, 2024

TomNicholas commented Oct 16, 2024

TomNicholas commented Oct 13, 2024 •

edited

Loading