ant_colony.py

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