Using fill_value = (fp[0], fp[-1]) with interp_like on Dataset

[jkoell]

I have several datasets of ocean data from different climate models that I'm trying to interpolate to a common vertical grid using xarray.interp_like. The depth of the surface cell isn't the same on each model, so I'm trying to interpolate in a way that mirrors the behavior of numpy.interp, i.e. filling in with the first and last array values, fp[0] and fp[-1], outside the range of the input data. But since interp_like uses the SciPy interpolator it instead fills with np.nan. The scipy documentation has an example of how to implement the same behavior in scipy, but that syntax really only works if you do it directly in scipy rather than passing it through xarray.

I've tried a few different things for the fill_value kwarg in interp_like but couldn't come up with anything that gives me the desired behavior. The documentation for interp_like mentions numpy is used in some cases, but it's not clear to me if there is a way to choose whether numpy or scipy is used.

So my question is, is there some way to tell xarray interp_like to use the constant extrapolation for values outside the range? This seems like it should be a fairly common use case, but I couldn't find anything suggesting this is currently possible

Example input:

import numpy as np
import xarray as xr

ds0 = xr.Dataset(coords={'z': [0, 5, 10, 20]})

ds = xr.Dataset(
    data_vars=dict(
        tmp=(['z'], [10, 5, 4, 4]),
        sal=(['z'], [35, 35.1, 35.1, 35.1]),
    ),
    coords=dict(
        z = [2.5, 5, 10, 20]
    ),
)

Output:

<xarray.Dataset> Size: 96B
Dimensions:  (z: 4)
Coordinates:
  * z        (z) int64 32B 0 5 10 20
Data variables:
    tmp      (z) float64 32B nan 5.0 4.0 4.0
    sal      (z) float64 32B nan 35.1 35.1 35.1

Desired behavior:

I'd like to interpolate the same way as numpy would, e.g. for the tmp variable:

np.interp(ds0['z'], ds['z'], ds['tmp'])
Out[3]: array([10.,  5.,  2.,  2.])

Extrapolation in interp_like:

If I instead use kwargs={'fill_value': 'extrapolate'}, the values are much different from the constant fill value, since it uses linear rather than constant extrapolation:

ds.interp_like(ds0, kwargs={'fill_value': 'extrapolate'})
Out[342]: 
<xarray.Dataset> Size: 96B
Dimensions:  (z: 4)
Coordinates:
  * z        (z) int64 32B 0 5 10 20
Data variables:
    tmp      (z) float64 32B 15.0 5.0 2.0 2.0
    sal      (z) float64 32B 34.9 35.1 35.1 35.1

[tayfuryldz]

For a single DataArray, SciPy's two-element fill_value=(left, right) does work through xarray:

ds["tmp"].interp_like(
    ds0,
    kwargs={"fill_value": (ds["tmp"].isel(z=0).item(),
                           ds["tmp"].isel(z=-1).item())},
)

The complication with your Dataset is that kwargs is one set of interpolation options applied while xarray interpolates the variables; the desired endpoint pair is different for tmp and sal. A single tuple therefore cannot mean "take each variable's own first and last value".

For a Dataset, one straightforward solution is to interpolate each data variable with its own endpoint tuple and combine them again:

result = xr.Dataset(
    {
        name: da.interp_like(
            ds0,
            kwargs={
                "fill_value": (
                    da.isel(z=0).item(),
                    da.isel(z=-1).item(),
                )
            },
        )
        for name, da in ds.data_vars.items()
    },
    coords=ds0.coords,
)

With the example in the question this gives constant endpoint extrapolation independently for tmp and sal, rather than linear extrapolation.

I tested the relevant call against current xarray before posting: a two-element fill_value is passed through successfully for a 1-D DataArray and produces constant left/right values. So you don't need to force xarray to select NumPy's interpolator; using SciPy's documented tuple form is sufficient. The per-variable loop is needed because the endpoint constants are variable-specific in a Dataset.

For large/lazy arrays, avoid .item() if forcing eager computation is undesirable; obtain or construct the endpoint scalars in a way appropriate for that data source instead.