mixture.coco

"""
The mixture backend. Lets you specify a distribution over different possible algorithms.
"""

from bbopt import constants
from bbopt.util import convert_match_errors
from bbopt.registry import alg_registry
from bbopt.backends.util import (
    Backend,
    get_backend,
    get_cum_probs_for,
    random_from_cum_probs,
)


# Backend:

class MixtureBackend(Backend):
    """Mixture backend. Takes in a distribution over different possible algorithms
    of the form [(algorithm, weight)]. The properties selected_alg and selected_backend
    can be used to retrieve which alg/backend is currently being used."""

    backend_name = "mixture"
    request_backend_store = True
    remove_erroring_algs = None

    @override
    @convert_match_errors
    match def attempt_update(self, examples, params, distribution, remove_erroring_algs=False, *, _backend_store):
        """Special method that allows fast updating of the backend."""
        self.use_distribution(distribution, force=remove_erroring_algs != self.remove_erroring_algs)

        self.examples = examples
        self.params = params
        self.remove_erroring_algs = remove_erroring_algs
        self.backend_store = _backend_store

        self.select_new_backend()
        return True

    def use_distribution(self, distribution, force=False):
        """Set the distribution to the given distribution."""
        distribution = tuple(
            (alg, weight() if callable(weight) else weight)
            for alg, weight in distribution
        )

        if force or distribution != self.distribution:
            self.cum_probs = get_cum_probs_for(distribution)
            self.distribution = distribution

    def select_new_backend(self):
        """Randomly select a new backend."""
        # randomly select algorithm
        self.selected_alg = random_from_cum_probs(self.cum_probs)
        if self.selected_alg is None:
            raise ValueError(f"could not select backend from distribution: {self.distribution}")

        # initialize backend
        self.selected_backend, options = alg_registry[self.selected_alg]
        try:
            self.current_backend = get_backend(
                self.backend_store,
                self.selected_backend,
                self.examples,
                self.params,
                **options,
            )
        except constants.erroring_backend_errs:
            if not self.remove_erroring_algs:
                raise
            self.reselect_backend()

    def reselect_backend(self):
        """Choose a new backend when the current one errors."""
        new_distribution = []
        for alg, weight in self.distribution:
            if alg != self.selected_alg:
                new_distribution.append((alg, weight))
        self.cum_probs = get_cum_probs_for(new_distribution)
        self.select_new_backend()

    @override
    def param(self, name, func, *args, **kwargs):
        """Defer parameter selection to the selected backend."""
        try:
            return self.current_backend.param(name, func, *args, **kwargs)
        except constants.erroring_backend_errs:
            if not self.remove_erroring_algs:
                raise
            self.reselect_backend()
        return self.param(name, func, *args, **kwargs)

    @classmethod
    def register_safe_alg_for(cls, base_alg, new_alg_name=None, fallback_alg=None):
        """Register a version of base_alg that defaults to the fallback if base_alg fails."""
        if new_alg_name is None:
            new_alg_name = "safe_" + base_alg
        if fallback_alg is None:
            fallback_alg = constants.safe_fallback_alg
        cls.register_alg(
            new_alg_name,
            distribution=(
                (base_alg, float("inf")),
                (fallback_alg, 1),
            ),
            remove_erroring_algs=True,
        )

    @classmethod
    def register_epsilon_exploration_alg_for(cls, base_alg, new_alg_name=None, eps=None):
        """Register a version of base_alg with epsilon greedy exploration."""
        if new_alg_name is None:
            new_alg_name = "epsilon_" + base_alg
        cls.register_alg(
            new_alg_name,
            distribution=(
                # we defer evaluation here so that constants.eps_greedy_explore_prob
                #  can be modified to change the epsilon
                (base_alg, -> 1 - (eps ?? constants.eps_greedy_explore_prob)),
                ("random", -> eps ?? constants.eps_greedy_explore_prob),
            ),
        )




# Registered names:

MixtureBackend.register()

MixtureBackend.register_epsilon_exploration_alg_for("max_greedy")

MixtureBackend.register_safe_alg_for("gaussian_process")
MixtureBackend.register_safe_alg_for("random_forest")
MixtureBackend.register_safe_alg_for("extra_trees")
MixtureBackend.register_safe_alg_for("gradient_boosted_regression_trees")