複素数 numpy ndarray の abs()**2 を計算する最もメモリ効率の良い方法 質問する

複素数 numpy ndarray の abs()**2 を計算する最もメモリ効率の良い方法 質問する

複雑なnumpy ndarrayの絶対値の2乗を計算する最もメモリ効率の良い方法を探しています

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

まさにそれを実行する ufunc が見つかりませんでしたnp.abs()**2

そのサイズとタイプの配列は約 0.5 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 を上書きしたくないので、out として使用することはできません。

ベストアンサー1

感謝numba.vectorizenumba の最近のバージョンでは、このタスク用の numpy ユニバーサル関数を作成するのは非常に簡単です。

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

私のマシンでは、中間配列を作成する純粋な NumPy バージョンと比較して 3 倍の速度向上が見られました。

>>> 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

おすすめ記事