pytme 0.2.9.post1__cp311-cp311-macosx_15_0_arm64.whl → 0.3.0__cp311-cp311-macosx_15_0_arm64.whl

scripts/preprocess.py

@@ -1,16 +1,17 @@
 #!python3
-""" Preprocessing routines for template matching.
+"""Preprocessing routines for template matching.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 import argparse
 import numpy as np
 
 from tme import Density, Structure
+from tme.cli import print_entry
 from tme.backends import backend as be
-from tme.filters import BandPassFilter
+from tme.filters import BandPassReconstructed
 
 
 def parse_args():
@@ -23,7 +24,6 @@ def parse_args():
     io_group.add_argument(
         "-m",
         "--data",
-        dest="data",
         type=str,
         required=True,
         help="Path to a file in PDB/MMCIF, CCP4/MRC, EM, H5 or a format supported by "
@@ -33,7 +33,6 @@ def parse_args():
     io_group.add_argument(
         "-o",
         "--output",
-        dest="output",
         type=str,
         required=True,
         help="Path the output should be written to.",
@@ -41,38 +40,34 @@ def parse_args():
 
     box_group = parser.add_argument_group("Box")
     box_group.add_argument(
-        "--box_size",
-        dest="box_size",
+        "--box-size",
         type=int,
         required=False,
         help="Box size of the output. Defaults to twice the required box size.",
     )
     box_group.add_argument(
-        "--sampling_rate",
-        dest="sampling_rate",
+        "--sampling-rate",
         type=float,
         required=True,
         help="Sampling rate of the output file.",
     )
     box_group.add_argument(
-        "--input_sampling_rate",
-        dest="input_sampling_rate",
+        "--input-sampling-rate",
         type=float,
         required=False,
-        help="Sampling rate of the input file.",
+        help="Sampling rate of the input file. Defaults to header for volume "
+        "and to --sampling_rate for atomic structures.",
     )
 
     modulation_group = parser.add_argument_group("Modulation")
     modulation_group.add_argument(
-        "--invert_contrast",
-        dest="invert_contrast",
+        "--invert-contrast",
         action="store_true",
         required=False,
         help="Inverts the template contrast.",
     )
     modulation_group.add_argument(
         "--lowpass",
-        dest="lowpass",
         type=float,
         required=False,
         default=None,
@@ -80,14 +75,12 @@ def parse_args():
         "A value of 0 disables the filter.",
     )
     modulation_group.add_argument(
-        "--no_centering",
-        dest="no_centering",
+        "--no-centering",
         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(),
@@ -96,15 +89,13 @@ def parse_args():
 
     alignment_group = parser.add_argument_group("Modulation")
     alignment_group.add_argument(
-        "--align_axis",
-        dest="align_axis",
+        "--align-axis",
         type=int,
         required=False,
         help="Align template to given axis, e.g. 2 for z-axis.",
     )
     alignment_group.add_argument(
-        "--align_eigenvector",
-        dest="align_eigenvector",
+        "--align-eigenvector",
         type=int,
         required=False,
         default=0,
@@ -112,8 +103,7 @@ def parse_args():
         "with numerically largest eigenvalue.",
     )
     alignment_group.add_argument(
-        "--flip_axis",
-        dest="flip_axis",
+        "--flip-axis",
         action="store_true",
         required=False,
         help="Align the template to -axis instead of axis.",
@@ -125,6 +115,7 @@ def parse_args():
 
 def main():
     args = parse_args()
+    print_entry()
 
     try:
         data = Structure.from_file(args.data)
@@ -182,19 +173,18 @@ def main():
     bpf_mask = 1
     lowpass = 2 * args.sampling_rate if args.lowpass is None else args.lowpass
     if args.lowpass != 0:
-        bpf_mask = BandPassFilter(
+        bpf_mask = BandPassReconstructed(
             lowpass=lowpass,
             highpass=None,
             use_gaussian=True,
             return_real_fourier=True,
-            shape_is_real_fourier=False,
             sampling_rate=data.sampling_rate,
         )(shape=data.shape)["data"]
-        bpf_mask = be.to_numpy_array(bpf_mask)
+        bpf_mask = be.to_backend_array(bpf_mask)
 
-    data_ft = np.fft.rfftn(data.data, s=data.shape)
-    data_ft = np.multiply(data_ft, bpf_mask, out=data_ft)
-    data.data = np.fft.irfftn(data_ft, s=data.shape).real
+        data_ft = be.rfftn(be.to_backend_array(data.data), s=data.shape)
+        data_ft = be.multiply(data_ft, bpf_mask, out=data_ft)
+        data.data = be.to_numpy_array(be.irfftn(data_ft, s=data.shape).real)
 
     data = data.resample(args.sampling_rate, method="spline", order=3)
 

scripts/preprocessor_gui.py

@@ -1,9 +1,10 @@
 #!python3
-""" GUI for identifying adequate template matching filter and masks.
+"""
+GUI for identifying suitable masks and analyzing template matchign results.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 import inspect
 import argparse
@@ -18,14 +19,14 @@ from magicgui import widgets
 from numpy.typing import NDArray
 from napari.layers import Image
 from scipy.fft import next_fast_len
-from qtpy.QtWidgets import QFileDialog
 from napari.utils.events import EventedList
+from qtpy.QtWidgets import QFileDialog, QMessageBox
 
 from tme.backends import backend
 from tme.rotations import align_vectors
-from tme.filters import BandPassFilter, CTF
 from tme import Preprocessor, Density, Orientations
 from tme.matching_utils import create_mask, load_pickle
+from tme.filters import BandPassReconstructed, CTFReconstructed
 
 preprocessor = Preprocessor()
 SLIDER_MIN, SLIDER_MAX = 0, 25
@@ -42,17 +43,16 @@ def bandpass_filter(
     hard_edges: bool = False,
     sampling_rate=None,
 ) -> NDArray:
-    bpf = BandPassFilter(
+    bpf = BandPassReconstructed(
         lowpass=lowpass_angstrom,
         highpass=highpass_angstrom,
         sampling_rate=np.max(sampling_rate),
         use_gaussian=not hard_edges,
-        shape_is_real_fourier=True,
         return_real_fourier=True,
     )
     template_ft = np.fft.rfftn(template, s=template.shape)
 
-    mask = bpf(shape=template_ft.shape)["data"]
+    mask = bpf(shape=template.shape)["data"]
     np.multiply(template_ft, mask, out=template_ft)
     return np.fft.irfftn(template_ft, s=template.shape).real
 
@@ -71,15 +71,14 @@ def ctf_filter(
     fast_shape = [next_fast_len(x) for x in np.multiply(template.shape, 2)]
     template_pad = backend.topleft_pad(template, fast_shape)
     template_ft = np.fft.rfftn(template_pad, s=template_pad.shape)
-    ctf = CTF(
-        angles=[0],
+    ctf = CTFReconstructed(
         shape=fast_shape,
         defocus_x=[defocus_angstrom],
         acceleration_voltage=acceleration_voltage * 1e3,
         spherical_aberration=spherical_aberration * 1e7,
         amplitude_contrast=amplitude_contrast,
-        phase_shift=[phase_shift],
-        defocus_angle=[defocus_angle],
+        phase_shift=phase_shift,
+        defocus_angle=defocus_angle,
         sampling_rate=np.max(sampling_rate),
         return_real_fourier=True,
         flip_phase=flip_phase,
@@ -393,7 +392,13 @@ class FilterWidget(widgets.Container):
             metadata = selected_layer_metadata.copy()
             if "filter_parameters" not in metadata:
                 metadata["filter_parameters"] = []
-            metadata["filter_parameters"].append({filter_name: kwargs.copy()})
+
+            payload = {filter_name: kwargs.copy()}
+            if isinstance(metadata["filter_parameters"], dict):
+                metadata["filter_parameters"].update(payload)
+            else:
+                metadata["filter_parameters"].append(payload)
+
             metadata["used_filter"] = filter_name
             new_layer.metadata = metadata
 
@@ -456,6 +461,28 @@ def box_mask(
     )
 
 
+def membrane_mask(
+    template: NDArray,
+    symmetry_axis: int,
+    center_x: float,
+    center_y: float,
+    center_z: float,
+    radius: float,
+    thickness: float = 1,
+    separation: float = 3,
+    sigma_decay: float = 0,
+    **kwargs,
+) -> NDArray:
+    return create_mask(
+        mask_type="membrane",
+        shape=template.shape,
+        radius=radius,
+        thickness=thickness,
+        separation=separation,
+        sigma_decay=sigma_decay,
+    )
+
+
 def tube_mask(
     template: NDArray,
     symmetry_axis: int,
@@ -585,6 +612,7 @@ class MaskWidget(widgets.Container):
             "Ellipsoid": ellipsod_mask,
             "Tube": tube_mask,
             "Box": box_mask,
+            "Membrane": membrane_mask,
             "Wedge": wedge_mask,
             "Threshold": threshold_mask,
             "Lowpass": lowpass_mask,
@@ -1141,6 +1169,8 @@ class MatchingWidget(widgets.Container):
         self.viewer = viewer
         self.dataframes = {}
 
+        self.load_target_checkbox = widgets.CheckBox(text="Load Target", value=False)
+        self.append(self.load_target_checkbox)
         self.import_button = widgets.PushButton(name="Import", text="Import Pickle")
         self.import_button.clicked.connect(self._get_load_path)
 
@@ -1160,6 +1190,26 @@ class MatchingWidget(widgets.Container):
         data = load_pickle(filename)
 
         fname = basename(filename).replace(".pickle", "")
+
+        if self.load_target_checkbox.value:
+            try:
+                target = Density.from_file(data[-1][-1].target)
+                _ = self.viewer.add_image(
+                    data=target.data,
+                    name=f"{fname}_target",
+                    metadata={
+                        "origin": target.origin,
+                        "sampling_rate": target.sampling_rate,
+                    },
+                )
+            except Exception as e:
+                msg = QMessageBox(self.native)
+                msg.setIcon(QMessageBox.Warning)
+                msg.setWindowTitle("Loading Error")
+                msg.setText(str(e))
+                msg.setStandardButtons(QMessageBox.Ok)
+                msg.exec_()
+
         if data[0].ndim == data[2].ndim:
             metadata = {"origin": data[-1][1], "sampling_rate": data[-1][2]}
             _ = self.viewer.add_image(
@@ -1175,11 +1225,8 @@ class MatchingWidget(widgets.Container):
                 metadata=metadata,
             )
             return None
-            detail = np.zeros_like(data[2])
-        else:
-            detail = data[3]
 
-        point_properties = {"score": data[2], "detail": detail}
+        point_properties = {"score": data[2], "detail": data[3]}
         point_properties["score_scaled"] = np.log1p(
             point_properties["score"] - point_properties["score"].min()
         )
@@ -1192,8 +1239,26 @@ class MatchingWidget(widgets.Container):
         )
 
 
+class CustomNapariViewer(napari.Viewer):
+    """
+    Custom viewer to ensure 3D image layers are by default shown as xy projection.
+    """
+
+    def add_image(self, data, **kwargs):
+        viewer_ndim = len(self.dims.order)
+        layer = super().add_image(data, **kwargs)
+
+        try:
+            # Set to xy view the first time data is opened
+            if viewer_ndim != 3 and data.ndim == 3:
+                self.dims.order = (2, 0, 1)
+        except Exception:
+            pass
+        return layer
+
+
 def main():
-    viewer = napari.Viewer()
+    viewer = CustomNapariViewer()
 
     filter_widget = FilterWidget(preprocessor, viewer)
     mask_widget = MaskWidget(viewer)
@@ -1216,7 +1281,8 @@ def main():
 
 def parse_args():
     parser = argparse.ArgumentParser(
-        description="GUI for preparing and analyzing template matching runs."
+        description="GUI for preparing and analyzing template matching runs.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
     args = parser.parse_args()
     return args