Elegant way to complete a parallel operation on two arrays of unequal lengths

Viewed 85

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.

2 Answers
Related