多个 numpy 数组上的多步线性代数运算

Question

我有 3 个 numpy 数组：

import numpy
arr_a = numpy.random.random((300, 300))
arr_b = numpy.random.random((300, 300))
arr_c = numpy.random.random((300, 300))

我想从 3 个数组的组合中创建第 4 个数组 (arr_d)。规则如下：

如果 arr_a 网格单元值 > 0.2 且 arr_b 值 0.6，则用 1 填充 arr_d

如果 arr_a 网格单元值 > 0.3 且 arr_b 值 0.6，则用 2 填充 arr_d

如果 arr_a 网格单元值 > 0.1 且 arr_b 值 0.5，则用 3 填充 arr_d

在所有其他情况下，用 4 填充 arr_d

我可以使用嵌套循环来做到这一点，但这很慢而且不是很pythonic。另外，这是一个测试用例，实际数组的大小为 1000 * 1000，所以我想要一个可扩展的解决方案，最好是可并行的。

Answer 1

使用纯 Python 和 for 循环绝对不是要走的路。您可以在 NumPy 中使用数组操作编写程序，在 C 中有效地执行循环，极大地加快代码速度。但是，这会为您的每个规则实例化一个全新的数组，每个规则的大小都与您的数据相同。相反，您可以使用 Numba 之类的东西，例如使用 Python 的 Anaconda 发行版。使用 Numba，您可以使用循环编写代码，但不会造成时间损失（它将您的代码编译为本机机器指令）。此外，不需要额外的大型数组，使其内存效率比 NumPy 高得多。 Numba 也恰好更快，如下例所示：

import numpy, numba, time

def using_numpy(shape):
    arr_a = numpy.random.random(shape)
    arr_b = numpy.random.random(shape)
    arr_c = numpy.random.random(shape)
    mask1 = numpy.logical_and(numpy.logical_and((arr_a > 0.2), (arr_b < 0.4)), (arr_c > 0.6))
    mask2 = numpy.logical_and(numpy.logical_and((arr_a > 0.3), (arr_b < 0.5)), (arr_c > 0.6))
    mask3 = numpy.logical_and(numpy.logical_and((arr_a > 0.1), (arr_b < 0.2)), (arr_c > 0.5))
    result = numpy.ones(arr_a.shape)*4
    result[mask1] = 1
    result[mask2] = 2
    result[mask3] = 3
    return result

@numba.jit
def using_numba(shape):
    arr_a = numpy.random.random(shape)
    arr_b = numpy.random.random(shape)
    arr_c = numpy.random.random(shape)
    result = numpy.empty(shape)
    for i in range(result.shape[0]):
        for j in range(result.shape[1]):
            if arr_a[i, j] > 0.2 and arr_b[i, j] < 0.4 and arr_c[i, j] > 0.6:
                result[i, j] = 1
            elif arr_a[i, j] > 0.3 and arr_b[i, j] < 0.5 and arr_c[i, j] > 0.6:
                result[i, j] = 2
            elif arr_a[i, j] > 0.1 and arr_b[i, j] < 0.2 and arr_c[i, j] > 0.5:
                result[i, j] = 3
            else:
                result[i, j] = 4
    return result
# Compile the using_numba function
using_numba((0, 0))

t0 = time.time()
result = using_numpy((3000, 3000))
print('NumPy took', time.time() - t0, 'seconds')

t0 = time.time()
result = using_numba((3000, 3000))
print('Numba took', time.time() - t0, 'seconds')

这里我使用了(3000, 3000) 数组。在我的机器上，使用 NumPy 需要 0.47 秒，而使用 Numba 需要 0.29 秒。

Answer 2

一种方法是使用布尔映射

condition_1 = numpy.logical_and(numpy.logical_and((arr_a > 0.2), (arr_b < 0.4)), (arr_c > 0.6))
condition_2 = numpy.logical_and(numpy.logical_and((arr_a > 0.3), (arr_b < 0.5)), (arr_c > 0.6))
condition_3 = numpy.logical_and(numpy.logical_and((arr_a > 0.1), (arr_b < 0.2)), (arr_c > 0.5))
result = numpy.ones((300, 300)) * 4
result[numpy.where(condition_3)] = 3
result[numpy.where(condition_2)] = 2
result[numpy.where(condition_1)] = 1

它避免了嵌套循环，但分配了三个专用数组并进行了很多多余的分配。必须有更优化的方法...