如何在Python中加速多个数组的计算

Question

假设我们有三个同维度的数组a、b、c，需要基于每列数据进行计算。一个例子如下：

import numpy as np 

col = 10
row = 1000000
a = np.random.normal(size=(row, col))
b = np.random.normal(size=(row, col))
c = np.random.normal(size=(row, col))

def my_func(a, b, c):
    if a[0] + b[0] + b[-1] > c[0]:
        return a * b * c
    else:
        return a * (b[1] + b[-1]) + c[-1]

for i in range(num):
    my_func(a[:, i], b[:, i], c[:, i])

首先考虑使用numpy，然而由于计算不是同质的并且涉及多个数组，numpy.apply_along_axes不能使用。
在这种情况下，有没有什么建议可以加速python中的计算，而不是使用c？

- exteral

---

num是什么？它是col吗？ - Mad Physicist

你是否担心只有10次迭代？这只是一个玩具示例吗？如果是后者，我会交换维度以使紧迫性更加明显。 - Mad Physicist

apply_along_axis是一个速度工具。当处理3D（或更高维度）数据且您的fn仅操作1D时，它只是一种方便的方法。 - hpaulj

Answer 1

当使用numpy时，您需要考虑将计算向量化以加快my_func的速度。在这种情况下，您可以尝试以下内容：
当使用numpy时，您需要考虑将计算向量化以加快代码速度。在本例中，您可以尝试类似于以下方式：

cond = np.broadcast_to(a[0] + b[0] + b[-1] > c[0], a.shape)
result = np.where(cond, a * b * c, a * (b[1] + b[-1]) + c[-1])

在这里，计算像a [0] + b [0]在两个数组上执行，而不是逐个元素执行。最初的回答。

- GZ0

Answer 2

也许可以使用Numba。它非常简单，只需在函数上添加Python装饰器@jit，可能存在兼容性问题，但通常只要是简单的Python或Numpy代码就可以正常工作。

- Gustavo

Answer 3

在你的问题中，你问如何加速Python的计算，"而不是使用C"。如果只是C，但你一般来说对其他环境开放，这可能对你有用。

"如果你希望你的代码运行得更快，你应该使用PyPy。" —— Python创始人Guido van Rossum https://www.pypy.org/index.html 你将需要安装他们的解释器，但你不需要重写任何东西。PyPy也支持numpy。

- JSON Brody

Answer 4

你可以使用numpy加速循环，即使在正常情况下，Python的for循环通常不会因为十次迭代而变得过于昂贵。

第一次优化是计算my_func中的两个分支，然后使用类似np.where的方法来为您进行选择。这种方法在空间上有些浪费，这意味着它还会浪费时间在不必要的分配上，更不用说为所有元素计算两个结果了。

更好的方法是制作一个掩码：

def my_func_2D(a, b, c):
    mask = (a[0, :] + b[0, :] + b[-1, :] > c[0, :])
    result = np.empty_like(a)
    result[:, mask] = a[:, mask] * b[:, mask] * c[:, mask]
    result[:, ~mask] = a[:, ~mask] * (b[1, ~mask] + b[-1, ~mask]) + c[-1, ~mask]    
    return result

这只执行必要的计算，但将中间计算结果存储在可能不必要的缓冲区中。

更好的方法是使用掩码数组。它们可以让您原地执行所需的计算。使用np.broadcast_to, 您甚至可以避免制作全尺寸的掩码：

def my_func_2D(a, b, c):
    mask = (a[0, :] + b[0, :] + b[-1, :] > c[0, :])
    result = ma.masked_array(a.copy(), mask=np.broadcast_to(mask, a.shape))
    bm = ma.masked_array(b, np.broadcast_to(mask, b.shape))
    cm = ma.masked_array(c, np.broadcast_to(mask, c.shape))
    result *= bm[1, :] + bm[-1, :]
    result += cm[-1, :]
    np.logical_not(mask, out=mask)
    result *= b
    result *= c
    return result.data

我们创建的掩码数组（除结果外）都是廉价视图。结果在原地计算，掩码元素不受影响。我们创建的唯一临时缓冲区是mask和bm[1, :] + bm[-1, :]，它们沿着单个维度，因此相对较便宜。

- Mad Physicist

---

第一个函数中有几个错别字：a.empty_like 和 c[c, mask]。在第二个函数中，使用 np 1.17 时，result掩码是只读的。 - hpaulj

我还没有让你的掩码数组起作用，但是@GZD的答案比OP的更快，而你的第一个答案则较慢。正如你所评论的那样，在numpy中进行10次迭代并不是一个缓慢的过程。 - hpaulj

@hpaulj。感谢您的指正和基准测试。我现在正在移动设备上胡乱瞎搞。 - Mad Physicist

@hpaulj。我想修改基础数组无论掩码是否只读都不应该是问题。我会在有机会时进行调查。也许需要在某个地方执行mask.flags['WRITEABLE'] = True。 - Mad Physicist

在我之前的实验中，当我尝试将Numpy内置的掩码数组应用于类似于这样的SO问题时，numpy.ma模块中的掩码数组的性能明显比常规的Numpy数组差很多。 - GZ0