Merge pull request #75 from y0z/feature/benchmarks

Introduce `ConstrainedMixin`

Author: nabenabe0928

docs/source/benchmarks.rst

@@ -8,4 +8,5 @@ optunahub.benchmarks
     :nosignatures:
     :template: custom_summary.rst
 
-    optunahub.benchmarks.BaseProblem
+    optunahub.benchmarks.BaseProblem
+    optunahub.benchmarks.ConstrainedMixin

optunahub/benchmarks/__init__.py

@@ -1,6 +1,8 @@
 from ._base_problem import BaseProblem
+from ._constrained_mixin import ConstrainedMixin
 
 
 __all__ = [
     "BaseProblem",
+    "ConstrainedMixin",
 ]

optunahub/benchmarks/_constrained_mixin.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+import optuna
+
+
+class ConstrainedMixin:
+    """Mixin class for constrained optimization problems.
+
+    Example:
+        You can define a constrained optimization problem by inheriting this class and implementing
+        the :meth:`evaluate_constraints` method as follows.
+
+        ::
+
+            import optuna
+            import optunahub
+
+            class BinAndKorn(optunahub.benchmarks.ConstrainedMixin, optunahub.benchmarks.BaseProblem):
+                def evaluate(self, params: dict[str, float]) -> tuple[float]:
+                    x = params["x"]
+                    y = params["y"]
+
+                    v0 = 4 * x**2 + 4 * y**2
+                    v1 = (x - 5)**2 + (y - 5)**2
+
+                    return v0, v1
+
+                def evaluate_constraints(self, params: dict[str, float]) -> tuple[float]:
+                    x = params["x"]
+                    y = params["y"]
+
+                    # Constraints which are considered feasible if less than or equal to zero.
+                    # The feasible region is basically the intersection of a circle centered at (x=5, y=0)
+                    # and the complement to a circle centered at (x=8, y=-3).
+                    c0 = (x - 5)**2 + y**2 - 25
+                    c1 = -((x - 8)**2) - (y + 3)**2 + 7.7
+
+                    return c0, c1
+
+                @property
+                def search_space(self) -> dict[str, optuna.distributions.BaseDistribution]:
+                    return {
+                        "x": optuna.distributions.FloatDistribution(low=-15, high=30),
+                        "y": optuna.distributions.FloatDistribution(low=-15, high=30)
+                    }
+
+                @property
+                def directions(self) -> list[optuna.study.StudyDirection]:
+                    return [optuna.study.StudyDirection.MINIMIZE, optuna.study.StudyDirection.MINIMIZE]
+
+            problem = BinAndKorn()
+            sampler = optuna.samplers.TPESampler(constraints_func=problem.constraints_func)
+            study = optuna.create_study(sampler=sampler, directions=problem.directions)
+            study.optimize(problem, n_trials=20)
+    """
+
+    def constraints_func(self, trial: optuna.trial.FrozenTrial) -> Sequence[float]:
+        """Evaluate the constraint functions.
+
+        Args:
+            trial: Optuna trial object.
+        Returns:
+            List of the constraint values.
+        """
+        return self.evaluate_constraints(trial.params.copy())
+
+    def evaluate_constraints(self, params: dict[str, Any]) -> Sequence[float]:
+        """Evaluate the constraint functions.
+
+        Args:
+            params: Dictionary of input parameters.
+        Returns:
+            List of the constraint values.
+        """
+        raise NotImplementedError

tests/test_benchmarks.py

@@ -1,24 +1,41 @@
 from __future__ import annotations
 
 import optuna
+from optuna.samplers._base import _CONSTRAINTS_KEY
 
 import optunahub
 
 
-def test_base_problem() -> None:
-    class TestProblem(optunahub.benchmarks.BaseProblem):
-        def evaluate(self, params: dict[str, float]) -> float:
-            x = params["x"]
-            return x**2
+class TestProblem(optunahub.benchmarks.BaseProblem):
+    def evaluate(self, params: dict[str, float]) -> float:
+        x = params["x"]
+        return x**2
+
+    @property
+    def search_space(self) -> dict[str, optuna.distributions.BaseDistribution]:
+        return {"x": optuna.distributions.FloatDistribution(low=-1, high=1)}
 
-        @property
-        def search_space(self) -> dict[str, optuna.distributions.BaseDistribution]:
-            return {"x": optuna.distributions.FloatDistribution(low=-1, high=1)}
+    @property
+    def directions(self) -> list[optuna.study.StudyDirection]:
+        return [optuna.study.StudyDirection.MINIMIZE]
 
-        @property
-        def directions(self) -> list[optuna.study.StudyDirection]:
-            return [optuna.study.StudyDirection.MINIMIZE]
 
+def test_base_problem() -> None:
     problem = TestProblem()
     study = optuna.create_study(directions=problem.directions)
     study.optimize(problem, n_trials=20)  # verify no error occurs
+
+
+def test_constrained_mixin() -> None:
+    class ConstrainedTestProblem(optunahub.benchmarks.ConstrainedMixin, TestProblem):
+        def evaluate_constraints(self, params: dict[str, float]) -> list[float]:
+            return [params["x"]]
+
+    problem = ConstrainedTestProblem()
+    sampler = optuna.samplers.TPESampler(constraints_func=problem.constraints_func)
+    study = optuna.create_study(sampler=sampler, directions=problem.directions)
+    study.optimize(problem, n_trials=20)  # verify no error occurs
+
+    # Check if constraints are stored in trials
+    for t in study.trials:
+        assert _CONSTRAINTS_KEY in study._storage.get_trial_system_attrs(t._trial_id)