XRT KB Mirror Demo

For ophyd beamline setup see:

• https://github.com/NSLS-II/blop/blob/main/src/blop/sim/xrt_beamline.py

• https://github.com/NSLS-II/blop/blob/main/src/blop/sim/xrt_kb_model.py

The picture below displays beam from geometric source propagating through a pair of toroidal mirrors focusing the beam on screen. Simulation of a KB setup.

import time
from datetime import datetime
import logging

import plotly.io as pio

pio.renderers.default = "notebook"
import bluesky.plan_stubs as bps  # noqa F401
import bluesky.plans as bp  # noqa F401
import databroker  # type: ignore[import-untyped]
import matplotlib.pyplot as plt
import tiled.client.container
from bluesky.callbacks import best_effort
from bluesky.callbacks.tiled_writer import TiledWriter
from bluesky.run_engine import RunEngine
from databroker import Broker
from ophyd.utils import make_dir_tree  # type: ignore[import-untyped]
from tiled.client import from_uri  # type: ignore[import-untyped]
from tiled.server import SimpleTiledServer

from blop import DOF, Objective
from blop.ax import Agent
from blop.sim import HDF5Handler
from blop.sim.xrt_beamline import DatabrokerBeamline, TiledBeamline

# Suppress noisy logs from httpx 
logging.getLogger("httpx").setLevel(logging.WARNING)

DETECTOR_STORAGE = "/tmp/blop/sim"

[WARNING 11-06 21:31:08] ax.service.utils.with_db_settings_base: Ax currently requires a sqlalchemy version below 2.0. This will be addressed in a future release. Disabling SQL storage in Ax for now, if you would like to use SQL storage please install Ax with mysql extras via `pip install ax-platform[mysql]`.

tiled_server = SimpleTiledServer(readable_storage=[DETECTOR_STORAGE])
tiled_client = from_uri(tiled_server.uri)
tiled_writer = TiledWriter(tiled_client)


def setup_re_env(db_type="default", root_dir="/default/path", method="tiled"):
    RE = RunEngine({})
    bec = best_effort.BestEffortCallback()
    RE.subscribe(bec)
    _ = make_dir_tree(datetime.now().year, base_path=root_dir)

    if method == "tiled":
        RE.subscribe(tiled_writer)
        return {"RE": RE, "db": tiled_client, "bec": bec}

    elif method == "databroker":
        db = Broker.named(db_type)
        db.reg.register_handler("HDF5", HDF5Handler, overwrite=True)
        try:
            databroker.assets.utils.install_sentinels(db.reg.config, version=1)
        except Exception:
            pass
        RE.subscribe(db.insert)
        return {
            "RE": RE,
            "db": db,
            "bec": bec,
        }
    else:
        raise ValueError("The method for data storage used is not supported")


def register_handlers(db, handlers):
    for handler_spec, handler_class in handlers.items():
        db.reg.register_handler(handler_spec, handler_class, overwrite=True)


env = setup_re_env(db_type="temp", root_dir="/tmp/blop/sim", method="tiled")
globals().update(env)
bec.disable_plots()

2025-11-06 21:31:12.105 INFO: Subprocess stdout:
2025-11-06 21:31:12.106 INFO: Subprocess stderr: Database sqlite+aiosqlite:////tmp/tmpu3nq3t3l/catalog.db is new. Creating tables.
Database initialized.
Tiled version 0.2.0
2025-11-06 21:31:12.447 INFO: Tiled version 0.2.0
2025-11-06 21:31:12.453 INFO: Context impl SQLiteImpl.
2025-11-06 21:31:12.454 INFO: Will assume non-transactional DDL.

plt.ion()

h_opt = 0
dh = 5

R1, dR1 = 40000, 10000
R2, dR2 = 20000, 10000

if isinstance(db, tiled.client.container.Container):
    beamline = TiledBeamline(name="bl")
else:
    beamline = DatabrokerBeamline(name="bl")
time.sleep(1)
dofs = [
    DOF(movable=beamline.kbv_dsv, search_domain=(R1 - dR1, R1 + dR1)),
    DOF(movable=beamline.kbh_dsh, search_domain=(R2 - dR2, R2 + dR2)),
]

objectives = [
    Objective(name="bl_det_sum", target="max"),
    Objective(name="bl_det_wid_x", target="min"),
    Objective(name="bl_det_wid_y", target="min"),
]

agent = Agent(
    readables=[beamline.det],
    dofs=dofs,
    objectives=objectives,
    db=db,
)
agent.configure_experiment(
    name="xrt-blop-demo",
    description="A demo of the Blop agent with XRT simulated beamline",
    experiment_type="demo",
)

2025-11-06 21:31:13.717 INFO: Configuring optimization with objective: bl_det_sum, -bl_det_wid_x, -bl_det_wid_y and outcome constraints: []

