DAQ: Include JFJoch, first try

2025-06-16 12:22:54 +02:00
parent 63fc40af7f
commit 5fdd3a3327
4 changed files with 211 additions and 96 deletions
@@ -1,3 +1,4 @@
+import copy
 import time
 from math import ceil
 from typing import List, Tuple
@@ -6,7 +7,13 @@ import cv2
 import numpy as np
 import redis

+from aaredaq import workflows
 from aaredaq.autofocus import calculate_focus_measure
+from aaredaq.config import BeamlineConfig, ABR_POS_MOUNT
+from aaredaq.config import BeamlineStateEnum
+from aaredaq.devices import BeamlineDevices
+from aaredaq.mlbox import MlBox
+from aaredaqlib.beamline import MXBeamline
 from aaredaqlib.coordinate import Coordinate, SmargonCoordinate
 from aaredaqlib.diffraction_geometry import DiffractionGeometry
 from aaredaqlib.models import (
@@ -17,13 +24,7 @@ from aaredaqlib.models import (
 from aaredaqlib.raster_grid import RasterGridRequest, CompletedRasterGrid
 from aaredaqlib.rotation_scan import RotationScanRequest
 from aaredaqlib.sample_geometry import SampleGeometryModel
-
-from aaredaq import workflows
-from aaredaqlib.beamline import MXBeamline
-from aaredaq.config import BeamlineConfig, ABR_POS_MOUNT
-from aaredaq.config import BeamlineStateEnum
-from aaredaq.devices import BeamlineDevices
-from aaredaq.mlbox import MlBox
+from mxlibs3.jfjoch import JFJochWrapper


 class TransformationInvalidException(Exception):
@@ -45,6 +46,7 @@ class AareDAQ:
        self.__cfg = cfg
        self.__devs = BeamlineDevices(bl)
        self.__mlbox = MlBox()
+        self.__jfjoch = JFJochWrapper(bl)

    @property
    def state(self) -> BeamlineStateEnum:
@@ -239,9 +241,10 @@ class AareDAQ:
    def list_loaded_pucks(self) -> List[PuckLoadedInfo]:
        return self.__devs.tell.get_detected_pucks()

-    def __raster(self, r: RasterGridRequest):
+    def __raster(self, r: RasterGridRequest) -> CompletedRasterGrid:
        max_time = r.exp_time_s * r.n_y * r.n_x + 60
        self.__set_state(BeamlineStateEnum.DataCollection)
+        self.__jfjoch.measure_raster(r, self.diffraction_geometry, self.sample)
        self.__devs.aerotech.move(r.omega_deg, wait=True, speed=360.0)
        save_smargon_position = self.__devs.smargon.readback
        self.__devs.smargon.target = r.smargon
@@ -251,15 +254,17 @@ class AareDAQ:
        self.__devs.aerotech.wait_scan_done(max_time)
        self.__devs.aerotech.reset()
        self.__devs.smargon.target = save_smargon_position
+        result = self.__jfjoch.wait_till_done(5)
+        return CompletedRasterGrid(request=copy.deepcopy(r), result=result)

    def measure_raster(self, r: RasterGridRequest) -> CompletedRasterGrid:
        max_time = r.exp_time_s * r.n_y * r.n_x + 60
        self.__cfg.try_set_busy(timeout=ceil(max_time + 360))
        try:
-            self.__raster(r)
+            result = self.__raster(r)
            self.__set_state(BeamlineStateEnum.SampleAlignment)
            self.__cfg.state_busy = False
-            return CompletedRasterGrid(request=r)
+            return result
        except Exception as e:
            self.__cfg.state_busy = False
            raise e
@@ -271,6 +276,8 @@ class AareDAQ:
            self.__devs.smargon.target = request.start
            self.__devs.smargon.wait()

+        self.__jfjoch.measure_rotation(request, self.diffraction_geometry, self.sample)
+
        if request.screening:
            self.__devs.aerotech.measure_screening(request.start_omega_deg,
                                               request.wedge_omega_deg,
@@ -278,26 +285,27 @@ class AareDAQ:
                                               request.incr_omega_deg,
                                               request.steps)
            self.__devs.aerotech.wait_scan_done(request.exp_time_s * request.steps + 60)
-            return
+        else:

-        total_omega = request.incr_omega_deg * request.steps
-        total_time = request.exp_time_s * request.steps
+            total_omega = request.incr_omega_deg * request.steps
+            total_time = request.exp_time_s * request.steps

-        self.__devs.aerotech.measure_standard(
-            request.start_omega_deg, total_omega, total_time
-        )
+            self.__devs.aerotech.measure_standard(
+                request.start_omega_deg, total_omega, total_time
+            )

-        if request.start is not None and request.end is not None:
-            smargon_time_step = request.time_sec / float(request.steps)
-            pos_step = (request.end.sh_mm - request.start.sh_mm) * (1.0 / float(request.steps))
+            if request.start is not None and request.end is not None:
+                smargon_time_step = request.time_sec / float(request.steps)
+                pos_step = (request.end.sh_mm - request.start.sh_mm) * (1.0 / float(request.steps))

-            for i in range(request.steps):
-                self.__devs.smargon.target = SmargonCoordinate(
-                    sh_mm=request.start.sh_mm + pos_step * i
-                )
-                time.sleep(smargon_time_step)
+                for i in range(request.steps):
+                    self.__devs.smargon.target = SmargonCoordinate(
+                        sh_mm=request.start.sh_mm + pos_step * i
+                    )
+                    time.sleep(smargon_time_step)

-        self.__devs.aerotech.wait_scan_done(request.total_time_sec + 60)
+            self.__devs.aerotech.wait_scan_done(request.total_time_sec + 60)
+        self.__jfjoch.wait_till_done(5)
        self.__devs.aerotech.reset()

    def measure_rotation(self, request: RotationScanRequest):
@@ -722,10 +730,11 @@ class AareDAQ:

    @property
    def diffraction_geometry(self) -> DiffractionGeometry:
+        det_cfg = self.__jfjoch.detector()
        return DiffractionGeometry(
            energy_keV=self.__devs.energy_kev,
            dtz_mm=self.__devs.dtz.value,
-            detector_size_pxl=(1553,1630),
+            detector_size_pxl=(det_cfg.width, det_cfg.height),
            pixel_size_mm=0.150, #PILATUS 4
            beam_center_pxl=self.__cfg.beam_center
        )
@@ -1,13 +0,0 @@
-from aaredaqlib.beamline import MXBeamline
-from jfjoch_client.api.default_api import DefaultApi
-from jfjoch_client.configuration import Configuration
-
-class JfjochWrapper:
-    def __init__(
-        self,
-        bl: MXBeamline,
-        host: str = "https://sls-gpu-001:8080",
-    ):
-        config = Configuration(host=host)
-        self.__client = DefaultApi(config)
-
@@ -0,0 +1,109 @@
+import math
+
+import jfjoch_client
+
+from aaredaqlib.beamline import MXBeamline
+from aaredaqlib.diffraction_geometry import DiffractionGeometry
+from aaredaqlib.models import SampleShortInfo
+from aaredaqlib.raster_grid import RasterGridRequest
+from aaredaqlib.rotation_scan import RotationScanRequest
+
+
+class JFJochWrapper:
+    def __init__(self, bl: MXBeamline):
+        match bl:
+            case MXBeamline.X06DA:
+                self.__url = "http://sls-gpu-001:8080"
+            case MXBeamline.SIMULATED:
+                self.__url = "http://localhost:8080"
+            case _:
+                raise Exception("unknown beamline")
+
+        self.__client = jfjoch_client.ApiClient(jfjoch_client.Configuration(host=self.__url))
+        self.__api = jfjoch_client.DefaultApi(self.__client)
+
+    def initialize(self):
+        self.__api.initialize_post()
+
+    def is_idle(self) -> bool:
+        status = self.__api.status_get()
+        return status.state == 'Idle'
+
+    def measure_rotation(self,
+                         r: RotationScanRequest,
+                         d: DiffractionGeometry,
+                         s: SampleShortInfo | None) -> None:
+        if s is None:
+            pgroup = "p16371"
+            sample = None
+        else:
+            pgroup = s.user
+            sample = s.sample_name
+
+        goniometer_settings = jfjoch_client.RotationAxis(
+            step=r.incr_omega_deg,
+            start=r.start_omega_deg,
+            name="omega",
+            vector= [0,1,0]
+        )
+
+        dataset_settings = jfjoch_client.DatasetSettings(
+            beam_x_pxl=d.beam_center_pxl[0],
+            beam_y_pxl=d.beam_center_pxl[1],
+            ntrigger= 1,
+            images_per_trigger=r.steps,
+            detector_distance_mm=d.dtz_mm,
+            file_prefix=f"{pgroup}/raw/{r.file_prefix}",
+            incident_energy_keV=d.energy_keV,
+            goniometer=goniometer_settings,
+            sample_name=sample,
+            image_time_us=round(r.exp_time_s * 1e6),
+            experiment_group=pgroup
+        )
+        self.__api.start_post(dataset_settings=dataset_settings)
+
+    def measure_raster(self,
+                       r: RasterGridRequest,
+                       d: DiffractionGeometry,
+                       s: SampleShortInfo | None):
+        if s is None:
+            pgroup = "p16371"
+            sample = None
+        else:
+            pgroup = s.user
+            sample = s.sample_name
+
+        grid_settings = jfjoch_client.GridScan(
+            n_fast= r.n_x,
+            snake=True,
+            step_x_um=r.grid_size_mm.x * 1000.0,
+            step_y_um=r.grid_size_mm.y * 1000.0,
+            vertical=False
+        )
+
+        dataset_settings = jfjoch_client.DatasetSettings(
+            beam_x_pxl=d.beam_center_pxl[0],
+            beam_y_pxl=d.beam_center_pxl[1],
+            ntrigger= r.n_y,
+            images_per_trigger=r.n_x,
+            detector_distance_mm=d.dtz_mm,
+            file_prefix=f"{pgroup}/raw/{r.file_prefix}",
+            incident_energy_keV=d.energy_keV,
+            grid_scan=grid_settings,
+            sample_name=sample,
+            image_time_us=round(r.exp_time_s * 1e6),
+            experiment_group=pgroup
+        )
+        self.__api.start_post(dataset_settings=dataset_settings)
+
+    def wait_till_done(self, timeout : int | float) -> jfjoch_client.models.ScanResult:
+        self.__api.wait_till_done_post_with_http_info(timeout=math.ceil(timeout))
+        return self.__api.result_scan_get()
+
+    def detector(self) -> jfjoch_client.models.DetectorListElement:
+        l = self.__api.config_select_detector_get()
+
+        if len(l.detectors) == 0:
+            raise Exception("no detectors configured")
+
+        return l.detectors[l.current_id]
@@ -6,87 +6,97 @@
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
-     "end_time": "2025-05-19T14:02:45.131974Z",
-     "start_time": "2025-05-19T14:02:44.020840Z"
+     "end_time": "2025-06-16T10:22:09.086440Z",
+     "start_time": "2025-06-16T10:22:08.146343Z"
    }
   },
   "source": [
-    "from aaredaq.devices import BeamlineDevices\n",
-    "from aaredaq.beamline import MXBeamline"
+    "from aaredaqlib.coordinate import Coordinate, SmargonCoordinate\n",
+    "from aaredaqlib.diffraction_geometry import DiffractionGeometry\n",
+    "from aaredaqlib.models import DAQStatusModel\n",
+    "from aaredaqlib.raster_grid import RasterGridRequest\n",
+    "from aaredaqlib.sample_geometry import SampleGeometryModel\n",
+    "from mxlibs3.jfjoch import JFJochWrapper\n",
+    "from aaredaqlib.beamline import MXBeamline"
   ],
   "outputs": [],
