pytme 0.3b0.post1__cp311-cp311-macosx_15_0_arm64.whl → 0.3.1.dev20250731__cp311-cp311-macosx_15_0_arm64.whl

scripts/postprocess.py

@@ -17,7 +17,7 @@ from scipy.special import erfcinv
 
 from tme import Density, Structure, Orientations
 from tme.cli import sanitize_name, print_block, print_entry
-from tme.matching_utils import load_pickle, centered_mask, write_pickle
+from tme.matching_utils import load_pickle, center_slice, write_pickle
 from tme.matching_optimization import create_score_object, optimize_match
 from tme.rotations import euler_to_rotationmatrix, euler_from_rotationmatrix
 from tme.analyzer import (
@@ -87,6 +87,11 @@ def parse_args():
         help="Output prefix. Defaults to basename of first input. Extension is "
         "added with respect to chosen output format.",
     )
+    output_group.add_argument(
+        "--angles-clockwise",
+        action="store_true",
+        help="Report Euler angles in clockwise format expected by RELION.",
+    )
     output_group.add_argument(
         "--output-format",
         choices=[
@@ -112,7 +117,7 @@ def parse_args():
     peak_group.add_argument(
         "--peak-caller",
         choices=list(PEAK_CALLERS.keys()),
-        default="PeakCallerScipy",
+        default="PeakCallerMaximumFilter",
         help="Peak caller for local maxima identification.",
     )
     peak_group.add_argument(
@@ -183,7 +188,7 @@ def parse_args():
     )
     additional_group.add_argument(
         "--n-false-positives",
-        type=int,
+        type=float,
         default=None,
         required=False,
         help="Number of accepted false-positives picks to determine minimum score.",
@@ -313,11 +318,7 @@ def normalize_input(foregrounds: Tuple[str], backgrounds: Tuple[str]) -> Tuple:
         data = load_matching_output(foreground)
         scores, _, rotations, rotation_mapping, *_ = data
 
-        # We could normalize to unit sdev, but that might lead to unexpected
-        # results for flat background distributions
-        scores -= scores.mean()
         indices = tuple(slice(0, x) for x in scores.shape)
-
         indices_update = scores > scores_out[indices]
         scores_out[indices][indices_update] = scores[indices_update]
 
@@ -364,9 +365,8 @@ def normalize_input(foregrounds: Tuple[str], backgrounds: Tuple[str]) -> Tuple:
     scores_norm = np.full(out_shape_norm, fill_value=0, dtype=np.float32)
     for background in backgrounds:
         data_norm = load_matching_output(background)
+        scores, _, rotations, rotation_mapping, *_ = data_norm
 
-        scores = data_norm[0]
-        scores -= scores.mean()
         indices = tuple(slice(0, x) for x in scores.shape)
         indices_update = scores > scores_norm[indices]
         scores_norm[indices][indices_update] = scores[indices_update]
@@ -375,8 +375,10 @@ def normalize_input(foregrounds: Tuple[str], backgrounds: Tuple[str]) -> Tuple:
     update = tuple(slice(0, int(x)) for x in np.minimum(out_shape, scores.shape))
     scores_out = np.full(out_shape, fill_value=0, dtype=np.float32)
     scores_out[update] = data[0][update] - scores_norm[update]
-    scores_out[update] = np.divide(scores_out[update], 1 - scores_norm[update])
     scores_out = np.fmax(scores_out, 0, out=scores_out)
+    scores_out[update] += scores_norm[update].mean()
+
+    # scores_out[update] = np.divide(scores_out[update], 1 - scores_norm[update])
     data[0] = scores_out
 
     fg, bg = simple_stats(data[0]), simple_stats(scores_norm)
@@ -478,15 +480,21 @@ def main():
     if orientations is None:
         translations, rotations, scores, details = [], [], [], []
 
-        # Data processed by normalize_input is guaranteed to have this shape
-        scores, offset, rotation_array, rotation_mapping, meta = data
+        var = None
+        # Data processed by normalize_input is guaranteed to have this shape)
+        scores, _, rotation_array, rotation_mapping, *_ = data
+        if len(data) == 6:
+            scores, _, rotation_array, rotation_mapping, var, *_ = data
 
         cropped_shape = np.subtract(
             scores.shape, np.multiply(args.min_boundary_distance, 2)
         ).astype(int)
 
         if args.min_boundary_distance > 0:
-            scores = centered_mask(scores, new_shape=cropped_shape)
+            _scores = np.zeros_like(scores)
+            subset = center_slice(scores.shape, cropped_shape)
+            _scores[subset] = scores[subset]
+            scores = _scores
 
         if args.n_false_positives is not None:
             # Rickgauer et al. 2017
@@ -499,17 +507,20 @@ def main():
             )
             args.n_false_positives = max(args.n_false_positives, 1)
             n_correlations = np.size(scores[cropped_slice]) * len(rotation_mapping)
+            std = np.std(scores[cropped_slice])
+            if var is not None:
+                std = np.asarray(np.sqrt(var)).reshape(())
+
             minimum_score = np.multiply(
                 erfcinv(2 * args.n_false_positives / n_correlations),
-                np.sqrt(2) * np.std(scores[cropped_slice]),
+                np.sqrt(2) * std,
             )
-            print(f"Determined minimum score cutoff: {minimum_score}.")
-            minimum_score = max(minimum_score, 0)
-            args.min_score = minimum_score
+            print(f"Determined cutoff --min-score {minimum_score}.")
+            args.min_score = max(minimum_score, 0)
 
         args.batch_dims = None
-        if hasattr(cli_args, "target_batch"):
-            args.batch_dims = cli_args.target_batch
+        if hasattr(cli_args, "batch_dims"):
+            args.batch_dims = cli_args.batch_dims
 
         peak_caller_kwargs = {
             "shape": scores.shape,
@@ -517,8 +528,8 @@ def main():
             "min_distance": args.min_distance,
             "min_boundary_distance": args.min_boundary_distance,
             "batch_dims": args.batch_dims,
-            "minimum_score": args.min_score,
-            "maximum_score": args.max_score,
+            "min_score": args.min_score,
+            "max_score": args.max_score,
         }
 
         peak_caller = PEAK_CALLERS[args.peak_caller](**peak_caller_kwargs)
@@ -611,6 +622,9 @@ def main():
             orientations.rotations[index] = angles
             orientations.scores[index] = score * -1
 
+    if args.angles_clockwise:
+        orientations.rotations *= -1
+
     if args.output_format in ("orientations", "relion4", "relion5"):
         file_format, extension = "text", "tsv"
 

scripts/preprocessor_gui.py

@@ -19,13 +19,13 @@ 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.backends import backend as be
 from tme.rotations import align_vectors
-from tme import Preprocessor, Density, Orientations
 from tme.matching_utils import create_mask, load_pickle
+from tme import Preprocessor, Density, Orientations
 from tme.filters import BandPassReconstructed, CTFReconstructed
 
 preprocessor = Preprocessor()
@@ -69,7 +69,7 @@ def ctf_filter(
     flip_phase: bool = False,
 ) -> NDArray:
     fast_shape = [next_fast_len(x) for x in np.multiply(template.shape, 2)]
-    template_pad = backend.topleft_pad(template, fast_shape)
+    template_pad = be.topleft_pad(template, fast_shape)
     template_ft = np.fft.rfftn(template_pad, s=template_pad.shape)
     ctf = CTFReconstructed(
         shape=fast_shape,
@@ -85,7 +85,7 @@ def ctf_filter(
     )
     np.multiply(template_ft, ctf()["data"], out=template_ft)
     template_pad = np.fft.irfftn(template_ft, s=template_pad.shape).real
-    template = backend.topleft_pad(template_pad, template.shape)
+    template = be.topleft_pad(template_pad, template.shape)
     return template
 
 
@@ -392,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
 
@@ -450,7 +456,30 @@ def box_mask(
         mask_type="box",
         shape=template.shape,
         center=(center_x, center_y, center_z),
-        height=(height_x, height_y, height_z),
+        size=(height_x, height_y, height_z),
+        sigma_decay=sigma_decay,
+    )
+
+
+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(
+        center=(center_x, center_y, center_z),
+        mask_type="membrane",
+        shape=template.shape,
+        radius=radius,
+        thickness=thickness,
+        separation=separation,
         sigma_decay=sigma_decay,
     )
 
@@ -471,7 +500,7 @@ def tube_mask(
         mask_type="tube",
         shape=template.shape,
         symmetry_axis=symmetry_axis,
-        base_center=(center_x, center_y, center_z),
+        center=(center_x, center_y, center_z),
         inner_radius=inner_radius,
         outer_radius=outer_radius,
         height=height,
@@ -584,6 +613,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,
@@ -817,7 +847,7 @@ class AlignmentWidget(widgets.Container):
         principal_eigenvector = eigenvectors[:, np.argmax(eigenvalues)]
 
         rotation_matrix = align_vectors(principal_eigenvector, alignment_axis)
-        rotated_data, _ = backend.rigid_transform(
+        rotated_data, _ = be.rigid_transform(
             arr=active_layer.data,
             rotation_matrix=rotation_matrix,
             use_geometric_center=False,
@@ -953,7 +983,6 @@ class PointCloudWidget(widgets.Container):
             if not isinstance(layer, napari.layers.Points):
                 continue
 
-            layer.face_color = "white"
             if event == "Label":
                 if len(layer.properties.get("detail", ())) == 0:
                     continue
@@ -970,9 +999,7 @@ class PointCloudWidget(widgets.Container):
                 layer.face_color = "score_scaled"
                 layer.face_colormap = "turbo"
                 layer.face_color_mode = "colormap"
-
             layer.refresh_colors()
-
         return None
 
     def _set_positive(self, event):
@@ -1140,9 +1167,18 @@ class MatchingWidget(widgets.Container):
         self.viewer = viewer
         self.dataframes = {}
 
+        option_container = widgets.Container(layout="horizontal")
+        self.load_target_checkbox = widgets.CheckBox(text="Load Target", value=False)
+        self.load_rotations_checkbox = widgets.CheckBox(
+            text="Load Rotations", value=False
+        )
+        option_container.append(self.load_target_checkbox)
+        option_container.append(self.load_rotations_checkbox)
+
         self.import_button = widgets.PushButton(name="Import", text="Import Pickle")
         self.import_button.clicked.connect(self._get_load_path)
 
+        self.append(option_container)
         self.append(self.import_button)
 
     def _get_load_path(self, event):
@@ -1150,7 +1186,7 @@ class MatchingWidget(widgets.Container):
             self.native,
             "Open Pickle File...",
             "",
-            "Pickle Files (*.pickle);;All Files (*)",
+            "Pickle Files (*.pickle *pickle.gz);;All Files (*)",
         )
         if filename:
             self._load_data(filename)
@@ -1159,14 +1195,35 @@ 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(
-                data=data[2],
-                name=f"{fname}_rotations",
-                colormap="orange",
-                metadata=metadata,
-            )
+            if self.load_rotations_checkbox.value:
+                _ = self.viewer.add_image(
+                    data=data[2],
+                    name=f"{fname}_rotations",
+                    colormap="orange",
+                    metadata=metadata,
+                )
             _ = self.viewer.add_image(
                 data=data[0],
                 name=f"{fname}_scores",
@@ -1174,11 +1231,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()
         )
@@ -1191,8 +1245,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)