Selecting from large array based on broadcasted smaller array - numba?

FooBar 0
Viewed 31

I have a data matrix of dimensions (a, b, c), and a selecting matrix of that is broadcaste-able (let's say, dimensions (a, c)).

import numpy as np
dim1, dim2, dim3 = 4, 5, 6
d1 = np.ones((dim1,))
d2 = np.ones((dim2,))
d3 = np.arange(dim3)
f1 = np.arange(dim1)
f2 = np.arange(dim3) + 100
D1, D2, D3 = np.meshgrid(d1, d2, d3, indexing='ij')
data = D1 + D2 + D3
F1, F2 = np.meshgrid(f1, f2, indexing='ij')
select = F1 + F2 > 101

I want a long array of elements from data, where data[i, j, k] is selected if

  • select[i, k] == True

In python, I can comfortably broadcast and select:

data[np.broadcast_to(select[:, None, :], (dim1, dim2, dim3))]

which, in this case, delivers a 105-long element. In my usecase, this is very slow, and unnecessarily so -- I shouldn't have to expand my select matrix to do this.

I was thinking about doing this in numba, but it wasn't clear to me how to do this: I can simply loop through all elements:

@numba.jit(nopython=True)
def select_large_array_using_small_array(large_array, small_array):
    shape1, shape2, shape3,  = large_array.shape
    large_array_flat = large_array.reshape(-1,)
    select_scaled = []
    for idx_1 in range(shape1):
        for idx_2 in range(shape2):
            for idx_3 in range(shape3):
                if small_array[idx_1, idx_3]:
                    index = idx_3 + idx_2 * shape3 + idx_1 * shape2 * shape3
                    select_scaled.append(large_array_flat[index])
    return np.array(select_scaled)

, but I have several issues with this. First, select_scaled grows dynamically, not great. Second, I would expect that this does not really use multiple cores well.

Is there any better way to attack this problem?

I'm looking for solutions that are appropriate for when data has millions of elements.

array-broadcasting numpy numba python
1 Answers

The usual solution to speed up this kind of problem is to find the size of the output array first and then fill it (possibly in parallel when it is possible and this worth it). Here is an example (untested):

@numba.jit(nopython=True)
def select_large_array_using_small_array(large_array, small_array):
    shape1, shape2, shape3,  = large_array.shape
    large_array_flat = large_array.reshape(-1,)
    size = 0
    for idx_1 in range(shape1):
        for idx_3 in range(shape3):
            size += small_array[idx_1, idx_3]
    size *= shape2
    select_scaled = np.empty_like(size, large_array_flat.dtype)
    cur = 0
    for idx_1 in range(shape1):
        for idx_2 in range(shape2):
            for idx_3 in range(shape3):
                if small_array[idx_1, idx_3]:
                    index = idx_3 + idx_2 * shape3 + idx_1 * shape2 * shape3
                    select_scaled[cur] = large_array_flat[index]
                    cur += 1
    return select_scaled

If small_array is predictable by the processor (ie. clearly not random), then the computation will be memory bound and so using multiple cores will barely speed up the computation (and make it significantly more complex). If you still want to try a parallel implementation, then there are few solutions.

If the order of the output items does not matter, then you can compute the idx_2-based loop in parallel and each thread can work on a subset of the array. Something like this:

    for idx_2 in nb.prange(shape2):
        local_cur = idx_2 * (select_scaled.size // shape2)
        for idx_1 in range(shape1):
            for idx_3 in range(shape3):
                if small_array[idx_1, idx_3]:
                    index = idx_3 + idx_2 * shape3 + idx_1 * shape2 * shape3
                    select_scaled[local_cur] = large_array_flat[index]
                    local_cur += 1

If the order matters, then this is significantly more complex since the location of the final items is dependent of the selection of the previous items that are meant to be computed in parallel. A general solution to this problem is to use a parallel scan algorithm so to find the location of each item. That being said, this method is pretty expensive and is only useful when the number of core is sufficiently big while your problem should not scale. An alternative solution is to pre-compute the sum of small_array along the rows and then compute the offsets based on the sum using a cumulative sum so to finally compute multiple segments of select_scaled in parallel.

Jérôme Richard0
Related