map_over_datasets: skip empty nodes

[mathause]

• Closes map_over_datasets throws error on nodes without datasets #9693
• Closes datatree gets dis-aligned in binary op #10013
• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst
• New functions/methods are listed in api.rst

---

• misses tests and docs but I'd like to get some feedback first
• needs some add. logic to only check the output on non-empty nodes and to ensure multi-output functions are correct
• no good way for a proper deprecation without a keyword

[Illviljan]

A interpolation use case that doesn't crash with this PR:

import numpy as np

import xarray as xr

number_of_files = 700
number_of_groups = 5
number_of_variables = 10

datasets = {}
for f in range(number_of_files):
    for g in range(number_of_groups):
        # Create random data
        time = np.linspace(0, 50 + f, 1 + 1000 * g)
        y = f * time + g

        # Create dataset:
        ds = xr.Dataset(
            data_vars={
                f"temperature_{g}{i}": ("time", y)
                for i in range(number_of_variables // number_of_groups)
            },
            coords={"time": ("time", time)},
        ).chunk()

        # Prepare for xr.DataTree:
        name = f"file_{f}/group_{g}"
        datasets[name] = ds
dt = xr.DataTree.from_dict(datasets)

# %% Interpolate to same time coordinate
def ds_interp(ds, *args, **kwargs):
    return ds.interp(*args, **kwargs)


new_time = np.linspace(0, 100, 50)
dt_interp = dt.map_over_datasets(
    ds_interp, kwargs=dict(time=new_time, assume_sorted=True)
)

[mathause]

Thanks for the example. This PR would also close #10013. This would be a huge plus for me. Not being able to subtract a ds from a datatree makes it extremely cumbersome. However, this implies that the binary ops are implemented using map_over_datasets and means there is a considerable behavior change.

[mathause]

@TomNicholas do you see any chance this PR might get merged (after adding tests etc. obviously)? Are there discussions beside #9693 that I am missing?

[TomNicholas]

Hey @mathause - sorry for forgetting about this - I've been busy.

I think something like this should get merged, but there are various small and fairly arbitrary choices to quibble over. They are basically all already mentioned in #9693 though.

this implies that the binary ops are implemented using map_over_datasets and means there is a considerable behavior change.

I don't understand this statement though - aren't binary ops already implemented using map_over_datasets?

xarray/xarray/core/datatree.py

Line 1590 in 5ea1e81

return map_over_datasets(ds_binop, self, other)

We're changing the behaviour, but changing it to be closer to the old datatree, which is what a lot of users expect anyway.

[mathause]

I think this is ready for review

---

Hey @mathause - sorry for forgetting about this - I've been busy.

No worries! Thanks for considering this PR!

I think something like this should get merged, but there are various small and fairly arbitrary choices to quibble over. They are basically all already mentioned in #9693 though.

The one unclear choice from #9693 was the comment by @shoyer #9693 (comment):

I'm not sure whether or not to call the mapped over function for nodes that only define coordinates. Certainly I would not blindly copy coordinates from otherwise empty nodes onto the result, because those coordinates may no longer be relevant on the result.

I currently use DataTree.has_data, this includes nodes that only have coords (although nodes which inherit coords are excluded (I think)). I don't see a way to be clever about these nodes.

this implies that the binary ops are implemented using map_over_datasets and means there is a considerable behavior change.

I don't understand this statement though - aren't binary ops already implemented using map_over_datasets?

Yes sorry that was not clear. I just wanted to say that binary ops are also affected.

[mathause]

gentle ping @TomNicholas

(apologies for bothering you again - I am unfortunately currently blocked by this in another project. Or is there someone else who could potentially review this?)

[TomNicholas]

no worries, replied here #9693 (comment)

[peanutfun]

Any updates on this? #10013 makes binary operations on datatrees quite unfeasible, and a fix would be very welcome!

[shoyer]

I would check every argument here, instead of just the first one.

[shoyer]

Assigning a single dataset as the result implicitly assumes that the result of map_over_datasets is a single Dataset, which may not be the case.

In cases of at least one non-empty node, we can infer the number of outputs, but this isn't possible in general, e.g., consider map_over_datasets(np.divmod, Datatree(), Datatree()).

I think we now need an explicit num_return_values argument which can be supplied to disambiguate this.

[mathause]

Thanks a lot for the review! I updated the function to check all passed datatree nodes and added default_num_return_values - happy to get another review!