Allow Dataset in numpy array with dtype=object

[telearis]

Discussed in #10043

^{Originally posted by telearis February 12, 2025}

Situation

xarray.Dataset explicitly restricts being put into a numpy.ndarray even if one sets dtype=object:

import numpy as np
import xarray as xr

a = np.array([xr.Dataset({'a': [1, 2,3]})], dtype=object)

This code fails with error:
"TypeError: cannot directly convert an xarray.Dataset into a numpy array. Instead, create an xarray.DataArray first, either with indexing on the Dataset or by invoking the to_dataarray() method."

However, this works:

a = np.empty((1,), dtype=object)
a[0] = xr.Dataset({'a': [1, 2,3]})

Proposal:

xarray.Dataset should not care about being put into numpy.ndarray if dtype=object.

Reason

• Using np.array([ <whatever> ], dtype=object) should work for any objects put into a numpy array.
• Assignment via index is possible (see above). Hence, this behavior is inconsistent.

Application

I came across this when using xarray.DataArray to store results from ray tracing. The data contains (among other data) points where rays have been reflected or diffracted. The number of points is variable. So I wanted to store the information of the reflection/refraction points into a separate xarray.Dataset and store this in an xarray.DataArray with dtype=object (via detour of a numpy array).

Workaround

I currently work around this limitation by subclassing xarray.Dataset:

class MyDataset(xr.Dataset):
    __slots__ = ()
    
    def __array__(self, dtype=None, copy=None):
        assert dtype == object
        assert (copy is None) or (not copy)
        
        x = np.array(None, dtype=object)
        x.flat[0] = self
        
        return x
``´</div>

[welcome]

Thanks for opening your first issue here at xarray! Be sure to follow the issue template!
If you have an idea for a solution, we would really welcome a Pull Request with proposed changes.
See the Contributing Guide for more.
It may take us a while to respond here, but we really value your contribution. Contributors like you help make xarray better.
Thank you!

[keewis]

you may be interested in awkward, which explicitly supports this sort of use case. There's also ragged, which should allow putting such an array into xarray objects (but as far as I remember, how much of xarray's API works is still untested).

[shoyer]

I agree, if a user explicitly writes dtype=object it is probably safe to convert the dataset into a scalar NumPy array.

[TomNicholas]

To be clear, the previous two comments are suggesting alternate ways to handle your use case given that you said that:

The number of points is variable.

This requirement is the same as what motivates awkward and its subset ragged. It's unfortunately not something that xarray has well-established support for, but there's long discussions with context as to what has been tried out on that front here #7988.

To go further if you have a small number of large arrays of inconsistent length you might want to look at xarray.DataTree. But if you have a large number of arrays of inconsistent length (as I suspect you would when ray tracing) then awkward is a better fit.

However your actual issue request - putting xarray objects inside numpy arrays - is not explicitly supported nor recommended. That's because xarray objects wrap numpy arrays and implement all sorts of special methods so that syntax like np.sin(ds) can work. Numpy wrapping xarray wrapping numpy causes all sorts of things to go wrong, and it's not something we're trying to support.

I would therefore suggest closing this issue whilst you try out awkward / ragged instead, and you open new issues if that doesn't work out for you.

[telearis]

Thank you for your suggestions!

I was researching about possible ways to store data along with their coordinates and units and found xarray to be the best fit.

Most of the data are usual numeric xarray.DataArray objects (with common dimensions and individual dimensions) stored inside an xarray.Dataset. For the application that I mentioned, the "ragged" part is only a single small part. And currently it works fine to have that data stored in multiple xarray.Dataset objects which are then put into an xarray.DataArray object having the same common dimensions as all the other data.

In general, I do not see why xarray.Dataset cares about being put in an array-like object with dtype==object. To my understanding, the whole purpose of using this object as dtype is that the array shall not care about the object.

BTW: The detour of using a numpy.ndarray instead of directly creating a xarray.DataArray is because xarray.DataArray does not have a dtype parameter.

This is how numpy behaves:

u = np.random.rand(3,4)
v = np.random.rand(3,4)   # same shape as v
w = np.random.rand(4,5)   # different shape than v

x = np.array([u, v], dtype=object);
y = np.array([u, w], dtype=object)

