DAQ: zoom_manager function added to models.py and zoom_settings condensed into one function in config.py. This will be adapted when using redis and yaml files.

2025-09-30 15:44:07 +02:00
parent b8b3e1b900
commit e7f1603a76
4 changed files with 167 additions and 106 deletions
@@ -9,6 +9,8 @@ from aaredaqlib.coordinate import Coordinate
 from aaredaqlib.diffraction_geometry import DiffractionGeometry
 from aaredaqlib.sample_geometry import SampleGeometryModel

+from aaredaqlib.beamline import MXBeamline
+
 class StagePositionEnum(Enum):
    MEASURE = 0
    PARK = 1
@@ -427,35 +429,6 @@ class BoundingBoxModel(BaseModel):
    bottom_x: float
    bottom_y: float

-
-class LoopCenteringZoomModelElem(BaseModel):
-    zoom_value: float
-    sam_cam_gain: float
-    sam_cam_exp: float
-
-
-class LoopCenteringZoomModel(BaseModel):
-        z: List[LoopCenteringZoomModelElem]
-
-# zoom gain exp (X06DA)
-#    1  50 0.05
-#  200  50 0.05
-#  500  70 0.05
-#  700 150 0.05
-#  800 100 0.10
-# 1000 200 0.10
-
-
-def def_loop_centering_zoom() -> LoopCenteringZoomModel:
-    return LoopCenteringZoomModel(
-        z=[
-            LoopCenteringZoomModelElem(
-                zoom_value=1, sam_cam_gain=0, sam_cam_exp=0.05,
-            )
-        ]
-    )
-
-
 class BeamlineStateEnum(Enum):
    Maintenance = 1
    SampleExchange = 2
@@ -478,6 +451,10 @@ class SampleCameraSettings(BaseModel):
    gain: float
    exposure: float

+class ZoomModeEnum(Enum):
+    User = 1
+    LoopCenter = 2
+    BeamLocation = 3

 class ZoomModel(BaseModel):
    z:dict[float, SampleCameraSettings]
@@ -519,29 +496,75 @@ class ZoomModel(BaseModel):
        elem = self.z[closest]
        return SampleCameraSettings(gain=elem.gain, exposure=elem.exposure)

