Add plan for automatically finding the beam centre on the OAV

pyproject.toml

@@ -48,7 +48,7 @@ dependencies = [
     "ophyd >= 1.10.5",
     "ophyd-async >= 0.14.0",
     "bluesky >= 1.14.6",
-    "dls-dodal @ git+https://github.com/DiamondLightSource/dodal.git@62960e0e587bf86943ce1b581848fa131ef884d5",
+    "dls-dodal @ git+https://github.com/DiamondLightSource/dodal.git@1726_beam_centre_device_post_testing",
 ]
 
 

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 (
+    find_beam_centres,
+    optimise_transmission_with_oav,
     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_transmission_with_oav",
+    "find_beam_centres",
 ]

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

@@ -6,8 +6,10 @@
 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.oav.oav_detector import OAV, ZoomControllerWithBeamCentres
 from dodal.devices.robot import BartRobot, PinMounted
 from dodal.devices.scintillator import InOut, Scintillator
 from dodal.devices.xbpm_feedback import XBPMFeedback
@@ -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,
+        shutter,
+        initial_wait_group,
     )
-
+    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,10 +88,11 @@ 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
+    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.
     """
     pin_mounted = yield from bps.rd(robot.gonio_pin_sensor)
@@ -91,6 +109,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 +130,174 @@ 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 _max_pixel_at_transmission(
+    max_pixel: MaxPixel,
+    attenuator: BinaryFilterAttenuator,
+    xbpm_feedback: XBPMFeedback,
+    transmission: float,
+):
+    yield from bps.trigger(xbpm_feedback, wait=True)
+    yield from bps.mv(attenuator, transmission)
+    yield from bps.trigger(max_pixel, wait=True)
+    return (yield from bps.rd(max_pixel.max_pixel_val))
+
+
+def optimise_transmission_with_oav(
+    upper_bound: float = 100,
+    lower_bound: float = 0,
+    target_brightness_fraction: 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 (target_brightness_fraction) 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 this.
+    Args:
+        upper_bound: Maximum transmission which will be searched. In percent.
+        lower_bound: Minimum transmission which will be searched. In percent.
+        target_brightness_fraction: Fraction of the brightest pixel at 100%
+                    transmission which should be used as the target max pixel brightness.
+        tolerance: Amount the brightness can be off by and still find a match.
+        max_iterations: Maximum amount of iterations.
+    """
+
+    if upper_bound < lower_bound:
+        raise ValueError(
+            f"Upper bound ({upper_bound}) must be higher than lower bound {lower_bound}"
+        )
+
+    brightest_pixel_at_full_beam = yield from _max_pixel_at_transmission(
+        max_pixel, attenuator, xbpm_feedback, 1
+    )
+
+    if brightest_pixel_at_full_beam == 0:
+        raise ValueError("No beam found at full transmission")
+
+    target_pixel_brightness = brightest_pixel_at_full_beam * target_brightness_fraction
+    LOGGER.info(
+        f"Optimising until max pixel in image has a value of {target_pixel_brightness}"
+    )
+
+    iterations = 0
+
+    while iterations < max_iterations:
+        mid = round((upper_bound + lower_bound) / 2, 2)  # limit to 2 dp
+        LOGGER.info(f"On iteration {iterations}")
+
+        brightest_pixel = yield from _max_pixel_at_transmission(
+            max_pixel, attenuator, xbpm_feedback, mid / 100
+        )
+
+        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_brightness - tolerance
+            <= brightest_pixel
+            <= target_pixel_brightness + tolerance
+        ):
+            mid = round(mid, 0)
+            LOGGER.info(f"\nOptimal transmission found: {mid}")
+            return mid
+
+        # condition for too low so want to try higher
+        elif brightest_pixel < target_pixel_brightness - tolerance:
+            LOGGER.info("Result: Too low \n")
+            lower_bound = mid
+
+        # condition for too high so want to try lower
+        elif brightest_pixel > target_pixel_brightness + tolerance:
+            LOGGER.info("Result: Too high \n")
+            upper_bound = mid
+        iterations += 1
+    raise StopIteration("Max iterations reached")
+
+
+def _get_all_zoom_levels(
+    zoom_controller: ZoomControllerWithBeamCentres,
+) -> MsgGenerator[list[str]]:
+    zoom_levels = []
+    level_signals = [
+        centring_device.level_name
+        for centring_device in zoom_controller.beam_centres.values()
+    ]
+    for signal in level_signals:
+        level_name = yield from bps.rd(signal)
+        if level_name:
+            zoom_levels.append(level_name)
+    return zoom_levels
+
+
+def find_beam_centres(
+    zoom_levels_to_centre: list[str] | None = None,
+    zoom_levels_to_optimise_transmission: list[str] | None = None,
+    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"),
+    zoom_controller: ZoomControllerWithBeamCentres = inject("zoom_controller"),
+    shutter: ZebraShutter = inject("sample_shutter"),
+) -> MsgGenerator:
+    """
+    zoom_levels: The levels to do centring at, by default runs at all known zoom levels.
+    """
+    if zoom_levels_to_optimise_transmission is None:
+        zoom_levels_to_optimise_transmission = ["1.0x", "7.5x"]
+
+    if zoom_levels_to_centre is None:
+        zoom_levels_to_centre = yield from _get_all_zoom_levels(zoom_controller)
+
+    LOGGER.info("Preparing beamline for images...")
+    yield from _prepare_beamline_for_scintillator_images(
+        robot,
+        beamstop,
+        backlight,
+        scintillator,
+        xbpm_feedback,
+        shutter,
+        initial_wait_group,
+    )
+
+    for centring_device in zoom_controller.beam_centres.values():
+        zoom_name = yield from bps.rd(centring_device.level_name)
+        if zoom_name in zoom_levels_to_centre:
+            LOGGER.info(f"Moving to zoom level {zoom_name}")
+            yield from bps.abs_set(zoom_controller, zoom_name, wait=True)
+            if zoom_name in zoom_levels_to_optimise_transmission:
+                LOGGER.info(f"Optimising transmission at zoom level {zoom_name}")
+                yield from optimise_transmission_with_oav(
+                    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)
+            centre_x = round(centre_x)
+            centre_y = round(centre_y)
+            LOGGER.info(f"Writing centre values ({centre_x}, {centre_y}) to OAV PVs")
+            yield from bps.mv(
+                centring_device.x_centre, centre_x, centring_device.y_centre, centre_y
+            )
+
+    LOGGER.info("Done!")

