Interpolate for multiple data variables with a single interpolate call

[mcleantom]

Is your feature request related to a problem? Please describe.
I want to interpolate for multiple variables on the same grid (My grib data holds wind speed, wind direction U, wind direction V etc). Currently, I need to make an interpolator for each variable like:

interpolators = {
    'speed': pyinterp.backends.xarray.Grid3D(ds.speed),
    'u': pyinterp.backends.xarray.Grid3D(ds.u),
    'v': pyinterp.backends.xarray.Grid3D(ds.v),
}

So when I want to interpolate for given lat, lon, time I have to interpolate the grid three times even though they are at the same point in the 3D grid. I think it would be more efficient to interpolate each of the three variables simultaneously.

Describe the solution you'd like
Be able to make an interpolator from an XArray dataset, rather than a DataArray.

Example:

import numpy
import pyinterp.backends.xarray
import pyinterp.tests

ds = pyinterp.tests.load_grid3d()
ds['something_else'] = ds.time.astype(int)/1e16 + ds.latitude + ds.longitude
interpolator = pyinterp.backends.xarray.Grid3D(ds)

mx, my, mz = numpy.meshgrid(numpy.arange(-180, 180, 0.25) + 1 / 3.0,
                            numpy.arange(-80, 80, 0.25) + 1 / 3.0,
                            numpy.array(['2002-07-02T15:00:00'],
                                        dtype='datetime64'),
                            indexing='ij')

trivariate = interpolator.trivariate(
    dict(longitude=mx.ravel(), latitude=my.ravel(), time=mz.ravel()))

Which will result in something like:

{'tcw': array([ 1.00482318,  1.01251311,  1.02020304, ..., 11.12754145,
       11.10233118, 11.07064156]),
 'something_else': array([203.22886667, 203.47886667, 203.72886667, ..., 362.22886667,
       362.47886667, 362.72886667])
}

Describe alternatives you've considered
I tried to stack the data axis and use a 4d interpolator:

interpolator = pyinterp.backends.xarray.Grid4D(ds.to_array())

However, the variable axis is not supported (it is not np.float64)

I tried dropping the variable axis:

interpolator = pyinterp.backends.xarray.Grid4D(ds.to_array().drop("variable"))

Which makes the interpolator, but I cant seem to interpolate it anywhere:

interpolator.quadrivariate(dict(longitude=mx.ravel(), latitude=my.ravel(), time=mz.ravel()))

Traceback (most recent call last):
  File "/home/tom.mclean/miniconda3/envs/routingserver/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3378, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-38-859473664d87>", line 1, in <module>
    tmp.quadrivariate(dict(longitude=mx.ravel(), latitude=my.ravel(), time=mz.ravel()))
  File "/home/tom.mclean/miniconda3/envs/routingserver/lib/python3.10/site-packages/pyinterp/backends/xarray.py", line 390, in quadrivariate
    self, *_coords(coords, self._dims, self._datetime64), *args,
  File "/home/tom.mclean/miniconda3/envs/routingserver/lib/python3.10/site-packages/pyinterp/backends/xarray.py", line 129, in _coords
    raise IndexError('too many indices for array')
IndexError: too many indices for array

Or:

tmp.quadrivariate(dict(longitude=mx.ravel(), latitude=my.ravel(), time=mz.ravel(), variable="tcw"))
Traceback (most recent call last):
  File "/home/tom.mclean/miniconda3/envs/routingserver/lib/python3.10/site-packages/IPython/core/interactiveshell.py", line 3378, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-39-3ee1e6de957d>", line 1, in <module>
    tmp.quadrivariate(dict(longitude=mx.ravel(), latitude=my.ravel(), time=mz.ravel(), variable="tcw"))
  File "/home/tom.mclean/miniconda3/envs/routingserver/lib/python3.10/site-packages/pyinterp/backends/xarray.py", line 389, in quadrivariate
    return interpolator.quadrivariate(
  File "/home/tom.mclean/miniconda3/envs/routingserver/lib/python3.10/site-packages/pyinterp/interpolator/quadrivariate.py", line 55, in quadrivariate
    return getattr(core, function)(instance,
TypeError: quadrivariate_float64(): incompatible function arguments. The following argument types are supported:
    1. (grid: pyinterp.core.Grid4DFloat64, x: numpy.ndarray[numpy.float64], y: numpy.ndarray[numpy.float64], z: numpy.ndarray[numpy.float64], u: numpy.ndarray[numpy.float64], interpolator: pyinterp.core.BivariateInterpolator3D, z_method: Optional[str] = None, u_method: Optional[str] = None, bounds_error: bool = False, num_threads: int = 0) -> numpy.ndarray[numpy.float64]
    2. (grid: pyinterp.core.TemporalGrid4DFloat64, x: numpy.ndarray[numpy.float64], y: numpy.ndarray[numpy.float64], z: numpy.ndarray[numpy.int64], u: numpy.ndarray[numpy.float64], interpolator: pyinterp.core.TemporalBivariateInterpolator3D, z_method: Optional[str] = None, u_method: Optional[str] = None, bounds_error: bool = False, num_threads: int = 0) -> numpy.ndarray[numpy.float64]
Invoked with: <pyinterp.core.TemporalGrid4DFloat64 object at 0x7feb47eb4d70>, array([-179.66666667, -179.66666667, -179.66666667, ...,  180.08333333,
        180.08333333,  180.08333333]), array([-79.66666667, -79.41666667, -79.16666667, ...,  79.58333333,
        79.83333333,  80.08333333]), array(['2002-07-02T15:00:00.000000000', '2002-07-02T15:00:00.000000000',
       '2002-07-02T15:00:00.000000000', ...,
       '2002-07-02T15:00:00.000000000', '2002-07-02T15:00:00.000000000',
       '2002-07-02T15:00:00.000000000'], dtype='datetime64[ns]'), array('tcw', dtype='<U3'), <pyinterp.core.TemporalBilinear3D object at 0x7feb8023b6f0>; kwargs: z_method='linear', u_method='linear', bounds_error=False, num_threads=0

[fbriol]

The library allows to use, for the moment, one interpolator per variable. A variable corresponds to a grid (2D, 3D or 4D) with its different axes. In the xarray terminology, it corresponds to a DataArray.

Indeed, it could be useful to interpolate several variables with a single call to a function. But this has not been implemented. For that, we have to take a Dataset as arguments.

Maybe we will do this in the future. For that we will have to take into account the Dataset. One way is to use xtensor.