Diffusion Limited Aggregation - solution to Exercise 4d in Sixty North's Boost.Python workshop

diffusion.py

# Diffusion limited aggregation simulation
# as an example solution to Sixty North's
# Boost.Python workshop.
#
# Usage:
# 
#  python diffusion.py 128 128 4096 diffusion.bmp
#
# To produce a 128x128 image with 4096 sticky
# 'grains' diffused into it.

import random
from colorsys import hsv_to_rgb

import bitmap

FOREGROUND_COLOR = (255, 0, 0)
BACKGROUND_COLOR = (0, 0, 0)

def fill_image(bmp, rgb):
    for y in range(bmp.height()):
        fill_row(bmp, y, rgb)

def fill_row(bmp, y, rgb):
    for x in range(bmp.width()):
        bmp[x, y] = rgb

def clamp_coord(v, stop):
    return max(0, min(v, stop-1))

def color_map(bmp, x, y):
    h = y / (bmp.height() - 1)
    rgb = hsv_to_rgb(h, 1.0, 1.0)
    return tuple(int(channel * 255) for channel in rgb)

def aggregable(bmp, x, y):
    for dx, dy in DIRECTIONS:
        nx = clamp_coord(x + dx, bmp.width())
        ny = clamp_coord(y + dy, bmp.height())
        if bmp[nx, ny] != BACKGROUND_COLOR:
            return True
    return False

DIRECTIONS = (
    (-1, 0), (-1, +1), (0, +1), (+1, +1), (+1, 0), (+1, -1), (0, -1), (-1, -1)
)

def diffuse_grain(bmp):
    x = random.randrange(bmp.width())
    y = bmp.height() - 1
    while not aggregable(bmp, x, y):
        dx, dy = random.choice(DIRECTIONS)
        x = clamp_coord(x + dx, bmp.width())
        y = clamp_coord(y + dy, bmp.height())
    bmp[x, y] = color_map(bmp, x, y)

def diffuse(width, height, num_grains):
    bmp = bitmap.Bitmap()
    bmp.set_width(width)
    bmp.set_height(height)

    fill_image(bmp, BACKGROUND_COLOR)
    fill_row(bmp, 0, FOREGROUND_COLOR)

    for i in range(num_grains):
        print("grain {} of {}".format(i, num_grains))
        diffuse_grain(bmp)

    return bmp

def diffusion_image(width, height, num_grains, filename):
    bmp = diffuse(width, height, num_grains)
    bmp.save_image(filename)

def main(argv=None):
    import argparse
    parser = argparse.ArgumentParser(description='Diffusion limited aggregation')
    parser.add_argument(dest='width', type=int, help='Image width in pixels')
    parser.add_argument(dest='height', type=int, help='Image height in pixels')
    parser.add_argument(dest='num_grains', type=int, help='Number of grains to aggregate')
    parser.add_argument(dest='filename', type=str, help="Filename of BMP image to write")
    args = parser.parse_args(argv)
    diffusion_image(args.width, args.height, args.num_grains, args.filename)

if __name__ == '__main__':
    import sys
    main(sys.argv[1:])