From f8f76b8e41c9d2be4ff0682cb19ed14bae65b982 Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Wed, 8 May 2024 14:58:53 +0200
Subject: [PATCH 01/21] add poli MT MOO

---
 .../core/utils/bo_pr/poli_objective_in_pr.py  | 64 ++++++++++++++++++-
 1 file changed, 62 insertions(+), 2 deletions(-)

diff --git a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
index 1f4babb..4bdbe32 100644
--- a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
+++ b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
@@ -9,14 +9,15 @@
 """
 
 from __future__ import annotations
-
+from typing import List
 
 import numpy as np
 import torch
+from botorch.utils.torch import BufferDict
 
 from poli.core.abstract_black_box import AbstractBlackBox
 
-from discrete_mixed_bo.problems.base import DiscreteTestProblem
+from discrete_mixed_bo.problems.base import DiscreteTestProblem, MultiObjectiveTestProblem
 
 
 class PoliObjective(DiscreteTestProblem):
@@ -54,3 +55,62 @@ def evaluate_true(self, X: torch.Tensor):
 
         # 2. evaluate the black box
         return torch.from_numpy(self.black_box(np.array(x_str)))
+
+
+class PoliMultiObjective(DiscreteTestProblem, MultiObjectiveTestProblem):
+    """
+    A bridge between poli black boxes and PR.
+    """
+    num_objectives: int
+    _ref_point: List[float]
+    _discrete_values: dict
+    _bounds: list
+
+    def __init__(
+        self,
+        black_box: AbstractBlackBox,
+        sequence_length: int,
+        alphabet: List[str] = None,
+        noise_std: float = None,
+        negate: bool = False,
+        integer_indices = None,
+        integer_bounds = None,
+        ref_point: List[float] = None,
+    ) -> None:
+        self._bounds = [(0, len(alphabet) - 1) for _ in range(sequence_length)]
+        self.dim = sequence_length
+        self.black_box = black_box
+        alphabet = alphabet or self.black_box.info.alphabet
+        self._ref_point = ref_point  #  NOTE: this assumes maximization of all objectives.
+        self.num_objectives = ref_point.shape[-1]
+        if alphabet is None:
+            raise ValueError("Alphabet must be provided.")
+
+        self.alphabet_s_to_i = {s: i for i, s in enumerate(alphabet)}
+        self.alphabet_i_to_s = {i: s for i, s in enumerate(alphabet)}
+        MultiObjectiveTestProblem.__init__(
+            self,
+            noise_std=noise_std,
+            negate=negate,
+        )
+        self._setup(integer_indices=integer_indices)
+        self.discrete_values = BufferDict()
+        self._discrete_values = {f"pos_{i}": list(self.alphabet_s_to_i.values()) for i in range(sequence_length)}
+        # for k, v in self._discrete_values.items():
+        #     self.discrete_values[k] = torch.tensor(v, dtype=torch.float)
+        #     self.discrete_values[k] /= self.discrete_values[k].max()
+        # super().__init__(noise_std, negate, categorical_indices=list(range(self.dim)))
+        for v in self._discrete_values.values():
+            self._bounds.append((0, len(alphabet)))
+
+    def evaluate_true(self, X: torch.Tensor):
+        # Evaluate true seems to be expecting
+        # a tensor of integers.
+        if X.ndim == 1:
+            X = X.unsqueeze(0)
+
+        # 1. transform to a list of strings
+        x_str = [[self.alphabet_i_to_s[int(i)] for i in x_i] for x_i in X.numpy(force=True)]
+
+        # 2. evaluate the black box
+        return torch.from_numpy(self.black_box(np.array(x_str)))
\ No newline at end of file

From 7299c8c52e716314dd76e934b6d9a5af8a47cbec Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Wed, 8 May 2024 15:03:14 +0200
Subject: [PATCH 02/21] black

---
 .../core/utils/bo_pr/poli_objective_in_pr.py  | 25 +++++++++++++------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
index 4bdbe32..eca1fa4 100644
--- a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
+++ b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
@@ -17,7 +17,10 @@
 
 from poli.core.abstract_black_box import AbstractBlackBox
 
-from discrete_mixed_bo.problems.base import DiscreteTestProblem, MultiObjectiveTestProblem
+from discrete_mixed_bo.problems.base import (
+    DiscreteTestProblem,
+    MultiObjectiveTestProblem,
+)
 
 
 class PoliObjective(DiscreteTestProblem):
@@ -61,6 +64,7 @@ class PoliMultiObjective(DiscreteTestProblem, MultiObjectiveTestProblem):
     """
     A bridge between poli black boxes and PR.
     """
+
     num_objectives: int
     _ref_point: List[float]
     _discrete_values: dict
@@ -73,15 +77,17 @@ def __init__(
         alphabet: List[str] = None,
         noise_std: float = None,
         negate: bool = False,
-        integer_indices = None,
-        integer_bounds = None,
+        integer_indices=None,
+        integer_bounds=None,
         ref_point: List[float] = None,
     ) -> None:
         self._bounds = [(0, len(alphabet) - 1) for _ in range(sequence_length)]
         self.dim = sequence_length
         self.black_box = black_box
         alphabet = alphabet or self.black_box.info.alphabet
-        self._ref_point = ref_point  #  NOTE: this assumes maximization of all objectives.
+        self._ref_point = (
+            ref_point  #  NOTE: this assumes maximization of all objectives.
+        )
         self.num_objectives = ref_point.shape[-1]
         if alphabet is None:
             raise ValueError("Alphabet must be provided.")
@@ -95,7 +101,10 @@ def __init__(
         )
         self._setup(integer_indices=integer_indices)
         self.discrete_values = BufferDict()
-        self._discrete_values = {f"pos_{i}": list(self.alphabet_s_to_i.values()) for i in range(sequence_length)}
+        self._discrete_values = {
+            f"pos_{i}": list(self.alphabet_s_to_i.values())
+            for i in range(sequence_length)
+        }
         # for k, v in self._discrete_values.items():
         #     self.discrete_values[k] = torch.tensor(v, dtype=torch.float)
         #     self.discrete_values[k] /= self.discrete_values[k].max()
@@ -110,7 +119,9 @@ def evaluate_true(self, X: torch.Tensor):
             X = X.unsqueeze(0)
 
         # 1. transform to a list of strings
-        x_str = [[self.alphabet_i_to_s[int(i)] for i in x_i] for x_i in X.numpy(force=True)]
+        x_str = [
+            [self.alphabet_i_to_s[int(i)] for i in x_i] for x_i in X.numpy(force=True)
+        ]
 
         # 2. evaluate the black box
-        return torch.from_numpy(self.black_box(np.array(x_str)))
\ No newline at end of file
+        return torch.from_numpy(self.black_box(np.array(x_str)))

From f625960c40c539b529e46c60a3766ad069aa26bb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miguel=20Gonz=C3=A1lez=20Duque?= <miguelgondu@gmail.com>
Date: Wed, 8 May 2024 08:49:45 -0500
Subject: [PATCH 03/21] Reverts to the original interface

---
 .../core/utils/bo_pr/get_problem.py           | 40 +++++++++++++++++++
 .../core/utils/bo_pr/poli_objective_in_pr.py  | 25 ++++++++----
 .../core/utils/bo_pr/run_one_replication.py   | 13 +++---
 .../bayesian_optimization/pr/solver.py        | 11 +++--
 4 files changed, 72 insertions(+), 17 deletions(-)
 create mode 100644 src/poli_baselines/core/utils/bo_pr/get_problem.py

diff --git a/src/poli_baselines/core/utils/bo_pr/get_problem.py b/src/poli_baselines/core/utils/bo_pr/get_problem.py
new file mode 100644
index 0000000..397097b
--- /dev/null
+++ b/src/poli_baselines/core/utils/bo_pr/get_problem.py
@@ -0,0 +1,40 @@
+import torch
+
+from botorch.utils.multi_objective import infer_reference_point
+
+from discrete_mixed_bo.problems.base import DiscreteTestProblem
+
+from .poli_objective_in_pr import PoliObjective, PoliMultiObjective
+
+
+def get_problem(name: str, **kwargs) -> DiscreteTestProblem:
+    r"""Initialize the test function."""
+    if name == "poli":
+        return PoliObjective(
+            black_box=kwargs["black_box"],
+            alphabet=kwargs.get("alphabet", None),
+            sequence_length=kwargs.get("sequence_length", None),
+            negate=kwargs.get("negate", False),
+        )
+    elif name == "poli_moo":
+        alphabet = kwargs.get("alphabet", None)
+        s_len = kwargs.get("sequence_length", None)
+        if s_len is None:
+            raise RuntimeError("Sequence Length None!")
+        integer_bounds = torch.zeros(2, s_len)
+        integer_bounds[1, :] = len(alphabet)
+        problem = PoliMultiObjective(
+            black_box=kwargs["black_box"],
+            alphabet=alphabet,
+            sequence_length=kwargs.get("sequence_length", None),
+            negate=kwargs.get("negate", False),
+            ref_point=infer_reference_point(
+                kwargs.get("y0", None)
+            ),  # NOTE from infer_reference_point: this assumes maximization of all objectives.
+            integer_indices=list(range(s_len)),
+            integer_bounds=integer_bounds,
+        )
+        return problem
+
+    else:
+        raise ValueError(f"Unknown function name: {name}!")
diff --git a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
index 4bdbe32..eca1fa4 100644
--- a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
+++ b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
@@ -17,7 +17,10 @@
 
 from poli.core.abstract_black_box import AbstractBlackBox
 
-from discrete_mixed_bo.problems.base import DiscreteTestProblem, MultiObjectiveTestProblem
+from discrete_mixed_bo.problems.base import (
+    DiscreteTestProblem,
+    MultiObjectiveTestProblem,
+)
 
 
 class PoliObjective(DiscreteTestProblem):
@@ -61,6 +64,7 @@ class PoliMultiObjective(DiscreteTestProblem, MultiObjectiveTestProblem):
     """
     A bridge between poli black boxes and PR.
     """
+
     num_objectives: int
     _ref_point: List[float]
     _discrete_values: dict
@@ -73,15 +77,17 @@ def __init__(
         alphabet: List[str] = None,
         noise_std: float = None,
         negate: bool = False,
-        integer_indices = None,
-        integer_bounds = None,
+        integer_indices=None,
+        integer_bounds=None,
         ref_point: List[float] = None,
     ) -> None:
         self._bounds = [(0, len(alphabet) - 1) for _ in range(sequence_length)]
         self.dim = sequence_length
         self.black_box = black_box
         alphabet = alphabet or self.black_box.info.alphabet
-        self._ref_point = ref_point  #  NOTE: this assumes maximization of all objectives.
+        self._ref_point = (
+            ref_point  #  NOTE: this assumes maximization of all objectives.
+        )
         self.num_objectives = ref_point.shape[-1]
         if alphabet is None:
             raise ValueError("Alphabet must be provided.")
@@ -95,7 +101,10 @@ def __init__(
         )
         self._setup(integer_indices=integer_indices)
         self.discrete_values = BufferDict()
-        self._discrete_values = {f"pos_{i}": list(self.alphabet_s_to_i.values()) for i in range(sequence_length)}
+        self._discrete_values = {
+            f"pos_{i}": list(self.alphabet_s_to_i.values())
+            for i in range(sequence_length)
+        }
         # for k, v in self._discrete_values.items():
         #     self.discrete_values[k] = torch.tensor(v, dtype=torch.float)
         #     self.discrete_values[k] /= self.discrete_values[k].max()
@@ -110,7 +119,9 @@ def evaluate_true(self, X: torch.Tensor):
             X = X.unsqueeze(0)
 
         # 1. transform to a list of strings
-        x_str = [[self.alphabet_i_to_s[int(i)] for i in x_i] for x_i in X.numpy(force=True)]
+        x_str = [
+            [self.alphabet_i_to_s[int(i)] for i in x_i] for x_i in X.numpy(force=True)
+        ]
 
         # 2. evaluate the black box
-        return torch.from_numpy(self.black_box(np.array(x_str)))
\ No newline at end of file
+        return torch.from_numpy(self.black_box(np.array(x_str)))
diff --git a/src/poli_baselines/core/utils/bo_pr/run_one_replication.py b/src/poli_baselines/core/utils/bo_pr/run_one_replication.py
index 60d738f..74026d1 100644
--- a/src/poli_baselines/core/utils/bo_pr/run_one_replication.py
+++ b/src/poli_baselines/core/utils/bo_pr/run_one_replication.py
@@ -37,7 +37,6 @@
     generate_initial_data,
     get_acqf,
     get_exact_rounding_func,
-    get_problem,
     initialize_model,
 )
 from discrete_mixed_bo.input import OneHotToNumeric
@@ -48,7 +47,7 @@
 )
 from discrete_mixed_bo.trust_region import TurboState, update_state
 
-from .poli_objective_in_pr import PoliObjective
+from .get_problem import get_problem
 
 supported_labels = [
     "sobol",
@@ -82,16 +81,16 @@
 
 """
 We modify this implementation slightly
-by introducing poli black boxes instead of
-function names.
+by restricting get problems to our poli-defined
+ones in .get_problem 
 """
 
 
-def run_one_replication_on_poli_black_box(
+def run_one_replication(
     seed: int,
     label: str,
     iterations: int,
-    black_box: AbstractBlackBox,
+    function_name: str,
     batch_size: int,
     mc_samples: int,
     n_initial_points: Optional[int] = None,
@@ -140,7 +139,7 @@ def run_one_replication_on_poli_black_box(
     optimization_kwargs = optimization_kwargs or {}
     # TODO: use model list when there are constraints
     # or multiple objectives
-    base_function = PoliObjective(black_box, **problem_kwargs)
+    base_function = get_problem(name=function_name, **problem_kwargs)
     base_function.to(**tkwargs)
     binary_dims = base_function.integer_indices
     binary_mask = base_function.integer_bounds[1] - base_function.integer_bounds[0] == 1
diff --git a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
index 47d21da..df339f7 100644
--- a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
+++ b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
@@ -14,7 +14,7 @@
 
 from poli_baselines.core.abstract_solver import AbstractSolver
 from poli_baselines.core.utils.bo_pr.run_one_replication import (
-    run_one_replication_on_poli_black_box,
+    run_one_replication,
 )
 
 
@@ -127,11 +127,16 @@ def solve(
         else:
             Y_init = torch.from_numpy(self.y0)
 
-        run_one_replication_on_poli_black_box(
+        if isinstance(self.black_box, MultiObjectiveBlackBox):
+            function_name = "poli_moo"
+        else:
+            function_name = "poli"
+
+        run_one_replication(
             seed=self.seed,
             label=self.label,
             iterations=max_iter,
-            black_box=self.black_box,
+            function_name=function_name,
             batch_size=self.batch_size,
             mc_samples=self.mc_samples,
             n_initial_points=self.n_initial_points,

From d65b1ae53c3e9aeab6b03ac49cbc11ce73da58f3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miguel=20Gonz=C3=A1lez=20Duque?= <miguelgondu@gmail.com>
Date: Wed, 8 May 2024 08:55:01 -0500
Subject: [PATCH 04/21] Adds the blackbox to the problem kwargs

---
 src/poli_baselines/solvers/bayesian_optimization/pr/solver.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
index df339f7..e0654e1 100644
--- a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
+++ b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
@@ -11,6 +11,7 @@
 
 
 from poli.core.abstract_black_box import AbstractBlackBox
+from poli.core.multi_objective_black_box import MultiObjectiveBlackBox
 
 from poli_baselines.core.abstract_solver import AbstractSolver
 from poli_baselines.core.utils.bo_pr.run_one_replication import (
@@ -141,6 +142,7 @@ def solve(
             mc_samples=self.mc_samples,
             n_initial_points=self.n_initial_points,
             problem_kwargs={
+                "black_box": self.black_box,
                 "sequence_length": self.sequence_length,
                 "alphabet": self.alphabet,
                 "negate": False,

From 8e82495946893ccf55e1392ca2c097795ca6ca39 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Miguel=20Gonz=C3=A1lez=20Duque?= <miguelgondu@gmail.com>
Date: Wed, 8 May 2024 09:37:01 -0500
Subject: [PATCH 05/21] Adds more checks for mooness

---
 src/poli_baselines/solvers/bayesian_optimization/pr/solver.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
index e0654e1..51b71d0 100644
--- a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
+++ b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
@@ -125,10 +125,12 @@ def solve(
 
         if self.y0 is None:
             Y_init = None
+            is_moo = None
         else:
             Y_init = torch.from_numpy(self.y0)
+            is_moo = Y_init.shape[1] > 1
 
-        if isinstance(self.black_box, MultiObjectiveBlackBox):
+        if is_moo or isinstance(self.black_box, MultiObjectiveBlackBox):
             function_name = "poli_moo"
         else:
             function_name = "poli"

From fd8803a9575044a8484526b4f90e321f519c89a3 Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Tue, 14 May 2024 14:24:02 +0200
Subject: [PATCH 06/21] add tokenizers

---
 .../solvers/bayesian_optimization/pr/solver.py      | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
index 47d21da..40572ef 100644
--- a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
+++ b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
@@ -39,6 +39,7 @@ def __init__(
         alphabet: list[str] | None = None,
         noise_std: float | None = None,
         use_fixed_noise: bool = False,
+        tokenizer: object = None,
         label: Literal[
             "sobol",
             "cont_optim__round_after__ei",
@@ -88,6 +89,7 @@ def __init__(
         self.alphabet = alphabet_
         self.alphabet_s_to_i = {s: i for i, s in enumerate(self.alphabet)}
         self.alphabet_i_to_s = {i: s for i, s in enumerate(self.alphabet)}
+        self.tokenizer = tokenizer
 
         if isinstance(x0, np.ndarray):
             # Checking that it's of the form [_, L], where
@@ -114,11 +116,14 @@ def solve(
     ):
         if self.x0 is not None:
             # We need to transform it to a tensor of integers.
-            X_init_ = [[self.alphabet_s_to_i[s] for s in x_i] for x_i in self.x0]
+            if self.tokenizer is not None: # tokenize if one provided
+                X_init_ = [[self.alphabet_s_to_i[s] for s in [s for s in self.tokenizer("".join(x_i)) if s]] for x_i in self.x0]
+            else:
+                X_init_ = [[self.alphabet_s_to_i[s] for s in x_i] for x_i in self.x0]
             X_init = torch.Tensor(X_init_).long()
-            X_init = torch.nn.functional.one_hot(X_init, len(self.alphabet)).flatten(
-                start_dim=1
-            )
+            # X_init = torch.nn.functional.one_hot(X_init, len(self.alphabet)).flatten(
+            #     start_dim=1
+            # )
         else:
             X_init = None
 

From 007c81f54985e8686301a96bb652196880827b0c Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Wed, 15 May 2024 15:38:22 +0200
Subject: [PATCH 07/21] explicit get_problem import

---
 src/poli_baselines/core/utils/bo_pr/run_one_replication.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/poli_baselines/core/utils/bo_pr/run_one_replication.py b/src/poli_baselines/core/utils/bo_pr/run_one_replication.py
index 74026d1..edb33c3 100644
--- a/src/poli_baselines/core/utils/bo_pr/run_one_replication.py
+++ b/src/poli_baselines/core/utils/bo_pr/run_one_replication.py
@@ -47,7 +47,7 @@
 )
 from discrete_mixed_bo.trust_region import TurboState, update_state
 
-from .get_problem import get_problem
+from poli_baselines.core.utils.bo_pr.get_problem import get_problem
 
 supported_labels = [
     "sobol",

From c5cfa570e19bdaa868f85c266579adf029b2465c Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Wed, 15 May 2024 15:38:57 +0200
Subject: [PATCH 08/21] added observations in kwargs

---
 src/poli_baselines/solvers/bayesian_optimization/pr/solver.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
index 7cc3c21..caf4818 100644
--- a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
+++ b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
@@ -154,6 +154,7 @@ def solve(
                 "alphabet": self.alphabet,
                 "negate": False,
                 "noise_std": self.noise_std,
+                "y0": self.y0,
             },
             model_kwargs={
                 "use_fixed_noise": self.use_fixed_noise,

From 364ad5fec7cc3fd70dfb75527e4ef2b1681ef3de Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Wed, 15 May 2024 15:40:22 +0200
Subject: [PATCH 09/21] add dedicated discrete objective, alternative to OneHot

---
 .../core/utils/bo_pr/poli_objective_in_pr.py  | 55 +++++++++++++++++--
 1 file changed, 51 insertions(+), 4 deletions(-)

diff --git a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
index eca1fa4..f48cef5 100644
--- a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
+++ b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
@@ -60,6 +60,57 @@ def evaluate_true(self, X: torch.Tensor):
         return torch.from_numpy(self.black_box(np.array(x_str)))
 
 
+class PoliDiscreteObjective(DiscreteTestProblem):
+    """
+    A bridge between poli black boxes and PR. Strictly discrete single objective - no one-hot.
+    """
+    _discrete_values: dict
+    _bounds: list
+
+
+    def __init__(
+        self,
+        black_box: AbstractBlackBox,
+        sequence_length: int,
+        alphabet: list[str] | None = None,
+        noise_std: float | None = None,
+        negate: bool = False,
+        integer_indices=None,
+    ) -> None:
+        self.dim = sequence_length
+        self.black_box = black_box
+        alphabet = alphabet or self.black_box.info.alphabet
+        if alphabet is None:
+            raise ValueError("Alphabet must be provided.")
+        if integer_indices is None:
+            integer_indices = [i for i in range(sequence_length)]
+
+        self._bounds = [(0, len(alphabet) - 1) for _ in range(sequence_length)]
+        self.alphabet_s_to_i = {s: i for i, s in enumerate(alphabet)}
+        self.alphabet_i_to_s = {i: s for i, s in enumerate(alphabet)}
+        super().__init__(noise_std, negate, categorical_indices=list(range(sequence_length)))
+        self._setup(integer_indices=integer_indices)
+        self.discrete_values = BufferDict()
+        self._discrete_values = {
+            f"pos_{i}": list(self.alphabet_s_to_i.values())
+            for i in range(sequence_length)
+        }
+        for v in self._discrete_values.values():
+            self._bounds.append((0, len(alphabet)))
+
+    def evaluate_true(self, X: torch.Tensor):
+        # Evaluate true seems to be expecting
+        # a tensor of integers.
+        if X.ndim == 1:
+            X = X.unsqueeze(0)
+
+        # 1. transform to a list of strings
+        x_str = [[self.alphabet_i_to_s[i] for i in x_i] for x_i in X.numpy(force=True)]
+
+        # 2. evaluate the black box
+        return torch.from_numpy(self.black_box(np.array(x_str)))
+
+
 class PoliMultiObjective(DiscreteTestProblem, MultiObjectiveTestProblem):
     """
     A bridge between poli black boxes and PR.
