将N-dim数组广播到（N + 1）-dim数组并在除1个dim之外的所有值上求和

Question

Assume you have a numpy array with shape (a,b,c) and a boolean mask of shape (a,b,c,d). I would like to apply the mask to the array iterating over the last axis, sum the masked array along the first three axes, and obtain a list (or an array) of length/shape (d,). I tried to do this with a list comprehension:

Result = [np.sum(Array[Mask[:,:,:,i]], axis=(0,1,2)) for i in range(d)]

It works, but it does not look very pythonic and it is a bit slow as well. I also tried something like

Array = Array[:,:,:,np.newaxis]
Result = np.sum(Array[Mask], axis=(0,1,2))

but of course this doesn't work, since the dimension of the Mask along the last axis, d, is larger than the dimension of the last axis of the Array, 1. Also, consider that each axis could have dimension of order 100 or 200, so repeating the Array d times along a new last axis using np.repeat would be really memory intensive, and I would like to avoid this. Are there any other faster and more pythonic alternatives to the list comprehension?

Answer 1

How about

Array.reshape(-1)@Mask.reshape(-1,d)

Since you are summing over the first three axes anyway you may as well merge them after which it is easy to see that the operation can be written as matrix-vector product

Example:

a,b,c,d = 4,5,6,7
Mask = np.random.randint(0,2,(a,b,c,d),bool)
Array = np.random.randint(0,10,(a,b,c))
[np.sum(Array[Mask[:,:,:,i]]) for i in range(d)]
# [310, 237, 253, 261, 229, 268, 184]    
Array.reshape(-1)@Mask.reshape(-1,d)
# array([310, 237, 253, 261, 229, 268, 184])

Answer 2

The most straightforward way of broadcasting a N-dimensional array to a matching (N+1)-dimensional array is to use np.broadcast_to() :

import numpy as np


arr = np.random.randint(0, 100, (2, 3))
mask = np.random.randint(0, 2, (2, 3, 4), dtype=bool)
b_arr = np.broadcast_to(arr[..., None], mask.shape)
print(mask.shape == b_arr.shape)
# True

However, as @hpaulj already pointed out, you cannot use mask for slicing b_arr without loosing the dimensions.

---

Given that you want to just sum the elements together and summing zeroes "does not hurt", you could simply multiply element-wise your array and your mask so as to keep the correct dimension but the elements that are False in the mask are irrelevant for the subsequent sum of the corresponding array elements:

result = np.sum(b_arr * mask, axis=tuple(range(mask.ndim - 1)))

or, since * will do the broadcasting automatically:

result = np.sum(arr[..., None] * mask, axis=tuple(range(mask.ndim - 1)))

without the need to use np.broadcast_to() in the first place (but you still need to match the number of dimension, ie using arr[..., None] and not just arr ).

---

As @PaulPanzer already pointed out , since you want to sum up over all but one dimensions, this can be further simplified using np.matmul() / @ :

result2 = arr.ravel() @ mask.reshape(-1, mask.shape[-1])
print(np.all(result == result2))
# True

For fancier operations involving the summation, please have a look at np.einsum() .

---

EDIT

The catch with broadcasting is that it will create temporary arrays during the evaluation of your expressions.

With the number you seems to be dealing with, I simply cannot use the broadcasted arrays as I run into MemoryError , but time-wise the element-wise multiplication may still be a better approach than what you originally proposed.

Alternatively, if you are after speed, you could do this at a somewhat lower level with explicit looping in Cython or Numba.

Below you can find a couple of Numba-based solutions (working on ravel() -ed data):

• _vector_matrix_product() : does not use any temporary array
• _vector_matrix_product_mp() : some as above but using parallel execution
• _vector_matrix_product_sum() : uses np.sum() and parallel execution

import numpy as np
import numba as nb


@nb.jit(nopython=True)
def _vector_matrix_product(
        vect_arr,
        mat_arr,
        result_arr):
    rows, cols = mat_arr.shape
    if vect_arr.shape == result_arr.shape:
        for i in range(rows):
            for j in range(cols):
                result_arr[i] += vect_arr[j] * mat_arr[i, j]
    else:
        for i in range(rows):
            for j in range(cols):            
                result_arr[j] += vect_arr[i] * mat_arr[i, j]


