From e7f1603a7689c288fbe3b92239feba2448bcd128 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Tue, 30 Sep 2025 15:44:07 +0200 Subject: [PATCH] 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. --- common/src/aaredaqlib/models.py | 125 +++++++++++++++++++------------- daq/src/aaredaq/config.py | 71 ++++++++++-------- daq/src/aaredaq/daq.py | 62 ++++++++++------ daq/src/aaredaq/workflows.py | 15 +++- 4 files changed, 167 insertions(+), 106 deletions(-) diff --git a/common/src/aaredaqlib/models.py b/common/src/aaredaqlib/models.py index 7c3e0677..db7b1972 100644 --- a/common/src/aaredaqlib/models.py +++ b/common/src/aaredaqlib/models.py @@ -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 diff --git a/daq/src/aaredaq/config.py b/daq/src/aaredaq/config.py index 90a56129..b33ece9b 100644 --- a/daq/src/aaredaq/config.py +++ b/daq/src/aaredaq/config.py @@ -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: diff --git a/daq/src/aaredaq/daq.py b/daq/src/aaredaq/daq.py index d10a96ad..f37ad1af 100644 --- a/daq/src/aaredaq/daq.py +++ b/daq/src/aaredaq/daq.py @@ -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) diff --git a/daq/src/aaredaq/workflows.py b/daq/src/aaredaq/workflows.py index 9d21e713..54139d04 100644 --- a/daq/src/aaredaq/workflows.py +++ b/daq/src/aaredaq/workflows.py @@ -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)