avoiding nans from coordinate union in dataset

[tylerflex]

New xarray user here.
It seems when more than one xr.DataArray objects are used to initialize a xr.Dataset that the remaining dataset combines the coordinates and creates new data arrays with nan values where there was no data supplied. The resulting data arrays are difficult to work with because the nans get in the way of interpolation and greatly increase their storage space in my application.

import xarray as XR
import numpy as np
print(f'xarray version = {xr.__version__}')
x1 = [-1, 0, 1]
x2 = [-0.5, 0.5]
da1 = xr.DataArray(np.random.random(len(x1)), coords=dict(x=x1))
da2 = xr.DataArray(np.random.random(len(x2)), coords=dict(x=x2))
ds = xr.Dataset(dict(d1=da1, d2=da2))
print(f'dataset[d1].values = {ds.d1.values}')
print(f'interpolating (x=0) from original data array: {da2.interp(x=0).values:.3f}')
print(f'interpolating (x=0) from dataset: {ds.d2.interp(x=0).values}')

gives

xarray version = 0.20.1
dataset[d1].values = [0.66355207        nan 0.76426127        nan 0.22273888]
interpolating (x=0) from original data array: 0.447
interpolating (x=0) from dataset: nan

Is there a way to turn this behavior off?
If not, what is a better alternative. Most naive thought is to just create my own custom Dataset class that manages a dictionary of xr.DataArray objects.

Thank you 🙏

[andersy005]

Not 100% sure this is what you are looking for, but you could use xr.merge() in conjunction with join='exact or join='inner'

In [13]: import xarray as xr
    ...: print(f'xarray version = {xr.__version__}')
    ...: x1 = [-1, 0, 1]
    ...: x2 = [-0.5, 0.5]
    ...: ds1 = xr.DataArray(np.random.random(len(x1)), coords=dict(x=x1)).to_dataset(name='d1')
    ...: ds2 = xr.DataArray(np.random.random(len(x2)), coords=dict(x=x2)).to_dataset(name='d2')
xarray version = 2022.3.0

In [26]: ds1
Out[26]: 
<xarray.Dataset>
Dimensions:  (x: 3)
Coordinates:
  * x        (x) int64 -1 0 1
Data variables:
    d1       (x) float64 0.9725 0.184 0.6954

In [27]: ds2
Out[27]: 
<xarray.Dataset>
Dimensions:  (x: 2)
Coordinates:
  * x        (x) float64 -0.5 0.5
Data variables:
    d2       (x) float64 0.337 0.4074

• join='inner' will proceed, but you may end up with an empty dataset if there is no overlap

In [23]: xr.merge([ds1, ds2], join="inner")
Out[23]: 
<xarray.Dataset>
Dimensions:  (x: 0)
Coordinates:
  * x        (x) float64 
Data variables:
    d1       (x) float64 
    d2       (x) float64 

• join='exact' will raise an error if there is no overlap

In [24]: xr.merge([ds1, ds2], join="exact")
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [24], in <cell line: 1>()
----> 1 xr.merge([da1, da2], join="exact")

File ~/mambaforge/envs/playground/lib/python3.10/site-packages/xarray/core/merge.py:900, in merge(objects, compat, join, fill_value, combine_attrs)
    897     obj = obj.to_dataset(promote_attrs=True) if isinstance(obj, DataArray) else obj
    898     dict_like_objects.append(obj)
--> 900 merge_result = merge_core(
    901     dict_like_objects,
    902     compat,
    903     join,
    904     combine_attrs=combine_attrs,
    905     fill_value=fill_value,
    906 )
    907 return Dataset._construct_direct(**merge_result._asdict())

File ~/mambaforge/envs/playground/lib/python3.10/site-packages/xarray/core/merge.py:629, in merge_core(objects, compat, join, combine_attrs, priority_arg, explicit_coords, indexes, fill_value)
    626 _assert_compat_valid(compat)
    628 coerced = coerce_pandas_values(objects)
--> 629 aligned = deep_align(
    630     coerced, join=join, copy=False, indexes=indexes, fill_value=fill_value
    631 )
    632 collected = collect_variables_and_indexes(aligned)
    634 prioritized = _get_priority_vars_and_indexes(aligned, priority_arg, compat=compat)

File ~/mambaforge/envs/playground/lib/python3.10/site-packages/xarray/core/alignment.py:436, in deep_align(objects, join, copy, indexes, exclude, raise_on_invalid, fill_value)
    433     else:
    434         out.append(variables)
--> 436 aligned = align(
    437     *targets,
    438     join=join,
    439     copy=copy,
    440     indexes=indexes,
    441     exclude=exclude,
    442     fill_value=fill_value,
    443 )
    445 for position, key, aligned_obj in zip(positions, keys, aligned):
    446     if key is no_key:

File ~/mambaforge/envs/playground/lib/python3.10/site-packages/xarray/core/alignment.py:324, in align(join, copy, indexes, exclude, fill_value, *objects)
    316 if (
    317     any(
    318         not matching_indexes[0].equals(other)
   (...)
    321     or dim in unlabeled_dim_sizes
    322 ):
    323     if join == "exact":
--> 324         raise ValueError(f"indexes along dimension {dim!r} are not equal")
    325     joiner = _get_joiner(join, type(matching_indexes[0]))
    326     index = joiner(matching_indexes)

ValueError: indexes along dimension 'x' are not equal

[tylerflex]

Thanks @andersy005, indeed it seems like

• merge='inner' just returns an empty dataset because there is no overlap.
• merge='outer' does the same as my original attempt with xr.Dataset() init.
• merge='exact' fails because the coords are not the same.

I think what I'm looking for is a merge where the coordinates of each of the data arrays are preserved. Maybe this goes against the whole philosophy behind Dataset though.

[andersy005]

I think what I'm looking for is a merge where the coordinates of each of the data arrays are preserved.

it's my understanding that this is currently not possible due to the fact that datarrays within a dataset are allowed to share coordinates and your use case would result in a dataset with ragged arrays (which are currently not supported).

[tylerflex]

Makes sense, that's what I figured. Appreciate your help!

[dcherian]

There is this experimental package: https://github.com/TomNicholas/datatree

where each node is a Dataset and nodes need not be "aligned" with each other. SO you could throw each variable in a different node, if that worked for your application