CODEWITHSUNDEEP
pythondaskpython-xarray
TomCho
TomCho

Reputation: 3507

Interpolating an xarray DataArray for N points and getting a list of N interpolated using dask

Sorry if the title isn't very descriptive, but what I want is the following.

I have a DataArray with coordinates x, y and t. I also have a list of N coordinates and I'd like to interpolate to get a list of N interpolated values. However, I don't quite know how to do that with xarray while still taking advantage of the parallelism of dask. Here's an example with random values:

import numpy as np
import xarray as xr

x = np.linspace(0, 1, 10)
datar = xr.DataArray(np.random.randn(10,10,10), dims=('x', 'y', 't'), coords=dict(x=x,
                                                                                  y=x,
                                                                                  t=x))
datar = datar.chunk(dict(t=1))

points = np.array([(0.1, 0.1, 0.1),
                  (0.2, 0.3, 0.3),
                  (0.6, 0.6, 0.6),
                   ])

ivals = []
for point in points:
    x0, y0, t0 = point
    interp_val = datar.interp(x=x0, y=y0, t=t0)
    ivals.append(float(interp_val))
print(ivals)

This gives me the correct result of [-1.7047738779949937, 0.9568015637947849, 0.04437392968785547].

Is there any way to achieve the same result but taking advantage of dask?

If I naively pass lists to the interpolating function I get a 3 cubed matrix instead:

In [35]: x0s, y0s, t0s = points.T
    ...: print(datar.interp(x=x0s, y=y0s, t=t0s))
    ...: 
<xarray.DataArray (x: 3, y: 3, t: 3)>
dask.array<dask_aware_interpnd, shape=(3, 3, 3), dtype=float64, chunksize=(3, 3, 3), chunktype=numpy.ndarray>
Coordinates:
  * x        (x) float64 0.1 0.2 0.6
  * y        (y) float64 0.1 0.3 0.6
  * t        (t) float64 0.1 0.3 0.6

Upvotes: 3

Views: 821

Answers (1)

cyril
cyril

Reputation: 609

A bit late, but in order to interpolate the way you want, and not having a cube as a result, you should cast your coordinates as xarray DataArrays with a fictitious dimension points:

import numpy as np
import xarray as xr

np.random.seed(1234)

x = np.linspace(0, 1, 10)
datar = xr.DataArray(np.random.randn(10, 10, 10), dims=('x', 'y', 't'), coords=dict(x=x, y=x,  t=x))
datar = datar.chunk(dict(t=1))

points = np.array([(0.1, 0.1, 0.1),
                  (0.2, 0.3, 0.3),
                  (0.6, 0.6, 0.6)])

x = xr.DataArray(points[:, 0], dims="points")
y = xr.DataArray(points[:, 1], dims="points")
t = xr.DataArray(points[:, 2], dims="points")

datar.interp(x=x, y=y, t=t).values

It gives you the three values tou want. Two remarks :

  • you should time the executions of the two methods, your loop for and my solution, to check if xarray really takes advantage of the multiple points given to interp,
  • you give the correct values you expect, but they depend on your random data. You should fix the seed before in order to give reproducible examples ;)

Upvotes: 5

Related Questions