src/ahdcma/algorithms/base.py

"""Abstract optimizer interface and result/history dataclasses.

Every optimizer in :mod:`ahdcma.algorithms` (DOA, CMA-ES wrapper, PSO, GWO,
WOA, SCSO, GEGO, I-HAHO, AHD-CMA) extends :class:`Optimizer`. The contract
is intentionally narrow — a single :meth:`Optimizer.optimize` call returns
:class:`OptimizationResult`, which fully describes the run for downstream
statistics and figures.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass, field
from typing import Any

import numpy as np
from numpy.typing import NDArray

FitnessFn = Callable[[NDArray[np.float64]], float]


@dataclass
class SearchSpace:
    """Continuous search-space description.

    Optimizers operate in the unit hyper-cube ``[lower, upper]^d``. The
    ``names`` list is optional metadata for logging.
    """

    dim: int
    lower: NDArray[np.float64]
    upper: NDArray[np.float64]
    names: Sequence[str] | None = None

    def __post_init__(self) -> None:
        if self.lower.shape != (self.dim,) or self.upper.shape != (self.dim,):
            raise ValueError(
                f"lower/upper must have shape ({self.dim},); got "
                f"{self.lower.shape} and {self.upper.shape}"
            )
        if np.any(self.lower >= self.upper):
            raise ValueError("each lower bound must be strictly below the upper bound")

    @classmethod
    def unit_cube(cls, dim: int, names: Sequence[str] | None = None) -> SearchSpace:
        """Convenience constructor: ``[0, 1]^dim``."""
        return cls(
            dim=dim,
            lower=np.zeros(dim, dtype=np.float64),
            upper=np.ones(dim, dtype=np.float64),
            names=names,
        )

    def clip(self, x: NDArray[np.float64]) -> NDArray[np.float64]:
        """Project ``x`` (any shape ending in ``dim``) into the box."""
        return np.clip(x, self.lower, self.upper)


@dataclass
class History:
    """Per-generation run trace.

    Each list has length equal to the number of completed generations.
    ``populations`` and ``fitnesses`` are ragged in principle (CMA-ES may
    grow the population), but in practice all our algorithms hold it fixed.
    """

    generations: list[int] = field(default_factory=list)
    populations: list[NDArray[np.float64]] = field(default_factory=list)
    fitnesses: list[NDArray[np.float64]] = field(default_factory=list)
    best_fitness: list[float] = field(default_factory=list)
    mode_per_gen: list[str] = field(default_factory=list)
    entropy_per_gen: list[float] = field(default_factory=list)
    ruggedness_per_gen: list[float] = field(default_factory=list)

    def append(
        self,
        generation: int,
        population: NDArray[np.float64],
        fitness: NDArray[np.float64],
        *,
        mode: str = "",
        entropy: float = float("nan"),
        ruggedness: float = float("nan"),
    ) -> None:
        """Record a single generation."""
        self.generations.append(int(generation))
        self.populations.append(np.asarray(population, dtype=np.float64).copy())
        self.fitnesses.append(np.asarray(fitness, dtype=np.float64).copy())
        self.best_fitness.append(float(np.min(fitness)))
        self.mode_per_gen.append(mode)
        self.entropy_per_gen.append(float(entropy))
        self.ruggedness_per_gen.append(float(ruggedness))

    def __len__(self) -> int:
        return len(self.generations)


@dataclass
class OptimizationResult:
    """Final outcome of a single :meth:`Optimizer.optimize` call."""

    best_x: NDArray[np.float64]
    best_f: float
    history: History
    config_snapshot: Mapping[str, Any]
    run_id: str
    wall_time: float


class Optimizer(ABC):
    """Abstract base class. Subclasses implement :meth:`optimize`."""

    def __init__(
        self,
        config: Mapping[str, Any],
        fitness_fn: FitnessFn,
        search_space: SearchSpace,
        *,
        run_id: str = "anonymous",
    ) -> None:
        self.config = dict(config)
        self.fitness_fn = fitness_fn
        self.search_space = search_space
        self.run_id = run_id
        self._history = History()

    @abstractmethod
    def optimize(self) -> OptimizationResult:
        """Run the optimization loop and return the result."""

    def get_history(self) -> History:
        """Return the recorded :class:`History`."""
        return self._history