Lost in apply_ufunc land without a map! Trying to find a faster way to apply interp1d.

[rodgerduffett]

Dear All,
I'm using xarray and dask to process wind speed data. The dataset has dimensions for time, level and space as a grid. The important variables are wind speed as U and V components and height above ground for each level. I need to interpolate U and V for a target height at each of the grid cells and of course for each time step. I've been able to do this using for loops and a dask delayed function that wraps scipy's interp1d function.

It works reasonably well as long as a small enough subset is selected out in the nested loop structure. But it seems an inefficient way to do things.

I've been trying to understand apply_ufunc and map_blocks in this context without success! The main conceptual issue being that interp1d needs an x and y as equal length vectors. These exist in the data as the 'level' dimension for the height and U/V variables.
Some code may help explain. Simplified data looks like this:

import xarray as xr
from scipy import interpolate
target_height=15
time = pd.date_range("2019-09-06", periods=2)
heights=np.array([
    [
        [[10,10,10,10,10],[10,10,10,10,10],[10,10,10,10,10],[10,10,10,10,10]],
        [[20,20,20,20,20],[20,20,20,20,20],[20,20,20,20,20],[20,20,20,20,20]],
        [[30,30,30,30,30],[30,30,30,30,30],[30,30,30,30,30],[30,30,30,30,30]]
    ],
    [
        [[10,10,10,10,10],[10,10,10,10,10],[10,10,10,10,10],[10,10,10,10,10]],
        [[20,20,20,20,20],[20,20,20,20,20],[20,20,20,20,20],[20,20,20,20,20]],
        [[30,30,30,30,30],[30,30,30,30,30],[30,30,30,30,30],[30,30,30,30,30]]
    ]
])
u=np.array([
    [
        [[1,1,1,1,1],[1,1,1,1,1],[1,1,1,1,1],[1,1,1,1,1]],
        [[2,2,2,2,2],[2,2,2,2,2],[2,2,2,2,2],[2,2,2,2,2]],
        [[3,3,3,3,3],[3,3,3,3,3],[3,3,3,3,3],[3,3,3,3,3]]
    ],
    [
        [[1,1,1,1,1],[1,1,1,1,1],[1,1,1,1,1],[1,1,1,1,1]],
        [[2,2,2,2,2],[2,2,2,2,2],[2,2,2,2,2],[2,2,2,2,2]],
        [[3,3,3,3,3],[3,3,3,3,3],[3,3,3,3,3],[3,3,3,3,3]]
    ]
])
v=np.array([
    [
        [[3,3,3,3,3],[3,3,3,3,3],[3,3,3,3,3],[3,3,3,3,3]],
        [[4,4,4,4,4],[4,4,4,4,4],[4,4,4,4,4],[4,4,4,4,4]],
        [[5,5,5,5,5],[5,5,5,5,5],[5,5,5,5,5],[5,5,5,5,5]]
    ],
    [
       [[3,3,3,3,3],[3,3,3,3,3],[3,3,3,3,3],[3,3,3,3,3]],
        [[4,4,4,4,4],[4,4,4,4,4],[4,4,4,4,4],[4,4,4,4,4]],
        [[5,5,5,5,5],[5,5,5,5,5],[5,5,5,5,5],[5,5,5,5,5]]
    ]
])
uvh_ds=xr.Dataset()
uvh_ds[f"H"]=(xr.DataArray(data=heights,
                                   dims=['Time','level','i','j'],
                                   coords=dict(Time=time),
                                   attrs=dict(description=f"Height above ground level [m]")))
uvh_ds[f"U"]=(xr.DataArray(data=u,
                                   dims=['Time','level','i','j'],
                                   coords=dict(Time=time),
                                   attrs=dict(description=f"u component of wind speed [m/s]")))
uvh_ds[f"V"]=(xr.DataArray(data=v,
                                   dims=['Time','level','i','j'],
                                   coords=dict(Time=time),
                                   attrs=dict(description=f"v component of wind speed [m/s]")))

What would be the most efficient way to apply:

result=interpolate.interp1d(x,y)(target_height)

at each time step and each grid cell?

At the moment I use something like:

@delayed
def intrp(x,y,target_height):
    try:
        result=np.fromiter((interpolate.interp1d(x[ts],y[ts])(target_height) for ts in range(len(x))),dtype='float32')
    except Exception as err:
        print(f"Died with {err}\n")
    return result
for ii in range(0,len(uvh_ds.i)):
    for jj in range(0,len(uvh_ds.j)):
       u_futures.append(u=intrp(uvh_ds.H[:,:,ii,jj],uvh_ds.U[:,:,ii,jj],target_height))

The futures must be computed and the results reshaped back to time,i,j shape before saving to disk.

This works as long as the number of time steps submitted in each pass is kept reasonable and the size of the grid like wise reduced. For bigger datasets it means a lot of looping and selecting. In addition the results come back to the local compute node.

Advice on how to approach this problem in a better way will be much appreciated!

[dcherian]

See https://docs.xarray.dev/en/stable/examples/apply_ufunc_vectorize_1d.html for this exact problem. Let us know if it can be improved in any way

[rodgerduffett]

