Realizar operaciones en todos los valores de una matriz numpy, haciendo referencia a i y j

Estoy tratando de mejorar el rendimiento numpy mediante la aplicación de operaciones en una matriz 2d, el problema es que el valor de cada elemento de la matriz depende de la ubicación i, j de ese elemento.

Obviamente, la forma fácil de hacer esto es usar un bucle for nested, pero me preguntaba si podría haber una mejor manera al hacer referencia a np.indices o algo parecido. Aquí está mi código ‘estúpido’:

for J in range(1025): for I in range(1025): PSI[I][J] = A*math.sin((float(I+1)-.5)*DI)*math.sin((float(J+1)-.5)*DJ) P[I][J] = PCF*(math.cos(2.*float(I)*DI)+math.cos(2.*float(J)*DJ))+50000. 

Ya que estás haciendo la multiplicación entre tus dos arreglos, puedes usar la función externa , después de usar un arange para obtener arreglos de tu pecado / cos.

Algo como esto (usa las funciones trigonométricas de numpy, ya que están vectorizadas)

 PSI_i = numpy.sin((arange(1,1026)-0.5)*DI) PSI_j = numpy.sin((arange(1,1026)-0.5)*DJ) PSI = A*outer(PSI_i, PSI_j) P_i = numpy.cos(2.*arange(1,1026)*DI) P_j = numpy.cos(2.*arange(1,1026)*DJ) P = PCF*outer(P_i, P_j) + 50000 

Si su entorno está configurado from numpy import * o from pylab import * , entonces no necesita esos numpy. Prefijos antes de sus funciones trigonométricas. Los mantuve para distinguirlos de los math , lo que no funcionará para este enfoque.

Puede obtener una cuadrícula de los valores de índice con índices:

 I,J=np.indices(PSI.shape) #All constants set to one PSI2=np.sin(I+1-.5)*np.sin(J+1-.5) print PSI-PSI2 # should be zero. 

Hice algunos tiempos con ipython:

 import numpy as np import math A = 1 P = 1 DI = 1 DJ = 1 def a(): PSI=np.zeros((1025,1025)) for J in range(1025): for I in range(1025): PSI[I][J] = A*math.sin((float(I+1)-.5)*DI)*math.sin((float(J+1)-.5)*DJ) %timeit a() def b(): PSI=np.zeros((1025,1025)) for I,J in np.ndindex(*PSI.shape): PSI[I,J] = A*math.sin((float(I+1)-.5)*DI)*math.sin((float(J+1)-.5)*DJ) %timeit b() def c(): I,J=np.indices((1025, 1025)) P2=A*np.sin((I+1-.5)*DI)*np.sin((J+1-.5)*DJ) %timeit c() def d(): PSI_i = np.sin((np.arange(1,1026)-0.5)*DI) PSI_j = np.sin((np.arange(1,1026)-0.5)*DJ) PSI = A*np.outer(PSI_i, PSI_j) %timeit d() 

El resultado no es para nada sorprendente en mi máquina:

 1 loops, best of 3: 1.75 s per loop 1 loops, best of 3: 3.51 s per loop 10 loops, best of 3: 77.1 ms per loop 100 loops, best of 3: 7.16 ms per loop 

Pruebe la función ndenumerate de numpy, que devuelve el valor y los índices:

 >>> a array([[5, 5, 5], [1, 2, 3]]) >>> for index, value in numpy.ndenumerate(a): ... print index, value (0, 0) 5 (0, 1) 5 (0, 2) 5 (1, 0) 1 (1, 1) 2 (1, 2) 3