indexing of segments

[selipot]

Discussed in #208

^{Originally posted by selipot June 28, 2023}
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])?

[selipot]

Alternatively, I'd like to see a function that apply apply any function onto a specific row of a ragged array, perhaps called apply_row? As an example, plot the trajectory in the n-th row

apply_row(plt.plot,[ds["longitude"],ds["latitude"],ds["count"],n)

So far, one needs to do

traj_idx = np.insert(np.cumsum(ds["count"].values), 0, 0)
sli = slice(traj_idx[j], traj_idx[j+1])
plt.plot(ds.longitude[sli],ds.latitude[sli])

That's just too much.

[milancurcic]

apply_row seems useful but it also seems a special case of apply_ragged. Consider adding a keyword parameter rows to apply_ragged. rows can be None (default) which resolves to the current apply_ragged behavior. Passing an int to rows does what you described as apply_row. And passing an array-like of ints to rows applies the function on the requested rows only.

[selipot]

I like that suggestion @milancurcic, could you work on it?

[philippemiron]

But in that case, this is using Matplotlib API (plt.plot) in parallel, which should be avoided...

[milancurcic]

I don't understand. Where does Matplotlib come in?

[milancurcic]

Oh, Shane's example from the top. Yes, use it with Matplotlib at your own risk.

[milancurcic]

I think for an int value of rows it would be OK even with Matplotlib because the function would be applied only once, and thus in one thread. User can also disable concurrency when calling apply_ragged, by passing executor with max_workers of 1.

[selipot]

I think this would be a good thing. For testing purposes I find myself wanting to use apply_ragged for a single row, or a small subset of a ragged array. And don't tell me that I can use subset ... :)

[selipot]

We can add a warning that using this for plotting is maybe not a good idea.