假设您有一个形状为（a，b，c）的numpy数组和一个形状为（a，b，c，d）的布尔掩码。 我想将掩码应用于在最后一个轴上迭代的数组，将掩码的数组沿前三个轴求和，并获得长度/形状（d，）的列表（或数组）。 我试图通过列表理解来做到这一点：
Result = [np.sum(Array[Mask[:,:,:,i]], axis=(0,1,2)) for i in range(d)]
它可以工作，但是看起来不是很pythonic，也有点慢。 我也尝试过类似的
Array = Array[:,:,:,np.newaxis]
Result = np.sum(Array[Mask], axis=(0,1,2))
但是这当然不起作用，因为Mask沿最后一个轴d的尺寸大于Array的最后一个轴的尺寸1。
另外，请考虑到每个轴的维数可能为100或200，因此使用np.repeat沿着新的最后一个轴重复Array d次实际上会占用大量内存，我想避免这种情况。
列表推导还有其他更快，更Python的替代方法吗？
答案 0 :(得分：1)
怎么样
value由于无论如何都要对前三个轴求和，因此也可以将它们合并，然后很容易看出该运算可以写成矩阵向量乘积
示例：
code答案 1 :(得分：1)
将N维数组广播为匹配的（N + 1）维数组的最直接方法是使用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
但是，正如@hpaulj已经指出的那样，您不能使用mask来切片b_arr，而不会失去尺寸。
鉴于您只想将元素加在一起，然后将零加起来“不会有伤害”，您可以简单地逐个元素地将数组和掩码相乘，以保持正确的尺寸，但元素{{1}掩码中的}与相应数组元素的后续False无关：
sum或者，因为result = np.sum(b_arr * mask, axis=tuple(range(mask.ndim - 1)))
将自动进行广播：
*首先无需使用np.broadcast_to()（但是您仍然需要匹配尺寸数，即使用result = np.sum(arr[..., None] * mask, axis=tuple(range(mask.ndim - 1)))
而不仅仅是arr[..., None]）。
与@PaulPanzer already pointed out一样，由于您想对除一个维度之外的所有维度进行汇总，因此可以使用np.matmul()/@进一步简化：
arr对于涉及求和的更高级操作，请查看np.einsum()。
广播的问题是它将在评估表达式时创建临时数组。
有了您似乎要处理的数字，我碰到result2 = arr.ravel() @ mask.reshape(-1, mask.shape[-1])
print(np.all(result == result2))
# True
时就无法使用广播数组，但是按时间逐元素相乘可能仍然是比您最初建议的更好的方法
或者，如果您追求速度，则可以在Cython或Numba中进行显式循环，从而在较低的级别上进行此操作。
下面您可以找到几个基于Numba的解决方案（适用于MemoryError版数据）：
ravel()：不使用任何临时数组_vector_matrix_product()：如上所述，但使用并行执行_vector_matrix_product_mp()：使用_vector_matrix_product_sum()和并行执行np.sum()与所有基于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]
的基于解决方案的解决方案相比，计时得以改善：
list正如预期的那样，JIT加速版本是最快的，并且对代码强制执行并行性不会提高速度。 还要注意，使用元素逐次乘法的方法比切片更快（对于这些基准，速度大约是切片的两倍）。
遵循@ max9111的建议，首先按行循环，然后按cols循环，使最耗时的循环对连续数据运行，从而显着提高了速度。
没有这个技巧，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)
和_vector_matrix_product_sum()的运行速度基本相同。