Half working generated interpolation example

generated_interp.py

from math import floor
from numba import njit

####
####  Working
####

@njit
def native_index_1d(mat, vec):
    return mat[vec[0]]

@njit
def native_index_2d(mat, vec):
    return mat[vec[0], vec[1]]


@njit
def findex_1d(mat, vec):

    x = vec[0]

    N = mat.shape[0]

    p = floor(x)             # such that x ∈ [N-1),(p+1)[]
    p = max(0, min(p, N-2))  # to extrapolate linearly
    q = x - p                    # barycentric coefficient of x
    res = mat[p]*(1-q) + mat[p+1]*q

    return res

@njit
def findex_2d(mat, vec):

    x = vec[0]
    y = vec[1]

    N_0 = mat.shape[0]
    N_1 = mat.shape[1]

    p_0 = floor(x)                 # such that x ∈ [p_0,p_0+1[
    p_0 = max(0, min(p_0, N_0-1))  # to extrapolate linearly
    q_0 = x - p_0                    # barycentric coefficient of x

    p_1 = floor(y)                 # such that x ∈ [p_1,p_1+1[
    p_1 = max(0, min(p_1, N_1-1))  # to extrapolate linearly
    q_1 = y - p_1                  # barycentric coefficient of x

    res = ( mat[p_0,p_1]  *(1-q_0) + mat[p_0+1,p_1]  *q_0 ) * (1-q_1) + \
          ( mat[p_0,p_1+1]*(1-q_0) + mat[p_0+1,p_1+1]*q_0 ) * (q_1)

    return res

from generated import generated
from numba import int64, float64

@generated(nopython=True)
def findex(t_mat, t_vec):
    print(t_vec)
    if t_vec == int64:
        if t_mat.ndim==1:
            print("1d")
            return native_index_1d
        elif t_mat.ndim==2:
            return native_index_2d
            print("2d")
    elif t_mat.ndim==1:
        print("mat")
        return findex_1d
    else:
        return findex_2d  # function of 3 arguments


# example:

import numpy
xvec = numpy.linspace(0,10,11)
yvec = xvec
x2vec = (xvec[None,:]+xvec[:,None]).copy()
y2vec = x2vec

# yvec can be indexed with an integer i s.t. 0<=i<=9
print('yvec[3]: {}'.format(findex(yvec,[3])))

# now, index it with a float
print('yvec[3.5]: {}'.format(findex(yvec,[3.5])))

# it also work in 2 dimensions

# yvec can be indexed with an integer i s.t. 0<=i<=9
print('yvec[3,4]: {}'.format(findex(y2vec,[3,4])))

# now, index it with a float
print('yvec[3.5,4.3]: {}'.format(findex(y2vec,[3.5,4.3])))

xxvec = numpy.linspace(-0.1,10.1,1000)
yyvec = numpy.array( [findex_1d(yvec, [x]) for x in xxvec] )
# is equivalent to:
yyvec = numpy.array( [findex(yvec, [x]) for x in xxvec] )

exit()

####
####  Not working
####
@njit
def findex_1d(mat, x):
    # x is now given as argument as a float
    N = mat.shape[0]

    p = floor(x)             # such that x ∈ [p/(N-1),(p+1)/(N-1)]
    p = max(0, min(p, N-2))  # to extrapolate linearly
    q = x                    # barycentric coefficient of x
    res = mat[p]*(1-q) + mat[p+1]*q

    return res

@njit
def findex_2d(mat, x, y):
    # x,y are now given as arguments as floats

    N_0 = mat.shape[0]
    N_1 = mat.shape[1]

    p_0 = floor(x)                 # such that x ∈ [p/(N-1),(p+1)/(N-1)]
    p_0 = max(0, min(p_0, N_0-1))  # to extrapolate linearly
    q_0 = x - p                    # barycentric coefficient of x

    p_1 = floor(y)                 # such that x ∈ [p/(N-1),(p+1)/(N-1)]
    p_1 = max(0, min(p_1, N_1-1))  # to extrapolate linearly
    q_1 = y - p_1                  # barycentric coefficient of x

    res = ( mat[p_0,p_1]  *(1-q_0) + mat[p_0+1,p_1]  *q_0 ) * (1-q_1) + \
          ( mat[p_0,p_1+1]*(1-q_0) + mat[p_0+1,p_1+1]*q_0 ) * (q_1)

    return res

from generated import generated


# this doesn't work (don't know why)
@generated(nopython=True)
def findex(mat, *args):
    if mat.ndim==1:
        return findex_1d  # function of 2 arguments
    else:
        return findex_2d  # function of 3 arguments

# guvectorize is another matter
# it works very well for findex_1d and findex_2d (modulo the issue with scalar args)
# the generated functions may have different core dimensions:

@generated_guvectorize(nopython=True)
def findex(mat, *args):
    if mat.ndim==1:
        core = '(d1,d2),()->()'  # function of 2 arguments
        return core, findex_1d
    else:
        core = '(d1),(),()->()' # function of 3 arguments
        return core, findex_2d