Implement 2d sel via Xarray accessor?

[willirath]

We have https://git.geomar.de/python/xorca_lonlat2ij which basically

• puts all horizontal coordinates into 3d space referenced to the center of the earth.
• constructs a KD tree
• looks up closest indices for any given lat, lon pair

(xorca_lonlat2ij.get_ij is a good entry point to get an overview of the logic.)

Then, we usually use these indices for a vector .sel() as seen here https://github.com/ESM-VFC/esm-vfc-workflows/blob/7edd9bbc7cd3532c0537f89f667154a9dd27156d/meteor_m85.1_vs_NEMO_FOCI_hydrography/compare_m81.1_hydrography_with_ORCA05_FOCI_Test_data.ipynb

What I'd like to have is this implemented as an Xarray accessor or something that is similarly easy to use.

This would not act on the staggered C grids, we'll need XGCM for, but could be easily generalized to XGCM-compatible datasets if xgcm/xgcm#200 became reality at a later point.

(@benbovy This could be a first thing to work on.)

[willirath]

Here's a discussion about distributed KDTrees (https://gist.github.com/brendancol/a3dd4a35ecd94660411112999923d561)
which points here (https://doi.org/10.1111/j.1467-8659.2007.01062.x) for a more distributed implementation.

[willirath]

@benbovy Here's a gist (better view on nbviever) with a very simple ndim selection example.

[benbovy]

Thanks @willirath. Some thoughts:

• I wonder if we couldn't directly use lat/lon coordinates when building a data structure for indexing. I've been looking again at this blog post by Jake VanderPlas about the implementation of a ball tree algorithm compatible with any distance metric. We could then use the haversine metric, or maybe a simpler/faster formula considering that the computed distances are relatively short (<= NEMO/FESOM grid resolution). Not sure whether that would be faster or slower, though.

• The case in your gist does rather correspond to the worst case scenario where the selection points are scattered across the whole domain. I wonder if this really reflects the more real cases where the selection points have specific spatial patterns (e.g., ship track) and/or very limited extend. In those cases, there might be simple tricks to first shrink the amount of data used to build an indexing tree (e.g., first selection using a bounding box or convex hull). Depending on the computational cost of those operations, this might avoid the need of distributed trees, which seem much more complex to me. I haven't read the paper linked in your comment above yet, so I might be wrong, but based on the title/abstract it seems to address the parallel construction of kd‐trees, not their distributed storage.

• If we can afford storing the whole tree in memory for each published grid, that's even simpler, we could build those trees ahead of time (the computational cost doesn't really matters).

[willirath]

I think there's two aspects here, both of which we can understand and optimize separately (if we design this properly):

1. What's the external API?

# haversine metric
points["temp"] = ds.ndsel.sel(x=points["x"], y=points["y"], metric="haversine")["temp"]
# arbitrary metrix
points["temp"] = ds.ndsel.sel(x=points["x"], y=points["y"], metric=metric_func)["temp"]
# or even
points["temp"] = ds.ndsel.sel(points, metric="haversine")["temp"]

For this, we can afford really inefficient internals.

2. Internal optimizations: This is mostly about the below. But we should make sure to have abstractions that make it easy to swap in / out the ideas you outlined.

I think we shold go for the implementation of the API first, because that makes it easier to play with / compare different optimizations.

I wonder if we couldn't directly use lat/lon coordinates when building a data structure for indexing. I've been looking again at this blog post by Jake VanderPlas about the implementation of a ball tree algorithm compatible with any distance metric. We could then use the haversine metric, or maybe a simpler/faster formula considering that the computed distances are relatively short (<= NEMO/FESOM grid resolution). Not sure whether that would be faster or slower, though.

Thanks for pointing me to this blog post. Sounds like this is a good way to go for the trees. As Jake's version seems to support arbitrary Python functions as metric, we can generalize this later.

The case in your gist does rather correspond to the worst case scenario where the selection points are scattered across the whole domain. I wonder if this really reflects the more real cases where the selection points have specific spatial patterns (e.g., ship track) and/or very limited extend. In those cases, there might be simple tricks to first shrink the amount of data used to build an indexing tree (e.g., first selection using a bounding box or convex hull).

Our target positions can get quite messy. If you think of all the autonomous floats and drifters, you get points that are not really easy to order, because these devices are scattered all across the globe. Here's the latest positions of the ARGO array:

This (with completely new positions every ~10 days) is about the most difficult use case we'll have to handle.

Depending on the computational cost of those operations, this might avoid the need of distributed trees, which seem much more complex to me. I haven't read the paper linked in your comment above yet, so I might be wrong, but based on the title/abstract it seems to address the parallel construction of kd‐trees, not their distributed storage.

I think going for fully distributed tree construction does not really matter. But (for chunked coords and / (x)or chunked ship sections) we could do something like:

• for zipped blocks of the x, y, ... dask arrays calculate a tree
• for zipped blocks of the target positions use each of these trees and find candidates for indices
• --> for each target position, you get nblocks candidates of nearet points. the minimal distance tells you which one to use

If we can afford storing the whole tree in memory for each published grid, that's even simpler, we could build those trees ahead of time (the computational cost doesn't really matters).

Yes, we should think about this. It's probably difficult to do this in a way that survives on disk and across versions. But we could aim at only calculating trees once per session.

