2D interpolation

[jecampagne]

Hello,
I'm looking for a JAX 2D interpolation implementation but do not yet find one, an idea?
Thanks

[hawkinsp]

Does jax.scipy.ndimage.map_coordinates do what you need? If not, can you say more about what exactly you are trying to compute?

[jecampagne]

Thanks for your interest. Well I am looking at least to an implementation of bilinear interpolation while the "best" would be 2D-spline (order 3 for instance). Googling I have found this one here but it does not accept that x and y have different sizes. Although it seems that having different sizes can be miss leading so may be a more solid implementation would be nice. And of course the spline 2D would be very cool.

The idea of the API would be to have something like

xp=jnp.linspace(-1.,1.,20)  # N
yp=jnp.linspace(-1.,1.,10)  # M

X, Y = np.meshgrid(xp, yp)

def func(x,y):
    return jnp.exp(-0.5*(x**2+y**2))

#Z = func(X,Y)
#plt.contourf(xp,yp,Z)  #### here fp MxN for contourf

fp = Z.T    # N x M   with N: size of xp, M size of yp

x_star = jnp.array([-0.5,0.1, 0.5])
y_star = jnp.array([0.7,-0.7,0.0])

f_star = jax______.spline2d(x_star, y_star,xp,yp,fp) 

[YouJiacheng]

Do you want to interpolation on regular grid(i.e. same space along one axis, but each axis can have different space)?
It seems that interp2d of https://github.com/adam-coogan/jaxinterp2d accepts x and y with different sizes, but CartesianGrid doesn't.

[YouJiacheng]

Okay I see.

[YouJiacheng]

@jecampagne
xp and yp can have different size.
xstar and ystar should have same size, since point i coordinate is (xstar[i], ystar[i])!

[jecampagne]

@jecampagne xp and yp can have different size. xstar and ystar should have same size, since point i coordinate is (xstar[i], ystar[i])!

Ok :) for xp, yp, now for xstar,ystarthere are 2 use cases:

• the one you mention fp[i]=f(xstar[i],ystar[i]) i<Ni
• by one could also predict on a grid fp[i,j] = f(xstar[i],ystar[j]) with i<Ni, j<Nj

[jakevdp]

jnp.meshgrid may be helpful here to turn your different-sized inputs into a grid where each array has the same size.

[jecampagne]

Right @YouJiacheng point out this 'trick'.
So, now the numerical exercice is to perform higher order interpolation (not just linear) and spline (1D & 2D) of order 3.

[YouJiacheng]

@jecampagne It might takes lots of effort to implement scipy.interpolate.interp2d which support B-spline of order 1~5, but it is possible. Would you like to port this to JAX?
https://github.com/scipy/scipy/blob/97f6127299bb84e5bb67a62d272b9fa1a36832f2/scipy/interpolate/_fitpack_impl.py#L798-L1062

[jecampagne]

Starting with 1d and then 2d would be super cool. I will ask some friends working on Goe-science because I wander how they deal with 3d representation of their surfaces (landscape, mountains...
).

[jecampagne]

Hello, I'm come back with a solution I've cooked:

