indexing of segments

[selipot]

I am wondering if we could improve indexing within rows. Let's say we start from a ragged array dataset ds with a given rowsize, we first need to generate a row index, define a slice, then index that slice:

row_index = np.insert(np.cumsum(ds.rowsize.values), 0, 0) # define row index
j = 1
row = slice(row_index[j], row_index[j+1]) # define a slice for row 1
ds.lon[row][n] # access the first data point of that row

How can we simply this to something like ds.lon[j][n] which is basically a syntax/functionality that is supported by awkward arrays but not xarray.

Is one long haul solution to define new classes such xarray.RaggedDataArray and xarray.RaggedDataset for which we could have .isel and .sel methods for double indexing (ds.lon[j][n])?

[milancurcic]

Yes, I think indexing segments can be improved. Currently, I think the simplest (syntax-wise) way to index an segment is with unpack_ragged:

unpack_ragged(ds.lon, ds.count)[j][n]

However, this is quite inefficient if the ragged-array is large.

Is one long haul solution to define new classes such xarray.RaggedDataArray and xarray.RaggedDataset for which we could have .isel and .sel methods for double indexing (ds.lon[j][n])?

Yes, this would work if we extended the xr.Dataset and xr.DataArray and introduced our own ragged-array classes (inheritance).

We could also monkey-patch the xr.Dataset so that sel and isel do a different thing if the DataArray is detected as ragged. "Monkey-patching" means that we would import xr.Dataset and xr.DataArray and from clouddrift return their instances with sel() and isel() methods that are our own, rather than Xarray's.

A simpler and less intrusive way would be to define standalone functions isel and sel rather than bound methods:

from cd.analysis import isel

nth_segment = isel(ds.lon, rowsize, j)

[selipot]

ds.count will not work by the way