Correction runners as friends

Status

Done with what I had in mind for now. Sorry that the diff is so big. At least, looking at git diff --numstat master, the individual too long files are a bit more balanced; some of the complexity of BaseCorrection is now kept in BaseKernelRunner and PreviewFriend. Also, detector-specific correction classes are much simpler.

Will do some more testing and figure out how to review.

Goal

Simplify interaction between correction device and kernel runner. To this end, must make the kernel runner smarter and more responsible for handling its relevant subschemata. Some plans:

• Device passes most correction schema (changes) along to runner which has reconfigure which looks at hash (before: device figures out everything, uses loads of special methods to configure runner)
  • This, in particular, means that the "correction flags" (in the end passed to correction kernel) live
• Device gets simpler interaction for running corrections and generating previews
  • Related to flags, now the kernel runner must itself figure out what to do
  • Before: host device tells it to load, correct, do previews, maybe correct again etc.
  • Medium-long term plan (after this MR): should allow concurrent processing, meaning stateful "load this, now process whatever you've loaded" is out
• Other things like specialized _load_constant_to_runner goes away because each runner has uniform load_constant interface

Input handler flow

One of the big changes is the flow of the input handler. This section explains a solution to the issues I was pondering earlier (left as quote below). Before, BaseCorrection would have the main input handler which would extract (most, though not always all) relevant data from the hash and handle (most) quirks we tend to encounter before handing over to subclass-specific process_data. However, these different process_data implementations had a lot of overlap. Two obvious options: keep them separate, encapsulating all the things they do in small reusable functions or unify them and provide hooks / utility functions to fill in per subclass. Currently trying out the latter:

• The input handler (now ON_DATA again*) is only in BaseCorrection
• It uses _get_data_from_hash (may be subclass-specific) to get all the relevant data out of the hash
• It asks the kernel runner about expected_input_shape (needed to fix up wrong axis order from DAQ)
• It asks the kernel runner about expected_output_shape (needed to prepare the next shmem output slot)
• It then passes the extracted data to the kernel runner's correct method

And similarly, the BaseKernelRunner hosts the overarching correct method, letting subclass define certain parts:

• correct in BaseKernelRunner gets all the data from _get_data_from_hash
• It uses _make_output_buffers (may be subclass-specific) to set up some buffers the kernel will write to
• It uses the subclass-specific _correct to actually do the correction
  • This _correct gets all the data and the output buffers, so most of _get_data_from_hash and _make_output_buffers goes here
• Then, previews are computed, with _preview_data_views (may be subclass-specific) providing the views of the relevant ndarrays to preview
  • This gets the image data, any additional data from _get_data_from_hash, and the output buffers from _make_output_buffers and can return arbitrary list of buffer views
  • Each buffer view returned should match a preview output of the device; this allow preview computation to just be applying the same preview reduction function on each view and write it out

Although an improvement to the previous version, I'm not pleased that the different parts are still so tightly coupled. Not sure how to avoid this.

*A long time ago, we switched from ON_DATA to ON_INPUT, partly because - in case we'd get irrelevant sources coming in - the latter would allow us to not read content of certain input hashes unless relevant to us. ON_DATA is simpler, however, and as I understand from Gero's description of the surrounding C++ code changes, we don't decide whether or not to read the input any longer (the C++ code reads hashes and just gives them to us), so the potential performance edge of ON_INPUT should be gone. And we don't need trainmatched data, so it's not important to line up hashes on same input event.

Follow-up work

I've removed the frame filter functionality for now for two reasons

• It complicates BaseCorrection, configuration, and shape handling
• It slices data out before it gets previewable; could be an issue for detector monitoring
  • OTOH, the operator may want the nanmean of only the relevant / lit frames; but frame filter cannot accommodate both use cases

I know that MID tends to use the frame filter, so I'd just add a trivial frame selection kernel with the same logic.

Deployment notes

Having restructured the preview output infrastructure anyway, I also tied it to more node-local subschemata. So now each preview output can have its own settings about how to reduce (it was kind of silly to, say, look at mean of corrected and then also have to get mean of gain map). With a node for each output, preview.[preview_name], I figured we could move the outputs there, so preview.[preview_name].output. This means that managers all over will need an update to the preview layer mapping* and stored scenes referencing previews need updating.

*Technically, the manager can now figure out the preview outputs easily by looking (in the schema) at the keys under preview. Indeed, the correction_device_overview scene does exactly this to generate per-device preview scene links.

src/calng/corrections/JungfrauCorrection.py

@@ -74,8 +74,8 @@ class JungfrauBaseRunner(base_kernel_runner.BaseKernelRunner):
     _bad_pixel_constants = bad_pixel_constants
     _correction_flag_class = CorrectionFlags
     _correction_steps = correction_steps
-    num_pixels_ss = 1024
-    num_pixels_fs = 512
+    num_pixels_ss = 512
+    num_pixels_fs = 1024
 
     @property
     def _map_shape(self):
@@ -168,7 +168,7 @@ class JungfrauBaseRunner(base_kernel_runner.BaseKernelRunner):
             # note: always masking unmapped pixels (not respecting NON_STANDARD_SIZE)
 
     def _load_constant(self, constant, constant_data):
-        if constant_data.shape[0] == self.num_pixels_ss:
+        if constant_data.shape[0] == self.num_pixels_fs:
             constant_data = np.transpose(constant_data, (2, 1, 0, 3))
         else:
             constant_data = np.transpose(constant_data, (2, 0, 1, 3))
@@ -219,8 +219,6 @@ class JungfrauGpuRunner(JungfrauBaseRunner):
         source_module = self._xp.RawModule(
             code=base_kernel_runner.get_kernel_template("jungfrau_gpu.cu").render(
                 {
-                    "pixels_x": self.num_pixels_ss,
-                    "pixels_y": self.num_pixels_fs,
                     "constant_memory_cells": self._constant_memory_cells,
                     "output_data_dtype": utils.np_dtype_to_c_type(
                         self._output_dtype
@@ -443,7 +441,7 @@ class JungfrauCalcatFriend(base_calcat.BaseCalcatFriend):
 @KARABO_CLASSINFO("JungfrauCorrection", deviceVersion)
 class JungfrauCorrection(base_correction.BaseCorrection):
     _base_output_schema = schemas.jf_output_schema
-    _correction_flag_class = CorrectionFlags
+    _correction_steps = correction_steps
     _calcat_friend_class = JungfrauCalcatFriend
     _constant_enum_class = Constants
     _preview_outputs = [
@@ -545,25 +543,3 @@ class JungfrauCorrection(base_correction.BaseCorrection):
             return JungfrauCpuRunner
         else:
             return JungfrauGpuRunner
-
-    def postReconfigure(self):
-        super().postReconfigure()
-
-        if not hasattr(self, "_prereconfigure_update_hash"):
-            return
-
-        update = self._prereconfigure_update_hash
-
-        if update.has("corrections.strixel.enable"):
-            self._lock_and_update(self._update_frame_filter)
-
-        if update.has("corrections.badPixels.subsetToUse"):
-            self.log_status_info("Updating bad pixel maps based on subset specified")
-            # note: now just always reloading from cache for convenience
-            with self.calcat_friend.cached_constants_lock:
-                self.kernel_runner.flush_buffers(bad_pixel_constants)
-                for constant in bad_pixel_constants:
-                    if constant in self.calcat_friend.cached_constants:
-                        self._load_constant_to_runner(
-                            constant, self.calcat_friend.cached_constants[constant]
-                        )