print(x.shape)   # prints: (2, 3, 4)
print(y.shape)   # prints: (2,)
print(y[0].shape) # prints: (3,4) 

For x, numpy detects a compatible shape and creates a 3D object. (Actually, I would not have expected that.) For y, numpy just stores references to the arrays. xarray kind of subverts this behavior.

[TomNicholas]

Most of the data are usual numeric xarray.DataArray objects (with common dimensions and individual dimensions) stored inside an xarray.Dataset. For the application that I mentioned, the "ragged" part is only a single small part.

Could you elaborate on the "single small part"? Maybe you could be using xarray.DataTree.

---

u = np.random.rand(3,4)
v = np.random.rand(3,4)   # same shape as v

x = np.array([u, v], dtype=object);

print(x.shape)   # prints: (2, 3, 4)

This is exactly the kind of problem I'm talking about. You've wrapped an array (of arrays) [u, v], which should have shape (2,), as there are two objects in there. But the dtype=object approach instead unexpectedly thinks the shape is (2, 3, 4). This weird behaviour is because you're trying to put a np.ndarray object in there - compare to if I try and put some other type inside:

In [1]: import numpy as np

In [2]: class A:
   ...:     ...
   ...: 

In [3]: arr = np.array([A(), A()], dtype='object')

In [4]: arr
Out[4]: 
array([<__main__.A object at 0x1051e7950>,
       <__main__.A object at 0x1051e4ed0>], dtype=object)

In [5]: arr.shape
Out[5]: (2,)

That returns a shape (2,), as I would expect, but unlike what happens when you tried this with a np.ndarray type.

Numpy clearly doesn't behave predictably when you do this, so xarray won't either, which is why it's not recommended.

---

BTW: The detour of using a numpy.ndarray instead of directly creating a xarray.DataArray is because xarray.DataArray does not have a dtype parameter.

I'm not sure what you're getting at here. You can put any compatible array type inside a DataArray, and that array type can have a wide variety of dtypes. It could even be an object dtype, but as we saw above if you try to make those objects np.ndarrays then things will go wrong.

[telearis]

So, my data is organized as follows:

• The ray tracing data is given for different time instants, different path IDs, and different reception points (i.e., antenna locations). Time, path ID, and reception point ID are the common dimensions.
• I store data like the electric field vector (three components), delays, phases, fields strengths, and interactions.
• Everything except the interaction points can be included into an xarray.Dataset as usual xarray.DataArray objects.
• The interactions are a list of another data type. This data type includes information like the coordinates of the interaction, the type of interaction (reflection, diffraction, ...) etc. For every ray (of a certain path ID, at a certain time, at a certain reception point) the number of interactions can be different.

dims = ['time', 'path_id', 'rx_point_id']

# create some random data:
e_field = xr.DataArray(np.random.rand(10, 20, 3, 3), dims=dims + ['position'])
delay = e_field = xr.DataArray(np.random.rand(10, 20, 3), dims=dims)
... # more data

# create the data sets for interactions:
interactions = np.array((10, 20, 3), dtype=None)
for i0, i1, i2 in np.ndindex(interactions.shape):
    num_int = np.random.randint(10)  # random number of interactions
    dims = ['interaction_id']
    positions = xr.DataArray(np.random.rand(num_int, 3), dims=dims + ['position'])
    interaction_type = xr.DataArray(np.random.choice([0, 1], [num_int,]), dims=dims)
    ... # more data
    interactions[i0, i1, i2] = xr.MyDataset({'coordinates': positions, 'type': interaction_type, <....> })

final_data = xr.Dataset({'E': e_field, 'delay': delay, <...>, 'interactions': interactions})

(code untested)

I could include the interactions in the base Dataset and fill it with NaNs. However, since the number varies significantly, this doesn't seem right to me. Hence, for each ray I create another Dataset and store these as references in a DataArray (of common size).

[shoyer]

I tend to agree with @TomNicholas that there are likely better ways to organize your data than putting Xarray data structures into NumPy object dtype arrays, but if this is really what you want to do, I don't think Xarray needs to stop you.

The error raised from Dataset.__array__ exists to stop a naive user who might expect np.array(xarray_dataset) to stack together the contained NumPy arrays, similar to what happens on np.array(pandas_dataframe). If you're explicitly writing np.array(xarray_dataset, dtype=object) that is almost certainly not your expectation.