# Number of iterations can be increased to be more specific
RE(agent.learn(iterations=15))

Transient Scan ID: 1     Time: 2025-11-06 21:31:13
Persistent Unique Scan ID: '287a7873-96b7-4bc0-a792-4a4625d91d8e'
ToroidMirror2
center:
[0, 11000, np.float64(176.32838854298492)]
New stream: 'primary'
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|   seq_num |       time | bl_kbv_dsv | bl_kbh_dsh | bl_det_sum | bl_det_cen_x | bl_det_cen_y | bl_det_wid_x | bl_det_wid_y |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|         1 | 21:31:13.9 |  40000.000 |  20000.000 |  22072.057 |      202.167 |      145.747 |       17.675 |       38.519 |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
generator list_scan ['287a7873'] (scan num: 1)


Transient Scan ID: 2     Time: 2025-11-06 21:31:14
Persistent Unique Scan ID: '6a5c1669-cbf4-4cdf-b323-d7d4e224fef7'
New stream: 'primary'
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|   seq_num |       time | bl_kbv_dsv | bl_kbh_dsh | bl_det_sum | bl_det_cen_x | bl_det_cen_y | bl_det_wid_x | bl_det_wid_y |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|         1 | 21:31:14.7 |  33694.285 |  29875.351 |   2630.221 |      204.292 |      149.945 |      277.333 |       77.764 |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
generator list_scan ['6a5c1669'] (scan num: 2)


Transient Scan ID: 3     Time: 2025-11-06 21:31:15
Persistent Unique Scan ID: '96996d58-2d98-49a0-8926-02ce6f4effdb'
New stream: 'primary'
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|   seq_num |       time | bl_kbv_dsv | bl_kbh_dsh | bl_det_sum | bl_det_cen_x | bl_det_cen_y | bl_det_wid_x | bl_det_wid_y |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|         1 | 21:31:15.4 |  46249.450 |  10673.315 |     28.583 |      199.075 |      157.787 |      398.150 |      118.725 |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
generator list_scan ['96996d58'] (scan num: 3)


Transient Scan ID: 4     Time: 2025-11-06 21:31:16
Persistent Unique Scan ID: 'f232cd98-72c7-4397-979a-791ba4fe981c'
New stream: 'primary'
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|   seq_num |       time | bl_kbv_dsv | bl_kbh_dsh | bl_det_sum | bl_det_cen_x | bl_det_cen_y | bl_det_wid_x | bl_det_wid_y |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|         1 | 21:31:16.4 |  41033.396 |  21852.533 |  20443.620 |      205.154 |      157.480 |       43.841 |       55.031 |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
generator list_scan ['f232cd98'] (scan num: 4)


Transient Scan ID: 5     Time: 2025-11-06 21:31:16
Persistent Unique Scan ID: '5088a8ad-3742-4561-a519-bcf82a1d9243'
New stream: 'primary'
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|   seq_num |       time | bl_kbv_dsv | bl_kbh_dsh | bl_det_sum | bl_det_cen_x | bl_det_cen_y | bl_det_wid_x | bl_det_wid_y |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|         1 | 21:31:17.0 |  38558.778 |  17602.459 |   6069.032 |      208.708 |      149.092 |      133.254 |       17.667 |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
generator list_scan ['5088a8ad'] (scan num: 5)


Transient Scan ID: 6     Time: 2025-11-06 21:31:19
Persistent Unique Scan ID: '7edef234-6ecf-4d49-aa4e-c409da521332'
New stream: 'primary'
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|   seq_num |       time | bl_kbv_dsv | bl_kbh_dsh | bl_det_sum | bl_det_cen_x | bl_det_cen_y | bl_det_wid_x | bl_det_wid_y |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|         1 | 21:31:19.7 |  37809.592 |  20397.339 |  23738.547 |      194.678 |      147.032 |       15.444 |       17.590 |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
generator list_scan ['7edef234'] (scan num: 6)


Transient Scan ID: 7     Time: 2025-11-06 21:31:22
Persistent Unique Scan ID: '9630fa10-044c-405c-b3a4-50be8110e2a2'
New stream: 'primary'
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|   seq_num |       time | bl_kbv_dsv | bl_kbh_dsh | bl_det_sum | bl_det_cen_x | bl_det_cen_y | bl_det_wid_x | bl_det_wid_y |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|         1 | 21:31:23.0 |  37253.422 |  20670.114 |  23509.207 |      193.350 |      150.200 |       13.700 |       22.200 |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
generator list_scan ['9630fa10'] (scan num: 7)


