main.py

import numpy as np
from numpy.random import RandomState
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
import random

from genetic import genetic_algorithm
from utils import fitness, load_data

prng = RandomState(1234567)
random.seed(42)


def beam_search(expected_returns, cov_matrix, n_beams=10, n_iterations=10):
    """Beam search optimization algorithm."""
    n_assets = len(expected_returns)
    # Initialize beams
    beams = [prng.dirichlet(np.ones(n_assets)) for _ in range(n_beams)]
    history = []

    for _ in range(n_iterations):
        # Evaluate fitness of each beam
        fitness_scores = [(fitness(beam, expected_returns, cov_matrix), beam) for beam in beams]
        fitness_scores.sort(reverse=True, key=lambda x: x[0])
        beams = [score[1] for score in fitness_scores[:n_beams]]
        history.append(fitness_scores[0][0])

        # Generate new beams by combining existing ones
        new_beams = []
        for i in range(n_beams):
            for j in range(i + 1, n_beams):
                new_beam = (beams[i] + beams[j]) / 2
                new_beams.append(new_beam / new_beam.sum())  # Normalize

        beams.extend(new_beams)
        # beams = beams[:n_beams]  # Keep the best beams

    best_beam = max(beams, key=lambda x: fitness(x, expected_returns, cov_matrix))
    return best_beam, fitness(best_beam, expected_returns, cov_matrix), history


def random_local_beam_search(expected_returns, cov_matrix, n_solutions=50, n_iterations=100, n_best=15):
    """Random local beam search optimization algorithm."""
    n_assets = len(expected_returns)
    solutions = [prng.dirichlet(np.ones(n_assets)) for _ in range(n_solutions)]
    history = []

    for _ in range(n_iterations):
        # Calculate fitness of each solution
        fitness_scores = [(fitness(solution, expected_returns, cov_matrix), solution) for solution in solutions]

        # Select the best solutions
        fitness_scores.sort(reverse=True, key=lambda x: x[0])
        best_solutions = fitness_scores[:n_best]
        history.append(best_solutions[0][0])

        # Generate new solutions from the best ones by random mutation
        new_solutions = []
        for score, solution in best_solutions:
            for _ in range(n_solutions // n_best):
                mutation = solution + prng.normal(0, 0.05, n_assets)
                mutation = np.clip(mutation, 0, 1)  # Ensure weights are between 0 and 1
                mutation /= mutation.sum()  # Normalize to sum to 1
                new_solutions.append(mutation)

        solutions = new_solutions  # Update solutions

    best_solution = max(solutions, key=lambda x: fitness(x, expected_returns, cov_matrix))
    best_fitness = fitness(best_solution, expected_returns, cov_matrix)
    history.append(best_fitness)
    return best_solution, best_fitness, history


def simulated_annealing(expected_returns, cov_matrix, n_iterations=1000, temperature=1000000, cooling_rate=0.95):
    """Simulated annealing optimization algorithm."""
    n_assets = len(expected_returns)
    solution = prng.dirichlet(np.ones(n_assets))
    best_solution = solution
    best_fitness = fitness(solution, expected_returns, cov_matrix)
    history = [best_fitness]

    for i in range(n_iterations):
        # Create a new solution by small mutation
        new_solution = solution + prng.normal(0, 0.1, n_assets)
        new_solution = np.clip(new_solution, 0, 1)
        new_solution /= new_solution.sum()

        new_fitness = fitness(new_solution, expected_returns, cov_matrix)
        history.append(new_fitness)

        # Decide whether to accept the new solution
        if new_fitness > best_fitness or np.exp((new_fitness - best_fitness) / temperature) > random.random():
            solution = new_solution
            if new_fitness > best_fitness:
                best_solution, best_fitness = new_solution, new_fitness

        # Cool down
        temperature *= cooling_rate

    return best_solution, best_fitness, history


def plot_results(
    history_bs, best_fitness_bs, history_rlb, best_fitness_rlb, history_sa, best_fitness_sa, history_ga, best_fitness_ga
):
    """Plot the results of all optimization algorithms."""
    fig = plt.figure(figsize=(12, 12))
    gs = GridSpec(2, 2, figure=fig, hspace=0.4)  # Adjust hspace for row margin

    ax1 = fig.add_subplot(gs[0, 0])
    ax1.plot(history_bs, label="Beam Search")
    ax1.set_title(f"Beam Search (Heuristic) - Solution Fitness: {best_fitness_bs:.2f}")
    ax1.set_xlabel("Iteration")
    ax1.set_ylabel("Fitness")
    ax1.legend()
    ax1.grid()

    ax2 = fig.add_subplot(gs[0, 1])
    ax2.plot(
        history_rlb,
        label="Random Local Beam Search",
        color="orange",
    )
    ax2.set_title(f"Random Local Beam Search (Heuristic) - Solution Fitness: {best_fitness_rlb:.2f}")
    ax2.set_xlabel("Iteration")
    ax2.set_ylabel("Fitness")
    ax2.legend()
    ax2.grid()

    ax3 = fig.add_subplot(gs[1, 0])
    ax3.plot(
        history_sa,
        label="Simulated Annealing",
        color="green",
    )
    ax3.set_title(f"Simulated Annealing (Meta-Heuristic) - Solution Fitness: {best_fitness_sa:.2f}")
    ax3.set_xlabel("Iteration")
    ax3.set_ylabel("Fitness")
    ax3.legend()
    ax3.grid()

    ax4 = fig.add_subplot(gs[1, 1])
    ax4.plot(history_ga, label="Genetic Algorithm", color="red")
    ax4.set_title(f"Genetic Algorithm (Meta-Heuristic) - Solution Fitness: {best_fitness_ga:.2f}")
    ax4.set_xlabel("Iteration")
    ax4.set_ylabel("Fitness")
    ax4.legend()
    ax4.grid()

    plt.tight_layout()
    plt.show()


def main():
    expected_returns, cov_matrix = load_data()

    print("Running Portfolio Optimization Algorithms...")
    print(f"Number of assets: {len(expected_returns)}")
    print()

    print("Running Beam Search...")
    best_solution_bs, best_fitness_bs, history_bs = beam_search(expected_returns, cov_matrix)
    print("Beam Search Best Solution:", best_solution_bs)
    print("Best Fitness (Return/Risk):", best_fitness_bs)
    print()

    print("Running Random Local Beam Search...")
    best_solution_rlb, best_fitness_rlb, history_rlb = random_local_beam_search(expected_returns, cov_matrix)
    print("Random Local Beam Search Best Solution:", best_solution_rlb)
    print("Best Fitness (Return/Risk):", best_fitness_rlb)
    print()

    print("Running Simulated Annealing...")
    best_solution_sa, best_fitness_sa, history_sa = simulated_annealing(expected_returns, cov_matrix)
    print("Simulated Annealing Best Solution:", best_solution_sa)
    print("Best Fitness (Return/Risk):", best_fitness_sa)
    print()

    print("Running Genetic Algorithm...")
    best_solution_ga, best_fitness_ga, history_ga = genetic_algorithm(expected_returns, cov_matrix)
    print("Genetic Algorithm Best Solution:", best_solution_ga)
    print("Best Fitness (Return/Risk):", best_fitness_ga)
    print()

    plot_results(
        history_bs,
        best_fitness_bs,
        history_rlb,
        best_fitness_rlb,
        history_sa,
        best_fitness_sa,
        history_ga,
        best_fitness_ga,
    )


if __name__ == "__main__":
    main()