import bisect import itertools class World: uid = 0 def __init__(self, nodes, lfunc, **kwargs): self.uid = self.__class__.uid self.__class__.uid += 1 self.name = kwargs.get('name', 'world{}'.format(self.uid)) self.description = kwargs.get('description', None) self._nodes = nodes self.lfunc = lfunc self.edges = self.create_edges() @property def nodes(self): """Node IDs.""" return list(range(len(self._nodes))) def create_edges(self): edges = {} for m in self.nodes: for n in self.nodes: a, b = self.data(m), self.data(n) if a != b: edge = Edge(a, b, length=self.lfunc(a, b)) edges[m, n] = edge return edges def reset_pheromone(self, level=0.01): for edge in self.edges.values(): edge.pheromone = level def data(self, idx, idy=None): try: if idy is None: return self._nodes[idx] else: return self.edges[idx, idy] except IndexError: return None class Edge: def __init__(self, start, end, length=None, pheromone=None): self.start = start self.end = end self.length = 1 if length is None else length self.pheromone = 0.1 if pheromone is None else pheromone def __eq__(self, other): if isinstance(other, self.__class__): return self.__dict__ == other.__dict__ return False class Ant: uid = 0 def __init__(self, alpha=1, beta=3): self.uid = self.__class__.uid self.__class__.uid += 1 self.world = None self.alpha = alpha self.beta = beta self.start = None self.distance = 0 self.visited = [] self.unvisited = [] self.traveled = [] def initialize(self, world, start=None): self.world = world if start is None: self.start = random.randrange(len(self.world.nodes)) else: self.start = start self.distance = 0 self.visited = [self.start] self.unvisited = [n for n in self.world.nodes if n != self.start] self.traveled = [] return self def clone(self): ant = Ant(self.alpha, self.beta) ant.world = self.world ant.start = self.start ant.visited = self.visited[:] ant.unvisited = self.unvisited[:] ant.traveled = self.traveled[:] ant.distance = self.distance return ant def node(self): try: return self.visited[-1] except IndexError: return None def tour(self): return [self.world.data(i) for i in self.visited] def path(self): return [edge for edge in self.traveled] def __eq__(self, other): """Return ``True`` if the distance is equal to the other distance. :param Ant other: right-hand argument :rtype: bool """ return self.distance == other.distance def __lt__(self, other): """Return ``True`` if the distance is less than the other distance. :param Ant other: right-hand argument :rtype: bool """ return self.distance < other.distance def can_move(self): return len(self.traveled) != len(self.visited) def move(self): remaining = self.remaining_moves() choice = self.choose_move(remaining) return self.make_move(choice) def remaining_moves(self): return self.unvisited def choose_move(self, choices): if len(choices) == 0: return None if len(choices) == 1: return choices[0] # Find the weight of the edges that take us to each of the choices. weights = [] for move in choices: edge = self.world.edges[self.node, move] weights.append(self.weigh(edge)) # Choose one of them using a weighted probability. total = sum(weights) cumdist = list(itertools.accumulate(weights)) + [total] return choices[bisect.bisect(cumdist, random.random() * total)] def make_move(self, dest): ori = self.node if dest is None: if self.can_move() is False: return None dest = self.start # last move is back to the start else: self.visited.append(dest) self.unvisited.remove(dest) edge = self.world.edges[ori, dest] self.traveled.append(edge) self.distance += edge.length return edge def weigh(self, edge): pre = 1 / (edge.length or 1) post = edge.pheromone return post ** self.alpha * pre ** self.beta import random from copy import copy class Solver: def __init__(self, **kwargs): self.alpha = kwargs.get('alpha', 1) self.beta = kwargs.get('beta', 3) self.rho = kwargs.get('rho', 0.8) self.q = kwargs.get('Q', 1) self.t0 = kwargs.get('t0', .01) self.limit = kwargs.get('limit', 100) self.ant_count = kwargs.get('ant_count', 10) self.elite = kwargs.get('elite', .5) def create_colony(self, world): if self.ant_count < 1: return self.round_robin_ants(world, len(world.nodes)) return self.random_ants(world, self.ant_count) def reset_colony(self, colony): for ant in colony: ant.initialize(ant.world) def aco(self, colony): self.find_solutions(colony) self.global_update(colony) return sorted(colony)[0] def solve(self, world): world.reset_pheromone(self.t0) global_best = None colony = self.create_colony(world) for i in range(self.limit): self.reset_colony(colony) local_best = self.aco(colony) if global_best is None or local_best < global_best: global_best = copy(local_best) self.trace_elite(global_best) return global_best def solutions(self, world): world.reset_pheromone(self.t0) global_best = None colony = self.create_colony(world) for i in range(self.limit): self.reset_colony(colony) local_best = self.aco(colony) if global_best is None or local_best < global_best: global_best = copy(local_best) yield global_best self.trace_elite(global_best) def round_robin_ants(self, world, count): starts = world.nodes n = len(starts) return [ Ant(self.alpha, self.beta).initialize( world, start=starts[i % n]) for i in range(count) ] def random_ants(self, world, count, even=False): ants = [] starts = world.nodes n = len(starts) if even: if count > n: for i in range(self.ant_count // n): ants.extend([ Ant(self.alpha, self.beta).initialize( world, start=starts[j]) for j in range(n) ]) ants.extend([ Ant(self.alpha, self.beta).initialize( world, start=starts.pop(random.randrange(n - i))) for i in range(count % n) ]) else: # Just pick random nodes. ants.extend([ Ant(self.alpha, self.beta).initialize( world, start=starts[random.randrange(n)]) for i in range(count) ]) return ants def find_solutions(self, ants): ants_done = 0 while ants_done < len(ants): ants_done = 0 for ant in ants: if ant.can_move(): edge = ant.move() self.local_update(edge) else: ants_done += 1 def local_update(self, edge): edge.pheromone = max(self.t0, edge.pheromone * self.rho) def global_update(self, ants): ants = sorted(ants)[:len(ants) // 2] for a in ants: p = self.q / a.distance for edge in a.path: edge.pheromone = max( self.t0, (1 - self.rho) * edge.pheromone + p) def trace_elite(self, ant): if self.elite: p = self.elite * self.q / ant.distance for edge in ant.path: edge.pheromone += p