diff --git a/docs/source/tutorials/xrt-kb-mirrors.md b/docs/source/tutorials/xrt-kb-mirrors.md index 02d94874..2c42b2ef 100644 --- a/docs/source/tutorials/xrt-kb-mirrors.md +++ b/docs/source/tutorials/xrt-kb-mirrors.md @@ -350,6 +350,77 @@ plt.title("Optimized KB Mirror Beam") plt.show() ``` +## Solving the Same Problem with XoptOptimizer + +Blop also supports Xopt through `XoptOptimizer`, which plugs directly into the +same `optimize` plan used throughout the package. This gives you a lower-level +optimization interface while reusing the same devices, acquisition flow, and +evaluation function from above. + +For this section, we keep the exact same objective and constraint: + +- Minimize `fwhm` +- Require `intensity >= 10000` + +```{code-cell} ipython3 +from xopt.generators.bayesian import ExpectedImprovementGenerator +from xopt.vocs import VOCS + +from blop.plans import optimize +from blop.protocols import OptimizationProblem +from blop.xopt.optimizer import XoptOptimizer +``` + +Define the Xopt search space (`VOCS`) from the same mirror bounds and outcome names used earlier: + +```{code-cell} ipython3 +xopt_vocs = VOCS( + variables={ + kbv.radius.name: list(VERTICAL_BOUNDS), + kbh.radius.name: list(HORIZONTAL_BOUNDS), + }, + objectives={"fwhm": "MINIMIZE"}, + constraints={"intensity": ["GREATER_THAN", 10000.0]}, +) + +xopt_optimizer = XoptOptimizer(generator=ExpectedImprovementGenerator(vocs=xopt_vocs)) + +xopt_problem = OptimizationProblem( + optimizer=xopt_optimizer, + actuators=[kbv.radius, kbh.radius], + sensors=[det], + evaluation_function=DetectorEvaluation(tiled_client), +) +``` + +Run an optimization loop. As with the Ax-based flow, this is executed via the Bluesky `RunEngine`. + +```{code-cell} ipython3 +RE(optimize(xopt_problem, iterations=10, n_points=1)) +``` + +Visualize the GP model learned by the Xopt generator: + +```{code-cell} ipython3 +fig, _ = xopt_optimizer.generator.visualize_model( + output_names=["fwhm"], + variable_names=[kbv.radius.name, kbh.radius.name], + show_feasibility=True, +) +plt.show() +``` + +Inspect the best point found by Xopt and the collected trial data: + +```{code-cell} ipython3 +xopt_best_points = xopt_optimizer.get_best_points() +xopt_best_points +``` + +```{code-cell} ipython3 +xopt_optimizer.generator.data.tail() +``` + ```{code-cell} ipython3 tiled_server.close() ``` diff --git a/pyproject.toml b/pyproject.toml index 4c7359a8..0972acd4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,7 +36,8 @@ dynamic = ["version"] [project.optional-dependencies] ax = ["ax-platform>=1.3.0,<1.4", "botorch>=0.18.0,<0.19.0", "gpytorch", "torch"] queueserver = ["bluesky-queueserver-api"] -all = ["blop[ax,queueserver]"] +xopt = ["xopt>3.0.0"] +all = ["blop[ax,queueserver,xopt]"] dev = [ "pytest", "pytest-cov", diff --git a/src/blop/tests/xopt/test_optimizer.py b/src/blop/tests/xopt/test_optimizer.py new file mode 100644 index 00000000..63433308 --- /dev/null +++ b/src/blop/tests/xopt/test_optimizer.py @@ -0,0 +1,403 @@ +from collections.abc import Callable + +import numpy as np +import pandas as pd +import pytest +from bluesky.run_engine import RunEngine +from xopt.generators.bayesian import ExpectedImprovementGenerator +from xopt.generators.random import RandomGenerator +from xopt.vocs import VOCS + +from blop.plans import optimize +from blop.protocols import OptimizationProblem +from blop.tests.conftest import MovableSignal, ReadableSignal +from blop.xopt import optimizer as xopt_optimizer_module +from blop.xopt.optimizer import XoptOptimizer + + +def _random_optimizer(vocs: VOCS, *, checkpoint_path: str | None = None) -> XoptOptimizer: + return XoptOptimizer( + generator=RandomGenerator(vocs=vocs), + checkpoint_path=checkpoint_path, + ) + + +def _bo_optimizer(vocs: VOCS, *, checkpoint_path: str | None = None) -> XoptOptimizer: + return XoptOptimizer( + generator=ExpectedImprovementGenerator(vocs=vocs), + checkpoint_path=checkpoint_path, + ) + + +@pytest.fixture(params=[_random_optimizer, _bo_optimizer], ids=["random", "bo"]) +def optimizer_factory(request: pytest.FixtureRequest) -> Callable[[VOCS], XoptOptimizer]: + return request.param + + +def test_xopt_optimizer_init(optimizer_factory: Callable[[VOCS], XoptOptimizer]): + if optimizer_factory is _bo_optimizer: + vocs = VOCS( + variables={"x1": [-5.0, 5.0], "x2": [-5.0, 5.0], "x3": [0.0, 5.0]}, + objectives={"y1": "MINIMIZE"}, + constraints={"y1": ["LESS_THAN", 10.0]}, + ) + else: + vocs = VOCS( + variables={"x1": [-5.0, 5.0], "x2": [-5.0, 5.0], "x3": [0.0, 5.0]}, + objectives={"y1": "MAXIMIZE", "y2": "MINIMIZE"}, + constraints={"y1": ["GREATER_THAN", 0.0], "y2": ["LESS_THAN", 0.0]}, + ) + + optimizer = optimizer_factory(vocs) + assert optimizer.generator is not None + assert set(optimizer.vocs.variable_names) == {"x1", "x2", "x3"} + + +def test_xopt_optimizer_suggest_ids_and_keys(): + vocs = VOCS(variables={"x1": [-5.0, 5.0], "x2": [-5.0, 5.0], "x3": [0.0, 5.0]}, objectives={"y1": "MINIMIZE"}) + optimizer = _random_optimizer(vocs) + + suggestions = optimizer.suggest(num_points=2) + assert len(suggestions) == 2 + for i, suggestion in enumerate(suggestions): + assert suggestion["_id"] == i + assert "x1" in suggestion + assert "x2" in suggestion + assert "x3" in suggestion + + +def test_xopt_optimizer_suggest_defaults_to_single_point(): + vocs = VOCS(variables={"x": [0.0, 1.0]}, objectives={"y": "MINIMIZE"}) + optimizer = _random_optimizer(vocs) + + suggestions = optimizer.suggest() + assert len(suggestions) == 1 + assert suggestions[0]["_id"] == 0 + + +def test_xopt_optimizer_first_suggest_uses_vocs_random_inputs(monkeypatch: pytest.MonkeyPatch): + vocs = VOCS(variables={"x": [0.0, 1.0]}, objectives={"y": "MINIMIZE"}) + optimizer = _bo_optimizer(vocs) + + def _random_inputs(_vocs: VOCS, n: int | None = None, **_kwargs) -> list[dict]: + assert set(_vocs.variable_names) == {"x"} + assert n == 1 + return [{"x": 0.42}] + + monkeypatch.setattr(xopt_optimizer_module, "random_inputs", _random_inputs) + + suggestions = optimizer.suggest() + + assert suggestions == [{"_id": 0, "x": 0.42}] + + +def test_xopt_optimizer_ingest_multiple_columns(optimizer_factory: Callable[[VOCS], XoptOptimizer]): + if optimizer_factory is _bo_optimizer: + vocs = VOCS( + variables={"x1": [-5.0, 5.0], "x2": [-5.0, 5.0], "x3": [0.0, 5.0]}, + objectives={"y1": "MINIMIZE"}, + ) + else: + vocs = VOCS( + variables={"x1": [-5.0, 5.0], "x2": [-5.0, 5.0], "x3": [0.0, 5.0]}, + objectives={"y1": "MAXIMIZE", "y2": "MINIMIZE"}, + ) + optimizer = optimizer_factory(vocs) + + if optimizer_factory is _bo_optimizer: + optimizer.ingest( + [ + {"x1": 0.0, "x2": 0.0, "x3": 0.0, "y1": 1.0}, + {"x1": 0.1, "x2": 0.2, "x3": 1.0, "y1": 3.0}, + ] + ) + else: + optimizer.ingest( + [ + {"x1": 0.0, "x2": 0.0, "x3": 0.0, "y1": 1.0, "y2": 2.0}, + {"x1": 0.1, "x2": 0.2, "x3": 1.0, "y1": 3.0, "y2": 4.0}, + ] + ) + + data = optimizer.generator.data + assert data is not None + assert len(data) == 2 + assert np.allclose(data["x1"].to_numpy(dtype=float), [0.0, 0.1]) + assert np.allclose(data["x2"].to_numpy(dtype=float), [0.0, 0.2]) + assert np.allclose(data["x3"].to_numpy(dtype=float), [0.0, 1.0]) + assert np.allclose(data["y1"].to_numpy(dtype=float), [1.0, 3.0]) + if optimizer_factory is _bo_optimizer: + assert "y2" not in data.columns + else: + assert np.allclose(data["y2"].to_numpy(dtype=float), [2.0, 4.0]) + + +def test_xopt_optimizer_ingest_baseline_id(optimizer_factory: Callable[[VOCS], XoptOptimizer]): + vocs = VOCS(variables={"x1": [-5.0, 5.0]}, objectives={"y1": "MINIMIZE"}) + optimizer = optimizer_factory(vocs) + + optimizer.ingest([{"x1": 0.0, "y1": 1.0, "_id": "baseline"}]) + data = optimizer.generator.data + assert data is not None + assert len(data) == 1 + assert data.iloc[0]["_id"] == "baseline" + + +def test_xopt_optimizer_seeds_state_when_existing_data_lacks_id(): + vocs = VOCS(variables={"x": [0.0, 1.0]}, objectives={"y": "MINIMIZE"}) + generator = RandomGenerator(vocs=vocs) + generator.add_data(pd.DataFrame([{"x": 0.4, "y": 1.2}])) + + optimizer = XoptOptimizer(generator=generator) + + assert optimizer._params_by_id[0] == {"x": 0.4} + assert optimizer._next_id == 1 + + +def test_xopt_optimizer_suggest_ingest(): + vocs = VOCS(variables={"x1": [-5.0, 5.0], "x2": [-5.0, 5.0]}, objectives={"y1": "MINIMIZE", "y2": "MINIMIZE"}) + optimizer = _random_optimizer(vocs) + + suggestions = optimizer.suggest(num_points=2) + outcomes = [ + {"_id": suggestions[0]["_id"], "y1": 1.0, "y2": 2.0}, + {"_id": suggestions[1]["_id"], "y1": 3.0, "y2": 4.0}, + ] + optimizer.ingest(outcomes) + + data = optimizer.generator.data + assert data is not None + assert len(data) == 2 + assert np.allclose(data["y1"].to_numpy(dtype=float), [1.0, 3.0]) + assert np.allclose(data["y2"].to_numpy(dtype=float), [2.0, 4.0]) + + +def test_xopt_optimizer_register_failures(): + vocs = VOCS(variables={"x1": [-5.0, 5.0], "x2": [-5.0, 5.0]}, objectives={"y1": "MINIMIZE"}) + optimizer = _random_optimizer(vocs) + + suggestions = optimizer.suggest(num_points=5) + optimizer.register_failures(suggestions) + + assert all(suggestion["_id"] not in optimizer._params_by_id for suggestion in suggestions) + + +def test_xopt_optimizer_checkpoint_roundtrip(tmp_path): + vocs = VOCS(variables={"x": [0.0, 1.0]}, objectives={"y": "MINIMIZE"}) + checkpoint_path = tmp_path / "xopt_optimizer.pkl" + optimizer = _random_optimizer(vocs, checkpoint_path=str(checkpoint_path)) + + suggestions = optimizer.suggest(1) + optimizer.ingest([{"_id": suggestions[0]["_id"], "y": 0.5}]) + optimizer.checkpoint() + + recovered = XoptOptimizer.from_checkpoint(str(checkpoint_path)) + assert recovered.generator.data is not None + assert len(recovered.generator.data) == 1 + assert recovered.checkpoint_path == str(checkpoint_path) + + +def test_xopt_optimizer_checkpoint_no_path(): + vocs = VOCS(variables={"x1": [-5.0, 5.0]}, objectives={"y1": "MINIMIZE"}) + optimizer = _random_optimizer(vocs) + + with pytest.raises(ValueError): + optimizer.checkpoint() + + +def test_xopt_optimizer_get_best_points_single_objective_minimize(optimizer_factory: Callable[[VOCS], XoptOptimizer]): + vocs = VOCS(variables={"x": [0.0, 1.0]}, objectives={"y": "MINIMIZE"}) + optimizer = optimizer_factory(vocs) + + optimizer.ingest( + [ + {"x": 0.1, "y": 5.0}, + {"x": 0.2, "y": 1.0}, + {"x": 0.3, "y": 3.0}, + ] + ) + + best_points = optimizer.get_best_points() + assert len(best_points) == 1 + _, params, outcomes = best_points[0] + assert params["x"] == 0.2 + assert outcomes["y"] == 1.0 + + +def test_xopt_optimizer_get_best_points_single_objective_maximize(optimizer_factory: Callable[[VOCS], XoptOptimizer]): + vocs = VOCS(variables={"x": [0.0, 1.0]}, objectives={"y": "MAXIMIZE"}) + optimizer = optimizer_factory(vocs) + + optimizer.ingest( + [ + {"x": 0.1, "y": 5.0}, + {"x": 0.2, "y": 1.0}, + {"x": 0.3, "y": 3.0}, + ] + ) + + best_points = optimizer.get_best_points() + assert len(best_points) == 1 + _, params, outcomes = best_points[0] + assert params["x"] == 0.1 + assert outcomes["y"] == 5.0 + + +def test_xopt_optimizer_get_best_points_multi_objective(): + vocs = VOCS(variables={"x": [0.0, 10.0]}, objectives={"y1": "MAXIMIZE", "y2": "MAXIMIZE"}) + optimizer = _random_optimizer(vocs) + + optimizer.ingest( + [ + {"x": 1.0, "y1": 10.0, "y2": 1.0}, + {"x": 5.0, "y1": 1.0, "y2": 10.0}, + {"x": 3.0, "y1": 2.0, "y2": 2.0}, + ] + ) + + best_points = optimizer.get_best_points() + assert len(best_points) == 3 + for trial_id, params, metrics in best_points: + assert isinstance(trial_id, (int, float, str)) + assert "x" in params + assert "y1" in metrics + assert "y2" in metrics + + +def test_xopt_optimizer_get_best_points_returns_empty_when_all_infeasible(): + vocs = VOCS( + variables={"x": [0.0, 1.0]}, + objectives={"y": "MINIMIZE"}, + constraints={"c": ["LESS_THAN", 0.0]}, + ) + optimizer = _random_optimizer(vocs) + + optimizer.ingest( + [ + {"x": 0.1, "y": 5.0, "c": 1.0}, + {"x": 0.2, "y": 1.0, "c": 2.0}, + {"x": 0.3, "y": 3.0, "c": 3.0}, + ] + ) + + best_points = optimizer.get_best_points() + assert best_points == [] + + +def test_xopt_optimizer_get_best_points_selects_best_feasible_only(): + vocs = VOCS( + variables={"x": [0.0, 1.0]}, + objectives={"y": "MINIMIZE"}, + constraints={"c": ["LESS_THAN", 0.5]}, + ) + optimizer = _random_optimizer(vocs) + + optimizer.ingest( + [ + {"x": 0.1, "y": 10.0, "c": 0.1}, + {"x": 0.2, "y": 1.0, "c": 0.9}, + {"x": 0.3, "y": 2.0, "c": 0.2}, + ] + ) + + best_points = optimizer.get_best_points() + assert len(best_points) == 1 + _, params, outcomes = best_points[0] + assert params["x"] == 0.3 + assert outcomes["y"] == 2.0 + assert outcomes["c"] == 0.2 + + +def test_xopt_optimizer_get_best_points_empty_without_data(): + vocs = VOCS(variables={"x": [0.0, 1.0]}, objectives={"y": "MINIMIZE"}) + optimizer = _random_optimizer(vocs) + + assert optimizer.get_best_points() == [] + + +def test_xopt_optimizer_get_best_points_includes_observables(): + vocs = VOCS( + variables={"x": [0.0, 1.0]}, + objectives={"y": "MINIMIZE"}, + observables=["obs"], + ) + optimizer = _random_optimizer(vocs) + + optimizer.ingest( + [ + {"x": 0.1, "y": 2.0, "obs": 10.0}, + {"x": 0.2, "y": 1.0, "obs": 20.0}, + ] + ) + + best_points = optimizer.get_best_points() + assert len(best_points) == 1 + _, _, outcomes = best_points[0] + assert outcomes["y"] == 1.0 + assert outcomes["obs"] == 20.0 + + +def test_xopt_expected_improvement_runs_simple_minimization(): + vocs = VOCS(variables={"x": [0.0, 1.0]}, objectives={"y": "MINIMIZE"}) + optimizer = XoptOptimizer(generator=ExpectedImprovementGenerator(vocs=vocs)) + + # Seed EI with initial evaluations for model training. + optimizer.ingest( + [ + {"x": 0.0, "y": (0.0 - 0.25) ** 2}, + {"x": 0.5, "y": (0.5 - 0.25) ** 2}, + {"x": 1.0, "y": (1.0 - 0.25) ** 2}, + ] + ) + + for _ in range(3): + suggestion = optimizer.suggest(1)[0] + x_val = float(suggestion["x"]) + optimizer.ingest([{"_id": suggestion["_id"], "y": (x_val - 0.25) ** 2}]) + + assert optimizer.generator.data is not None + assert len(optimizer.generator.data) == 6 + + best_points = optimizer.get_best_points() + assert len(best_points) == 1 + _, _, outcomes = best_points[0] + assert outcomes["y"] <= 0.0625 + + +def test_xopt_inside_run_engine(): + vocs = VOCS(variables={"x": [0.0, 1.0]}, objectives={"y": "MINIMIZE"}) + optimizer = XoptOptimizer(generator=ExpectedImprovementGenerator(vocs=vocs)) + + actuator = MovableSignal("x", initial_value=0.5) + sensor = ReadableSignal("y") + + def evaluation_function(_uid: str, suggestions: list[dict]) -> list[dict]: + return [ + { + "_id": suggestion["_id"], + "y": (float(suggestion["x"]) - 0.25) ** 2, + } + for suggestion in suggestions + ] + + optimization_problem = OptimizationProblem( + optimizer=optimizer, + actuators=[actuator], + sensors=[sensor], + evaluation_function=evaluation_function, + ) + + RE = RunEngine({}) + RE(optimize(optimization_problem, iterations=4, n_points=1)) + + data = optimizer.generator.data + assert data is not None + assert len(data) == 4 + assert "x" in data.columns + assert "y" in data.columns + + best_points = optimizer.get_best_points() + assert len(best_points) == 1 + _, params, outcomes = best_points[0] + assert 0.0 <= float(params["x"]) <= 1.0 + assert float(outcomes["y"]) >= 0.0 diff --git a/src/blop/xopt/__init__.py b/src/blop/xopt/__init__.py new file mode 100644 index 00000000..959745c3 --- /dev/null +++ b/src/blop/xopt/__init__.py @@ -0,0 +1,3 @@ +from .optimizer import XoptOptimizer + +__all__ = ["XoptOptimizer"] diff --git a/src/blop/xopt/optimizer.py b/src/blop/xopt/optimizer.py new file mode 100644 index 00000000..2a0423db --- /dev/null +++ b/src/blop/xopt/optimizer.py @@ -0,0 +1,260 @@ +import json +from collections.abc import Mapping +from importlib import import_module +from pathlib import Path +from typing import Any + +import pandas as pd +from xopt import VOCS +from xopt.generator import Generator +from xopt.vocs import FeasibilityError, random_inputs, select_best + +from ..protocols import ID_KEY, CanRegisterSuggestions, Checkpointable, Optimizer, TrialFaultAware + + +def _normalize_trial_id(value: Any) -> int | str: + if isinstance(value, int): + return value + if isinstance(value, float) and value.is_integer(): + return int(value) + return str(value) + + +def _is_missing_scalar(value: Any) -> bool: + return value is None or (isinstance(value, float) and pd.isna(value)) + + +class XoptOptimizer(Optimizer, Checkpointable, CanRegisterSuggestions, TrialFaultAware): + """Adapter that exposes an arbitrary Xopt generator through blop's Optimizer protocol.""" + + def __init__( + self, + generator: Generator, + *, + checkpoint_path: str | None = None, + ): + # Keep API simple: caller provides a fully configured Xopt generator instance. + self._generator = generator + + # Internal state tracks IDs, pending/known parameterizations, and checkpoint metadata. + self._checkpoint_path = checkpoint_path + self._next_id = 0 + self._params_by_id: dict[int | str, dict[str, Any]] = {} + self._seed_state_from_existing_data() + + @classmethod + def from_checkpoint(cls, checkpoint_path: str) -> "XoptOptimizer": + # Restore all persistent adapter state from JSON payload. + path = Path(checkpoint_path) + with path.open("r", encoding="utf-8") as stream: + payload = json.load(stream) + + instance = object.__new__(cls) + + generator_class_name = payload["generator"]["class"] + generator_module_name = payload["generator"]["module"] + generator_state = payload["generator"]["state"] + generator_data = payload["generator"].get("data") + + # Dynamically import the generator class from its module. + generator_module = import_module(generator_module_name) + generator_class = getattr(generator_module, generator_class_name) + instance._generator = generator_class.model_validate(generator_state) + if generator_data is not None: + instance._generator.ingest(generator_data) + + instance._checkpoint_path = str(path) + instance._next_id = payload.get("next_id", 0) + instance._params_by_id = payload.get("params_by_id", {}) + instance._seed_state_from_existing_data() + return instance + + @property + def checkpoint_path(self) -> str | None: + return self._checkpoint_path + + @property + def generator(self) -> Generator: + """Return the underlying Xopt generator instance.""" + return self._generator + + @property + def vocs(self) -> VOCS: + return self._generator.vocs + + def _seed_state_from_existing_data(self) -> None: + # Recover known trial IDs/parameters from existing generator data when available. + data = self._generator.data + if not isinstance(data, pd.DataFrame) or data.empty: + return + + for _, row in data.iterrows(): + # Reuse stored IDs when present, otherwise allocate synthetic IDs. + if ID_KEY in row and not _is_missing_scalar(row[ID_KEY]): + trial_id = _normalize_trial_id(row[ID_KEY]) + else: + trial_id = self._next_id + self._next_id += 1 + + self._params_by_id[trial_id] = {name: row[name] for name in self.vocs.variable_names if name in row} + if isinstance(trial_id, int): + self._next_id = max(self._next_id, trial_id + 1) + + def _generator_has_data(self) -> bool: + data = self._generator.data + return isinstance(data, pd.DataFrame) and not data.empty + + def suggest(self, num_points: int | None = None) -> list[dict]: + # Default to single-point suggestion when caller does not specify cardinality. + if num_points is None: + num_points = 1 + + # Bootstrap first call with random VOCS inputs to avoid model-based generators requiring prior data. + first_suggest_call = self._next_id == 0 and not self._params_by_id + has_data = self._generator_has_data() + + if first_suggest_call and not has_data: + suggestions = random_inputs(self.vocs, n=num_points) + else: + suggestions = self._generator.suggest(num_points) + return self.register_suggestions(suggestions) + + def register_suggestions(self, suggestions: list[dict]) -> list[dict]: + # Attach stable blop trial IDs and cache suggested parameterizations by ID. + registered: list[dict] = [] + for suggestion in suggestions: + trial_id = self._next_id + self._next_id += 1 + + params = {name: suggestion[name] for name in self.vocs.variable_names if name in suggestion} + self._params_by_id[trial_id] = params + registered.append({ID_KEY: trial_id, **suggestion}) + + # add the registered suggestions to the generator's data + self._generator.ingest(registered) + + return registered + + def ingest(self, points: list[dict]) -> None: + if not points: + return + + # Convert outcomes into trial rows, keeping only the latest entry per trial ID. + variable_names = list(self.vocs.variable_names) + variable_name_set = set(variable_names) + rows_by_id: dict[int | str, dict[str, Any]] = {} + + for point in points: + # Preserve provided IDs when available, else allocate a new one. + raw_trial_id = point.get(ID_KEY) + if raw_trial_id is None: + trial_id: int | str = self._next_id + self._next_id += 1 + else: + trial_id = _normalize_trial_id(raw_trial_id) + + # Merge known suggested parameters with any explicit parameters in incoming point. + point_parameters = {name: point[name] for name in variable_names if name in point} + parameters = {**self._params_by_id.get(trial_id, {}), **point_parameters} + + self._params_by_id[trial_id] = parameters + # Everything not in variables and not _id is treated as measured output. + outcomes = {k: v for k, v in point.items() if k not in variable_name_set | {ID_KEY}} + rows_by_id[trial_id] = {ID_KEY: trial_id, **parameters, **outcomes} + + # Update existing trial rows in-place by ID; only append truly new IDs. + data = self._generator.data + if not (isinstance(data, pd.DataFrame) and not data.empty and ID_KEY in data.columns): + # Persist all observations when no existing data is present. + self._generator.ingest(list(rows_by_id.values())) + return + + # Build a mapping from trial ID to row indices in the existing generator data. + index_by_id: dict[int | str, list] = {} + for index, raw_trial_id in data[ID_KEY].items(): + trial_id = _normalize_trial_id(raw_trial_id) + index_by_id.setdefault(trial_id, []).append(index) + + # Update existing rows in-place and collect new rows to append. + rows_to_append: list[dict[str, Any]] = [] + for trial_id, row in rows_by_id.items(): + indices = index_by_id.get(trial_id) + if indices is None: + rows_to_append.append(row) + continue + + for key, value in row.items(): + data.loc[indices, key] = value + + # Append any new rows to the generator's data after in-place updates. + if rows_to_append: + self._generator.ingest(rows_to_append) + + def register_failures(self, suggestions: list[dict]) -> None: + # Remove failed suggestions from pending parameter cache. + for suggestion in suggestions: + trial_id = suggestion.get(ID_KEY) + if trial_id is not None: + self._params_by_id.pop(trial_id, None) + + def get_best_points(self) -> list[tuple[int | str, Mapping, Mapping]]: + # Return no points when no data has been ingested. + if not self._generator_has_data(): + return [] + + data = self._generator.data + if not isinstance(data, pd.DataFrame): + return [] + + objective_names = list(self.vocs.objective_names) + # For single-objective problems, return the single extremum according to direction. + if len(objective_names) == 1 and objective_names[0] in data: + try: + best_indices, _, _ = select_best(self.vocs, data, n=1) + best_index = best_indices[0] + selected = data.loc[[best_index]] + except FeasibilityError: + # If no feasible points exist, return an empty list. + return [] + else: + # For multi-objective and objective-less modes, return the available candidate set. + selected = data + + output_names = self.vocs.output_names + results: list[tuple[int | str, Mapping, Mapping]] = [] + for _, row in selected.iterrows(): + # Normalize IDs and split into parameter and outcome mappings. + trial_id = _normalize_trial_id(row[ID_KEY] if ID_KEY in row else _) + + parameterization = {name: row[name] for name in self.vocs.variable_names if name in row} + outcomes = {name: row[name] for name in output_names if name in row} + results.append((trial_id, parameterization, outcomes)) + + return results + + def checkpoint(self) -> None: + """Dump serialized optimizer state to the configured checkpoint JSON file.""" + # Enforce explicit checkpoint path configuration before writing state. + if not self._checkpoint_path: + raise ValueError("Checkpoint path is not set. Please set a checkpoint path when initializing the optimizer.") + + # Persist generator and adapter bookkeeping to a single JSON artifact. + payload = { + "generator": { + "class": self._generator.__class__.__name__, + "module": self._generator.__class__.__module__, + "state": self._generator.model_dump(), + "data": ( + self._generator.data.to_dict(orient="records") + if isinstance(self._generator.data, pd.DataFrame) + else self._generator.data + ), + }, + "next_id": self._next_id, + "params_by_id": self._params_by_id, + } + + # Write the payload to the configured checkpoint path. + path = Path(self._checkpoint_path) + with path.open("w", encoding="utf-8") as stream: + json.dump(payload, stream, default=str)