-def def_zoom() -> ZoomModel:
-    return ZoomModel(
-        z={
-            1: SampleCameraSettings(gain=0, exposure=0.05),
-            280: SampleCameraSettings(gain=0, exposure=0.05),
-            500: SampleCameraSettings(gain=0, exposure=0.05),
-            700: SampleCameraSettings(gain=0, exposure=0.05),
-            800: SampleCameraSettings(gain=0, exposure=0.10),
-            1000: SampleCameraSettings(gain=0, exposure=0.20)
-        }
-    )
+def def_zoom(beamline) -> ZoomModel:
+    print(f'user zoom for {beamline}')
+    if beamline == MXBeamline.X06DA:
+        return ZoomModel(
+            z={
+                1: SampleCameraSettings(gain=0, exposure=0.05),
+                280: SampleCameraSettings(gain=0, exposure=0.05),
+                500: SampleCameraSettings(gain=0, exposure=0.05),
+                700: SampleCameraSettings(gain=0, exposure=0.05),
+                800: SampleCameraSettings(gain=0, exposure=0.10),
+                1000: SampleCameraSettings(gain=0, exposure=0.20)
+            }
+        )
+    elif beamline == MXBeamline.X10SA:
+        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
+    elif beamline == MXBeamline.X06SA:
+        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
+    elif beamline == MXBeamline.SIMULATED:
+        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
+    else:
+        raise ValueError(f"Invalid beamline: {beamline}")
+
+def def_bl_zoom(beamline) -> ZoomModel:
+    if beamline == MXBeamline.X06DA:
+        return ZoomModel(
+            z={
+                1: SampleCameraSettings(gain=0, exposure=0.002),
+                280: SampleCameraSettings(gain=0, exposure=0.002),
+                500: SampleCameraSettings(gain=0, exposure=0.002),
+                700: SampleCameraSettings(gain=0, exposure=0.002),
+                800: SampleCameraSettings(gain=0, exposure=0.002),
+                1000: SampleCameraSettings(gain=0, exposure=0.005)
+            }
+        )
+    elif beamline == MXBeamline.X10SA:
+        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
+    elif beamline == MXBeamline.X06SA:
+        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
+    elif beamline == MXBeamline.SIMULATED:
+        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.002)})
+    else:
+        raise ValueError(f"Invalid beamline: {beamline}")
+
+def def_loop_centering_zoom(beamline) -> ZoomModel:
+    if beamline == MXBeamline.X06DA:
+        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.05)})
+    elif beamline == MXBeamline.X10SA:
+        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.05)})
+    elif beamline == MXBeamline.X06SA:
+        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.05)})
+    elif beamline == MXBeamline.SIMULATED:
+        return ZoomModel(z={1: SampleCameraSettings(gain=0, exposure=0.05)})
+    else:
+        raise ValueError(f"Invalid beamline: {beamline}")
+
+def zoom_manager(mode: ZoomModeEnum = ZoomModeEnum.User, beamline: MXBeamline = None) -> ZoomModel | None:
+    if beamline is None or beamline == MXBeamline.SIMULATED:
+        print('SIMULATED')
+        return def_zoom(beamline)
+    if mode == ZoomModeEnum.BeamLocation:
+        return def_bl_zoom(beamline)
+    elif mode == ZoomModeEnum.LoopCenter:
+        return def_loop_centering_zoom(beamline)
+    elif mode == ZoomModeEnum.User:
+        return def_zoom(beamline)
+    else:
+        raise ValueError(f"Invalid zoom mode: {mode}")
+

-def def_bl_zoom() -> ZoomModel:
-    return ZoomModel(
-        z={
-            1: SampleCameraSettings(gain=0, exposure=0.002),
-            280: SampleCameraSettings(gain=0, exposure=0.002),
-            500: SampleCameraSettings(gain=0, exposure=0.002),
-            700: SampleCameraSettings(gain=0, exposure=0.002),
-            800: SampleCameraSettings(gain=0, exposure=0.002),
-            1000: SampleCameraSettings(gain=0, exposure=0.005)
-        }
-    )

 class AutofocusSettings(BaseModel):
    center_x_pxl: float
@@ -10,11 +10,11 @@ from aaredaqlib.coordinate import Coordinate
 from aaredaqlib.models import (
    BeamlineSettingsModel,
    BeamMarkCoeffModel,
-    LoopCenteringZoomModel,
-    def_loop_centering_zoom,
+    ZoomModeEnum,
+    zoom_manager,
    SampleShortInfo,
    SessionStatus,
-    BeamlineStateEnum, SessionsStateEnum, SampleShortInfoList, CryojetSettingsModel, ZoomModel, def_zoom, def_bl_zoom,
+    BeamlineStateEnum, SessionsStateEnum, SampleShortInfoList, CryojetSettingsModel, ZoomModel
 )

 from aaredaqlib.beamline import MXBeamline
