[namedarray] split .set_dims() into .expand_dims() and broadcast_to()

[andersy005]

this is the initial phase of our plan to split the .set_dims() method into two distinct methods .expand_dims() and .broadcast_to() as outlined in the design document.

• add expand_dims()
• add broadcast_to()
• add .transpose() and .T to NamedArray
• update whats-new.rst

[dcherian]

I have frequently unwrapped to access and compute on the underlying data container, which does indeed care about axis order.

[Illviljan]

Designing an API around not using the class seems odd to me.

I agree that once you have retrieved the underlying data you will indeed have to deal with the axis order. But didn't you feel like you were hacking around a problem with xarray? That's what it sounds like to me.

This pattern is not realistic since we support any Hashable as a dimension name:

from numpy import array_api as nxp
import xarray as xr
from xarray.namedarray import _array_api as xrp

na = xr.namedarray.core.NamedArray(
    dims=("a", None, ("c",)), data=nxp.reshape(nxp.arange(4 * 4 * 4), (4, 4, 4))
)

na.expand_dims(x=-2)
na.expand_dims(100=-2) # Error, not a valid identifier, str(100).isidentifier() == False

Is this what you want then?

from numpy import array_api as nxp
import xarray as xr
from xarray.namedarray import _array_api as xrp

na = xr.namedarray.core.NamedArray(
    dims=("a", None, ("c",)), data=nxp.reshape(nxp.arange(4 * 4 * 4), (4, 4, 4))
)

na.expand_dims(dim=100, axis=-2)

[andersy005]

i'm revisiting this PR, and i've noticed that there is no clear agreement on the desired API spec for expand_dims(), is my observation accurate?

i understand that

na.expand_dims(100=-2)

# or 

na.expand_dims(None=0)

won;t work but the following works

na.expand_dims(dim={100:-2})

# or 

na.expand_dims(dim={None=0})

should we then remove support for dim_kwarg if we want to allow users to control the position of the new dimensions?

am i missing something?

[andersy005]

This pattern is not realistic since we support any Hashable as a dimension name:

from numpy import array_api as nxp
import xarray as xr
from xarray.namedarray import _array_api as xrp

na = xr.namedarray.core.NamedArray(
    dims=("a", None, ("c",)), data=nxp.reshape(nxp.arange(4 * 4 * 4), (4, 4, 4))
)

na.expand_dims(x=-2)
na.expand_dims(100=-2) # Error, not a valid identifier, str(100).isidentifier() == False

thinking more about this, we don't seem to restrict user from assigning hashables like (None, 100) to dims which results in the following edge cases that don't work when using keyword arguments like:

In [33]: v = xr.Variable(data=np.zeros(6).reshape(2, 3), dims=[100, None])

In [34]: v
Out[34]: 
<xarray.Variable (100: 2, None: 3)>
array([[0., 0., 0.],
       [0., 0., 0.]])

In [35]: v.sel(100=1)
  Cell In[35], line 1
    v.sel(100=1)
          ^
SyntaxError: expression cannot contain assignment, perhaps you meant "=="?

so, i'm not sure we should drop dim_kwarg for expand_dims

na.expand_dims(b=2)

[dcherian]

na.expand_dims(dim="a", axis=2) would be fine since we'd only support one at a time.

THis is a fairly minor point so lets just go with na.expand_dims(dim="a") and revisit if someone asks for axis support.

[Illviljan]

thinking more about this, we don't seem to restrict user from assigning hashables like (None, 100) to dims which results in the following edge cases that don't work when using keyword arguments like:

In [33]: v = xr.Variable(data=np.zeros(6).reshape(2, 3), dims=[100, None])

In [34]: v
Out[34]: 
<xarray.Variable (100: 2, None: 3)>
array([[0., 0., 0.],
       [0., 0., 0.]])

In [35]: v.sel(100=1)
  Cell In[35], line 1
    v.sel(100=1)
          ^
SyntaxError: expression cannot contain assignment, perhaps you meant "=="?```

Yeah, I'm thinking v.sel(100=1) was a mistake considering xarray in general allows any Hashable as dimension name, a string like "2023_x" wont work either.
Either we significantly restrict the dimension names to strings passing .isidentifier() only or don't allow v.sel(100=1).

[andersy005]

@dcherian, when you get a chance, this is ready for another round of review