Add plan for automatically finding the beam centre on the OAV

src/mx_bluesky/Getting started.ipynb

@@ -52,6 +52,7 @@
     "import importlib\n",
     "\n",
     "from dodal.utils import collect_factories\n",
+    "\n",
     "beamline = \"i02_2\"\n",
     "module_name = f\"dodal.beamlines.{beamline}\"\n",
     "beamline_module = importlib.import_module(module_name)\n",

src/mx_bluesky/beamlines/i04/__init__.py

@@ -2,6 +2,8 @@
     i04_default_grid_detect_and_xray_centre,
 )
 from mx_bluesky.beamlines.i04.oav_centering_plans.oav_imaging import (
+    automated_centring,
+    optimise_oav_transmission_binary_search,
     take_oav_image_with_scintillator_in,
 )
 from mx_bluesky.beamlines.i04.thawing_plan import (
@@ -16,4 +18,6 @@
     "i04_default_grid_detect_and_xray_centre",
     "thaw_and_murko_centre",
     "take_oav_image_with_scintillator_in",
+    "optimise_oav_transmission_binary_search",
+    "automated_centring",
 ]

src/mx_bluesky/beamlines/i04/oav_centering_plans/oav_imaging.py

@@ -6,6 +6,8 @@
 from dodal.common import inject
 from dodal.devices.attenuator.attenuator import BinaryFilterAttenuator
 from dodal.devices.backlight import Backlight
+from dodal.devices.i04.beam_centre import CentreEllipseMethod
+from dodal.devices.i04.max_pixel import MaxPixel
 from dodal.devices.mx_phase1.beamstop import Beamstop, BeamstopPositions
 from dodal.devices.oav.oav_detector import OAV
 from dodal.devices.robot import BartRobot, PinMounted
@@ -19,6 +21,7 @@
 from ophyd_async.core import InOut as core_INOUT
 
 from mx_bluesky.common.utils.exceptions import BeamlineStateError
+from mx_bluesky.common.utils.log import LOGGER
 
 initial_wait_group = "Wait for scint to move in"
 
