import datetime import io import os from typing import List, Optional import aareDBclient import cv2 import numpy as np import requests from aareDBclient import ( SetTellPosition, SampleEventCreate, SetTellPositionRequest, RotationParameters, GridScanParameters, Detector as DetectorParameters, BeamlineParameters, ExperimentParametersCreate) from aaredaqlib.diffraction_geometry import DiffractionGeometry from aaredaqlib.models import ( SampleShortInfo, PuckLoadedInfo, DewarAddress, SampleShortInfoList, DAQStatusModel, ) from aaredaqlib.beamline import MXBeamline from aaredaqlib.raster_grid import RasterGridRequest, CompletedRasterGrid from aaredaqlib.rotation_scan import RotationScanRequest from aaredaqlib.sample_geometry import SampleGeometryModel class AareWrapper: def __init__( self, bl: MXBeamline, host: str = "https://mx-db-01.psi.ch/dispatcher", ): configuration = aareDBclient.Configuration(host=host) configuration.verify_ssl = False # Disable SSL verification self.client = aareDBclient.ApiClient(configuration) self.client.default_headers["X-Shared-Password"] = os.getenv("AAREDB_SHARED_PASSWORD") self.__host = host self.__tell_api = aareDBclient.TellsRunnerApi(self.client) self.__sample_api = aareDBclient.SamplesRunnerApi(self.client) self.__proc_api = aareDBclient.ProcessingsRunnerApi(self.client) self.__bl = bl def set_pucks_beamline(self, input_list: List[PuckLoadedInfo]): o = [] for i in input_list: t = SetTellPosition( puck_name=i.puck_name, segment=i.location.segment, puck_in_segment=i.location.pos, ) o.append(t) payload = SetTellPositionRequest(pucks = o, tell=self.__bl.value.upper()) ret = self.__tell_api.set_tell_positions( set_tell_position_request=payload, ) print(ret) def get_sample_info(self) -> SampleShortInfoList: sample_list = ( self.__tell_api.get_pucks_with_tell_position(tell=self.__bl.value.upper()) ) ret = [] for p in sample_list: for s in p.samples: dewar_address = None if p.tell_position is not None and len(p.tell_position) == 2: dewar_address = DewarAddress( segment=p.tell_position[0], pos=p.tell_position[1] ) ret.append( SampleShortInfo( db_id=s.id, puck_name=p.puck_name, dewar_name="Unknown" if p.dewar_name is None else p.dewar_name, sample_name=s.sample_name, user=p.pgroup, pin=s.position, location=dewar_address, priority=1.0 if s.priority is None else s.priority, comment=s.comments, mount_count=s.mount_count or 0 ) ) return SampleShortInfoList(s=ret) def sample_mounted(self, s: Optional[SampleShortInfo]): if s is not None: try: self.__sample_api.create_sample_event( sample_id=s.db_id, sample_event_create=SampleEventCreate(event_type="Mounted"), ) except Exception as e: print(e) def sample_unmounted(self, s: Optional[SampleShortInfo]): if s is not None: try: self.__sample_api.create_sample_event( sample_id=s.db_id, sample_event_create=SampleEventCreate(event_type="Unmounted"), ) except Exception as e: print(e) def sample_centered(self, s: Optional[SampleShortInfo]): if s is not None: try: self.__sample_api.create_sample_event( sample_id=s.db_id, sample_event_create=SampleEventCreate(event_type="Centered"), ) except Exception as e: print(e) def sample_collected(self, s: Optional[SampleShortInfo]): if s is None: return try: self.__sample_api.create_sample_event( sample_id=s.db_id, sample_event_create=SampleEventCreate(event_type="Collected"), ) except Exception as e: print(e) def upload_image(self, sample_id: int, filename: str, bgr_image: np.ndarray): _, buffer = cv2.imencode('.jpg', bgr_image) jpeg_bytes = io.BytesIO(buffer) url = f"{self.__host}/samples/{sample_id}/upload-images" headers = { "accept": "application/json" } response = requests.post(url, files={'uploaded_file': (filename + ".jpg", jpeg_bytes, "image/jpeg")}, verify=False, headers=headers) print(response.status_code) def create_rotation_run(self, s: Optional[SampleShortInfo], r:RotationScanRequest, d:DAQStatusModel, sg:SampleGeometryModel): if s is None: return try: rotation = RotationParameters( omegaStart_deg=r.start_omega_deg, omegaStep=r.incr_omega_deg, phi=r.start.phi_deg, chi=r.start.chi_deg, numberOfImages=r.steps, exposureTime_s=r.exp_time_s, ) detector_data = DetectorParameters( manufacturer="DECTRIS", model=d.diffraction.detector_description, type="photon-counting", serialNumber=d.detector_serial_number, detectorDistance_mm=r.dtz, beamCenterX_px=d.diffraction.beam_center_pxl[0], beamCenterY_px=d.diffraction.beam_center_pxl[1], pixelSizeX_um=d.diffraction.pixel_size_mm * 1000, pixelSizeY_um=d.diffraction.pixel_size_mm * 1000, ) beamline_params = BeamlineParameters( synchrotron="Swiss Light Source", beamline=d.bl.name, detector=detector_data, energy_keV=d.diffraction.energy_keV, wavelength_Ang=d.diffraction.wavelength_angstrom, ringCurrent_mA=d.bl.ring_current_mA, ringMode="Beamline Development", monochromator="Si111", transmission=r.transmission, focusingOptic="Kirkpatrick-Baez", beamlineFluxAtSample_ph_s=d.flux_ph_s, beamSizeWidth=sg.beam_size_mm * 100, beamSizeHeight=sg.beam_size_mm * 100, cryojetTemperature_K=d.bl.cryojet_K, rotation=rotation ) experiment_params_payload = ExperimentParametersCreate( type="standard", beamline_parameters=beamline_params, sample_id=s.db_id ) response = self.__sample_api.create_experiment_parameters_for_sample( sample_id=s.db_id, experiment_parameters_create=experiment_params_payload ) print("Experiment parameters created:", response) except Exception as e: print(e) def create_gridscan_run(self, s: Optional[SampleShortInfo], r:RasterGridRequest, d:DAQStatusModel, sg:SampleGeometryModel): if s is None: return try: gridscan = GridScanParameters( #xStart=90.0, xStep=0.1, #yStart=0.0, yStep= 0.1, x_col=r.n_x, y_row=r.n_y, omegaStart_deg=r.omega_deg, numberOfImages=r.get_image_number, exposureTime_s=r.exp_time_s, ) detector_data = DetectorParameters( manufacturer="DECTRIS", model=d.diffraction.detector_description, type="photon-counting", serialNumber=d.detector_serial_number, detectorDistance_mm=r.dtz, beamCenterX_px=d.diffraction.beam_center_pxl[0], beamCenterY_px=d.diffraction.beam_center_pxl[1], pixelSizeX_um=d.diffraction.pixel_size_mm * 1000, pixelSizeY_um=d.diffraction.pixel_size_mm * 1000, ) beamline_params = BeamlineParameters( synchrotron="Swiss Light Source", beamline=d.bl.name, detector=detector_data, energy_keV=d.diffraction.energy_keV, wavelength_Ang=d.diffraction.wavelength_angstrom, ringCurrent_mA=d.bl.ring_current_mA, ringMode="Beamline Development", monochromator="Si111", transmission=r.transmission, focusingOptic="Kirkpatrick-Baez", beamlineFluxAtSample_ph_s=d.flux_ph_s, beamSizeWidth=sg.beam_size_mm * 100, beamSizeHeight=sg.beam_size_mm * 100, cryojetTemperature_K=d.bl.cryojet_K, gridScan=gridscan ) experiment_params_payload = ExperimentParametersCreate( type="standard", beamline_parameters=beamline_params, sample_id=s.db_id ) response = self.__sample_api.create_experiment_parameters_for_sample( sample_id=s.db_id, experiment_parameters_create=experiment_params_payload ) print("Experiment parameters created:", response) except Exception as e: print(e) def confirm_image_written(self, s: Optional[SampleShortInfo], r:RotationScanRequest): if s is None: return try: response = self.__sample_api.update_dataset_for_experiment_run( filepath=r.file_prefix, #need to check if this is correct status="written", written_at=datetime.datetime.now().isoformat(), ) print("Experiment run updated:", response) except Exception as e: print(e) #def send_gridscan_result(self, s: Optional[SampleShortInfo], c: CompletedRasterGrid): # if s is None: # return # # try: # response = self.__proc_api.create_result( # # )