不同大小数组的元素操作

Question

在不对较小数组进行过采样的情况下，对不同大小的数组执行逐元素操作的最快和最 Python 的方法是什么？

例如： 我有一个大数组，A 1000x1000 和一个小数组 B 10x10 我希望 B 中的每个元素都响应数组 B 中的 100x100 元素。不需要任何插值，只需对 B 中的所有 10000 个操作使用相同的元素答：

我可以调整两个数组的大小，使 A 的形状是 B 的倍数。通常它们在所有维度上都是 1:10000 或 1:1000。这些数组表示具有不同分辨率但范围相同的数据样本。

我知道我可以对数组 B 进行过采样，例如通过使用 Kronecker 产品，但保持数组 B 较小会更好，特别是当我的一些数组变得非常大以处理和存储时。我正在使用 xarray 和 dask，但任何 numpy 操作也可以。

我希望这个 sn-p 解释了我想要做什么：

import numpy as np

A = np.random.rand(10,10)
B = np.random.rand(1000,1000)

res = np.shape(B)[0]//np.shape(A)[0]

#I want to add A and B so that each element in A is added to 100x100 elements in B.
#This doesn't work of obvious reasons:
#C = A+B

#This solution sacrifices the resolution of B:
C = A+B[::res,::res]

#These solutions creates an unnecessary large array for the operation(don't they?):
K = np.ones((res,res))
C = np.kron(A, K) + B

C = np.repeat(np.repeat(A,res, axis=0), res, axis=1)+B

我有一种感觉，这个问题一定是以前出现过的，但我找不到任何适用于这种特殊情况的答案。

Answer 1

通过广播，我们可以“复制”一个小数组来处理更大的数组。例如，如果

a = np.random.rand(10)
b = np.random.rand(1000).reshape(10,100)
a[:,None]+b

这里的技巧是将A 的每个元素与B 的(100,100) 块配对。为此，我们需要重塑和转置。

In [3]: A = np.random.rand(10,10)
   ...: B = np.random.rand(1000,1000)
   ...: res = np.shape(B)[0]//np.shape(A)[0]

你的目标：

In [4]: K = np.ones((res,res))
   ...: C = np.kron(A, K) + B
   ...: 
   ...: C = np.repeat(np.repeat(A,res, axis=0), res, axis=1)+B

In [5]: C.shape
Out[5]: (1000, 1000)

将B 分成这些 (100,100) 个块：

In [7]: B1 = B.reshape(10,100,10,100).transpose(0,2,1,3)

In [8]: B1.shape
Out[8]: (10, 10, 100, 100)

现在我们可以通过广播添加

In [9]: C1 = B1 + A[:,:,None,None]

In [10]: C1.shape
Out[10]: (10, 10, 100, 100)

重塑回原来的B 形状：

In [11]: C1=C1.transpose(0,2,1,3).reshape(B.shape)

它们匹配：

In [12]: np.allclose(C,C1)
Out[12]: True

Answer 2

我遇到了类似的问题，最终解决了这个问题：

import numpy as np
import numba as nb
import time

@nb.jit(nb.void(nb.float64[:,:], nb.float64[:,:]), nopython=True, cache=True)
def func(A, B):
    # Assume different resolution along the different axes
    res_row = B.shape[0]//A.shape[0]
    res_col = B.shape[1]//A.shape[1]

    for i in range(A.shape[0]):
        for j in range(A.shape[1]):
            # Example to add, can be changed to any other function
            B[i * res_row : (i+1) * res_row, j * res_col : (j+1) * res_col] += A[i,j]


a = np.random.rand(1000,1000)
b = np.random.rand(10000,10000)

t1 = time.time()
func(a,b)
print("Time taken = {time_taken}".format(time_taken = time.time() - t1))
# with jitting on 8GB, 8 cores MacBook Pro = 0.45
# without jitting = 5.46

虽然，numba 与这个问题无关，但这是我解决问题的方法，使用 JIT 编译可以显着提高性能。