@@ -48,21 +51,35 @@ def take_oav_image_with_scintillator_in(
                     defaults are always correct.
     """
 
+    LOGGER.info("prearing beamline")
     yield from _prepare_beamline_for_scintillator_images(
-        robot, beamstop, backlight, scintillator, xbpm_feedback, initial_wait_group
+        robot,
+        beamstop,
+        backlight,
+        scintillator,
+        xbpm_feedback,
+        initial_wait_group,
+        shutter,
     )
-
+    LOGGER.info("setting transmission")
     yield from bps.abs_set(attenuator, transmission, group=initial_wait_group)
 
     if image_name is None:
         image_name = f"{time.time_ns()}ATT{transmission * 100}"
-
+    LOGGER.info(f"using image name {image_name}")
+    LOGGER.info("Waiting for initial_wait_group...")
     yield from bps.wait(initial_wait_group)
 
+    LOGGER.info("Opening shutter...")
+
     yield from bps.abs_set(shutter.control_mode, ZebraShutterControl.MANUAL, wait=True)
     yield from bps.abs_set(shutter, ZebraShutterState.OPEN, wait=True)
 
-    take_and_save_oav_image(file_path=image_path, file_name=image_name, oav=oav)
+    LOGGER.info("Taking image...")
+
+    yield from take_and_save_oav_image(
+        file_path=image_path, file_name=image_name, oav=oav
+    )
 
 
 def _prepare_beamline_for_scintillator_images(
@@ -71,11 +88,16 @@ def _prepare_beamline_for_scintillator_images(
     backlight: Backlight,
     scintillator: Scintillator,
     xbpm_feedback: XBPMFeedback,
+    shutter: ZebraShutter,
     group: str,
 ) -> MsgGenerator:
     """
     Prepares the beamline for oav image by making sure the pin is NOT mounted and
     the beam is on (feedback check). Finally, the scintillator is moved in.
+
+     Args:
+        devices: These are the specific ophyd-devices used for the plan, the
+                    defaults are always correct.
     """
     pin_mounted = yield from bps.rd(robot.gonio_pin_sensor)
     if pin_mounted == PinMounted.PIN_MOUNTED:
@@ -91,6 +113,9 @@ def _prepare_beamline_for_scintillator_images(
 
     yield from bps.abs_set(scintillator.selected_pos, InOut.IN, group=group)
 
+    yield from bps.abs_set(shutter.control_mode, ZebraShutterControl.MANUAL, wait=True)
+    yield from bps.abs_set(shutter, ZebraShutterState.OPEN, wait=True)
+
 
 def take_and_save_oav_image(
     file_name: str,
@@ -109,7 +134,178 @@ def take_and_save_oav_image(
     if not os.path.exists(full_file_path):
         yield from bps.abs_set(oav.snapshot.filename, file_name, group=group)
         yield from bps.abs_set(oav.snapshot.directory, file_path, group=group)
-        yield from bps.wait(group)
+        yield from bps.wait(group, timeout=60)
         yield from bps.trigger(oav.snapshot, wait=True)
     else:
         raise FileExistsError("OAV image file path already exists")
+
+
+def _get_max_pixel_from_100_transmission(
+    max_pixel: MaxPixel,
+    attenuator: BinaryFilterAttenuator,
+):
+    yield from bps.mv(attenuator, 1)  # 100 % transmission
+    yield from bps.trigger(max_pixel, wait=True)
+    target_brightest_pixel = yield from bps.rd(max_pixel.max_pixel_val)
+    return target_brightest_pixel
+
+
+def optimise_oav_transmission_binary_search(
+    upper_bound: float = 100,  # in percent
+    lower_bound: float = 0,  # in percent
+    frac_of_max: float = 0.75,
+    tolerance: int = 5,
+    max_iterations: int = 10,
+    max_pixel: MaxPixel = inject("max_pixel"),
+    attenuator: BinaryFilterAttenuator = inject("attenuator"),
+    xbpm_feedback: XBPMFeedback = inject("xbpm_feedback"),
+) -> MsgGenerator:
+    """
+    Plan to find the optimal oav transmission. First the brightest pixel at 100%
+    transmission is taken. A fraction of this (frac_of_max) is taken as the target -
+    as in the optimal transmission will have it's max pixel as the set target.
+    A binary search is used to reach the target.
+    Args:
+        upper_bound: Maximum transmission which will be searched.
+        lower_bound: Minimum transmission which will be searched.
+        frac_of_max: Fraction of the brightest pixel at 100% transmission which should be
+                     used as the target max pixel brightness.
+        tolerance: Amount the search can be off by and still find a match.
+        max_iterations: Maximum amount of iterations.
+    """
+    brightest_pixel_sat = yield from _get_max_pixel_from_100_transmission(
+        max_pixel, attenuator
+    )
+    target_pixel_l = brightest_pixel_sat * frac_of_max
+    LOGGER.info(f"~~Target luminosity: {target_pixel_l}~~\n")
+
+    iterations = 0
+
+    while iterations < max_iterations:
+        mid = round((upper_bound + lower_bound) / 2, 2)  # limit to 2 dp
+        LOGGER.info(f"on iteration {iterations}")
+
+        yield from bps.mv(attenuator, mid / 100)
+        yield from bps.trigger(xbpm_feedback, wait=True)
+        yield from bps.trigger(max_pixel, wait=True)
+        brightest_pixel = yield from bps.rd(max_pixel.max_pixel_val)
+
+        # brightest_pixel = get_max_pixel_value_from_transmission(transmission=mid)
+        LOGGER.info(f"Upper bound is: {upper_bound}, Lower bound is: {lower_bound}")
+        LOGGER.info(
+            f"Testing transmission {mid}, brightest pixel found {brightest_pixel}"
+        )
+
+        if target_pixel_l - tolerance < brightest_pixel < target_pixel_l + tolerance:
+            mid = round(mid, 0)
+            LOGGER.info(f"\nOptimal transmission found: {mid}")
+            yield from bps.trigger(xbpm_feedback)
+            return mid
+
+        # condition for too low so want to try higher
+        elif brightest_pixel < target_pixel_l - tolerance:
+            LOGGER.info("Result: Too low \n")
+            lower_bound = mid
+
+        # condition for too high so want to try lower
+        elif brightest_pixel > target_pixel_l + tolerance:
+            LOGGER.info("Result: Too high \n")
+            upper_bound = mid
+        iterations += 1
+    raise StopIteration("Max iterations reached")
+
+
+def automated_centring(
+    zoom_levels: list[str] = [
+        "1.0x",
+        "1.5x",
+        "2.0x",
+        "2.5x",
+        "3.0x",
+        "5.0x",
+        "7.5x",
+        "10.0x",
+    ],
+    robot: BartRobot = inject("robot"),
+    beamstop: Beamstop = inject("beamstop"),
+    backlight: Backlight = inject("backlight"),
+    scintillator: Scintillator = inject("scintillator"),
+    xbpm_feedback: XBPMFeedback = inject("xbpm_feedback"),
+    max_pixel: MaxPixel = inject("max_pixel"),
+    centre_ellipse: CentreEllipseMethod = inject("beam_centre"),
+    attenuator: BinaryFilterAttenuator = inject("attenuator"),
+    oav: OAV = inject("oav"),
+    shutter: ZebraShutter = inject("sample_shutter"),
+) -> MsgGenerator:
+    zoom_level_to_dict = {
+        "7.5x": [
+            oav.zoom_controller.x_placeholder_zoom_7,
+            oav.zoom_controller.y_placeholder_zoom_7,
+        ],
+        "1.0x": [
+            oav.zoom_controller.x_placeholder_zoom_1,
+            oav.zoom_controller.y_placeholder_zoom_1,
+        ],
+        "1.5x": [
+            oav.zoom_controller.x_placeholder_zoom_2,
+            oav.zoom_controller.y_placeholder_zoom_2,
+        ],
+        "2.0x": [
+            oav.zoom_controller.x_placeholder_zoom_3,
+            oav.zoom_controller.y_placeholder_zoom_3,
+        ],
+        "2.5x": [
+            oav.zoom_controller.x_placeholder_zoom_4,
+            oav.zoom_controller.y_placeholder_zoom_4,
+        ],
+        "3.0x": [
+            oav.zoom_controller.x_placeholder_zoom_5,
+            oav.zoom_controller.y_placeholder_zoom_5,
+        ],
+        "5.0x": [
+            oav.zoom_controller.x_placeholder_zoom_6,
+            oav.zoom_controller.y_placeholder_zoom_6,
+        ],
+        "10.0x": [
+            oav.zoom_controller.x_placeholder_zoom_8,
+            oav.zoom_controller.y_placeholder_zoom_8,
+        ],
+    }
+
+    LOGGER.info("Preparing beamline for images...")
+    yield from _prepare_beamline_for_scintillator_images(
+        robot,
+        beamstop,
+        backlight,
+        scintillator,
+        xbpm_feedback,
+        shutter,
+        initial_wait_group,
+    )
+
+    for zoom in zoom_levels:
+        LOGGER.info(f"Moving to zoom level {zoom}")
+        yield from bps.abs_set(oav.zoom_controller, zoom, wait=True)
+        yield from bps.sleep(1)
+        if zoom == "7.5x" or zoom == "1.0x":
+            LOGGER.info(f"Optimising transmission (zoom level {zoom})")
+            yield from optimise_oav_transmission_binary_search(
+                100,
+                0,
+                max_pixel=max_pixel,
+                attenuator=attenuator,
+                xbpm_feedback=xbpm_feedback,
+            )
+
+        yield from bps.trigger(centre_ellipse, wait=True)
+        centre_x = yield from bps.rd(centre_ellipse.center_x_val)
+        centre_y = yield from bps.rd(centre_ellipse.center_y_val)
+        LOGGER.info(f"Centre X: {centre_x}, Centre Y: {centre_y}")
+        centre_x = round(centre_x)
+        centre_y = round(centre_y)
+        x_signal = zoom_level_to_dict[zoom][0]
+        y_signal = zoom_level_to_dict[zoom][1]
+        LOGGER.info("Writing centre values to OAV PVs")
+        yield from bps.mv(x_signal, centre_x, y_signal, centre_y)
+
+    LOGGER.info("Done!")