@nb.jit(nopython=True, parallel=True)
def _vector_matrix_product_mp(
        vect_arr,
        mat_arr,
        result_arr):
    rows, cols = mat_arr.shape
    if vect_arr.shape == result_arr.shape:
        for i in nb.prange(rows):
            for j in nb.prange(cols):
                result_arr[i] += vect_arr[j] * mat_arr[i, j]
    else:
        for i in nb.prange(rows):
            for j in nb.prange(cols):        
                result_arr[j] += vect_arr[i] * mat_arr[i, j]


@nb.jit(nopython=True, parallel=True)
def _vector_matrix_product_sum(
        vect_arr,
        mat_arr,
        result_arr):
    rows, cols = mat_arr.shape
    if vect_arr.shape == result_arr.shape:
        for i in nb.prange(rows):
            result_arr[i] = np.sum(vect_arr * mat_arr[i, :])
    else:
        for j in nb.prange(cols):
            result_arr[j] = np.sum(vect_arr * mat_arr[:, j])


def vector_matrix_product(
        vect_arr,
        mat_arr,
        swap=False,
        dtype=None,
        mode=None):
    rows, cols = mat_arr.shape
    if not dtype:
        dtype = (vect_arr[0] * mat_arr[0, 0]).dtype
    if not swap:
        result_arr = np.zeros(cols, dtype=dtype)
    else:
        result_arr = np.zeros(rows, dtype=dtype)
    if mode == 'sum':
        _vector_matrix_product_sum(vect_arr, mat_arr, result_arr)
    elif mode == 'mp':
        _vector_matrix_product_mp(vect_arr, mat_arr, result_arr)
    else:
        _vector_matrix_product(vect_arr, mat_arr, result_arr)
    return result_arr


np.random.seed(0)
arr = np.random.randint(0, 100, (2, 3, 4))
mask = np.random.randint(0, 2, (2, 3, 4, 5), dtype=bool)
target = arr.ravel() @ mask.reshape(-1, mask.shape[-1])
print(target)
# [820 723 861 486 408]
result1 = vector_matrix_product(arr.ravel(), mask.reshape(-1, mask.shape[-1]))
print(result1)
# [820 723 861 486 408]
result2 = vector_matrix_product(arr.ravel(), mask.reshape(-1, mask.shape[-1]), mode='mp')
print(result2)
# [820 723 861 486 408]
result3 = vector_matrix_product(arr.ravel(), mask.reshape(-1, mask.shape[-1]), mode='sum')
print(result3)
# [820 723 861 486 408]

with improved timing over any list -comprehension-based solutions:

arr = np.random.randint(0, 100, (256, 256, 256))
mask = np.random.randint(0, 2, (256, 256, 256, 128), dtype=bool)


%timeit np.sum(arr[..., None] * mask, axis=tuple(range(mask.ndim - 1)))
# MemoryError

%timeit arr.ravel() @ mask.reshape(-1, mask.shape[-1])
# MemoryError

%timeit np.array([np.sum(arr * mask[..., i], axis=tuple(range(mask.ndim - 1))) for i in range(mask.shape[-1])])
# 24.1 s ± 105 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit np.array([np.sum(arr[mask[..., i]]) for i in range(mask.shape[-1])])
# 46 s ± 119 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit vector_matrix_product(arr.ravel(), mask.reshape(-1, mask.shape[-1]))
# 408 ms ± 2.12 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit vector_matrix_product(arr.ravel(), mask.reshape(-1, mask.shape[-1]), mode='mp')
# 1.63 s ± 3.58 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit vector_matrix_product(arr.ravel(), mask.reshape(-1, mask.shape[-1]), mode='sum')
# 7.17 s ± 258 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

As expected, the JIT accelerated version is the fastest, and enforcing parallelism on the code does not result in improved speed-ups. Note also that the approach with element-wise multiplication is faster than slicing (approx. twice as fast for these benchmarks).

---

EDIT 2

Following @max9111 suggestion, looping first by rows and then by cols cause the most time-consuming loop to run on contiguous data, resulting in significant speed-up. Without this trick, _vector_matrix_product_sum() and _vector_matrix_product_mp() would run at essentially the same speed.