-   "execution_count": 1
-  },
-  {
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2025-05-19T14:02:47.337605Z",
-     "start_time": "2025-05-19T14:02:45.686407Z"
-    }
-   },
-   "cell_type": "code",
-   "source": "d = BeamlineDevices(MXBeamline.X06DA)",
-   "id": "4deb7964c8ca4964",
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Connecting TELL p-shell service at http://x06da-tell.psi.ch:22222 ..."
-     ]
-    }
-   ],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
-     "end_time": "2025-05-19T13:44:40.713609Z",
-     "start_time": "2025-05-19T13:44:40.710429Z"
+     "end_time": "2025-06-16T10:22:16.804934Z",
+     "start_time": "2025-06-16T10:22:16.802426Z"
    }
   },
   "cell_type": "code",
-   "source": "d.shutter",
-   "id": "8a2651e47eea2441",
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "False"
-      ]
-     },
-     "execution_count": 5,
-     "metadata": {},
-     "output_type": "execute_result"
+   "source": "j = JFJochWrapper(bl=MXBeamline.X06DA)",
+   "id": "4deb7964c8ca4964",
+   "outputs": [],
+   "execution_count": 4
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2025-06-16T09:38:54.734836Z",
+     "start_time": "2025-06-16T09:38:54.731027Z"
    }
