diff --git a/extra_foam/pipeline/processors/image_assembler.py b/extra_foam/pipeline/processors/image_assembler.py index c547aef46..2bd2be942 100644 --- a/extra_foam/pipeline/processors/image_assembler.py +++ b/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]: diff --git a/extra_foam/pipeline/processors/tests/test_image_assembler.py b/extra_foam/pipeline/processors/tests/test_image_assembler.py index 78b9d5ff8..e5c0c4603 100644 --- a/extra_foam/pipeline/processors/tests/test_image_assembler.py +++ b/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) diff --git a/extra_foam/pipeline/processors/tests/test_image_roi.py b/extra_foam/pipeline/processors/tests/test_image_roi.py index 64896d329..110c98db9 100644 --- a/extra_foam/pipeline/processors/tests/test_image_roi.py +++ b/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"):