Самый эффективный для памяти способ вычисления abs() ** 2 сложного numpy ndarray

Я ищу наиболее эффективный для памяти способ вычисления абсолютного квадрата значения сложного numpy ndarray

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

Я не нашел ufunc, который сделал бы точно np.abs()**2.

Поскольку массив такого размера и типа занимает около половины ГБ, я ищу в первую очередь способ, способствующий экономии памяти.

Мне также хотелось бы, чтобы он был переносимым, поэтому идеально сочетается с 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, поэтому я не могу использовать его как вне.

Ответ 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

На моей машине я нахожу трехкратное ускорение по сравнению с версией pure-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

Ответ 2

arr.real и arr.imag - это только виды в сложный массив. Таким образом, не выделяется дополнительная память.

Ответ 3

Если основной целью является сохранение памяти, NumPy ufuncs использует необязательный параметр 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 содержит значения, которые мы хотим.

Ответ 4

EDIT: это решение имеет в два раза минимальное требование к памяти, и оно просто немного быстрее. Однако обсуждение в комментариях полезно для справки.

Здесь более быстрое решение, с результатом, хранящимся в res:

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

где мы использовали свойство abs комплексного числа, т.е. abs(z) = sqrt(z*z.conjugate), так что abs(z)**2 = z*z.conjugate