Search code examples
numpylinear-interpolation

NumPy: How to calulate piecewise linear interpolant on multiple axes


Given the following ndarray t -

In [26]: t.shape                                                                                     
Out[26]: (3, 3, 2)

In [27]: t                                                                                           
Out[27]: 
array([[[ 0,  1],
        [ 2,  3],
        [ 4,  5]],

       [[ 6,  7],
        [ 8,  9],
        [10, 11]],

       [[12, 13],
        [14, 15],
        [16, 17]]])

this piecewise linear interpolant for the points t[:, 0, 0] can evaluated for [0 , 0.66666667, 1.33333333, 2.] as follows using numpy.interp -

In [38]: x = np.linspace(0, t.shape[0]-1, 4)                                                         

In [39]: x                                                                                           
Out[39]: array([0.        , 0.66666667, 1.33333333, 2.        ])

In [30]: xp = np.arange(t.shape[0])                                                                  

In [31]: xp                                                                                          
Out[31]: array([0, 1, 2])

In [32]: fp = t[:,0,0]                                                                               

In [33]: fp                                                                                          
Out[33]: array([ 0,  6, 12])

In [40]: np.interp(x, xp, fp)                                                                        
Out[40]: array([ 0.,  4.,  8., 12.])

How can all the interpolants be efficiently calculated and returned together for all values of fp -

array([[[ 0,  1],
        [ 2,  3],
        [ 4,  5]],

       [[ 4,  5],
        [ 6,  7],
        [ 8,  9]],

       [[ 8,  9],
        [10, 11],
        [12, 13]],

       [[12, 13],
        [14, 15],
        [16, 17]]])

Solution

  • As the interpolation is 1d with changing y values it must be run for each 1d slice of t. It's probably faster to loop explicitly but neater to loop using np.apply_along_axis

    import numpy as np
    
    t = np.arange( 18 ).reshape(3,3,2)
    
    x = np.linspace( 0, t.shape[0]-1, 4)
    xp = np.arange(t.shape[0])
    
    
    def interfunc( arr ):
        """ Function interpolates a 1d array. """
        return np.interp( x, xp, arr )
    
    np.apply_along_axis( interfunc, 0, t ) # apply function along axis 0
    
    """  Result
    array([[[ 0.,  1.],
            [ 2.,  3.],
            [ 4.,  5.]],
    
           [[ 4.,  5.],
            [ 6.,  7.],
            [ 8.,  9.]],
    
           [[ 8.,  9.],
            [10., 11.],
            [12., 13.]],
    
           [[12., 13.],
            [14., 15.],
            [16., 17.]]]) """
    

    With explicit loops

    result = np.zeros((4,3,2))
    for c in range(t.shape[1]):
        for p in range(t.shape[2]):
           result[:,c,p] = np.interp( x, xp, t[:,c,p])
    

    On my machine the second option runs in half the time.

    Edit to use np.nditer

    As the result and the parameter have different shapes I seem to have to create two np.nditer objects one for the parameter and one for the result. This is my first attempt to use nditer for anything so it could be over complicated.

    def test( t ):                                                              
        ts = t.shape              
        result = np.zeros((ts[0]+1,ts[1],ts[2]))
        param = np.nditer( [t], ['external_loop'], ['readonly'],  order = 'F')
        with np.nditer( [result], ['external_loop'], ['writeonly'],  order = 'F') as res:       
            for p, r in zip( param, res ):
                r[:] = interfunc(p)
        return result
    

    It's slightly slower than the explicit loops and less easy to follow than either of the other solutions.