ACO_timeAware.py

import numpy as np
import random
from typing import Any, List
from tdp import TravelingDeliverymanProblem
from visualizer import Visualizer

#Section IV-A of Dorigo, M., Maniezzo, V., & Colorni, A. (1996). "Ant system: optimization by a colony of cooperating agents." IEEE Transactions on Systems, Man, and Cybernetics, Part B (Cybernetics), 26(1), 29-41.

def dist_matrix(tdp_obj, i, j):
    """
    Function to "simulate" the distance matrix lookup. Needed for fair comparison with other
    algorithm that don't utilize a cache/dist-matrix

    :param tdp_obj: TravelingDeliverymanProblem
    :param i: index of the node 1 (not node id)
    :param j: index of the node 2 (not node id)
    """
    if i != j:
        d = tdp_obj.distance(tdp_obj.node_id_by_idx(i), tdp_obj.node_id_by_idx(j))
        return d
    else:
        return 0

def ant_colony_optimization_time_aware(
        tdp: TravelingDeliverymanProblem,
        n_ants: int = 15,
        alpha: float = 1.0,
        beta: float = 4.0,
        rho: float = 0.1,
        q: float = 100.0,
        verbose: bool = True,
        seed: int = 42,
        max_fe: int = 10000
) -> tuple[list[int] | Any, float | Any]:
    """
    Constraint-Aware Ant Colony Optimization (ACO) for TDP with Deadlines.
    This version do not use an dist_matrix cache, because this would be an unfair cost advantage over the other non-cached algorithms.

    Key Features:
    - Dynamic Heuristic: Combines Distance and 'Slack Time' (Urgency).
    - Lookahead Filter: Prevents ants from choosing nodes that violate deadlines immediately.

    Implementation is based on:
    M. Dorigo, V. Maniezzo, and A. Colorni, “Ant System: Optimization by a colony of cooperating agents. IEEE Trans Syst Man Cybernetics - Part B,” IEEE transactions on systems, man, and cybernetics. Part B, Cybernetics : a publication of the IEEE Systems, Man, and Cybernetics Society, vol. 26, pp. 29–41, Feb. 1996, doi: 10.1109/3477.484436.

    :param tdp: TravelingDeliverymanProblem
    :param n_ants: number of ants to simulate
    :param alpha: how much the pheromone (experience) influences the movements of the ants
    :param beta: how much the heuristic (reverse-distance) influences the movement of the ants#
    :param rho: Evaporation
    :param q: q value, refer to cited paper for detailed explanation
    :param verbose: Print logs
    :param seed: random seed
    :param max_fe: maximum number of FE-checks (FE = Function Evaluations, aka checks of a full tour)

    :return: tuple(best_solution, best_fitness) // best solution is a list-object representing the different tour indices, refer to tdp class.
    """
    random.seed(seed)
    np.random.seed(seed)

    n_cities = tdp.dimension
    start_node_id = tdp.start_node
    start_node_idx = tdp.idx_by_node_id(start_node_id)

    # =========================================================================
    # 1. Initialization
    # =========================================================================
    pheromone = np.ones((n_cities, n_cities)) * 0.1

    # Global Best
    best_solution = None
    best_fitness = -float('inf')

    # Initial FE check
    if tdp.get_num_fe() >= max_fe:
        dummy = tdp.get_unique_tours(1)[0]
        return dummy, tdp.objective_func(dummy)

    if verbose:
        print(f"ACO initialized. Ants: {n_ants}, Alpha: {alpha}, Beta: {beta}")

    # =========================================================================
    # Main Loop
    # =========================================================================
    iteration = 0
    while tdp.get_num_fe() < max_fe:

        ant_tours = []

        # ---------------------------------------------------------------------
        # Step 1: Construct Solutions (with Time Awareness)
        # ---------------------------------------------------------------------
        for _ in range(n_ants):
            if tdp.get_num_fe() >= max_fe: break

            # --- Initialize Ant ---
            current_idx = start_node_idx
            tour = [current_idx]
            unvisited = list(range(n_cities))
            unvisited.remove(current_idx)

            # Track Time for this specific ant
            # Start node has deadline too, but we assume we start at t=0
            current_time = 0.0
            valid_tour = True

            while unvisited:
                candidates = np.array(unvisited)

                # --- A. Calculate Travel Times & Deadlines ---
                # Get distances to all candidates
                dists = np.array([dist_matrix(tdp, current_idx, candidate) for candidate in candidates])

                # Convert Distance to Time (Minutes)
                # Formula from your TDP class: (km / speed) * 60
                travel_times = (dists / tdp.recovery_speed) * 60.0

                arrival_times = current_time + travel_times
                candidate_deadlines = np.array([tdp.get_deadline_by_id(tdp.node_id_by_idx(i)) or float('inf') for i in candidates])

                # --- B. Feasibility Filter (The "Lookahead") ---
                # Only keep candidates where arrival BEFORE the deadline
                feasible_mask = arrival_times <= candidate_deadlines

                if not np.any(feasible_mask):
                    # DEAD END: No reachable cities left.
                    # This ant dies. We stop building this tour.
                    valid_tour = False
                    break

                # Filter our arrays to keep only feasible options
                valid_candidates = candidates[feasible_mask]
                valid_dists = dists[feasible_mask]
                valid_arrival_times = arrival_times[feasible_mask]
                valid_deadlines = candidate_deadlines[feasible_mask]

                # --- C. Dynamic Heuristic (Urgency) ---
                # Slack = Time left until deadline after arrival
                # We add epsilon to avoid division by zero
                slack = valid_deadlines - valid_arrival_times

                # Heuristic: Favor Short Distance AND Low Slack (Urgent)
                # eta = (1/dist) * (1/(slack + 1))
                # Note: We prioritize Urgency heavily
                eta_values = (1.0 / (valid_dists + 1e-6)) * (1.0 / (slack + 1.0))

                # --- D. Probability & Selection ---
                p_vals = pheromone[current_idx, valid_candidates]

                scores = (p_vals ** alpha) * (eta_values ** beta)

                if scores.sum() == 0:
                    probs = np.ones(len(scores)) / len(scores)
                else:
                    probs = scores / scores.sum()

                next_city = np.random.choice(valid_candidates, p=probs)

                # Update State
                tour.append(next_city)
                unvisited.remove(next_city)

                # Update Time (Find index of selected city in the valid arrays to get exact time)
                # Optimization: We re-calculate simple edge weight to be safe/clean
                time_step = (dist_matrix(tdp, current_idx, next_city) / tdp.recovery_speed) * 60.0
                current_time += time_step
                current_idx = next_city

            # Only save the tour if it visited ALL cities (valid complete tour)
            if valid_tour and len(tour) == n_cities:
                ant_tours.append(tour)

        # ---------------------------------------------------------------------
        # Step 2: Evaluation & Ant-Level Pheromone Deposit
        # ---------------------------------------------------------------------
        # We only evaluate tours that survived the construction phase
        if not ant_tours:
            # If ALL ants died, we might need to reduce Beta or check if graph is solvable
            # But we just continue to next iter (pheromone evaporation might help exploration)
            pass

        pheromone *= (1.0 - rho)  # Evaporation

        iter_best_solution = None
        iter_best_fitness = -float('inf')

        for tour in ant_tours:
            if tdp.get_num_fe() < max_fe:
                fit = tdp.objective_func(tour)

                # Ant-level pheromone deposit (proportional to fitness)
                deposit = q * fit if fit > 0 else q * 0.01
                for i in range(len(tour) - 1):
                    u, v = tour[i], tour[i + 1]
                    pheromone[u, v] += deposit
                    pheromone[v, u] += deposit

                # Wrap-around edge
                u, v = tour[-1], tour[0]
                pheromone[u, v] += deposit
                pheromone[v, u] += deposit

                # Track iteration best
                if fit > iter_best_fitness:
                    iter_best_fitness = fit
                    iter_best_solution = tour

                # Track global best
                if fit > best_fitness:
                    best_fitness = fit
                    best_solution = list(tour)
            else:
                break

        # ---------------------------------------------------------------------
        # Step 3: Elitist Pheromone Boost
        # ---------------------------------------------------------------------
        # Elitist Update (Global Best gets extra reinforcement)
        if best_solution is not None:
            elitist_deposit = q * best_fitness
            for i in range(len(best_solution) - 1):
                u, v = best_solution[i], best_solution[i + 1]
                pheromone[u, v] += elitist_deposit
                pheromone[v, u] += elitist_deposit

            u, v = best_solution[-1], best_solution[0]
            pheromone[u, v] += elitist_deposit
            pheromone[v, u] += elitist_deposit

        iteration += 1
        if verbose and iteration % 10 == 0:
            print(f"Iteration {iteration}: Best Fitness = {best_fitness:.2f}")

    if verbose:
        print(f"\nFinal Best Fitness: {best_fitness:.2f}")

    return best_solution, best_fitness


if __name__ == "__main__":
    DATASET = "datasets/custom/kochi_small.tdp"
    #DATASET = "datasets/SolomonPotvinBengio_tdp/rc_202.2.tdp"
    tdp = TravelingDeliverymanProblem(DATASET, "ACO (time)")

    best_tour, best_obj_val = ant_colony_optimization_time_aware(
        tdp,
        n_ants=15,
        alpha=1.0,
        beta=3.0,
        rho=0.1,
        max_fe=10000
    )

    visualizer = Visualizer(tdp)
    visualizer.show_convergence_plot()
    visualizer.save_convergence_plot()
    visualizer.show_solution_plot()
    visualizer.save_solution_plot()

    print(f"\nBest tour found: {best_tour}")
    print(f"Best fitness: {best_obj_val:.2f}")