Files
AareDAQ/daq/src/aaredaq/aaredb.py
T
2025-06-18 09:02:58 +02:00

284 lines
10 KiB
Python

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(
#
# )