Python Scipy for 2D extrapolated spline function?

Question

I'd like to write an extrapolated spline function for a 2D matrix. What I have now is an extrapolated spline function for 1D arrays as below. scipy.interpolate.InterpolatedUnivariateSpline() is used.

import numpy as np 
import scipy as sp 

def extrapolated_spline_1D(x0,y0):
    x0 = np.array(x0)
    y0 = np.array(y0)
    assert x0.shape == y.shape 

    spline = sp.interpolate.InterpolatedUnivariateSpline(x0,y0)
    def f(x, spline=spline):
        return np.select(
            [(x<x0[0]),              (x>x0[-1]),              np.ones_like(x,dtype='bool')], 
            [np.zeros_like(x)+y0[0], np.zeros_like(x)+y0[-1], spline(x)])

    return f

It takes x0, which is where the function is defined, and y0, which is the according values. When x < x0[0], y = y0[0]; and when x > x0[-1], y = y0[-1]. Here, assuming x0 is in an ascending order.

I want to have a similar extrapolated spline function for dealing with 2D matrices using np.select() as in extrapolated_spline_1D. I thought scipy.interpolate.RectBivariateSpline() might help, but I'm not sure how to do it.

For reference, my current version of the extrapolated_spline_2D is very slow. The basic idea is:

(1) first, given 1D arrays x0, y0 and 2D array z2d0 as input, making nx0 extrapolated_spline_1D functions, y0_spls, each of which stands for a layer z2d0 defined on y0;

(2) second, for a point (x,y) not on the grid, calculating nx0 values, each equals to y0_spls[i](y);

(3) third, fitting (x0, y0_spls[i](y)) with extrapolated_spline_1D to x_spl and returning x_spl(x) as the final result.

def extrapolated_spline_2D(x0,y0,z2d0): 
    '''    
    x0,y0 : array_like, 1-D arrays of coordinates in strictly monotonic order. 
    z2d0  : array_like, 2-D array of data with shape (x.size,y.size).
    '''    
    nx0 = x0.shape[0]
    ny0 = y0.shape[0]
    assert z2d0.shape == (nx0,ny0)

    # make nx0 splines, each of which stands for a layer of z2d0 on y0 
    y0_spls = [extrapolated_spline_1D(y0,z2d0[i,:]) for i in range(nx0)]

    def f(x, y):     
        '''
        f takes 2 arguments at the same time --> x, y have the same dimention
        Return: a numpy ndarray object with the same shape of x and y
        '''
        x = np.array(x,dtype='f4')
        y = np.array(y,dtype='f4') 
        assert x.shape == y.shape        
        ndim = x.ndim 

        if ndim == 0:    
            '''
            Given a point on the xy-plane. 
            Make ny = 1 splines, each of which stands for a layer of new_xs on x0
            ''' 
            new_xs = np.array([y0_spls[i](y) for i in range(nx0)]) 
            x_spl  = extrapolated_spline_1D(x0,new_xs)
            result = x_spl(x)

        elif ndim == 1:
            '''
            Given a 1-D array of points on the xy-plane. 
            '''
            ny     = len(y)            
            new_xs = np.array([y0_spls[i](y)                 for i in range(nx0)]) # new_xs.shape = (nx0,ny)       
            x_spls = [extrapolated_spline_1D(x0,new_xs[:,i]) for i in range(ny)]
            result = np.array([x_spls[i](x[i])               for i in range(ny)])

        else:
            '''
            Given a multiple dimensional array of points on the xy-plane.  
            '''
            x_flatten = x.flatten()
            y_flatten = y.flatten()     
            ny = len(y_flatten)       
            new_xs = np.array([y0_spls[i](y_flatten)         for i in range(nx0)])         
            x_spls = [extrapolated_spline_1D(x0,new_xs[:,i]) for i in range(ny)]
            result = np.array([x_spls[i](x_flatten[i])       for i in range(ny)]).reshape(y.shape)
        return result      
    return f

Answer 1

I've done a similar work called GlobalSpline2D here, and it works perfectly under either liner, cubic, or quintic splines.

Basically it inherits interp2d, and promoting the usage to 2D extrapolation by InterpolatedUnivariateSpline. Both of them are scipy internal functions.

Its usage should be referred to the document as well as the call method of interp2d.

Answer 2

I think I've come up with an answer myself, which utilizes scipy.interpolate.RectBivariateSpline() and is over 10 times faster than my old one.

Here is the function extrapolated_spline_2D_new.

def extrapolated_spline_2D_new(x0,y0,z2d0):
    '''    
    x0,y0 : array_like，1-D arrays of coordinates in strictly ascending order. 
    z2d0  : array_like，2-D array of data with shape (x.size,y.size).
    ''' 
    assert z2d0.shape == (x0.shape[0],y0.shape[0])

    spline = scipy.interpolate.RectBivariateSpline(x0,y0,z2d0,kx=3,ky=3)
    '''
    scipy.interpolate.RectBivariateSpline
    x,y : array_like, 1-D arrays of coordinates in strictly ascending order.
    z   : array_like, 2-D array of data with shape (x.size,y.size).
    '''  
    def f(x,y,spline=spline):
        '''
        x and y have the same shape with the output. 
        '''
        x = np.array(x,dtype='f4')
        y = np.array(y,dtype='f4') 
        assert x.shape == y.shape 
        ndim = x.ndim   
        # We want the output to have the same dimension as the input, 
        # and when ndim == 0 or 1, spline(x,y) is always 2D. 
        if   ndim == 0: result = spline(x,y)[0][0]
        elif ndim == 1: 
            result = np.array([spline(x[i],y[i])[0][0] for i in range(len(x))])
        else:           
            result = np.array([spline(x.flatten()[i],y.flatten()[i])[0][0] for i in range(len(x.flatten()))]).reshape(x.shape)         
        return result
    return f

Note: In the above version, I calculate the value one by one instead of using the codes beneath.

def f(x,y,spline=spline):
    '''
    x and y have the same shape with the output. 
    '''
    x = np.array(x,dtype='f4')
    y = np.array(y,dtype='f4') 
    assert x.shape == y.shape 
    ndim = x.ndim
    if   ndim == 0: result = spline(x,y)[0][0]
    elif ndim == 1: 
         result = spline(x,y).diagonal()
    else:           
         result = spline(x.flatten(),y.flatten()).diagonal().reshape(x.shape)       
    return result

Because when I tried to do the calculation with the codes beneath, it sometimes give the error message as:

<ipython-input-65-33285fd2319d> in f(x, y, spline)
 29         if   ndim == 0: result = spline(x,y)[0][0]
 30         elif ndim == 1:
---> 31             result = spline(x,y).diagonal()
 32         else:
 33             result = spline(x.flatten(),y.flatten()).diagonal().reshape(x.shape)

/usr/local/lib/python2.7/site-packages/scipy/interpolate/fitpack2.pyc in __call__(self, x, y, mth, dx, dy, grid)
826                 z,ier = dfitpack.bispev(tx,ty,c,kx,ky,x,y)
827                 if not ier == 0:
--> 828                     raise ValueError("Error code returned by bispev: %s" % ier)
829         else:
830             # standard Numpy broadcasting

ValueError: Error code returned by bispev: 10

I don't know what it means.