@@ -105,10 +156,6 @@ def __init__(
             f"pos_{i}": list(self.alphabet_s_to_i.values())
             for i in range(sequence_length)
         }
-        # for k, v in self._discrete_values.items():
-        #     self.discrete_values[k] = torch.tensor(v, dtype=torch.float)
-        #     self.discrete_values[k] /= self.discrete_values[k].max()
-        # super().__init__(noise_std, negate, categorical_indices=list(range(self.dim)))
         for v in self._discrete_values.values():
             self._bounds.append((0, len(alphabet)))
 

From e66c8b2d91935c44923da848387ecbcf1318da42 Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Wed, 15 May 2024 15:40:35 +0200
Subject: [PATCH 10/21] add padding element if required, example:RFP

---
 .../solvers/bayesian_optimization/pr/solver.py               | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
index caf4818..9c2ef29 100644
--- a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
+++ b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
@@ -3,6 +3,7 @@
 """
 
 from __future__ import annotations
+import logging
 
 from typing import Literal
 
@@ -87,7 +88,11 @@ def __init__(
             raise ValueError(
                 f"For this specific black box ({self.black_box.info.name}), an alphabet must be provided."
             )
+        self.add_padding_element = any(["" in x for x in x0])
         self.alphabet = alphabet_
+        if self.add_padding_element:
+            logging.warn("PADDING ADDED! Element found in x0 and added to alphabet\n THIS MAY BE UNDESIRED BEHAVIOR")
+            self.alphabet = [""] + alphabet
         self.alphabet_s_to_i = {s: i for i, s in enumerate(self.alphabet)}
         self.alphabet_i_to_s = {i: s for i, s in enumerate(self.alphabet)}
         self.tokenizer = tokenizer

From 492b5dbfdb3b8f4872a7adeebfcce93442f4a33a Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Wed, 15 May 2024 15:40:50 +0200
Subject: [PATCH 11/21] infer reference point from tensor

---
 src/poli_baselines/core/utils/bo_pr/get_problem.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/poli_baselines/core/utils/bo_pr/get_problem.py b/src/poli_baselines/core/utils/bo_pr/get_problem.py
index 397097b..6f952c8 100644
--- a/src/poli_baselines/core/utils/bo_pr/get_problem.py
+++ b/src/poli_baselines/core/utils/bo_pr/get_problem.py
@@ -29,7 +29,7 @@ def get_problem(name: str, **kwargs) -> DiscreteTestProblem:
             sequence_length=kwargs.get("sequence_length", None),
             negate=kwargs.get("negate", False),
             ref_point=infer_reference_point(
-                kwargs.get("y0", None)
+                torch.from_numpy(kwargs.get("y0", None))
             ),  # NOTE from infer_reference_point: this assumes maximization of all objectives.
             integer_indices=list(range(s_len)),
             integer_bounds=integer_bounds,

From a8f40677b28b700b8d339a492bfba1f38e8b69bb Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Wed, 15 May 2024 15:41:09 +0200
Subject: [PATCH 12/21] replace PoliObjective with PoliDiscreteObjective

---
 src/poli_baselines/core/utils/bo_pr/get_problem.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/poli_baselines/core/utils/bo_pr/get_problem.py b/src/poli_baselines/core/utils/bo_pr/get_problem.py
index 6f952c8..a7275ea 100644
--- a/src/poli_baselines/core/utils/bo_pr/get_problem.py
+++ b/src/poli_baselines/core/utils/bo_pr/get_problem.py
@@ -4,13 +4,16 @@
 
 from discrete_mixed_bo.problems.base import DiscreteTestProblem
 
-from .poli_objective_in_pr import PoliObjective, PoliMultiObjective
+from .poli_objective_in_pr import PoliObjective, PoliMultiObjective, PoliDiscreteObjective
 
 
 def get_problem(name: str, **kwargs) -> DiscreteTestProblem:
     r"""Initialize the test function."""
     if name == "poli":
-        return PoliObjective(
+        # test dimensionality if solvable:
+        dim = len(kwargs.get("alphabet", None)) * kwargs.get("sequence_length", None)
+        # objective = PoliObjective if dim < 1000 else PoliDiscreteObjective
+        return PoliDiscreteObjective(
             black_box=kwargs["black_box"],
             alphabet=kwargs.get("alphabet", None),
             sequence_length=kwargs.get("sequence_length", None),

From 7396ac87f3c6137a19a0ece6718cf27a6b560efd Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Wed, 15 May 2024 15:45:11 +0200
Subject: [PATCH 13/21] black compliant

---
 src/poli_baselines/core/utils/bo_pr/get_problem.py |  6 +++++-
 .../core/utils/bo_pr/poli_objective_in_pr.py       |  6 ++++--
 .../solvers/bayesian_optimization/pr/solver.py     | 14 +++++++++++---
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/src/poli_baselines/core/utils/bo_pr/get_problem.py b/src/poli_baselines/core/utils/bo_pr/get_problem.py
index a7275ea..667795e 100644
--- a/src/poli_baselines/core/utils/bo_pr/get_problem.py
+++ b/src/poli_baselines/core/utils/bo_pr/get_problem.py
@@ -4,7 +4,11 @@
 
 from discrete_mixed_bo.problems.base import DiscreteTestProblem
 
-from .poli_objective_in_pr import PoliObjective, PoliMultiObjective, PoliDiscreteObjective
+from .poli_objective_in_pr import (
+    PoliObjective,
+    PoliMultiObjective,
+    PoliDiscreteObjective,
+)
 
 
 def get_problem(name: str, **kwargs) -> DiscreteTestProblem:
diff --git a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
index f48cef5..f459637 100644
--- a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
+++ b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
@@ -64,10 +64,10 @@ class PoliDiscreteObjective(DiscreteTestProblem):
     """
     A bridge between poli black boxes and PR. Strictly discrete single objective - no one-hot.
     """
+
     _discrete_values: dict
     _bounds: list
 
-
     def __init__(
         self,
         black_box: AbstractBlackBox,
@@ -88,7 +88,9 @@ def __init__(
         self._bounds = [(0, len(alphabet) - 1) for _ in range(sequence_length)]
         self.alphabet_s_to_i = {s: i for i, s in enumerate(alphabet)}
         self.alphabet_i_to_s = {i: s for i, s in enumerate(alphabet)}
-        super().__init__(noise_std, negate, categorical_indices=list(range(sequence_length)))
+        super().__init__(
+            noise_std, negate, categorical_indices=list(range(sequence_length))
+        )
         self._setup(integer_indices=integer_indices)
         self.discrete_values = BufferDict()
         self._discrete_values = {
diff --git a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
index 9c2ef29..a1385f1 100644
--- a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
+++ b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
@@ -91,7 +91,9 @@ def __init__(
         self.add_padding_element = any(["" in x for x in x0])
         self.alphabet = alphabet_
         if self.add_padding_element:
-            logging.warn("PADDING ADDED! Element found in x0 and added to alphabet\n THIS MAY BE UNDESIRED BEHAVIOR")
+            logging.warn(
+                "PADDING ADDED! Element found in x0 and added to alphabet\n THIS MAY BE UNDESIRED BEHAVIOR"
+            )
             self.alphabet = [""] + alphabet
         self.alphabet_s_to_i = {s: i for i, s in enumerate(self.alphabet)}
         self.alphabet_i_to_s = {i: s for i, s in enumerate(self.alphabet)}
@@ -122,8 +124,14 @@ def solve(
     ):
         if self.x0 is not None:
             # We need to transform it to a tensor of integers.
-            if self.tokenizer is not None: # tokenize if one provided
-                X_init_ = [[self.alphabet_s_to_i[s] for s in [s for s in self.tokenizer("".join(x_i)) if s]] for x_i in self.x0]
+            if self.tokenizer is not None:  # tokenize if one provided
+                X_init_ = [
+                    [
+                        self.alphabet_s_to_i[s]
+                        for s in [s for s in self.tokenizer("".join(x_i)) if s]
+                    ]
+                    for x_i in self.x0
+                ]
             else:
                 X_init_ = [[self.alphabet_s_to_i[s] for s in x_i] for x_i in self.x0]
             X_init = torch.Tensor(X_init_).long()

From c6c92bfe82dc108c7c1e27b450e786210836a082 Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Wed, 15 May 2024 16:14:55 +0200
Subject: [PATCH 14/21] add x0 to problem

---
 src/poli_baselines/core/utils/bo_pr/get_problem.py            | 1 +
 src/poli_baselines/solvers/bayesian_optimization/pr/solver.py | 1 +
 2 files changed, 2 insertions(+)

diff --git a/src/poli_baselines/core/utils/bo_pr/get_problem.py b/src/poli_baselines/core/utils/bo_pr/get_problem.py
index 667795e..7025f30 100644
--- a/src/poli_baselines/core/utils/bo_pr/get_problem.py
+++ b/src/poli_baselines/core/utils/bo_pr/get_problem.py
@@ -40,6 +40,7 @@ def get_problem(name: str, **kwargs) -> DiscreteTestProblem:
             ),  # NOTE from infer_reference_point: this assumes maximization of all objectives.
             integer_indices=list(range(s_len)),
             integer_bounds=integer_bounds,
+            x0=kwargs.get("x0", None),
         )
         return problem
 
diff --git a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
index a1385f1..c9bb30f 100644
--- a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
+++ b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
@@ -168,6 +168,7 @@ def solve(
                 "negate": False,
                 "noise_std": self.noise_std,
                 "y0": self.y0,
+                "x0": self.x0,
             },
             model_kwargs={
                 "use_fixed_noise": self.use_fixed_noise,

From 1af9342da0fbcc95fac78158feab959e6aac8741 Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Wed, 15 May 2024 16:15:41 +0200
Subject: [PATCH 15/21] RFP compliance: exclude pad sampling, length check,
 valid length outputs

---
 .../core/utils/bo_pr/poli_objective_in_pr.py  | 27 ++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
index f459637..2f2ed63 100644
--- a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
+++ b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
@@ -126,6 +126,7 @@ class PoliMultiObjective(DiscreteTestProblem, MultiObjectiveTestProblem):
     def __init__(
         self,
         black_box: AbstractBlackBox,
+        x0: np.ndarray,
         sequence_length: int,
         alphabet: List[str] = None,
         noise_std: float = None,
@@ -133,8 +134,13 @@ def __init__(
         integer_indices=None,
         integer_bounds=None,
         ref_point: List[float] = None,
+        preserve_len: bool = True,
     ) -> None:
         self._bounds = [(0, len(alphabet) - 1) for _ in range(sequence_length)]
+        if "" == alphabet[0]:
+            self._bounds = [
+                (1, len(alphabet) - 1) for _ in range(sequence_length)
+            ]  # eliminate pad symbol from sampling
         self.dim = sequence_length
         self.black_box = black_box
         alphabet = alphabet or self.black_box.info.alphabet
@@ -142,6 +148,9 @@ def __init__(
             ref_point  #  NOTE: this assumes maximization of all objectives.
         )
         self.num_objectives = ref_point.shape[-1]
+        self.sequence_length = sequence_length
+        self.Ls = [len(x[x != ""]) for x in x0]
+        self.preserve_len = preserve_len
         if alphabet is None:
             raise ValueError("Alphabet must be provided.")
 
@@ -161,6 +170,21 @@ def __init__(
         for v in self._discrete_values.values():
             self._bounds.append((0, len(alphabet)))
 
+    def _consistent_length(self, x: List[str]):
+        valid_x = []
+        for _x in x:
+            cand_len = len(_x[_x != ""])
+            if cand_len not in self.Ls:
+                closest_len = min(
+                    self.Ls, key=lambda x: abs(x - cand_len)
+                )  # clip to closest length
+                valid_x.append(
+                    list(_x[:closest_len]) + [""] * (self.sequence_length - closest_len)
+                )
+            else:
+                valid_x.append(_x)
+        return np.vstack(valid_x)
+
     def evaluate_true(self, X: torch.Tensor):
         # Evaluate true seems to be expecting
         # a tensor of integers.
@@ -171,6 +195,7 @@ def evaluate_true(self, X: torch.Tensor):
         x_str = [
             [self.alphabet_i_to_s[int(i)] for i in x_i] for x_i in X.numpy(force=True)
         ]
-
+        if self.preserve_len:
+            x_str = self._consistent_length(x_str)
         # 2. evaluate the black box
         return torch.from_numpy(self.black_box(np.array(x_str)))

From bcead9e72985c518552306d1a3b1fd6d3f3e977c Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Thu, 16 May 2024 14:21:13 +0200
Subject: [PATCH 16/21] add tokenizer to poliobjective

---
 src/poli_baselines/core/utils/bo_pr/get_problem.py    |  3 +++
 .../core/utils/bo_pr/poli_objective_in_pr.py          | 11 +++++++----
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/poli_baselines/core/utils/bo_pr/get_problem.py b/src/poli_baselines/core/utils/bo_pr/get_problem.py
index 7025f30..c83e9f4 100644
--- a/src/poli_baselines/core/utils/bo_pr/get_problem.py
+++ b/src/poli_baselines/core/utils/bo_pr/get_problem.py
@@ -21,7 +21,10 @@ def get_problem(name: str, **kwargs) -> DiscreteTestProblem:
             black_box=kwargs["black_box"],
             alphabet=kwargs.get("alphabet", None),
             sequence_length=kwargs.get("sequence_length", None),
+            integer_indices=list(range(kwargs.get("sequence_length", None))),
             negate=kwargs.get("negate", False),
+            tokenizer=kwargs.get("tokenizer"),
+            # categorical_indices=list(range(kwargs.get("sequence_length", None))),
         )
     elif name == "poli_moo":
         alphabet = kwargs.get("alphabet", None)
diff --git a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
index 2f2ed63..3b50a31 100644
--- a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
+++ b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
@@ -9,7 +9,7 @@
 """
 
 from __future__ import annotations
