382 lines
14 KiB
Python
382 lines
14 KiB
Python
import datetime
|
|
import io
|
|
import json
|
|
import os
|
|
from typing import List, Optional
|
|
|
|
import aareDBclient
|
|
import cv2
|
|
import numpy as np
|
|
import requests
|
|
from aareDBclient import (
|
|
SetTellPosition,
|
|
SampleEventCreate,
|
|
SetTellPositionRequest,
|
|
CharacterizationParameters,
|
|
RotationParameters,
|
|
GridScanParameters,
|
|
Datasets,
|
|
Detector as DetectorParameters,
|
|
BeamlineParametersInput,
|
|
ExperimentParametersCreate)
|
|
|
|
from aaredaqlib.diffraction_geometry import DiffractionGeometry
|
|
from aaredaqlib.models import (
|
|
SampleShortInfo,
|
|
PuckLoadedInfo,
|
|
DewarAddress,
|
|
SampleShortInfoList,
|
|
DAQStatusModel, SessionStatus,
|
|
)
|
|
|
|
from aaredaqlib.beamline import MXBeamline
|
|
from aaredaqlib.raster_grid import RasterGridRequest, CompletedRasterGrid
|
|
from aaredaqlib.rotation_scan import RotationScanRequest
|
|
from aaredaqlib.sample_geometry import SampleGeometryModel
|
|
|
|
from jfjoch_client.models import ScanResult
|
|
|
|
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.__raster_api = aareDBclient.GridscanRunnerApi(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 create_manual_sample(self, s: SampleShortInfo):
|
|
from aareDBclient.models import ManualSampleCreate
|
|
|
|
manual_sample = ManualSampleCreate(
|
|
pgroup=s.user,
|
|
sample_name=s.sample_name,
|
|
data_collection_parameters=s.aaredb_params
|
|
)
|
|
|
|
try:
|
|
s.db_id = self.__sample_api.insert_sample(manual_sample).id
|
|
except Exception as e:
|
|
print(f"Error inserting sample: {e}")
|
|
|
|
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 sample_failed(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="Failed"),
|
|
)
|
|
except Exception as e:
|
|
print(e)
|
|
|
|
def axc_failed(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="AXCFailed"),
|
|
)
|
|
except Exception as e:
|
|
print(e)
|
|
|
|
def alc_failed(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="ALCFailed"),
|
|
)
|
|
except Exception as e:
|
|
print(e)
|
|
|
|
def sample_lost(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="Lost"),
|
|
)
|
|
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}/protected_router/sample_runner/{sample_id}/upload-images"
|
|
headers = {
|
|
"accept": "application/json",
|
|
"X-Shared-Password": os.getenv("AAREDB_SHARED_PASSWORD")
|
|
}
|
|
response = requests.post(url,
|
|
files={'uploaded_file': (filename + ".jpg", jpeg_bytes, "image/jpeg")},
|
|
verify=False, headers=headers)
|
|
print(f"Response status code: {response.status_code}")
|
|
|
|
def create_rotation_run(self, s: Optional[SampleShortInfo], r:RotationScanRequest, d:DAQStatusModel):
|
|
if s is None:
|
|
return
|
|
|
|
try:
|
|
if r.screening:
|
|
characterization = CharacterizationParameters(
|
|
omegaStart_deg=r.start_omega_deg,
|
|
omegaStep=r.incr_omega_deg,
|
|
phi=d.geom.smargon.phi_deg,
|
|
chi=d.geom.smargon.chi_deg,
|
|
numberOfImages=r.steps,
|
|
exposureTime_s=r.exp_time_s,
|
|
oscillation_deg=r.wedge_omega_deg
|
|
)
|
|
rotation = None
|
|
else:
|
|
rotation = RotationParameters(
|
|
omegaStart_deg=r.start_omega_deg,
|
|
omegaStep=r.incr_omega_deg,
|
|
phi=d.geom.smargon.phi_deg,
|
|
chi=d.geom.smargon.chi_deg,
|
|
numberOfImages=r.steps,
|
|
exposureTime_s=r.exp_time_s,
|
|
)
|
|
characterization = None
|
|
dataset = Datasets(
|
|
filepath=r.file_prefix,
|
|
status="written",
|
|
written_at=datetime.datetime.now()
|
|
)
|
|
detector_data = DetectorParameters(
|
|
manufacturer="DECTRIS",
|
|
model=d.diffraction.detector_description,
|
|
type="photon-counting",
|
|
serialNumber=d.diffraction.detector_serial_number,
|
|
detectorDistance_mm=d.diffraction.dtz_mm,
|
|
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,
|
|
dataset=dataset
|
|
)
|
|
beamline_params = BeamlineParametersInput(
|
|
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=d.bl.transmission,
|
|
focusingOptic="Kirkpatrick-Baez",
|
|
beamlineFluxAtSample_ph_s=d.bl.flux_ph_s,
|
|
beamSizeWidth=d.geom.beam_size_mm.x * 100,
|
|
beamSizeHeight=d.geom.beam_size_mm.y * 100,
|
|
cryojetTemperature_K=d.bl.cryojet_K,
|
|
rotation=rotation,
|
|
characterization=characterization
|
|
)
|
|
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):
|
|
if s is None:
|
|
return
|
|
|
|
try:
|
|
gridscan = GridScanParameters(
|
|
#xStart=90.0,
|
|
xStep=r.grid_size_mm.x,
|
|
#yStart=0.0,
|
|
yStep= r.grid_size_mm.y,
|
|
x_col=r.n_x,
|
|
y_row=r.n_y,
|
|
omegaStart_deg=r.omega_deg,
|
|
numberOfImages=r.n_x * r.n_y,
|
|
exposureTime_s=r.exp_time_s,
|
|
)
|
|
dataset = Datasets(
|
|
filepath=r.file_prefix,
|
|
status="written",
|
|
written_at=datetime.datetime.now()
|
|
)
|
|
detector_data = DetectorParameters(
|
|
manufacturer="DECTRIS",
|
|
model=d.diffraction.detector_description,
|
|
type="photon-counting",
|
|
serialNumber=d.diffraction.detector_serial_number,
|
|
detectorDistance_mm=d.diffraction.dtz_mm,
|
|
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,
|
|
dataset=dataset
|
|
)
|
|
beamline_params = BeamlineParametersInput(
|
|
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=d.bl.transmission,
|
|
focusingOptic="Kirkpatrick-Baez",
|
|
beamlineFluxAtSample_ph_s=d.bl.flux_ph_s,
|
|
beamSizeWidth=d.geom.beam_size_mm.x * 100,
|
|
beamSizeHeight=d.geom.beam_size_mm.y * 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 _create_default_geometry(self):
|
|
"""Create default geometry configuration for raster grid ingestion."""
|
|
return {
|
|
'beam_location_pxl': {'x': 512, 'y': 384, 'z': 0},
|
|
'pixel_in_mm': 0.01,
|
|
'aerotech': {'x': 0, 'y': 0, 'z': 0},
|
|
'aerotech_meas': {'x': 0, 'y': 0, 'z': 0},
|
|
'smargon': {
|
|
'sh_mm': {'x': 0, 'y': 0, 'z': 0},
|
|
'phi_deg': 0,
|
|
'chi_deg': 0
|
|
},
|
|
'omega_deg': 0,
|
|
'beam_size_mm': {'x': 0.02, 'y': 0.02, 'z': 0},
|
|
}
|
|
|
|
def ingest_gridscan(self, s: Optional[SampleShortInfo], raster_result: ScanResult, r: RasterGridRequest):
|
|
|
|
if s is None:
|
|
return
|
|
|
|
payload = self.format_gridscan_payload(s, raster_result, r)
|
|
|
|
if payload is None:
|
|
return
|
|
|
|
url = f"{self.__host}/protected_router/gridscan_runner/ingest"
|
|
headers = {
|
|
"accept": "application/json",
|
|
"X-Shared-Password": os.getenv("AAREDB_SHARED_PASSWORD")
|
|
}
|
|
response = requests.post(url,
|
|
auth=(os.getenv("AAREDB_USERNAME"), os.getenv("AAREDB_PASSWORD")),
|
|
headers=headers, data=json.dumps(payload), timeout=30, verify=False)
|
|
response.raise_for_status()
|
|
|
|
print(f"Response status code: {response.status_code}")
|
|
|
|
|
|
def format_gridscan_payload(self, s: Optional[SampleShortInfo], raster_result:ScanResult, r:RasterGridRequest):
|
|
|
|
try:
|
|
|
|
geometry = self._create_default_geometry()
|
|
|
|
payload = {
|
|
"request": r.model_dump(), # RasterGridRequest as dict
|
|
"result": raster_result.model_dump(), # ScanResult as dict
|
|
"geometry": geometry, # SampleGeometryModel as dict
|
|
"sample_id": s.db_id,
|
|
"attach_image": True,
|
|
}
|
|
|
|
return payload
|
|
|
|
except Exception as e:
|
|
print(e)
|
|
raise e
|
|
|
|
|