Poisson Disc

main.py

from math import pi, sin, cos, hypot, floor
import cairocffi as cairo
import random

R = 0.012

class Grid(object):
    def __init__(self, r):
        self.r = r
        self.size = r / 2 ** 0.5
        self.cells = {}
    def points(self):
        return self.cells.values()
    def normalize(self, x, y):
        i = int(floor(x / self.size))
        j = int(floor(y / self.size))
        return (i, j)
    def nearby(self, x, y):
        result = []
        i, j = self.normalize(x, y)
        for p in xrange(i - 2, i + 3):
            for q in xrange(j - 2, j + 3):
                if (p, q) in self.cells:
                    result.append(self.cells[(p, q)])
        return result
    def nearest(self, x, y):
        def func((bx, by)):
            return hypot(x - bx, y - by)
        return sorted(self.nearby(x, y), key=func)
    def insert(self, x, y):
        for bx, by in self.nearby(x, y):
            if hypot(x - bx, y - by) < self.r:
                return False
        i, j = self.normalize(x, y)
        self.cells[(i, j)] = (x, y)
        return True

def poisson_disc(x1, y1, x2, y2, r, n):
    x = x1 + (x2 - x1) / 2.0
    y = y1 + (y2 - y1) / 2.0
    active = [(x, y)]
    grid = Grid(r)
    grid.insert(x, y)
    while active:
        ax, ay = random.choice(active)
        for i in xrange(n):
            a = random.random() * 2 * pi
            d = random.random() * r + r
            x = ax + cos(a) * d
            y = ay + sin(a) * d
            if x < x1 or y < y1 or x > x2 or y > y2:
                continue
            if not grid.insert(x, y):
                continue
            active.append((x, y))
            break
        else:
            active.remove((ax, ay))
    return grid

def render(grid):
    scale = 1024
    width = height = scale
    surface = cairo.ImageSurface(cairo.FORMAT_RGB24, width, height)
    dc = cairo.Context(surface)
    dc.set_line_cap(cairo.LINE_CAP_ROUND)
    dc.set_line_join(cairo.LINE_JOIN_ROUND)
    dc.scale(scale, scale)
    dc.set_source_rgb(0, 0, 0)
    dc.paint()
    for x1, y1 in grid.points():
        if hypot(x1 - 0.5, y1 - 0.5) > 0.45:
            continue
        for x2, y2 in grid.nearest(x1, y1)[1:2]:
            if hypot(x2 - 0.5, y2 - 0.5) > 0.45:
                continue
            dc.move_to(x1, y1)
            dc.line_to(x2, y2)
    dc.set_source_rgb(0.996, 0.961, 0.922)
    dc.set_line_width(R * 0.8)
    dc.stroke_preserve()
    dc.set_source_rgb(0.192, 0.243, 0.290)
    dc.set_line_width(R * 0.4)
    dc.stroke_preserve()
    return surface

def main():
    grid = poisson_disc(0, 0, 1, 1, R, 32)
    surface = render(grid)
    surface.write_to_png('output.png')

if __name__ == '__main__':
    main()