genetic_algorithm.py

"""
Genetic Algorithm for Neural Networks
=====================================

This module extends the training framework to support genetic algorithms for
neural network evolution. It manages populations, fitness evaluation, selection,
mutation, and crossover operations.

Usage:
    python genetic_algorithm.py config_genetic.json
"""

import sys
import os
import json
import numpy as np
import torch
import torch.nn as nn
import copy
import random
from pathlib import Path
import argparse

# Import the train module and run_train utilities
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import train
import run_train

# ============================================================================
# GENETIC ALGORITHM CONFIGURATION
# ============================================================================

# Population Configuration
POPULATION_SIZE = 100          # Number of models in each generation
NUM_GENERATIONS = 50           # Number of generations to evolve
ELITE_COUNT = 10               # Number of best models to keep unchanged
MUTATION_RATE = 0.1            # Probability of mutating each weight
MUTATION_STRENGTH = 0.1        # Standard deviation for weight mutations
CROSSOVER_RATE = 0.7           # Probability of crossover vs mutation
SELECTION_METHOD = "tournament"  # "tournament", "roulette", "rank"
TOURNAMENT_SIZE = 5           # Size of tournament for selection
FITNESS_METRIC = "R2"          # Metric to use as fitness (R2, loss, MAE, etc.)

# Generation Tracking
GENERATION_SAVE_PATH = "generations"  # Where to save generation data
SAVE_BEST_N = 5               # Save top N models from each generation
SAVE_ALL = False               # Save all models (can be very large!)

# Training Configuration
TRAIN_BEFORE_EVALUATION = False  # If True, train each model before evaluating fitness
TRAINING_EPOCHS = 10             # Number of epochs to train each model (if TRAIN_BEFORE_EVALUATION=True)

# ============================================================================
# GENETIC OPERATIONS
# ============================================================================

def mutate_weights(model, mutation_rate=MUTATION_RATE, mutation_strength=MUTATION_STRENGTH):
    """
    Mutate model weights by adding random noise.
    
    Args:
        model: PyTorch model to mutate
        mutation_rate: Probability of mutating each parameter
        mutation_strength: Standard deviation of mutation noise
    
    Returns:
        Mutated model (new instance)
    """
    mutated_model = copy.deepcopy(model)
    
    with torch.no_grad():
        for param in mutated_model.parameters():
            # Create mutation mask
            mutation_mask = torch.rand_like(param) < mutation_rate
            # Generate random mutations
            mutations = torch.randn_like(param) * mutation_strength
            # Apply mutations
            param.data += mutations * mutation_mask.float()
    
    return mutated_model

def crossover_weights(parent1, parent2):
    """
    Create offspring by combining weights from two parent models.
    
    Args:
        parent1: First parent model
        parent2: Second parent model
    
    Returns:
        Offspring model (new instance)
    """
    offspring = copy.deepcopy(parent1)
    
    with torch.no_grad():
        for (name1, param1), (name2, param2), (name_off, param_off) in zip(
            parent1.named_parameters(),
            parent2.named_parameters(),
            offspring.named_parameters()
        ):
            # Create random crossover mask (50/50 split)
            crossover_mask = torch.rand_like(param1) < 0.5
            # Combine weights from both parents
            param_off.data = param1.data * crossover_mask.float() + param2.data * (1 - crossover_mask.float())
    
    return offspring

def initialize_population(population_size=POPULATION_SIZE):
    """
    Initialize a population of random models.
    
    Args:
        population_size: Number of models to create
    
    Returns:
        List of model instances
    """
    population = []
    for i in range(population_size):
        model = train.create_model()
        population.append(model)
    return population

