Skip to content

Instantly share code, notes, and snippets.

@albop
Last active Jan 21, 2016
Embed
What would you like to do?
Half working generated interpolation example
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment