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

Question

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

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

我还没有找到完全可以 np.abs()**2.

的 ufunc

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

我也希望它是可移植的，所以最好是 ufunc 的一些组合。

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

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 属性也是 return 一个不连续的数组（但这不太重要）。

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

编辑 1：不好意思没说清楚，我不想覆盖arr，所以不能当out用。

Answer 1

arr.real 和 arr.imag 只是复杂数组的视图。所以没有分配额外的内存。

Answer 2

如果您的主要目标是节省内存，NumPy 的 ufunc 会采用一个可选的 out 参数，让您可以将输出定向到您选择的数组。当您想就地执行操作时，它会很有用。

如果您对第一个方法稍作修改，那么您就可以完全原地对 arr 执行操作：

np.abs(arr, out=arr)
arr **= 2

一种只使用 一点点 额外内存的复杂方法可能是就地修改 arr，计算新的实数值数组，然后恢复 arr.

这意味着存储有关符号的信息（除非您知道您的复数都具有正实部和虚部）。每个实数或虚数的符号只需要一位，因此这使用 1/16 + 1/16 == 1/8 arr 的内存（除了您创建的新浮点数数组）。

>>> signs_real = np.signbit(arr.real) # store information about the signs
>>> signs_imag = np.signbit(arr.imag)
>>> arr.real **= 2 # square the real and imaginary values
>>> arr.imag **= 2
>>> result = arr.real + arr.imag
>>> arr.real **= 0.5 # positive square roots of real and imaginary values
>>> arr.imag **= 0.5
>>> arr.real[signs_real] *= -1 # restore the signs of the real and imagary values
>>> arr.imag[signs_imag] *= -1

以存储符号位为代价，arr不变，result保留我们想要的值。

Answer 3

编辑：此解决方案的最低内存要求是其两倍，而且速度稍快。不过评论里的讨论还是可以参考的

这是一个更快的解决方案，结果存储在 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

Answer 4

感谢 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 5

如果你不想sqrt（什么东西应该比乘法重很多），那就不要abs.

如果不想双倍内存，那就没有real**2 + imag**2

那么你可以试试这个（使用索引技巧）

N0 = 23
np0 = (np.random.randn(N0) + 1j*np.random.randn(N0)).astype(np.complex128)
ret_ = np.abs(np0)**2
tmp0 = np0.view(np.float64)
ret0 = np.matmul(tmp0.reshape(N0,1,2), tmp0.reshape(N0,2,1)).reshape(N0)
assert np.abs(ret_-ret0).max()<1e-7

无论如何，我更喜欢numba解决方案

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

Most memory-efficient way to compute abs()**2 of complex numpy ndarray

python

numpy

complex-numbers

memory-efficient

numpy-ufunc