Thank-you! I will work through that and report back!

[rodgerduffett]

The tutorial is nicely written!
My case is slightly different as the 'x' values for the interp are stored in a variable - uvh_ds.H. The 'y' is in uvh_ds.U. So it is not interpolating along a dimension. Perhaps it is not possible to use apply_ufunc in this way and I must look to try improve the loops?

[dcherian]

then pass H but set the core dimension appropriately. In the tutorial both the variable and core dimension are lat but this isn't necessary. This would be a nice addition to the tutorial

[rodgerduffett]

Am still working on this and will report back when I have made progress.

[rodgerduffett]

Finally got back to working on this!

A working minimal code example is:

import xarray as xr
from scipy import interpolate
target_height=15
time = pd.date_range("2019-09-06", periods=2)
heights=np.array([
    [
        [[10,10,10,10,10],[10,10,10,10,10],[10,10,10,10,10],[10,10,10,10,10]],
        [[20,20,20,20,20],[20,20,20,20,20],[20,20,20,20,20],[20,20,20,20,20]],
        [[30,30,30,30,30],[30,30,30,30,30],[30,30,30,30,30],[30,30,30,30,30]]
    ],
    [
        [[10,10,10,10,10],[10,10,10,10,10],[10,10,10,10,10],[10,10,10,10,10]],
        [[20,20,20,20,20],[20,20,20,20,20],[20,20,20,20,20],[20,20,20,20,20]],
        [[30,30,30,30,30],[30,30,30,30,30],[30,30,30,30,30],[30,30,30,30,30]]
    ]
])
u=np.array([
    [
        [[1,1,1,1,1],[1,1,1,1,1],[1,1,1,1,1],[1,1,1,1,1]],
        [[2,2,2,2,2],[2,2,2,2,2],[2,2,2,2,2],[2,2,2,2,2]],
        [[3,3,3,3,3],[3,3,3,3,3],[3,3,3,3,3],[3,3,3,3,3]]
    ],
    [
        [[11,1,1,1,1],[1,1,1,1,1],[1,1,1,1,1],[1,1,1,1,1]],
        [[12,2,2,2,2],[2,2,2,2,2],[2,2,2,2,2],[2,2,2,2,2]],
        [[13,3,3,3,3],[3,3,3,3,3],[3,3,3,3,3],[3,3,3,3,3]]
    ]
])
v=np.array([
    [
        [[3,3,3,3,3],[3,3,3,3,3],[3,3,3,3,3],[3,3,3,3,3]],
        [[4,4,4,4,4],[4,4,4,4,4],[4,4,4,4,4],[4,4,4,4,4]],
        [[5,5,5,5,5],[5,5,5,5,5],[5,5,5,5,5],[5,5,5,5,5]]
    ],
    [
       [[3,3,3,3,3],[3,3,3,3,3],[3,3,3,3,3],[3,3,3,3,3]],
        [[4,4,4,4,4],[4,4,4,4,4],[4,4,4,4,4],[4,4,4,4,4]],
        [[5,5,5,5,5],[5,5,5,5,5],[5,5,5,5,5],[5,5,5,5,5]]
    ]
])
uvh_ds=xr.Dataset()
uvh_ds[f"H"]=(xr.DataArray(data=heights,
                                   dims=['Time','bottom_top','i','j'],
                                   coords=dict(Time=time),
                                   attrs=dict(description=f"Height above ground level [m]")))

uvh_ds[f"U"]=(xr.DataArray(data=u,
                                   dims=['Time','bottom_top','i','j'],
                                   coords=dict(Time=time),
                                   attrs=dict(description=f"u component of wind speed [m/s]")))
uvh_ds[f"V"]=(xr.DataArray(data=v,
                                   dims=['Time','bottom_top','i','j'],
                                   coords=dict(Time=time),
                                   attrs=dict(description=f"v component of wind speed [m/s]")))

def intrp(x,y,target_height):
    print(f"intrp with:- x: {x},y: {y}, target:{target_height}")
    try:     
        result=interpolate.interp1d(x,y)(target_height)
    except Exception as err:
        print(f"Died with {err}\n")
    return result

target_height=np.array([15,])
print(f"target_height: {target_height.shape}")
interped=xr.apply_ufunc(intrp,uvh_ds.H,uvh_ds.U,target_height,
                input_core_dims=[['bottom_top'],['bottom_top'],["new_height"]],
                output_core_dims=[["new_height"]],  # returned data has one dimension
                exclude_dims=set(("bottom_top",)),  # dimensions allowed to change size. Must be a set!
                vectorize=True,  # loop over non-core dims
               )

This will return an array of shape [Time,i,j,1]

For parallel execution via dask:

interped=xr.apply_ufunc(intrp,uvh_ds.H,uvh_ds.U,target_height,
                input_core_dims=[['bottom_top'],['bottom_top'],["new_height"]],
                output_core_dims=[["new_height"]],  # returned data has one dimension
                exclude_dims=set(("bottom_top",)),  # dimensions allowed to change size. Must be a set!
                vectorize=True,  # loop over non-core dims
                dask="parallelized",
                output_dtypes=[uvh_ds.H.dtype],
               )