From c0164645a992c5f6061eedf1a07cbd9a2c38d9de Mon Sep 17 00:00:00 2001
From: Filip Leonarski <filip.leonarski@psi.ch>
Date: Thu, 13 Nov 2025 11:42:52 +0100
Subject: [PATCH] DAQ: Save Fluorescence for rotation scan

---
 daq/src/aaredaq/config.py | 18 +++++++++++++++++-
 daq/src/aaredaq/daq.py    |  6 +++++-
 daq/src/mxlibs3/jfjoch.py | 15 ++++++++++++---
 3 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/daq/src/aaredaq/config.py b/daq/src/aaredaq/config.py
index 6b363943..535a7422 100644
--- a/daq/src/aaredaq/config.py
+++ b/daq/src/aaredaq/config.py
@@ -14,7 +14,8 @@ from aaredaqlib.models import (
     zoom_manager,
     SampleShortInfo,
     SessionStatus,
-    BeamlineStateEnum, SessionsStateEnum, SampleShortInfoList, CryojetSettingsModel, ZoomModel
+    BeamlineStateEnum, SessionsStateEnum, SampleShortInfoList, CryojetSettingsModel, ZoomModel,
+    FluorescenceSpectrumOutputModel
 )
 
 from aaredaqlib.beamline import MXBeamline
@@ -446,6 +447,21 @@ class BeamlineConfig:
     def dtz(self, dtz: float):
         self.__client.set(f"{self.__bl}:dtz", dtz)
 
+    @property
+    def xrf(self) -> FluorescenceSpectrumOutputModel | None:
+        tmp = self.__client.get(f"{self.__bl}:xrf")
+        if tmp is None:
+            return None
+        data_dict = json.loads(tmp)
+        return FluorescenceSpectrumOutputModel(**data_dict)
+
+    @xrf.setter
+    def xrf(self, data: FluorescenceSpectrumOutputModel | None):
+        if data is None:
+            self.__client.delete(f"{self.__bl}:xrf")
+        else:
+            self.__client.set(f"{self.__bl}:xrf", data.model_dump_json())
+
     def clear_mark_beam(self):
         self.__client.delete(f"{self.__bl}:beam_mark")
 
diff --git a/daq/src/aaredaq/daq.py b/daq/src/aaredaq/daq.py
index bbd36ad9..1864771e 100644
--- a/daq/src/aaredaq/daq.py
+++ b/daq/src/aaredaq/daq.py
@@ -396,6 +396,7 @@ class AareDAQ:
             self.crystal_size = CrystalSize(x=0, y=0, z=0)
             self.last_best_b_factor = None
             self.last_best_res = None
+            self.__cfg.xrf = None
 
             if curr_sample is not None and curr_sample.location is None:
                 self.__cfg.current_sample = None
@@ -626,7 +627,7 @@ class AareDAQ:
             self.__devs.smargon.wait()
 
         status = self.status
-        self.__jfjoch.measure_rotation(request, status)
+        self.__jfjoch.measure_rotation(request, status, self.__cfg.xrf)
         self.__aare.create_rotation_run(self.sample, request, status)
 
         if self.sample is not None and self.sample.db_id is not None:
@@ -1545,6 +1546,9 @@ class AareDAQ:
                                                           bkg=self.__devs.fluorimeter.get_current_background(),
                                                           energy_eV=energy,
                                                           average_dead_time=self.__devs.fluorimeter.average_dead_time() / 100.0)
+            if self.sample is not None:
+                self.__cfg.xrf = fluo_output
+
             self.__set_state(BeamlineStateEnum.SampleAlignment)
             self.__cfg.state_busy = False
             return fluo_output
diff --git a/daq/src/mxlibs3/jfjoch.py b/daq/src/mxlibs3/jfjoch.py
index d1d862e0..a4003b4b 100644
--- a/daq/src/mxlibs3/jfjoch.py
+++ b/daq/src/mxlibs3/jfjoch.py
@@ -4,7 +4,7 @@ import jfjoch_client
 
 from aaredaqlib.beamline import MXBeamline
 from aaredaqlib.diffraction_geometry import DiffractionGeometry
-from aaredaqlib.models import SampleShortInfo, DAQStatusModel
+from aaredaqlib.models import SampleShortInfo, DAQStatusModel, FluorescenceSpectrumOutputModel
 from aaredaqlib.raster_grid import RasterGridRequest
 from aaredaqlib.rotation_scan import RotationScanRequest
 
@@ -34,7 +34,8 @@ class JFJochWrapper:
 
     def measure_rotation(self,
                          r: RotationScanRequest,
-                         s: DAQStatusModel) -> None:
+                         s: DAQStatusModel,
+                         f: FluorescenceSpectrumOutputModel | None = None) -> None:
         if s.sample is None:
             pgroup = "p16371"
             sample = "unknown_sample"
@@ -62,6 +63,13 @@ class JFJochWrapper:
             images = r.steps
             trigger = 1
 
+        xrf = None
+        if f is not None:
+            xrf = jfjoch_client.DatasetSettingsXrayFluorescenceSpectrum(
+                energy_eV=f.energy_eV,
+                data=f.spectrum
+            )
+
         dataset_settings = jfjoch_client.DatasetSettings(
             beam_x_pxl=s.diffraction.beam_center_pxl[0],
             beam_y_pxl=s.diffraction.beam_center_pxl[1],
@@ -81,7 +89,8 @@ class JFJochWrapper:
             poni_rot2_rad=s.diffraction.poni_rot2_rad,
             total_flux=s.bl.flux_ph_s,
             max_spot_count=1000,
-            detect_ice_rings=True
+            detect_ice_rings=True,
+            xray_fluorescence_spectrum=xrf
         )
         self.__api.start_post(dataset_settings=dataset_settings)