Skip to content

Instantly share code, notes, and snippets.

Created June 8, 2015 15:06
Show Gist options
  • Star 2 You must be signed in to star a gist
  • Fork 1 You must be signed in to fork a gist
  • Save StuartGordonReid/09abf6a7dc9cfaf85506 to your computer and use it in GitHub Desktop.
Save StuartGordonReid/09abf6a7dc9cfaf85506 to your computer and use it in GitHub Desktop.
import os
import math
import time
import numpy
import pandas
import random
import matplotlib
import numpy.random as nrand
import matplotlib.pylab as plt
from sklearn.preprocessing import normalize
class AntColonyOptimization:
def __init__(self):
class Grid:
def __init__(self, height, width, path, rand_test=True):
This method initializes a grid object. A grid is basically just a 2D array of Datum objects
:param height: this is the height of the grid
:param width: this is the weight of the grid
self.path = path
# Store the dimensions of the grid
self.dim = numpy.array([height, width])
# Initialize an empty numpy matrix of type Datum
self.grid = numpy.empty((height, width), dtype=Datum)
if rand_test:
# This is used to fill the grid randomly
# This makes the plot redraw
plt.figure(figsize=(10, 10))
self.max_d = 0.001
def rand_grid(self, sparse):
A method for randomly initializing the grid
:param sparse: the percentage of the grid to fill up
for y in range(self.dim[0]):
for x in range(self.dim[1]):
if random.random() <= sparse:
r = random.randint(0, 1)
if r == 0:
self.grid[y][x] = Datum(nrand.normal(5, 0.25, 10))
elif r == 1:
self.grid[y][x] = Datum(nrand.normal(-5, 0.25, 10))
def matrix_grid(self):
This method condenses the grid (2D array of Datum objects) to a matrix which can be visualized
:return: matrix of the grid
matrix = numpy.empty((self.dim[0], self.dim[1]))
for y in range(self.dim[0]):
for x in range(self.dim[1]):
if self.grid[y][x] is not None:
matrix[y][x] = self.get_grid()[y][x].condense()
return matrix
def plot_grid(self, name="", save_figure=True):
This plots the 2D representation of the grid
:param name: the name of the image to save
plt.matshow(self.matrix_grid(), cmap="RdBu", fignum=0)
# Option to save images
if save_figure:
plt.savefig(self.path + name + '.png')
# plt.draw()
def get_grid(self):
return self.grid
def get_probability(self, d, y, x, n, c):
This gets the probability of drop / pickup for any given Datum, d
:param d: the datum
:param x: the x location of the datum / ant carrying datum
:param y: the y location of the datum / ant carrying datum
:param n: the size of the neighbourhood function
:param c: constant for convergence control
:return: the probability of
# Starting x and y locations
y_s = y - n
x_s = x - n
total = 0.0
# For each neighbour
for i in range((n*2)+1):
xi = (x_s + i) % self.dim[0]
for j in range((n*2)+1):
# If we are looking at a neighbour
if j != x and i != y:
yj = (y_s + j) % self.dim[1]
# Get the neighbour, o
o = self.grid[xi][yj]
# Get the similarity of o to x
if o is not None:
s = d.similarity(o)
total += s
# Normalize the density by the max seen distance to date
md = total / (math.pow((n*2)+1, 2) - 1)
if md > self.max_d:
self.max_d = md
density = total / (self.max_d * (math.pow((n*2)+1, 2) - 1))
density = max(min(density, 1), 0)
t = math.exp(-c * density)
probability = (1-t)/(1+t)
return probability
class Ant:
def __init__(self, y, x, grid):
This initializes an ant object. This Ant class is just a dumb ant with no memory
:param y: the y location it is initialized to
:param x: the x location it is initialized to
:param grid: a reference to the grid
self.loc = numpy.array([y, x])
self.carrying = grid.get_grid()[y][x]
self.grid = grid
def move(self, n, c):
A recursive function for making ants move around the grid
:param step_size: the size of each step
step_size = random.randint(1, 25)
# Add some vector (-1,+1) * step_size to the ants location
self.loc += nrand.randint(-1 * step_size, 1 * step_size, 2)
# Mod the new location by the grid size to prevent overflow
self.loc = numpy.mod(self.loc, self.grid.dim)
# Get the object at that location on the grid
o = self.grid.get_grid()[self.loc[0]][self.loc[1]]
# If the cell is occupied, move again
if o is not None:
# If the ant is not carrying an object
if self.carrying is None:
# Check if the ant picks up the object
if self.p_pick_up(n, c) >= random.random():
# Pick up the object and rem from grid
self.carrying = o
self.grid.get_grid()[self.loc[0]][self.loc[1]] = None
# If not then move
self.move(n, c)
# If carrying an object then just move
self.move(n, c)
# If on an empty cell
if self.carrying is not None:
# Check if the ant drops the object
if self.p_drop(n, c) >= random.random():
# Drop the object at the empty location
self.grid.get_grid()[self.loc[0]][self.loc[1]] = self.carrying
self.carrying = None
def p_pick_up(self, n, c):
Returns the probability of picking up an object
:param n: the neighbourhood size
:return: probability of picking up
ant = self.grid.get_grid()[self.loc[0]][self.loc[1]]
return 1 - self.grid.get_probability(ant, self.loc[0], self.loc[1], n, c)
def p_drop(self, n, c):
Returns the probability of dropping an object
:return: probability of dropping
ant = self.carrying
return self.grid.get_probability(ant, self.loc[0], self.loc[1], n, c)
class Datum:
def __init__(self, data):
A Datum object is basically just a ND vector
:param data: the ND vector
""" = data
def similarity(self, datum):
Returns the sum-squared distance between this datum and some other datum
:param datum: the other datum
:return: sum squared distance
diff = numpy.abs( -
return numpy.sum(diff**2)
def condense(self):
A method for condensing ND into 1D for visualization ... many options exist for this
:return: the 1D representation of the vector
return numpy.mean(
def optimize(height, width, ants, sims, n, c, freq=500, path="image"):
Main method for running the algorithm
# Initialize the grid
grid = Grid(height, width, path)
# Create the ants
ant_agents = []
for i in range(ants):
ant = Ant(random.randint(0, height - 1), random.randint(0, width - 1), grid)
for i in range(sims):
for ant in ant_agents:
ant.move(n, c)
if i % freq == 0:
s = "img" + str(i).zfill(6)
if __name__ == '__main__':
optimize(500, 500, 250, 500000, 25, 10, freq=500, path="Video 8/")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment