Add zip_subtrees for paired iteration over DataTrees

[shoyer]

This should be used for implementing DataTree arithmetic inside map_over_datasets, so the result does not depend on the order in which child nodes are defined.

I have also added a minimal implementation of breadth-first-search with an explicit queue the current recursion based solution in xarray.core.iterators (which has been removed). The new implementation is also slightly faster in my microbenchmark:

In [1]: import xarray as xr

In [2]: tree = xr.DataTree.from_dict({f"/x{i}": None for i in range(100)})

In [3]: %timeit _ = list(tree.subtree)
# on main
87.2 μs ± 394 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

# with this branch
55.1 μs ± 294 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

• Tests added

[shoyer]

I made a pass at re-implementing map_over_datasets using zip_subtree in shoyer#2

[headtr1ck]

Only some minor typing remarks, the rest looks good!

[headtr1ck]

Why are these additional type hints required?
Can mypy not resolve it?

[headtr1ck]

I think this should be fixed by changing the class definition to NamedNode(TreeNode[Tree])

[shoyer]

mypy does seem to struggle here -- it raises an error about missing type annotations.

I have not been able to precisely reproduce the mypy setup from CI on my local machine, so I'm going to save this existing issue for someone else to look into.

[headtr1ck]

I can only see the issue of not passing the generic type to the parent class like I said before. But ofc, we can keep this open by now.

[headtr1ck]

Note, that theoretically it is possible to pass no arguments to this function. Then trees and here active_nodes is an empty tuple.

Maybe add something like this at the start:

if len(trees) < 2:
    yield trees
    return

[shoyer]

Good point, I added an error to catch this case.

[TomNicholas]

so the result does not depend on the order in which child nodes are defined.

I'm not sure I understand this. Surely zip_subtrees is zipping nodes according to the order they appear in .children, which is the order they are defined?

[TomNicholas]

Replacing the entire iterators.py file with 6 lines is so clever it's almost rude 🤣

[shoyer]

so the result does not depend on the order in which child nodes are defined.

I'm not sure I understand this. Surely zip_subtrees is zipping nodes according to the order they appear in .children, which is the order they are defined?

Let me restate this: zip_subtrees allows for zipping together multiple DataTree objects even if child nodes on different trees are defined in different orders, as long as the sets of each node's children match.

[TomNicholas]

And when you say "match" you mean the set of names of the children on tree A match the set of names of the children on tree B?

[shoyer]

And when you say "match" you mean the set of names of the children on tree A match the set of names of the children on tree B?

Exactly, matching is based on relative path from each root.

[TomNicholas]

matching is based on relative path from each root.

So I think this will basically create a breaking change relative to how map_over_subtree (and hence all arithmetic) used to work in xarray-contrib/datatree. In the old model the definition of isomorphic was such that you could actually multiply two trees with the same structure but differently-named nodes, e.g.

dt1 = DataTree.from_dict({'a': ..., 'a/b': ..., 'a/c': ...})
dt2 = DataTree.from_dict({'e': ..., 'e/f': ..., 'e/g': ...})

dt1 * dt2  # would return a tree with names following dt1

I was never really sure if that generality was actually necessary though.

Your new definition of isomorphic is sort of like having strict_names=True always, except that the as long as the set of names match then their names will be used to determine their corresponding nodes in the other tree, regardless of order.

Therefore if we're relaxing this then I think there is no longer any need to think of the data model of datatree as being an ordered set of children, as no mapping behaviour will depend on that order any longer.

I think it's okay to change this behaviour, including getting rid of the strict_names kwarg. After all I have never once needed the generality above! I just wanted to really spell out the implications to make sure I and anyone reading is following.

The last sentence of this section of the docs will need changing, and the change would merit an entry in the migration guide.

[TomNicholas]

Now that I think I understand the behaviour change here, the implementation is great! I really like the idea of zip_subtrees as a primitive.

[shoyer]

The last sentence of this section of the docs will need changing, and the change would merit an entry in the migration guide.

I'll do this in the next PR when I migrate over the map_over_datasets stuff.