Updated version to 0.2.4. Moved wedge mask from tme.preprocessor to t…

…me.preprocessing

Updated version to 0.2.4. Moved wedge mask from tme.preprocessor to tme.preprocessing

doc/quickstart/preprocessing/filters.rst

@@ -197,7 +197,7 @@ The high-pass and band-pass filtered images exhibit sharp pronounced peaks and a
 Spectral Whitening
 ^^^^^^^^^^^^^^^^^^
 
-Spectral whitening normalizes each frequency by dividing the amplitude of each frequency by its magnitude. This can bring out subtle features that might be overshadowed otherwise. Spectral Whitening is particulary useful when analyzing large heterogeneous datasets, because the assumptions made are fairly weak. Although neither new nor specific to cryogenic electron microscopy, spectral whitening is a fairly popular [2]_ [3]_ approach and graphically explained `here <https://github.com/ZauggGroup/DeePiCt/blob/main/spectrum_filter/tomo-matcher.svg>`_.
+Spectral whitening normalizes each frequency by dividing the amplitude of each frequency by its magnitude. This can bring out subtle features that might be overshadowed otherwise. Spectral Whitening is particulary useful when analyzing large heterogeneous datasets, because the assumptions made are fairly weak. Although not specific to cryogenic electron microscopy, spectral whitening is a fairly popular [2]_ [3]_ approach and graphically explained `here <https://github.com/ZauggGroup/DeePiCt/blob/main/spectrum_filter/tomo-matcher.svg>`_.
 
 
 .. plot::
@@ -306,14 +306,13 @@ Broadly speaking, |project| distinguishes between continuous, discrete and weigh
 	from tme import Density
 	from tme.preprocessing.tilt_series import WedgeReconstructed
 
-	target = Density.from_file("../../_static/examples/preprocessing_target.png")
-
 	wedge = WedgeReconstructed(
 	    angles = [60,60],
 	    opening_axis = 0,
 	    tilt_axis = 1,
 	    create_continuous_wedge = True,
-	    weight_wedge = False
+	    weight_wedge = False,
+	    reconstruction_filter = "cosine"
 	)
 
 	fig, axs = plt.subplots(nrows=1, ncols=3, figsize=(10, 5), constrained_layout=True)

doc/reference/preprocessing/preprocessor.rst

@@ -17,8 +17,8 @@ Filters
 .. autosummary::
    :toctree: api/
 
-   Preprocessor.wedge_mask
    Preprocessor.continuous_wedge_mask
+   Preprocessor.step_wedge_mask
    Preprocessor.gaussian_filter
    Preprocessor.difference_of_gaussian_filter
    Preprocessor.bandpass_filter

pyproject.toml

@@ -7,7 +7,7 @@ name="pytme"
 authors = [
   { name = "Valentin Maurer", email = "[email protected]" },
 ]
-version="0.2.3"
+version="0.2.4"
 description="Python Template Matching Engine"
 readme="README.md"
 requires-python = ">=3.11"

scripts/match_template.py

@@ -186,8 +186,10 @@ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Com
         ReconstructFromTilt,
     )
 
+    needs_reconstruction = False
     template_filter, target_filter = [], []
     if args.tilt_angles is not None:
+        needs_reconstruction = args.tilt_weighting is not None
         try:
             wedge = Wedge.from_file(args.tilt_angles)
             wedge.weight_type = args.tilt_weighting
@@ -232,11 +234,12 @@ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Com
                     angles=tilt_angles,
                     weight_wedge=args.tilt_weighting == "angle",
                     create_continuous_wedge=create_continuous_wedge,
+                    reconstruction_filter=args.reconstruction_filter,
                 )
                 wedge_target = WedgeReconstructed(
-                    angles=tilt_angles,
+                    angles=(np.abs(np.min(tilt_angles)), np.abs(np.max(tilt_angles))),
                     weight_wedge=False,
-                    create_continuous_wedge=create_continuous_wedge,
+                    create_continuous_wedge=True,
                 )
                 target_filter.append(wedge_target)
 
@@ -330,9 +333,6 @@ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Com
         template_filter.append(whitening_filter)
         target_filter.append(whitening_filter)
 
-    needs_reconstruction = any(
-        [isinstance(t, ReconstructFromTilt) for t in template_filter]
-    )
     if needs_reconstruction and args.reconstruction_filter is None:
         warnings.warn(
             "Consider using a --reconstruction_filter such as 'ramp' to avoid artifacts."

scripts/preprocess.py

@@ -79,6 +79,14 @@ def parse_args():
         action="store_true",
         help="Assumes the template is already centered and omits centering.",
     )
+    modulation_group.add_argument(
+        "--backend",
+        dest="backend",
+        type=str,
+        default=None,
+        choices=be.available_backends(),
+        help="Determines more suitable box size for the given compute backend.",
+    )
     args = parser.parse_args()
     return args
 
@@ -95,11 +103,18 @@ def main():
     if not args.no_centering:
         data, _ = data.centered(0)
 