[benbovy]

For the public API, I was thinking of something like

ds_selected = ds.xsomething.sel_nearest(lat=lat_points, lon=lon_points)

• I think it would be better to provide an accessor that is not too generic (i.e., nearest neighbor lookup in a n-d space), but general enough to support NEMO, FESOM2 and possibly other model grids working with lat/lon coordinates.

• Hence, do we really need to expose a metric parameter? If the haversine formula is not too costly then we could stick with it. Otherwise, if the more simple alternative is a good approximation for 99% of the cases, we could use it instead.

The API should be compatible with Xarray's advanced indexing using DataArray objects. I think it will be compatible if we use .isel internally.

We could also provide a convenient method for the case where we have a set of points (observations) already in a Dataset or DataArray object:

ds_selected = ds.xsomething.reindex_like(ds_points)

Not sure reindex_like is a good name, though, as the behavior is different than Xarray's reindex_like.

[benbovy]

Our target positions can get quite messy. If you think of all the autonomous floats and drifters, you get points that are not really easy to order, because these devices are scattered all across the globe.

Ah that's good to know, I didn't thought about this case!

But (for chunked coords and / (x)or chunked ship sections) we could do something like ...

So, do we want to support the case where the number of observation points is very high? I haven't thought about that either.

[benbovy]

I quickly tested sklearn's BallTree implementation here, it works pretty well (fast and low memory footprint, for 1e6 points).

[willirath]

Hence, do we really need to expose a metric parameter?

While our use case is positions on the globe, I'd strongly favor some implementation that uses arbitrary ndim indexing with an arbitrary metric and a simple wrapper for georeferenced data.

Firstly, there's use cases that are very close to global and realistic Ocean model grids (although they're not really applicable to observational data from the real world): Idealized models with rectangular boxes are often used by the same people that are also interested in the realistic data.

Secondly, the ndim selection could easily add value for people that are not Ocean or climate modellers without really adding complexity on our side. (I'm thinking of people who do stuff like plasma simulations in torodial geometry.)

Thirdly (not sure people say this...), clearly separating out the metric ensures we design and code in a generalizable way, while just hard coding dim names lat and lon somewhere deep in the stack is an easy slip through if we rely on being on the Earth too much.

[willirath]

I quickly tested sklearn's BallTree implementation here, it works pretty well (fast and low memory footprint, for 1e6 points).

Looks good. I was initially feeling a little hesitant to rely on sklearn rather than something that is lower in the stack (like scipy). But this is only a gut feeling.

[willirath]

I think it would be better to provide an accessor that is not too generic (i.e., nearest neighbor lookup in a n-d space), but general enough to support NEMO, FESOM2 and possibly other model grids working with lat/lon coordinates.

xgeosel?

(Although I'm still liking the idea of going for a wider use case. But this is probably something that will eventually be built elsewhere and closer to the Xarray core devs?)

[willirath]

geondsel?

[benbovy]

I'd also be in favor of a generic solution for indexing n-d coordinates! However, I wouldn't spend too much time trying to settle the API now.

We are going to work very soon on Xarray's flexible indexes refactoring. Once indexes will be considered as first-class citizens of the Xarray data model, the recommended approach for addressing our problem will be to reuse/provide a custom index class and then directly use Dataset's or DataArray's .sel() method.

Using an accessor should be a temporary solution before flexible indexes are implemented (hopefully within the next 12 months), so probably best is to avoid generalizing it too much. Ideally, the accessor methods should mimic the current .sel methods to facilitate smooth transition in the future.

Providing generic index classes for some use-cases including n-d irregular data indexing is also part of the Xarray development roadmap. I think this will require input from the broader community.

there's use cases that are very close to global and realistic Ocean model grids (although they're not really applicable to observational data from the real world): Idealized models with rectangular boxes are often used by the same people that are also interested in the realistic data.

If this is something we want to support as soon as possible (< 12 months), and if those idealized models use non-regular grids, then flexible metrics make perfect sense. But this is something we could still easily support later, if needed.

[benbovy]

For the accessor API, we'll need to think how to handle index construction and index queries. There are two options:

A. implicit index construction (cache)

# 1st time call: triggers index construction 
ds.balltree.sel(lat=[...], lon=[...])

# reuse the built index
ds.baltree.sel(lat=[...], lon=[...])

B. explicit index construction

ds.balltree.set_index(['lat', 'lon'], metric='haversine')

ds.balltree.sel(lat=[...], lon=[...])

I'm in favor of option B, which is closer to what it will (most probably) look like after the indexes refactor in Xarray. Additionally, we could provide a ds.balltree.index property for accessing the underlying index.

[benbovy]

If we are going for a more generic solution (which is probably not a bad idea after all :-) ), then I would rather choose a descriptive index name for the name of the accessor, e.g., balltree or btree.

I was initially feeling a little hesitant to rely on sklearn rather than something that is lower in the stack (like scipy). But this is only a gut feeling.

I was too, and I'm still a bit hesitant of depending on the whole scikit-learn library just for reusing it's balltree implementation. I'll l have a look at the code to see if it would make sense to copy it instead.

[willirath]

Are you available for a quick (5 min) skype call? I'll be in meetings for most of the rest of the day, then.