genetic adaption for travelling salesman

travelling.py

import math
import random
import matplotlib.pyplot as plt

import pickle
FNAME = "data.obj"

# generate 30 towns in a 1000*1000 grid
TOWN = 30 
GRID = 1000

# size of population pool
POPULATION = 50

# how many of nodes to shift over
CROSSOVER_RATE = 0.5

### genetic stuff 
def fitness(path, towns):
    # calculate distance of the whole trip (back to original town)
    total, past = 0, None
    for town in path+[path[0]]:
        if past is not None:
            dx = towns[town][0] - towns[past][0]
            dy = towns[town][1] - towns[past][1]
            total += math.sqrt(dx**2 + dy**2)
        past = town 

    return total

def crossover(path1, path2):
    start = random.randint(1, TOWN-1)
    cnt = random.randint(1, TOWN*CROSSOVER_RATE)

    # clone from path1
    child = [-1]*TOWN
    for i in range(cnt):
        idx = (start+i)%TOWN
        child[idx] = path1[idx]
    
    # copy the rest of the path2 over
    idx = 0
    for i in range(TOWN):
        if child[i]>-1: continue
        while child.count(path2[idx]): idx+=1
        child[i]=path2[idx]

    return child

def mutate(path):
    # swapping one of the towns
    idx1 = random.randint(0, TOWN-1)
    idx2 = random.randint(0, TOWN-1)
    
    child = path[:]
    child[idx1], child[idx2] = path[idx2], path[idx1]
    return child

def populate(paths):
    # generating the children 
    children = []
    for i in range(len(paths)/2):
        children.append(crossover(paths[i], paths[i+1]))
    
    # mutating for new child
    for path in paths:
        children.append(mutate(path))

    return paths+children

def cull(paths, towns):
    # keeping only the top few 
    weighted = sorted([(fitness(p, towns), p) for p in paths])
    return [p[1] for p in weighted[:POPULATION]]

### scaffolding
def gen():
    towns = []
    for i in range(TOWN):
        towns.append((random.randint(1, GRID), random.randint(1, GRID)))
    return towns

def showGraph(towns, path=None):
    plt.plot([t[0] for t in towns], [t[1] for t in towns], "ro")
    plt.axis([0, 1000, 0, 1000])

    if path is not None:
        best = [towns[t] for t in path+[path[0]]]
        plt.plot([t[0] for t in best], [t[1] for t in best], "b-")

    plt.show()

def main():
    towns = gen()
    showGraph(towns) 
    
    # default paths (forward and backwards) 
    paths = []
    paths.append(list(range(TOWN)))
    paths.append(list(range(TOWN-1,-1,-1)))
   
    # letting time do its thing
    for i in range(501):
        paths = populate(paths)
        paths = cull(paths, towns)

        if i%50==0:
            showGraph(towns, paths[0])
            print "best @", i, ":", fitness(paths[0], towns)

    # storing it
    pickle.dump({
        "towns": towns,
        "paths": paths
    }, open(FNAME, "wb"))

if __name__=="__main__":
    main()