Support n-dimensional coordinates

[prisae]

Carrying over from the discussion on SWUNG.

Expand the spline kernel to allow n-dimensional data. In particular, data in 3D (so interpolate [x, y, z, data] to arbitrary points [xi, yi, zi]).

Something like scipy.interpolate.interpn with method=spline, which is only supported for 1D and 2D in scipy.

Loosely related to #138

[prisae]

This does what I had in mind: https://github.com/prisae/tmp-share/blob/master/interp_cubic_spline_3d.py

import numpy as np
from scipy import interpolate, ndimage

def interp_cubic_spline_3d(points, values, xi):
    """3D cubic spline interpolation.
    Zeroes are returned for interpolated values requested outside of the domain
    of the input data.
    """

    # Replicate the same expansion of xi as used in
    # RegularGridInterpolator, so the input xi can be quite flexible.
    xi = interpolate.interpnd._ndim_coords_from_arrays(xi, ndim=3)
    xi_shape = xi.shape
    xi = xi.reshape(-1, 3)

    # map_coordinates uses the indices of the input data as coordinates. We
    # have therefore to transform our desired output coordinates to this
    # artificial coordinate system too.
    params = {'kind': 'cubic',
              'bounds_error': False,
              'fill_value': 'extrapolate'}
    x = interpolate.interp1d(
            points[0], np.arange(len(points[0])), **params)(xi[:, 0])
    y = interpolate.interp1d(
            points[1], np.arange(len(points[1])), **params)(xi[:, 1])
    z = interpolate.interp1d(
            points[2], np.arange(len(points[2])), **params)(xi[:, 2])
    coords = np.vstack([x, y, z])

    # map_coordinates only works for real data; split it up if complex.
    if 'complex' in values.dtype.name:
        real = ndimage.map_coordinates(values.real, coords, order=3)
        imag = ndimage.map_coordinates(values.imag, coords, order=3)
        result = real + 1j*imag
    else:
        result = ndimage.map_coordinates(values, coords, order=3)

    return result.reshape(xi_shape[:-1])

It might (or might not) be of interest for verde to add such a functionality.

[leouieda]

@prisae this is of interest for sure 👍 We would need to have #138 before we can proceed with this, though.

After that is done, the way forward would be to use all given coordinates instead of just coordinates[:2] in the kernels of verde/spline.py. There might be some issues with this down the line or when using Chain with BlockReduce (which is inherently 2D).

Would you be interested in taking a stab at #138 (or know someone who might be)?

[prisae]

I don't think I have the bandwidth at the moment @leouieda , sorry. Particularly as my problem is solved with the above solution. But let's keep it and if more demand is coming from others we can come back to it!