Transient Scan ID: 8     Time: 2025-11-06 21:31:28
Persistent Unique Scan ID: '2c4c7a98-4a04-41e8-90d6-3b1795eeb145'
New stream: 'primary'
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|   seq_num |       time | bl_kbv_dsv | bl_kbh_dsh | bl_det_sum | bl_det_cen_x | bl_det_cen_y | bl_det_wid_x | bl_det_wid_y |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|         1 | 21:31:28.5 |  37837.182 |  19223.562 |  23153.228 |      205.650 |      145.560 |       75.480 |       21.664 |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
generator list_scan ['2c4c7a98'] (scan num: 8)


Transient Scan ID: 9     Time: 2025-11-06 21:31:32
Persistent Unique Scan ID: '5e739de0-90e3-42f2-a363-c44a87c99218'
New stream: 'primary'
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|   seq_num |       time | bl_kbv_dsv | bl_kbh_dsh | bl_det_sum | bl_det_cen_x | bl_det_cen_y | bl_det_wid_x | bl_det_wid_y |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|         1 | 21:31:32.9 |  37710.892 |  19929.899 |  23490.637 |      200.755 |      144.436 |       24.472 |       19.242 |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
generator list_scan ['5e739de0'] (scan num: 9)


Transient Scan ID: 10     Time: 2025-11-06 21:31:38
Persistent Unique Scan ID: '6bc60ff9-f8f6-483d-a973-93ab41c2e9d1'
New stream: 'primary'
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|   seq_num |       time | bl_kbv_dsv | bl_kbh_dsh | bl_det_sum | bl_det_cen_x | bl_det_cen_y | bl_det_wid_x | bl_det_wid_y |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|         1 | 21:31:38.3 |  37692.572 |  14272.157 |   1782.050 |      199.867 |      148.974 |      352.533 |       22.831 |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
generator list_scan ['6bc60ff9'] (scan num: 10)


Transient Scan ID: 11     Time: 2025-11-06 21:31:44
Persistent Unique Scan ID: '8a44d2fc-a2cf-4b87-9357-fc22748c54b7'
New stream: 'primary'
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|   seq_num |       time | bl_kbv_dsv | bl_kbh_dsh | bl_det_sum | bl_det_cen_x | bl_det_cen_y | bl_det_wid_x | bl_det_wid_y |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|         1 | 21:31:44.3 |  37627.898 |  20205.847 |  23667.335 |      196.825 |      145.851 |       16.495 |       18.990 |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
generator list_scan ['8a44d2fc'] (scan num: 11)


Transient Scan ID: 12     Time: 2025-11-06 21:31:48
Persistent Unique Scan ID: '5bda4347-c2e8-4ad1-978e-61137f1dabe4'
New stream: 'primary'
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|   seq_num |       time | bl_kbv_dsv | bl_kbh_dsh | bl_det_sum | bl_det_cen_x | bl_det_cen_y | bl_det_wid_x | bl_det_wid_y |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|         1 | 21:31:48.7 |  37602.849 |  22920.856 |  22157.994 |      204.510 |      151.850 |       97.280 |       21.974 |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
generator list_scan ['5bda4347'] (scan num: 12)


Transient Scan ID: 13     Time: 2025-11-06 21:31:54
Persistent Unique Scan ID: 'dd89c09f-4347-4db3-a212-34578c73c4ce'
New stream: 'primary'
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|   seq_num |       time | bl_kbv_dsv | bl_kbh_dsh | bl_det_sum | bl_det_cen_x | bl_det_cen_y | bl_det_wid_x | bl_det_wid_y |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|         1 | 21:31:54.5 |  38081.794 |  20005.243 |  23653.156 |      198.333 |      145.035 |       25.241 |       19.760 |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
generator list_scan ['dd89c09f'] (scan num: 13)


Transient Scan ID: 14     Time: 2025-11-06 21:32:00
Persistent Unique Scan ID: '35ba91a8-dbd8-447a-94e8-0ea5cab4a004'
New stream: 'primary'
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|   seq_num |       time | bl_kbv_dsv | bl_kbh_dsh | bl_det_sum | bl_det_cen_x | bl_det_cen_y | bl_det_wid_x | bl_det_wid_y |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|         1 | 21:32:01.0 |  37457.284 |  20116.746 |  23504.136 |      198.171 |      145.713 |       18.430 |       19.960 |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
generator list_scan ['35ba91a8'] (scan num: 14)


Transient Scan ID: 15     Time: 2025-11-06 21:32:09
Persistent Unique Scan ID: '651053fb-fd17-435d-9a88-efbc51f00e9b'
New stream: 'primary'
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|   seq_num |       time | bl_kbv_dsv | bl_kbh_dsh | bl_det_sum | bl_det_cen_x | bl_det_cen_y | bl_det_wid_x | bl_det_wid_y |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|         1 | 21:32:09.5 |  30000.000 |  20940.185 |    648.749 |      193.912 |      167.250 |       10.575 |       87.500 |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
generator list_scan ['651053fb'] (scan num: 15)

