Using python generators   Leave a comment

I just made a generator in python.

Wow. Just, Wow.

Generators are really awesome.
I’ve been having a couple issues with my IFS algorithm, because it’s not quite as fast as I would like, and I can’t seem to do everything I want just by passing callback functions. Here’s the original algorithm:

def run_IFS(iters, skip, rules, draw_method, speak_method=None, interval = 10000):
    """Calculate the points in an iterated fractal system, drawing them with draw_method,
    and announcing results with speak_method every interval iterations."""
    probs = [r[PROB_I] for r in rules]
    rules = [r[RULE_I] for r in rules]

    x = random_float()
    y = random_float()

    while i < iters + skip
        chosen = choose_rule(probs)
        x, y = calculate_transform(x, y, rules[chosen])
        if i > skip:
            draw_method(x, y, chosen)
        #Announce the results
        if not speak_method: pass
        elif i%interval == 0:
            speak_method(x, y, i)
        i += 1

This isn’t bad code, and it’s not particularly difficult to use. I know, because it was very easy for me to add a tkinter interface in less than ten minutes. However, when I wanted to write a function to calculate the best size settings for a fractal, I found it to be rather messy:


def old_calculate_best_settings(rules, scale, iters = 10000, skip = 100):
    """Return a tuple of width and height and x_off and y_off settings
((w,h),(x_off, y_off))"""
    #note, I know this is ugly, I don't know how else to do it though
    global h_x, h_y, l_x, l_y
    h_x = h_y = l_x = l_y = 0

    def draw_method(x, y, chosen):
        global h_x, h_y, l_x, l_y
        if x*scale > h_x: h_x = x*scale
        elif x*scale < l_x: l_x = x*scale
        if y*scale > h_y: h_y = y*scale
        elif y*scale < l_y: l_y = y*scale

    run_IFS(iters, skip, rules, draw_method)
    width = int(abs(l_x)+abs(h_x))
    height = int(abs(l_y)+abs(h_y))
    x_off = int(abs(l_x))
    y_off = int(abs(l_y))

    return ((width, height), (x_off, y_off))

Yes, that’s right.

    Global Variables


Because of an interesting feature in python, it’s not possible for draw_method to change the values of h_x (and friends) without resorting to globals. This is because you can access variables from outside scope, but assigning to them makes a new local variable with the same name, which is not really what we want (thanks to bob2 from #python on freenode.net for the explanation). Of course, I could have made these variables classes, and then called some __set__ method to change their values. Purge your mind friend. The very fact that we’re thinking of such roundabout, convoluted methods is a sure sign that something is not being done right.

Enter generators

Python, of course, has a nice solution to this problem. We can see that this function is basically iterating and producing points. A possible solution would be to store the points in a list, and then iterate over these to see what the low/high points are. This is halfway there, however, when we’re talking millions of points, memory quickly becomes an issue (at least on my 486 :). Instead, we can use a generator, which behaves like a list, and makes for more “pythonic” code. Here’s a simple example:


def generate_x():
    i = 0
    while i < 10:
        print i
        yield i
        i += 1

for x in generate_x():
    print x

When you run that code (or on codepad) you can see that the prints are interleaved in the generator and the for loop, demonstrating how it works.
To change a function to a generator, we just rearrange the function a bit and make it yield rather than return the points.
Here’s the run_IFS function as a generator:


def IFS_generator(iters, skip, rules):

    probs = [r[PROB_I] for r in rules]
    rules = [r[RULE_I] for r in rules]

    x = random_float()
    y = random_float()
    i = 0

    while i < iters + skip:         chosen = choose_rule(probs)         x, y, = calculate_transform(x, y, rules[chosen])         if i > skip:
            yield (x, y, chosen)
        i += 1

That’s a lot simpler, isn’t it? We can now modify the calculate_best_settings function to use the generator:


def calculate_best_settings(rules, scale, iters = 10000, skip = 100):
    """Return a tuple of width and height and x_off and y_off settings
((w,h),(x_off, y_off))"""
    h_x = l_x = h_y = l_y = 0

    for points in IFS_generator(iters, skip, rules):
        x, y, _ = points
        if (x*scale) > h_x: h_x = x*scale
        elif (x*scale) < l_x: l_x = x*scale         if (y*scale) > h_y: h_y = y*scale
        if (y*scale) < l_y: l_y = y*scale

    width = int(abs(l_x)+abs(h_x))
    height = int(abs(l_y)+abs(h_y))
    x_off = int(abs(l_x))
    y_off = int(abs(l_y))

    return ((width, height), (x_off, y_off))

See how this is now nice, pythonic code? No more messing with global variables and callbacks. Of course, there’s also a nice speed bonus in getting rid of those callbacks.

After timing both functions, I got these results:
Old:
[4.5442321300506592, 4.4281911849975586, 4.3002359867095947] min: 4.30023598671
New:
[5.3141071796417236, 4.0443418025970459, 3.9916889667510986] min: 3.99168896675

This is a pretty significant speed up for such a small change! Note that the calculate_best_settings function only runs ~10000 iterations. When you generate a fractal with 1000000 iterations this difference becomes even greater. The code gets a boost once we’re rid of draw_method and speak_method. Of course, the next step would be to implement this function as a C extension.

Further reading:
Pep 255 – Simple Generators

Advertisements

Posted 19/04/2010 by Emmanuel Jacyna in Code, Python

Tagged with , , ,

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s

%d bloggers like this: