perl_d
0135129c89
CI / lint (pull_request) Failing after 33s
CI / test (3.11) (pull_request) Skipped
CI / test (3.12) (pull_request) Skipped
CI / test (3.13) (pull_request) Skipped
CI / test-with-beamline-plugins (pxi_bec) (pull_request) Skipped
CI / test-with-beamline-plugins (pxii_bec) (pull_request) Skipped
CI / test-with-beamline-plugins (pxiii_bec) (pull_request) Skipped
CI / test-with-coverage (pull_request) Skipped
460 lines
17 KiB
Python
460 lines
17 KiB
Python
import datetime
|
|
import io
|
|
import json
|
|
import os
|
|
|
|
import aareDB
|
|
import cv2
|
|
import numpy as np
|
|
import requests
|
|
from aarecommon.config.logger import setup_logger
|
|
from aarecommon.config.logger_events import log_timing
|
|
from aarecommon.math.coordinate import Coordinate
|
|
from aarecommon.math.find_xtal import compute_crystal_score_array
|
|
from aarecommon.math.sample_geometry import SampleGeometryModel
|
|
from aarecommon.models.beamline import MXBeamline
|
|
from aarecommon.models.models import (
|
|
DAQStatusModel,
|
|
PuckLoadedInfo,
|
|
SampleShortInfo,
|
|
ScanResultPayloadModel,
|
|
)
|
|
from aarecommon.models.raster_grid import CenterOfMassModel, RasterGridRequest, RasterPayloadModel
|
|
from aarecommon.models.rotation_scan import RotationScanRequest
|
|
from aareDB import (
|
|
BeamlineParametersInput,
|
|
CharacterizationParameters,
|
|
Datasets,
|
|
ExperimentParametersCreate,
|
|
GridScanParameters,
|
|
RotationParameters,
|
|
SampleEventCreate,
|
|
SampleEventType,
|
|
SetTellPosition,
|
|
SetTellPositionRequest,
|
|
)
|
|
from aareDB import Detector as DetectorParameters
|
|
from jfjoch_client.models import ScanResult
|
|
from pydantic import StrictInt
|
|
|
|
logger = setup_logger("aareDAQ")
|
|
|
|
|
|
class AareWrapper:
|
|
def __init__(self, bl: MXBeamline, host: str = "https://mx-aaredb-dmz-01.psi.ch/dispatcher"):
|
|
configuration = aareDB.Configuration(host=host)
|
|
|
|
# --- mTLS & SSL CONFIGURATION ---
|
|
# 1. Trust the Server (CA that signed mx-aaredb-dmz-01)
|
|
configuration.verify_ssl = True
|
|
configuration.ssl_ca_cert = "/etc/ssl/certs/secrets/mx-aaredb-dmz-01_Full_Chain_CA.pem"
|
|
|
|
# 2. Present Machine Identity (The certs that worked in curl)
|
|
beamline_name = bl.value.lower()
|
|
configuration.cert_file = (
|
|
f"/etc/ssl/certs/secrets/mx-{beamline_name}-queue-01_from_dmz-01.crt"
|
|
)
|
|
configuration.key_file = (
|
|
f"/etc/ssl/certs/secrets/mx-{beamline_name}-queue-01_from_dmz-01.key"
|
|
)
|
|
|
|
# 3. Initialize the Client with this config
|
|
self.client = aareDB.ApiClient(configuration)
|
|
|
|
# Identity Forwarding (Optional now that mTLS is active, but safe to keep)
|
|
self.client.default_headers["X-Shared-Password"] = os.getenv("AAREDB_SHARED_PASSWORD")
|
|
self._host = host
|
|
self._tell_api = aareDB.TellsRunnerApi(self.client)
|
|
self._sample_api = aareDB.SamplesRunnerApi(self.client)
|
|
self._proc_api = aareDB.ProcessingsRunnerApi(self.client)
|
|
self._raster_api = aareDB.GridscanRunnerApi(self.client)
|
|
self._bl = bl
|
|
|
|
self._ssl_ca_cert = configuration.ssl_ca_cert
|
|
self._cert_file = configuration.cert_file
|
|
self._key_file = configuration.key_file
|
|
|
|
@log_timing(logger, "AareDB call")
|
|
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)
|
|
logger.debug(ret)
|
|
|
|
@log_timing(logger, "AareDB call")
|
|
def create_manual_sample(self, s: SampleShortInfo):
|
|
from aareDB.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:
|
|
logger.exception("Error inserting sample")
|
|
|
|
@log_timing(logger, "AareDB call")
|
|
def send_sample_event(
|
|
self, sample_id: StrictInt, event_type: SampleEventType, comment: str | None = None
|
|
) -> None:
|
|
if sample_id is None or sample_id < 0:
|
|
if sample_id is None:
|
|
logger.debug(f"Sample db_id is None, skipping sample event {event_type!s}")
|
|
elif sample_id < 0:
|
|
logger.debug(
|
|
f"Sample db_id is invalid ({sample_id}), skipping sample event {event_type!s}"
|
|
)
|
|
return
|
|
try:
|
|
self._sample_api.create_sample_event(
|
|
sample_id=sample_id,
|
|
sample_event_create=SampleEventCreate(event_type=event_type, comment=comment),
|
|
)
|
|
except Exception:
|
|
logger.exception(f"Error sending sample event {event_type!s} to db")
|
|
|
|
@log_timing(logger, "AareDB call")
|
|
def upload_image(
|
|
self, sample_id: int, filename: str, bgr_image: np.ndarray, message: str | None = None
|
|
):
|
|
_, 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"),
|
|
}
|
|
|
|
request_kwargs = {
|
|
"files": {"uploaded_file": (filename + ".jpg", jpeg_bytes, "image/jpeg")},
|
|
"verify": self._ssl_ca_cert,
|
|
"cert": (self._cert_file, self._key_file),
|
|
"headers": headers,
|
|
}
|
|
if message is not None:
|
|
request_kwargs["data"] = {"comment": message}
|
|
response = requests.post(url, **request_kwargs)
|
|
logger.debug(f"Response status code: {response.status_code}")
|
|
|
|
@log_timing(logger, "AareDB call")
|
|
def upload_jpg(self, sample_id: int, filename: str, jpg_image, message: str | None = None):
|
|
logger.debug(f"jppg_image of type: {type(jpg_image)}")
|
|
url = f"{self._host}/protected_router/sample_runner/{sample_id}/upload-images"
|
|
headers = {
|
|
"accept": "application/json",
|
|
"X-Shared-Password": os.getenv("AAREDB_SHARED_PASSWORD"),
|
|
}
|
|
request_kwargs = {
|
|
"files": {"uploaded_file": (filename + ".jpg", jpg_image, "image/jpeg")},
|
|
"verify": self._ssl_ca_cert,
|
|
"cert": (self._cert_file, self._key_file),
|
|
"headers": headers,
|
|
}
|
|
if message is not None:
|
|
request_kwargs["data"] = {"comment": message}
|
|
response = requests.post(url, **request_kwargs)
|
|
logger.debug(f"Response status code: {response.status_code}")
|
|
|
|
@log_timing(logger, "AareDB call")
|
|
def create_rotation_run(
|
|
self, s: SampleShortInfo | None, r: RotationScanRequest, d: DAQStatusModel
|
|
):
|
|
if s is None:
|
|
return
|
|
|
|
try:
|
|
if r.screening:
|
|
characterization = CharacterizationParameters(
|
|
omegaStart_deg=round(r.start_omega_deg, 3),
|
|
omegaStep=r.incr_omega_deg,
|
|
phi=round(d.geom.smargon.phi_deg, 3),
|
|
chi=round(d.geom.smargon.chi_deg, 3),
|
|
numberOfImages=r.steps,
|
|
exposureTime_s=r.exp_time_s,
|
|
oscillation_deg=r.wedge_omega_deg,
|
|
)
|
|
rotation = None
|
|
else:
|
|
rotation = RotationParameters(
|
|
omegaStart_deg=round(r.start_omega_deg, 3),
|
|
omegaStep=r.incr_omega_deg,
|
|
phi=round(d.geom.smargon.phi_deg, 3),
|
|
chi=round(d.geom.smargon.chi_deg, 3),
|
|
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=r.dtz,
|
|
resolution_at_edge_Ang=d.diffraction.max_resolution_angstrom,
|
|
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 * 1000,
|
|
beamSizeHeight=d.geom.beam_size_mm.y * 1000,
|
|
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
|
|
)
|
|
self._sample_api.create_experiment_parameters_for_sample(
|
|
sample_id=s.db_id, experiment_parameters_create=experiment_params_payload
|
|
)
|
|
# logger.debug("Experiment parameters created:", response)
|
|
except Exception:
|
|
logger.exception("Error creating experiment parameters for rotation run")
|
|
|
|
@log_timing(logger, "AareDB call")
|
|
def create_gridscan_run(
|
|
self, s: SampleShortInfo | None, 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=round(r.omega_deg, 2),
|
|
numberOfImages=round(r.n_x * r.n_y, 0),
|
|
exposureTime_s=round(r.exp_time_s, 4),
|
|
)
|
|
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=round(r.dtz, 3),
|
|
resolution_at_edge_Ang=round(d.diffraction.max_resolution_angstrom, 3),
|
|
beamCenterX_px=round(d.diffraction.beam_center_pxl[0], 2),
|
|
beamCenterY_px=round(d.diffraction.beam_center_pxl[1], 2),
|
|
pixelSizeX_um=round(d.diffraction.pixel_size_mm * 1000, 2),
|
|
pixelSizeY_um=round(d.diffraction.pixel_size_mm * 1000, 2),
|
|
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=round(d.bl.ring_current_mA, 3),
|
|
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 * 1000,
|
|
beamSizeHeight=d.geom.beam_size_mm.y * 1000,
|
|
cryojetTemperature_K=d.bl.cryojet_K,
|
|
gridScan=gridscan,
|
|
)
|
|
experiment_params_payload = ExperimentParametersCreate(
|
|
type="standard", beamline_parameters=beamline_params, sample_id=s.db_id
|
|
)
|
|
self._sample_api.create_experiment_parameters_for_sample(
|
|
sample_id=s.db_id, experiment_parameters_create=experiment_params_payload
|
|
)
|
|
# logger.info("Experiment parameters created:", response)
|
|
except Exception as e:
|
|
logger.debug(e, exc_info=True)
|
|
|
|
@log_timing(logger, "AareDB call")
|
|
def ingest_gridscan(
|
|
self,
|
|
sample: SampleShortInfo | None,
|
|
raster_result: ScanResult,
|
|
raster_request: RasterGridRequest,
|
|
geom: SampleGeometryModel,
|
|
com: CenterOfMassModel | None,
|
|
beam_mark_pxl: tuple[float, float],
|
|
):
|
|
|
|
if sample is None:
|
|
return
|
|
|
|
payload_model = self.format_gridscan_payload(
|
|
sample, raster_result, raster_request, geom, com, beam_mark_pxl
|
|
)
|
|
if payload_model is None:
|
|
return
|
|
payload = payload_model.model_dump()
|
|
|
|
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=self._ssl_ca_cert,
|
|
cert=(self._cert_file, self._key_file),
|
|
)
|
|
response.raise_for_status()
|
|
|
|
logger.info(f"Response status code: {response.status_code}")
|
|
|
|
def format_gridscan_payload(
|
|
self,
|
|
sample: SampleShortInfo | None,
|
|
raster_result: ScanResult,
|
|
raster_request: RasterGridRequest,
|
|
geom: SampleGeometryModel,
|
|
com: CenterOfMassModel | None,
|
|
beam_mark_pxl: tuple[float, float],
|
|
) -> RasterPayloadModel | None:
|
|
|
|
try:
|
|
cell_size_pxl = Coordinate(
|
|
x=raster_request.grid_size_mm.x / geom.pixel_in_mm,
|
|
y=raster_request.grid_size_mm.y / geom.pixel_in_mm,
|
|
)
|
|
|
|
if raster_request.smargon_top_left is None:
|
|
beam_loc = geom.smargon.sh_mm
|
|
else:
|
|
beam_loc = raster_request.smargon_top_left.sh_mm
|
|
|
|
start_pxl = geom.smargon_to_picture(beam_loc)
|
|
|
|
if com:
|
|
com_grid_pxl = com.get_com_pxl(raster_request, geom)
|
|
x = start_pxl.x + com_grid_pxl.x
|
|
y = start_pxl.y + com_grid_pxl.y
|
|
center_pxl = Coordinate(x=x, y=y)
|
|
else:
|
|
center_pxl = None
|
|
|
|
try:
|
|
score_arr = compute_crystal_score_array(raster_result.images)
|
|
score = [
|
|
float(score_arr[img.nx, img.ny])
|
|
if img.nx is not None and img.ny is not None
|
|
else None
|
|
for img in raster_result.images
|
|
]
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"raster score computation failed, sending null score: {e}", exc_info=True
|
|
)
|
|
score = None
|
|
|
|
payload = RasterPayloadModel(
|
|
request=raster_request,
|
|
result=raster_result,
|
|
sample_id=sample.db_id,
|
|
attach_image=True,
|
|
centre_of_mass=com,
|
|
raster_score=score, # TODO: Check whether AareDB needs to change for taking this input
|
|
start_pxl=start_pxl,
|
|
center_pxl=center_pxl,
|
|
cell_size_pxl=cell_size_pxl,
|
|
beam_mark_pxl=beam_mark_pxl,
|
|
beam_size_mm=geom.beam_size_mm,
|
|
)
|
|
|
|
return payload
|
|
|
|
except Exception as e:
|
|
logger.error(e)
|
|
raise
|
|
|
|
@log_timing(logger, "AareDB call")
|
|
def ingest_scan(
|
|
self,
|
|
sample: SampleShortInfo | None,
|
|
result: ScanResult,
|
|
geom: SampleGeometryModel,
|
|
beam_mark_pxl: tuple[float, float],
|
|
):
|
|
|
|
if sample is None:
|
|
return
|
|
|
|
payload_model = self.format_scan_payload(sample, result, geom, beam_mark_pxl)
|
|
if payload_model is None:
|
|
return
|
|
payload = payload_model.model_dump()
|
|
|
|
url = f"{self._host}/protected_router/scan_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=self._ssl_ca_cert,
|
|
cert=(self._cert_file, self._key_file),
|
|
)
|
|
response.raise_for_status()
|
|
|
|
logger.info(f"Response status code: {response.status_code}")
|
|
|
|
def format_scan_payload(
|
|
self,
|
|
sample: SampleShortInfo | None,
|
|
result: ScanResult,
|
|
geom: SampleGeometryModel,
|
|
beam_mark_pxl: tuple[float, float],
|
|
) -> ScanResultPayloadModel | None:
|
|
try:
|
|
payload = ScanResultPayloadModel(
|
|
result=result,
|
|
sample_id=sample.db_id,
|
|
attach_image=True,
|
|
beam_mark_pxl=beam_mark_pxl,
|
|
beam_size_mm=geom.beam_size_mm,
|
|
)
|
|
|
|
return payload
|
|
|
|
except Exception as e:
|
|
logger.error(e)
|
|
raise
|