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): pass 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 self.rand_grid(0.25) # This makes the plot redraw plt.ion() 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])) matrix.fill(0) 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 :return: """ 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 else: self.move(n, c) # If carrying an object then just move else: self.move(n, c) # If on an empty cell else: 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 """ self.data = 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(self.data - datum.data) 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(self.data) 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) ant_agents.append(ant) for i in range(sims): for ant in ant_agents: ant.move(n, c) if i % freq == 0: print(i) s = "img" + str(i).zfill(6) grid.plot_grid(s) if __name__ == '__main__': optimize(500, 500, 250, 500000, 25, 10, freq=500, path="Video 8/")