计算复杂 numpy ndarray 的 abs()**2 的最节省内存的方法

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

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

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

由于该大小和类型的数组占用大约半 GB，因此我正在寻找一种主要节省内存的方法。

我还希望它是可移植的，所以最好是 ufunc 的某种组合。

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

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

它将不必要地计算(**0.5)**2，但应该计算**2到位。总而言之，峰值内存需求只是原始数组大小 + 结果数组大小，它应该是 1.5 * 原始数组大小，因为结果是真实的。

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

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

但如果我没记错的话，这意味着我必须分配内存both实部和虚部计算，因此峰值内存使用量将为 2.0 * 原始数组大小。这arr.real属性还返回一个不连续的数组（但这不太重要）。

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

EDIT 1： 很抱歉没有说清楚，我不想覆盖 arr，所以我不能将其用作 out。

谢谢numba.vectorize http://numba.pydata.org/numba-doc/latest/user/vectorize.html在最新版本的 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