The xr.Dataset.expand_dims() method can be used to add new dimensions to a dataset. The axis parameter is used to specify where to insert the new dimension in the dataset. However, it's worth noting that the axis parameter only works when expanding along a 1D coordinate, not when expanding along a multi-dimensional array.

Here's an example to illustrate how to use the axis parameter to expand a dataset along a 1D coordinate:

import xarray as xr

create a sample dataset

data = xr.DataArray([[1, 2], [3, 4]], dims=('x', 'y')) ds = xr.Dataset({'foo': data})

add a new dimension along the 'x' coordinate using the 'axis' parameter

ds_expanded = ds.expand_dims({'z': [1]}, axis='x')

In this example, we create a 2D array with dimensions x and y, and then add a new dimension along the x coordinate using the axis='x' parameter.

However, if you try to use the axis parameter to expand a dataset along a multi-dimensional array, you may encounter an error. This is because expanding along a multi-dimensional array would result in a dataset with non-unique dimension names, which is not allowed in xarray.

Here's an example to illustrate this issue:

import xarray as xr

create a sample dataset with a 2D array

data = xr.DataArray([[1, 2], [3, 4]], dims=('x', 'y')) ds = xr.Dataset({'foo': data})

add a new dimension along the 'x' and 'y' coordinates using the 'axis' parameter

ds_expanded = ds.expand_dims({'z': [1]}, axis=('x', 'y'))

In this example, we try to use the axis=('x', 'y') parameter to add a new dimension along both the x and y coordinates. However, this results in a ValueError because the resulting dataset would have non-unique dimension names.

To add a new dimension along a multi-dimensional array, you can instead use the xr.concat() function to concatenate the dataset with a new data array along the desired dimension:

import xarray as xr

create a sample dataset with a 2D array

data = xr.DataArray([[1, 2], [3, 4]], dims=('x', 'y')) ds = xr.Dataset({'foo': data})

add a new dimension along the 'x' and 'y' coordinates using xr.concat

ds_expanded = xr.concat([ds, xr.DataArray([1], dims=('z'))], dim='z')

In this example, we use the xr.concat() function to concatenate the original dataset with a new data array that has a single value along the new dimension z. The dim='z' parameter is used to specify that the new dimension should be named z.

EDIT: Lots of confusion below about nothing, plz disregard

Okay, regardless of expected behavior here, my particular use-case requires that I transpose these dimensions. Can someone show me a way to do this? I tried to explain the xarray point of view to Keras, but Keras is really not interested ;)

I tried something like ds.expand_dims("sample").transpose('sample','nlat','nlon') to complete futility, probably something to do with the Frozen stuff if I had to guess.

I was thinking something like this:

da = xr.DataArray([[1,2,3],[4,5,6],[7,8,9]], coords={'a':[1,2,3], 'b':[1,2,3]}) ds = xr.Dataset({'da':da}) ds1 = ds.expand_dims('yomama', axis=0) print(ds1[0].dims) ds2 = ds.expand_dims('yomama', axis=2) print(ds2[0].dims)

...but this throws an error (like it should). I think I must be reading my code wrong lol

Okay I think I get the philosophy now. However, indexing a DataSet with an integer actually does work. If performance is the goal, shouldn't something like ds[0] throw a warning or an error?

Can you share your code for this? I would interpret that as meaning you have a variable in your dataset mapped to an integer key, which is allowed as a hashable type but can cause problems with downstream packages.

Okay I think I get the philosophy now. However, indexing a DataSet with an integer actually does work. If performance is the goal, shouldn't something like ds[0] throw a warning or an error?

I'm not an xarray developer, but my guess is that your argument is why positional indexing/slicing is not available for datasets.

As for the specific case of using axis parameter of expand_dims, I think this is useful for the case in which the user is either confident about the axis order in each DataArray or will use label based operations such that axis order doesn’t matter. I was curious so I did a quick comparison of the speed for using this parameter versus a subsequent transpose operation: shape = (10, 50, 100, 200) ds = xr.Dataset( { "foo": (["time", "x", "y", "z"], np.random.rand(*shape)), "bar": (["time", "x", "y", "z"], np.random.randint(0, 10, shape)), }, { "time": (["time"], np.arange(shape[0])), "x": (["x"], np.arange(shape[1])), "y": (["y"], np.arange(shape[2])), "z": (["z"], np.arange(shape[3])), }, ) %%timeit -r 4 ds1 = ds.expand_dims("sample", axis=1) 38.1 µs ± 76 ns per loop (mean ± std. dev. of 4 runs, 10,000 loops each) %%timeit -r 4 ds2 = ds.expand_dims("sample").transpose("time", "sample", "x", "y", "z") 172 µs ± 612 ns per loop (mean ± std. dev. of 4 runs, 10,000 loops each)

Nvm, my use case isn't what I thought it was, but I'll push the issue a bit.

So I'm not disputing anything about what these functions actually do now, the issue I have is that the functions here treat the dimension order of a DataSet as if it's arbitrary, but calling [] on a DataSet slices it in a decidedly non-arbitrary way. It turns out that [] actually does care about which axis you select if you call expand_dims first, and you index with an integer like [0]. I think this inconsistency is what's confusing to me atm.

Okay, regardless of expected behavior here, my particular use-case requires that I transpose these dimensions. Can someone show me a way to do this? I tried to explain the xarray point of view to Keras, but Keras is really not interested ;)

I tried something like ds.expand_dims("sample").transpose('sample','nlat','nlon') to complete futility, probably something to do with the Frozen stuff if I had to guess.

The transpose method should change the dimension order on each array in the dataset. One particularly important component from Kai's comment above is that ds.dims does not tell you information about the axis order for the DataArrays in the Dataset. Can you please describe how the DataArray dimension order reported by the code below differs from your expectations? for var in ds.data_vars: print(ds[var].sizes)

Yeah, I could put something together. It'll probably have to wait until next week though.

the documentation is very misleading

Updating the docstring would be a fairly easy and impactful PR if you're up for it!

I mean that's fine, but in that case, the documentation is very misleading

I think it might be enough to describe this thoroughly with examples in the docstring., though I do like the solution of recommending transpose.

I wonder if we shouldn't recommend using expand_dims without axis plus a transpose afterwards if we care about dimension order? Most of xarray's functions work without making assumptions about the dimension order, and I don't think expand_dims should, either (though I might be missing something, of course)

@cmdupuis3 dimensions are not given in a particular order in the Dataset. There could be two DataArray's which have reversed dimensions for instance.

You would need to inspect the DataArray's:

python da = xr.DataArray([[1,2,3],[4,5,6],[7,8,9]], coords={'a':[1,2,3], 'b':[1,2,3]}) ds = xr.Dataset({'da':da}) ds0 = ds.expand_dims('yomama', axis=0) print(ds0.dims) print(ds0.da.dims) ds1 = ds.expand_dims('yomama', axis=1) print(ds1.dims) print(ds1.da.dims) ds2 = ds.expand_dims('yomama', axis=2) print(ds2.dims) print(ds2.da.dims) python Frozen({'a': 3, 'b': 3, 'yomama': 1}) ('yomama', 'a', 'b') Frozen({'a': 3, 'b': 3, 'yomama': 1}) ('a', 'yomama', 'b') Frozen({'a': 3, 'b': 3, 'yomama': 1}) ('a', 'b', 'yomama')