add graphs/ant_colony_optimization_algorithms.py (#11163)

* add ant_colonyant_colony_optimization_algorithms.py

* Modify details

* Modify type annotation

* Add tests for KeyError, IndexError, StopIteration, etc.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Christian Clauss <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

Diff for: graphs/ant_colony_optimization_algorithms.py

@@ -0,0 +1,226 @@
+"""
+Use an ant colony optimization algorithm to solve the travelling salesman problem (TSP)
+which asks the following question:
+"Given a list of cities and the distances between each pair of cities, what is the
+ shortest possible route that visits each city exactly once and returns to the origin
+ city?"
+
+https://en.wikipedia.org/wiki/Ant_colony_optimization_algorithms
+https://en.wikipedia.org/wiki/Travelling_salesman_problem
+
+Author: Clark
+"""
+
+import copy
+import random
+
+cities = {
+    0: [0, 0],
+    1: [0, 5],
+    2: [3, 8],
+    3: [8, 10],
+    4: [12, 8],
+    5: [12, 4],
+    6: [8, 0],
+    7: [6, 2],
+}
+
+
+def main(
+    cities: dict[int, list[int]],
+    ants_num: int,
+    iterations_num: int,
+    pheromone_evaporation: float,
+    alpha: float,
+    beta: float,
+    q: float,  # Pheromone system parameters Q，which is a constant
+) -> tuple[list[int], float]:
+    """
+    Ant colony algorithm main function
+    >>> main(cities=cities, ants_num=10, iterations_num=20,
+    ...      pheromone_evaporation=0.7, alpha=1.0, beta=5.0, q=10)
+    ([0, 1, 2, 3, 4, 5, 6, 7, 0], 37.909778143828696)
+    >>> main(cities={0: [0, 0], 1: [2, 2]}, ants_num=5, iterations_num=5,
+    ...      pheromone_evaporation=0.7, alpha=1.0, beta=5.0, q=10)
+    ([0, 1, 0], 5.656854249492381)
+    >>> main(cities={0: [0, 0], 1: [2, 2], 4: [4, 4]}, ants_num=5, iterations_num=5,
+    ...      pheromone_evaporation=0.7, alpha=1.0, beta=5.0, q=10)
+    Traceback (most recent call last):
+      ...
+    IndexError: list index out of range
+    >>> main(cities={}, ants_num=5, iterations_num=5,
+    ...      pheromone_evaporation=0.7, alpha=1.0, beta=5.0, q=10)
+    Traceback (most recent call last):
+      ...
+    StopIteration
+    >>> main(cities={0: [0, 0], 1: [2, 2]}, ants_num=0, iterations_num=5,
+    ...      pheromone_evaporation=0.7, alpha=1.0, beta=5.0, q=10)
+    ([], inf)
+    >>> main(cities={0: [0, 0], 1: [2, 2]}, ants_num=5, iterations_num=0,
+    ...      pheromone_evaporation=0.7, alpha=1.0, beta=5.0, q=10)
+    ([], inf)
+    >>> main(cities={0: [0, 0], 1: [2, 2]}, ants_num=5, iterations_num=5,
+    ...      pheromone_evaporation=1, alpha=1.0, beta=5.0, q=10)
+    ([0, 1, 0], 5.656854249492381)
+    >>> main(cities={0: [0, 0], 1: [2, 2]}, ants_num=5, iterations_num=5,
+    ...      pheromone_evaporation=0, alpha=1.0, beta=5.0, q=10)
+    ([0, 1, 0], 5.656854249492381)
+    """
+    # Initialize the pheromone matrix
+    cities_num = len(cities)
+    pheromone = [[1.0] * cities_num] * cities_num
+
+    best_path: list[int] = []
+    best_distance = float("inf")
+    for _ in range(iterations_num):
+        ants_route = []
+        for _ in range(ants_num):
+            unvisited_cities = copy.deepcopy(cities)
+            current_city = {next(iter(cities.keys())): next(iter(cities.values()))}
+            del unvisited_cities[next(iter(current_city.keys()))]
+            ant_route = [next(iter(current_city.keys()))]
+            while unvisited_cities:
+                current_city, unvisited_cities = city_select(
+                    pheromone, current_city, unvisited_cities, alpha, beta
+                )
+                ant_route.append(next(iter(current_city.keys())))
+            ant_route.append(0)
+            ants_route.append(ant_route)
+
+        pheromone, best_path, best_distance = pheromone_update(
+            pheromone,
+            cities,
+            pheromone_evaporation,
+            ants_route,
+            q,
+            best_path,
+            best_distance,
+        )
+    return best_path, best_distance
+
+
+def distance(city1: list[int], city2: list[int]) -> float:
+    """
+    Calculate the distance between two coordinate points
+    >>> distance([0, 0], [3, 4] )
+    5.0
+    >>> distance([0, 0], [-3, 4] )
+    5.0
+    >>> distance([0, 0], [-3, -4] )
+    5.0
+    """
+    return (((city1[0] - city2[0]) ** 2) + ((city1[1] - city2[1]) ** 2)) ** 0.5
+
+
+def pheromone_update(
+    pheromone: list[list[float]],
+    cities: dict[int, list[int]],
+    pheromone_evaporation: float,
+    ants_route: list[list[int]],
+    q: float,  # Pheromone system parameters Q，which is a constant
+    best_path: list[int],
+    best_distance: float,
+) -> tuple[list[list[float]], list[int], float]:
+    """
+    Update pheromones on the route and update the best route
+    >>>
+    >>> pheromone_update(pheromone=[[1.0, 1.0], [1.0, 1.0]],
+    ...                  cities={0: [0,0], 1: [2,2]}, pheromone_evaporation=0.7,
+    ...                  ants_route=[[0, 1, 0]], q=10, best_path=[],
+    ...                  best_distance=float("inf"))
+    ([[0.7, 4.235533905932737], [4.235533905932737, 0.7]], [0, 1, 0], 5.656854249492381)
+    >>> pheromone_update(pheromone=[],
+    ...                  cities={0: [0,0], 1: [2,2]}, pheromone_evaporation=0.7,
+    ...                  ants_route=[[0, 1, 0]], q=10, best_path=[],
+    ...                  best_distance=float("inf"))
+    Traceback (most recent call last):
+      ...
+    IndexError: list index out of range
+    >>> pheromone_update(pheromone=[[1.0, 1.0], [1.0, 1.0]],
+    ...                  cities={}, pheromone_evaporation=0.7,
+    ...                  ants_route=[[0, 1, 0]], q=10, best_path=[],
+    ...                  best_distance=float("inf"))
+    Traceback (most recent call last):
+      ...
+    KeyError: 0
+    """
+    for a in range(len(cities)):  # Update the volatilization of pheromone on all routes
+        for b in range(len(cities)):
+            pheromone[a][b] *= pheromone_evaporation
+    for ant_route in ants_route:
+        total_distance = 0.0
+        for i in range(len(ant_route) - 1):  # Calculate total distance
+            total_distance += distance(cities[ant_route[i]], cities[ant_route[i + 1]])
+        delta_pheromone = q / total_distance
+        for i in range(len(ant_route) - 1):  # Update pheromones
+            pheromone[ant_route[i]][ant_route[i + 1]] += delta_pheromone
+            pheromone[ant_route[i + 1]][ant_route[i]] = pheromone[ant_route[i]][
+                ant_route[i + 1]
+            ]
+
+        if total_distance < best_distance:
+            best_path = ant_route
+            best_distance = total_distance
+
+    return pheromone, best_path, best_distance
+
+
+def city_select(
+    pheromone: list[list[float]],
+    current_city: dict[int, list[int]],
+    unvisited_cities: dict[int, list[int]],
+    alpha: float,
+    beta: float,
+) -> tuple[dict[int, list[int]], dict[int, list[int]]]:
+    """
+    Choose the next city for ants
+    >>> city_select(pheromone=[[1.0, 1.0], [1.0, 1.0]], current_city={0: [0, 0]},
+    ...             unvisited_cities={1: [2, 2]}, alpha=1.0, beta=5.0)
+    ({1: [2, 2]}, {})
+    >>> city_select(pheromone=[], current_city={0: [0,0]},
+    ...             unvisited_cities={1: [2, 2]}, alpha=1.0, beta=5.0)
+    Traceback (most recent call last):
+      ...
+    IndexError: list index out of range
+    >>> city_select(pheromone=[[1.0, 1.0], [1.0, 1.0]], current_city={},
+    ...             unvisited_cities={1: [2, 2]}, alpha=1.0, beta=5.0)
+    Traceback (most recent call last):
+      ...
+    StopIteration
+    >>> city_select(pheromone=[[1.0, 1.0], [1.0, 1.0]], current_city={0: [0, 0]},
+    ...             unvisited_cities={}, alpha=1.0, beta=5.0)
+    Traceback (most recent call last):
+      ...
+    IndexError: list index out of range
+    """
+    probabilities = []
+    for city in unvisited_cities:
+        city_distance = distance(
+            unvisited_cities[city], next(iter(current_city.values()))
+        )
+        probability = (pheromone[city][next(iter(current_city.keys()))] ** alpha) * (
+            (1 / city_distance) ** beta
+        )
+        probabilities.append(probability)
+
+    chosen_city_i = random.choices(
+        list(unvisited_cities.keys()), weights=probabilities
+    )[0]
+    chosen_city = {chosen_city_i: unvisited_cities[chosen_city_i]}
+    del unvisited_cities[next(iter(chosen_city.keys()))]
+    return chosen_city, unvisited_cities
+
+
+if __name__ == "__main__":
+    best_path, best_distance = main(
+        cities=cities,
+        ants_num=10,
+        iterations_num=20,
+        pheromone_evaporation=0.7,
+        alpha=1.0,
+        beta=5.0,
+        q=10,
+    )
+
+    print(f"{best_path = }")
+    print(f"{best_distance = }")