大多数内存有效的方法来计算复杂numpy ndarray的abs()**2

Question

我正在寻找最节省内存的方法来计算复杂numpy ndarray的绝对平方值

arr = np.empty((250000, 150), dtype='complex128')  # common size

我还没有找到一个完全可以做到的ufunc np.abs()**2.

由于该大小和类型的数组占用大约半GB,我正在寻找一种主要的内存效率方式.

我也希望它是可移植的,所以理想情况下是ufuncs的一些组合.

到目前为止,我的理解是,这应该是最好的

result = np.abs(arr)
result **= 2

它将不必要地计算(**0.5)**2,但应该**2就地计算.总的来说,峰值内存要求只是原始数组大小+结果数组大小,应该是1.5*原始数组大小,因为结果是真实的.

如果我想摆脱无用的**2电话,我必须做这样的事情

result = arr.real**2
result += arr.imag**2

但如果我没有记错,这意味着我将不得不为分配内存都实部和虚部的计算,因此峰值内存使用量是2.0*原始数组的大小.这些arr.real属性还返回一个非连续的数组(但这个问题不太重要).

有什么我想念的吗？有没有更好的方法来做到这一点？

编辑1:我很抱歉没有说清楚,我不想覆盖arr,所以我不能用它作为出来.

Answer 1

感谢numba.vectorize最新版本的numba,为任务创建一个numpy通用函数非常简单:

@numba.vectorize([numba.float64(numba.complex128),numba.float32(numba.complex64)])
def abs2(x):
    return x.real**2 + x.imag**2

在我的机器上,与创建中间数组的纯numpy版本相比,我发现速度提高了三倍:

>>> x = np.random.randn(10000).view('c16')
>>> y = abs2(x)
>>> np.all(y == x.real**2 + x.imag**2)   # exactly equal, being the same operation
True
>>> %timeit np.abs(x)**2
10000 loops, best of 3: 81.4 µs per loop
>>> %timeit x.real**2 + x.imag**2
100000 loops, best of 3: 12.7 µs per loop
>>> %timeit abs2(x)
100000 loops, best of 3: 4.6 µs per loop

Answer 2

编辑：此解决方案的最低内存要求是最低内存要求的两倍，并且速度稍快一些。不过评论里的讨论还是很有参考价值的。

这是一个更快的解决方案，结果存储在res：

import numpy as np
res = arr.conjugate()
np.multiply(arr,res,out=res)

我们利用了复数的绝对值的性质，即abs(z) = sqrt(z*z.conjugate)，因此abs(z)**2 = z*z.conjugate