Use genetic algorithms to create graphical rainbows.

genetic_rainbows.py

#!/usr/bin/env python

from Tkinter import *
import colorsys
import optparse
import random

"""
Use genetic algorithms to create graphical rainbows.

Example usages:

python genetic-rainbows.py # use all defaults

python genetic-rainbows.py -w 8 -r 3 # make an 8x3 rainbow

# Use a population of 5000, and evolve for 1000 generations
python genetic-rainbows.py -p 5000 -g 1000

# use crossover and mutation probabilities of .9 and .1, respectively
python genetic-rainbows.py -c .9 -m .1

Each chromosome is an array of hsv descriptions; each element colors a section
of a GUI grid. To make sure we get bright colors, we fix the saturation and
value and only randomize on the hue.

There's something particularly pleasing about a one-row rainbow, but any
rectangle with a width of at least 3 works.
"""

def rainbowFitness(chromosome):
    """
    Determine the fitness for a chromosome.

    The most desirable color for an element in the chromosome (rainbow) is based
    on its x offset in the grid, eg the left-most cells should be red, the
    right-most cells should be purple.

    The closer each element is to its most desirable color, the more positive
    the fitness.
    """
    f = 0
    for i in range(len(chromosome)):
        f = f - abs(chromosome[i][0] - 1.0*(i % gui.width)/gui.width)
    return f

def seedGeneration(populationSize, numGenerations, chromosomeLength):
    """
    Allocate the chromosome space for and seed the first generation.
    """
    # allocate
    generation = [[] for i in range(populationSize)]

    # seed
    for i in range(populationSize):
        generation[i] = [(random.random(), 1, 1) for j in range(chromosomeLength)]

    return generation

def getCandidates(generation, populationSize):
    """
    Given a generation of chromosomes, pick two to become candidate parents for
    a child.
    """
    p1 = random.randint(0, populationSize - 1)
    candidate1 = generation[p1]
    p2 = random.randint(0, populationSize - 1)
    # make sure you have two different parents
    while p2 == p1:
        p2 = random.randint(0, populationSize - 1)
    candidate2 = generation[p2]
    return candidate1, candidate2

def decideParent(fitness, candidate1, candidate2):
    """
    Given two candidates, pick which gets to be the parent in the case of a
    clone. For now, just choose whichever is most fit.
    """
    if fitness(candidate1) > fitness(candidate2):
        parent = candidate1
    else:
        parent = candidate2
    return parent

def produceOffspring(crossoverProb, candidate1, candidate2, parent):
    """
    Produce a child that is either a crossover combination of its candidates, or
    a clone of one of the candidates.
    """
    if random.random() < crossoverProb:
        # Perform crossover on parents to produce new children.
        # Don't allow crossovers at the far ends of the chromosomes,
        # or that's just cloning.
        crossoverPoint = random.randint(1, len(candidate1) - 2)
        return candidate1[0:crossoverPoint] + candidate2[crossoverPoint:]
    else:
        # clone
        return parent

def mutateChild(mutationProb, chromosome):
    """
    Mutation one of the elements (one hue) in the chromosome.
    """
    if random.random() < mutationProb:
        # choose a point to mutate
        mutationPoint = random.randint(0, len(chromosome) - 1)
        mutated_h = chromosome[mutationPoint][0] + random.random()
        mutated_h = mutated_h % 1
        return chromosome[:mutationPoint] + [(mutated_h, 1, 1)] + chromosome[mutationPoint + 1:]
    else:
        return chromosome

def mostFit(generation, fitness):
    """
    Return the most fit chromosome in its generation.
    """
    maxFit = -1000
    fitIndex = 0
    for i in range(len(generation)):
        fit = fitness(generation[i])
        if fit > maxFit:
            maxFit = fit
            fitIndex = i
    return generation[fitIndex]

def displayMostFit(chromosome):
    """
    Display the rainbow based on the hues in the most fit chromosome.
    """
    gui.pixels = chromosome
    gui.draw()

