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(omega*time - 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

The only changes that come to my mind is to move some operations out of the loop:

for i in xrange(width):
    if i%(width/10)==0:
        print i,    
    if i%20==0:
        print '.',
    arri = arr[i]
    is1x = i - s1x
    is2x = i - s2x
    for j in xrange(height):
        d1 = hy(is1x,j-s1y)
        d2 = hy(is2x,j-s2y)
        arri[j] = abs(d1-d2)

The improvement, if any, will probably be minor though.

Answer 2

List comprehensions are much faster than loops. For example, instead of

for j in xrange(height):
        d1 = hy(i-s1x,j-s1y)
        d2 = hy(i-s2x,j-s2y)
        arr[i][j] = abs(d1-d2)

You'd write

arr[i] = [abs(hy(i-s1x,j-s1y) - hy(i-s2x,j-s2y)) for j in xrange(height)]

On the other hand, if you're really trying to "optimize", then you might want to reimplement this algorithm in C, and use SWIG or the like to call it from python.

Answer 3

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.

Answer 4

If you use array operations instead of loops, it is much, much faster. For me, the image generation is now what takes so long time. Instead of your two i,j loops, I have this:

I,J = np.mgrid[0:width,0:height]
D1 = np.hypot(I - s1x, J - s1y)
D2 = np.hypot(I - s2x, J - s2y)

arr = np.abs(D1-D2)
# triplicate into 3 layers
arr = np.array((arr, arr, arr)).transpose(1,2,0)
# .. continue program

The basics that you want to remember for the future is: this is not about optimization; using array forms in numpy is just using it like it is supposed to be used. With experience, your future projects should not go the detour over python loops, the array forms should be the natural form.

What we did here was really simple. Instead of math.hypot we found numpy.hypot and used it. Like all such numpy functions, it accepts ndarrays as arguments, and does exactly what we want.