Ax Integration

We use Ax as the primary backend for optimization and experiment tracking. This enables us to offload the optimization and experiment tracking processes and focus on delivering a great user experience for optimization using Bluesky.

Note

The use of Ax is optional. You can implement your own optimizer using the blop.protocols.Optimizer protocol. However, the blop.ax.Agent class is designed to be a nicer interface for using Ax only.

Wrapping Ax

Generally, we want to avoid wrapping the Ax API and allow users to use the full power of Ax themselves. However, there are some cases where an API wrapper is useful:

• DOFs: We wrap Ax parameters to link them with Bluesky Actuators, enabling automatic movement during acquisition (see blop.ax.dof.DOF).

• Constraints: We provide blop.ax.dof.DOFConstraint and blop.ax.objective.OutcomeConstraint with a readable syntax that maps to Ax constraints.

• Agent: The blop.ax.Agent class provides a higher-level interface for common optimization workflows.

Blop handles the following aspects of the Ax API for you:

• Experiment creation and configuration

• Optimization configuration (objectives, constraints)

• Trial suggestion and ingestion

We chose this minimal approach because it encompasses the basic usage of Ax for optimization. For more complex setups (e.g. multi-fidelity optimization, custom model fitting), use the Ax API directly.

Using the Ax API directly

You can access the underlying ax.Client instance via blop.ax.Agent.ax_client.

You can learn all about the various Ax features in the Ax documentation. Some notable features that are not used by Blop by default are:

• Generation strategy configuration

  • Influencing Ax’s default generation strategy: https://ax.dev/docs/recipes/influence-gs-choice

  • Using custom BoTorch models: https://ax.dev/docs/tutorials/modular_botorch/ (see Using custom generation strategies for an example in Blop)

  • Using external generators: https://ax.dev/docs/tutorials/external_generation_node/

• Configuring early stopping: https://ax.dev/docs/tutorials/early_stopping/ (first-class support for early stopping is coming soon, see https://github.com/NSLS-II/blop/issues/129)

• Analyzing the optimization results and model fit: https://ax.dev/docs/tutorials/analyses/ (see blop.ax.Agent.plot_objective() for an example in Blop)

• Configuring tracking metrics: https://ax.dev/docs/tutorials/tracking_metrics/

• Saving and loading experiments

  • To JSON: https://ax.dev/docs/recipes/experiment-to-json

  • To SQLite: https://ax.dev/docs/recipes/experiment-to-sqlite

• Summarizing the experiment: https://ax.dev/docs/tutorials/quickstart/