+   },
+   "cell_type": "code",
+   "source": "j.initialize()",
+   "id": "8a2651e47eea2441",
+   "outputs": [],
+   "execution_count": 4
+  },
+  {
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2025-06-16T10:22:19.022173Z",
+     "start_time": "2025-06-16T10:22:17.836161Z"
+    }
+   },
+   "cell_type": "code",
+   "source": [
+    "r = RasterGridRequest(file_prefix=\"test\",exp_time_s=0.01, n_x=100, n_y=1, grid_size_mm=Coordinate(x=0.01, y=0.01), smargon=SmargonCoordinate())\n",
+    "diffraction = DiffractionGeometry(\n",
+    "    energy_keV=12.4,\n",
+    "    dtz_mm=100.0,\n",
+    "    detector_size_pxl=(1553,1630),\n",
+    "    pixel_size_mm=0.150, #PILATUS 4\n",
+    "    beam_center_pxl=(750, 750)\n",
+    ")\n",
+    "\n",
+    "\n",
+    "j.measure_raster(r, diffraction, None)"
   ],
+   "id": "f83e9d8a7ef278e8",
+   "outputs": [],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
-     "end_time": "2025-05-19T13:44:49.376359Z",
-     "start_time": "2025-05-19T13:44:49.374050Z"
+     "end_time": "2025-06-16T10:22:05.636129Z",
+     "start_time": "2025-06-16T10:22:05.474053Z"
    }
   },
   "cell_type": "code",
