Modo più memoria-efficiente per calcolare abs() ** 2 di NumPy complesso ndarray

Question

sto cercando il modo più efficiente di memoria per calcolare il valore assoluto quadrato di un complesso NumPy ndarray

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

Non ho trovato un ufunc che faccia esattamente lo np.abs()**2.

Come una matrice di tale dimensione e tipo occupa circa mezzo GB, sto cercando un modo principalmente efficiente di memoria.

Vorrei anche che fosse portatile, così idealmente una combinazione di ufuncs.

Finora la mia comprensione è che questo dovrebbe essere di circa il migliore

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

Sarà inutilmente calcolare (**0.5)**2, ma dovrebbe calcolare **2 sul posto. Complessivamente, il requisito di memoria di picco è solo la dimensione dell'array originale + la dimensione dell'array del risultato, che dovrebbe essere 1,5 * la dimensione originale dell'array in quanto il risultato è reale.

Se avessi voluto eliminare l'inutile **2 chiamata avrei dovuto fare qualcosa di simile

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

ma se non mi sbaglio, questo significa che dovrò allocare memoria per sia il il calcolo della parte reale e immaginaria, quindi l'utilizzo della memoria di picco sarebbe pari a 2.0 * della dimensione originale dell'array. Le proprietà arr.real restituiscono anche un array non contiguo (ma ciò è di minore importanza).

C'è qualcosa che mi manca? Ci sono modi migliori per farlo?

EDIT 1: Mi dispiace per non mettendo in chiaro, non voglio sovrascrivere arr, quindi non posso usarlo come fuori.

Answer 1

Grazie a numba.vectorize in versioni recenti di numba, creando una funzione universale numpy per l'attività è molto semplice:

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

Sulla mia macchina, trovo un triplice aumento di velocità rispetto ad una versione pura-NumPy che crea gli array intermedi:

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

arr.real e arr.imag sono solo viste nella matrice complessa. Quindi non viene allocata memoria aggiuntiva.

Answer 3

Se il tuo obiettivo principale è quello di risparmiare memoria, gli ufunc di NumPy accettano un parametro opzionale out che consente di indirizzare l'output a un array di tua scelta. Può essere utile quando si desidera eseguire operazioni sul posto.

Se fate questa piccola modifica al tuo primo metodo, quindi è possibile eseguire l'operazione su arr completamente a posto:

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

---

Un modo contorto che utilizza solo un po 'di memoria extra poteva nel modificare arr in luogo, calcolare la nuova matrice di valori reali e quindi ripristinare arr.

Ciò significa memorizzare le informazioni sui segni (a meno che non si sa che i numeri complessi sono tutte parti reale e immaginaria positivi). è necessario solo un singolo bit per il segno di ogni valore reale o immaginario, quindi questo utilizza 1/16 + 1/16 == 1/8 la memoria di arr (oltre alla nuova gamma di carri si 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 scapito di signbits memorizzazione, arr è invariato e result detiene i valori che vogliamo.

Answer 4

EDIT: questa soluzione ha il doppio del requisito di memoria minimo ed è leggermente più veloce. La discussione nei commenti è buona per riferimento comunque.

Ecco una soluzione più veloce, con il risultato memorizzato in res:

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

dove abbiamo sfruttato la proprietà del abs di un numero complesso, vale a dire abs(z) = sqrt(z*z.conjugate), in modo che abs(z)**2 = z*z.conjugate