Question

I am trying to improve numpy performance by applying operations on a 2d array, the problem is that the value at each element in the array depends on the i,j location of that element.

Obviously the easy way to do this is to use a nested for-loop, but I was wondering if there might be a better way by referencing np.indices or something along those lines? Here is my ‘stupid’ code:

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.

Answer 1

Since you’re doing multiplication among your two arrays, you can use the outer function, after using arange to get arrays of your sin/cos.

Something like this (use numpy’s trig functions, since they’re vectorized)

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

If your environment is set up using from numpy import * or from pylab import *, then you don’t need those numpy. prefixes before your trig functions. I kept them in to distinguish them from the math ones, which won’t work for this approach.