def evaluate_fitness(model, train_loader, test_loader, criterion, fitness_metric=FITNESS_METRIC,
                    train_before_eval=False, training_epochs=10):
    """
    Evaluate model fitness on data.
    
    Args:
        model: Model to evaluate
        train_loader: DataLoader with training data (for training if train_before_eval=True)
        test_loader: DataLoader with test data
        criterion: Loss function
        fitness_metric: Metric to use as fitness
        train_before_eval: If True, train the model before evaluating
        training_epochs: Number of epochs to train (if train_before_eval=True)
    
    Returns:
        Fitness score (higher is better)
    """
    # Optionally train the model first
    if train_before_eval:
        optimizer = train.get_optimizer(model)
        for epoch in range(training_epochs):
            train.train_epoch(model, train_loader, criterion, optimizer)
    
    # Evaluate on test data
    model.eval()
    all_predictions = []
    all_targets = []
    
    with torch.no_grad():
        for batch_X, batch_y in test_loader:
            batch_X = batch_X.to(train.DEVICE)
            batch_y = batch_y.to(train.DEVICE)
            predictions = model(batch_X)
            all_predictions.append(predictions.cpu())
            all_targets.append(batch_y.cpu())
    
    all_predictions = torch.cat(all_predictions, dim=0)
    all_targets = torch.cat(all_targets, dim=0)
    
    # Calculate fitness based on specified metric
    if fitness_metric == "R2":
        ss_res = torch.sum((all_targets - all_predictions) ** 2).item()
        ss_tot = torch.sum((all_targets - torch.mean(all_targets)) ** 2).item()
        fitness = 1 - (ss_res / (ss_tot + 1e-8))
    elif fitness_metric == "loss" or fitness_metric == "MSE":
        fitness = -criterion(all_predictions, all_targets).item()  # Negative because lower is better
    elif fitness_metric == "MAE":
        fitness = -torch.mean(torch.abs(all_predictions - all_targets)).item()  # Negative
    else:
        # Default: use R2
        ss_res = torch.sum((all_targets - all_predictions) ** 2).item()
        ss_tot = torch.sum((all_targets - torch.mean(all_targets)) ** 2).item()
        fitness = 1 - (ss_res / (ss_tot + 1e-8))
    
    return fitness

def select_parents(population, fitness_scores, selection_method=SELECTION_METHOD, tournament_size=TOURNAMENT_SIZE):
    """
    Select parents for reproduction.
    
    Args:
        population: List of models
        fitness_scores: List of fitness scores
        selection_method: "tournament", "roulette", or "rank"
        tournament_size: Size of tournament (for tournament selection)
    
    Returns:
        Selected parent model
    """
    if selection_method == "tournament":
        # Tournament selection
        tournament_indices = random.sample(range(len(population)), min(tournament_size, len(population)))
        tournament_fitness = [fitness_scores[i] for i in tournament_indices]
        winner_idx = tournament_indices[np.argmax(tournament_fitness)]
        return population[winner_idx]
    
    elif selection_method == "roulette":
        # Roulette wheel selection (proportional to fitness)
        # Normalize fitness scores to be positive
        min_fitness = min(fitness_scores)
        normalized_fitness = [f - min_fitness + 1e-8 for f in fitness_scores]
        total_fitness = sum(normalized_fitness)
        probabilities = [f / total_fitness for f in normalized_fitness]
        selected_idx = np.random.choice(len(population), p=probabilities)
        return population[selected_idx]
    
    elif selection_method == "rank":
        # Rank-based selection
        sorted_indices = sorted(range(len(fitness_scores)), key=lambda i: fitness_scores[i], reverse=True)
        ranks = [0] * len(population)
        for rank, idx in enumerate(sorted_indices):
            ranks[idx] = len(population) - rank
        total_rank = sum(ranks)
        probabilities = [r / total_rank for r in ranks]
        selected_idx = np.random.choice(len(population), p=probabilities)
        return population[selected_idx]
    
    else:
        # Default: random selection
        return random.choice(population)

def evolve_generation(population, fitness_scores, elite_count=ELITE_COUNT, 
                     crossover_rate=CROSSOVER_RATE, mutation_rate=MUTATION_RATE,
                     mutation_strength=MUTATION_STRENGTH):
    """
    Create next generation through selection, crossover, and mutation.
    
    Args:
        population: Current generation of models
        fitness_scores: Fitness scores for each model
        elite_count: Number of best models to keep unchanged
        crossover_rate: Probability of crossover vs mutation
        mutation_rate: Probability of mutation
        mutation_strength: Strength of mutations
    
    Returns:
        New generation of models
    """
    # Sort by fitness (best first)
    sorted_indices = sorted(range(len(fitness_scores)), key=lambda i: fitness_scores[i], reverse=True)
    
    new_population = []
    
    # Keep elite models unchanged
    for i in range(elite_count):
        elite_idx = sorted_indices[i]
        new_population.append(copy.deepcopy(population[elite_idx]))
    
    # Generate rest of population through reproduction
    while len(new_population) < len(population):
        if random.random() < crossover_rate and len(population) >= 2:
            # Crossover: select two parents and create offspring
            parent1 = select_parents(population, fitness_scores)
            parent2 = select_parents(population, fitness_scores)
            offspring = crossover_weights(parent1, parent2)
            # Mutate the offspring
            offspring = mutate_weights(offspring, mutation_rate, mutation_strength)
            new_population.append(offspring)
        else:
            # Mutation only: select one parent and mutate
            parent = select_parents(population, fitness_scores)
            offspring = mutate_weights(parent, mutation_rate, mutation_strength)
            new_population.append(offspring)
    
    return new_population

