diff --git a/docs/source/tutorials.rst b/docs/source/tutorials.rst index 502efbe3..6339da1a 100644 --- a/docs/source/tutorials.rst +++ b/docs/source/tutorials.rst @@ -5,6 +5,7 @@ Tutorials :maxdepth: 1 tutorials/simple-experiment.md + tutorials/gradient-optimization.md tutorials/queueserver.md tutorials/xrt-demo.md tutorials/xrt-kb-mirrors.md diff --git a/docs/source/tutorials/gradient-optimization.md b/docs/source/tutorials/gradient-optimization.md new file mode 100644 index 00000000..15ba7807 --- /dev/null +++ b/docs/source/tutorials/gradient-optimization.md @@ -0,0 +1,244 @@ +--- +jupytext: + text_representation: + extension: .md + format_name: myst + format_version: 0.13 + jupytext_version: 1.17.3 +kernelspec: + display_name: dev + language: python + name: python3 +--- + +# Your first Scipy optimization with Blop + +In this tutorial, you will learn the three core concepts of Blop: **DOFs** (the parameters you can adjust), **objectives** (what you want to optimize), and the **Agent** (which coordinates the optimization). We'll optimize a simple mathematical function using simulated devices—the same patterns apply to real hardware. + +## Setup + +First, let's import what we need and start the data infrastructure: + +```{code-cell} ipython3 +import logging +import time +from typing import Any + +from bluesky.protocols import HasHints, HasParent, Hints, NamedMovable, Readable, Status +from bluesky.run_engine import RunEngine +from bluesky_tiled_plugins import TiledWriter +from tiled.client import from_uri +from tiled.client.container import Container +from tiled.server import SimpleTiledServer + +from blop.scipy import SCP, ScipyCFG, Objective, RangeDOF, Scipy + +# Suppress noisy logs from httpx +logging.getLogger("httpx").setLevel(logging.WARNING) +``` + +```{code-cell} ipython3 +# Start a local Tiled server for data storage +tiled_server = SimpleTiledServer() + +# Set up the Bluesky RunEngine and connect it to Tiled +RE = RunEngine({}) +tiled_client = from_uri(tiled_server.uri) +tiled_writer = TiledWriter(tiled_client) +RE.subscribe(tiled_writer) +``` + +## Creating simulated devices + +Bluesky controls devices through protocols. For this tutorial, we create simple simulated "movable" devices. In real experiments, you would use [Ophyd](https://blueskyproject.io/ophyd-async) devices or similar—the code below is just boilerplate to simulate hardware: + +```{code-cell} ipython3 +class AlwaysSuccessfulStatus(Status): + def add_callback(self, callback) -> None: + callback(self) + + def exception(self, timeout=0.0): + return None + + @property + def done(self) -> bool: + return True + + @property + def success(self) -> bool: + return True + + +class ReadableSignal(Readable, HasHints, HasParent): + def __init__(self, name: str) -> None: + self._name = name + self._value = 0.0 + + @property + def name(self) -> str: + return self._name + + @property + def hints(self) -> Hints: + return {"fields": [self._name], "dimensions": [], "gridding": "rectilinear"} + + @property + def parent(self) -> Any | None: + return None + + def read(self): + return {self._name: {"value": self._value, "timestamp": time.time()}} + + def describe(self): + return {self._name: {"source": self._name, "dtype": "number", "shape": []}} + + +class MovableSignal(ReadableSignal, NamedMovable): + def __init__(self, name: str, initial_value: float = 0.0) -> None: + super().__init__(name) + self._value: float = initial_value + + def set(self, value: float) -> Status: + self._value = value + return AlwaysSuccessfulStatus() +``` + +## Defining DOFs and objectives + +**DOFs** (degrees of freedom) are the parameters the optimizer can adjust. **Objectives** are what you want to optimize. Here we define two DOFs (`x1` and `x2`) that can range from -5 to 5, and one objective (the Himmelblau function) that we want to minimize: + +```{code-cell} ipython3 +x1 = MovableSignal("x1", initial_value=0.1) +x2 = MovableSignal("x2", initial_value=0.23) + +dofs = [ + RangeDOF(actuator=x1, bounds=(-5, 5), parameter_type="float"), + RangeDOF(actuator=x2, bounds=(-5, 5), parameter_type="float"), +] +objectives = [ + Objective(name="himmelblau_2d", minimize=True), +] +sensors = [] +``` + +## Writing the evaluation function + +The **evaluation function** computes objective values from experimental data. After each run, Blop calls this function with the run's unique ID and the suggestions that were tried. It returns the computed objective values: + +```{code-cell} ipython3 +class Himmelblau2DEvaluation: + def __init__(self, tiled_client: Container): + self.tiled_client = tiled_client + + def __call__(self, uid: str, suggestions: list[dict]) -> list[dict]: + run = self.tiled_client[uid] + outcomes = [] + reordered_suggestions = run.start["blop_suggestions"] + x1_data = run["primary/x1"].read() + x2_data = run["primary/x2"].read() + + print( + "[Himmelblau] evaluating suggestions: ", + [s["_id"] for s in suggestions], + " reordered to: ", + [s["_id"] for s in reordered_suggestions], + ) + for index, suggestion in enumerate(reordered_suggestions): + # Special key to identify a suggestion + suggestion_id = suggestion["_id"] + x1 = x1_data[index] + x2 = x2_data[index] + # Himmelblau function: has four global minima where value = 0 + outcomes.append({"himmelblau_2d": (x1**2 + x2 - 11) ** 2 + (x1 + x2**2 - 7) ** 2, "_id": suggestion_id}) + + return outcomes +``` + +## Running the optimization + +The **Agent** brings everything together. Create one with your DOFs, objectives, and evaluation function, then run the optimization: + +```{code-cell} ipython3 +agent = Scipy.Agent( + sensors=sensors, + dofs=dofs, + objectives=objectives, + evaluation_function=Himmelblau2DEvaluation(tiled_client=tiled_client), + name="simple-experiment", + description="A simple experiment optimizing the Himmelblau function", +) + +RE(agent.optimize(10)) +``` + +## Configuring the optimization + +Sometimes a default **Agent** optimization may not do all that you'd like. We expose a configuration object called ScipyCFG and a pure scipy interface so that the classic parameters of scipy minimize can be tweaked (and some multipoint sampling can be used). + +```{code-cell} ipython3 +config = ScipyCFG(dofs=dofs, objective=objectives[0], optimizer=SCP.DUAL_ANNEALING, threads=4, max_iter=2, eps=0.1) +agent = Scipy( + sensors=sensors, + config=config, + evaluation_function=Himmelblau2DEvaluation(tiled_client=tiled_client), + name="test_experiment", +) +res_uid = RE(agent.optimize(20, n_points=2)) +``` + +## Viewing the results + +Scipy is a local optimizer so it doesn't have internal point tracking, but we can to grab it from our datastore. + +```{code-cell} ipython3 +import numpy as np +import matplotlib.pyplot as plt +import pandas as pd + +res_client = tiled_client[res_uid[0]] +data = res_client["primary/internal"].read() +cols = ["suggestion_ids", "x1", "x2", "himmelblau_2d"] +vec = data[cols] +res = [] +for _, row in vec.iterrows(): + vic = [row.suggestion_ids, row.x1, row.x2, row.himmelblau_2d] + vic = [x.strip("[]").split() for x in vic] + for id, x, y, obj in zip(*vic, strict=True): + if id != "''": + res.append([int(id.strip("'")), float(x), float(y), float(obj)]) +res = np.array(res) + +fig, ax = plt.subplots(figsize=(12, 8)) + +xb, yb = np.random.uniform(-5, 5, (2, 1000)) +ax.tripcolor(xb, yb, (xb**2 + yb - 11) ** 2 + (xb + yb**2 - 7) ** 2, shading="gouraud") + +i, x, y, z = res.T +ps = ax.scatter(x, y, c=range(len(x)), cmap="plasma", s=50) +plt.colorbar(ps).set_label("sample index") +plt.title("Visualizing Scipy's traversal of Himmelblau") +``` + +Seeing the sample history + +```{code-cell} ipython3 +pd.DataFrame(data=res, columns=cols) +``` + +```{code-cell} ipython3 +print(agent.get_best_points()) +``` + +The Himmelblau function has four global minima (all with value 0). The `summarize` output shows which one(s) the optimizer found. + +## What you learned + +You now understand the three core concepts of Blop: + +- **DOFs**: The parameters the optimizer adjusts (here, `x1` and `x2` with bounds) +- **Objectives**: What you're optimizing (here, minimizing the Himmelblau function) +- **Agent**: Coordinates the optimization loop between Bluesky and the evaluation function + +## Next steps + +For a more comprehensive tutorial with multiple objectives and diagnostic tools, see [Optimizing KB Mirrors](./xrt-kb-mirrors.md). diff --git a/src/blop/scipy/__init__.py b/src/blop/scipy/__init__.py new file mode 100644 index 00000000..6b0c7389 --- /dev/null +++ b/src/blop/scipy/__init__.py @@ -0,0 +1,18 @@ +"""Scipy Backend for Pertubative gradient and in house global optimizers.""" + +from .configs import SCP, Objective, RangeDOF, ScipyCFG +from .inverter import InteractiveOptimizer +from .normalizers import SHGO, DualAnnealing, Minimize +from .scipy import Scipy + +__all__ = [ + "SCP", + "ScipyCFG", + "Scipy", + "DualAnnealing", + "Minimize", + "SHGO", + "InteractiveOptimizer", + "Objective", + "RangeDOF", +] diff --git a/src/blop/scipy/configs.py b/src/blop/scipy/configs.py new file mode 100644 index 00000000..d8b6ca4a --- /dev/null +++ b/src/blop/scipy/configs.py @@ -0,0 +1,151 @@ +"""Collection of data and configuration objects used by scipy package.""" + +from collections.abc import Sequence +from dataclasses import dataclass +from enum import StrEnum +from typing import Literal, cast + +from scipy.optimize import Bounds + +from blop.protocols import Actuator + + +@dataclass(frozen=True, kw_only=True, eq=False) +class RangeDOF: + """ + A degree of freedom that is a continuous range. + + Use this class for continuous parameters that can take any value within + specified bounds, such as motor positions, voltages, or temperatures. + + Attributes + ---------- + bounds : tuple[float, float] + The search domain of the DOF as (lower_bound, upper_bound). + parameter_type : Literal["float", "int"] + The data type of the DOF. Use "float" for continuous values or "int" for integer values. + step_size : float | None, optional + The step size of the DOF. If provided, the optimizer will only suggest values + at multiples of this step size. + scaling : Literal["linear", "log"] | None, optional + The scaling of the DOF. Use "log" for parameters that span orders of magnitude. + + Examples + -------- + Define a continuous DOF with a name (for non-actuator parameters): + + >>> from blop.scipy.configs import RangeDOF + >>> dof = RangeDOF(name="voltage", bounds=(-10.0, 10.0), parameter_type="float") + + Define an integer DOF with a step size: + + >>> dof = RangeDOF(name="num_exposures", bounds=(1, 100), parameter_type="int", step_size=1) + + For examples with actuators, see :doc:`/tutorials/simple-experiment`. + """ + + name: str | None = None + actuator: Actuator | str | None = None + bounds: tuple[float, float] + parameter_type: Literal["float", "int"] + step_size: float | None = None + scaling: Literal["linear", "log"] | None = None + + @property + def parameter_name(self) -> str: + """The parameter name used internally by Ax.""" + if isinstance(self.actuator, Actuator): + param_name = self.actuator.name + elif isinstance(self.actuator, str): + param_name = self.actuator + else: + param_name = cast(str, self.name) + return param_name + + def to_scipy_bounds(self) -> Bounds: + """Convert DOF to the Scipy equivalent Bounds.""" + return Bounds(lb=self.bounds[0], ub=self.bounds[1]) + + +@dataclass(frozen=True, kw_only=True) +class Objective: + """ + An objective to optimize. + + An objective represents a measurable outcome that you want to optimize. + The optimizer will try to minimize or maximize this outcome based on the + acquired data and evaluation function. + + Attributes + ---------- + name : str + The name of the objective. This must match the key returned by the + evaluation function for this outcome. + minimize : bool + Whether to minimize or maximize the objective. Set to True for minimization + (e.g., reducing beam width) or False for maximization (e.g., increasing intensity). + + Examples + -------- + Define an objective to maximize beam intensity: + + >>> from blop.scipy.objective import Objective + >>> objective = Objective(name="beam_intensity", minimize=False) + + Define an objective to minimize beam width: + + >>> objective = Objective(name="beam_width", minimize=True) + """ + + name: str + minimize: bool + + +class SCP(StrEnum): + """Enumeration of all optimizers currently supported/tested. + + #TODO all commented optimizers require jacobian and are currently suspended in impl until its clear that external + # gradient sampling is necesary and clearly cross implementable. likely necesary for noisy opts but this is clearly + # a usage defined addition. + """ + + Default = "L-BFGS-B" + + NELDER_MEAD = "Nelder-Mead" + POWELL = "Powell" + CG = "CG" + BFGS = "BFGS" + # Newton_CG = "Newton-CG" + LBFGS = "L-BFGS-B" + TNC = "TNC" + COBYLA = "COBYLA" + COBYQA = "COBYQA" + SLSQP = "SLSQP" + TRUST_CONSTR = "trust-constr" + # DOGLEG = "dogleg" + # Trust_NCG = "trust-ncg" + # Trust_Exact = "trust-exact" + # Trust_Krylov = "trust-krylov" + + DUAL_ANNEALING = "dual annealing" + SHGO = "SHGO" + + +@dataclass +class ScipyCFG: + """ + Configuration dataclass that encompasses the core optimization problem and extra parameters within Scipy. + + Used as the optimizer/generation function is not injectable like in Ax + """ + + dofs: Sequence[RangeDOF] + objective: Objective + # dof_constraints: Sequence[DOFConstraint] | None = None + # outcome_constraints: Sequence[OutcomeConstraint] | None = None + optimizer: SCP = SCP.Default + initial: Sequence[float] | None = None + rescale: Sequence[float] | float | None = None + max_iter: int | None = 100 + eps: float | None = None + threads: int | None = None diff --git a/src/blop/scipy/inverter.py b/src/blop/scipy/inverter.py new file mode 100644 index 00000000..a1a82ed6 --- /dev/null +++ b/src/blop/scipy/inverter.py @@ -0,0 +1,219 @@ +"""Core Scipy optimizer porting scipy algorithms.""" + +from collections import OrderedDict +from collections.abc import Mapping +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass +from threading import Event, Thread +from typing import Any, cast + +import numpy as np +from scipy.optimize import OptimizeResult + +from blop.protocols import ID_KEY, Optimizer +from blop.scipy.configs import SCP, Objective, ScipyCFG +from blop.scipy.normalizers import InnerOptimizer, ScipyResult + + +@dataclass +class _Request: + args: tuple + future: Future + + +class InteractiveOptimizer(Optimizer): + """An optimizer object to supply an interactive interface for the scipy optimizers, with some caveats.""" + + def __init__(self, optimizer: InnerOptimizer, config: ScipyCFG | None = None, timeout: int | None = 200): + self.optimizer = optimizer + self.session(config=config if config else optimizer.config, timeout=timeout) + + def session(self, config: ScipyCFG, timeout: int | None = None): + """ + Through path for initialization and stateful reinitialization of optimization. + + derived so that mutiple initializations and lifetimes can be used for optimization. + Such as the standard ScipyOptimizer(...) call or a following "with" + """ + self._params: list[str] = [dof.parameter_name for dof in config.dofs] + self._increment: int = 0 + self._objective: Objective = config.objective + self.force_resiliance = False # kinda hidden for now + self._scale = np.ones(len(config.dofs)) + self._active: dict[int, _Request] = OrderedDict() + self.intermediate: OptimizeResult | ScipyResult | None = None + self.final: OptimizeResult | ScipyResult | None = None + self.SUGGESTION_TIMEOUT = timeout + self.thread_monitor = Future() + self.thread_start = Event() + + if config.rescale is not None: + if isinstance(config.rescale, list): + self._scale = config.rescale + else: + self._scale *= config.rescale + + def cost(x): # thread safety needs timeout so there is not infinite hang on programs + """Cooperative thread that defers evaluation of cost call by scipy to the run engine.""" + req = _Request(args=x, future=Future()) + self._active[self._increment] = req + self.thread_start.set() + self._increment += 1 + res = req.future.result(timeout=self.SUGGESTION_TIMEOUT) + if res is None: + raise ValueError("return value is not present") + return res + + kw: dict = {} + self._thread_pool = None + if config.max_iter is not None: + if config.optimizer is not SCP.TRUST_CONSTR: + kw["max_iter"] = config.max_iter + else: + kw["maxiter"] = config.max_iter + if config.eps is not None: + kw["eps"] = config.eps + + def default_callback(intermediate_result: OptimizeResult): + if self.intermediate and self.intermediate.fun < intermediate_result.fun: + return + self.intermediate = intermediate_result + self.intermediate.nit = self._increment + + def mini_worker(): + try: + if config.threads: + with ThreadPoolExecutor(max_workers=config.threads) as pool: + kw["workers"] = pool.map + res = self.optimizer.call(cost, default_callback, kws=kw) + else: + res = self.optimizer.call(cost, default_callback, kws=kw) + self.thread_monitor.set_result(res) + + except (KeyboardInterrupt, TimeoutError): + # have to have timeout, so made it that it can be restored to its state on agent auto reboot + if self.final: + ... + if self.intermediate: + self.final = self.intermediate + else: + self.final = ScipyResult(list(self.optimizer.x0), np.nan, nit=self._increment) + # self.thread_monitor.set_result(self.final) + return + except Exception as e: + self.thread_monitor.set_exception(e) + + self._t = Thread(target=mini_worker, name="optimizer") + self._t.start() + if not self.thread_start.wait(timeout=0.1): + try: + err = self.thread_monitor.exception(timeout=0.01) + if err: + raise err + except TimeoutError: + ... + return self + + def __enter__(self): + """Magic convenience to use "with" to better control thread lifetime.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Lifetime threads when using with.""" + self.close() + + def suggest(self, num_points: int | None = None) -> list[dict]: + """ + Provide a set of points in the input space, to be evaulated next. + + The "_id" key is optional and can be used to identify suggested trials for later evaluation + and ingestion. + + Parameters + ---------- + num_points : int | None, optional + The number of points to suggest. If not provided, will default to 1. + + Returns + ------- + list[dict] + A list of dictionaries, each containing a parameterization of a point to evaluate next. + Each dictionary must contain a unique "_id" key to identify each parameterization. + """ + try: + self.final = self.thread_monitor.result(timeout=0.01) + if not self.force_resiliance: + print(self.final) + raise RuntimeError("The optimizer has suspended or reached convergence") + except TimeoutError: + ... + + if self.final is not None: + vector = [x_n * s for s, x_n in zip(self._scale, self.final.x, strict=True)] + suggestion = dict(zip(self._params, vector, strict=True)) + suggestion[ID_KEY] = self.final.nit + return [suggestion] + + suggestions = [] + for id in list(self._active.keys())[: num_points if num_points is not None else 1]: + x = self._active[id].args + vector = [x_n * s for s, x_n in zip(self._scale, x, strict=True)] + + suggestion = dict(zip(self._params, vector, strict=True)) + suggestion[ID_KEY] = id + suggestions.append(suggestion) + return suggestions + + def ingest(self, points: list[dict]) -> None: + """ + Ingest a set of points into the experiment. Either from previously suggested points or from an external source. + + The "_id" key is optional. + + Parameters + ---------- + points : list[dict] + A list of dictionaries, each containing the outcomes of each suggested parameterization. + """ + for res in points: + y = res[self._objective.name] + if res[ID_KEY] not in self._active: + if not self.force_resiliance: + raise ValueError("optimizer did not expect to receive an update") + continue + self._active.pop(res[ID_KEY]).future.set_result(y) + + def get_best_points(self) -> list[tuple[Any, Mapping, Mapping]]: + """ + Get a list of the optimal point found during optimization. + + Returns + ------- + list[tuple[int, TParameterization, TOutcome]] + Each element in the list is a tuple of: + - trial index (int) + - parameter values (dict) + - metric values (dict, where values may be (value, sem) tuples) + + See Also + -------- + navigate_to_best : Plan stub to move actuators to a best point. + """ + result = self.intermediate + if self.final is not None: + result = self.final + if (result is None) or (self._objective is None): + raise ValueError("no optimization epoch has been recorded") + + vector = [x_n * s for s, x_n in zip(self._scale, result.x, strict=True)] + cart = [ + result.nit - 1, + cast(Mapping, dict(zip(self._params, vector, strict=True))), + cast(Mapping, {self._objective.name: result.fun}), + ] + return cart + + def close(self): + """Clear out futures to allow cleanup of threads.""" + for ind in list(self._active.keys()): + self._active.pop(ind).future.set_exception(KeyboardInterrupt("Execution has been suspended")) diff --git a/src/blop/scipy/normalizers.py b/src/blop/scipy/normalizers.py new file mode 100644 index 00000000..08779835 --- /dev/null +++ b/src/blop/scipy/normalizers.py @@ -0,0 +1,245 @@ +"""Normalized SciPy optimizer wrappers used by the cooperative optimization loop.""" + +from dataclasses import dataclass +from typing import Any + +import numpy as np +from scipy.optimize import OptimizeResult, dual_annealing, minimize, shgo + +from blop.scipy.configs import SCP, ScipyCFG + + +@dataclass +class ScipyResult: + """Class to unify Optimize Result and other Scipy Results.""" + + x: list[float | int] + fun: float + nit: int + status: int = 2 + + +class InnerOptimizer: + """Protocol for SciPy optimizer wrappers used by the suggest/ingest loop. + + Subclasses adapt optimizer-specific call signatures into a shared + ``call(cost, callback, kws)`` interface. This keeps optimizer internals + decoupled from the cooperative optimization loop that requests suggestions, + evaluates them externally, and ingests outcomes. + """ + + def __init__(self, config: ScipyCFG, base_args: dict | None = None) -> None: + """Store normalized configuration varaibles. + + Parameters + ---------- + config : ScipyCFG + Normalized optimizer configuration, including default bounds, + initial values, and selected method. + base_args : dict, optional + Extra keyword arguments always forwarded to wrapped optimizer. + """ + self.config = config + self.base_args = base_args + self._bounds: list[tuple[Any, Any]] = [] + scale = np.ones(len(config.dofs)) + + if config.rescale is not None: + if isinstance(config.rescale, list): + scale = config.rescale + else: + scale *= config.rescale + + for ind, dof in enumerate(config.dofs): + self._bounds.append(tuple(np.array(dof.bounds) / scale[ind])) + + self.x0 = np.mean(self._bounds, axis=1) + if config.initial is not None: + self.x0 = np.array(config.initial) / scale + + def call(self, cost, callback, kws=None) -> ScipyResult | OptimizeResult: + """Run the wrapped optimizer. + + Parameters + ---------- + cost : callable + Objective function evaluated by the optimizer. + callback : callable + Progress callback invoked by the underlying optimizer. + kws : dict, optional + Optimizer-specific options and temporary overrides. + + Returns + ------- + ScipyResult | OptimizeResult + Result object from the wrapped SciPy optimizer. + + Raises + ------ + NotImplementedError + Raised by the base protocol class when no implementation is + provided. + """ + raise NotImplementedError("Optimizer implementation not provided") + + +class Minimize(InnerOptimizer): + """Normalized wrapper around ``scipy.optimize.minimize``. + + This adapter reads default bounds and initial conditions from ``ScipyCFG`` + and forwards them to ``minimize`` using the common ``InnerOptimizer`` + interface. + """ + + def call(self, cost, callback, kws=None) -> ScipyResult: + """Execute ``scipy.optimize.minimize`` with normalized defaults. + + Parameters + ---------- + cost : callable + Objective function consumed by SciPy. + callback : callable + Callback passed directly to ``minimize``. + kws : dict, optional + Temporary call-time overrides. ``bounds`` and ``x0`` are extracted + from this dictionary when present; remaining values are passed as + ``options``. + + Returns + ------- + ScipyResult + SciPy optimization result (runtime type is ``OptimizeResult``). + """ + bounds = kws.pop("bounds", self._bounds) if kws else self._bounds + x0 = kws.pop("x0", self.x0) if kws else self.x0 + return minimize( + fun=cost, + x0=x0, + method=self.config.optimizer if self.config.optimizer != SCP.Default else None, + bounds=bounds, + callback=callback, + options=kws, + **self.base_args if self.base_args else {}, + ) + + +class DualAnnealing(InnerOptimizer): + """Normalized wrapper around ``scipy.optimize.dual_annealing``. + + ``dual_annealing`` uses a callback signature different from + ``scipy.optimize.minimize``. This adapter normalizes callback payloads so + the outer loop can handle intermediate results consistently. + """ + + def __init__(self, config: ScipyCFG, base_args: dict | None = None, inner_args: dict | None = None) -> None: + """Store normalized configuration for ``dual_annealing``. + + Parameters + ---------- + config : ScipyCFG + Normalized optimizer configuration, including bounds and initial + values. + base_args : dict, optional + Extra keyword arguments forwarded directly to + ``scipy.optimize.dual_annealing``. + inner_args : dict, optional + Additional values merged into ``minimizer_kwargs`` for the local + minimizer stage. + """ + self.inner_args = inner_args + super().__init__(config=config, base_args=base_args) + + def dual_callback(self, x, f, context): + """Convert dual-annealing callback values into a unified result type.""" + return ScipyResult(x, f, -1, context) + + def call(self, cost, callback, kws=None): + """Execute ``dual_annealing`` with normalized bounds and callbacks. + + Parameters + ---------- + cost : callable + Objective function consumed by SciPy. + callback : callable + Outer-loop callback expecting a normalized result object. + kws : dict, optional + Temporary call-time overrides. ``bounds`` and ``x0`` are extracted + when present; remaining keys are forwarded as local-minimizer + ``options``. + + Returns + ------- + OptimizeResult + Final SciPy result from ``dual_annealing``. + """ + bounds = kws.pop("bounds", self._bounds) if kws else self._bounds + x0 = kws.pop("x0", self.x0) if kws else self.x0 + opt = self.inner_args["options"] if self.inner_args else {} + return dual_annealing( + func=cost, + x0=x0, + bounds=bounds, + # Adapt SciPy's (x, f, context) callback to the outer callback + # contract that expects a normalized result object. + callback=lambda x, f, c: callback(self.dual_callback(x, f, c)), + minimizer_kwargs=self.inner_args + if self.inner_args + else {} | {"callback": callback, "bounds": bounds, "options": opt | kws if kws else {}}, + **self.base_args if self.base_args else {}, + ) + + +class SHGO(InnerOptimizer): + """Normalized wrapper around ``scipy.optimize.shgo``. + + This adapter forwards globally optimized search settings while preserving + the shared ``InnerOptimizer`` call contract. + """ + + def __init__(self, config: ScipyCFG, base_args: dict | None = None, inner_args: dict | None = None) -> None: + """Store normalized configuration for ``scipy.optimize.shgo``. + + Parameters + ---------- + config : ScipyCFG + Normalized optimizer configuration, including bounds. + base_args : dict, optional + Extra keyword arguments forwarded directly to ``shgo``. + inner_args : dict, optional + Additional values merged into ``minimizer_kwargs`` for the local + minimizer phase. + """ + self.inner_args = inner_args + super().__init__(config=config, base_args=base_args) + + def call(self, cost, callback, kws=None): + """Execute ``shgo`` with normalized bounds and minimizer options. + + Parameters + ---------- + cost : callable + Objective function consumed by SciPy. + callback : callable + Callback forwarded to the local minimizer configuration. + kws : dict, optional + Temporary call-time overrides. ``bounds`` and ``workers`` are + extracted when present; remaining keys are forwarded as + local-minimizer ``options``. + + Returns + ------- + OptimizeResult + Final SciPy result from ``shgo``. + """ + bounds = kws.pop("bounds", self._bounds) if kws else self._bounds + workers = kws.pop("workers", 1) if kws else 1 + opt = self.inner_args["options"] if self.inner_args else {} + return shgo( + func=cost, + bounds=bounds, + minimizer_kwargs=self.inner_args + if self.inner_args + else {} | {"callback": callback, "bounds": bounds, "options": opt | kws if kws else {}}, + **self.base_args if self.base_args else {}, + workers=workers, + ) diff --git a/src/blop/scipy/scipy.py b/src/blop/scipy/scipy.py new file mode 100644 index 00000000..85143609 --- /dev/null +++ b/src/blop/scipy/scipy.py @@ -0,0 +1,299 @@ +"""Scipy optimization power class for fast start QOL and Ax like agent behavior.""" + +from collections.abc import Mapping, Sequence +from typing import Any, cast + +import bluesky.preprocessors as bpp +from bluesky.callbacks import CallbackBase + +from blop.callbacks.logger import OptimizationLogger +from blop.callbacks.router import OptimizationCallbackRouter +from blop.plans import optimize +from blop.protocols import ( + AcquisitionPlan, + Actuator, + EvaluationFunction, + OptimizationProblem, + Sensor, +) +from blop.scipy.configs import SCP, Objective, RangeDOF, ScipyCFG +from blop.scipy.inverter import InteractiveOptimizer +from blop.scipy.normalizers import SHGO, DualAnnealing, Minimize +from blop.utils import InferredReadable + + +class Scipy: + """ + A convenience interface associated with running optimizations with Scipy, providing similar syntax to the Ax Agent + (allowing drop in swapping as much as possible). + + Useful as a cover in for all the QOL provided by the Agent object. + """ # noqa: D205 + + def __init__( + self, + sensors: Sequence[Sensor], + config: ScipyCFG, + evaluation_function: EvaluationFunction, + acquisition_plan: AcquisitionPlan | None = None, + **kwargs: Any, + ): + if config.optimizer not in list(SCP): + raise ValueError(f"optimizer {config.optimizer} not in supported optimizers:{list(SCP)}") + + match config.optimizer: + case SCP.DUAL_ANNEALING: + self.inner = DualAnnealing(config) + case SCP.SHGO: + self.inner = SHGO(config) + case _: + self.inner = Minimize(config) + + self.config = config + self._sensors = sensors + self._actuators = [cast(Actuator, dof.actuator) for dof in config.dofs if dof.actuator is not None] + self._evaluation_function = evaluation_function + self._acquisition_plan = acquisition_plan + self.timeout = kwargs.pop("timeout", 200) + self.optimizer = InteractiveOptimizer(self.inner, timeout=self.timeout) + self.optimizer.force_resiliance = self.resiliance = kwargs.pop("resiliance", True) + self._readable_cache: dict[str, InferredReadable] = {} + self._callbacks: list[CallbackBase] = [OptimizationLogger()] + self._callback_router = OptimizationCallbackRouter(self._callbacks) + self.sessioning = kwargs.pop("sessioning", True) + + @classmethod + def Agent( + cls, + sensors: Sequence[Sensor], + dofs: Sequence[RangeDOF], + objectives: Sequence[Objective], + evaluation_function: EvaluationFunction, + acquisition_plan: AcquisitionPlan | None = None, + optimizer: SCP = SCP.Default, + # dof_constraints: Sequence[DOFConstraint] | None = None, #implemented in future iterations? make to match ax? + # outcome_constraints: Sequence[OutcomeConstraint] | None = None, + **kwargs: Any, + ): + """ + An emcompassing interface to provide strong interoperability with Ax agent formalism. + + Parameters + ---------- + sensors : Sequence[Sensor] + The sensors to use for acquisition. These should be the minimal set + of sensors that are needed to compute the objectives. + dofs : Sequence[DOF] + The degrees of freedom that the agent can control, which determine the search space. + objectives : Sequence[Objective] + The objectives which the agent will try to optimize. + evaluation_function : EvaluationFunction + The function to evaluate acquired data and produce outcomes. + acquisition_plan : AcquisitionPlan | None, optional + The acquisition plan to use for acquiring data from the beamline. If not provided, + :func:`blop.plans.default_acquire` will be used. + **kwargs : Any + Additional keyword arguments to configure the Ax experiment. + + See Also + -------- + blop.ax.Agent + + Notes + ----- + This is a nearly drop in replacement for Ax agent sans dof + outcome constraints and checkpointing + + + """ # noqa: D401 + if len(objectives) > 1: + raise ValueError("Multiple Objectives are not supported for gradient optimizers") + config = ScipyCFG( + dofs=dofs, + objective=objectives[0], + optimizer=optimizer, + max_iter=kwargs.get("max_iter", None), + eps=kwargs.get("eps", None), + rescale=kwargs.get("scale", None), + ) + return cls(sensors, config, evaluation_function, acquisition_plan, **kwargs) + + @property + def sensors(self) -> Sequence[Sensor]: + """The sensors used for data acquisition.""" + return self._sensors + + @property + def actuators(self) -> Sequence[Actuator]: + """The actuators that control the degrees of freedom.""" + return self._actuators + + @property + def evaluation_function(self) -> EvaluationFunction: + """The function used to evaluate acquired data and produce outcomes.""" + return self._evaluation_function + + @property + def acquisition_plan(self) -> AcquisitionPlan | None: + """The acquisition plan for acquiring data, or ``None`` if using the default.""" + return self._acquisition_plan + + @property + def callbacks(self) -> list[CallbackBase]: + """The list of active optimization callbacks. + + Callbacks in this list receive documents from ``"optimize"`` and + ``"sample_suggestions"`` runs. The default list contains an + :class:`~blop.callbacks.logger.OptimizationLogger`. + + The list can be mutated directly, or use :meth:`subscribe` / + :meth:`unsubscribe` for convenience. + """ + return self._callbacks + + def subscribe(self, callback: CallbackBase) -> None: + """Subscribe a callback to receive optimization run documents. + + Parameters + ---------- + callback : CallbackBase + A Bluesky callback instance. + + Raises + ------ + ValueError + If *callback* is already subscribed. + """ + if callback in self._callbacks: + raise ValueError(f"Callback {callback!r} is already subscribed.") + self._callbacks.append(callback) + + def unsubscribe(self, callback: CallbackBase) -> None: + """Unsubscribe a previously subscribed callback. + + Parameters + ---------- + callback : CallbackBase + The callback instance to remove. + + Raises + ------ + ValueError + If *callback* is not subscribed. + """ + self._callbacks.remove(callback) + + def to_optimization_problem(self) -> OptimizationProblem: + """ + Construct an optimization problem from the Scipy Base class. + + Creates an immutable :class:`blop.protocols.OptimizationProblem` that + encapsulates all components needed for optimization. This is typically + used internally by optimization plans. + + Returns + ------- + OptimizationProblem + An immutable optimization problem that can be deployed via Bluesky. + + See Also + -------- + blop.protocols.OptimizationProblem : The optimization problem dataclass. + blop.plans.optimize : Uses the optimization problem to run optimization. + """ + return OptimizationProblem( + optimizer=self.optimizer, + actuators=self._actuators, + sensors=self._sensors, + evaluation_function=self._evaluation_function, + acquisition_plan=self._acquisition_plan, + ) + + def suggest(self, num_points: int = 1) -> list[dict]: + """ + Get the next point(s) to evaluate in the search space. + + Uses the Bayesian optimization algorithm to suggest promising points based + on all previously acquired data. Each suggestion includes an "_id" key for + tracking. + + Parameters + ---------- + num_points : int, optional + The number of points to suggest. Default is 1. Higher values enable + batch optimization but may reduce optimization efficiency per iteration. + + Returns + ------- + list[dict] + A list of dictionaries, each containing a parameterization of a point to + evaluate next. Each dictionary includes an "_id" key for identification. + """ + return self.optimizer.suggest(num_points) + + def ingest(self, points: list[dict]) -> None: + """ + Ingest evaluation results into the optimizer. + + Updates the optimizer's model with new data. Can ingest both suggested points + (with "_id" key) and external data (without "_id" key). + + Parameters + ---------- + points : list[dict] + A list of dictionaries, each containing outcomes for a trial. For suggested + points, include the "_id" key. For external data, include DOF names and + objective values, and omit "_id". + + Notes + ----- + This method is typically called automatically by :meth:`optimize`. Manual usage + is only needed for custom workflows or when ingesting external data. + + For complete examples, see :doc:`/how-to-guides/attach-data-to-experiments`. + """ + self.optimizer.ingest(points) + + def optimize(self, iterations=10, n_points=1): + """Optimization plan wrapper used by the agent interface.""" + if self.optimizer.final is not None: + self.config.initial = self.optimizer.final.x + self.optimizer = InteractiveOptimizer(self.inner, timeout=self.timeout) + self.optimizer.force_resiliance = self.resiliance + optimize_plan = optimize( + self.to_optimization_problem(), + iterations=iterations, + n_points=n_points, + readable_cache=self._readable_cache, + ) + + if self._callbacks: + optimize_plan = bpp.subs_wrapper( + optimize_plan, + self._callback_router, + ) + if self.sessioning: + with self.optimizer: + yield from optimize_plan + else: + yield from optimize_plan + + def get_best_points(self) -> list[tuple[Any, Mapping, Mapping]]: + """ + Get a list of the optimal points found during optimization. + + For single-objective optimization, returns a single best point. + For multi-objective optimization, returns the Pareto-optimal set. + + Returns + ------- + list[tuple[int, TParameterization, TOutcome]] + Each element in the list is a tuple of: + - trial index (int) + - parameter values (dict) + - metric values (dict, where values may be (value, sem) tuples) + + See Also + -------- + navigate_to_best : Plan stub to move actuators to a best point. + """ + return self.optimizer.get_best_points() diff --git a/src/blop/tests/scipy/__init__.py b/src/blop/tests/scipy/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/blop/tests/scipy/test_optimizer.py b/src/blop/tests/scipy/test_optimizer.py new file mode 100644 index 00000000..95a928f6 --- /dev/null +++ b/src/blop/tests/scipy/test_optimizer.py @@ -0,0 +1,429 @@ +from unittest.mock import MagicMock + +import pytest +from scipy.optimize import OptimizeResult + +from blop.ax import Objective, RangeDOF +from blop.protocols import ID_KEY, AcquisitionPlan, EvaluationFunction +from blop.scipy.configs import SCP, ScipyCFG +from blop.scipy.inverter import InteractiveOptimizer +from blop.scipy.normalizers import SHGO, DualAnnealing, InnerOptimizer, Minimize, ScipyResult + +from ..conftest import MovableSignal + + +@pytest.fixture(scope="function") +def mock_evaluation_function(): + return MagicMock(spec=EvaluationFunction) + + +@pytest.fixture(scope="function") +def mock_acquisition_plan(): + return MagicMock(spec=AcquisitionPlan) + + +@pytest.fixture(scope="function") +def optimizer_prep(): + movable1 = MovableSignal(name="test_movable1") + movable2 = MovableSignal(name="test_movable2") + dof1 = RangeDOF(actuator=movable1, bounds=(0, 10), parameter_type="float") + dof2 = RangeDOF(actuator=movable2, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + config = ScipyCFG(dofs=[dof1, dof2], objective=objective, threads=4, rescale=[2.0, 3.0], eps=0.5) + inner = Minimize(config) + return InteractiveOptimizer(inner, timeout=1) + + +# ============================================================================ +# PHASE 2: Optimizer Algorithm Variations Tests +# ============================================================================ + + +@pytest.mark.parametrize("optimizer", list(SCP)[:10]) +def test_scipy_optimizer_algorithms(mock_evaluation_function, mock_acquisition_plan, optimizer): + """Test ScipyOptimizer with different SCP algorithms.""" + movable1 = MovableSignal(name="test_movable1") + movable2 = MovableSignal(name="test_movable2") + dof1 = RangeDOF(actuator=movable1, bounds=(0, 10), parameter_type="float") + dof2 = RangeDOF(actuator=movable2, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + + config = ScipyCFG( + dofs=[dof1, dof2], + objective=objective, + optimizer=optimizer, + max_iter=10, + ) + + inner = Minimize(config) + opt = InteractiveOptimizer(inner, timeout=1) + assert opt._active is not None + opt.close() + + +def test_scipy_optimizer_bfgs_specific(mock_evaluation_function, mock_acquisition_plan): + """Test ScipyOptimizer explicitly with BFGS.""" + movable1 = MovableSignal(name="test_movable1") + movable2 = MovableSignal(name="test_movable2") + dof1 = RangeDOF(actuator=movable1, bounds=(0, 10), parameter_type="float") + dof2 = RangeDOF(actuator=movable2, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + + config = ScipyCFG( + dofs=[dof1, dof2], + objective=objective, + optimizer=SCP.BFGS, + max_iter=10, + ) + + inner = Minimize(config) + opt = InteractiveOptimizer(inner, timeout=1) + assert opt.final is None # No optimization run yet + opt.close() + + +def test_scipy_optimizer_dual_annealing_specific(mock_evaluation_function, mock_acquisition_plan): + """Test ScipyOptimizer explicitly with Dual_Annealing.""" + movable = MovableSignal(name="test_movable") + dof = RangeDOF(actuator=movable, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + + config = ScipyCFG( + dofs=[dof], + objective=objective, + optimizer=SCP.DUAL_ANNEALING, + ) + + inner = DualAnnealing(config) + opt = InteractiveOptimizer(inner, timeout=1) + assert opt.final is None # No optimization run yet + opt.close() + + +def test_scipy_optimizer_SHGO_specific(mock_evaluation_function, mock_acquisition_plan): + """Test ScipyOptimizer explicitly with Dual_Annealing.""" + movable = MovableSignal(name="test_movable") + dof = RangeDOF(actuator=movable, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + + config = ScipyCFG( + dofs=[dof], + objective=objective, + optimizer=SCP.SHGO, + ) + + inner = SHGO(config) + opt = InteractiveOptimizer(inner, timeout=1) + assert opt.final is None # No optimization run yet + opt.close() + + +def test_scipy_optimizer_threads_none(mock_evaluation_function, mock_acquisition_plan): + """Test ScipyOptimizer with threads=None (no parallelization).""" + movable = MovableSignal(name="test_movable") + dof = RangeDOF(actuator=movable, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + + config = ScipyCFG( + dofs=[dof], + objective=objective, + threads=None, + ) + + inner = Minimize(config) + opt = InteractiveOptimizer(inner, timeout=1) + assert opt._thread_pool is None # No thread pool when threads=None + opt.close() + + +def test_scipy_optimizer_threads_multiple(mock_evaluation_function, mock_acquisition_plan): + """Test ScipyOptimizer with multiple threads.""" + movable = MovableSignal(name="test_movable") + dof = RangeDOF(actuator=movable, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + + config = ScipyCFG( + dofs=[dof], + objective=objective, + threads=2, + ) + + inner = Minimize(config) + opt = InteractiveOptimizer(inner, timeout=1) + # Configuration accepted + opt.close() + + +# ============================================================================ +# PHASE 3: Rescaling & Parameter Handling Tests +# ============================================================================ + + +def test_rescaling_suggest_output(optimizer_prep): + """Test suggest() respects rescaling (scaled input → unscaled output).""" + suggestions = optimizer_prep.suggest(1) + assert len(suggestions) == 1 + + # Suggested values should be in original (unscaled) space + # Bounds are (0, 10), so suggestions should be in [0, 10] + assert 0 <= suggestions[0]["test_movable1"] <= 10 + assert 0 <= suggestions[0]["test_movable2"] <= 10 + optimizer_prep.close() + + +def test_rescaling_ingest_parameters(optimizer_prep): + """Test ingest() processes scaled parameters correctly.""" + # Suggest first + optimizer_prep.suggest(1) + # Ingest with outcome + optimizer_prep.ingest([{"test_movable1": 2.5, "test_movable2": 5.0, "test_objective": 0.8, ID_KEY: 0}]) + optimizer_prep.close() + + +def test_get_best_points_scaling(optimizer_prep): + """Test get_best_points() with scaling works (verify basic structure).""" + # Set final result manually (simulate completed optimization) + optimizer_prep.final = ScipyResult( + x=[2.5, 3.0], # Scaled values + fun=0.85, + nit=15, + status=0, + ) + + best = optimizer_prep.get_best_points() + assert isinstance(best, list) + assert len(best) == 3 # (trial_idx, params_dict, metrics_dict) + assert best[0] == 14 # nit - 1 + assert "test_movable1" in best[1] + assert "test_movable2" in best[1] + optimizer_prep.close() + + +# ============================================================================ +# PHASE 4: Error Handling & Validation Tests +# ============================================================================ + + +def test_scipy_optimizer_internal_startup_error(mock_evaluation_function, mock_acquisition_plan): + """Test ScipyOptimizer passes error from scipy internal (optimizer not defined)""" + movable1 = MovableSignal(name="test_movable1") + movable2 = MovableSignal(name="test_movable2") + dof1 = RangeDOF(actuator=movable1, bounds=(0, 10), parameter_type="float") + dof2 = RangeDOF(actuator=movable2, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + + config = ScipyCFG( + dofs=[dof1, dof2], + objective=objective, + max_iter=10, + ) + inner = InnerOptimizer(config) + with pytest.raises(NotImplementedError): + InteractiveOptimizer(inner, timeout=1) + + +def test_scipy_optimizer_internal_runtime_error(mock_evaluation_function, mock_acquisition_plan): + """Test ScipyOptimizer passes error from scipy internal (optimizer raises operational error)""" + movable1 = MovableSignal(name="test_movable1") + movable2 = MovableSignal(name="test_movable2") + dof1 = RangeDOF(actuator=movable1, bounds=(0, 10), parameter_type="float") + dof2 = RangeDOF(actuator=movable2, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + + config = ScipyCFG( + dofs=[dof1, dof2], + objective=objective, + max_iter=10, + ) + + class ErrFun(InnerOptimizer): + def call(self, cost, callback, kws=None) -> ScipyResult | OptimizeResult: + cost([1, 2]) + raise RuntimeError("This should be caught by main thread") + + inner = ErrFun(config) + opt = InteractiveOptimizer(inner, timeout=1) + with pytest.raises(RuntimeError): + sg = opt.suggest() + opt.ingest([sg[0] | {objective.name: -1}]) + opt.suggest() + + +def test_ingest_raises_on_unknown_id(optimizer_prep): + """Test ingest() raises ValueError when ID not in _active requests.""" + # Try to ingest with unknown ID + with pytest.raises(ValueError, match="optimizer did not expect to receive an update"): + optimizer_prep.ingest([{"test_movable1": 5.0, "test_movable2": 5.0, "test_objective": 0.5, ID_KEY: 999}]) + + optimizer_prep.close() + + +def test_ingest_force_resiliance_skips_unknown_id(optimizer_prep): + """Test ingest() with force_resiliance=True skips unknown IDs.""" + # Enable resiliance + optimizer_prep.force_resiliance = True + + # This should NOT raise (unknown IDs are skipped) + optimizer_prep.ingest([{"test_movable1": 5.0, "test_movable2": 5.0, "test_objective": 0.5, ID_KEY: 999}]) + + optimizer_prep.close() + + +def test_ingest_missing_objective_name(optimizer_prep): + """Test ingest() raises ValueError if objective name missing from data.""" + # Suggest first to create an active request + optimizer_prep.suggest(1) + + # Try to ingest without objective value + with pytest.raises(KeyError): + optimizer_prep.ingest([{"test_movable1": 5.0, "test_movable2": 5.0, ID_KEY: 0}]) # Missing "test_objective" + + optimizer_prep.close() + + +def test_optimizer_close_cancels_futures(optimizer_prep): + """Test ScipyOptimizer.close() cancels all active futures.""" + # Suggest to create active futures + optimizer_prep.suggest(2) + active_count = len(optimizer_prep._active) + assert active_count > 0 + + # Close should cancel all futures + optimizer_prep.close() + # All futures should now have exceptions set + for future_wrapper in optimizer_prep._active.values(): + assert future_wrapper.future.done() + + +def test_suggest_before_optimization(optimizer_prep): + """Test suggest() before any optimization returns expected state.""" + # Suggest before any ingest + suggestions = optimizer_prep.suggest(1) + assert len(suggestions) == 1 + assert "_id" in suggestions[0] + optimizer_prep.close() + + +# ============================================================================ +# PHASE 6: State Management & Context Manager Tests +# ============================================================================ + + +def test_scipy_optimizer_context_manager(mock_evaluation_function, mock_acquisition_plan): + """Test ScipyOptimizer context manager protocol (__enter__/__exit__).""" + movable = MovableSignal(name="test_movable") + dof = RangeDOF(actuator=movable, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + + config = ScipyCFG(dofs=[dof], objective=objective) + + inner = Minimize(config) + with InteractiveOptimizer(inner, timeout=1) as opt: + assert opt is not None + suggestions = opt.suggest(1) + assert len(suggestions) == 1 + + +def test_scipy_optimizer_session_reinit(mock_evaluation_function, mock_acquisition_plan): + """Test ScipyOptimizer.session() reinitializes state.""" + movable = MovableSignal(name="test_movable") + dof = RangeDOF(actuator=movable, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + + config = ScipyCFG(dofs=[dof], objective=objective) + + inner = Minimize(config) + opt = InteractiveOptimizer(inner, timeout=1) + opt.suggest(1) + + # Call session to reinitialize + opt.session(config, timeout=1) + # State should be reset + assert opt._increment == 1 + assert len(opt._active) == 1 + opt.close() + + +def test_get_best_points_intermediate_only(optimizer_prep): + """Test get_best_points() with only intermediate results (final=None).""" + # Set intermediate result manually (simulate partway through optimization) + optimizer_prep.intermediate = ScipyResult( + x=[5.0, -5.0], + fun=0.7, + nit=5, + status=0, + ) + + best = optimizer_prep.get_best_points() + assert len(best) == 3 + assert best[0] == 4 # nit - 1 + optimizer_prep.close() + + +def test_get_best_points_final_preferred(optimizer_prep): + """Test get_best_points() prefers final over intermediate.""" + # Set both intermediate and final + optimizer_prep.intermediate = ScipyResult( + x=[5.0, -5.0], + fun=0.7, + nit=5, + status=0, + ) + optimizer_prep.final = ScipyResult( + x=[7.0, -5.0], + fun=0.9, + nit=10, + status=0, + ) + + best = optimizer_prep.get_best_points() + # best[2] is a dict with objective as key; get the first (only) value + objective_value = list(best[2].values())[0] + assert objective_value == 0.9 # Uses final result + optimizer_prep.close() + + +def test_get_best_points_no_optimization_raises(optimizer_prep): + """Test get_best_points() raises ValueError if no optimization run.""" + # No optimization run: both intermediate and final are None + with pytest.raises(ValueError, match="no optimization epoch has been recorded"): + optimizer_prep.get_best_points() + + optimizer_prep.close() + + +def test_callback_updates(): + """Test ScipyOptimizer recovers optimization reporting from optimizer""" + movable1 = MovableSignal(name="test_movable1") + movable2 = MovableSignal(name="test_movable2") + dof1 = RangeDOF(actuator=movable1, bounds=(0, 10), parameter_type="float") + dof2 = RangeDOF(actuator=movable2, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + + config = ScipyCFG( + dofs=[dof1, dof2], + objective=objective, + max_iter=10, + ) + + first = ScipyResult([1, 1], 1, nit=99) + best = ScipyResult([1, 1], 0, nit=99) + worst = ScipyResult([1, 1], 2, nit=99) + + class CallbackFun(InnerOptimizer): + def call(self, cost, callback, kws=None) -> ScipyResult | OptimizeResult: + callback(first) + cost([1, 1]) + callback(best) + cost([1, 1]) + callback(worst) + cost([1, 1]) + + inner = CallbackFun(config) + opt = InteractiveOptimizer(inner, timeout=1) + assert opt.intermediate is first + opt.ingest([opt.suggest()[0] | {objective.name: -1}]) + opt.ingest([opt.suggest()[0] | {objective.name: -1}]) + opt.ingest([opt.suggest()[0] | {objective.name: -1}]) + assert opt.intermediate is best # make sure it only keeps the best + opt.close() diff --git a/src/blop/tests/scipy/test_scipy.py b/src/blop/tests/scipy/test_scipy.py new file mode 100644 index 00000000..46b136d3 --- /dev/null +++ b/src/blop/tests/scipy/test_scipy.py @@ -0,0 +1,430 @@ +import types +from unittest.mock import MagicMock + +import pytest + +from blop.ax import Objective, RangeDOF +from blop.protocols import ID_KEY, AcquisitionPlan, EvaluationFunction +from blop.scipy.configs import ScipyCFG +from blop.scipy.inverter import InteractiveOptimizer +from blop.scipy.normalizers import ScipyResult +from blop.scipy.scipy import Scipy + +from ..conftest import MovableSignal, ReadableSignal + + +@pytest.fixture(scope="function") +def mock_evaluation_function(): + return MagicMock(spec=EvaluationFunction) + + +@pytest.fixture(scope="function") +def mock_acquisition_plan(): + return MagicMock(spec=AcquisitionPlan) + + +# agent.optimizer.close() is called so the standard timeout doesnt make the testing take forever + + +@pytest.fixture(scope="function") +def agent_prep(mock_evaluation_function, mock_acquisition_plan): + movable1 = MovableSignal(name="test_movable1") + movable2 = MovableSignal(name="test_movable2") + readable = ReadableSignal(name="test_readable") + dof1 = RangeDOF(actuator=movable1, bounds=(0, 10), parameter_type="float") + dof2 = RangeDOF(actuator=movable2, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + config = ScipyCFG(dofs=[dof1, dof2], objective=objective, threads=4) + agent = Scipy( + sensors=[readable], + config=config, + evaluation_function=mock_evaluation_function, + acquisition_plan=mock_acquisition_plan, + name="test_experiment", + timeout=5, + ) + return agent + + +@pytest.fixture(scope="function") +def secoundary_agent_prep(mock_evaluation_function, mock_acquisition_plan): + movable = MovableSignal(name="test_movable") + readable = ReadableSignal(name="test_readable") + dof = RangeDOF(actuator=movable, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + config = ScipyCFG(dofs=[dof], objective=objective) + agent = Scipy( + sensors=[readable], + config=config, + evaluation_function=mock_evaluation_function, + acquisition_plan=mock_acquisition_plan, + timeout=5, + ) + return agent + + +def test_general_init(mock_evaluation_function, mock_acquisition_plan): + """Test that the simple Scipy can be initialized.""" + movable1 = MovableSignal(name="test_movable1") + movable2 = MovableSignal(name="test_movable2") + readable = ReadableSignal(name="test_readable") + dof1 = RangeDOF(actuator=movable1, bounds=(0, 10), parameter_type="float") + dof2 = RangeDOF(actuator=movable2, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + config = ScipyCFG( + dofs=[dof1, dof2], + objective=objective, + ) + agent = Scipy( + sensors=[readable], + config=config, + evaluation_function=mock_evaluation_function, + acquisition_plan=mock_acquisition_plan, + name="test_experiment", + timeout=5, + ) + assert agent.sensors == [readable] + assert agent.actuators == [dof1.actuator, dof2.actuator] + assert agent.evaluation_function == mock_evaluation_function + assert agent.acquisition_plan == mock_acquisition_plan + agent.optimizer.close() + + +def test_agent_init(mock_evaluation_function, mock_acquisition_plan): + """Test that the agent can be initialized.""" + movable1 = MovableSignal(name="test_movable1") + movable2 = MovableSignal(name="test_movable2") + readable = ReadableSignal(name="test_readable") + dof1 = RangeDOF(actuator=movable1, bounds=(0, 10), parameter_type="float") + dof2 = RangeDOF(actuator=movable2, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + agent = Scipy.Agent( + sensors=[readable], + dofs=[dof1, dof2], + objectives=[objective], + evaluation_function=mock_evaluation_function, + acquisition_plan=mock_acquisition_plan, + name="test_experiment", + timeout=5, + ) + assert agent.sensors == [readable] + assert agent.actuators == [dof1.actuator, dof2.actuator] + assert agent.evaluation_function == mock_evaluation_function + assert agent.acquisition_plan == mock_acquisition_plan + agent.optimizer.close() + + +def test_agent_to_optimization_problem(mock_evaluation_function): + """Test that the agent can be converted to an optimization problem.""" + movable1 = MovableSignal(name="test_movable1") + movable2 = MovableSignal(name="test_movable2") + dof1 = RangeDOF(actuator=movable1, bounds=(0, 10), parameter_type="float") + dof2 = RangeDOF(actuator=movable2, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + agent = Scipy.Agent( + sensors=[], dofs=[dof1, dof2], objectives=[objective], evaluation_function=mock_evaluation_function, timeout=5 + ) + optimization_problem = agent.to_optimization_problem() + assert optimization_problem.evaluation_function == mock_evaluation_function + assert optimization_problem.actuators == [movable1, movable2] + assert optimization_problem.sensors == [] + assert isinstance(optimization_problem.optimizer, InteractiveOptimizer) + assert optimization_problem.acquisition_plan is None + agent.optimizer.close() + + +def test_agent_suggest(agent_prep): + parameterizations = agent_prep.suggest(1) + assert len(parameterizations) == 1 + assert parameterizations[0]["_id"] == 0 + assert "test_movable1" in parameterizations[0] + assert "test_movable2" in parameterizations[0] + assert isinstance(parameterizations[0]["test_movable1"], (int, float)) + assert isinstance(parameterizations[0]["test_movable2"], (int, float)) + agent_prep.optimizer.close() + + +def test_agent_ingest(agent_prep): + agent_prep.suggest() + agent_prep.ingest([{"test_movable1": 0.1, "test_movable2": 0.2, "test_objective": 0.3, ID_KEY: 0}]) + agent_prep.optimizer.close() + + +def test_agent_multithread(agent_prep): + agent_prep.suggest(1) + agent_prep.ingest([{"test_movable1": 0.1, "test_movable2": 0.2, "test_objective": 0.3, ID_KEY: 0}]) + params = agent_prep.suggest(4) + print(agent_prep.optimizer._active) + assert len(params) > 1 + agent_prep.optimizer.close() + + +# ============================================================================ +# PHASE 1: Configuration & Initialization Tests +# ============================================================================ + + +def test_scipy_cfg_rescaling_scalar(mock_evaluation_function, mock_acquisition_plan): + """Test ScipyCFG with scalar rescaling.""" + movable = MovableSignal(name="test_movable") + dof = RangeDOF(actuator=movable, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + + config = ScipyCFG( + dofs=[dof], + objective=objective, + rescale=2.0, + ) + + agent = Scipy( + sensors=[], + config=config, + evaluation_function=mock_evaluation_function, + acquisition_plan=mock_acquisition_plan, + timeout=5, + ) + + # Verify rescaling was applied + assert agent.optimizer._scale[0] == 2.0 + agent.optimizer.close() + + +def test_scipy_cfg_rescaling_list(mock_evaluation_function, mock_acquisition_plan): + """Test ScipyCFG with list rescaling per parameter.""" + movable1 = MovableSignal(name="test_movable1") + movable2 = MovableSignal(name="test_movable2") + dof1 = RangeDOF(actuator=movable1, bounds=(0, 10), parameter_type="float") + dof2 = RangeDOF(actuator=movable2, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + + config = ScipyCFG( + dofs=[dof1, dof2], + objective=objective, + rescale=[2.0, 3.0], + ) + + agent = Scipy( + sensors=[], + config=config, + evaluation_function=mock_evaluation_function, + acquisition_plan=mock_acquisition_plan, + timeout=5, + ) + + # Verify rescaling per DOF + assert agent.optimizer._scale[0] == 2.0 + assert agent.optimizer._scale[1] == 3.0 + agent.optimizer.close() + + +def test_scipy_cfg_initial_parameters(mock_evaluation_function, mock_acquisition_plan): + """Test ScipyCFG with initial parameter values.""" + movable1 = MovableSignal(name="test_movable1") + movable2 = MovableSignal(name="test_movable2") + dof1 = RangeDOF(actuator=movable1, bounds=(0, 10), parameter_type="float") + dof2 = RangeDOF(actuator=movable2, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + + initial_params = [2.5, 7.5] + config = ScipyCFG( + dofs=[dof1, dof2], + objective=objective, + initial=initial_params, + ) + + agent = Scipy( + sensors=[], + config=config, + evaluation_function=mock_evaluation_function, + acquisition_plan=mock_acquisition_plan, + timeout=5, + ) + + # Verify initial parameters are set + agent.optimizer.close() + + +def test_scipy_cfg_max_iter_and_eps(mock_evaluation_function, mock_acquisition_plan): + """Test ScipyCFG with max_iter and eps parameters.""" + movable = MovableSignal(name="test_movable") + dof = RangeDOF(actuator=movable, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + + config = ScipyCFG( + dofs=[dof], + objective=objective, + max_iter=50, + eps=1e-6, + ) + + assert config.max_iter == 50 + assert config.eps == 1e-6 + + +def test_agent_invalidoptimizer_enum(mock_evaluation_function, mock_acquisition_plan): + """Test Scipy.Agent raises ValueError for invalid optimizer.""" + movable = MovableSignal(name="test_movable") + dof = RangeDOF(actuator=movable, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + readable = ReadableSignal(name="test_readable") + + with pytest.raises((ValueError, NotImplementedError), match="optimizer.*not in supported optimizers"): + Scipy.Agent( + sensors=[readable], + dofs=[dof], + objectives=[objective], + evaluation_function=mock_evaluation_function, + optimizer="invalidoptimizer", + ) + + +def test_agent_multiple_objectives_not_supported(mock_evaluation_function, mock_acquisition_plan): + """Test Scipy.Agent raises ValueError for multiple objectives.""" + movable = MovableSignal(name="test_movable") + dof = RangeDOF(actuator=movable, bounds=(0, 10), parameter_type="float") + objective1 = Objective(name="test_objective_1", minimize=False) + objective2 = Objective(name="test_objective_2", minimize=False) + readable = ReadableSignal(name="test_readable") + + with pytest.raises(ValueError, match="Multiple Objectives are not supported"): + Scipy.Agent( + sensors=[readable], + dofs=[dof], + objectives=[objective1, objective2], + evaluation_function=mock_evaluation_function, + ) + + +# ============================================================================ +# PHASE 5: Callback Management Tests +# ============================================================================ + + +def test_subscribe_callback(secoundary_agent_prep): + """Test subscribe() adds callback to list.""" + callback = MagicMock() + initial_count = len(secoundary_agent_prep.callbacks) + secoundary_agent_prep.subscribe(callback) + + assert len(secoundary_agent_prep.callbacks) == initial_count + 1 + assert callback in secoundary_agent_prep.callbacks + secoundary_agent_prep.optimizer.close() + + +def test_subscribe_duplicate_raises(secoundary_agent_prep): + """Test subscribe() raises ValueError on duplicate callback.""" + callback = MagicMock() + secoundary_agent_prep.subscribe(callback) + + with pytest.raises(ValueError, match="already subscribed"): + secoundary_agent_prep.subscribe(callback) + + secoundary_agent_prep.optimizer.close() + + +def test_unsubscribe_callback(secoundary_agent_prep): + """Test unsubscribe() removes callback from list.""" + callback = MagicMock() + secoundary_agent_prep.subscribe(callback) + assert callback in secoundary_agent_prep.callbacks + + secoundary_agent_prep.unsubscribe(callback) + assert callback not in secoundary_agent_prep.callbacks + secoundary_agent_prep.optimizer.close() + + +def test_unsubscribe_not_subscribed_raises(secoundary_agent_prep): + """Test unsubscribe() raises ValueError if not subscribed.""" + callback = MagicMock() + + with pytest.raises(ValueError): + secoundary_agent_prep.unsubscribe(callback) + + secoundary_agent_prep.optimizer.close() + + +# ============================================================================ +# PHASE 7: Edge Cases & Boundary Conditions Tests +# ============================================================================ + + +def test_scipy_secoundary(secoundary_agent_prep): + """Test Scipy with single DOF (one parameter).""" + suggestions = secoundary_agent_prep.suggest(1) + assert len(suggestions) == 1 + assert "test_movable" in suggestions[0] + secoundary_agent_prep.optimizer.close() + + +def test_scipy_large_rescale_factors(mock_evaluation_function, mock_acquisition_plan): + """Test Scipy with large rescale factors (extreme scaling).""" + movable1 = MovableSignal(name="test_movable1") + movable2 = MovableSignal(name="test_movable2") + dof1 = RangeDOF(actuator=movable1, bounds=(0, 10), parameter_type="float") + dof2 = RangeDOF(actuator=movable2, bounds=(0, 10), parameter_type="float") + objective = Objective(name="test_objective", minimize=False) + readable = ReadableSignal(name="test_readable") + + config = ScipyCFG( + dofs=[dof1, dof2], + objective=objective, + rescale=[0.001, 1000.0], # Extreme scaling + ) + + agent = Scipy( + sensors=[readable], + config=config, + evaluation_function=mock_evaluation_function, + acquisition_plan=mock_acquisition_plan, + timeout=5, + ) + + suggestions = agent.suggest(1) + assert len(suggestions) == 1 + # Values should still be in original bounds + assert 0 <= suggestions[0]["test_movable1"] <= 10 + assert 0 <= suggestions[0]["test_movable2"] <= 10 + agent.optimizer.close() + + +def test_suggest_after_final_optimization(secoundary_agent_prep): + """Test suggest() after final optimization returns final result parameterization.""" + # Set final optimization result + secoundary_agent_prep.optimizer.final = ScipyResult( + x=[7.0], + fun=0.95, + nit=20, + status=0, + ) + + suggestions = secoundary_agent_prep.optimizer.suggest() + assert len(suggestions) == 1 + assert suggestions[0]["test_movable"] == 7.0 + assert suggestions[0][ID_KEY] == 20 + secoundary_agent_prep.optimizer.close() + + +def test_optimize(secoundary_agent_prep): + """Test agent setup of optimize plane""" + # Set final optimization result + + iter = secoundary_agent_prep.optimize(0) + assert isinstance(iter, types.GeneratorType) + next(iter) + secoundary_agent_prep.optimizer.close() + + +def test_optimize_with_prev_final(secoundary_agent_prep): + """Test agent setup of optimize plane""" + # Set final optimization result + secoundary_agent_prep.optimizer.final = ScipyResult( + x=[7.0], + fun=0.95, + nit=20, + status=0, + ) + + iter = secoundary_agent_prep.optimize(0) + assert isinstance(iter, types.GeneratorType) + next(iter) + secoundary_agent_prep.optimizer.close()