{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simulated KB Mirror Demo using Ax and Blop\n",
    "\n",
    "This notebook introduces the use of the [Ax Adaptive Experimentation Platform](https://ax.dev) with integrations for Blop.\n",
    "\n",
    "Blop integrates the following into Ax:\n",
    "- Running Bluesky plans using the run engine\n",
    "- Using devices as parameters\n",
    "- Using detectors to produce data\n",
    "- Retrieving the results from databroker\n",
    "\n",
    "These features make it simple to optimize your beamline using both the [Bluesky ecosystem](https://blueskyproject.io) and Ax."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing a test environment\n",
    "\n",
    "Here we prepare the `RunEngine` and data service `Databroker`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from blop.utils import prepare_re_env  # noqa\n",
    "\n",
    "%run -i $prepare_re_env.__file__ --db-type=temp\n",
    "bec.disable_plots()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulated beamline with KB mirror pair\n",
    "\n",
    "Here we describe an analytical simulated beamline with a [KB mirror](https://en.wikipedia.org/wiki/Kirkpatrick%E2%80%93Baez_mirror) pair."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from blop.sim import Beamline\n",
    "\n",
    "beamline = Beamline(name=\"bl\")\n",
    "beamline.det.noise.put(False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a Blop-Ax experiment\n",
    "\n",
    "Now we can define the experiment we plan to run.\n",
    "\n",
    "This involves setting 4 parameters that simulate motor positions controlling two KB mirrors. The objectives of the experiment are to maximize the beam intensity while minimizing the area of the beam."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ax.service.ax_client import AxClient, ObjectiveProperties\n",
    "from blop.integrations.ax import create_blop_experiment\n",
    "\n",
    "ax_client = AxClient()\n",
    "create_blop_experiment(\n",
    "    ax_client,\n",
    "    parameters=[\n",
    "        {\n",
    "            \"movable\": beamline.kbv_dsv,\n",
    "            \"type\": \"range\",\n",
    "            \"bounds\": [-5.0, 5.0],\n",
    "        },\n",
    "        {\n",
    "            \"movable\": beamline.kbv_usv,\n",
    "            \"type\": \"range\",\n",
    "            \"bounds\": [-5.0, 5.0],\n",
    "        },\n",
    "        {\n",
    "            \"movable\": beamline.kbh_dsh,\n",
    "            \"type\": \"range\",\n",
    "            \"bounds\": [-5.0, 5.0],\n",
    "        },\n",
    "        {\n",
    "            \"movable\": beamline.kbh_ush,\n",
    "            \"type\": \"range\",\n",
    "            \"bounds\": [-5.0, 5.0],\n",
    "        },\n",
    "    ],\n",
    "    objectives={\n",
    "        \"beam_intensity\": ObjectiveProperties(minimize=False, threshold=200.0),\n",
    "        \"beam_area\": ObjectiveProperties(minimize=True, threshold=1000.0),\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an evaluation function\n",
    "\n",
    "Now that we have setup the experiment, we need to define how to compute the objective values.\n",
    "\n",
    "In this example, the `RunEngine` produces readings from the detector that are retrieved from `Databroker` and transformed into a Pandas `DataFrame`. Using the image produced from this, we can compute some statistics from the image to produce the beam intensity and beam area (our objectives).\n",
    "\n",
    "Ax expects a `tuple[float, float]` for each objective representing the mean value and standard error, respectively. For a single image, the average intensity is just the intensity (same for the area), and we assume no uncertainty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from blop.utils import get_beam_stats\n",
    "\n",
    "\n",
    "def evaluate(results_df: pd.DataFrame) -> dict[str, tuple[float, float]]:\n",
    "    stats = get_beam_stats(results_df[\"bl_det_image\"].iloc[0])\n",
    "    area = stats[\"wid_x\"] * stats[\"wid_y\"]\n",
    "    return {\n",
    "        \"beam_intensity\": (stats[\"sum\"], None),\n",
    "        \"beam_area\": (area, None),\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create an evaluator\n",
    "\n",
    "We need a Bluesky evaluator that actually launches the experiment using the `RunEngine` and retreives the result using `Databroker`. Here we need to specify which detectors will produce the image as well as which motors we will be moving. Also, we pass the evaulation function here to produce the objective values.\n",
    "\n",
    "This evaluator will be used to produce the raw data needed by Ax to optimize the parameters to satisfy our objectives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from blop.integrations.ax import create_bluesky_evaluator\n",
    "\n",
    "\n",
    "evaluator = create_bluesky_evaluator(\n",
    "    RE, db, [beamline.det], [beamline.kbv_dsv, beamline.kbv_usv, beamline.kbh_dsh, beamline.kbh_ush], evaluate\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optimize!\n",
    "\n",
    "Finally, with all of our experimental setup done, we can optimize the parameters to satisfy our objectives.\n",
    "\n",
    "For this example, Ax will optimize the 4 motor positions to produce the greatest intensity beam with the smallest beam width and height (smallest area). It does this by first running a couple of `Trial`s which are random samples, then the remainder using Bayesian optimization through BoTorch.\n",
    "\n",
    "A single Ax `Trial` represents the training and evaluation of BoTorch models to produce a suggested next `Arm`. An `Arm` in Ax is a single parameterization to be evaluated while a `Trial` can consist of many `Arm`s. In this demo, we have a single `Arm` per `Trial`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for _ in range(25):\n",
    "    parameterization, trial_index = ax_client.get_next_trial()\n",
    "    ax_client.complete_trial(trial_index=trial_index, raw_data=evaluator(parameterization))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analyze Results\n",
    "\n",
    "Below we will show how we can use Ax to visualize the results and retrieve each step of the experiment that was run. This is where Ax becomes extremely useful for beamline optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax_client.experiment.to_df()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Viewing slices of parameter space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ax.utils.notebook.plotting import render\n",
    "from ax.plot.slice import plot_slice\n",
    "\n",
    "model = ax_client.generation_strategy.model\n",
    "render(plot_slice(model, \"bl_kbv_dsv\", \"beam_intensity\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "render(plot_slice(model, \"bl_kbv_dsv\", \"beam_area\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Viewing each arm's objective values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ax.plot.scatter import interact_fitted\n",
    "\n",
    "render(interact_fitted(model, rel=False))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualizing the optimal beam\n",
    "\n",
    "Below we get the optimal parameters, move the motors to their optimal positions, and observe the resulting beam."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimal_arm = next(iter(ax_client.get_pareto_optimal_parameters()))\n",
    "optimal_parameters = ax_client.get_trial(optimal_arm).arm.parameters\n",
    "optimal_parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bluesky.plans import list_scan\n",
    "\n",
    "scan_motor_params = []\n",
    "for motor in [beamline.kbv_dsv, beamline.kbv_usv, beamline.kbh_dsh, beamline.kbh_ush]:\n",
    "    scan_motor_params.append(motor)\n",
    "    scan_motor_params.append([optimal_parameters[motor.name]])\n",
    "RE(list_scan([beamline.det], *scan_motor_params))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "plt.imshow(db[-1].table(fill=True)[\"bl_det_image\"].iloc[0])\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "blop-dev",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}