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.