@@ -65,6 +65,7 @@ class BeamlineConfig:
    """

    def __init__(self, bl: MXBeamline):
+        self.__mxb = bl
        self.__bl = bl.value.lower()
        if bl is MXBeamline.SIMULATED:
            host = "localhost"
@@ -294,7 +295,6 @@ class BeamlineConfig:
        self.__client.set(f"{self.__bl}:cryojet_settings", data.model_dump_json())

    def get_alc_bkg(self, zoom: float, exp: float, gain: float) -> np.ndarray | None:
-        print(f"{self.__bl}:bkg{zoom:.1f}_{exp:.2f}_{gain:.1f}")
        return base64_to_numpy(self.__client.get(f"{self.__bl}:bkg{zoom:.1f}_{exp:.2f}_{gain:.1f}"))

    def put_alc_bkg(self, zoom: float, exp: float, gain: float, data: np.ndarray):
@@ -354,43 +354,54 @@ class BeamlineConfig:
        self.__client.set(f"{self.__bl}:beam_center_camera", data.model_dump_json())

 #TODO tidy up zoom functions
-
    @property
-    def alc_zoom_settings(self) -> LoopCenteringZoomModel:
-        tmp = self.__client.get(f"{self.__bl}:alc_zoom_settings")
-        if tmp is None:
-            return def_loop_centering_zoom()
+    def zoom_mode(self) -> ZoomModeEnum:
+        raw_value = self.__client.get(f"{self.__bl}:zoom_mode")
+        if raw_value is None:
+            print('no zoom mode given, defaulting to user mode')
+            return ZoomModeEnum.User
+        try:
+            int_value = int(raw_value)  # Ensure it's an integer
+            return ZoomModeEnum(int_value)  # Convert to ZoomModeEnum
+        except (ValueError, KeyError):
+            raise ValueError(
+                f"Invalid 'mx_zoom_mode' value: {raw_value}. Expected integer corresponding to a ZoomModeEnum."
+            )

-        data_dict = json.loads(tmp)
-        return LoopCenteringZoomModel(**data_dict)
+    @zoom_mode.setter
+    def zoom_mode(self,  mode: ZoomModeEnum):
+        self.__client.set(f"{self.__bl}:zoom_mode", mode.value)

-    @alc_zoom_settings.setter
-    def alc_zoom_settings(self, data: LoopCenteringZoomModel):
-        self.__client.set(f"{self.__bl}:alc_zoom_settings", data.model_dump_json())
+    @staticmethod
+    def zoom_setting_string(mode: ZoomModeEnum = ZoomModeEnum.User) -> str:
+        if mode == ZoomModeEnum.BeamLocation:
+            setting_string = "bl_zoom_settings"
+        elif mode == ZoomModeEnum.LoopCenter:
+            setting_string = "alc_zoom_settings"
+        elif mode == ZoomModeEnum.User:
+            setting_string = "zoom_settings"
+        else:
+            print("invalid zoom_settings mode using manual user settings")
+            setting_string = "zoom_settings"
+        return setting_string

    @property
    def zoom_settings(self) -> ZoomModel:
-        tmp = self.__client.get(f"{self.__bl}:zoom_settings")
+        mode = self.zoom_mode
+        if not mode or not isinstance(mode, ZoomModeEnum):
+            raise Exception("incorrect zoom settings mode used")
+        tmp = self.__client.get(f"{self.__bl}:{self.zoom_setting_string(mode)}")
        if tmp is None:
-            return def_zoom()
+            return zoom_manager(mode, self.__mxb)
        data_dict = json.loads(tmp)
        return ZoomModel(**data_dict)

    @zoom_settings.setter
-    def zoom_settings(self, data: ZoomModel):
-        self.__client.set(f"{self.__bl}:zoom_settings", data.model_dump_json())
-
-    @property
-    def bl_zoom_settings(self) -> ZoomModel:
-        tmp = self.__client.get(f"{self.__bl}:bl_zoom_settings")
-        if tmp is None:
-            return def_bl_zoom()
-        data_dict = json.loads(tmp)
-        return ZoomModel(**data_dict)
-
-    @zoom_settings.setter
-    def bl_zoom_settings(self, data: ZoomModel):
-        self.__client.set(f"{self.__bl}:bl_zoom_settings", data.model_dump_json())
+    def zoom_settings(self,data: ZoomModel):
+        mode = self.zoom_mode
+        if not mode or not isinstance(mode, ZoomModeEnum):
+            raise Exception("incorrect zoom settings mode used")
+        self.__client.set(f"{self.__bl}:{self.zoom_setting_string(mode)}", data.model_dump_json())

    @property
    def abr_meas_pos(self) -> Coordinate:
@@ -27,7 +27,7 @@ from aaredaqlib.models import (
    PuckLoadedInfo,
    SampleShortInfoList,
    DAQStatusModel, BeamlineStatus, SessionStatus, SampleCameraSettings, AutofocusSettings, BoundingBoxModel,
-    LoopCenteringZoomModelElem)
+    ZoomModeEnum)
 from aaredaqlib.raster_grid import RasterGridRequest, CompletedRasterGrid, CompletedRasterGridElem
 from aaredaqlib.rotation_scan import RotationScanRequest, CompletedRotationScan
 from aaredaqlib.sample_geometry import SampleGeometryModel
@@ -105,11 +105,11 @@ class AareDAQ:
        self.__devs.zoom = val
        print(f"{self.__cfg.state}")
        if self.__cfg.state == BeamlineStateEnum.BeamLocation:
-            print(f"Setting BL zoom {val} {self.__cfg.bl_zoom_settings.get_camera_settings(val)}")
-            self.samcam_settings = self.__cfg.bl_zoom_settings.get_camera_settings(val)
+            self.__cfg.zoom_mode = ZoomModeEnum.BeamLocation
        else:
-            zoom_settings = self.__cfg.zoom_settings.get_camera_settings(val)
-            print(f"Setting zoom {val} {zoom_settings}")
+            self.__cfg.zoom_mode = ZoomModeEnum.User
+        zoom_settings = self.__cfg.zoom_settings.get_camera_settings(val)
+        print(f"Setting zoom {val} {zoom_settings}")
        self.samcam_settings = zoom_settings
        #elf.__cfg.state_busy = False

@@ -778,16 +778,18 @@ class AareDAQ:
        self.__set_state(BeamlineStateEnum.SampleAlignment)
        try:
            self.__devs.lamp_light = 2.5
-
-            for s in self.__cfg.alc_zoom_settings.z:
-                self.__devs.samcam_settings = SampleCameraSettings(exposure=s.sam_cam_exp, gain=s.sam_cam_gain)
-                self.__devs.zoom_sync(s.zoom_value)
+            self.__cfg.zoom_mode = ZoomModeEnum.LoopCenter
+            zoom_settings = self.__cfg.zoom_settings.z
+            for zoom_value in zoom_settings:
+                exposure = zoom_settings[zoom_value].exposure
+                gain = zoom_settings[zoom_value].gain
+                self.__devs.samcam_settings = SampleCameraSettings(exposure=exposure, gain=gain)
+                self.__devs.zoom_sync(zoom_value)
                time.sleep(5.0) # Wait for settings to stabilize
-
                image = self.__devs.sample_cam.get_image(gray=False)
-                self.__cfg.put_alc_bkg(s.zoom_value, s.sam_cam_exp, s.sam_cam_gain,  image)
+                self.__cfg.put_alc_bkg(zoom_value, exposure, gain,  image)
                bgr_array = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-                cv2.imwrite(f"bkg{s.zoom_value:.0f}_{s.sam_cam_exp*1000:.0f}_{s.sam_cam_gain:.0f}.jpg", bgr_array)
+                cv2.imwrite(f"bkg{zoom_value:.0f}_{exposure*1000:.0f}_{gain:.0f}.jpg", bgr_array)
            self.__cfg.state_busy = False
            self.__devs.sample_cam.set_auto()
        except Exception as e:
@@ -985,12 +987,12 @@ class AareDAQ:
            self.__cfg.state_busy = False
            raise

-    def __loop_center(self, s: LoopCenteringZoomModelElem, filename: str | None = "", iteration: int = 0) -> SmargonCoordinate:
+    def __loop_center(self, s, filename: str | None = "", iteration: int = 0) -> SmargonCoordinate:
        # TODO: tidy up!
        curr_image = self.camera_image

        gray_without_feature = cv2.cvtColor(
-            self.__cfg.get_alc_bkg(self.zoom, s.sam_cam_exp, s.sam_cam_gain), cv2.COLOR_RGB2GRAY
+            self.__cfg.get_alc_bkg(self.zoom, s.exposure, s.gain), cv2.COLOR_RGB2GRAY
        )
        gray_with_feature = cv2.cvtColor(curr_image, cv2.COLOR_RGB2GRAY)

@@ -1044,7 +1046,11 @@ class AareDAQ:
        self.__devs.lamp_light = 2.5
        try:
            i = 0
-            for s in self.__cfg.alc_zoom_settings.z:
+            self.__cfg.zoom_mode = ZoomModeEnum.LoopCenter
+            zoom_settings = self.__cfg.zoom_settings.z
+            for zoom_value in zoom_settings:
+                exposure = zoom_settings[zoom_value].exposure
+                gain = zoom_settings[zoom_value].gain
                max_attempt = 2
                attempt = 0
                found = 0
@@ -1052,9 +1058,9 @@ class AareDAQ:

                while attempt < max_attempt:
                    print('new loop center settings')
-                    self.__devs.samcam_settings = SampleCameraSettings(exposure=s.sam_cam_exp, gain=s.sam_cam_gain)
-                    self.__devs.zoom_sync(s.zoom_value)
-                    print(f'zoom={s.zoom_value},gain={s.sam_cam_gain}, exp={s.sam_cam_exp}')
+                    self.__devs.samcam_settings = SampleCameraSettings(exposure=exposure, gain=gain)
+                    self.__devs.zoom_sync(zoom_value)
+                    print(f'zoom={zoom_value},gain={gain}, exp={exposure}')

                    if  sample_id is not None:
                        self.save_screenshot_db(sample_id, f"pre_alc")
@@ -1063,19 +1069,19 @@ class AareDAQ:
                    found_flag = False
                    found_angle = None
                    targets_found_this_attempt = 0
-                    exp = int(s.sam_cam_exp * 1000)
-                    gain = int(s.sam_cam_gain)
+                    exp = int(exposure * 1000)
+                    gn = int(gain)

                    for loop, angle in enumerate(angles):
                        print(f"Moving to new omega: {angle}")
                        if sample_id is not None and i == 0:
-                            self.save_screenshot_db(sample_id, f"pre_alc_{sample_id}_{angle}_{s.zoom_value:.0f}_{exp}_{gain}")
+                            self.save_screenshot_db(sample_id, f"pre_alc_{sample_id}_{angle}_{zoom_value:.0f}_{exp}_{gn}")

                        self.__devs.aerotech.move(angle, wait=True)
                        time.sleep(1)
                        filename = None
                        if sample_id is not None:
-                            filename = f"{sample_id}_{angle}_{s.zoom_value:.0f}_{exp}_{gain}"
+                            filename = f"{sample_id}_{angle}_{zoom_value:.0f}_{exp}_{gn}"

                        target = self.__ml_loop_centre_box(sample_id, filename)

@@ -1091,7 +1097,7 @@ class AareDAQ:
                                print(sample_id)
                                self.__aare.sample_centered(self.__cfg.current_sample)
                                time.sleep(0.1)
-                                self.save_screenshot_db(sample_id, f"{sample_id}_{angle}_{s.zoom_value:.0f}_{exp}_{gain}")
+                                self.save_screenshot_db(sample_id, f"{sample_id}_{angle}_{zoom_value:.0f}_{exp}_{gn}")


                    if targets_found_this_attempt == 0:
@@ -1124,10 +1130,16 @@ class AareDAQ:
                self.__aare.alc_failed(self.sample)
                raise LoopCenteringFailed
            self.__cfg.state_busy = False
+
        except Exception:
+            self.__cfg.zoom_mode = ZoomModeEnum.User
+            self.__devs.samcam_settings = self.__cfg.zoom_settings.get_camera_settings(self.zoom)
            self.__cfg.state_busy = False
            raise

+        self.__cfg.zoom_mode = ZoomModeEnum.User
+        self.__devs.samcam_settings = self.__cfg.zoom_settings.get_camera_settings(self.zoom)
+
        end = time.perf_counter()
        return end - start

@@ -1195,6 +1207,8 @@ class AareDAQ:
                grid_size_mm=Coordinate(x=geom.beam_size_mm.x * 0.5, y=geom.beam_size_mm.y * 0.5),
                omega_deg=0
            )):
+                self.__cfg.zoom_mode = ZoomModeEnum.User
+                self.__devs.samcam_settings = self.__cfg.zoom_settings.get_camera_settings(self.zoom)
                self.__rotation( RotationScanRequest(start_omega_deg=0,
                                                     dtz=det_dist,
                                                     file_prefix=sample_prefix + f"_{hex_string}",
@@ -1204,6 +1218,8 @@ class AareDAQ:
            else:
                self.__aare.axc_failed(sample)
            self.zoom = 1
+            self.__cfg.zoom_mode = ZoomModeEnum.User
+            self.__devs.samcam_settings = self.__cfg.zoom_settings.get_camera_settings(self.zoom)
            self.__cfg.state_busy = False
        except Exception:
            self.__aare.sample_failed(sample)
@@ -6,7 +6,7 @@ from aaredaq.config import ABR_POS_MOUNT, ABR_OMEGA_MOUNT
 from aaredaq.devices import BeamlineDevices
 from aaredaq.config import BeamlineConfig
 from aaredaqlib.coordinate import Coordinate
-from aaredaqlib.models import StagePositionEnum, SampleCameraSettings
+from aaredaqlib.models import StagePositionEnum, SampleCameraSettings, ZoomModeEnum
 from mxlibs3.mx_lib import pv_wait


@@ -73,6 +73,9 @@ def m2se(devs: BeamlineDevices, cfg: BeamlineConfig):
    hub = cfg.settings
    hub_cryo = cfg.cryojet_settings

+    cfg.zoom_mode = ZoomModeEnum.User
+    devs.samcam_settings = cfg.zoom_settings.get_camera_settings(devs.zoom)
+
    devs.collimator = StagePositionEnum.PARK
    devs.scintillator = StagePositionEnum.DOWN

@@ -141,6 +144,8 @@ def se2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    devs.detector_cover.put(2) # Open detector cover
    print(time.ctime(), " mchanging settings")
    hub = cfg.settings
+    cfg.zoom_mode = ZoomModeEnum.User
+    devs.samcam_settings = cfg.zoom_settings.get_camera_settings(devs.zoom)
    hub_cryo = cfg.cryojet_settings
    _cryo = hub_cryo.cryojet_in_use
    cryo_meas = hub_cryo.cryojet_measurement_position
@@ -278,7 +283,8 @@ def ba2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
 def sa2bl(devs: BeamlineDevices, cfg: BeamlineConfig):
    val=1000
    devs.zoom = val
-    devs.samcam_settings = cfg.bl_zoom_settings.get_camera_settings(val)
+    cfg.zoom_mode = ZoomModeEnum.BeamLocation
+    devs.samcam_settings = cfg.zoom_settings.get_camera_settings(val)
    devs.lamp_light = 2.5

    devs.abr_pos = cfg.abr_meas_pos + Coordinate(x=-10)
@@ -306,6 +312,11 @@ def bl2sa(devs: BeamlineDevices, cfg: BeamlineConfig):
    devs.scintillator = StagePositionEnum.DOWN
    devs.collimator = StagePositionEnum.PARK

+    val=1
+    devs.zoom = val
+    cfg.zoom_mode = ZoomModeEnum.User
+    devs.samcam_settings = cfg.zoom_settings.get_camera_settings(val)
+
    if _cryo:
        devs.cryo.move(cryo_meas)