2025-11-06 21:31:13.742 INFO: Executing plan <generator object Agent.learn at 0x7f4dd8973840>
2025-11-06 21:31:13.743 INFO: Change state on <bluesky.run_engine.RunEngine object at 0x7f4dd8ad3200> from 'idle' -> 'running'
2025-11-06 21:32:10.074 INFO: Change state on <bluesky.run_engine.RunEngine object at 0x7f4dd8ad3200> from 'running' -> 'idle'
2025-11-06 21:32:10.075 INFO: Cleaned up from plan <generator object Agent.learn at 0x7f4dd8973840>

('287a7873-96b7-4bc0-a792-4a4625d91d8e',
 '6a5c1669-cbf4-4cdf-b323-d7d4e224fef7',
 '96996d58-2d98-49a0-8926-02ce6f4effdb',
 'f232cd98-72c7-4397-979a-791ba4fe981c',
 '5088a8ad-3742-4561-a519-bcf82a1d9243',
 '7edef234-6ecf-4d49-aa4e-c409da521332',
 '9630fa10-044c-405c-b3a4-50be8110e2a2',
 '2c4c7a98-4a04-41e8-90d6-3b1795eeb145',
 '5e739de0-90e3-42f2-a363-c44a87c99218',
 '6bc60ff9-f8f6-483d-a973-93ab41c2e9d1',
 '8a44d2fc-a2cf-4b87-9357-fc22748c54b7',
 '5bda4347-c2e8-4ad1-978e-61137f1dabe4',
 'dd89c09f-4347-4db3-a212-34578c73c4ce',
 '35ba91a8-dbd8-447a-94e8-0ea5cab4a004',
 '651053fb-fd17-435d-9a88-efbc51f00e9b')

_ = agent.plot_objective(x_dof_name="bl_kbh_dsh", y_dof_name="bl_kbv_dsv", objective_name="bl_det_sum")

bl_kbh_dsh, bl_kbv_dsv vs. bl_det_sum

The contour plot visualizes the predicted outcomes for bl_det_sum across a two-dimensional parameter space, with other parameters held fixed at their status_quo value (or mean value if status_quo is unavailable). This plot helps in identifying regions of optimal performance and understanding how changes in the selected parameters influence the predicted outcomes. Contour lines represent levels of constant predicted values, providing insights into the gradient and potential optima within the parameter space.

Visualizing the optimal beam

Below we get the optimal parameters, move the motors to their optimal positions, and observe the resulting beam.

optimal_parameters = next(iter(agent.client.get_pareto_frontier()))[0]
optimal_parameters

[INFO 11-06 21:32:13] ax.service.utils.best_point: Using inferred objective thresholds: [ObjectiveThreshold(bl_det_sum >= 23396.818968197644), ObjectiveThreshold(bl_det_wid_x <= 26.219958766413697), ObjectiveThreshold(bl_det_wid_y <= 22.416038599510138)], as objective thresholds were not specified as part of the optimization configuration on the experiment.

{'bl_kbv_dsv': 37809.592020963195, 'bl_kbh_dsh': 20397.338541652087}

from bluesky.plans import list_scan

scan_motor_params = []
for motor in [beamline.kbv_dsv, beamline.kbh_dsh]:
    scan_motor_params.append(motor)
    scan_motor_params.append([optimal_parameters[motor.name]])
uid = RE(list_scan([beamline.det], *scan_motor_params))

Transient Scan ID: 16     Time: 2025-11-06 21:32:13
Persistent Unique Scan ID: '244f9224-49d1-4db7-8dcc-db37bacfca2a'
New stream: 'primary'
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|   seq_num |       time | bl_kbv_dsv | bl_kbh_dsh | bl_det_sum | bl_det_cen_x | bl_det_cen_y | bl_det_wid_x | bl_det_wid_y |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
|         1 | 21:32:13.4 |  37809.592 |  20397.339 |  23722.704 |      194.369 |      147.373 |       15.603 |       18.955 |
+-----------+------------+------------+------------+------------+--------------+--------------+--------------+--------------+
generator list_scan ['244f9224'] (scan num: 16)

2025-11-06 21:32:13.268 INFO: Executing plan <generator object list_scan at 0x7f4d9c25aa40>
2025-11-06 21:32:13.269 INFO: Change state on <bluesky.run_engine.RunEngine object at 0x7f4dd8ad3200> from 'idle' -> 'running'
2025-11-06 21:32:13.661 INFO: Change state on <bluesky.run_engine.RunEngine object at 0x7f4dd8ad3200> from 'running' -> 'idle'
2025-11-06 21:32:13.662 INFO: Cleaned up from plan <generator object list_scan at 0x7f4d9c25aa40>

import matplotlib.pyplot as plt

image = db[uid[0]]["primary"]["bl_det_image"].read().squeeze()
plt.imshow(image)
plt.colorbar()
plt.show()