Refactor stacking for reuse and overlappability

Currently, the stacking feature adds a lot of code to various parts of ShmemTrainMatcher. It would be nice if this could be cleaner, for:

• more maintainable ShmemTrainMatcher
• reusability: the DetectorAssembler also does stacking except in its own less good way
• reusability: we will at some point probably split ShmemTrainMatcher and some of the resulting matching-and-processing devices will want to stack
• improvability: the current implementation is pretty naive and though the handle_source is intended to be run in a thread pool, only one train can be processed at a time; this requires rethinking of buffer handling to change
  • bonus concern: the stacking buffers, if we start overlapping processing and decoupling this from sending, would not satisfy the safeNDArray flag in the current version

Plan for this MR:

• Move stacking functionality to "friend" class
• Get rid of global stacking buffers; idea: create, per train, a stacking context with fresh buffers
• Add some testing of stacking (not just utils, but the friend class working on tables)
• (Out of scope) Test that it works in a overlapping trains setting
• (Out of scope) Start overlapping train processing
• Refactor DetectorAssembler to use StackingFriend to stack before assembly
• Change the configuration that the manager passes to DetectorAssembler correspondingly
• Nevermind, DetectorAssembler is not a great fit actually; maybe use xarray and exploit extra-geom's partial data assembly functionality (Out of scope; see !74 (closed))
• Test everything with actual Karabo (pending node allocation)

src/tests/test_stacking.py

+import concurrent.futures
+
+from karabo.bound import Hash
+import numpy as np
+import pytest
+
+from calng import stacking_utils
+
+
+class NotALog:
+    def __init__(self, parent):
+        self.parent = parent
+
+    def WARN(self, s):
+        print(f"Warning: {s}")
+        self.parent.warnings.append(s)
+
+
+class NotADevice:
+    def __init__(self):
+        self.log = NotALog(self)
+        self.warnings = []
+
+
+datas = [np.arange(i * 100, i * 100 + 8).reshape(4, 2) for i in range(3)]
+
+
+class CommonTestFixtureGuy:
+    @pytest.fixture(params=[False, True])
+    def thread_pool(self, request):
+        if request.param:
+            with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool:
+                yield pool
+        else:
+            yield None
+
+
+class TestSourceStacking(CommonTestFixtureGuy):
+    @pytest.fixture
+    def friend(self):
+        source_table = [
+            {
+                "select": True,
+                "source": f"source{i}@device{i}:channel",
+            }
+            for i in range(3)
+        ]
+
+        merge_rules = [
+            {
+                "select": True,
+                "sourcePattern": "source\\d+",
+                "keyPattern": "keyToStack",
+                "replacement": "newSource",
+                "groupType": "sources",
+                "mergeMethod": "stack",
+                "axis": 1,
+                "missingValue": "0",
+            }
+        ]
+        device = NotADevice()
+        return stacking_utils.StackingFriend(device, source_table, merge_rules)
+
+    @pytest.fixture
+    def sources(self):
+        return {
+            f"source{i}": (Hash("keyToStack", data), None)
+            for i, data in enumerate(datas)
+        }
+
+    def test_simple(self, friend, sources, thread_pool):
+        friend.process(sources, thread_pool)
+        assert "newSource" in sources
+        stacked = sources["newSource"][0]["keyToStack"]
+        assert stacked.shape == (4, 3, 2)
+        for i, data in enumerate(datas):
+            assert np.array_equal(stacked[:, i], data, equal_nan=True)
+        assert not friend._device.warnings
+
+    def test_missing_source(self, friend, sources, thread_pool):
+        del sources["source0"]
+        friend.process(sources, thread_pool)
+        assert "newSource" in sources
+        stacked = sources["newSource"][0]["keyToStack"]
+        assert np.all(stacked[:, 0] == 0)
+        for i, data in enumerate(datas[1:], start=1):
+            assert np.array_equal(stacked[:, i], data, equal_nan=True)
+        assert not friend._device.warnings
+
+    def test_missing_data(self, friend, sources, thread_pool):
+        sources["source0"][0].erase("keyToStack")
+        friend.process(sources, thread_pool)
+        assert "newSource" in sources
+        stacked = sources["newSource"][0]["keyToStack"]
+        assert np.all(stacked[:, 0] == 0)
+        for i, data in enumerate(datas[1:], start=1):
+            assert np.array_equal(stacked[:, i], data, equal_nan=True)
+        assert not friend._device.warnings
+
+    def test_source_stacking_no_sources(self, friend, sources, thread_pool):
+        sources = {}
+        friend.process(sources, thread_pool)
+        assert friend._device.warnings
+        assert not sources
+
+    def test_source_stacking_erroneous_data(self, friend, sources, thread_pool):
+        sources["source1"][0][
+            "keyToStack"
+        ] = "and now for something completely different"
+        friend.process(sources, thread_pool)
+        assert "newSource" in sources
+        stacked = sources["newSource"][0]["keyToStack"]
+        for i, data in enumerate(datas):
+            if i == 1:
+                assert np.all(stacked[:, i] == 0)
+            else:
+                assert np.array_equal(stacked[:, i], data, equal_nan=True)
+        assert friend._device.warnings
+
+
+class TestKeyStacking(CommonTestFixtureGuy):
+    @pytest.fixture
+    def friend(self):
+        source_table = [
+            {
+                "select": True,
+                "source": "source@device:channel",
+            }
+        ]
+
+        merge_rules = [
+            {
+                "select": True,
+                "sourcePattern": "source",
+                "keyPattern": "key\\d+",
+                "replacement": "newKey",
+                "groupType": "keys",
+                "mergeMethod": "stack",
+                "axis": 1,
+                "missingValue": "0",
+            }
+        ]
+        device = NotADevice()
+        return stacking_utils.StackingFriend(device, source_table, merge_rules)
+
+    @pytest.fixture
+    def sources(self):
+        h = Hash()
+        for i, data in enumerate(datas):
+            h[f"key{i}"] = data
+        return {"source": (h, None)}
+
+    def test_simple(self, friend, sources, thread_pool):
+        friend.process(sources, thread_pool)
+        assert sources["source"][0].has("newKey")
+        stacked = sources["source"][0]["newKey"]
+        for i, data in enumerate(datas):
+            assert np.array_equal(stacked[:, i], data, equal_nan=True)
+        assert not friend._device.warnings