-    recommended_box = be.compute_convolution_shapes([args.box_size], [1])[1][0]
-    if recommended_box != args.box_size:
-        warnings.warn(
-            f"Consider using --box_size {recommended_box} instead of {args.box_size}."
-        )
+    for name in be.available_backends():
+        be.change_backend(name, device="cpu")
+        box = be.compute_convolution_shapes([args.box_size], [1])[1][0]
+        if box != args.box_size and args.backend is None:
+            print(f"Consider --box_size {box} instead of {args.box_size} for {name}.")
+
+    if args.backend is not None:
+        be.change_backend(name, device="cpu")
+        box = be.compute_convolution_shapes([args.box_size], [1])[1][0]
+        if box != args.box_size:
+            print(f"Changed --box_size from {args.box_size} to {box}.")
+        args.box_size = box
 
     data.pad(
         np.multiply(args.box_size, np.divide(args.sampling_rate, data.sampling_rate)),

scripts/preprocessor_gui.py

@@ -147,9 +147,7 @@ def wedge(
     tilt_step: float = 0,
     opening_axis: int = 0,
     tilt_axis: int = 1,
-    gaussian_sigma: float = 0,
     omit_negative_frequencies: bool = True,
-    extrude_plane: bool = True,
     infinite_plane: bool = True,
 ) -> NDArray:
     template_ft = np.fft.rfftn(template)
@@ -161,9 +159,7 @@ def wedge(
             tilt_axis=tilt_axis,
             opening_axis=opening_axis,
             shape=template.shape,
-            sigma=gaussian_sigma,
             omit_negative_frequencies=omit_negative_frequencies,
-            extrude_plane=extrude_plane,
             infinite_plane=infinite_plane,
         )
         np.multiply(template_ft, wedge_mask, out=template_ft)
@@ -177,7 +173,6 @@ def wedge(
         tilt_step=tilt_step,
         opening_axis=opening_axis,
         shape=template.shape,
-        sigma=gaussian_sigma,
         omit_negative_frequencies=omit_negative_frequencies,
     )
     np.multiply(template_ft, wedge_mask, out=template_ft)
@@ -483,10 +478,9 @@ def wedge_mask(
     tilt_step: float = 0,
     opening_axis: int = 0,
     tilt_axis: int = 2,
-    gaussian_sigma: float = 0,
     omit_negative_frequencies: bool = False,
-    extrude_plane: bool = True,
-    infinite_plane: bool = True,
+    infinite_plane: bool = False,
+    weight_angle: bool = False,
     **kwargs,
 ) -> NDArray:
     if tilt_step <= 0:
@@ -496,22 +490,23 @@ def wedge_mask(
             tilt_axis=tilt_axis,
             opening_axis=opening_axis,
             shape=template.shape,
-            sigma=gaussian_sigma,
             omit_negative_frequencies=omit_negative_frequencies,
-            extrude_plane=extrude_plane,
             infinite_plane=infinite_plane,
         )
         wedge_mask = np.fft.fftshift(wedge_mask)
         return wedge_mask
 
+    weights = None
+    tilt_angles = np.arange(-tilt_start, tilt_stop + tilt_step, tilt_step)
+    if weight_angle:
+        weights = np.cos(np.radians(tilt_angles))
+
     wedge_mask = preprocessor.step_wedge_mask(
-        start_tilt=tilt_start,
-        stop_tilt=tilt_stop,
+        tilt_angles=tilt_angles,
         tilt_axis=tilt_axis,
-        tilt_step=tilt_step,
         opening_axis=opening_axis,
         shape=template.shape,
-        sigma=gaussian_sigma,
+        weights=weights,
         omit_negative_frequencies=omit_negative_frequencies,
     )
 

tests/preprocessing/test_preprocessor.py

@@ -126,19 +126,11 @@ def test_rank_filter(self, rank):
             rank=rank,
         )
 
-    def test_wedge_mask(self):
-        angles = np.zeros((3, 10))
-        angles[0, :] = np.random.rand(angles.shape[1])
-        _ = self.preprocessor.wedge_mask(
-            shape=self.structure_density.data.shape,
-            tilt_angles=angles,
-        )
-
-    @pytest.mark.parametrize("extrude_plane", [False, True])
-    def test_continuous_wedge_mask(self, extrude_plane):
+    @pytest.mark.parametrize("infinite_plane", [False, True])
+    def test_continuous_wedge_mask(self, infinite_plane):
         _ = self.preprocessor.continuous_wedge_mask(
             start_tilt=50,
             stop_tilt=-40,
             shape=(50, 50, 50),
-            extrude_plane=extrude_plane,
+            infinite_plane=infinite_plane,
         )

tme/__version__.py

@@ -1 +1 @@
-__version__ = "0.2.3"
+__version__ = "0.2.4"

tme/matching_utils.py

@@ -645,6 +645,7 @@ def get_rotation_matrices(
         dets = np.linalg.det(ret)
         neg_dets = dets < 0
         ret[neg_dets, :, -1] *= -1
+        ret[0] = np.eye(dim, dtype = ret.dtype)
     return ret
 
 

tme/preprocessing/frequency_filters.py

@@ -99,16 +99,15 @@ def discrete_bandpass(
         grid = be.to_backend_array(grid)
         sampling_rate = be.to_backend_array(sampling_rate)
 
-        lowpass = 0 if lowpass is None else lowpass
-        highpass = 1e10 if highpass is None else highpass
-
         highcut = grid.max()
-        if lowpass > 0:
+        if lowpass is not None:
             highcut = be.max(2 * sampling_rate / lowpass)
-        lowcut = be.max(2 * sampling_rate / highpass)
 
-        bandpass_filter = ((grid <= highcut) & (grid >= lowcut)) * 1.0
+        lowcut = 0
+        if highpass is not None:
+            lowcut = be.max(2 * sampling_rate / highpass)
 
+        bandpass_filter = ((grid <= highcut) & (grid >= lowcut)) * 1.0
         bandpass_filter = shift_fourier(
             data=bandpass_filter, shape_is_real_fourier=shape_is_real_fourier
         )