Files
AareDAQ/gui/src/aaregui/threads/daq_worker.py
T

364 lines
13 KiB
Python

import copy
import random
from PySide6.QtCore import Signal, QUrl, Slot, QTimer, QObject, QByteArray
from PySide6.QtNetwork import QNetworkAccessManager, QNetworkRequest, QNetworkReply
from jfjoch_client import ScanResult, ScanResultImagesInner
from aaredaqlib.coordinate import SmargonCoordinate, Coordinate
from aaredaqlib.models import DAQStatusModel, SampleShortInfoList, SampleShortInfo, SampleCameraSettings, \
AutofocusSettings
from aaredaqlib.raster_grid import RasterGridRequest, CompletedRasterGrid
from aaredaqlib.rotation_scan import RotationScanRequest, CompletedRotationScan
SPREADHSEET_FREQUENCY = 25 # Every 5 seconds
class DAQWorker(QObject):
update = Signal(DAQStatusModel)
spreadsheet = Signal(SampleShortInfoList)
http_error = Signal(str)
auth_error = Signal()
automated_scan_done = Signal(int, bool) # sample ID, success
run_number_incremented = Signal()
raster_scan_completed = Signal(CompletedRasterGrid)
standard_scan_completed = Signal(CompletedRotationScan)
raster_generated_by_ml = Signal(RasterGridRequest)
def __init__(self, base_url: str | None, token: str, parent=None):
super().__init__(parent)
self.__token = token
self.__base_url = base_url
self.__net_manager = QNetworkAccessManager()
self.__timer = QTimer()
self.__timer.setInterval(200)
self.__timer.timeout.connect(self.regular_update)
self.__timer.start()
self.__counter = 0
@Slot()
def regular_update(self):
if self.__counter % SPREADHSEET_FREQUENCY == 0:
self.load_spreadsheet()
self.__counter = (self.__counter + 1) % SPREADHSEET_FREQUENCY
self.send_status_request()
def send_status_request(self):
if self.__base_url is None:
return
request = QNetworkRequest(QUrl(f"{self.__base_url}/status"))
request.setRawHeader(b"Authorization", f"Bearer {self.__token}".encode("utf-8"))
reply = self.__net_manager.get(request)
reply.finished.connect(lambda: self.handle_status_response(reply))
@staticmethod
def handle_response(reply: QNetworkReply):
if reply.error() == QNetworkReply.NetworkError.NoError:
response_data = reply.readAll().data().decode("utf-8")
reply.deleteLater()
return response_data
else:
reply.deleteLater()
raise RuntimeError(reply.errorString())
@Slot(QNetworkReply)
def handle_status_response(self, reply: QNetworkReply):
try:
response_data = self.handle_response(reply)
parsed_response = DAQStatusModel.model_validate_json(response_data)
self.update.emit(parsed_response)
except Exception as e:
print(e)
self.http_error.emit(str(e))
@Slot(QNetworkReply)
def handle_spreadsheet_response(self, reply: QNetworkReply):
try:
response_data = self.handle_response(reply)
parsed_response = SampleShortInfoList.model_validate_json(response_data)
self.spreadsheet.emit(parsed_response)
except Exception as e:
print(e)
self.http_error.emit(str(e))
def handle_req_response(self, reply: QNetworkReply):
if reply.error() != QNetworkReply.NetworkError.NoError:
if reply.attribute(QNetworkRequest.Attribute.HttpStatusCodeAttribute) == 401:
print(reply.errorString())
self.auth_error.emit()
else:
print(reply.errorString())
self.http_error.emit(reply.errorString())
reply.deleteLater()
def generic_post(self, url: str, body: str = ""):
if self.__base_url is None:
print(f"POST /{url}: {body}")
return
request = QNetworkRequest(QUrl(f"{self.__base_url}/{url}"))
request.setRawHeader(b"Authorization", f"Bearer {self.__token}".encode("utf-8"))
if str:
request.setRawHeader(b"Content-Type", b"application/json")
reply = self.__net_manager.post(request, QByteArray(body.encode("utf-8")))
reply.finished.connect(lambda: self.handle_req_response(reply))
def generic_put(self, url: str, body: str = ""):
if self.__base_url is None:
print(f"PUT /{url}: {body}")
return
request = QNetworkRequest(QUrl(f"{self.__base_url}/{url}"))
request.setRawHeader(b"Authorization", f"Bearer {self.__token}".encode("utf-8"))
if str:
request.setRawHeader(b"Content-Type", b"application/json")
reply = self.__net_manager.put(request, QByteArray(body.encode("utf-8")))
reply.finished.connect(lambda: self.handle_req_response(reply))
def generic_delete(self, url: str):
if self.__base_url is None:
print(f"DELETE /{url}")
return
request = QNetworkRequest(QUrl(f"{self.__base_url}/{url}"))
request.setRawHeader(b"Authorization", f"Bearer {self.__token}".encode("utf-8"))
reply = self.__net_manager.deleteResource(request)
reply.finished.connect(lambda: self.handle_req_response(reply))
@Slot(float)
def set_omega(self, f: float):
self.generic_put(f"beamline/omega?val={f:.3f}")
@Slot(float)
def zoom(self, f: float):
self.generic_put(f"beamline/zoom?val={f:.3f}")
@Slot(int)
def light(self, v: int):
self.generic_put(f"beamline/light?val={v:d}")
@Slot()
def close_shutter(self):
self.generic_post(f"beamline/shutter?val=false")
@Slot()
def open_shutter(self):
self.generic_post(f"beamline/shutter?val=true")
@Slot()
def alc_background(self):
self.generic_post("alc/background")
@Slot()
def center_loop(self):
self.generic_post("alc/center_loop")
@Slot()
def force_session(self):
self.generic_post("access/force_current_session")
@Slot()
def end_session(self):
self.generic_post("access/end_session")
@Slot()
def dewar_exchange(self):
self.generic_post("state/dewar_exchange")
@Slot()
def sample_exchange(self):
self.generic_post("state/sample_exchange")
@Slot()
def sample_alignment(self):
self.generic_post("state/sample_alignment")
@Slot()
def beam_location(self):
self.generic_post("state/beam_location")
@Slot(str)
def set_pgroup(self, val: str):
if val == "":
self.generic_delete("access/pgroup")
else:
self.generic_put(f"access/pgroup?val={val}")
@Slot(SampleCameraSettings)
def samcam_settings(self, s: SampleCameraSettings):
self.generic_put("beamline/samcam", s.model_dump_json())
@Slot(AutofocusSettings)
def autofocus(self, f: AutofocusSettings):
self.generic_post("samcam/autofocus", f.model_dump_json())
@Slot(SmargonCoordinate)
def move_smargon(self, coord: SmargonCoordinate):
self.generic_put("beamline/smargon", coord.model_dump_json())
def handle_rotation_scan_response(self, reply: QNetworkReply):
try:
response_data = self.handle_response(reply)
parsed_response = CompletedRotationScan.model_validate_json(response_data)
self.standard_scan_completed.emit(parsed_response)
except Exception as e:
print(e)
self.http_error.emit(str(e))
@Slot(RotationScanRequest)
def standard_scan(self, r: RotationScanRequest):
self.run_number_incremented.emit()
request = QNetworkRequest(QUrl(f"{self.__base_url}/scan/rotation"))
request.setRawHeader(b"Authorization", f"Bearer {self.__token}".encode("utf-8"))
request.setRawHeader(b"Content-Type", b"application/json")
body = r.model_dump_json()
reply = self.__net_manager.post(request, QByteArray(body.encode("utf-8")))
reply.finished.connect(lambda: self.handle_rotation_scan_response(reply))
def handle_raster_scan_response(self, reply: QNetworkReply):
try:
response_data = self.handle_response(reply)
parsed_response = CompletedRasterGrid.model_validate_json(response_data)
self.raster_scan_completed.emit(parsed_response)
except Exception as e:
print(e)
self.http_error.emit(str(e))
@Slot(RasterGridRequest)
def raster_scan(self, r: RasterGridRequest):
self.run_number_incremented.emit()
if self.__base_url is None:
print(f"POST /scan/raster: {r.model_dump_json()}")
image_number = r.get_image_number()
new_copy = copy.deepcopy(r)
images = []
for i in range(image_number):
images.append(ScanResultImagesInner(
number=i,
efficiency=1.0,
bkg = random.gauss(3.0, 0.1),
spots= random.randint(0, 250),
index= random.randint(0, 1),
mos = random.uniform(0, 0.1),
b= random.uniform(15.0, 80.0)
))
reply = CompletedRasterGrid(request = new_copy,
result = ScanResult(file_prefix=r.file_prefix, images=images))
self.raster_scan_completed.emit(reply)
return
request = QNetworkRequest(QUrl(f"{self.__base_url}/scan/raster"))
request.setRawHeader(b"Authorization", f"Bearer {self.__token}".encode("utf-8"))
request.setRawHeader(b"Content-Type", b"application/json")
body = r.model_dump_json()
reply = self.__net_manager.post(request, QByteArray(body.encode("utf-8")))
reply.finished.connect(lambda: self.handle_raster_scan_response(reply))
@Slot()
def load_spreadsheet(self):
if self.__base_url is None:
print(f"GET /sample/spreadsheet")
return
request = QNetworkRequest(QUrl(f"{self.__base_url}/sample/spreadsheet"))
request.setRawHeader(b"Authorization", f"Bearer {self.__token}".encode("utf-8"))
reply = self.__net_manager.get(request)
reply.finished.connect(lambda: self.handle_spreadsheet_response(reply))
def handle_auto_scan_response(self, reply, sample_id: int):
if reply.error() == QNetworkReply.NetworkError.NoError:
resp = reply.readAll().data().decode("utf-8")
print(f"Sample time {resp} s")
self.automated_scan_done.emit(sample_id, True)
else:
if reply.attribute(QNetworkRequest.Attribute.HttpStatusCodeAttribute) == 401:
print(reply.errorString())
self.auth_error.emit()
else:
print(reply.errorString())
self.http_error.emit(reply.errorString())
self.automated_scan_done.emit(sample_id, False)
reply.deleteLater()
@Slot(SampleShortInfo)
def automated_scan(self, s: SampleShortInfo):
if self.__base_url is None:
print(f"POST /scan/auto: {s.model_dump_json()}")
return
request = QNetworkRequest(QUrl(f"{self.__base_url}/scan/auto"))
request.setRawHeader(b"Authorization", f"Bearer {self.__token}".encode("utf-8"))
request.setRawHeader(b"Content-Type", b"application/json")
body = s.model_dump_json()
reply = self.__net_manager.post(request, QByteArray(body.encode("utf-8")))
reply.finished.connect(lambda: self.handle_auto_scan_response(reply, s.db_id))
@Slot(Coordinate)
def abr_tweak(self, c: Coordinate):
self.generic_post("beamline/tweak_abr_meas_pos", c.model_dump_json())
@Slot()
def abr_save(self):
self.generic_post("beamline/save_abr_meas_pos")
@Slot()
def abr_goto_meas(self):
self.generic_post("beamline/goto_abr_meas_pos")
@Slot(float, float)
def beam_mark_add(self, x: float, y: float):
self.generic_post(f"beam_mark/add?x={x}&y={y}")
@Slot()
def beam_mark_clear(self):
self.generic_post(f"beam_mark/clear")
@Slot(float, float)
def beam_center(self, x: float, y: float):
self.generic_post(f"beamline/beam_center?x={x}&y={y}")
@Slot(float, float)
def beam_size_mm(self, x: float, y: float):
self.generic_post(f"beamline/beam_size_mm?x={x}&y={y}")
@Slot()
def unmount(self):
self.generic_post("sample/unmount")
@Slot(SampleShortInfo)
def mount(self, s: SampleShortInfo):
self.generic_post(f"sample/mount?dbid={s.db_id}")
@Slot(SampleShortInfo)
def sample_manual(self, s: SampleShortInfo):
self.generic_post(f"sample/manual", s.model_dump_json())
@Slot()
def cancel(self):
self.generic_post("scan/cancel")
def handle_ml_box_response(self, reply):
try:
response_data = self.handle_response(reply)
if response_data != "":
parsed_response = RasterGridRequest.model_validate_json(response_data)
self.raster_generated_by_ml.emit(parsed_response)
except Exception as e:
print(e)
self.http_error.emit(str(e))
@Slot()
def ml_bounding_box(self):
if self.__base_url is None:
print(f"POST /alc/ml_bounding_box")
return
request = QNetworkRequest(QUrl(f"{self.__base_url}/alc/ml_bounding_box"))
request.setRawHeader(b"Authorization", f"Bearer {self.__token}".encode("utf-8"))
reply = self.__net_manager.post(request, QByteArray(b""))
reply.finished.connect(lambda: self.handle_ml_box_response(reply))