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

[owenlittlejohns]

What is your issue?

Originally posted by @TomNicholas in xarray-contrib/datatree#146

Arithmetic involving one DataTree and one Dataset is supposed to be node-wise, i.e. the operation involving the single Dataset is automatically applied to every node of the tree individually, returning a new tree, like this:

In [8]: ds = xr.Dataset({"a": 1})

In [9]: ds
Out[9]: 
<xarray.Dataset>
Dimensions:  ()
Data variables:
    a        int64 1

In [10]: dt = DataTree(data=ds)

In [11]: dt * ds
Out[11]: 
DataTree('None', parent=None)
    Dimensions:  ()
    Data variables:
        a        int64 1

However there is a bad bug:

In [12]: ds * dt
Out[12]: 
<xarray.Dataset>
Dimensions:  ()
Data variables:
    a        int64 1

This didn't return a tree, so arithmetic between Datasets and DataTrees currently doesn't respect commutativity :(

Interestingly this does work fine for python scalars

In [27]: dt * 2
Out[27]: 
DataTree('None', parent=None)
    Dimensions:  ()
    Data variables:
        a        int64 2

In [28]: 2 * dt
Out[28]: 
DataTree('None', parent=None)
    Dimensions:  ()
    Data variables:
        a        int64 2

and for numpy arrays:

In [30]: dt * np.array(2)
Out[30]: 
DataTree('None', parent=None)
    Dimensions:  ()
    Data variables:
        a        int64 2

In [31]: np.array(2) * dt
Out[31]: 
DataTree('None', parent=None)
    Dimensions:  ()
    Data variables:
        a        int64 2

I do have tests for this arithmetic behaviour but it looks like I didn't think to check commutativity in any of those tests!

I haven't looked into this deeply yet but my hunch is that it's something to do with __mul__/__rmul__ on Dataset competing for priority with __mul__/__rmul__ on DataTree. If that is the case then it might only be possible to fix it upstream in xarray by changing the behaviour of Dataset to defer to DataTree...

[TomNicholas]

In the meeting last week @shoyer and I discussed just making this behaviour stricter, so that there was no automatic "broadcasting" of a single-node tree across a multi-node tree.

However, we could still allow DataTree * Dataset, even if we disallow DataTree<many-nodes> * DataTree<single-node>.

[TomNicholas]

Closed by #9619