def save_generation(generation_num, population, fitness_scores, generation_path=GENERATION_SAVE_PATH):
    """
    Save generation data and top models.
    
    Args:
        generation_num: Generation number
        population: List of models
        fitness_scores: List of fitness scores
        generation_path: Path to save generation data
    """
    Path(generation_path).mkdir(parents=True, exist_ok=True)
    
    # Sort by fitness
    sorted_indices = sorted(range(len(fitness_scores)), key=lambda i: fitness_scores[i], reverse=True)
    
    # Save generation statistics
    stats = {
        'generation': generation_num,
        'population_size': len(population),
        'best_fitness': fitness_scores[sorted_indices[0]],
        'worst_fitness': fitness_scores[sorted_indices[-1]],
        'mean_fitness': np.mean(fitness_scores),
        'std_fitness': np.std(fitness_scores),
        'top_n_indices': sorted_indices[:SAVE_BEST_N]
    }
    
    stats_path = os.path.join(generation_path, f"generation_{generation_num}_stats.json")
    with open(stats_path, 'w') as f:
        json.dump(stats, f, indent=2)
    
    # Save top N models
    for i, idx in enumerate(sorted_indices[:SAVE_BEST_N]):
        model = population[idx]
        fitness = fitness_scores[idx]
        model_name = f"gen{generation_num}_rank{i+1}_fitness{fitness:.6f}"
        
        # Temporarily set model name for saving
        original_name = train.MODEL_NAME
        train.MODEL_NAME = model_name
        train.save_model(model, generation_num, fitness, {'fitness': fitness}, suffix="")
        train.MODEL_NAME = original_name
    
    # Save all models if requested
    if SAVE_ALL:
        for i, (model, fitness) in enumerate(zip(population, fitness_scores)):
            model_name = f"gen{generation_num}_model{i}_fitness{fitness:.6f}"
            original_name = train.MODEL_NAME
            train.MODEL_NAME = model_name
            train.save_model(model, generation_num, fitness, {'fitness': fitness}, suffix="")
            train.MODEL_NAME = original_name

# ============================================================================
# GENETIC ALGORITHM MAIN LOOP
# ============================================================================

