Add ignore tomo mask option (#209)

* update version number

* add --ignore_tomogram_mask option to entry points and extraction code

* add extraction code test

* also make sure warning is logged

* actually call the expected code branch...

* exclude possible fail state from coverage

* Update tests/test_tmjob.py

Co-authored-by: Marten Chaillet <[email protected]>

---------

Co-authored-by: Marten Chaillet <[email protected]>

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pytom-match-pick"
-version = "0.7.1"
+version = "0.7.2"
 description = "PyTOM's GPU template matching module as an independent package"
 readme = "README.md"
 license = {file = "LICENSE"}

src/pytom_tm/entry_points.py

@@ -364,6 +364,14 @@ def estimate_roc(argv=None):
         action=ParseLogging,
         help="Can be set to `info` or `debug`",
     )
+    parser.add_argument(
+        "--ignore_tomogram_mask",
+        action="store_true",
+        default=False,
+        required=False,
+        help="Flag to ignore the TM job tomogram mask. "
+        "Useful if the scores mrc looks reasonable, but this finds 0 particles",
+    )
 
     # ---8<--- [end:estimate_roc_usage]
 
@@ -378,6 +386,7 @@ def estimate_roc(argv=None):
         args.number_of_particles,
         cut_off=0,
         create_plot=False,
+        ignore_tomogram_mask=args.ignore_tomogram_mask,
     )
 
     score_volume = read_mrc(
@@ -433,6 +442,14 @@ def extract_candidates(argv=None):
         "be used to extract annotations only within a specific cellular region."
         "If the job was run with a tomogram mask, this file will be used instead of the job mask",
     )
+    parser.add_argument(
+        "--ignore_tomogram_mask",
+        action="store_true",
+        default=False,
+        required=False,
+        help="Flag to ignore the input and TM job tomogram mask. "
+        "Useful if the scores mrc looks reasonable, but this finds 0 particles to extract",
+    )
     parser.add_argument(
         "-n",
         "--number-of-particles",
@@ -525,6 +542,7 @@ def extract_candidates(argv=None):
         tophat_filter=args.tophat_filter,
         tophat_connectivity=args.tophat_connectivity,
         relion5_compat=args.relion5_compat,
+        ignore_tomogram_mask=args.ignore_tomogram_mask,
     )
 
     # write out as a RELION type starfile

src/pytom_tm/extract.py

@@ -148,6 +148,7 @@ def extract_particles(
     create_plot: bool = True,
     tophat_connectivity: int = 1,
     relion5_compat: bool = False,
+    ignore_tomogram_mask: bool = False,
 ) -> tuple[pd.DataFrame, list[float, ...]]:
     """
     Parameters
@@ -174,6 +175,10 @@ def extract_particles(
     relion5_compat: bool, default False
         relion5 compatibility writes coordinates relative to center and in Angstrom
         center definition should be: tomo_shape / 2 - 1
+    ignore_tomogram_mask: bool, default False
+        Debug option to force the code to ignore job.tomogram_mask and input mask.
+        Allows for re-exctraction without rerunning the TM job
+        (assuming the scores volume seems reasonable)
 
     Returns
     -------
@@ -203,7 +208,9 @@ def extract_particles(
 
     # apply tomogram mask if provided
     tomogram_mask = None
-    if tomogram_mask_path is not None:
+    if ignore_tomogram_mask:
+        logging.warn("Ignoring tomogram mask")
+    elif tomogram_mask_path is not None:
         tomogram_mask = read_mrc(tomogram_mask_path)
     elif job.tomogram_mask is not None:
         tomogram_mask = read_mrc(job.tomogram_mask)

tests/test_tmjob.py

@@ -643,6 +643,28 @@ def test_extraction(self):
             msg="Length of returned list should be 0 after applying mask where the "
             "object is not in the region of interest.",
         )
+        # test if all masks are ignored if ignore_tomogram_mask=True and that a warning is raised
+        with self.assertLogs(level="WARNING") as cm:
+            df, scores = extract_particles(
+                job,
+                5,
+                100,
+                tomogram_mask_path=TEST_EXTRACTION_MASK_OUTSIDE,
+                create_plot=False,
+                ignore_tomogram_mask=True,
+            )
+        # Test if expected warning is logged
+        for o in cm.output:
+            if "Ignoring tomogram mask" in o:
+                break
+        else:
+            # break is not hit
+            self.fail("expected warning is not logged")  # pragma: no cover
+        self.assertNotEqual(
+            len(scores),
+            0,
+            msg="We would expect some annotations if all tomogram masks are ignored",
+        )
 
         # test mask that covers the particle and should override the one now attached to the job
         df, scores = extract_particles(