Simulated annealing algorithm

simulated_annealing.py

import math
import random
from functools import partial
from typing import Callable, Sequence


class System:
    def __init__(self) -> None:
        pass

    def generate_initial_solution(self) -> list[float]:
        return [3]

    def energy(self, w: list[float]) -> float:
        return w[0] ** 2


def arithmetic_cooling(
    max_temperature: float, iterations: int, alpha: float = 0.01
) -> float:
    return max_temperature - (iterations * alpha)


def geometric_cooling(
    max_temperature: float, iterations: int, alpha: float = 0.01
) -> float:
    return max_temperature * (alpha**iterations)


class SimulatedAnnealing:
    def __init__(
        self,
        model: System,
        temperatures: Sequence[float] = (0, 300),
        perturbation_alpha: float = 0.1,
        temperature_fraction: float = 1,
        cooling_schedule: Callable = arithmetic_cooling,
    ) -> None:
        self.model = model
        self.min_temperature = temperatures[0]
        self.max_temperature = temperatures[1]
        self.perturbation_alpha = perturbation_alpha
        self.temperature_fraction = temperature_fraction
        self.cooling_schedule = cooling_schedule

    def compute_energy(self, solution: list[float]) -> float:
        return self.model.energy(solution)

    def generate_neighbor(self, solution: list[float]) -> list[float]:
        neighbor = solution[:]
        for i in range(len(solution)):
            neighbor[i] = solution[i] + random.uniform(
                -self.perturbation_alpha, self.perturbation_alpha
            )
        return neighbor

    def metropolis(self, diff_energy: float, temperature: float) -> bool:
        if diff_energy < 0:
            return True
        return random.uniform(0, 1) < math.exp(-diff_energy / temperature)

    def optimize(self) -> list[float]:
        temperature = self.max_temperature
        solution = self.model.generate_initial_solution()
        energy = self.compute_energy(solution)
        iteration = 1
        while temperature > self.min_temperature:
            new_solution = self.generate_neighbor(solution)
            new_energy = self.compute_energy(new_solution)
            diff_energy = new_energy - energy
            if self.metropolis(diff_energy, temperature):
                solution = new_solution[:]
                energy = new_energy
            temperature = self.cooling_schedule(self.max_temperature, iteration)
            print(f"[it={iteration}] energy = {energy}, temperature = {temperature}")
            iteration += 1
        return solution


if __name__ == "__main__":
    model = System()
    sa = SimulatedAnnealing(
        model,
        temperatures=(0, 1),
        cooling_schedule=partial(geometric_cooling, alpha=0.8),
    )
    final_solution: list[float] = sa.optimize()
    print("final_solution", final_solution)