-from typing import List
+from typing import List, Optional
 
 import numpy as np
 import torch
@@ -75,15 +75,18 @@ def __init__(
         alphabet: list[str] | None = None,
         noise_std: float | None = None,
         negate: bool = False,
-        integer_indices=None,
+        integer_indices: Optional[List[int]] = None,
+        categorical_indices: Optional[List[int]] = None,
+        tokenizer: object = None,
     ) -> None:
         self.dim = sequence_length
         self.black_box = black_box
+        self.tokenizer = tokenizer
         alphabet = alphabet or self.black_box.info.alphabet
         if alphabet is None:
             raise ValueError("Alphabet must be provided.")
-        if integer_indices is None:
-            integer_indices = [i for i in range(sequence_length)]
+        # if integer_indices is None:
+        #     integer_indices = [i for i in range(sequence_length)]
 
         self._bounds = [(0, len(alphabet) - 1) for _ in range(sequence_length)]
         self.alphabet_s_to_i = {s: i for i, s in enumerate(alphabet)}

From bcc5ff4458eec852a35eca4e7e4832078d7d82d0 Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Thu, 16 May 2024 14:21:23 +0200
Subject: [PATCH 17/21] account for categorical indices

---
 src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
index 3b50a31..48c0613 100644
--- a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
+++ b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
@@ -94,7 +94,7 @@ def __init__(
         super().__init__(
             noise_std, negate, categorical_indices=list(range(sequence_length))
         )
-        self._setup(integer_indices=integer_indices)
+        self._setup(integer_indices=integer_indices, categorical_indices=categorical_indices)
         self.discrete_values = BufferDict()
         self._discrete_values = {
             f"pos_{i}": list(self.alphabet_s_to_i.values())

From 345c76dc6d85db1a0e399d494f589e50301caa56 Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Thu, 16 May 2024 14:21:34 +0200
Subject: [PATCH 18/21] correct alphabet lookup index

---
 src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
index 48c0613..7b9a779 100644
--- a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
+++ b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
@@ -110,7 +110,7 @@ def evaluate_true(self, X: torch.Tensor):
             X = X.unsqueeze(0)
 
         # 1. transform to a list of strings
-        x_str = [[self.alphabet_i_to_s[i] for i in x_i] for x_i in X.numpy(force=True)]
+        x_str = [[self.alphabet_i_to_s[int(i)] for i in x_i] for x_i in X.numpy(force=True)]
 
         # 2. evaluate the black box
         return torch.from_numpy(self.black_box(np.array(x_str)))

From 19f61deded44f50b06dd21af86eb702f0f03f884 Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Thu, 16 May 2024 14:22:10 +0200
Subject: [PATCH 19/21] add tokenizer to solver parameters

---
 src/poli_baselines/solvers/bayesian_optimization/pr/solver.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
index c9bb30f..64e21fd 100644
--- a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
+++ b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
@@ -169,6 +169,7 @@ def solve(
                 "noise_std": self.noise_std,
                 "y0": self.y0,
                 "x0": self.x0,
+                "tokenizer": self.tokenizer,
             },
             model_kwargs={
                 "use_fixed_noise": self.use_fixed_noise,

From 5a2155303f0ca733e17c01b6b4892d3fb2a9e88b Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Thu, 16 May 2024 14:22:21 +0200
Subject: [PATCH 20/21] pad for inconsistent x0 shapes

---
 src/poli_baselines/solvers/bayesian_optimization/pr/solver.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
index 64e21fd..c84deab 100644
--- a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
+++ b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
@@ -134,6 +134,9 @@ def solve(
                 ]
             else:
                 X_init_ = [[self.alphabet_s_to_i[s] for s in x_i] for x_i in self.x0]
+            if not all(len(x) == len(X_init_[0]) for x in X_init_): # unequal length due to pad skip
+                max_len = max([len(x) for x in X_init_])
+                X_init_ = np.vstack([list(x) + [self.alphabet_s_to_i[""]]*int(max_len-len(x)) for x in X_init_])
             X_init = torch.Tensor(X_init_).long()
             # X_init = torch.nn.functional.one_hot(X_init, len(self.alphabet)).flatten(
             #     start_dim=1

From 4940ca4f26925c1c112381930fd6caca288ded49 Mon Sep 17 00:00:00 2001
From: Richard Michael <r.michael@posteo.net>
Date: Thu, 16 May 2024 14:22:42 +0200
Subject: [PATCH 21/21] black formatting

---
 .../core/utils/bo_pr/poli_objective_in_pr.py          |  8 ++++++--
 .../solvers/bayesian_optimization/pr/solver.py        | 11 +++++++++--
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
index 7b9a779..85bcd36 100644
--- a/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
+++ b/src/poli_baselines/core/utils/bo_pr/poli_objective_in_pr.py
@@ -94,7 +94,9 @@ def __init__(
         super().__init__(
             noise_std, negate, categorical_indices=list(range(sequence_length))
         )
-        self._setup(integer_indices=integer_indices, categorical_indices=categorical_indices)
+        self._setup(
+            integer_indices=integer_indices, categorical_indices=categorical_indices
+        )
         self.discrete_values = BufferDict()
         self._discrete_values = {
             f"pos_{i}": list(self.alphabet_s_to_i.values())
@@ -110,7 +112,9 @@ def evaluate_true(self, X: torch.Tensor):
             X = X.unsqueeze(0)
 
         # 1. transform to a list of strings
-        x_str = [[self.alphabet_i_to_s[int(i)] for i in x_i] for x_i in X.numpy(force=True)]
+        x_str = [
+            [self.alphabet_i_to_s[int(i)] for i in x_i] for x_i in X.numpy(force=True)
+        ]
 
         # 2. evaluate the black box
         return torch.from_numpy(self.black_box(np.array(x_str)))
diff --git a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
index c84deab..3d3539a 100644
--- a/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
+++ b/src/poli_baselines/solvers/bayesian_optimization/pr/solver.py
@@ -134,9 +134,16 @@ def solve(
                 ]
             else:
                 X_init_ = [[self.alphabet_s_to_i[s] for s in x_i] for x_i in self.x0]
-            if not all(len(x) == len(X_init_[0]) for x in X_init_): # unequal length due to pad skip
+            if not all(
+                len(x) == len(X_init_[0]) for x in X_init_
+            ):  # unequal length due to pad skip
                 max_len = max([len(x) for x in X_init_])
-                X_init_ = np.vstack([list(x) + [self.alphabet_s_to_i[""]]*int(max_len-len(x)) for x in X_init_])
+                X_init_ = np.vstack(
+                    [
+                        list(x) + [self.alphabet_s_to_i[""]] * int(max_len - len(x))
+                        for x in X_init_
+                    ]
+                )
             X_init = torch.Tensor(X_init_).long()
             # X_init = torch.nn.functional.one_hot(X_init, len(self.alphabet)).flatten(
             #     start_dim=1