hyperopt.coco

"""
The hyperopt backend. Does black box optimization using hyperopt.
"""

import sys

import numpy as np

from hyperopt import (
    hp,
    FMinIter,
    tpe,
    anneal,
)
from hyperopt.pyll import as_apply
from hyperopt.base import (
    Domain,
    Trials,
    STATUS_OK,
    STATUS_RUNNING,
    JOB_STATE_DONE,
    spec_from_misc,
)

from bbopt.util import sorted_items
from bbopt.backends.util import (
    StandardBackend,
    negate_objective,
    get_names_and_features,
)


# Utilities:

def create_space(name, func, *args):
    """Create a hyperopt space for the given parameter."""
    match func, args:
        case "choice", (choices,):
            return hp.choice(name, choices)
        case "randrange", (start, stop, step):
            if step != 1:
                raise ValueError("the hyperopt backend only supports a randrange step size of 1 (use bb.choice(name, range(start, stop, step)) instead)")
            # despite being called randint, hp.randint is exclusive
            return start + hp.randint(name, stop - start)
        case "uniform", (a, b):
            return hp.uniform(name, a, b)
        case "normalvariate", (mu, sigma):
            return hp.normal(name, mu, sigma)
    raise TypeError(f"invalid parameter {name}")


def examples_to_trials(examples, params):
    """Create hyperopt trials from the given examples."""
    trials = []
    NA = object()  # used to mark missing values

    for tid, ex in enumerate(examples):

        match {"gain": gain, **_} in ex:
            loss = negate_objective(gain)
        else:
            loss = ex["loss"]
        result = {
            "status": STATUS_OK,
            "loss": loss,
        }

        vals = {}
        idxs = {}
        for k, v in get_names_and_features(
            ex["values"],
            params,
            fallback_func=(name, func, *args, **kwargs) -> NA,
            converters={
                "choice": (val, choices) -> choices.index(val),
                "randrange": (val, start, stop, step) -> val - start,
            },
            convert_fallback=False,
        ):
            vals[k] = [v] if v is not NA else []
            idxs[k] = [tid] if v is not NA else []

        misc = {
            "tid": tid,
            "idxs": idxs,
            "vals": vals,
            "cmd": None,
        }

        trials.append({
            "tid": tid,
            "result": result,
            "misc": misc,
            "spec": spec_from_misc(misc),
            "state": JOB_STATE_DONE,
            "owner": None,
            "book_time": None,
            "refresh_time": None,
            "exp_key": None,
        })

    return trials


# Backend:

class HyperoptBackend(StandardBackend):
    """The hyperopt backend uses hyperopt for black box optimization."""
    backend_name = "hyperopt"
    implemented_funcs = (
        # should match create_space above
        "choice",
        "randrange",
        "uniform",
        "normalvariate",
    )

    @override
    def setup_backend(self, params, algo=tpe.suggest, rstate=None, show_progressbar=False, **options):
        """Special method to initialize the backend from params."""
        if rstate is None:
            try:
                rstate = np.random.default_rng()
            except AttributeError:
                rstate = np.random.RandomState()
        self.params = params

        space = {
            name: create_space(name, func, *args)
            for name, (func, args, kwargs) in sorted_items(params)
        } |> as_apply

        domain = Domain(self.set_current_values, space)

        self.trials = Trials()

        self.fmin_iter = FMinIter(
            algo,
            domain,
            self.trials,
            rstate,
            show_progressbar=show_progressbar,
            **options,
        )

    @override
    def tell_examples(self, new_examples):
        """Special method that allows fast updating of the backend with new examples."""
        trial_list = examples_to_trials(new_examples, self.params)
        self.trials.insert_trial_docs(trial_list)
        self.trials.refresh()

        # run one iteration of hyperparameter optimization, with values saved
        #  to the self.set_current_values callback passed to Domain
        next(self.fmin_iter)

        assert self.current_values is not None, self.current_values
        assert set(self.current_values.keys()) == set(self.params), self.current_values

    def set_current_values(self, values):
        """Callback to set the values for this run."""
        assert isinstance(values, dict), values
        self.current_values = values
        return {
            "status": STATUS_RUNNING,
        }


# Registered names:

HyperoptBackend.register()

HyperoptBackend.register_alg("tree_structured_parzen_estimator", algo=tpe.suggest)
HyperoptBackend.register_alg("annealing", algo=anneal.suggest)
if sys.version_info >= (3,):
    from hyperopt import atpe
    HyperoptBackend.register_alg("adaptive_tpe", algo=atpe.suggest)

HyperoptBackend.register_meta_for_all_algs("any_hyperopt")