From c4ea87e4c7ca4233b37fb3783188b5a884171926 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 18 Dec 2025 11:37:27 -0500 Subject: [PATCH 1/2] add generation strategies with custom models --- src/blop/ax/generation_strategy/__init__.py | 57 +++++++++++++++++++ src/blop/bayesian/kernels.py | 4 +- src/blop/bayesian/models.py | 10 ++++ .../ax/test_ax_generation_strageties.py | 55 ++++++++++++++++++ 4 files changed, 124 insertions(+), 2 deletions(-) create mode 100644 src/blop/ax/generation_strategy/__init__.py create mode 100644 src/blop/tests/integration/ax/test_ax_generation_strageties.py diff --git a/src/blop/ax/generation_strategy/__init__.py b/src/blop/ax/generation_strategy/__init__.py new file mode 100644 index 00000000..093c2867 --- /dev/null +++ b/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, + }, + }, + ), + ], + ), + ], +) diff --git a/src/blop/bayesian/kernels.py b/src/blop/bayesian/kernels.py index f86230eb..56a2e2f3 100644 --- a/src/blop/bayesian/kernels.py +++ b/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] = tuple( 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)) diff --git a/src/blop/bayesian/models.py b/src/blop/bayesian/models.py index 2dc8baa1..69af5ff4 100644 --- a/src/blop/bayesian/models.py +++ b/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, diff --git a/src/blop/tests/integration/ax/test_ax_generation_strageties.py b/src/blop/tests/integration/ax/test_ax_generation_strageties.py new file mode 100644 index 00000000..96810760 --- /dev/null +++ b/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)) From 6ac9f1d42420fd357043e1dfcfa637b60013e060 Mon Sep 17 00:00:00 2001 From: Thomas Morris Date: Thu, 18 Dec 2025 16:51:37 -0500 Subject: [PATCH 2/2] explicit tuple --- src/blop/bayesian/kernels.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/blop/bayesian/kernels.py b/src/blop/bayesian/kernels.py index 56a2e2f3..3c5cb408 100644 --- a/src/blop/bayesian/kernels.py +++ b/src/blop/bayesian/kernels.py @@ -50,7 +50,7 @@ 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] = tuple( + 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)], )