main.py

from py5 import Sketch
import numpy as np
import csv

"""
Traveling Salesman Problem (TSP) solver using genetic algorithms.

Applies concepts from genetic algorithms to find an approximate solution to the TSP in a shortened amount of time.
"""

# Problem Parameters
NUM_CITIES = 20 # The number of cities

# Genetic Algorithm Parameters
POPULATION_SIZE = 200 # Number of individuals in the population
MUTATION_RATE = 0.05 # Probability of mutation
CROSSOVER_RATE = 0.20 # Probability of crossover
NUM_ELITES = 20
MAX_GENERATIONS = 1000  # Maximum number of generations to run

CROSSOVER_METHOD = 'obx'

OUTPUT_FILE_NAME = "data-obx-case.csv"

def obx_crossover(parent1, parent2):
    """Order Based Crossover (OBX)"""
    # Select three random nodes for crossover
    nodes = sorted(np.random.choice(parent1, 3, replace=False))

    order_in_parent1 = []
    for node in parent1:
        if node in nodes:
            order_in_parent1.append(node)

    order_in_parent2 = []
    for node in parent2:
        if node in nodes:
            order_in_parent2.append(node)

    child1 = parent1.copy()
    for i, node in enumerate(child1):
        if node in nodes:
            child1[i] = order_in_parent2.pop(0)

    child2 = parent2.copy()
    for i, node in enumerate(child2):
        if node in nodes:
            child2[i] = order_in_parent1.pop(0)

    return child1, child2


def ox_crossover(parent1, parent2):
    """Order Crossover (OX)"""
    # Choose two crossover points
    x_points = sorted(np.random.choice(len(parent1), 2))

    # Create children with empty contents
    child1 = [-1] * len(parent1)
    child2 = [-1] * len(parent2)

    # Copy parent 1 substring to child 1
    child1[x_points[0]:x_points[1]] = parent1[x_points[0]:x_points[1]]

    # Copy parent 2 substring to child 2
    child2[x_points[0]:x_points[1]] = parent2[x_points[0]:x_points[1]]

    # Copy remaining nodes from parent 2 to child 1
    i = x_points[1] # index for child 1
    for node in parent2[x_points[1]:] + parent2[:x_points[1]]:
        if node not in child1:
            child1[i] = node
            i = (i + 1) % len(child1)

    # Copy remaining nodes from parent 1 to child 2
    i = x_points[1] # index for child2
    for node in parent1[x_points[1]:] + parent1[:x_points[1]]:
        if node not in child2:
            child2[i] = node
            i = (i + 1) % len(child2)

    return child1, child2


class GeneticAlgorithm:
    """Genetic Algorithm for TSP."""

    def __init__(self, cities, population):
        self.cities = cities
        self.population = population

    def next_generation(self, crossover_strategy):
        # Step #1: Evaluate fitness
        self.evaluate_fitness()

        # Step #2: Reproduce
        self.reproduce(crossover_strategy)
    
    def reproduce(self, crossover_strategy):
        new_population = []

        # Elites selection
        items = [{'idx': i, 'fitness': f} for i, f in enumerate(self.fitness_scores)]
        items.sort(key=lambda x: x['fitness'], reverse=True)
        for i in range(NUM_ELITES):
            idx = items[i]['idx']
            new_population.append(self.population[idx])

        # Proportional selection
        fitness_sum = sum(self.fitness_scores)
        probabilities = [score / fitness_sum for score in self.fitness_scores]
        while len(new_population) < POPULATION_SIZE:
            # Select parents based on fitness probabilities
            parent1 = self.population[np.random.choice(len(self.population), p=probabilities)]
            parent2 = self.population[np.random.choice(len(self.population), p=probabilities)]

            # Crossover
            if np.random.rand() < CROSSOVER_RATE:
                child1, child2 = crossover_strategy(parent1, parent2)
            else:
                child1, child2 = parent1, parent2 # Skip crossover

            # Mutation
            child1 = self.mutate(child1)
            child2 = self.mutate(child2)

            # Add children to new population
            new_population.append(child1)
            new_population.append(child2)
        self.population = new_population

    def evaluate_fitness(self):
        self.fitness_scores = []
        self.lengths = []
        for individual in self.population:
            total_distance = self.total_distance(individual)
            self.lengths.append(total_distance)
            self.fitness_scores.append(1.0 / (total_distance / 1000.0))

    def get_individual_info(self, idx: int):
        return {
            'genome': self.population[idx],
            'fitness': self.fitness_scores[idx],
            'length': self.lengths[idx],
        }

    def mutate(self, individual):
        if np.random.rand() < MUTATION_RATE:
            if np.random.rand() < 0.5:
                return self.swap_mutation(individual)
            else:
                return self.remove_and_insert_mutation(individual)

        return individual

    def swap_mutation(self, individual):
        indices = np.random.choice(len(individual), 2, replace=False)
        temp = individual[indices[0]]
        individual[indices[0]] = individual[indices[1]]
        individual[indices[1]] = temp
        return individual

    def remove_and_insert_mutation(self, individual):
        index = np.random.choice(range(len(individual)))
        node = individual.pop(index)
        insert_index = np.random.choice(range(len(individual)))
        individual.insert(insert_index, node)
        return individual

    def distance(self, p1, p2):
        return np.linalg.norm(np.array(p1) - np.array(p2))

    def total_distance(self, individual):
        total_distance = 0
        for i in range(len(individual)):
            p1 = self.cities[individual[i]]
            p2 = self.cities[individual[(i + 1) % len(individual)]]
            total_distance += self.distance(p1, p2)
        return total_distance


