diff --git a/src/eaa/task_managers/tuning/bo.py b/src/eaa/task_managers/tuning/bo.py new file mode 100644 index 0000000..9e5733a --- /dev/null +++ b/src/eaa/task_managers/tuning/bo.py @@ -0,0 +1,120 @@ +import logging +from typing import Optional, Callable + +import torch + +from eaa.task_managers.base import BaseTaskManager +from eaa.tools.base import BaseTool +from eaa.tools.bo import BayesianOptimizationTool +from eaa.api.llm_config import LLMConfig + +logger = logging.getLogger(__name__) + + +class BayesianOptimizationTaskManager(BaseTaskManager): + + def __init__( + self, + llm_config: LLMConfig = None, + tools: list[BaseTool] = [], + bayesian_optimization_tool: BayesianOptimizationTool = None, + initial_points: Optional[torch.Tensor] = None, + n_initial_points: int = 20, + objective_function: Callable = None, + message_db_path: Optional[str] = None, + build: bool = True, + *args, **kwargs + ) -> None: + """Bayesian optimization task manager. + + Parameters + ---------- + llm_config : LLMConfig, optional + The configuration for the LLM. + tools : list[BaseTool], optional + A list of tools for the agent. This should NOT include the + `BayesianOptimizationTool`. + bayesian_optimization_tool : BayesianOptimizationTool + The Bayesian optimization tool to use. + initial_points : torch.Tensor, optional + A (n_points, n_features) tensor giving the initial points where + the objective function should be evaluated to initialize the + Gaussian process model. If None, random initial points will be + generated. + n_initial_points : int, optional + The number of initial points to generate if `initial_points` is None. + objective_function : Callable + The objective function to be maximized. This function should take + a single argument, which is a (n_points, n_features) tensor of + points to evaluate the objective function at. It should return + a (n_points, n_objectives) tensor of objective function values. + message_db_path : Optional[str] + If provided, the entire chat history will be stored in + a SQLite database at the given path. This is essential + if you want to use the WebUI, which polls the database + for new messages. + build : bool, optional + Whether to build the internal state of the task manager. + """ + if bayesian_optimization_tool is None: + raise ValueError( + "Bayesian optimization tool should be explicitly passed to " + "`bayesian_optimization_tool`." + ) + if objective_function is None: + raise ValueError("`objective_function` is required.") + + self.bayesian_optimization_tool = bayesian_optimization_tool + + for tool in tools: + if isinstance(tool, BayesianOptimizationTool): + raise ValueError( + "`BayesianOptimizationTool` should not be included in `tools`. " + "Instead, pass it to `bayesian_optimization_tool`." + ) + + self.objective_function = objective_function + + self.initial_points = initial_points + self.n_initial_points = n_initial_points + + super().__init__( + llm_config=llm_config, + tools=tools, + message_db_path=message_db_path, + build=build, + *args, **kwargs + ) + + def run( + self, + n_iterations: int = 50, + *args, **kwargs + ) -> None: + """Run Bayesian optimization. Upon the second or later call, + this function continues from the last iteration. + + Parameters + ---------- + n_iterations : int, optional + The number of iterations to run. + """ + if len(self.bayesian_optimization_tool.xs_untransformed) == 0: + if self.initial_points is None: + xs_init = self.bayesian_optimization_tool.get_random_initial_points(n_points=self.n_initial_points) + else: + xs_init = self.initial_points + logger.info(f"Initial points (shape: {xs_init.shape}):\n{xs_init}") + + for x in xs_init: + x = x[None, :] + y = self.objective_function(x) + self.bayesian_optimization_tool.update(x, y) + self.bayesian_optimization_tool.build() + + for i in range(n_iterations): + candidates = self.bayesian_optimization_tool.suggest(n_suggestions=1) + logger.info(f"Candidate suggested: {candidates[0]}") + y = self.objective_function(candidates) + logger.info(f"Objective function value: {y.item()}") + self.bayesian_optimization_tool.update(candidates, y) diff --git a/src/eaa/task_managers/tuning/bo_mic_optics.py b/src/eaa/task_managers/tuning/bo_mic_optics.py new file mode 100644 index 0000000..4ce5fc1 --- /dev/null +++ b/src/eaa/task_managers/tuning/bo_mic_optics.py @@ -0,0 +1,148 @@ +import logging +from typing import Optional + +import torch +import numpy as np +from PIL import Image + +from eaa.task_managers.tuning.bo import BayesianOptimizationTaskManager +from eaa.task_managers.imaging.feature_tracking import FeatureTrackingTaskManager +from eaa.tools.base import BaseTool +from eaa.tools.bo import BayesianOptimizationTool +from eaa.api.llm_config import LLMConfig + +logger = logging.getLogger(__name__) + + +class MicroscopyOpticsTuningBOTaskManager( + BayesianOptimizationTaskManager, + FeatureTrackingTaskManager +): + def __init__( + self, + llm_config: LLMConfig = None, + image_acquisition_tool: BaseTool = None, + parameter_setting_tool: BaseTool = None, + bayesian_optimization_tool: BayesianOptimizationTool = None, + initial_points: Optional[torch.Tensor] = None, + n_initial_points: int = 20, + image_acquisition_kwargs: dict = {}, + feature_tracking_kwargs: dict = {}, + message_db_path: Optional[str] = None, + *args, **kwargs + ): + """The Bayesian optimization task manager for microscopy optics tuning. + + Parameters + ---------- + llm_config : LLMConfig, optional + The configuration for the LLM. + bayesian_optimization_tool : BayesianOptimizationTool + The Bayesian optimization tool to use. + initial_points : torch.Tensor, optional + A (n_points, n_features) tensor giving the initial points where + the objective function should be evaluated to initialize the + Gaussian process model. If None, random initial points will be + generated. + n_initial_points : int, optional + The number of initial points to generate if `initial_points` is None. + image_acquisition_kwargs : dict, optional + The arguments of the image acquisition tool that should be used + when acquiring images for evaluating the objective function. + feature_tracking_kwargs : dict, optional + The arguments of the feature tracking task manager's `run` method. + message_db_path : Optional[str] + If provided, the entire chat history will be stored in + a SQLite database at the given path. This is essential + if you want to use the WebUI, which polls the database + for new messages. + build : bool, optional + Whether to build the internal state of the task manager. + """ + if bayesian_optimization_tool is None: + raise ValueError( + "Bayesian optimization tool should be explicitly passed to " + "`bayesian_optimization_tool`." + ) + if image_acquisition_tool is None: + raise ValueError( + "Image acquisition tool should be explicitly passed to " + "`image_acquisition_tool`." + ) + if parameter_setting_tool is None: + raise ValueError( + "Parameter setting tool should be explicitly passed to " + "`parameter_setting_tool`." + ) + + self.image_acquisition_tool = image_acquisition_tool + self.parameter_setting_tool = parameter_setting_tool + self.image_acquisition_kwargs = image_acquisition_kwargs + self.feature_tracking_kwargs = feature_tracking_kwargs + + BayesianOptimizationTaskManager.__init__( + self, + llm_config=llm_config, + tools=[], + bayesian_optimization_tool=bayesian_optimization_tool, + initial_points=initial_points, + n_initial_points=n_initial_points, + objective_function=self.objective_function, + message_db_path=message_db_path, + build=False, + *args, **kwargs + ) + FeatureTrackingTaskManager.__init__( + self, + llm_config=llm_config, + tools=[image_acquisition_tool], + build=True, + *args, **kwargs + ) + + def objective_function(self, x: torch.Tensor, *args, **kwargs): + """Calculate the objective function value. + + Parameters + ---------- + x : torch.Tensor + A (n_points, n_features) tensor of points to evaluate the + objective function at. + + Returns + ------- + torch.Tensor + A (n_points, 1) tensor of objective function values. + """ + if x.ndim != 2: + raise ValueError( + "`x` should be a 2D tensor of shape (n_points, n_features)." + ) + + objective_values = torch.zeros(x.shape[0], 1, device=x.device) + + for i, x_i in enumerate(x): + # Acquire an image with the current parameters. It will be used + # as the reference image for feature tracking. + acquired_image_path = self.image_acquisition_tool.acquire_image( + **self.image_acquisition_kwargs + ) + + # Apply parameters. + self.parameter_setting_tool.set_parameters(x_i) + + # Now the original feature will have drifted. Run feature tracking + # to bring it back. + self.run_feature_tracking( + **self.feature_tracking_kwargs + ) + + # Get a new image after feature tracking. + acquired_image_path = self.image_acquisition_tool.acquire_image( + **self.image_acquisition_kwargs + ) + image = Image.open(acquired_image_path) + image = np.array(image) + + objective_values[i, 0] = np.std(image) + return objective_values diff --git a/tests/test_bo.py b/tests/test_bo.py index 46a88cf..6beaf87 100644 --- a/tests/test_bo.py +++ b/tests/test_bo.py @@ -8,6 +8,7 @@ import torch from eaa.tools.bo import BayesianOptimizationTool +from eaa.task_managers.tuning.bo import BayesianOptimizationTaskManager import test_utils as tutils @@ -85,6 +86,41 @@ def objective_function(x: torch.Tensor) -> torch.Tensor: final_suggestion = candidates[0] assert torch.allclose(final_suggestion.float(), torch.tensor([1.0, 2.0]), rtol=0.1) + + def test_bo_task_manager(self): + def objective_function(x: torch.Tensor) -> torch.Tensor: + # Expected input shape: (n_samples, n_features) + # Maximum: x = [1, 2] + y = torch.exp(-((x[:, 0] - 1) ** 2 + (x[:, 1] - 2) ** 2) / (2 * 10 ** 2)) + return y[:, None] + + tutils.set_seed(42) + + bo_tool = BayesianOptimizationTool( + bounds=([-10, -10], [10, 10]), + acquisition_function_class=botorch.acquisition.LogExpectedImprovement, + acquisition_function_kwargs={ + "best_f": -100 + }, + model_class=botorch.models.SingleTaskGP, + model_kwargs={ + "covar_module": gpytorch.kernels.MaternKernel( + nu=2.5, + ) + }, + optimization_function=botorch.optim.optimize_acqf, + ) + + task_manager = BayesianOptimizationTaskManager( + llm_config=None, + bayesian_optimization_tool=bo_tool, + n_initial_points=20, + objective_function=objective_function, + ) + task_manager.run(n_iterations=20) + + final_suggestion = task_manager.bayesian_optimization_tool.xs_untransformed[-1] + assert torch.allclose(final_suggestion.float(), torch.tensor([1.0, 2.0]), rtol=0.1) if __name__ == '__main__': @@ -94,3 +130,4 @@ def objective_function(x: torch.Tensor) -> torch.Tensor: tester = TestBayesianOptimization() tester.setup_method(name="", generate_data=False, generate_gold=False, debug=True) tester.test_bayesian_optimization() + tester.test_bo_task_manager()