Merge branch 'dev/generalize-litpixels-addon' into 'master'

Generalized LitPixelCounter

See merge request !134

setup.py

@@ -51,7 +51,7 @@ setup(
         ],
         "calng.correction_addon": [
             "IntegratedIntensity = calng.correction_addons.integrated_intensity:IntegratedIntensity",  # noqa
-            "LitPixelCounter = calng.correction_addons.litpixel_counter:LitPixelCounter [agipd]",  # noqa
+            "LitPixelCounter = calng.correction_addons.litpixel_counter:LitPixelCounter",  # noqa
             "Peakfinder9 = calng.correction_addons.peakfinder9:Peakfinder9",  # noqa
             "RandomFrames = calng.correction_addons.random_frames:RandomFrames",
             "SaturationMonitor = calng.correction_addons.saturation_monitor:SaturationMonitor",# noqa

src/calng/correction_addons/litpixel_counter.py

 import numpy as np
-
-from karabo.bound import NODE_ELEMENT, NDARRAY_ELEMENT, DOUBLE_ELEMENT
+from karabo.bound import (
+    DOUBLE_ELEMENT, NDARRAY_ELEMENT, NODE_ELEMENT, UINT32_ELEMENT)
 
 from .base_addon import BaseCorrectionAddon
 
 
+def maybe_get(a):
+    # TODO: proper check for cupy
+    if hasattr(a, "get"):
+        return a.get()
+    return a
+
+
+def factors(n):
+    if n is None:
+        return np.array([1])
+    f = [i for i in range(2, int(np.sqrt(n)) + 1) if (n % i) == 0]
+    f.extend(n // i for i in reversed(f))
+    return np.array(f + [n])
+
+
 class LitPixelCounter(BaseCorrectionAddon):
     def __init__(self, device, prefix, config):
         super().__init__(device, prefix, config)
         global cupy
         import cupy
 
-        self._threshold = config["threshold"]
+        self._ss_size = device.kernel_runner.num_pixels_ss
+        self._ss_factors = factors(self._ss_size)
+        self._fs_size = device.kernel_runner.num_pixels_fs
+        self._fs_factors = factors(self._fs_size)
+
+        self.reconfigure(config)
+
+        # these properties may be changed to the factors
+        # of module dimensions due to protection
+        device.set(f"{prefix}.ssAsicSize", self._ss_asic_size)
+        device.set(f"{prefix}.fsAsicSize", self._fs_asic_size)
 
     def reconfigure(self, changed_config):
         if changed_config.has("threshold"):
             self._threshold = changed_config["threshold"]
+        if changed_config.has("ssAsicSize"):
+            size = changed_config["ssAsicSize"]
+            if size == 0:
+                self._ss_asic_size = self._ss_size
+            else:
+                self._ss_asic_size = int(self._ss_factors[
+                    np.argmin(np.abs(self._ss_factors - size))])
+            changed_config["ssAsicSize"] = self._ss_asic_size
+        if changed_config.has("fsAsicSize"):
+            size = changed_config["fsAsicSize"]
+            if size == 0:
+                self._fs_asic_size = self._fs_size
+            else:
+                self._fs_asic_size = int(self._fs_factors[
+                    np.argmin(np.abs(self._fs_factors - size))])
+            changed_config["fsAsicSize"] = self._fs_asic_size
 
     @staticmethod
     def extend_output_schema(schema):
@@ -30,9 +71,19 @@ class LitPixelCounter(BaseCorrectionAddon):
             .commit(),
 
             NDARRAY_ELEMENT(schema)
-            .key("litpixels.unmasked")
+            .key("litpixels.intensity")
             .dtype('FLOAT')
-            .commit()
+            .commit(),
+
+            NDARRAY_ELEMENT(schema)
+            .key("litpixels.unmasked")
+            .dtype('UINT32')
+            .commit(),
+
+            NDARRAY_ELEMENT(schema)
+            .key("litpixels.total")
+            .dtype('UINT32')
+            .commit(),
         )
 
     @staticmethod
@@ -44,15 +95,46 @@ class LitPixelCounter(BaseCorrectionAddon):
             .assignmentOptional()
             .defaultValue(6.0)
             .reconfigurable()
-            .commit()
+            .commit(),
+
+            UINT32_ELEMENT(schema)
+            .key(f"{prefix}.ssAsicSize")
+            .tags("managed")
+            .assignmentOptional()
+            .defaultValue(0)
+            .reconfigurable()
+            .commit(),
+
+            UINT32_ELEMENT(schema)
+            .key(f"{prefix}.fsAsicSize")
+            .tags("managed")
+            .assignmentOptional()
+            .defaultValue(0)
+            .reconfigurable()
+            .commit(),
         )
 
     def post_correction(self, tid, data, cell_table, pulse_table, output_hash):
-        n_cells, n_x, n_y = data.shape
-        per_asic_data = data.reshape(n_cells, 64, n_x // 64, 64, n_y // 64)
+        if np.ndim(data) == 2:
+            n_cells, fs_size = data.shape
+            axis = 1
+            shape = (n_cells,
+                     self._fs_asic_size, fs_size // self._fs_asic_size)
+        else:
+            n_cells, ss_size, fs_size = data.shape
+            axis = (1, 3)
+            shape = (n_cells,
+                     self._ss_asic_size, ss_size // self._ss_asic_size,
+                     self._fs_asic_size, fs_size // self._fs_asic_size)
+
+        per_asic_data = data.reshape(shape)
 
-        lit_pixels = np.sum(per_asic_data > self._threshold, axis=(1, 3))
-        unmasked_pixels = np.isfinite(per_asic_data).sum(axis=(1, 3))
+        lit_pixels = np.sum(per_asic_data > self._threshold, axis=axis)
+        intensity = np.nansum(per_asic_data, axis=axis)
+        unmasked_pixels = np.isfinite(per_asic_data).sum(axis=axis)
 
-        output_hash["litpixels.count"] = lit_pixels.get()
-        output_hash["litpixels.unmasked"] = unmasked_pixels.get()
+        output_hash["litpixels.count"] = maybe_get(lit_pixels)
+        output_hash["litpixels.intensity"] = maybe_get(intensity)
+        output_hash["litpixels.unmasked"] = maybe_get(unmasked_pixels)
+        output_hash["litpixels.total"] = maybe_get(np.full_like(
+            lit_pixels, self._ss_asic_size * self._fs_asic_size))