I want to write a function in numba that runs a math operation on 2 arrays, and accommodate for when both arrays don't have the same element count.
So for example: lets say I want a function that adds each element of array a to the elements of array b with these 3 possible scenarios:
1) Both a and b have the same number of items, do c[ii]=a[ii]+b[ii]
2) a has more items than b: do c[ii]=a[ii]+b[ii] until b's upper limit, and complete with c[ii]=a[ii]+b[-1]
3) a has fewer items than b: do c[ii]=a[ii]+b[ii] until a's upper limit, and complete with c[ii]=a[-1]+b[ii]
For this I wrote the code below, which works fine and is fast when dealing with millions of values, but I can clearly see three nearly identical code blocks which feel terribly wasteful. Plus the if/else running in a loop also feels terrible.
from numba import jit, float64, int32
@jit(float64[:](float64[:], float64[:]), nopython=True)
def add(a, b):
# Both shapes are equal: add between a[i] and b[i]
if a.shape[0] == b.shape[0]:
c = np.empty(a.shape)
for i in range(a.shape[0]):
c[i] = a[i] + b[i]
return c
# a has more entries than b: add between a[i] and b[i] until b.shape[0]-1 is reached.
# finish the loop with add between a[i] and b[-1]
elif a.shape[0] > b.shape[0]:
c = np.empty(a.shape)
i_ = b.shape[0]-1 # upper limit of b's shape
for i in range(a.shape[0]):
if i < b.shape[0]:
c[i] = a[i] + b[i]
else:
c[i] = a[i] + b[i_]
return c
# b has more entries than a: add between a[i] and b[i] until a.shape[0]-1 is reached.
# finish the loop with add between a[-1] and b[i]
else:
c = np.empty(b.shape)
i_ = a.shape[0]-1 # upper limit of a's shape
for i in range(b.shape[0]):
if i < a.shape[0]:
c[i] = a[i] + b[i]
else:
c[i] = a[i_] + b[i]
return c
I'm new to numba and jit python code compilation so this may just be "the most efficient way" to do what I want.
But if there is a more elegant way to do this without sacrificing speed I would love to know how.