def run_genetic_algorithm(config_path):
    """
    Run genetic algorithm for neural network evolution.
    
    Args:
        config_path: Path to JSON configuration file
    """
    print("="*60)
    print("Genetic Algorithm for Neural Networks")
    print("="*60)
    
    # Load base configuration
    with open(config_path, 'r') as f:
        config = json.load(f)
    
    # Load genetic algorithm specific config
    ga_config = config.get('genetic_algorithm', {})
    global POPULATION_SIZE, NUM_GENERATIONS, ELITE_COUNT, MUTATION_RATE
    global MUTATION_STRENGTH, CROSSOVER_RATE, SELECTION_METHOD, TOURNAMENT_SIZE
    global FITNESS_METRIC, GENERATION_SAVE_PATH, SAVE_BEST_N, SAVE_ALL
    global TRAIN_BEFORE_EVALUATION, TRAINING_EPOCHS
    
    POPULATION_SIZE = ga_config.get('population_size', POPULATION_SIZE)
    NUM_GENERATIONS = ga_config.get('num_generations', NUM_GENERATIONS)
    ELITE_COUNT = ga_config.get('elite_count', ELITE_COUNT)
    MUTATION_RATE = ga_config.get('mutation_rate', MUTATION_RATE)
    MUTATION_STRENGTH = ga_config.get('mutation_strength', MUTATION_STRENGTH)
    CROSSOVER_RATE = ga_config.get('crossover_rate', CROSSOVER_RATE)
    SELECTION_METHOD = ga_config.get('selection_method', SELECTION_METHOD)
    TOURNAMENT_SIZE = ga_config.get('tournament_size', TOURNAMENT_SIZE)
    FITNESS_METRIC = ga_config.get('fitness_metric', FITNESS_METRIC)
    GENERATION_SAVE_PATH = ga_config.get('generation_save_path', GENERATION_SAVE_PATH)
    SAVE_BEST_N = ga_config.get('save_best_n', SAVE_BEST_N)
    SAVE_ALL = ga_config.get('save_all', SAVE_ALL)
    TRAIN_BEFORE_EVALUATION = ga_config.get('train_before_evaluation', False)
    TRAINING_EPOCHS = ga_config.get('training_epochs', 10)
    
    # Apply base configuration to train module
    run_train.apply_config_to_train(config)
    
    # Setup
    train.setup_seed(train.SEED)
    train.setup_directories()
    logger = train.setup_logging()
    
    logger.info("="*60)
    logger.info("Genetic Algorithm Started")
    logger.info("="*60)
    logger.info(f"Population Size: {POPULATION_SIZE}")
    logger.info(f"Generations: {NUM_GENERATIONS}")
    logger.info(f"Elite Count: {ELITE_COUNT}")
    logger.info(f"Fitness Metric: {FITNESS_METRIC}")
    
    # Load data
    print("\nLoading data...")
    X, y = run_train.load_data_from_config(config)
    train_loader, test_loader, _ = train.prepare_data(X, y)
    
    # Setup loss function
    criterion = train.get_loss_function()
    
    # Initialize population
    print(f"\nInitializing population of {POPULATION_SIZE} models...")
    population = initialize_population(POPULATION_SIZE)
    
    # Evolution loop
    best_fitness_history = []
    mean_fitness_history = []
    
    for generation in range(NUM_GENERATIONS):
        print(f"\n{'='*60}")
        print(f"Generation {generation + 1}/{NUM_GENERATIONS}")
        print(f"{'='*60}")
        
        # Evaluate fitness for all models
        print("Evaluating fitness...")
        fitness_scores = []
        for i, model in enumerate(population):
            if (i + 1) % 10 == 0:
                print(f"  {'Training and ' if TRAIN_BEFORE_EVALUATION else ''}Evaluating model {i+1}/{POPULATION_SIZE}...")
            fitness = evaluate_fitness(
                model, train_loader, test_loader, criterion, FITNESS_METRIC,
                TRAIN_BEFORE_EVALUATION, TRAINING_EPOCHS
            )
            fitness_scores.append(fitness)
        
        # Statistics
        best_fitness = max(fitness_scores)
        mean_fitness = np.mean(fitness_scores)
        best_fitness_history.append(best_fitness)
        mean_fitness_history.append(mean_fitness)
        
        print(f"\nGeneration {generation + 1} Results:")
        print(f"  Best Fitness: {best_fitness:.6f}")
        print(f"  Mean Fitness: {mean_fitness:.6f}")
        print(f"  Worst Fitness: {min(fitness_scores):.6f}")
        
        logger.info(f"Generation {generation + 1}: Best={best_fitness:.6f}, Mean={mean_fitness:.6f}")
        
        # Save generation
        save_generation(generation + 1, population, fitness_scores)
        
        # Evolve to next generation (except for last generation)
        if generation < NUM_GENERATIONS - 1:
            print("Evolving to next generation...")
            population = evolve_generation(
                population, fitness_scores, ELITE_COUNT,
                CROSSOVER_RATE, MUTATION_RATE, MUTATION_STRENGTH
            )
    
    # Final results
    print(f"\n{'='*60}")
    print("Genetic Algorithm Completed!")
    print(f"{'='*60}")
    print(f"Best fitness achieved: {max(best_fitness_history):.6f}")
    print(f"Final mean fitness: {mean_fitness_history[-1]:.6f}")
    print(f"\nResults saved in: {GENERATION_SAVE_PATH}/")
    
    logger.info("="*60)
    logger.info("Genetic Algorithm Completed")
    logger.info(f"Best fitness: {max(best_fitness_history):.6f}")
    logger.info("="*60)
    
    return population, fitness_scores, best_fitness_history

def main():
    """Main function for genetic algorithm."""
    parser = argparse.ArgumentParser(description='Run genetic algorithm for neural networks')
    parser.add_argument('config_file', help='Path to JSON configuration file')
    args = parser.parse_args()
    
    if not os.path.exists(args.config_file):
        print(f"Error: Configuration file not found: {args.config_file}")
        sys.exit(1)
    
    run_genetic_algorithm(args.config_file)

if __name__ == "__main__":
    main()