@jit
def bispline_interp(xnew,ynew,xp,yp,zp):
    """
    (xnew,ynew): two 1D vector  of same size where to perform predictions  f(xnew[i],ynew[i])
    (xp,yp): original grid points 1D vector
    zp: original values of functions  zp[i,j] = value at xp[i], yp[j]
    """
    
    
    M = 1./16 * jnp.array([[0, 0, 0, 0, 0, 16, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 
                           [0, 0, 0, 0, -8, 0, 8, 0, 0, 0, 0, 0, 0, 0, 0, 0], 
                           [0, 0, 0, 0, 16, -40, 32, -8, 0, 0, 0, 0, 0, 0, 0, 0], 
                           [0, 0, 0, 0, -8, 24, -24, 8, 0, 0, 0, 0, 0, 0, 0, 0],
                           [0, -8, 0, 0, 0, 0, 0, 0, 0, 8, 0, 0, 0, 0, 0, 0], 
                           [4, 0, -4, 0, 0, 0, 0, 0, -4, 0, 4, 0, 0, 0, 0, 0], 
                           [-8, 20, -16, 4, 0, 0, 0, 0, 8, -20, 16, -4, 0, 0, 0, 0],
                           [4, -12, 12, -4, 0, 0, 0, 0, -4, 12, -12, 4, 0, 0, 0, 0],
                           [0, 16, 0, 0, 0, -40, 0, 0, 0, 32, 0, 0, 0, -8, 0, 0], 
                           [-8, 0, 8, 0, 20, 0, -20, 0, -16, 0, 16, 0, 4, 0, -4, 0], 
                           [16, -40, 32, -8, -40, 100, -80, 20, 32, -80, 64, -16, -8, 20, -16, 4], 
                           [-8, 24, -24, 8, 20, -60, 60, -20, -16, 48, -48, 16, 4, -12, 12, -4], 
                           [0, -8, 0, 0, 0, 24, 0, 0, 0, -24, 0, 0, 0, 8, 0, 0], 
                           [4, 0, -4, 0, -12, 0, 12, 0, 12, 0, -12, 0, -4, 0, 4, 0], 
                           [-8, 20, -16, 4, 24, -60, 48, -12, -24, 60, -48, 12, 8, -20, 16, -4], 
                           [4, -12, 12, -4, -12, 36, -36, 12, 12, -36, 36, -12, -4, 12, -12, 4]]
                         )
    
    M1 = jnp.array([[1.,0.,0.,0.],
                    [-1.,1.,0.,0.],
                    [-1.,0.,1.,0.],
                    [1.,-1.,-1.,1.]])

    def built_Ivec(zp,ix,iy):
        return jnp.array([zp[ix+i,iy+j] for j in range(-1,3) for i in range(-1,3)])


    def built_Ivec1(zp,ix,iy):
        return jnp.array([zp[ix+i,iy+j] for j in range(0,2) for i in range(0,2)])

    
    
    def compute_basis(x,order=3):
        """
        x in [0,1]
        """ 
        return jnp.array([x**i for i in jnp.arange(0, order+1)])
    
    def tval(xnew,ix,xp):
        return (xnew-xp[ix-1])/(xp[ix]-xp[ix-1])
    
    ix = jnp.clip(jnp.searchsorted(xp, xnew, side="right"), 0, len(xp)-1)
    iy = jnp.clip(jnp.searchsorted(yp, ynew, side="right"), 0, len(yp)-1)

    def bilinear_interp(ix,iy):
        Iv = built_Ivec1(zp,ix-1,iy-1)
        av = M1 @ Iv
        amtx = av.reshape(2,2,-1)
        tx = tval(xnew,ix,xp)
        ty = tval(ynew,iy,yp)
        basis_x = compute_basis(tx,order=1)
        basis_y = compute_basis(ty,order=1)
        res = jnp.einsum("i...,ij...,j...",basis_y,amtx,basis_x)
        return res

    def bispline_interp(ix,iy):
        Iv = built_Ivec(zp,ix-1,iy-1)
        av = M @ Iv
        amtx = av.reshape(4,4,-1)
        tx = tval(xnew,ix,xp)
        ty = tval(ynew,iy,yp)
        basis_x = compute_basis(tx)
        basis_y = compute_basis(ty)
        res = jnp.einsum("i...,ij...,j...",basis_y,amtx,basis_x)
        return res
    
    condx = jnp.logical_and(ix>=2, ix<=len(xp)-2)
    condy = jnp.logical_and(iy>=2, iy<=len(yp)-2)
    
    cond = jnp.logical_and(condx,condy)
    return jnp.where(cond,
             bispline_interp(ix,iy),
             bilinear_interp(ix,iy))

A possible usage

Nx=40
Ny=80
xp = jnp.linspace(0,8.,Nx)
yp = jnp.linspace(0.,6.,Ny)

xgp, ygp = jnp.meshgrid(xp, yp)

def func(x,y):
    xn = (x-4)/2.
    yn = (y-3)/2.5
    return jnp.exp(-0.5*(xn**2 + yn**2))

data = func(xgp,ygp).T   



xnews = jnp.linspace(np.min(xp),np.max(xp),Nx*2)
ynews = jnp.linspace(np.min(yp),np.max(yp),Ny*2)
xnews, ynews = jnp.meshgrid(xnews, ynews)
xnews_flat = xnews.reshape((-1,))
ynews_flat = ynews.reshape((-1,))

interp_data = bispline_interp(xnews_flat,ynews_flat,xp,yp,data)
interp_data = interp_data.reshape(xnews.shape)

true_data = func(xnews, ynews)

# plot
fig= plt.figure(figsize=(20,6))
gs = fig.add_gridspec(1, 4, hspace=10, wspace=0)
axs = gs.subplots()
g0=axs[0].imshow(data.T, interpolation=None,);
axs[0].set_title("Low resol.")
g1=axs[1].imshow(interp_data, interpolation=None);
axs[1].set_title("Spline Interp.")
g2=axs[2].imshow(true_data, interpolation=None);
axs[2].set_title("Truth")
g3 = axs[3].imshow(interp_data-true_data, interpolation=None);

for ax in axs.flat:
    ax.label_outer()

fig.colorbar(g0, ax=axs[0])
fig.colorbar(g1, ax=axs[1])
fig.colorbar(g2, ax=axs[2])
fig.colorbar(g3, ax=axs[3])

This gives

Comparing with bilinear interpolation I get better results

I guess my code is not the cleanest one nor the best implementation considering JAX/JIT. Concerning the numerical side, one can certainly get better results for the sides & corner.

I hope people will improve it....

[davindicode]

What kind of cubic spline is this implementing? Natural, akima? Thanks for this!

[jecampagne]

Well, I do not think this is Akima spline. But let me try to explain how I have proceeded.
On a regular 2D grid (xi,yj), I can locally use this parametrisation

I need then beyond f(x,y) with (x,y) in {0,1},
f(x,y) = \sum_{i=0}^3 \sum_{j=0} a_{ij} x^i y^j
12 derivatives
f_x(x,y) = \partial_x f(x,y) = \sum_{i=0}^3 \sum_{j=0} a_{ij} i x^i y^j
f_y(x,y) = \partial_y f(x,y) = \sum_{i=0}^3 \sum_{j=0} a_{ij} j x^i y^j
f_{xy}(x,y) = \partial_x f(x,y) = \sum_{i=0}^3 \sum_{j=0} a_{ij} i j x^i y^j

Then, from the data 2D (noted I(x,y) like en image) I can compute
I(x,y) which is directly related to f(x,y) and for the derivative I use these finite difference approx:
f_x(x,y) <-> 1/2 [ I(x+1,y) - I(x-1,y)]
f_y(x,y) <-> 1/2 [ I(x,y+1) - I(x,y-1)]
f_{xy}(x,y) <-> 1/4 [ I(x+1,y+1) - I(x-1,y+1)-I(x+1,y-1) + I(x-1,y-1)]
which needs a 4x4 local grids as in the image above.

Now, if one collect of the (f(x,y), f_x(x,y), f_y(x,y), f_{xy}(x,y)}{(x,y)\in{0,1}} in a 1D vector noted F, then I can write
F = A a
with A a matrix and a the (a{i,j}) flatten as a 1D vector
Similarly, on gets
F = B I
with B another matrix and the I a 1D vector with the I(x,y) involved above.

Then, one gets finally
a = (A^-1 B) I := M I
the M matrix is the one I set at the beginning of the routine.

Now, concerning the edges, I do a simple linear fit, in the local square which do not need prescription as reflection/zeros or what ever padding of the data.

Well, I may be wrong at the math so if you have a benchmark, this would be great. Also, if one can improve the edges (sides & corners) that would be great too. Finally, if one find a way to perform non-regular sampling then it would be wonderful (I only keep-up with regular linear or regular log).

[jecampagne]

Need not to say that my JAX implementation may not be optimized as I am a beginner.

[jecampagne]

Does someone can make a code review form the point of view of JAX better coding? That could be great.

[zeeshan5885]

@jecampagne I really loved your solution, I tried to implement to use your code on my problem and it seems working. Did you written this interpolation in some published package. ?
I need to use this.

[f0uriest]

As another option, I've been developing interpax to try to recover as much of the functionality from scipy.interpolate as possible. Currently has options for linear and cubic splines in 1,2,and 3 dimensions, and I'm still working on building more of it out (contributions welcome!)

[jecampagne]

If you need I can try to give you some help.