tsp/test/demo2.py

import numpy as np

class AntColonyOptimization:
    def __init__(self, distances, num_ants=10, num_iterations=100, alpha=1, beta=1, evaporation=0.5, q0=0.9):
        self.distances = distances
        self.num_ants = num_ants
        self.num_iterations = num_iterations
        self.alpha = alpha
        self.beta = beta
        self.evaporation = evaporation
        self.q0 = q0

    def optimize(self):
        num_cities = len(self.distances)
        pheromone = np.ones((num_cities, num_cities))
        best_path = None
        best_length = float('inf')

        for _ in range(self.num_iterations):
            paths = []
            lengths = []

            for _ in range(self.num_ants):
                path = self.construct_path(pheromone)
                length = self.calculate_length(path)
                paths.append(path)
                lengths.append(length)

                if length < best_length:
                    best_length = length
                    best_path = path

            self.update_pheromone(pheromone, paths, lengths)

        return best_path

    def construct_path(self, pheromone):
        num_cities = len(pheromone)
        path = [0]  # 起始城市为0
        visited = np.zeros(num_cities, dtype=bool)
        visited[0] = True

        for _ in range(num_cities - 1):
            probs = self.calculate_probs(pheromone, path[-1], visited)
            next_city = self.choose_next_city(probs)

            path.append(next_city)
            visited[next_city] = True

        return path

    def calculate_probs(self, pheromone, curr_city, visited):
        num_cities = len(pheromone)
        probs = np.zeros(num_cities)

        denominator = 0
        for city in range(num_cities):
            if not visited[city]:
                denominator += pheromone[curr_city, city] ** self.alpha + (1.0 / self.distances[curr_city, city]) ** self.beta

        for city in range(num_cities):
            if not visited[city]:
                numerator = pheromone[curr_city, city] ** self.alpha + (1.0 / self.distances[curr_city, city]) ** self.beta
                probs[city] = numerator / denominator

        return probs

    def choose_next_city(self, probs):
        q = np.random.rand()

        if q < self.q0:
            next_city = np.argmax(probs)
        else:
            next_city = np.random.choice(len(probs), p=probs)

        return next_city

    def calculate_length(self, path):
        length = 0
        for i in range(len(path) - 1):
            length += self.distances[path[i], path[i+1]]

        return length

    def update_pheromone(self, pheromone, paths, lengths):
        pheromone *= self.evaporation

        for i, path in enumerate(paths):
            for j in range(len(path) - 1):
                city1 = path[j]
                city2 = path[j+1]
                pheromone[city1, city2] += 1.0 / lengths[i]
                pheromone[city2, city1] += 1.0 / lengths[i]

distances = np.array([
    [np.inf, 50, 99, np.inf, 64, np.inf],
    [50, np.inf, 12, 45, 74, 13],
    [99, 12, np.inf, 25, np.inf, 61],
    [np.inf, 45, 25, np.inf, 45, 47],
    [64, 74, np.inf, 45, np.inf, np.inf],
    [np.inf, 13, 61, 47, np.inf, np.inf]
])

aco = AntColonyOptimization(distances)
best_path = aco.optimize()
print(best_path)