[Feature Request] iteration equivalent numpy's nditer or ndenumerate

[AdrianSosic]

Hi folks, is there any simple way to iterate over all elements of an xarray together with its labels? What I am looking for is an equivalent to numpy's nditer or rather ndenumerate method:

https://docs.scipy.org/doc/numpy/reference/generated/numpy.nditer.html
https://docs.scipy.org/doc/numpy/reference/generated/numpy.ndenumerate.html

Ideally, the iterator should return both the current data element and its coordinates, potentially in the form of an (ordered) dictionary. Is there any direct possibility to achieve this with the current functionality of the package?

[shoyer]

You could convert your data into pandas and use .itertuples(), e.g.,

import xarray
import itertools

ds = xarray.tutorial.open_dataset('air_temperature')
records = ds.to_dataframe().reset_index().itertuples(index=False, name='Record')
print(list(itertools.islice(records, 5)))

Outputs:

[Record(lat=75.0, lon=200.0, time=Timestamp('2013-01-01 00:00:00'), air=241.1999969482422),
 Record(lat=75.0, lon=200.0, time=Timestamp('2013-01-01 06:00:00'), air=242.09999084472656),
 Record(lat=75.0, lon=200.0, time=Timestamp('2013-01-01 12:00:00'), air=242.29998779296875),
 Record(lat=75.0, lon=200.0, time=Timestamp('2013-01-01 18:00:00'), air=241.88999938964844),
 Record(lat=75.0, lon=200.0, time=Timestamp('2013-01-02 00:00:00'), air=243.1999969482422)]

[AdrianSosic]

Hi shoyer, many thanks for your quick reply. Converting the xarray to a DataFrame indeed does the job and I will use this solution for the time being. Nevertheless, to me the approach seems rather like an ad-hoc solution since it requires a series of conversions / function calls and I feel like there should be some built-in solution from xarray. In particular, in the above solution, you lose track of what are the coordinates and what is the actual data (all is stored in a single NamedTuple), which requires an additional step to separate the two data structures. Anyway, thanks for your help! If you find another solution, please let me know!

[lanougue]

Hi guys,

For now, when I want to iterate over all my dataset I use the simple (but dangerous I believe) workaround:

for i in np.ndindex(ds.shape):
        res = ds[{d:j for d,j in zip(ds.dims,i)}]

but, I am not sure that ndindex will iterate in the good order relatively to the ds.dims return.

Is there any news on this topic ?

Many thanks !

[lanougue]

Hello guys,

While waiting for a integrated solution. Here is a function that should do the job in a safe way. It returns an iterator

def xndindex(ds, dims=None):
    if dims is None:
        dims = ds.dims
    elif type(dims) is str:
        dims=[dims]
    else:
        pass
        
    for d in dims:
        if d not in ds.dims:
            raise ValueError("Invalid dimension '{}'. Available dimensions {}".format(d, ds.dims))
            
    iter_dict = {k:v for k,v in ds.sizes.items() if k in dims}
    for d,k in zip(repeat(tuple(iter_dict.keys())),zip(np.ndindex(tuple(iter_dict.values())))):
        yield {k:l for k,l in zip(d,k[0])}
  

Example of use

a = xr.DataArray(np.random.rand(4,3), dims=['x','y'], coords={'x':np.arange(4), 'y':np.arange(3)})
for i in xndindex(a):
    print(a[i])

[gauteh]

Hi,

I would also greatly appreciate this. Currently this is possible with groupby, assuming the values are unique. This would make for a ergonomic syntax, in line with other methods in xarray.

This is what I would do now for instance:

#%% Computer temporal spectrum for each point
def welch(ds):
    Z = ds.Z[:,0]/1000.
    Z = Z - np.mean(Z)
    f, Z = sc.signal.welch(Z, fs=4., nperseg=4096*4/52, nfft=2*1024)
    return xr.DataArray(Z, dims=('f'), coords={'f': f})

pxc = xc.groupby('points').apply(welch)
pxs = xs.groupby('points').apply(welch)
print(pxc)

I would really like to have the exact same syntax, but with a method that is called iter or map or something, and doesn't group along the dimension: just iterates and maps.

Regards, Gaute

[dcherian]

@gauteh Your example would be much faster if you used the axis kwarg to welch: https://tutorial.xarray.dev/advanced/apply_ufunc/core-dimensions.html

In general, you should avoid iteration as much as possible.

[gauteh]

Yes, wouldn't it make sense to have apply_ufunc as a method on the dataset?

[gauteh]

Here's essentially the same use case, iterating over a specific dimension without having to transpose first: https://stackoverflow.com/a/72037652/377927 . Even if it is faster to apply_ufunc, this makes for a much simpler interface.