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?

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.

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

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