Source code for optilab.metamodels.approximate_ranking_metamodel

"""
Approximate ranking metamodel based on lmm-CMA-ES.
"""

import copy

from ..data_classes import PointList
from ..functions import ObjectiveFunction
from ..functions.surrogate import SurrogateObjectiveFunction




class ApproximateRankingMetamodel:
    """Approximate ranking metamodel based on lmm-CMA-ES"""



    def __init__(
        self,
        population_size: int,
        mu: int,
        objective_function: ObjectiveFunction,
        surrogate_function: SurrogateObjectiveFunction,
        *,
        buffer_size: int | None = None,
    ) -> None:
        """
        Class constructor.

        Args:
            population_size: The population size (lambda).
            mu: The number of points compared in the ranking procedure (mu).
            objective_function: The objective function that's being optimized.
            surrogate_function: Surrogate function used to estimate
                the optimized function.
            buffer_size: Number of last evaluated samples to use for training the
                surrogate function. If None, all samples are used.
        """
        self.population_size = population_size
        self.mu = mu

        self.n_init: int
        self.n_step: int
        self.init_n()
        self.train_set = PointList(points=[])

        self.objective_function = objective_function
        self.surrogate_function = surrogate_function

        self.buffer_size = buffer_size




    def init_n(self) -> None:
        """
        Initializes n_init and n_step values based on the population size.
        """
        self.n_init = self.population_size
        self.n_step = max(1, self.population_size // 10)




    def _update_n(self, num_iters: int) -> None:
        """
        Updates n_init and n_step values base on the number of algorithm iterations.

        Args:
            num_iters: The number of iterations done by the algorithm.
        """
        if num_iters > 2:
            self.n_init = min(
                self.n_init + self.n_step, self.population_size - self.n_step
            )
        elif num_iters < 2:
            self.n_init = max(self.n_step, self.n_init - self.n_step)




    def __call__(self, points: PointList) -> PointList:
        """
        Approximates the values of provided points with surrogate objective function.

        Args:
            points: List of points to evaluate.

        Returns:
            List of evaluated points.
        """
        return PointList(points=[self.surrogate_function(x) for x in points])




    def train_surrogate(self) -> None:
        """
        Retrain the surrogate function with samples from the training set.
        """
        if self.buffer_size:
            self.surrogate_function.train(
                PointList(self.train_set[-self.buffer_size :])
            )
        else:
            self.surrogate_function.train(self.train_set)




    def evaluate(self, xs: PointList) -> PointList:
        """
        Evaluate provided point with the objective function and append results to the training
        set and retrain the surrogate function on new training data.

        Args:
            xs: List of point to evaluate with the objective function.

        Returns:
            List of evaluated points.
        """
        result = PointList(
            points=[self.objective_function(point) for point in xs.points]
        )
        self.train_set.extend(result)

        self.train_surrogate()

        return result




    def get_log(self) -> PointList:
        """
        Get the list of points evaluated by the objective function.

        Returns:
            List of points evaluated by the objective function.
        """
        return self.train_set




    def adapt(self, xs: PointList) -> None:
        """
        Perform another loop of the optimization on new data.

        Args:
            xs: Solution candidates generated by the optimizer.

        Raises:
            ValueError: When number of provided points mismatches the expected input size.
        """
        if not len(xs) == self.population_size:
            raise ValueError(
                f"The number of provided points is different than expected."
                f"Expected {self.population_size}, got {len(xs)}."
            )

        if len(self.train_set) < self.population_size:
            self.evaluate(xs)
            return

        # points not evaluated in this run
        not_evaluated = copy.deepcopy(xs)

        # points from current and previous loop
        items_previous = None
        items_current = self(xs)

        # rank items
        items_current.rank()

        # evaluate first n_init items
        not_evaluated = self(not_evaluated)
        not_evaluated.rank()
        self.evaluate(not_evaluated[: self.n_init])
        not_evaluated = not_evaluated[self.n_init :]

        num_iter = 0
        for _ in range((self.population_size - self.n_init) // self.n_step):
            # start new loop
            num_iter += 1
            items_previous = items_current
            items_current = self(xs)

            # rank items
            items_current.rank()

            # check if the mu ranking changed
            if all(
                (
                    new_pt == pt
                    for new_pt, pt in zip(
                        items_previous[: self.mu],
                        items_current[: self.mu],
                    )
                )
            ):
                break

            # else evaluate n_step next items
            not_evaluated = self(not_evaluated)
            not_evaluated.rank()
            self.evaluate(not_evaluated[: self.n_step])
            not_evaluated = not_evaluated[self.n_step :]

        self._update_n(num_iter)