Merge pull request #653 from DiamondLightSource/datachecker

Data checker auto insertion

Author: dkazanc

httomo/runner/task_runner.py

@@ -96,11 +96,20 @@ def _execute_section(self, section: Section, section_index: int = 0):
             f"Section {section_index} (pattern={section.methods[0].pattern.name}) with the following methods:",
             level=logging.INFO,
         )
-        methods_info = [
-            f"    {method.method_name} ({method.package_name})\n"
-            for method in section.methods
-        ]
+
+        methods_info = []
+        for method in section.methods:
+            if method.method_name == "data_checker":
+                methods_info.append(
+                    f"    --- output check with data_checker --- ({method.package_name})\n"
+                )
+            else:
+                methods_info.append(
+                    f"    {method.method_name} ({method.package_name})\n"
+                )
+
         methods_info[-1] = methods_info[-1].rstrip("\n")
+
         self._log_pipeline(methods_info, level=logging.INFO)
 
         slicing_dim_section: Literal[0, 1] = _get_slicing_dim(section.pattern) - 1  # type: ignore

httomo/transform_layer.py

@@ -37,6 +37,8 @@ def __init__(
 
     def transform(self, pipeline: Pipeline) -> Pipeline:
         pipeline = self.insert_data_reducer(pipeline)
+        pipeline = self.insert_data_checker(pipeline)
+
         pipeline_is_sweep = _check_if_pipeline_has_a_sweep(pipeline)
 
         if pipeline_is_sweep:
@@ -89,6 +91,70 @@ def insert_data_reducer(self, pipeline: Pipeline) -> Pipeline:
             methods.append(m)
         return Pipeline(loader, methods)
 
+    def insert_data_checker(self, pipeline: Pipeline) -> Pipeline:
+        """This will insert CPU or GPU data checker method AFTER most of the methods in the pipeline"""
+        loader = pipeline.loader
+        methods = []
+        methods.append(
+            GenericMethodWrapper(
+                self._repo,
+                "httomolib.misc.utils",
+                "data_checker",
+                comm=self._comm,
+                save_result=False,
+                task_id="datachecker_0",
+                infsnans_correct=False,  # the input (raw) data is 16bit
+                zeros_warning=True,  # we count the zeros if the data is too sparse
+                data_to_method_name="Data Loader",
+            ),
+        )
+        for index, m in enumerate(pipeline):
+            methods.append(m)
+            # handling some exceptions here after which we don't need to insert the data checker
+            if (
+                m.method_name
+                not in [
+                    "data_reducer",
+                    "data_checker",
+                    "calculate_stats",
+                    "rescale_to_int",
+                ]
+                and "rotation" not in m.module_path
+                and index < len(pipeline._methods) - 1
+            ):
+                if m.is_cpu:
+                    # add the CPU checker method
+                    methods.append(
+                        GenericMethodWrapper(
+                            self._repo,
+                            "httomolib.misc.utils",
+                            "data_checker",
+                            comm=self._comm,
+                            save_result=False,
+                            task_id=f"datachecker_{m.task_id}",
+                            infsnans_correct=True,
+                            zeros_warning=True,
+                            data_to_method_name=m.method_name,
+                        ),
+                    )
+                else:
+                    # add the GPU checker method
+                    methods.append(
+                        GenericMethodWrapper(
+                            self._repo,
+                            "httomolibgpu.misc.utils",
+                            "data_checker",
+                            comm=self._comm,
+                            save_result=False,
+                            task_id=f"datachecker_{m.task_id}",
+                            infsnans_correct=True,
+                            zeros_warning=False,
+                            data_to_method_name=m.method_name,
+                        ),
+                    )
+
+        return Pipeline(loader, methods)
+
     def insert_save_images_after_sweep(self, pipeline: Pipeline) -> Pipeline:
         """For sweep methods we add image saving method."""
         loader = pipeline.loader

tests/test_transform_layer.py

@@ -186,6 +186,82 @@ def test_insert_image_save_after_sweep(mocker: MockerFixture, tmp_path: Path):
     assert pipeline[3].config_params["axis"] == 1
 
 
+def test_insert_data_checker(mocker: MockerFixture, tmp_path: Path):
+    comm = MPI.COMM_SELF
+    repo = make_mock_repo(mocker)
+    loader = mocker.create_autospec(
+        LoaderInterface,
+        instance=True,
+    )
+    pipeline = Pipeline(
+        loader=loader,
+        methods=[
+            make_test_method(
+                mocker,
+                method_name="remove_outlier",
+                module_path="httomolibgpu.misc.corr",
+                save_result=False,
+                task_id="t1",
+            ),
+            make_test_method(
+                mocker,
+                method_name="find_center_vo",
+                module_path="httomolibgpu.recon.rotation",
+                save_result=False,
+                task_id="t2",
+            ),
+            make_test_method(
+                mocker,
+                method_name="normalize",
+                module_path="httomolibgpu.prep.normalize",
+                save_result=False,
+                task_id="t3",
+            ),
+            make_test_method(
+                mocker,
+                method_name="FBP3d_tomobar",
+                module_path="httomolibgpu.recon.algorithm",
+                save_result=False,
+                sweep=True,
+                task_id="t4",
+            ),
+            make_test_method(
+                mocker,
+                method_name="calculate_stats",
+                module_path="httomo.methods",
+                save_result=False,
+                task_id="t5",
+            ),
+            make_test_method(
+                mocker,
+                method_name="rescale_to_int",
+                module_path="httomolib.misc.rescale",
+                save_result=False,
+                task_id="t6",
+            ),
+            make_test_method(
+                mocker,
+                method_name="save_to_images",
+                module_path="httomolib.misc.images",
+                save_result=False,
+                task_id="t7",
+            ),
+        ],
+    )
+    trans = TransformLayer(comm, repo=repo, save_all=False, out_dir=tmp_path)
+    pipeline = trans.insert_data_checker(pipeline)
+
+    assert len(pipeline) == 11
+    assert pipeline[0].method_name == "data_checker"
+    assert pipeline[2].method_name == "data_checker"
+    assert pipeline[4].method_name == "normalize"
+    assert pipeline[5].method_name == "data_checker"
+    assert pipeline[6].method_name == "FBP3d_tomobar"
+    assert pipeline[7].method_name == "data_checker"
+    assert pipeline[9].method_name == "rescale_to_int"
+    assert pipeline[10].method_name == "save_to_images"
+
+
 def test_insert_image_save_after_sweep2(mocker: MockerFixture, tmp_path: Path):
     comm = MPI.COMM_SELF
     repo = make_mock_repo(mocker)
@@ -220,7 +296,7 @@ def test_insert_image_save_after_sweep2(mocker: MockerFixture, tmp_path: Path):
             ),
             make_test_method(
                 mocker,
-                method_name="FBP",
+                method_name="FBP3d_tomobar",
                 module_path="httomolibgpu.recon.algorithm",
                 save_result=False,
                 task_id="t4",
@@ -237,7 +313,7 @@ def test_insert_image_save_after_sweep2(mocker: MockerFixture, tmp_path: Path):
 
     assert pipeline[5].method_name == "save_to_images"
     assert pipeline[5].task_id == "saveimage_sweep_t4"
-    assert pipeline[5].config_params["subfolder_name"] == "images_sweep_FBP"
+    assert pipeline[5].config_params["subfolder_name"] == "images_sweep_FBP3d_tomobar"
     assert pipeline[5].config_params["axis"] == 1
 
 
@@ -285,13 +361,13 @@ def test_insert_paganin_not_last_sweep(mocker: MockerFixture, tmp_path: Path):
     trans = TransformLayer(comm, repo=repo, save_all=False, out_dir=tmp_path)
     pipeline = trans.transform(pipeline)
 
-    assert len(pipeline) == 7
-    assert pipeline[4].method_name == "save_to_images"
-    assert pipeline[4].task_id == "saveimage_sweep_t3"
-    assert pipeline[4].config_params["subfolder_name"] == "images_sweep_paganin_filter"
-    assert pipeline[5].method_name == "FBP3d_tomobar"
-    assert pipeline[6].task_id == "saveimage_sweep_t4"
-    assert pipeline[6].config_params["subfolder_name"] == "images_sweep_FBP3d_tomobar"
+    assert len(pipeline) == 11
+    assert pipeline[7].method_name == "save_to_images"
+    assert pipeline[7].task_id == "saveimage_sweep_t3"
+    assert pipeline[7].config_params["subfolder_name"] == "images_sweep_paganin_filter"
+    assert pipeline[9].method_name == "FBP3d_tomobar"
+    assert pipeline[10].task_id == "saveimage_sweep_t4"
+    assert pipeline[10].config_params["subfolder_name"] == "images_sweep_FBP3d_tomobar"
 
 
 def test_insert_paganin_is_last_sweep(mocker: MockerFixture, tmp_path: Path):
@@ -331,10 +407,10 @@ def test_insert_paganin_is_last_sweep(mocker: MockerFixture, tmp_path: Path):
     trans = TransformLayer(comm, repo=repo, save_all=False, out_dir=tmp_path)
     pipeline = trans.transform(pipeline)
 
-    assert len(pipeline) == 5
-    assert pipeline[4].method_name == "save_to_images"
-    assert pipeline[4].task_id == "saveimage_sweep_t3"
-    assert pipeline[4].config_params["subfolder_name"] == "images_sweep_paganin_filter"
+    assert len(pipeline) == 8
+    assert pipeline[7].method_name == "save_to_images"
+    assert pipeline[7].task_id == "saveimage_sweep_t3"
+    assert pipeline[7].config_params["subfolder_name"] == "images_sweep_paganin_filter"
 
 
 def test_insert_denoise_last_after_FBP_sweep(mocker: MockerFixture, tmp_path: Path):
@@ -374,7 +450,7 @@ def test_insert_denoise_last_after_FBP_sweep(mocker: MockerFixture, tmp_path: Pa
     trans = TransformLayer(comm, repo=repo, save_all=False, out_dir=tmp_path)
     pipeline = trans.transform(pipeline)
 
-    assert len(pipeline) == 5
-    assert pipeline[4].method_name == "save_to_images"
-    assert pipeline[4].task_id == "saveimage_sweep_t3"
-    assert pipeline[4].config_params["subfolder_name"] == "images_sweep_median_filter"
+    assert len(pipeline) == 8
+    assert pipeline[7].method_name == "save_to_images"
+    assert pipeline[7].task_id == "saveimage_sweep_t3"
+    assert pipeline[7].config_params["subfolder_name"] == "images_sweep_median_filter"