tests/conftest.py

@@ -10,7 +10,7 @@
 from pathlib import Path
 from types import ModuleType
 from typing import Any
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import MagicMock, patch
 
 import numpy
 import pydantic
@@ -510,13 +510,7 @@ def oav(test_config_files):
     )
     oav = i03.oav.build(mock=True, connect_immediately=True, params=parameters)
 
-    zoom_levels_list = ["1.0x", "3.0x", "5.0x", "7.5x", "10.0x", "15.0x"]
-    oav.zoom_controller._get_allowed_zoom_levels = AsyncMock(
-        return_value=zoom_levels_list
-    )
-    # Equivalent to previously set values for microns and beam centre
     set_mock_value(oav.zoom_controller.level, "5.0x")
-
     set_mock_value(oav.grid_snapshot.x_size, 1024)
     set_mock_value(oav.grid_snapshot.y_size, 768)
 

tests/system_tests/conftest.py

@@ -201,10 +201,6 @@ async def trigger_with_test_image(self):
     ):
         mock_get.return_value.__aenter__.return_value = empty_response
         set_mock_value(oav.zoom_controller.level, "1.0x")
-        zoom_levels_list = ["1.0x", "3.0x", "5.0x", "7.5x", "10.0x", "15.0x"]
-        oav.zoom_controller._get_allowed_zoom_levels = AsyncMock(
-            return_value=zoom_levels_list
-        )
         yield oav