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/saved_outputs%2Fcompare-for-xarray-folks%2FREADME.md , 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

[shoyer]

Hmm! We might need specialized dask logic for efficient vectorized indexing. Vectorized indexing is inherently very hard to parallelize in dask's data model, because the communication patterns can be so variable. Some indexing patterns inherently require all-to-all communication, whereas others are embarrassingly parallel. This information doesn't get preserved in dask.array's data model.

[VeckoTheGecko]

I agree that improving dask here would be really beneficial. In the meantime, I was also wondering whether it would be possible to bypass it directly

I did some profiling without Xarray (i.e. just numpy and Zarr) and the results are really good for our usecase https://github.com/Parcels-code/xarray-interpolation/blob/main/saved_outputs%2Fcompare-from-raw-zarr%2FREADME.md .

Is there a way that we can get Xarray to wrap Zarr arrays directly by specifying a custom Zarr backend? This would forgo a lot of the lazy functionality - but would be awesome if the format on disk already matches our expectations.

If this is possible, I can see a "best of both multiple worlds" solution for our users.

Though this is of course dependent on how flexible Xarray is for wrapping arrays (i.e. if it assumes array API, or is also happy with a subset of it).

[shoyer]

If you open a Zarr with open_zarr(..., chunks=None) it just uses NumPy, no Dask. This should have comparable performance to using NumPy/Zarr alone, and if not, we can and should fix it!

[VeckoTheGecko]

I think the message of our problem is getting lost a bit here, so I'll try word it a bit differently :) :

Dask has been prohibitively difficult to work with due to the slow 'out of the box' performance of the vectorized interpolation. My exploration in #11236 (reply in thread) was looking at what happens if we forget about Dask and Xarray for a moment, and just assume that the data exists on disk in a specific format - what does vectorized interpolation look like on that going from the "raw Zarr"?

As we were hoping, vectorized interpolation using the zarr-python is very performant (also having the benefit of not having to load everything into memory).

We are working with large climate model output (i.e., multiple fields from a single - or even multiple - climate model), but we don't need all the data in memory since:

• we are stepping through time using a for loop - temporally the particles are all in pretty much the same chunks
• we don't necessarily have full spatial coverage - particles aren't necessarily in the full spatial extent

Having all the data in memory would be a huge waste of memory and be prohibitive for users.

I think ultimately we want better vectorized interpolation with Dask. But (assuming the data on disk is in exactly the format that we expect internally - which is a significant assumption) supporting workflows where people work directly from Zarr using zarr-python might also be a nice alternative.

The main thing I was wondering about is if Xarray can wrap zarr.Array objects without Dask and without loading into Numpy - simplifying our API somewhat since we just use dataset.isel . I assume that this isn't currently possible since zarr.Array isn't, nor tries to be, array API compatible (it doesn't even support addition of arrays - there's no lazy-execution backend as that extends beyond the scope of the package). I also assume that this "read only" view into Zarr datasets via Xarray is on the roadmap, or (understandably) out of scope for Xarray?

We might need specialized dask logic for efficient vectorized indexing.

Do you know of good resources about working with this dask logic?

[shoyer]

Right -- so Xarray does in fact have lazy loading machinery for data being read from disk. It only supports indexing operations, and data gets automatically converted in NumPy when you do anything fancier. It is used automatically every time you open a dataset in Xarray, even if you use Dask, which just becomes an optional layer on top.