Enable transposition from weird DAQ axis order on GPU

src/calng/JungfrauCorrection.py

@@ -3,6 +3,7 @@ import enum
 import cupy
 import numpy as np
 from karabo.bound import (
+    BOOL_ELEMENT,
     DOUBLE_ELEMENT,
     KARABO_CLASSINFO,
     OUTPUT_CHANNEL,
@@ -98,10 +99,14 @@ class JungfrauGpuRunner(base_gpu.BaseGpuRunner):
     def _get_gain_map_for_preview(self):
         return self.input_gain_map_gpu
 
-    def load_data(self, image_data, input_gain_map, cell_table):
+    def load_data(self, image_data, input_gain_map, cell_table, daq_transpose=False):
         """Experiment: loading all three in one function as they are tied"""
-        self.input_data_gpu.set(image_data)
-        self.input_gain_map_gpu.set(input_gain_map)
+        if daq_transpose:
+            self.input_data_gpu[:] = cupy.asarray(image_data).transpose()[0]
+            self.input_gain_map_gpu[:] = cupy.asarray(input_gain_map).transpose()[0]
+        else:
+            self.input_data_gpu.set(image_data)
+            self.input_gain_map_gpu.set(input_gain_map)
         if self.burst_mode:
             self.cell_table_gpu.set(cell_table)
 
@@ -253,6 +258,7 @@ class JungfrauCorrection(BaseCorrection):
     _calcat_friend_class = JungfrauCalcatFriend
     _constant_enum_class = JungfrauConstants
     _managed_keys = BaseCorrection._managed_keys.copy()
+    _schema_cache_fields = BaseCorrection._schema_cache_fields.copy()
     _image_data_path = "data.adc"
     _cell_table_path = "data.memoryCell"
 
@@ -281,6 +287,22 @@ class JungfrauCorrection(BaseCorrection):
             .commit(),
         )
 
+        (
+            BOOL_ELEMENT(expected)
+            .key("dataFormat.daqTranspose")
+            .displayedName("Transpose axes from DAQ")
+            .description(
+                "Data on daqOutput channel has interesting axis order. In online "
+                "deployments, this means that a transpose is needed before correction."
+            )
+            .assignmentOptional()
+            .defaultValue(True)
+            .reconfigurable()
+            .commit(),
+        )
+        JungfrauCorrection._schema_cache_fields.add("dataFormat.daqTranspose")
+        JungfrauCorrection._managed_keys.add("dataFormat.daqTranspose")
+
         (
             OUTPUT_CHANNEL(expected)
             .key("preview.outputGainMap")
@@ -341,7 +363,10 @@ class JungfrauCorrection(BaseCorrection):
             cell_table = cell_table[np.newaxis]
         try:
             self.kernel_runner.load_data(
-                image_data, data_hash.get("data.gain"), cell_table
+                image_data,
+                data_hash.get("data.gain"),
+                cell_table,
+                daq_transpose=self._schema_cache["dataFormat.daqTranspose"]
             )
         except ValueError as e:
             self.log_status_warn(f"Failed to load data: {e}")