openevolve_database.py

"""
Program management for OpenEvolve
"""

import logging
import random
import uuid
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

from openevolve.config import Config
from openevolve.utils.code_utils import apply_diff

logger = logging.getLogger(__name__')


class Program:
    """Stores a program and its metadata"""

    def __init__(
        self,
        id: str,
        code: str,
        language: str,
        metrics: Dict[str, float],
        parent_id: Optional[str] = None,
        iteration_found: int = 0,
        changes: str = '',
        parent_metrics: str = '',
        key_features: Optional[List[str]] = None,
    ):
        self.id = id
        self.code = code
        self.language = language
        self.metrics = metrics
        self.parent_id = parent_id
        self.iteration_found = iteration_found
        self.changes = changes
        self.parent_metrics = parent_metrics or {}
        self.key_features = key_features or []

    def to_dict(self) -> Dict[str, Any]:
        """Convert program to dictionary"""
        return {
            "id": self.id,
            "code": self.code,
            "language": self.language,
            "metrics": self.metrics,
            "parent_id": self.parent_id,
            "iteration_found": self.iteration_found,
            "changes": self.changes,
            "parent_metrics": self.parent_metrics,
            "key_features": self.key_features,
        }


class ProgramDatabase:
    """Manages the program population and archive"""

    def __init__(self, config: Config):
        self.config = config
        self.programs: Dict[str, Program] = {}
        self.archive: Dict[str, Program] = {}
        self.island_populations: Dict[int, List[str]] = {
            i: [] for i in range(config.database.num_islands)
        }
        self.last_iteration = 0
        self._program_hashes: Dict[str, str] = {}  # Maps program hash to ID

        if not config.database.in_memory:
            logger.warning("Non-in-memory database not yet implemented")

    def add(self, program: Program) -> bool:
        """Add a program to the database if it doesn't already exist"""
        # Compute a simple hash of the code (for deduplication)
        program_hash = str(hash(program.code))

        if program_hash in self._program_hashes:
            logger.debug(f"Duplicate program detected: {program.id}")
            return False

        # Add to main programs dictionary
        self.programs[program.id] = program
        self._program_hashes[program_hash] = program.id

        # Add to island population
        island = random.randint(0, self.config.database.num_islands - 1)
        self.island_populations[island].append(program.id)

        # Maintain population size
        while len(self.island_populations[island]) > self.config.database.population_size:
            program_to_remove = self.island_populations[island].pop(0)
            if program_to_remove in self.programs:
                del self._program_hashes[str(hash(self.programs[program_to_remove].code))]
                del self.programs[program_to_remove]

        # Add to archive if it's among the best
        if len(self.archive) < self.config.database.archive_size:
            self.archive[program.id] = program
        else:
            worst_in_archive = min(
                self.archive.values(),
                key=lambda p: sum(p.metrics.values()) / max(1, len(p.metrics)),
            )
            if sum(program.metrics.values()) / max(1, len(program.metrics)) > sum(
                worst_in_archive.metrics.values()
            ) / max(1, len(worst_in_archive.metrics)):
                del self.archive[worst_in_archive.id]
                self.archive[program.id] = program

        return True

    def get(self, program_id: str) -> Optional[Program]:
        """Retrieve a program by ID"""
        return self.programs.get(program_id)

    def get_best_program(self) -> Optional[Program]:
        """Get the best performing program"""
        if not self.programs:
            return None
        return max(
            self.programs.values(),
            key=lambda p: sum(p.metrics.values()) / max(1, len(p.metrics)),
        )

    def sample(self) -> Tuple[Program, List[Program]]:
        """Sample a parent program and inspirations from the database."""
        import random

        mutation_rate = getattr(self.config.database, "mutation_rate", 0.1)  # Default to 0.1
        parent = None
        inspirations = []

        # Select parent from elite programs
        if self.programs and random.random() > mutation_rate:
            elite = sorted(
                self.programs.values(),
                key=lambda p: sum(p.metrics.values()) / max(1, len(p.metrics)),
                reverse=True,
            )[: int(len(self.programs) * self.config.database.elite_selection_ratio)]
            parent = random.choice(elite) if elite else random.choice(list(self.programs.values()))
        else:
            # Mutate a random program
            parent = random.choice(list(self.programs.values()))
            parent = self._mutate_program(parent)  # Apply mutation

        # Select inspirations
        inspirations = random.sample(
            list(self.programs.values()),
            min(self.config.prompt.num_top_programs, len(self.programs)),
        )

        return parent, inspirations

    def _mutate_program(self, program: Program) -> Program:
        """Apply a simple mutation to a program."""
        import random

        # Example: Add a comment to a random line
        code_lines = program.code.split("\n")
        if code_lines and random.random() < 0.5:
            line_idx = random.randint(0, len(code_lines) - 1)
            code_lines[line_idx] += " # Mutated"
        new_code = "\n".join(code_lines)
        return Program(
            id=f"{program.id}_mutated_{random.randint(1000, 9999)}",
            code=new_code,
            language=program.language,
            metrics=program.metrics,
            iteration_found=self.last_iteration + 1,
        )

    def get_programs_by_iteration(self, iteration: int) -> List[Program]:
        """Get all programs from a specific iteration"""
        return [p for p in self.programs.values() if p.iteration_found == iteration]

    def save_checkpoint(self, output_dir: str) -> None:
        """Save the current database state to disk"""
        # Placeholder for checkpoint saving
        pass

    def load_checkpoint(self, checkpoint_dir: str) -> None:
        """Load database state from a checkpoint"""
        # Placeholder for checkpoint loading
        pass