Add beamline-friendly generation strategies

src/blop/ax/generation_strategy/__init__.py

@@ -0,0 +1,57 @@
+from ax.generation_strategy.generation_node import GenerationNode
+from ax.generation_strategy.generation_strategy import GenerationStrategy
+from ax.generation_strategy.model_spec import GeneratorSpec
+from ax.generation_strategy.transition_criterion import MinTrials
+from ax.modelbridge.registry import Generators
+from ax.models.torch.botorch_modular.surrogate import ModelConfig, SurrogateSpec
+from botorch.acquisition.logei import qLogNoisyExpectedImprovement
+from botorch.models.transforms.outcome import Log
+
+from blop.bayesian.models import LatentGP
+
+default_generation_strategy = GenerationStrategy(
+    name="Custom Generation Strategy",
+    nodes=[
+        GenerationNode(
+            node_name="Sobol",
+            model_specs=[
+                GeneratorSpec(model_enum=Generators.SOBOL, model_kwargs={"seed": 0}),
+            ],
+            transition_criteria=[
+                MinTrials(
+                    threshold=16,
+                    transition_to="LatentGP",
+                    use_all_trials_in_exp=True,
+                ),
+            ],
+        ),
+        GenerationNode(
+            node_name="LatentGP",
+            model_specs=[
+                GeneratorSpec(
+                    model_enum=Generators.BOTORCH_MODULAR,
+                    model_kwargs={
+                        "surrogate_spec": SurrogateSpec(
+                            model_configs=[
+                                ModelConfig(
+                                    botorch_model_class=LatentGP,
+                                    input_transform_classes=None,
+                                    model_options={},
+                                    outcome_transform_classes=[Log],
+                                ),
+                            ],
+                        ),
+                        "botorch_acqf_class": qLogNoisyExpectedImprovement,
+                        "acquisition_options": {},
+                    },
+                    model_gen_kwargs={
+                        "optimizer_kwargs": {
+                            "num_restarts": 10,
+                            "sequential": True,
+                        },
+                    },
+                ),
+            ],
+        ),
+    ],
+)

src/blop/bayesian/kernels.py

@@ -50,10 +50,10 @@ def __init__(
                 i = dim * torch.ones(j.shape).long()
                 skew_group_submatrix_indices.append(torch.cat((i, j, k), dim=0))
 
-        self.diag_matrix_indices: list[torch.Tensor] = [
+        self.diag_matrix_indices: list[torch.Tensor] = (
             torch.kron(torch.arange(self.num_outputs), torch.ones(self.num_inputs)).long(),
             *2 * [torch.arange(self.num_inputs).repeat(self.num_outputs)],
-        ]
+        )
 
         self.skew_matrix_indices: tuple[torch.Tensor, ...] = (
             tuple(torch.cat(skew_group_submatrix_indices, dim=1))

src/blop/bayesian/models.py

@@ -6,6 +6,11 @@
 import torch  # type: ignore[import-untyped]
 from botorch.models.gp_regression import SingleTaskGP  # type: ignore[import-untyped]
 from botorch.models.multitask import MultiTaskGP  # type: ignore[import-untyped]
+from botorch.models.transforms.input import InputTransform
+from botorch.models.transforms.outcome import OutcomeTransform
+from botorch.utils.types import DEFAULT, _DefaultType
+from gpytorch.likelihoods.likelihood import Likelihood
+from torch import Tensor
 
 from . import kernels
 
@@ -137,6 +142,11 @@ def __init__(
         self,
         train_X: torch.Tensor,
         train_Y: torch.Tensor,
+        train_Tvar: torch.Tensor = None,
+        train_Yvar: Tensor | None = None,
+        likelihood: Likelihood | None = None,
+        outcome_transform: OutcomeTransform | _DefaultType | None = DEFAULT,
+        input_transform: InputTransform | None = None,
         skew_dims: bool | list[tuple[int, ...]] = True,
         *args: Any,
         **kwargs: Any,

src/blop/tests/integration/ax/test_ax_generation_strageties.py

@@ -0,0 +1,55 @@
+from blop.ax.agent import Agent
+from blop.ax.dof import RangeDOF
+from blop.ax.generation_strategy import default_generation_strategy
+from blop.ax.objective import Objective
+from blop.sim.beamline import TiledBeamline
+
+
+def test_ax_agent_sim_beamline(RE, setup):
+    beamline = TiledBeamline(name="bl")
+    beamline.det.noise.put(False)
+
+    dofs = [
+        RangeDOF(actuator=beamline.kbv_dsv, bounds=(-5.0, 5.0), parameter_type="float"),
+        RangeDOF(actuator=beamline.kbv_usv, bounds=(-5.0, 5.0), parameter_type="float"),
+        RangeDOF(actuator=beamline.kbh_dsh, bounds=(-5.0, 5.0), parameter_type="float"),
+        RangeDOF(actuator=beamline.kbh_ush, bounds=(-5.0, 5.0), parameter_type="float"),
+    ]
+
+    objectives = [
+        Objective(name="bl_det_sum", minimize=False),
+        Objective(name="bl_det_wid_x", minimize=True),
+        Objective(name="bl_det_wid_y", minimize=True),
+    ]
+
+    def evaluation_function(uid: str, suggestions: list[dict]) -> list[dict]:
+        run = setup[uid]
+
+        bl_det_sums = run["primary/bl_det_sum"].read()
+        bl_det_wid_x = run["primary/bl_det_wid_x"].read()
+        bl_det_wid_y = run["primary/bl_det_wid_y"].read()
+
+        trial_ids = [suggestion["_id"] for suggestion in run.metadata["start"]["blop_suggestions"]]
+        outcomes = []
+        for suggestion in suggestions:
+            idx = trial_ids.index(suggestion["_id"])
+            outcome = {
+                "_id": suggestion["_id"],
+                "bl_det_sum": bl_det_sums[idx],
+                "bl_det_wid_x": bl_det_wid_x[idx],
+                "bl_det_wid_y": bl_det_wid_y[idx],
+            }
+            outcomes.append(outcome)
+
+        return outcomes
+
+    agent = Agent(
+        sensors=[beamline.det],
+        dofs=dofs,
+        objectives=objectives,
+        evaluation=evaluation_function,
+    )
+
+    agent.ax_client.set_generation_strategy(default_generation_strategy)
+
+    RE(agent.optimize(iterations=12, n_points=1))