Add gain map preview layer to JungfrauCorrection

src/calng/JungfrauCorrection.py

@@ -5,6 +5,7 @@ import numpy as np
 from karabo.bound import (
     DOUBLE_ELEMENT,
     KARABO_CLASSINFO,
+    OUTPUT_CHANNEL,
     OVERWRITE_ELEMENT,
     STRING_ELEMENT,
     VECTOR_STRING_ELEMENT,
@@ -12,7 +13,7 @@ from karabo.bound import (
 
 from . import base_gpu, calcat_utils, utils
 from ._version import version as deviceVersion
-from .base_correction import BaseCorrection, add_correction_step_schema
+from .base_correction import BaseCorrection, add_correction_step_schema, preview_schema
 
 
 _pretend_pulse_table = np.arange(16, dtype=np.uint8)
@@ -58,8 +59,8 @@ class JungfrauGpuRunner(base_gpu.BaseGpuRunner):
         # TODO: avoid superclass creating cell table with wrong dtype first
         self.cell_table_gpu = cupy.empty(self.memory_cells, dtype=cupy.uint8)
         self.input_gain_map_gpu = cupy.empty(self.input_shape, dtype=cupy.uint8)
+        self.preview_buffer_getters.append(self._get_gain_map_for_preview)
         self.map_shape = (self.constant_memory_cells, self.pixels_y, self.pixels_x, 3)
-        # is jungfrau stuff gain mapped?
         self.offset_map_gpu = cupy.zeros(self.map_shape, dtype=cupy.float32)
         self.rel_gain_map_gpu = cupy.ones(self.map_shape, dtype=cupy.float32)
         self.bad_pixel_map_gpu = cupy.zeros(self.map_shape, dtype=cupy.uint32)
@@ -93,6 +94,9 @@ class JungfrauGpuRunner(base_gpu.BaseGpuRunner):
     def _get_corrected_for_preview(self):
         return self.processed_data_gpu
 
+    def _get_gain_map_for_preview(self):
+        return self.input_gain_map_gpu
+
     def load_data(self, image_data, input_gain_map, cell_table):
         """Experiment: loading all three in one function as they are tied"""
         self.input_data_gpu.set(image_data)
@@ -270,6 +274,13 @@ class JungfrauCorrection(BaseCorrection):
             .commit(),
         )
 
+        (
+            OUTPUT_CHANNEL(expected)
+            .key("preview.outputGainMap")
+            .dataSchema(preview_schema)
+            .commit(),
+        )
+
         JungfrauCalcatFriend.add_schema(expected, JungfrauCorrection._managed_keys)
         add_correction_step_schema(
             expected,
@@ -352,9 +363,11 @@ class JungfrauCorrection(BaseCorrection):
                 _pretend_pulse_table,
                 warn_func=self.log_status_warn,
             )
-            preview_raw, preview_corrected = self.kernel_runner.compute_previews(
-                preview_slice_index
-            )
+            (
+                preview_raw,
+                preview_corrected,
+                preview_gain_map
+            ) = self.kernel_runner.compute_previews(preview_slice_index)
 
         # reusing input data hash for sending
         data_hash.set(self._image_data_path, buffer_handle)
@@ -367,6 +380,7 @@ class JungfrauCorrection(BaseCorrection):
                 (
                     ("preview.outputRaw", preview_raw),
                     ("preview.outputCorrected", preview_corrected),
+                    ("preview.outputGainMap", preview_gain_map),
                 ),
                 train_id,
                 source,