python - vectorize - La forma más eficiente en la memoria de calcular abs()** 2 de ndarray de número complejo

numpy array

Question

Estoy buscando la forma más eficiente de calcular la memoria del valor cuadrado absoluto de un ndarray de números complejos

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

No he encontrado un ufunc que haría exactamente np.abs()**2 .

Como una matriz de ese tamaño y tipo ocupa aproximadamente la mitad de un GB, estoy buscando una forma de memoria principalmente eficiente.

También me gustaría que fuera portátil, así que idealmente alguna combinación de ufuncs.

Hasta ahora, mi entendimiento es que esto debería ser lo mejor

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

Se calculará innecesariamente (**0.5)**2 , pero se debe calcular **2 en el lugar. En conjunto, el requisito de memoria máxima es solo el tamaño de la matriz original + el tamaño de la matriz del resultado, que debe ser 1.5 * el tamaño de la matriz original, ya que el resultado es real

Si quisiera deshacerme de lo inútil **2 llamada, tendría que hacer algo como esto

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

pero si no me equivoco, esto significa que tendré que asignar memoria para el cálculo de la parte real e imaginaria, por lo que el uso máximo de la memoria sería de 2.0 * tamaño de matriz original. Las propiedades arr.real también devuelven una matriz no contigua (pero eso es menos preocupante).

¿Hay algo que me falta? ¿Hay alguna forma mejor de hacer esto?

EDITAR 1 : lo siento por no haberlo dejado en claro, no quiero sobrescribir arr, así que no puedo usarlo como no.

Answer 1

EDITAR: esta solución tiene el doble del requisito de memoria mínimo, y es ligeramente más rápida. La discusión en los comentarios es buena para referencia sin embargo.

Aquí hay una solución más rápida, con el resultado almacenado en res :

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

donde explotamos la propiedad del abs de un número complejo, es decir, abs(z) = sqrt(z*z.conjugate) , de modo que abs(z)**2 = z*z.conjugate

Answer 2

Gracias a numba.vectorize en versiones recientes de numba, crear una función universal numpy para la tarea es muy fácil:

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

En mi máquina, encuentro una aceleración triple en comparación con una versión de número puro que crea arreglos intermedios:

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

Si su objetivo principal es conservar la memoria, los ufuncs de NumPy toman un parámetro de out opcional que le permite dirigir la salida a una matriz de su elección. Puede ser útil cuando desee realizar operaciones en el lugar.

Si realiza esta pequeña modificación en su primer método, puede realizar la operación completamente en su lugar:

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

Una forma complicada que solo usa un poco de memoria adicional podría ser modificar arr en su lugar, calcular la nueva matriz de valores reales y luego restaurarlos.

Esto significa almacenar información sobre los signos (a menos que sepa que todos sus números complejos tienen partes reales e imaginarias positivas). Solo se necesita un solo bit para el signo de cada valor real o imaginario, por lo que utiliza 1/16 + 1/16 == 1/8 la memoria de arr (además de la nueva matriz de flotadores que crea).

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

A expensas de almacenar los bits de firma, arr se modifica y el result mantiene los valores que queremos.

Answer 4

arr.real y arr.imag son solo vistas en el conjunto complejo. Así que no se asigna memoria adicional.