Merge pull request #668 from DiamondLightSource/remove_redundant_in_sweep

adding function to remove redundant methods in sweep pipeline

Author: dkazanc

httomo/transform_layer.py

@@ -36,11 +36,13 @@ def __init__(
         self._out_dir = out_dir if out_dir is not None else httomo.globals.run_out_dir
 
     def transform(self, pipeline: Pipeline) -> Pipeline:
+        pipeline_is_sweep = _check_if_pipeline_has_a_sweep(pipeline)
+
         pipeline = self.insert_data_reducer(pipeline)
+        if pipeline_is_sweep:
+            pipeline = self.remove_redundant_method_in_sweep(pipeline)
         pipeline = self.insert_data_checker(pipeline)
 
-        pipeline_is_sweep = _check_if_pipeline_has_a_sweep(pipeline)
-
         if pipeline_is_sweep:
             pipeline = self.insert_save_images_after_sweep(pipeline)
         else:
@@ -95,30 +97,18 @@ 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
+            exceptions_methods = [
+                "data_reducer",
+                "data_checker",
+                "calculate_stats",
+                "rescale_to_int",
+                "save_to_images",
+            ]
             if (
-                m.method_name
-                not in [
-                    "data_reducer",
-                    "data_checker",
-                    "calculate_stats",
-                    "rescale_to_int",
-                ]
+                m.method_name not in exceptions_methods
                 and "rotation" not in m.module_path
                 and index < len(pipeline._methods) - 1
             ):
@@ -152,7 +142,6 @@ def insert_data_checker(self, pipeline: Pipeline) -> Pipeline:
                             data_to_method_name=m.method_name,
                         ),
                     )
-
         return Pipeline(loader, methods)
 
     def insert_save_images_after_sweep(self, pipeline: Pipeline) -> Pipeline:
@@ -177,3 +166,13 @@ def insert_save_images_after_sweep(self, pipeline: Pipeline) -> Pipeline:
                 )
                 sweep_before = True
         return Pipeline(loader, methods)
+
+    def remove_redundant_method_in_sweep(self, pipeline: Pipeline) -> Pipeline:
+        """Remove "redundant" methods in the sweep pipeline that were inserted by the user."""
+        redundant_methods = ["calculate_stats", "rescale_to_int", "save_to_images"]
+        loader = pipeline.loader
+        methods = []
+        for m in pipeline:
+            if m.method_name not in redundant_methods:
+                methods.append(m)
+        return Pipeline(loader, methods)

tests/test_transform_layer.py

@@ -186,6 +186,70 @@ def test_insert_image_save_after_sweep(mocker: MockerFixture, tmp_path: Path):
     assert pipeline[3].config_params["axis"] == 1
 
 
+def test_remove_redundant_method_in_sweep(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="dark_flat_field_correction",
+                module_path="httomolibgpu.prep.normalize",
+                save_result=False,
+                task_id="t2",
+            ),
+            make_test_method(
+                mocker,
+                method_name="FBP3d_tomobar",
+                module_path="httomolibgpu.recon.algorithm",
+                save_result=False,
+                sweep=True,
+                task_id="t3",
+            ),
+            make_test_method(
+                mocker,
+                method_name="calculate_stats",
+                module_path="httomo.methods",
+                save_result=False,
+                task_id="t4",
+            ),
+            make_test_method(
+                mocker,
+                method_name="rescale_to_int",
+                module_path="httomolib.misc.rescale",
+                save_result=False,
+                task_id="t5",
+            ),
+            make_test_method(
+                mocker,
+                method_name="save_to_images",
+                module_path="httomolib.misc.images",
+                save_result=False,
+                task_id="t6",
+            ),
+        ],
+    )
+    trans = TransformLayer(comm, repo=repo, save_all=False, out_dir=tmp_path)
+    pipeline = trans.remove_redundant_method_in_sweep(pipeline)
+
+    assert len(pipeline) == 3
+    assert pipeline[0].method_name == "remove_outlier"
+    assert pipeline[1].method_name == "dark_flat_field_correction"
+    assert pipeline[2].method_name == "FBP3d_tomobar"
+
+
 def test_insert_data_checker(mocker: MockerFixture, tmp_path: Path):
     comm = MPI.COMM_SELF
     repo = make_mock_repo(mocker)
@@ -251,15 +315,14 @@ def test_insert_data_checker(mocker: MockerFixture, tmp_path: Path):
     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"
+    assert len(pipeline) == 10
+    assert pipeline[1].method_name == "data_checker"
+    assert pipeline[3].method_name == "normalize"
+    assert pipeline[4].method_name == "data_checker"
+    assert pipeline[5].method_name == "FBP3d_tomobar"
+    assert pipeline[6].method_name == "data_checker"
+    assert pipeline[8].method_name == "rescale_to_int"
+    assert pipeline[9].method_name == "save_to_images"
 
 
 def test_insert_image_save_after_sweep2(mocker: MockerFixture, tmp_path: Path):
@@ -361,13 +424,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) == 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"
+    assert len(pipeline) == 10
+    assert pipeline[6].method_name == "save_to_images"
+    assert pipeline[6].task_id == "saveimage_sweep_t3"
+    assert pipeline[6].config_params["subfolder_name"] == "images_sweep_paganin_filter"
+    assert pipeline[8].method_name == "FBP3d_tomobar"
+    assert pipeline[9].task_id == "saveimage_sweep_t4"
+    assert pipeline[9].config_params["subfolder_name"] == "images_sweep_FBP3d_tomobar"
 
 
 def test_insert_paganin_is_last_sweep(mocker: MockerFixture, tmp_path: Path):
@@ -407,10 +470,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) == 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"
+    assert len(pipeline) == 7
+    assert pipeline[6].method_name == "save_to_images"
+    assert pipeline[6].task_id == "saveimage_sweep_t3"
+    assert pipeline[6].config_params["subfolder_name"] == "images_sweep_paganin_filter"
 
 
 def test_insert_denoise_last_after_FBP_sweep(mocker: MockerFixture, tmp_path: Path):
@@ -450,7 +513,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) == 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"
+    assert len(pipeline) == 7
+    assert pipeline[6].method_name == "save_to_images"
+    assert pipeline[6].task_id == "saveimage_sweep_t3"
+    assert pipeline[6].config_params["subfolder_name"] == "images_sweep_median_filter"