single module operation of a multi-module detector

extra_foam/pipeline/processors/image_assembler.py

@@ -309,6 +309,11 @@ def _assemble(self, modules):
             """
             image_dtype = config["SOURCE_PROC_IMAGE_DTYPE"]
             if self._geom is not None:
+                n_modules = modules.shape[1]
+                if n_modules == 1:
+                    # single module operation
+                    return modules.astype(image_dtype).squeeze(axis=1)
+
                 n_pulses = modules.shape[0]
                 if self._out_array is None or self._out_array.shape[0] != n_pulses:
                     self._out_array = self._geom.output_array_for_position_fast(
@@ -380,12 +385,15 @@ def process(self, data):
 
                 # check number of modules
                 if ndim >= 3 and shape[-3] != n_modules:
+                    n_modules_actual = shape[-3]
                     if config["DETECTOR"] != "JungFrauPR":
-                        raise ValueError(f"Expected {n_modules} modules, but get "
-                                         f"{shape[0]} instead!")
-                    elif shape[-3] > 2:
+                        # allow single module operation
+                        if n_modules_actual != 1:
+                            raise ValueError(f"Expected {n_modules} modules, but get "
+                                             f"{n_modules_actual} instead!")
+                    elif n_modules_actual > 2:
                         raise ValueError(f"Expected 1 or 2 modules, but get "
-                                         f"{shape[0]} instead!")
+                                         f"{n_modules_actual} instead!")
 
                 # check number of memory cells
                 if ndim == 4 and not shape[0]:

extra_foam/pipeline/processors/tests/test_image_assembler.py

@@ -39,6 +39,17 @@ def _check_assembled_result(data, src):
     assert assembled_shape[2] >= 1024
 
 
+def _check_single_module_result(data, src, module_shape):
+    # test the module keys have been deleted
+    assert data['raw'][src] is None
+
+    assembled = data['assembled']['data']
+    assert 3 == assembled.ndim
+    assembled_shape = assembled.shape
+    assert 4 == assembled_shape[0]
+    assert assembled_shape[-2:] == module_shape
+
+
 _tmp_cfg_dir = tempfile.mkdtemp()
 
 
@@ -167,7 +178,6 @@ def testAssembleBridge(self):
             data['raw'][src] = np.ones((0, 16, 512, 128), dtype=_IMAGE_DTYPE)
             self._assembler.process(data)
 
-        expected_shape = (4, 16, 512, 128)
         # (modules, fs, ss, memory cells)
         data['raw'][src] = np.ones((16, 128, 512, 4), dtype=_IMAGE_DTYPE)
         self._assembler.process(data)
@@ -178,6 +188,18 @@ def testAssembleBridge(self):
         self._assembler.process(data)
         _check_assembled_result(data, src)
 
+        # test single module
+
+        # (modules, fs, ss, memory cells)
+        data['raw'][src] = np.ones((1, 128, 512, 4), dtype=_IMAGE_DTYPE)
+        self._assembler.process(data)
+        _check_single_module_result(data, src, config["MODULE_SHAPE"])
+
+        # (memory cells, modules, ss, fs)
+        data['raw'][src] = np.ones((4, 1, 512, 128), dtype=_IMAGE_DTYPE)
+        self._assembler.process(data)
+        _check_single_module_result(data, src, config["MODULE_SHAPE"])
+
 
 class TestLpdAssembler:
     @classmethod
@@ -323,6 +345,11 @@ def testAssembleBridge(self, assembler_type):
         self._assembler.process(data)
         _check_assembled_result(data, src)
 
+        # test single module
+        data['raw'][src] = np.ones((1, 256, 256, 4), dtype=_IMAGE_DTYPE)
+        self._assembler.process(data)
+        _check_single_module_result(data, src, config["MODULE_SHAPE"])
+
     @pytest.mark.parametrize("assembler_type", [GeomAssembler.EXTRA_GEOM, GeomAssembler.OWN])
     def testOutArray(self, assembler_type):
         self._load_geometry(assembler_type)
@@ -513,6 +540,11 @@ def testAssembleBridge(self, assembler_type):
         self._assembler.process(data)
         _check_assembled_result(data, src)
 
+        # test single module
+        data['raw'][src] = np.ones((1, 512, 128, 4), dtype=_IMAGE_DTYPE)
+        self._assembler.process(data)
+        _check_single_module_result(data, src, config["MODULE_SHAPE"])
+
     def testOutArray(self):
         key_name = 'image.data'
         src, catalog = self._create_catalog('SCS_CDIDET_DSSC/CAL/APPEND_CORRECTED', key_name)

extra_foam/pipeline/processors/tests/test_image_roi.py

@@ -542,11 +542,11 @@ def testRoiProjPumpProbe(self, error, direct, axis):
             np.testing.assert_array_almost_equal(y_on_gt, processed.pp.y_on)
             np.testing.assert_array_almost_equal(y_off_gt, processed.pp.y_off)
             np.testing.assert_array_almost_equal(y_on_gt - y_off_gt, processed.pp.y)
-            assert np.abs(y_on_gt - y_off_gt).sum() == pytest.approx(processed.pp.fom, rel=1e-4)
+            assert np.abs(y_on_gt - y_off_gt).sum() == pytest.approx(processed.pp.fom, rel=1e-3)
             # test abs_difference == False
             processed.pp.abs_difference = False
             proc.process(data)
-            assert (y_on_gt - y_off_gt).sum() == pytest.approx(processed.pp.fom, rel=1e-4)
+            assert (y_on_gt - y_off_gt).sum() == pytest.approx(processed.pp.fom, rel=1e-3)
             # test when ROI2 has different shape from ROI1
             processed.roi.geom2.geometry = [1, 0, 1, 3]
             with patch.object(proc, "_process_proj"):