From 4e031585b2d90b290d6bcd1dca5ccd50547a448a Mon Sep 17 00:00:00 2001
From: Contramundum <contramundum53@outlook.com>
Date: Thu, 14 Sep 2023 13:40:19 +0900
Subject: [PATCH] Support all-categorical cases

---
 optuna_dashboard/preferential/samplers/gp.py | 42 +++++++++++++++-----
 1 file changed, 31 insertions(+), 11 deletions(-)

diff --git a/optuna_dashboard/preferential/samplers/gp.py b/optuna_dashboard/preferential/samplers/gp.py
index 8349b1a4..349da03a 100644
--- a/optuna_dashboard/preferential/samplers/gp.py
+++ b/optuna_dashboard/preferential/samplers/gp.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import math
 from typing import Any
 from typing import Callable
+from typing import cast
 
 import botorch.acquisition.analytic
 import botorch.models.model
@@ -16,6 +17,8 @@ import optuna
 import optuna._transform
 import torch
 from torch import Tensor
+from optuna.distributions import CategoricalDistribution
+import itertools
 
 from .._system_attrs import get_preferences
 
@@ -310,7 +313,7 @@ class PreferentialGPSampler(optuna.samplers.BaseSampler):
         search_space: dict[str, optuna.distributions.BaseDistribution],
     ) -> dict[str, Any]:
         preferences = get_preferences(study.system_attrs)
-        if len(preferences) == 0:
+        if len(preferences) == 0 or len(search_space) == 0:
             return {}
 
         trials = study.get_trials(deepcopy=False)
@@ -355,16 +358,33 @@ class PreferentialGPSampler(optuna.samplers.BaseSampler):
                 best_f=torch.max(sampled_gp.posterior(params[:, None, :]).mean),
             )
 
-            # TODO: Make it possible to apply it on categorical variables
-            candidates, _ = botorch.optim.optimize_acqf(
-                acq_function=acqf,
-                bounds=torch.from_numpy(trans.bounds.T),
-                q=1,
-                num_restarts=10,
-                raw_samples=512,
-                options={"batch_limit": 5, "maxiter": 200},
-                sequential=True,
-            )
+            # TODO: Make it possible to apply it on mixed search space
+            if all(isinstance(dist, CategoricalDistribution) for dist in search_space.values()):
+                all_param_combinations = itertools.product(
+                    *[
+                        [(name, choice) for choice in cast(CategoricalDistribution, dist).choices]
+                        for name, dist in search_space.items()
+                    ]
+                )
+                choices = torch.tensor(
+                    np.array([trans.transform(dict(params)) for params in all_param_combinations]),
+                    dtype=torch.float64,
+                )
+                candidates, _ = botorch.optim.optimize_acqf_discrete(
+                    acq_function=acqf,
+                    choices=choices,
+                    q=1,
+                )
+            else:
+                candidates, _ = botorch.optim.optimize_acqf(
+                    acq_function=acqf,
+                    bounds=torch.from_numpy(trans.bounds.T),
+                    q=1,
+                    num_restarts=10,
+                    raw_samples=512,
+                    options={"batch_limit": 5, "maxiter": 200},
+                    sequential=True,
+                )
             next_x = trans.untransform(candidates[0].detach().numpy())
             return next_x