FIX: fixed issue with getting the masks from the probe positions. The simulated probe used for masks was being created incorrectly.

Author: hruth

src/pyxalign/api/options/projections.py

@@ -88,7 +88,7 @@ class VolumeWidthOptions:
 
 
 @dataclasses.dataclass
-class SimulatedProbe:
+class SimulatedProbeOptions:
     """
     Parameters for creating a gaussian probe
     """
@@ -114,7 +114,7 @@ class ProbePositionMaskOptions:
     use the probe in the Projections object.
     """
 
-    probe: SimulatedProbe = field(default_factory=SimulatedProbe)
+    probe: SimulatedProbeOptions = field(default_factory=SimulatedProbeOptions)
 
 
 @dataclasses.dataclass

src/pyxalign/data_structures/projections.py

@@ -35,7 +35,7 @@
 import pyxalign.gpu_utils as gpu_utils
 from pyxalign.gpu_wrapper import device_handling_wrapper
 from pyxalign.data_structures.volume import Volume
-from pyxalign.mask import build_masks_from_threshold
+from pyxalign.mask import build_masks_from_threshold, get_simulated_probe_for_masks
 from pyxalign.io.utils import load_list_of_arrays
 from pyxalign.io.save import save_generic_data_structure_to_h5
 
@@ -381,9 +381,15 @@ def center_projections(self):
 
     @timer()
     def get_masks_from_probe_positions(self):
+        if self.options.mask_from_positions.use_simulated_probe:
+            probe = get_simulated_probe_for_masks(
+                self.probe, self.options.mask_from_positions.probe
+            )
+        else:
+            probe = self.probe
         self.masks = build_masks_from_threshold(
             self.data.shape,
-            self.probe,
+            probe,
             self.probe_positions.data,
             self.options.mask_from_positions.threshold,
         )

src/pyxalign/interactions/mask.py

@@ -15,16 +15,12 @@
     QPushButton,
 )
 
-from pyxalign.api.enums import RoundType
 import pyxalign.data_structures.projections as p
 from pyxalign.interactions.utils.loading_display_tools import loading_bar_wrapper
 from pyxalign.interactions.utils.misc import switch_to_matplotlib_qt_backend
-from pyxalign.mask import place_patches_fourier_batch
+from pyxalign.mask import get_simulated_probe_for_masks, place_patches_fourier_batch
 from pyxalign.interactions.viewers.base import IndexSelectorWidget
 from pyxalign.mask import clip_masks
-from pyxalign.model_functions import symmetric_gaussian_2d
-from pyxalign.transformations.helpers import round_to_divisor
-from pyxalign.api import constants
 
 """
 Interactive mask threshold selector based on pyqtgraph and the shared
@@ -153,23 +149,14 @@ def __init__(
         self.projections = projections
         self.options = self.projections.options.mask_from_positions
 
-        # use simulated probe if specified by options; this typically
-        # gives better results
         if self.options.use_simulated_probe:
-            shape = self.projections.probe.shape
-            probe_width = round_to_divisor(
-                shape[0] * self.options.probe.fractional_width,
-                round_type=RoundType.NEAREST,
-                divisor=constants.divisor
-            )
-            probe = symmetric_gaussian_2d(shape, amplitude=1, sigma=probe_width)
+            probe = get_simulated_probe_for_masks(self.projections.probe, self.options.probe)
         else:
             probe = self.projections.probe
-
         # Precompute masks (floating-point values)
-        load_bar_func_wrapper = loading_bar_wrapper("Initializing masks...")(
-            place_patches_fourier_batch
-        )
+        load_bar_func_wrapper = loading_bar_wrapper(
+            "Initializing masks...", block_all_windows=True
+        )(place_patches_fourier_batch)
         masks = load_bar_func_wrapper(
             self.projections.data.shape,
             probe,

src/pyxalign/mask.py

@@ -11,13 +11,15 @@
 from tqdm import tqdm
 from contextlib import nullcontext
 from pyxalign import gpu_utils
+from pyxalign.api.enums import RoundType
 from pyxalign.api.options.device import DeviceOptions
+from pyxalign.api.options.projections import SimulatedProbeOptions
 from pyxalign.gpu_wrapper import device_handling_wrapper
 
 # from pyxalign.interactions.mask import ThresholdSelector, illum_map_threshold_plotter
-from pyxalign.transformations.helpers import is_array_real
+from pyxalign.model_functions import symmetric_gaussian_2d
+from pyxalign.transformations.helpers import is_array_real, round_to_divisor
 from IPython.display import display
-from PyQt5.QtWidgets import QApplication
 import plotly.graph_objects as go
 from plotly.subplots import make_subplots
 
@@ -26,8 +28,6 @@
 from pyxalign.timing.timer_utils import timer, InlineTimer
 from pyxalign.api.types import ArrayType, r_type
 
-from PyQt5.QtWidgets import QWidget
-
 
 @memory_releasing_error_handler
 @timer()
@@ -387,45 +387,16 @@ def build_masks_from_threshold(
     return clip_masks(masks, threshold)
 
 
-# class IlluminationMapMaskBuilder:
-#     """
-#     Class for building mask from the illumination map.
-#     """
-
-#     def get_mask_base(
-#         self,
-#         probe: np.ndarray,
-#         positions: list[np.ndarray],
-#         projections: np.ndarray,
-#         use_fourier: bool = True,
-#     ):
-#         # The base for building the mask is the illumination map
-#         if use_fourier:
-#             self.masks = place_patches_fourier_batch(projections.shape, probe, positions)
-#         else:
-#             for i in range(len(positions)):
-#                 self.masks = np.zeros_like(projections, dtype=r_type)
-#                 get_illumination_map(self.masks[i], probe, positions[i])
-
-#     def set_mask_threshold_interactively(self, projections: np.ndarray) -> float:
-#         # temporary bugfix: all windows need to be closed or else app.exec_() will 
-#         # hang indefinitely. I am putting this temporary solution (which I don't like
-#         # very much) in place, because any changes will be overwritten once merged with
-#         # interactive_pma_gui anyway.
-#         app = QApplication.instance() or QApplication([])
-#         app.closeAllWindows()
-
-#         # Use interactivity to decide mask threshold"
-#         self.threshold_selector = illum_map_threshold_plotter(
-#             self.masks, projections, init_thresh=0.01
-#         )
-#         self.threshold_selector.show()
-
-#         app.exec_()
-#         threshold = self.threshold_selector.threshold
-#         return threshold
-
-#     def clip_masks(self, thresh: Optional[float] = None):
-#         clip_idx = self.masks > thresh
-#         self.masks[:] = 0
-#         self.masks[clip_idx] = 1
+def get_simulated_probe_for_masks(
+    probe: np.ndarray,
+    simulated_probe_options: SimulatedProbeOptions,
+):
+    # simulated probe typically gives better results
+    shape = probe.shape
+    probe_width = round_to_divisor(
+        shape[0] * simulated_probe_options.fractional_width,
+        round_type=RoundType.NEAREST,
+        divisor=2,
+    )
+    probe = symmetric_gaussian_2d(shape, amplitude=1, sigma=probe_width)
+    return probe