Remove SearchSpace from NelderMeadSampler. (#396)

* Remove SearchSpace.

* Fix examples.

* Fix tests.

---------

Co-authored-by: Yoshiaki Bando <yoshipon@users.noreply.github.com>

aiaccel/hpo/optuna/samplers/nelder_mead_sampler.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import numpy.typing as npt
-from typing import Any, TypedDict
+from typing import Any
 
 from collections.abc import Sequence
 import math
@@ -19,23 +19,6 @@
 __all__ = ["NelderMeadSampler", "NelderMeadEmptyError"]
 
 
-class SearchSpaceRequired(TypedDict):
-    """
-    Required keys
-    """
-
-    low: int | float
-    high: int | float
-
-
-class SearchSpace(SearchSpaceRequired, total=False):
-    """
-    Optional keys
-    """
-
-    log: bool
-
-
 class NelderMeadSampler(optuna.samplers.BaseSampler):
     """Sampler using the NelderMead algorithm
 
@@ -103,24 +86,14 @@ class NelderMeadSampler(optuna.samplers.BaseSampler):
 
     def __init__(
         self,
-        search_space: dict[str, SearchSpace],
+        search_space: dict[str, tuple[int | float, int | float]],
         seed: int | None = None,
         rng: np.random.RandomState | None = None,
         coeff: NelderMeadCoefficient | None = None,
         block: bool = False,
         sub_sampler: optuna.samplers.BaseSampler | None = None,
     ) -> None:
-        self._search_space = {}
-        for key, value in search_space.items():
-            if "log" in value and value["log"]:
-                self._search_space[key] = {
-                    "low": math.log(value["low"]),
-                    "high": math.log(value["high"]),
-                    "log": value["log"],
-                }
-            else:
-                self._search_space[key] = value | {"log": False}  # type: ignore[assignment]
-
+        self._search_space = search_space
         _rng = rng if rng is not None else np.random.RandomState(seed) if seed is not None else None
 
         self.nm = NelderMeadAlgorism(
@@ -145,8 +118,7 @@ def sample_relative(
 
     def _get_params(self, study: Study, trial: FrozenTrial) -> npt.NDArray[np.float64] | None:
         try:
-            it = zip(self.nm.get_vertex(), self._search_space.values(), strict=False)
-            params = np.array([(space["high"] - space["low"]) * value + space["low"] for value, space in it])
+            params = self.nm.get_vertex()
         except NelderMeadEmptyError as e:
             if self.sub_sampler is None:
                 raise e
@@ -158,11 +130,11 @@ def _get_params(self, study: Study, trial: FrozenTrial) -> npt.NDArray[np.float6
     def _put_params(self, study: Study, trial: FrozenTrial, state: TrialState, values: Sequence[float] | None) -> None:
         if isinstance(values, list):
             system_attr = study._storage.get_trial_system_attrs(trial._trial_id)
-            raw_params = system_attr["params"] if "params" in system_attr else trial.params.values()
-            search_space = [(space["low"], space["high"]) for space in self._search_space.values()]
-            params = np.array(
-                [(value - low) / (high - low) for value, (low, high) in zip(raw_params, search_space, strict=False)]
-            )
+            if "params" in system_attr and "fixed_params" not in system_attr:
+                params = np.array(system_attr["params"])
+            else:  # sub_sampler or enqueued
+                it = zip(trial.params.values(), self._search_space.values(), strict=False)
+                params = np.array([(value - low) / (high - low) for value, (low, high) in it])
 
             self.nm.put_value(
                 params,
@@ -274,21 +246,10 @@ def sample_independent(
                 A parameter value.
 
         """
-        if (
-            isinstance(
-                param_distribution, optuna.distributions.IntDistribution | optuna.distributions.FloatDistribution
-            )
-            and param_distribution.log != self._search_space[param_name]["log"]
-        ):
-            raise ValueError(
-                f"Parameter {param_name} is set with log={self._search_space[param_name]['log']} "
-                f"but optuna.distributions.param_distribution.log={param_distribution.log}"
-            )
         system_attr = study._storage.get_trial_system_attrs(trial._trial_id)
         if "sub_trial" in system_attr and self.sub_sampler is not None:
             param_value = self.sub_sampler.sample_independent(study, trial, param_name, param_distribution)
-            value = math.log(param_value) if self._search_space[param_name]["log"] else param_value
-            if self._search_space[param_name]["low"] <= value <= self._search_space[param_name]["high"]:
+            if self._search_space[param_name][0] <= param_value <= self._search_space[param_name][1]:
                 return param_value
             else:
                 raise ValueError(
@@ -301,13 +262,20 @@ def sample_independent(
         param_index = list(self._search_space.keys()).index(param_name)
         param_value = system_attr["params"][param_index]
 
-        if self._search_space[param_name]["log"]:
-            param_value = math.exp(param_value)
+        # reverse normalization
+        assert hasattr(param_distribution, "high")
+        assert hasattr(param_distribution, "low")
+        assert hasattr(param_distribution, "log")
+        if param_distribution.log:  # log scale
+            high = math.log(param_distribution.high)
+            low = math.log(param_distribution.low)
+            param_value = math.exp((high - low) * param_value + low)
+        else:
+            param_value = (param_distribution.high - param_distribution.low) * param_value + param_distribution.low
 
         if isinstance(param_distribution, optuna.distributions.IntDistribution):
             param_value = int(param_value)
         if hasattr(param_distribution, "step") and param_distribution.step is not None:
-            assert hasattr(param_distribution, "low")
             param_value -= (param_value - param_distribution.low) % param_distribution.step
 
         contains = param_distribution._contains(param_distribution.to_internal_repr(param_value))

examples/hpo/samplers/coco/experiment_coco.py

@@ -13,19 +13,19 @@
 import cocoex
 import optuna
 
-from aiaccel.hpo.optuna.samplers.nelder_mead_sampler import NelderMeadEmptyError, NelderMeadSampler, SearchSpace
+from aiaccel.hpo.optuna.samplers.nelder_mead_sampler import NelderMeadEmptyError, NelderMeadSampler
 
 
 def _optimize_sequential(
-    study: optuna.Study, func: Callable[[list[float]], float], search_space: dict[str, SearchSpace]
+    study: optuna.Study, func: Callable[[list[float]], float], search_space: dict[str, tuple[int | float, int | float]]
 ) -> float | None:
     try:
         trial = study.ask()
     except NelderMeadEmptyError:
         return None
     param = []
     for name, distribution in search_space.items():
-        param.append(trial.suggest_float(name, distribution["low"], distribution["high"]))
+        param.append(trial.suggest_float(name, *distribution))
 
     result = func(param)
     time.sleep(0.1)
@@ -42,7 +42,7 @@ def _optimize_sequential_wrapper(args: list[Any]) -> float | None:
 def optimize(
     study: optuna.Study,
     func: Callable[[list[float]], float],
-    search_space: dict[str, SearchSpace],
+    search_space: dict[str, tuple[int | float, int | float]],
     result_csv_name: str,
     num_trial: int = 1000,
     num_parallel: int = 10,
@@ -105,9 +105,9 @@ def experiment_bbob() -> None:
     for problem in suite:  # this loop will take several minutes or longer
         problem.observe_with(observer)  # generates the data for cocopp post-processing
 
-        search_space: dict[str, SearchSpace] = {}
+        search_space: dict[str, tuple[int | float, int | float]] = {}
         for i in range(problem.dimension):
-            search_space[f"x{i}"] = {"low": -5.0, "high": 5.0}
+            search_space[f"x{i}"] = (-5.0, 5.0)
         print(search_space)
 
         if sampler_name == "nelder-mead":

examples/hpo/samplers/example.py

@@ -2,18 +2,18 @@
 
 import optuna
 
-from aiaccel.hpo.optuna.samplers.nelder_mead_sampler import NelderMeadSampler, SearchSpace
+from aiaccel.hpo.optuna.samplers.nelder_mead_sampler import NelderMeadSampler
 
-search_space: dict[str, SearchSpace] = {
-    "x": {"low": -10.0, "high": 10.0},
-    "y": {"low": -10.0, "high": 10.0},
+search_space = {
+    "x": (-10.0, 10.0),
+    "y": (-10.0, 10.0),
 }
 
 
 def sphere(trial: optuna.trial.Trial) -> float:
     params = []
     for name, distribution in search_space.items():
-        params.append(trial.suggest_float(name, distribution["low"], distribution["high"]))
+        params.append(trial.suggest_float(name, *distribution))
 
     return float(np.sum(np.asarray(params) ** 2))
 

examples/hpo/samplers/example_enqueue.py

@@ -4,11 +4,11 @@
 
 import optuna
 
-from aiaccel.hpo.optuna.samplers.nelder_mead_sampler import NelderMeadEmptyError, NelderMeadSampler, SearchSpace
+from aiaccel.hpo.optuna.samplers.nelder_mead_sampler import NelderMeadEmptyError, NelderMeadSampler
 
-search_space: dict[str, SearchSpace] = {
-    "x": {"low": -10.0, "high": 10.0},
-    "y": {"low": -10.0, "high": 10.0},
+search_space = {
+    "x": (-10.0, 10.0),
+    "y": (-10.0, 10.0),
 }
 
 
@@ -31,14 +31,14 @@ def sphere(params: list[float]) -> float:
                 except NelderMeadEmptyError:  # random sampling
                     study.enqueue_trial(
                         {
-                            "x": _rng.uniform(search_space["x"]["low"], search_space["x"]["high"]),
-                            "y": _rng.uniform(search_space["y"]["low"], search_space["y"]["high"]),
+                            "x": _rng.uniform(*search_space["x"]),
+                            "y": _rng.uniform(*search_space["y"]),
                         }
                     )
                     trial = study.ask()
 
-                x = trial.suggest_float("x", search_space["x"]["low"], search_space["x"]["high"])
-                y = trial.suggest_float("y", search_space["y"]["low"], search_space["y"]["high"])
+                x = trial.suggest_float("x", *search_space["x"])
+                y = trial.suggest_float("y", *search_space["y"])
 
                 trials.append(trial)
                 params.append([x, y])

examples/hpo/samplers/example_parallel.py

@@ -4,11 +4,11 @@
 
 import optuna
 
-from aiaccel.hpo.optuna.samplers.nelder_mead_sampler import NelderMeadSampler, SearchSpace
+from aiaccel.hpo.optuna.samplers.nelder_mead_sampler import NelderMeadSampler
 
-search_space: dict[str, SearchSpace] = {
-    "x": {"low": -10.0, "high": 10.0},
-    "y": {"low": -10.0, "high": 10.0},
+search_space = {
+    "x": (-10.0, 10.0),
+    "y": (-10.0, 10.0),
 }
 
 
@@ -17,7 +17,7 @@ def sphere(trial: optuna.trial.Trial) -> float:
     time.sleep(0.01)
 
     for name, distribution in search_space.items():
-        params.append(trial.suggest_float(name, distribution["low"], distribution["high"]))
+        params.append(trial.suggest_float(name, *distribution))
 
     return float(np.sum(np.asarray(params) ** 2))
 

examples/hpo/samplers/example_sub_sampler.py

@@ -4,11 +4,11 @@
 
 import optuna
 
-from aiaccel.hpo.optuna.samplers.nelder_mead_sampler import NelderMeadSampler, SearchSpace
+from aiaccel.hpo.optuna.samplers.nelder_mead_sampler import NelderMeadSampler
 
-search_space: dict[str, SearchSpace] = {
-    "x": {"low": -10.0, "high": 10.0},
-    "y": {"low": -10.0, "high": 10.0},
+search_space = {
+    "x": (-10.0, 10.0),
+    "y": (-10.0, 10.0),
 }
 
 
@@ -17,7 +17,7 @@ def sphere(trial: optuna.trial.Trial) -> float:
     time.sleep(0.01)
 
     for name, distribution in search_space.items():
-        params.append(trial.suggest_float(name, distribution["low"], distribution["high"]))
+        params.append(trial.suggest_float(name, *distribution))
 
     return float(np.sum(np.asarray(params) ** 2))