def evolve(populationSize, numGenerations, crossoverProb, mutationProb, fitness, chromosomeLength):
    """
    Evolve and display rainbows until numGenerations have passed.
    """
    generation = seedGeneration(populationSize, numGenerations, chromosomeLength)

    for gen in range(1, numGenerations):
        displayMostFit(mostFit(generation, fitness))
        newGeneration = []
        for i in range(0, populationSize):
            candidate1, candidate2 = getCandidates(generation, populationSize)
            parent = decideParent(fitness, candidate1, candidate2)
            offspring = produceOffspring(crossoverProb, candidate1, candidate2, parent)
            offspring = mutateChild(mutationProb, offspring)
            newGeneration.append(offspring)
        generation = newGeneration

class GUI(object):
    MIN_RED = MIN_GREEN = MIN_BLUE = 0x0
    MAX_RED = MAX_GREEN = MAX_BLUE = 0xFF

    PIXEL_WIDTH = 50

    def __init__(self, width, height):
        self.width = width
        self.height = height
        self.tk_init()
        self.pixels = []

    def tk_init(self):
        self.root = Tk()
        self.root.title("Wall %d x %d" % (self.width, self.height))
        self.root.resizable(0, 0)
        self.frame = Frame(self.root, bd=5, relief=SUNKEN)
        self.frame.pack()

        self.canvas = Canvas(self.frame,
                             width=self.PIXEL_WIDTH * self.width,
                             height=self.PIXEL_WIDTH * self.height,
                             bd=0, highlightthickness=0)
        self.canvas.pack()
        self.root.update()

    def hsv_to_rgb(self, hsv):
        rgb = colorsys.hsv_to_rgb(*hsv)
        red = self.MAX_RED * rgb[0]
        green = self.MAX_GREEN * rgb[1]
        blue = self.MAX_BLUE * rgb[2]
        return (red, green, blue)

    def get_rgb(self, hsv):
        red, green, blue = self.hsv_to_rgb(hsv)
        red = int(float(red) / (self.MAX_RED - self.MIN_RED) * 0xFF)
        green = int(float(green) / (self.MAX_GREEN - self.MIN_GREEN) * 0xFF)
        blue = int(float(blue) / (self.MAX_BLUE - self.MIN_BLUE) * 0xFF)
        return (red, green, blue)

    def draw(self):
        for x in range(len(self.pixels)):
            x_0 = (x % self.width) * self.PIXEL_WIDTH
            y_0 = (x / self.width) * self.PIXEL_WIDTH
            x_1 = x_0 + self.PIXEL_WIDTH
            y_1 = y_0 + self.PIXEL_WIDTH
            hue = "#%02x%02x%02x" % self.get_rgb(self.pixels[x])
            self.canvas.create_rectangle(x_0, y_0, x_1, y_1, fill=hue)
        self.canvas.update()

if __name__ == "__main__":
    parser = optparse.OptionParser("""Usage: %prog [options]""")
    parser.add_option("-p", "--population", type="int",
                      action="store", dest="populationSize",
                      default=500, help="size of population")
    parser.add_option("-g", "--num-generations", type="int",
                      action="store", dest="numGenerations",
                      default=150, help="number of generations to run")
    parser.add_option("-c", "--crossover-prob", type="float",
                      action="store", dest="crossoverProb",
                      default=.8, help="crossover probability")
    parser.add_option("-m", "--mutation-prob", type="float",
                      action="store", dest="mutationProb",
                      default=.05, help="mutation probability")
    parser.add_option("-w", "--rainbow-width", type="int",
                      action="store", dest="rainbowWidth",
                      default=10, help="rainbow width")
    parser.add_option("-r", "--rainbow-height", type="int",
                      action="store", dest="rainbowHeight",
                      default=1, help="rainbow height")
    (opts, args) = parser.parse_args(sys.argv[1:])

    if opts.rainbowWidth < 3:
        print >>sys.stderr, "Rainbows have a minimum width of 3. Please choose a larger width."
        sys.exit(1)

    global gui
    gui = GUI(opts.rainbowWidth, opts.rainbowHeight)

    evolve(opts.populationSize, opts.numGenerations, opts.crossoverProb,
           opts.mutationProb, rainbowFitness, gui.width * gui.height)

    raw_input("Press a key to close")