Built in Index for 2D (or higher) coord

[ianhi]

Is your feature request related to a problem?

If I have a 2D coord like abs_time below

<xarray.Dataset> Size: 28kB
Dimensions:   (trial: 3, rel_time: 500)
Coordinates:
  * trial     (trial) <U8 96B 'cosine' 'square' 'sawtooth'
  * rel_time  (rel_time) float64 4kB 0.0 0.01 0.02 0.03 ... 4.96 4.97 4.98 4.99
    abs_time  (trial, rel_time) float64 12kB 0.0 0.01 0.02 ... 14.97 14.98 14.99
Data variables:
    data      (trial, rel_time) float64 12kB 1.0 0.9995 0.998 ... -0.02 -0.01

which looks like this:

I cannot do set_xindex on it to make a selection such as sel(abs_time=slice(0,2.5))

Describe the solution you'd like

Xarray having a built in NDIndex object to handle this case would be fantastic. There may be some tough choices to make for what the right behavior is for some sel operations but I think overall it would be a benefit.

I've been messing around and have something that is most of what I want here: https://ianhuntisaak.com/xarray-linked-indexes/nd-index-demo/ (🤖 fair warning this me heavily using claude - i've checked lots, but definitely not all)

Describe alternatives you've considered

I thought about using NDPointIndex here but I think it's not quite the right application
ds.set_xindex('abs_time', xr.indexes.NDPointIndex) gives

ValueError: the number of variables 1 doesn't match the number of dimensions 2

and if i understand correctly i solving for a conceptually different problem. One where euclidean distance is meaningfully defined over the space of (x, y). In this case that's not meaningful inwhere the space is over (trials, time).

Additional context

No response

[keewis]

One where euclidean distance is meaningfully defined over the space of (x, y).

That's not necessarily the case: the sklearn-based trees that have been (or are in the process of being added) to xoak support metrics other than euclidean.

If I understand correctly, your abs_time coordinate defines a 2.5D surface (2 dimensions plus height), and with sel you'd like to find the points where a horizontal plane intersects that surface.

If that's correct, we might need to find a different data structure that supports this kind of search efficiently (so far we'd use where, which does an exhaustive search). Also note that in theory in may be possible to stack the coordinate into 3 coordinate arrays and create a nd-point index, but not sure how efficient that would be.

[ianhi]

One where euclidean distance is meaningfully defined over the space of (x, y).

That's not necessarily the case: the sklearn-based trees that have been (or are in the process of being added) to xoak support metrics other than euclidean.

I meant in the sense that "trial" is not a physically meaningful dimension. So could there even be a reasonable distance metric here?

coordinate defines a 2.5D surface (2 dimensions plus height)

I'm not sure I fully grasp the meaning of 2.5D here. 2D + height sounds 3d to me? But after reading https://geohubkenya.wordpress.com/2020/11/25/dimensions-0d-1d-2d-2-5d-3d-4d/ I think I understand? This is a 2D surface warped into the third dimension, but not a

If I understand correctly, your abs_time coordinate defines a 2.5D surface (2 dimensions plus height), and with sel you'd like to find the points where a horizontal plane intersects that surface.

I think that's one use case. I'd also like slices (which I know get trickier to decide the right behavior)

f that's correct, we might need to find a different data structure that supports this kind of search efficiently (so far we'd use where, which does an exhaustive search)

Is this sort of selection not possible with KDTree or similar? I admit some ignorance here. I think this is maybe because I am asking for all values for the first two dims and only constraining on the 3rd one that are managed by the tree?

Re using np.where. At the scales I'm thinking about (e.g. 10 x 100000) this doesn't seem to be a big issue. Or rather it's "slow" in the sense of 10ms, but I think I'm ok with that given that currently to make these involves a human thinking for at least a minute about how make the selection they want.

edit: you can go wayyyy faster if the coord is sorted via the power of np.searchsorted on the flattened coord.

Also note that in theory in may be possible to stack the coordinate into 3 coordinate arrays and create a nd-point index, but not sure how efficient that would be.

I'll give it a shot!