Add SFX hitfinder arbiter kernel

setup.py

@@ -22,8 +22,8 @@ setup(
     packages=find_packages("src"),
     entry_points={
         "calng.arbiter_kernel": [
-            "HitFinderSPI = sfx_addons.hitfinder_spi:HitFinderSPI",
             "SPIhitfinder = sfx_addons.spi_hitfinder:SPIhitfinder",
+            "SFXhitfinder = sfx_addons.sfx_hitfinder:SFXhitfinder",
         ],
     },
     requires=[],

src/sfx_addons/sfx_hitfinder.py

+from collections import deque
+
+import numpy as np
+from calng.arbiter_kernels.base_kernel import BaseArbiterKernel
+from karabo.bound import (
+    DOUBLE_ELEMENT, INT32_ELEMENT, NDARRAY_ELEMENT, NODE_ELEMENT,
+    OVERWRITE_ELEMENT, VECTOR_UINT32_ELEMENT, Hash)
+
+from .utils import guess_module_number
+
+
+class Peaks:
+    def __init__(self):
+        self._reset_parts()
+
+    def _reset_parts(self):
+        self._parts = dict(x=[], y=[], intensity=[], mod=[], num=[])
+
+    def add_module(self, mod_no, data):
+        self._parts["num"].append(data["numPeaks"])
+        self._parts["mod"].append(mod_no)
+        self._parts["x"].append(data["peakX"])
+        self._parts["y"].append(data["peakY"])
+        self._parts["intensity"].append(data["peakIntensity"])
+
+    def compress(self):
+        self.num = np.sum(self._parts["num"], axis=0)
+        max_peaks = np.max(self.num)
+        num_frames = len(self.num)
+
+        self.mod = np.zeros([num_frames, max_peaks], np.uint16)
+        self.x = np.zeros([num_frames, max_peaks], np.float32)
+        self.y = np.zeros([num_frames, max_peaks], np.float32)
+        self.intensity = np.zeros([num_frames, max_peaks], np.float32)
+
+        num_current = np.zeros(num_frames, dtype=int)
+        for mod_ix, mod_no in enumerate(self._parts["mod"]):
+            num_on_module = self._parts["num"][mod_ix]
+            for frame_ix, (f0, n) in enumerate(zip(num_current, num_on_module)):
+                fN = f0 + n
+                self.x[frame_ix, f0:fN] = self._parts["x"][mod_ix][frame_ix, :n]
+                self.y[frame_ix, f0:fN] = self._parts["y"][mod_ix][frame_ix, :n]
+                self.mod[frame_ix, f0:fN] = mod_no
+                self.intensity[frame_ix, f0:fN] = self._parts["intensity"][mod_ix][frame_ix, :n]
+
+            num_current += num_on_module
+
+        self._reset_parts()
+
+    def to_hash(self, out=Hash()):
+        out.set("number", list(map(float, self.num)))
+        out.set("mod", self.mod)
+        out.set("x", self.x)
+        out.set("y", self.y)
+        out.set("intensity", self.intensity)
+        return out
+
+    @staticmethod
+    def extend_schema(schema, prefix):
+        (
+            VECTOR_UINT32_ELEMENT(schema)
+            .key(f"{prefix}.number")
+            .assignmentOptional()
+            .defaultValue([])
+            .commit(),
+
+            NDARRAY_ELEMENT(schema)
+            .key(f"{prefix}.mod")
+            .dtype("UINT32")
+            .commit(),
+
+            NDARRAY_ELEMENT(schema)
+            .key(f"{prefix}.x")
+            .dtype("FLOAT")
+            .commit(),
+
+            NDARRAY_ELEMENT(schema)
+            .key(f"{prefix}.y")
+            .dtype("FLOAT")
+            .commit(),
+
+            NDARRAY_ELEMENT(schema)
+            .key(f"{prefix}.intensity")
+            .dtype("FLOAT")
+            .commit(),
+        )
+
+
+class SFXhitfinder(BaseArbiterKernel):
+    _node_name = "sfxHitfinder"
+
+    def __init__(self, device, name, config):
+        self._history_hitrate = deque([], 100)
+        super().__init__(device, name, config)
+
+    @staticmethod
+    def extend_device_schema(schema, prefix):
+        (
+            OVERWRITE_ELEMENT(schema)
+            .key(prefix)
+            .setNewDescription(
+                "This kernel selects the frames by comparing the number"
+                "of peaks to the threshold."
+            )
+            .commit(),
+
+            INT32_ELEMENT(schema)
+            .key(f"{prefix}.minPeaks")
+            .assignmentOptional()
+            .defaultValue(10)
+            .reconfigurable()
+            .commit(),
+
+            INT32_ELEMENT(schema)
+            .key(f"{prefix}.maxPeaks")
+            .assignmentOptional()
+            .defaultValue(2000)
+            .reconfigurable()
+            .commit(),
+
+            INT32_ELEMENT(schema)
+            .key(f"{prefix}.hitrateAverageWindow")
+            .assignmentOptional()
+            .defaultValue(100)
+            .reconfigurable()
+            .commit(),
+        )
+
+    @staticmethod
+    def extend_output_schema(schema, name):
+        (
+            NODE_ELEMENT(schema)
+            .key(f"{name}.peaks")
+            .commit(),
+
+            INT32_ELEMENT(schema)
+            .key(f"{name}.numberOfHits")
+            .assignmentOptional()
+            .defaultValue(0)
+            .commit(),
+
+            INT32_ELEMENT(schema)
+            .key(f"{name}.numberOfMiss")
+            .assignmentOptional()
+            .defaultValue(0)
+            .commit(),
+
+            DOUBLE_ELEMENT(schema)
+            .key(f"{name}.hitrate")
+            .assignmentOptional()
+            .defaultValue(0.0)
+            .commit(),
+
+            DOUBLE_ELEMENT(schema)
+            .key(f"{name}.averageHitrate")
+            .assignmentOptional()
+            .defaultValue(0.0)
+            .commit(),
+        )
+        Peaks.extend_schema(schema, f"{name}.peaks")
+
+    def reconfigure(self, config):
+        if config.has("minPeaks"):
+            self._min_peaks = config["minPeaks"]
+        if config.has("maxPeaks"):
+            self._max_peaks = config["maxPeaks"]
+        if config.has("hitrateAverageWindow"):
+            self._hitrate_window = config["hitrateAverageWindow"]
+            self._history_hitrate = deque(self._history_hitrate, self._hitrate_window)
+
+    def consider(self, train_id, sources, num_frames, mask, out_hash):
+        num_work_frames = int(np.sum(mask))
+
+        peaks = Peaks()
+        for source, (data, _) in sources.items():
+            modno = guess_module_number(source)
+            if modno < 0:
+                continue
+            if data.has("peakfinding"):
+                peaks.add_module(modno, data["peakfinding"])
+
+        peaks.compress()
+
+        hits = mask & (self._min_peaks <= peaks.num) & (peaks.num <= self._max_peaks)
+
+        num_hits = int(np.sum(hits))
+        num_miss = num_work_frames - num_hits
+        hitrate = num_hits / num_work_frames if num_work_frames > 0 else .0
+        self._history_hitrate.append(hitrate)
+
+        out_hash.set(f"{self._name}.peaks", peaks.to_hash())
+        out_hash.set(f"{self._name}.numberOfHits", num_hits)
+        out_hash.set(f"{self._name}.numberOfMiss", num_miss)
+        out_hash.set(f"{self._name}.hitrate", hitrate)
+        out_hash.set(f"{self._name}.averageHitrate", np.mean(self._history_hitrate))
+
+        return hits