Python optimization problem?

Question

Alright, i had this homework recently (don't worry, i've already done it, but in c++) but I got curious how i could do it in python. The problem is about 2 light sources that emit light. I won't get into details tho.

Here's the code (that I've managed to optimize a bit in the latter part):

import math, array
import numpy as np
from PIL import Image

size = (800,800)
width, height = size

s1x = width * 1./8
s1y = height * 1./8
s2x = width * 7./8
s2y = height * 7./8

r,g,b = (255,255,255)
arr = np.zeros((width,height,3))
hy = math.hypot
print 'computing distances (%s by %s)'%size,
for i in xrange(width):
    if i%(width/10)==0:
        print i,    
    if i%20==0:
        print '.',
    for j in xrange(height):
        d1 = hy(i-s1x,j-s1y)
        d2 = hy(i-s2x,j-s2y)
        arr[i][j] = abs(d1-d2)
print ''

arr2 = np.zeros((width,height,3),dtype="uint8")        
for ld in [200,116,100,84,68,52,36,20,8,4,2]:
    print 'now computing image for ld = '+str(ld)
    arr2 *= 0
    arr2 += abs(arr%ld-ld/2)*(r,g,b)/(ld/2)
    print 'saving image...'
    ar2img = Image.fromarray(arr2)
    ar2img.save('ld'+str(ld).rjust(4,'0')+'.png')
    print 'saved as ld'+str(ld).rjust(4,'0')+'.png'

I have managed to optimize most of it, but there's still a huge performance gap in the part with the 2 for-s, and I can't seem to think of a way to bypass that using common array operations... I'm open to suggestions :D

Edit: In response to Vlad's suggestion, I'll post the problem's details: There are 2 light sources, each emitting light as a sinusoidal wave: E1 = E0*sin(omega1*time+phi01) E2 = E0*sin(omega2*time+phi02) we consider omega1=omega2=omega=2*PI/T and phi01=phi02=phi0 for simplicity by considering x1 to be the distance from the first source of a point on the plane, the intensity of the light in that point is Ep1 = E0*sin(omega*time - 2*PI*x1/lambda + phi0) where lambda = speed of light * T (period of oscillation) Considering both light sources on the plane, the formula becomes Ep = 2*E0*cos(PI*(x2-x1)/lambda)sin(omegatime - PI*(x2-x1)/lambda + phi0) and from that we could make out that the intensity of the light is maximum when (x2-x1)/lambda = (2*k) * PI/2 and minimum when (x2-x1)/lambda = (2*k+1) * PI/2 and varies in between, where k is an integer

For a given moment of time, given the coordinates of the light sources, and for a known lambda and E0, we had to make a program to draw how the light looks IMHO i think i optimized the problem as much as it could be done...

Answer 1

Interference patterns are fun, aren't they?

So, first off this is going to be minor because running this program as-is on my laptop takes a mere twelve and a half seconds.

But let's see what can be done about doing the first bit through numpy array operations, shall we? We have basically that you want:

arr[i][j] = abs(hypot(i-s1x,j-s1y) - hypot(i-s2x,j-s2y))

For all i and j.

So, since numpy has a hypot function that works on numpy arrays, let's use that. Our first challenge is to get an array of the right size with every element equal to i and another with every element equal to j. But this isn't too hard; in fact, an answer below points my at the wonderful numpy.mgrid which I didn't know about before that does just this:

array_i,array_j = np.mgrid[0:width,0:height]

There is the slight matter of making your (width, height)-sized array into (width,height,3) to be compatible with your image-generation statements, but that's pretty easy to do:

arr = (arr * np.ones((3,1,1))).transpose(1,2,0)

Then we plug this into your program, and let things be done by array operations:

import math, array
import numpy as np
from PIL import Image

size = (800,800)
width, height = size

s1x = width * 1./8
s1y = height * 1./8
s2x = width * 7./8
s2y = height * 7./8

r,g,b = (255,255,255)

array_i,array_j = np.mgrid[0:width,0:height]

arr = np.abs(np.hypot(array_i-s1x, array_j-s1y) -
             np.hypot(array_i-s2x, array_j-s2y))

arr = (arr * np.ones((3,1,1))).transpose(1,2,0)

arr2 = np.zeros((width,height,3),dtype="uint8")
for ld in [200,116,100,84,68,52,36,20,8,4,2]:
    print 'now computing image for ld = '+str(ld)
    # Rest as before

And the new time is... 8.2 seconds. So you save maybe four whole seconds. On the other hand, that's almost exclusively in the image generation stages now, so maybe you can tighten them up by only generating the images you want.