class TSPSolver(Sketch):
    """TSP solver using genetic algorithm."""

    def __init__(self, file):
        super().__init__()
        self.file = file
        self.writer = csv.writer(file)
        self.writer.writerow([
            'Generation',
            'Best Fitness',
            'Average Fitness',
            'Worst Fitness',
            'Diversity Proxy',
            'Best Length',
            'Best Genome',
            'Fitness Delta',
            'Length Delta'
        ])
        self.initial_best = None
        self.curr_best = None
        self.overall_best = None
        self.curr_worst = None
        self.draw_best_overall = False

    def settings(self):
        self.size(800, 600)

    def setup(self):
        # Generate random cities
        # cities = np.random.rand(NUM_CITIES, 2) * [self.width, self.height]
        # with open('input.txt', 'w') as f:
        #     for city in cities:
        #         f.write(f"{city[0]} {city[1]}\n")

        cities = []
        with open('input.txt', 'r') as f:
            for line in f:
                values = line.split()
                x = np.float64(values[0])
                y = np.float64(values[1])
                cities.append([x, y])

        # Generate random population
        population = []
        for _ in range(POPULATION_SIZE):
            individual = np.random.permutation(len(cities)).tolist()
            population.append(individual)

        self.generations = 0
        
        # Initialize genetic algorithm
        self.ga = GeneticAlgorithm(cities, population)

    def draw(self):
        if self.generations < MAX_GENERATIONS:
            # Select crossover method
            if CROSSOVER_METHOD == 'obx':
                crossover = obx_crossover
            elif CROSSOVER_METHOD == 'ox':
                crossover = ox_crossover
            
            # Evolve to next generation
            self.ga.next_generation(crossover)
            self.generations += 1

            # Get info about best individual
            best_idx = np.argmax(self.ga.fitness_scores)
            self.curr_best = self.ga.get_individual_info(best_idx)
            if self.initial_best is None:
                self.initial_best = self.curr_best
            if self.overall_best is None or self.curr_best['fitness'] >= self.overall_best['fitness']:
                self.overall_best = self.curr_best

            # Get info about worst individual
            worst_idx = np.argmin(self.ga.fitness_scores)
            self.curr_worst = self.ga.get_individual_info(worst_idx)

            # Calculate some data
            self.avg_fitness = sum(self.ga.fitness_scores) / len(self.ga.fitness_scores)
            self.delta_f = self.curr_best['fitness'] - self.initial_best['fitness']
            self.delta_l = self.curr_best['length'] - self.initial_best['length']
            self.diversity_proxy = self.curr_best['fitness'] - self.curr_worst['fitness']

            # Write data to CSV
            self.writer.writerow([
                self.generations,
                self.curr_best['fitness'],
                self.avg_fitness,
                self.curr_worst['fitness'],
                self.diversity_proxy,
                self.curr_best['length'],
                self.curr_best['genome'],
                self.delta_f,
                self.delta_l,
            ])
        
        # Clear background
        self.background(255)

        # Draw the cities
        self.draw_cities()

        # Draw best individual
        self.stroke(255, 0, 0)
        self.stroke_weight(2)
        
        genome = self.curr_best['genome']
        if self.draw_best_overall:
            genome = self.overall_best['genome']
        
        for i in range(len(genome)):
            p1 = self.ga.cities[genome[i]]
            p2 = self.ga.cities[genome[(i + 1) % len(genome)]]
            self.line(p1[0], p1[1], p2[0], p2[1])

        # Draw data
        self.fill(0)
        self.text_size(16)
        self.text(f"Generation: {self.generations}", 10, 20)
        self.text(f'Best Fitness: {self.curr_best['fitness']:.6f}', 10, 40)
        self.text(f'Overall Best Fitness: {self.overall_best['fitness']:.6f}', 10, 60)

        self.fill(0, 0, 255)
        if self.draw_best_overall:
            self.text("Showing Best Overall", 10, 80)

    def draw_cities(self):
        self.no_stroke()
        self.fill(0)
        for x, y in self.ga.cities:
            self.ellipse(x, y, 8, 8)
    
    def mouse_clicked(self):
        self.draw_best_overall = not self.draw_best_overall


if __name__ == '__main__':
    with open(OUTPUT_FILE_NAME, 'w', newline='') as file:
        TSPSolver(file).run_sketch()