Concatenate across multiple dimensions with open_mfdataset

[TomNicholas]

Code Sample

# Create 4 datasets containing sections of contiguous (x,y) data
for i, x in enumerate([1, 3]):
    for j, y in enumerate([10, 40]):
        ds = xr.Dataset({'foo': (('x', 'y'), np.ones((2, 3)))},
                         coords={'x': [x, x+1],
                                 'y': [y, y+10, y+20]})

        ds.to_netcdf('ds.' + str(i) + str(j) + '.nc')

# Try to open them all in one go
ds_read = xr.open_mfdataset('ds.*.nc')
print(ds_read)

Problem description

Currently xr.open_mfdataset will detect a single common dimension and concatenate DataSets along that dimension. However a common use case is a set of NetCDF files which have two or more common dimensions that need to be concatenated along simultaneously (for example collecting the output of any large-scale simulation which parallelizes in more than one dimension simultaneously). For the behaviour of xr.open_mfdataset to be n-dimensional it should automatically recognise and concatenate along all common dimensions.

Expected Output

<xarray.Dataset>
Dimensions:  (x: 4, y: 6)
Coordinates:
  * x        (x) int64 1 2 3 4
  * y        (y) int64 10 20 30 40 50 60
Data variables:
    foo      (x, y) float64 dask.array<shape=(4, 6), chunksize=(2, 3)>

Current output of xr.open_mfdataset()

<xarray.Dataset>
Dimensions:  (x: 4, y: 12)
Coordinates:
  * x        (x) int64 1 2 3 4
  * y        (y) int64 10 20 30 40 50 60 10 20 30 40 50 60
Data variables:
    foo      (x, y) float64 dask.array<shape=(4, 12), chunksize=(4, 3)>

[chiaral]

Thanks for opening up this issue. This would be very helpful for the forecasting community as well, where we usually concatenate along Start time and Lead time dimensions.
Here, however, was mentioned that it is quite difficult to generalize it, and he suggested a workaround.
I know that some people did it for specific datasets, so maybe it would be helpful to add an example to the documentation that shows how this can be implemented on a case by case basis?

[jhamman]

Since you linked to my SO answer, I will add that I think it is quite possible for us to develop this functionality in xarray. My view is that it will take a concerted effort by an interested developer to come up with an approach to do this but it is possible.

Also, I seem to remember seeing this topic before in our issue tracker but I'm not finding it now.

[shoyer]

I agree with @jhamman that it would take effort from an interested developer to do this but in principle it's quite doable.

I think our logic in auto_combine (which powers open_mfdataset) could probably be extended to handle concatenation across multiple dimensions. The main implementation would need to look at coordinates along concatenated dimensions to break the operation into multiple calls xarray.concat()

[TomNicholas]

Another suggestion: as one of the obvious uses for this is in collecting the output from parallelized simulations, which always have ghost cells around the domain each processor computes on, would it be worth adding an option to throw those away as the mf dataset is loaded? Or is that a task better dealt with by slicing the resultant array after the fact?

[shoyer]

@TomNicholas I think you could use the existing preprocess argument to open_mfdataset() for that.

[TomNicholas]

@shoyer At the risk of going off on a tangent - I think that only works if the number of guard cells you want to remove can be determined from the data in the dataset you're loading, because preprocess doesn't accept any further arguments.

For example, say you want to remove all ghost cells except the ones at the edge of your simulation domain. If there's no information in each dataset which marks it as a dataset containing a simulation boundary region, then the preprocess function can't know to treat it differently without further arguments. I might be wrong though?

[aluhamaa]

Just wanted to add the same request;)

• But, I want to add that current behavior, where combined y = 10 20 30 40 50 60 10 20 30 40 50 60 is incorrect in the context of CF, which states that coordinates must be monotonic, and I cannot see many other use cases where such a result would work.
• Also, currently the documentation is not clear about what will happen if you have multiple dimensions with different ranges.

I also do not understand what is the real complexity of implementing it.
As I understand the problem, the initial full dataset is some sort of N-d hypercube and then it is being split into parts along any nr of dimensions. When reading multiple files, which are just parts of this hypercube, it should be enough to just find the possible dimension values, form a hypercube and place each files content into the correct slot? What am I missing here?

[shoyer]

@aluhamaa I don't think you're missing anything here. I agree that it would be pretty straightforward, it just would take a bit of work.

[aidanheerdegen]

👍 for this feature

[jnhansen]

I just had the exact same problem, and while I didn't yet have time to dig into the source code of xarray.open_mfdataset, I wrote my own function to achieve this:

https://gist.github.com/jnhansen/fa474a536201561653f60ea33045f4e2

Maybe it's helpful to some of you.

Note that I make the following assumptions (which are reasonable for my use case):

