Improving performance of vectorized interpolation in Xarray

[VeckoTheGecko]

In Parcels we do a lot of point-wise interpolation onto ocean model data. We do this by first (a) looking up of the barycentric particle positions on the staggered curvilinear mesh, then (b) using these positions to use xarrays vectorized interpolation to get the data.

We're quite invested in the performance of Xarray vector interpolation (doing some initial profiling it looks like that where the majority of performance is left on the table). @shoyer you were mentioning during the xarray dev meeting that you think the vectorized interpolation in xarray can be improved, could you provide more details?

I'm still getting into all this, so this attach plan might not be fully formed: To deal with (a), I'm thinking of looking into Xarrays flexible indices - to see if we can define indices that allow us to easily from particle lat/lon look up position in the staggered mesh. For (b), I'm hoping this discussion topic provides more insight - maybe there are lower level public xarray endpoints that we can hook into (or, longer term, improvements that we can bring back upstream into Xarray)

I'm happy to provide performance profiles/tracing as part of this.

Feedback/tips are welcome.

[shoyer]

My suspicion is that we can we something very efficient (at least for linear interpolation), by writing the operation ourselves in terms of shifted indexing operations.

Once you have the appropriate indexes calculated, implementing linear interpolation like np.interp is actually very straightforward in terms of vectorized operations. Here's the core implementation of interp(x, xp, fp) from JAX:

i = clip(searchsorted(xp, x, side='right', 1, len(xp) - 1)
df = fp[i] - fp[i - 1]
dx = xp[i] - xp[i - 1]
delta_x = x - x[i - 1]
f = fp[i] + (df / dx) * delta_x

For Xarray, we would replace indexing like fp[i] with fp.isel({dim: i}), and then I think vectorized operations would automatically support multi-dimensional inputs. Then I think interpolation along multiple dimensions could be done iteratively, because linear interpolation can be calculated in any order.

[dcherian]

I agree. I prototyped a vindex + xr.dot solution here: #6799 (comment) . We really should do this as an alternate "engine" or something.

[shoyer]

I'm not sure xr.dot() along a two-element long dimension is worth the trouble. That's not big enough to get any computational savings from using matrix multiplication kernels.

[dcherian]

Yeah, I was mostly curious if there were memory savings to be had; and whether there was any speedup for interpolating along multiple dimensions. I never got around to profiling those bits though.

[VeckoTheGecko]

Revisiting this, it seems like a lot of the problems we're encountering is due to Dask and how it handles vectorized interpolation - resulting in a huge slowdown. I've gone ahead and put up a reproduction repo, with my results here https://github.com/Parcels-code/xarray-interpolation/blob/main/ANALYSIS.md#cases , and the full context available in the README of that repo. I've gone ahead and tried making that repo as easy to navigate and reproducible as possible.

Basically the TLDR is that interpolating point clouds of data in Xarray using Dask is orders of magnitude slower since:

• There's overhead from Dask (e.g., using a single-threaded executor is actually faster since it doesn't have to wait for locks)
• The interactions between the experimental CacheStore from Zarr and Dask aren't clear - in fact, using the CacheStore seems to take longer
• I don't think its sensitive to the Dask chunking - quickly setting ds.chunk(-1) before profiling doesn't help

I would really appreciate someone taking the time to read the above repo and share any insights they have. I think efficient vectorized indexing is something people would benefit from.

(by the way, I mention vectorized indexing as I assume that interpolation in Xarray uses this. Besides, I think we're more interested in vectorized indexing since we don't necessarily want to use xarrays built-in interpolation. I should have been clearer in my initial post).

cc @dcherian @shoyer