不同大小数组的逐元素操作

Question

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

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

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

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

我希望这段代码解释了我想做的事情：

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

我遇到了类似的问题，最后是这样解决的：

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 编译可以显着提高性能。

Answer 2

通过广播，我们可以 'replicate' 一个小数组与一个大数组一起工作。例如，如果

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

不同大小数组的逐元素操作

Element-wise operations of arrays of different size

python

arrays

numpy

dask

python-xarray