• the data variables in each part are identical
• equality of the first element of two coordinate arrays is sufficient to assume equality of the two coordinate arrays

[TomNicholas]

Thanks @jnhansen ! I actually ended up writing my own, much lower level, version of this using the netcdf library. The reason I did that was because I was finding it hard to work out how to merge multiple datasets, then write the data out to a new netcdf file in chunks - I kept accidentally loading the entire merged dataset into memory at once. This might just be because I wasn't using the dask integration properly though.

Have you tried using your function to merge netcdf files, then write out a single file which is larger than RAM? Is that even possible in xarray?

[jnhansen]

Yes, xarray should support that very easily -- assuming you have dask installed:

ds = auto_merge('*.nc')
ds.to_netcdf('larger_than_memory.nc')

auto_merge conserves the chunk sizes resulting from the individual files. If the single files are still too large to fit into memory individually you can rechunk to smaller chunk sizes. The same goes of course for the original xarray.open_mfdataset.

I tested it on a ~25 GB dataset (on a machine with less memory than that).

Note: ds = auto_merge('*.nc') actually runs in a matter of milliseconds, as it merely provides a view of the merged dataset. Only once you call ds.to_netcdf('larger_than_memory.nc') all the disk I/O happens.

[TomNicholas]

I've been looking through the functions open_mfdataset, auto_combine, _auto_concat and concat to see how one might go about achieving this in general.

The current behaviour isn't completely explicit, and I would like to check my understanding with a few questions:

1. If you concat two datasets along a dimension which doesn't have a coordinate, then concat will not be able to know what order to concatenate them in, so it just does it in the order they were provided?

2. Although auto_combine can determine the common dimension to concatenate datasets over, it doesn't know anything about insertion order! Even if the datasets have dimension coordinates, the line

grouped = itertoolz.groupby(lambda ds: tuple(sorted(ds.data_vars)), datasets).values()

will only organise the datasets into groups according to the set of dimensions they have, it doesn't order the datasets within each group according to the values in the dimension coordinates?

We can show this because this (new) testcase fails:

@requires_dask
def test_auto_combine_along_coords(self):
    # drop the third dimension to keep things relatively understandable
    data = create_test_data()
    for k in list(data.variables):
        if 'dim3' in data[k].dims:
            del data[k]

    data_split1 = data.isel(dim2=slice(4))
    data_split2 = data.isel(dim2=slice(4, None))
    split_data = [data_split2, data_split1]  # Deliberately arrange datasets in wrong order
    assert_identical(data, auto_combine(split_data, 'dim2'))

with output

E       AssertionError: <xarray.Dataset>
E       Dimensions:  (dim1: 8, dim2: 9, dim3: 10, time: 20)
E       Coordinates:
E         * time     (time) datetime64[ns] 2000-01-01 2000-01-02 ... 2000-01-20
E         * dim2     (dim2) float64 0.0 0.5 1.0 1.5 2.0 2.5 3.0 3.5 4.0
E       Dimensions without coordinates: dim1, dim3
E       Data variables:
E           var1     (dim1, dim2) float64 1.473 1.363 -1.192 ... 0.2341 -0.3403 0.405
E           var2     (dim1, dim2) float64 -0.7952 0.7566 0.2468 ... -0.6822 1.455 0.7314
E       <xarray.Dataset>
E       Dimensions:  (dim1: 8, dim2: 9, time: 20)
E       Coordinates:
E         * time     (time) datetime64[ns] 2000-01-01 2000-01-02 ... 2000-01-20
E         * dim2     (dim2) float64 2.0 2.5 3.0 3.5 4.0 0.0 0.5 1.0 1.5
E       Dimensions without coordinates: dim1
E       Data variables:
E           var1     (dim1, dim2) float64 1.496 -1.834 -0.6588 ... 1.326 0.6805 -0.2999
E           var2     (dim1, dim2) float64 0.7926 -1.063 0.1062 ... -0.447 -0.8955

3. So the call to _auto_concat just assumes that the datasets are provided in the correct order:

concatenated = [_auto_concat(ds, dim=dim, data_vars=data_vars, coords=coords) for ds in grouped]

4. Therefore what needs to be done here is the groupby call needs to be replaced with something that actually orders the datasets according to the value in the dimension coordinates, works in N dimensions, and outputs a structure of datasets upon which _auto_concat can be called repeatedly, along every concatenation dimension?

Also, concat has a positions argument, which allows you to manually specify the concatenation order, but it's not used at all by auto_combine. In the main use case imagined here (concatenating the domains of multi-parallel simulation output) then the user will know the desired positions of each dataset, because it will correspond to how they divided up their domain in the first place. Perhaps an easier approach to providing for that use case would be to propagate the positions argument upwards so that the user can do something like

# User specifies how they split up their domain
domain_decomposition_structure = how_was_this_parallelized('output.*.nc')

