Prepare hack to sneak in gain counting

src/calng/base_correction.py

@@ -893,6 +893,10 @@ class BaseCorrection(PythonDevice):
             corrections, processed_buffer, previews = self.kernel_runner.correct(
                 image_data, cell_table, *additional_data
             )
+            # hack: get extra data out from AGIPD until addon API is updated
+            if isinstance(corrections, list):
+                corrections, extra_stuff = corrections
+                data_hash.merge(extra_stuff)
             data_hash["corrections"] = corrections
 
             # write previews first so addons cannot mess with them

src/calng/corrections/AgipdCorrection.py

@@ -6,10 +6,12 @@ from karabo.bound import (
     DOUBLE_ELEMENT,
     FLOAT_ELEMENT,
     KARABO_CLASSINFO,
+    NDARRAY_ELEMENT,
     OUTPUT_CHANNEL,
     OVERWRITE_ELEMENT,
     STRING_ELEMENT,
     UINT32_ELEMENT,
+    Hash,
     MetricPrefix,
     Unit,
 )
@@ -343,11 +345,6 @@ class AgipdBaseRunner(base_kernel_runner.BaseKernelRunner):
             self.num_pixels_fs,
         )
 
-    def extra_processing(
-        self, out_hash, processed_dat,a cell_table, pulse_table, gain_map
-    ):
-        out_hash["numPixelsPerGainStage"] = ...
-
     @property
     def _gm_map_shape(self):
         return self._map_shape + (3,)  # for gain-mapped constants
@@ -611,6 +608,15 @@ class AgipdGpuRunner(AgipdBaseRunner):
             ),
             name="common_mode_asic",
         )
+        self.gain_count_kernel = self._xp.RawKernel(
+            code=base_kernel_runner.get_kernel_template("gaincount.cu").render(
+                ss_dim=self.num_pixels_ss,
+                fs_dim=self.num_pixels_fs,
+                block_size=128,  # TODO: tune
+                num_gain_stages=3,
+            ),
+            name="count_pixels_per_gain_stage",
+        )
 
     def _correct(self, flags, image_data, cell_table, processed_data, gain_map):
         num_frames = self._xp.uint16(image_data.shape[0])
@@ -658,6 +664,59 @@ class AgipdGpuRunner(AgipdBaseRunner):
                 ),
             )
 
+    def correct(self, image_data, cell_table, *additional_data):
+        """Temporary override to utilize gain map for monitoring"""
+        num_frames = image_data.shape[0]
+        processed_buffers = self._make_output_buffers(
+            num_frames, self._correction_flag_preview
+        )
+        image_data = self._xp.asarray(image_data)
+        additional_data = [self._xp.asarray(data) for data in additional_data]
+        self._correct(
+            self._correction_flag_preview,
+            image_data,
+            cell_table,
+            *additional_data,
+            *processed_buffers,
+        )
+
+        preview_buffers = self._preview_data_views(
+            image_data, *additional_data, *processed_buffers
+        )
+
+        if self._correction_flag_preview != self._correction_flag_enabled:
+            processed_buffers = self._make_output_buffers(
+                num_frames, self._correction_flag_enabled
+            )
+            self._correct(
+                self._correction_flag_enabled,
+                image_data,
+                cell_table,
+                *additional_data,
+                *processed_buffers,
+            )
+        num_frames = buffers["image.data"].shape[0]
+        gain_count_buffer = self._xp.empty(
+            (num_frames, 3), dtype=np.uint32
+        )
+        self.gain_count_kernel(
+            (num_frames, 1, 1),
+            (1, 128, 1),
+            (
+                buffers["image.gain"],
+                num_frames,
+                gain_count_buffer,
+            )
+        )
+        return (
+            [
+                self._correction_applied_hash,
+                Hash("numPixelsPerGainStage", gain_count_buffer.get()),
+            ],
+            processed_buffers,
+            preview_buffers,
+        )
+
 
 class AgipdCalcatFriend(base_calcat.BaseCalcatFriend):
     _constant_enum_class = Constants
@@ -786,10 +845,20 @@ class AgipdCalcatFriend(base_calcat.BaseCalcatFriend):
         return res
 
 
+def custom_output_schema(use_shmem_handles):
+    res = schemas.xtdf_output_schema(use_shmem_handles)
+    (
+        NDARRAY_ELEMENT(res)
+        .key("numPixelsPerGainStage")
+        .commit(),
+    )
+    return res
+
+
 @KARABO_CLASSINFO("AgipdCorrection", deviceVersion)
 class AgipdCorrection(base_correction.BaseCorrection):
     # subclass *must* set these attributes
-    _base_output_schema = schemas.xtdf_output_schema
+    _base_output_schema = custom_output_schema
     _calcat_friend_class = AgipdCalcatFriend
     _constant_enum_class = Constants
     _correction_steps = correction_steps