Update sampler.py

optuna · Dec 4, 2024 · 8da9640 · 8da9640
1 parent 4417a72
commit 8da9640
Showing 1 changed file with 13 additions and 19 deletions.
diff --git a/package/samplers/hebo_base_sampler/sampler.py b/package/samplers/hebo_base_sampler/sampler.py
@@ -1,18 +1,20 @@
 from __future__ import annotations
 
+from hebo.design_space.design_space import DesignSpace
+from hebo.optimizers.hebo import HEBO
 import numpy as np
 import optuna
-import pandas as pd
-from optuna.distributions import (BaseDistribution, CategoricalDistribution,
-                                  FloatDistribution, IntDistribution)
+from optuna.distributions import BaseDistribution
+from optuna.distributions import CategoricalDistribution
+from optuna.distributions import FloatDistribution
+from optuna.distributions import 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
+from optuna.trial import FrozenTrial
+from optuna.trial import TrialState
+import pandas as pd
 
 
 class HEBOSampler(BaseSampler):  # type: ignore
@@ -24,9 +26,7 @@ def __init__(
     ) -> None:
         self._seed = seed
         self._intersection_search_space = IntersectionSearchSpace()
-        self._independent_sampler = (
-            independent_sampler or optuna.samplers.RandomSampler(seed=seed)
-        )
+        self._independent_sampler = independent_sampler or optuna.samplers.RandomSampler(seed=seed)
         self._constant_liar = constant_liar
 
     def sample_relative(
@@ -45,13 +45,9 @@ def sample_relative(
 
         # 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
-            )
+            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
-            )
+            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))
@@ -136,9 +132,7 @@ def _convert_to_hebo_design_space(
         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)
-        )
+        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(