# Feeds this info into open_mfdataset
full_domain = xr.open_mfdataset('output.*.nc', positions=domain_decomposition_structure)

This approach would be much less general but would dodge the issue of writing generalized N-D auto-concatenation logic.

Final point - this common use case also has the added complexity of having ghost or guard cells around every dataset, which should be thrown away. Clearly some user input is required here (ghost_cells_x=2, ghost_cells_y=2, ghost_cells_z=0, ...), but I'm really not sure what the best way to fit that kind of logic in is. Yet more arguments to open_mfdataset?

[shoyer]

@TomNicholas I think your analysis is correct here.

I suspect that in most cases we could figure out how to tile datasets by looking at 1D coordinates along each dimension (e.g., indexes for each dataset), e.g., to find a "chunk id" along each concatenated dimension.

These could be used to build something like a NumPy object array of xarray.Dataset/DataArray objects, which could split up into a bunch of 1D calls to xarray.concat().

I would rather avoid using the positions argument of concat. It doesn't really add any flexibility compared to reordering the inputs with xarray.core.nputils.inverse_permutation.

Final point - this common use case also has the added complexity of having ghost or guard cells around every dataset, which should be thrown away. Clearly some user input is required here (ghost_cells_x=2, ghost_cells_y=2, ghost_cells_z=0, ...), but I'm really not sure what the best way to fit that kind of logic in is. Yet more arguments to open_mfdataset?

We could potentially just encourage using the existing preprocess argument.

[TomNicholas]

I started having a go at writing the second half of this - the "n-dimensional-concatenation" function which would accept a grid of xarray.DataSet/DataArray objects (assumed to be in the correct order along all dimensions), and return a single merged dataset.

However, I think there's an issue with using

something like a NumPy object array of xarray.Dataset/DataArray objects.

My plan was to call np.apply_along_axis to apply the 1D xr.concat() function along each axis in turn, something like

from numpy import apply_along_axis
from xarray import concat

def concat_nd(obj_grid, concat_dims=None):
    """
    Concatenates a structured ndarray of xarray Datasets along multiple dimensions.

    Parameters
    ----------
    obj_grid : numpy array of Dataset and DataArray objects
        N-dimensional numpy object array containing xarray objects in the shape they
        are to be concatenated. Each object is expected to
        consist of variables and coordinates with matching shapes except for
        along the concatenated dimension.
    concat_dims : list of str or DataArray or pandas.Index
        Names of the dimensions to concatenate along. Each dimension in this argument
        is passed on to :py:func:`xarray.concat` along with the dataset objects.
        Should therefore be a list of valid dimension arguments to xarray.concat().
  
    Returns
    -------
    combined : xarray.Dataset
    """

    # Combine datasets along one dimension at a time
    # Start with last axis and finish with axis=0
    for axis in reversed(range(obj_grid.ndim)):
        obj_grid = apply_along_axis(concat, axis, arr=obj_grid, dim=concat_dims[axis])

    # Grid should now only contain one xarray object
    return obj_grid.item

However, testing this code with

def test_concat_1d(self):
    data = create_test_data()

    split_data = [data.isel(dim1=slice(3)), data.isel(dim1=slice(3, None))]

    # Will explain why I'm forced to create ndarray like this shortly
    split_data_grid = np.empty(shape=(2), dtype=np.object)
    split_data_grid[0] = split_data[0]
    split_data_grid[1] = split_data[1]

    reconstructed = concat_nd(split_data_grid, ['dim1'])

    xrt.assert_identical(data, reconstructed)

throws an error from within np.apply_along_axis

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_array()` method.

I think this is because even just the idea of having a ndarray containing xarray datasets seems to cause problems - if I do it with a single item then xarray thinks I'm trying to convert the Dataset into a numpy array and throws the same error:

data = create_test_data()
data_grid = np.array(data, dtype='object')

and if I do it with multiple items then numpy will dive down and extract the variables in the dataset instead of just storing a reference to the dataset:

data = create_test_data()
split_data = [data.isel(dim1=slice(3)), data.isel(dim1=slice(3, None))]
split_data_grid = np.array(split_data, dtype='object')
print(split_data_grid)

returns

[['time' 'dim2' 'dim3' 'var1' 'var2' 'var3' 'numbers']
 ['time' 'dim2' 'dim3' 'var1' 'var2' 'var3' 'numbers']]

when I expected something more like

numpy.array([<xarray.Dataset object at 0x7f94efa62278>, <xarray.Dataset object at 0x9f23phj10582>])

(This is why I had to create an empty array and then fill it afterwards in my example test further up.)

Is this the intended behaviour of xarray? Does this mean I can't use numpy arrays of xarray objects at all for this problem? If so then what structure do you think I should use instead (list of lists etc.)?