Add files via upload

package/samplers/hebo_base_sampler/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Hiroki Takizawa
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

package/samplers/hebo_base_sampler/__init__.py

@@ -0,0 +1,4 @@
+from .sampler import HEBOSampler
+
+
+__all__ = ["HEBOSampler"]

package/samplers/hebo_base_sampler/example.py

@@ -0,0 +1,24 @@
+import optuna
+import optunahub
+import time
+
+
+module = optunahub.load_module("samplers/hebo_base_sampler")
+HEBOSampler = module.HEBOSampler
+
+
+def objective(trial: optuna.trial.Trial) -> float:
+    x = trial.suggest_float("x", -10, 10)
+    y = trial.suggest_int("y", -10, 10)
+    time.sleep(1.0)
+    return x**2 + y**2
+
+
+if __name__ == "__main__":
+    sampler = HEBOSampler(constant_liar=True)
+    study = optuna.create_study(sampler=sampler)
+    study.optimize(objective, n_trials=100, n_jobs=2)
+    print(study.best_trial.params)
+
+    fig = optuna.visualization.plot_optimization_history(study)
+    fig.write_image("hebo_optimization_history.png")

package/samplers/hebo_base_sampler/requirements.txt

@@ -0,0 +1,3 @@
+optuna
+optunahub
+hebo@git+https://github.com/huawei-noah/[email protected]#subdirectory=HEBO

package/samplers/hebo_base_sampler/sampler.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+import numpy as np
+import optuna
+import pandas as pd
+from optuna.distributions import (BaseDistribution, CategoricalDistribution,
+                                  FloatDistribution, IntDistribution)
+from optuna.samplers import BaseSampler
+from optuna.search_space import IntersectionSearchSpace
+from optuna.study import Study
+from optuna.study._study_direction import StudyDirection
+from optuna.trial import FrozenTrial, TrialState
+
+from hebo.design_space.design_space import DesignSpace
+from hebo.optimizers.hebo import HEBO
+
+
+class HEBOSampler(BaseSampler):  # type: ignore
+    def __init__(
+        self,
+        seed: int | None = None,
+        constant_liar: bool = False,
+        independent_sampler: BaseSampler | None = None,
+    ) -> None:
+        self._seed = seed
+        self._intersection_search_space = IntersectionSearchSpace()
+        self._independent_sampler = (
+            independent_sampler or optuna.samplers.RandomSampler(seed=seed)
+        )
+        self._constant_liar = constant_liar
+
+    def sample_relative(
+        self,
+        study: Study,
+        trial: FrozenTrial,
+        search_space: dict[str, BaseDistribution],
+    ) -> dict[str, float]:
+        if self._constant_liar:
+            target_states = [TrialState.COMPLETE, TrialState.RUNNING]
+        else:
+            target_states = [TrialState.COMPLETE]
+        trials = study.get_trials(deepcopy=False, states=target_states)
+        if len(t for t in trials if t.state == TrialState.COMPLETE) < 1:
+            return {}
+
+        # Assume that the back-end HEBO implementation aims to minimize.
+        if study.direction == StudyDirection.MINIMIZE:
+            worst_values = max(
+                t.values for t in trials if t.state == TrialState.COMPLETE
+            )
+        else:
+            worst_values = min(
+                t.values for t in trials if t.state == TrialState.COMPLETE
+            )
+        sign = 1.0 if study.direction == StudyDirection.MINIMIZE else -1.0
+
+        hebo = HEBO(self._convert_to_hebo_design_space(search_space))
+        for t in trials:
+            hebo_params = {name: t.params[name] for name in search_space.keys()}
+            if t.state == TrialState.COMPLETE:
+                hebo.observe(pd.DataFrame([hebo_params]), np.asarray([t.values * sign]))
+            elif t.state == TrialState.RUNNING:
+                # If `constant_liar == True`, assume that the RUNNING params result in bad values,
+                # thus preventing the simultaneous suggestion of (almost) the same params
+                # during parallel execution.
+                hebo.observe(pd.DataFrame([hebo_params]), np.asarray([worst_values]))
+            else:
+                assert False
+        params_pd = hebo.suggest()
+        params = {}
+        for name in search_space.keys():
+            params[name] = params_pd[name].to_numpy()[0]
+        return params
+
+    def _convert_to_hebo_design_space(
+        self, search_space: dict[str, BaseDistribution]
+    ) -> DesignSpace:
+        design_space = []
+        for name, distribution in search_space.items():
+            if isinstance(distribution, FloatDistribution) and not distribution.log:
+                design_space.append(
+                    {
+                        "name": name,
+                        "type": "num",
+                        "lb": distribution.low,
+                        "ub": distribution.high,
+                    }
+                )
+            elif isinstance(distribution, FloatDistribution) and distribution.log:
+                design_space.append(
+                    {
+                        "name": name,
+                        "type": "pow",
+                        "lb": distribution.low,
+                        "ub": distribution.high,
+                    }
+                )
+            elif isinstance(distribution, IntDistribution) and distribution.log:
+                design_space.append(
+                    {
+                        "name": name,
+                        "type": "pow_int",
+                        "lb": distribution.low,
+                        "ub": distribution.high,
+                    }
+                )
+            elif isinstance(distribution, IntDistribution) and distribution.step:
+                design_space.append(
+                    {
+                        "name": name,
+                        "type": "step_int",
+                        "lb": distribution.low,
+                        "ub": distribution.high,
+                        "step": distribution.step,
+                    }
+                )
+            elif isinstance(distribution, IntDistribution):
+                design_space.append(
+                    {
+                        "name": name,
+                        "type": "int",
+                        "lb": distribution.low,
+                        "ub": distribution.high,
+                    }
+                )
+            elif isinstance(distribution, CategoricalDistribution):
+                design_space.append(
+                    {
+                        "name": name,
+                        "type": "cat",
+                        "categories": distribution.choices,
+                    }
+                )
+            else:
+                raise NotImplementedError(f"Unsupported distribution: {distribution}")
+        return DesignSpace().parse(design_space)
+
+    def infer_relative_search_space(self, study, trial):
+        return optuna.search_space.intersection_search_space(
+            study.get_trials(deepcopy=False)
+        )
+
+    def sample_independent(self, study, trial, param_name, param_distribution):
+        return self._independent_sampler.sample_independent(
+            study, trial, param_name, param_distribution
+        )