-   "source": "d.shutter = False",
-   "id": "f83e9d8a7ef278e8",
-   "outputs": [],
-   "execution_count": 7
-  },
-  {
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2025-05-19T13:46:27.545922Z",
-     "start_time": "2025-05-19T13:46:27.543624Z"
-    }
-   },
-   "cell_type": "code",
-   "source": "import cv2",
+   "source": "j.measure_raster()",
   "id": "9020825c2384f66f",
-   "outputs": [],
-   "execution_count": 8
+   "outputs": [
+    {
+     "ename": "NameError",
+     "evalue": "name 'j' is not defined",
+     "output_type": "error",
+     "traceback": [
+      "\u001B[31m---------------------------------------------------------------------------\u001B[39m",
+      "\u001B[31mNameError\u001B[39m                                 Traceback (most recent call last)",
+      "\u001B[36mCell\u001B[39m\u001B[36m \u001B[39m\u001B[32mIn[1]\u001B[39m\u001B[32m, line 1\u001B[39m\n\u001B[32m----> \u001B[39m\u001B[32m1\u001B[39m \u001B[43mj\u001B[49m.measure_raster()\n",
+      "\u001B[31mNameError\u001B[39m: name 'j' is not defined"
+     ]
+    }
+   ],
+   "execution_count": 1
  },
  {
   "metadata": {