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

pytme-0.3.1.dev20250731.data/scripts/estimate_ram_usage.py

@@ -0,0 +1,97 @@
+#!python
+""" Estimate RAM requirements for template matching jobs.
+
+    Copyright (c) 2023 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+import numpy as np
+import argparse
+from tme import Density
+from tme.matching_utils import estimate_ram_usage
+from tme.matching_exhaustive import MATCHING_EXHAUSTIVE_REGISTER
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Estimate RAM usage for template matching."
+    )
+    parser.add_argument(
+        "-m",
+        "--target",
+        dest="target",
+        type=str,
+        required=True,
+        help="Path to a target in CCP4/MRC format.",
+    )
+    parser.add_argument(
+        "-i",
+        "--template",
+        dest="template",
+        type=str,
+        required=True,
+        help="Path to a template in PDB/MMCIF or CCP4/MRC format.",
+    )
+    parser.add_argument(
+        "--matching_method",
+        required=False,
+        default=None,
+        help="Analyzer method to use.",
+    )
+    parser.add_argument(
+        "-s",
+        dest="score",
+        type=str,
+        default="FLCSphericalMask",
+        help="Template matching scoring function.",
+        choices=MATCHING_EXHAUSTIVE_REGISTER.keys(),
+    )
+    parser.add_argument(
+        "--ncores", type=int, help="Number of cores for parallelization.", required=True
+    )
+    parser.add_argument(
+        "--no_edge_padding",
+        dest="no_edge_padding",
+        action="store_true",
+        default=False,
+        help="Whether to pad the edges of the target. This is useful, if the target"
+        " has a well defined bounding box, e.g. a density map.",
+    )
+    parser.add_argument(
+        "--no_fourier_padding",
+        dest="no_fourier_padding",
+        action="store_true",
+        default=False,
+        help="Whether input arrays should be zero-padded to the full convolution shape"
+        " for numerical stability. When working with very large targets such as"
+        " tomograms it is safe to use this flag and benefit from the performance gain.",
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    target = Density.from_file(args.target)
+    template = Density.from_file(args.template)
+
+    target_box = target.shape
+    if not args.no_edge_padding:
+        target_box = np.add(target_box, template.shape)
+
+    template_box = template.shape
+    if args.no_fourier_padding:
+        template_box = np.ones(len(template_box), dtype=int)
+
+    result = estimate_ram_usage(
+        shape1=target_box,
+        shape2=template_box,
+        matching_method=args.score,
+        ncores=args.ncores,
+        analyzer_method="MaxScoreOverRotations",
+    )
+    print(result)
+
+
+if __name__ == "__main__":
+    main()

{pytme-0.3b0.post1.data → pytme-0.3.1.dev20250731.data}/scripts/match_template.py

@@ -12,8 +12,8 @@ from sys import exit
 from time import time
 from typing import Tuple
 from copy import deepcopy
-from os.path import exists
 from tempfile import gettempdir
+from os.path import exists, abspath
 
 import numpy as np
 
@@ -359,8 +359,8 @@ def parse_args():
         "--invert-target-contrast",
         action="store_true",
         default=False,
-        help="Invert the target's contrast for cases where templates to-be-matched have "
-        "negative values, e.g. tomograms.",
+        help="Invert the target contrast. Useful for matching on tomograms if the "
+        "template has not been inverted.",
     )
     io_group.add_argument(
         "--scramble-phases",
@@ -576,7 +576,7 @@ def parse_args():
         "'angles', or a single column file without header. Exposure will be taken from "
         "the input file , if you are using a tab-separated file, the column names "
         "'angles' and 'weights' need to be present. It is also possible to specify a "
-        "continuous wedge mask using e.g., -50,45.",
+        "continuous wedge mask using e.g., 50,45.",
     )
     filter_group.add_argument(
         "--tilt-weighting",
@@ -685,13 +685,6 @@ def parse_args():
         help="Useful if the target does not have a well-defined bounding box. Will be "
         "activated automatically if splitting is required to avoid boundary artifacts.",
     )
-    performance_group.add_argument(
-        "--pad-filter",
-        action="store_true",
-        default=False,
-        help="Pad the template filter to the shape of the target. Useful for fast "
-        "oscilating filters to avoid aliasing effects.",
-    )
     performance_group.add_argument(
         "--interpolation-order",
         required=False,
@@ -700,13 +693,6 @@ def parse_args():
         help="Spline interpolation used for rotations. Defaults to 3, and 1 for jax "
         "and pytorch backends.",
     )
-    performance_group.add_argument(
-        "--use-mixed-precision",
-        action="store_true",
-        default=False,
-        help="Use float16 for real values operations where possible. Not supported "
-        "for jax backend.",
-    )
     performance_group.add_argument(
         "--use-memmap",
         action="store_true",
@@ -742,6 +728,7 @@ def parse_args():
         args.interpolation_order = 3
         if args.backend in ("jax", "pytorch"):
             args.interpolation_order = 1
+            args.reconstruction_interpolation_order = 1
 
     if args.interpolation_order < 0:
         args.interpolation_order = None
@@ -779,6 +766,14 @@ def parse_args():
         )
         args.orientations = orientations
 
+    args.target = abspath(args.target)
+    if args.target_mask is not None:
+        args.target_mask = abspath(args.target_mask)
+
+    args.template = abspath(args.template)
+    if args.template_mask is not None:
+        args.template_mask = abspath(args.template_mask)
+
     return args
 
 
@@ -796,15 +791,23 @@ def main():
             sampling_rate=target.sampling_rate,
         )
 
+    if np.allclose(target.sampling_rate, 1):
+        warnings.warn(
+            "Target sampling rate is 1.0, which may indicate missing or incorrect "
+            "metadata. Verify that your target file contains proper sampling rate "
+            "information, as filters (CTF, BandPass) require accurate sampling rates "
+            "to function correctly."
+        )
+
     if target.sampling_rate.size == template.sampling_rate.size:
         if not np.allclose(
             np.round(target.sampling_rate, 2), np.round(template.sampling_rate, 2)
         ):
-            print(
-                f"Resampling template to {target.sampling_rate}. "
-                "Consider providing a template with the same sampling rate as the target."
+            warnings.warn(
+                f"Sampling rate mismatch detected: target={target.sampling_rate} "
+                f"template={template.sampling_rate}. Proceeding with user-provided "
+                f"values. Make sure this is intentional. "
             )
-            template = template.resample(target.sampling_rate, order=3)
 
     template_mask = load_and_validate_mask(
         mask_target=template, mask_path=args.template_mask
@@ -881,16 +884,13 @@ def main():
     print("\n" + "-" * 80)
 
     if args.scramble_phases:
-        template.data = scramble_phases(
-            template.data, noise_proportion=1.0, normalize_power=False
-        )
+        template.data = scramble_phases(template.data, noise_proportion=1.0)
 
     callback_class = MaxScoreOverRotations
-    if args.peak_calling:
-        callback_class = PeakCallerMaximumFilter
-
     if args.orientations is not None:
         callback_class = MaxScoreOverRotationsConstrained
+    elif args.peak_calling:
+        callback_class = PeakCallerMaximumFilter
 
     # Determine suitable backend for the selected operation
     available_backends = be.available_backends()
@@ -939,16 +939,6 @@ def main():
             args.use_gpu = False
             be.change_backend("pytorch", device=device)
 
-    # TODO: Make the inverse casting from complex64 -> float 16 stable
-    # if args.use_mixed_precision:
-    #     be.change_backend(
-    #         backend_name=args.backend,
-    #         float_dtype=be._array_backend.float16,
-    #         complex_dtype=be._array_backend.complex64,
-    #         int_dtype=be._array_backend.int16,
-    #         device=device,
-    #     )
-
     available_memory = be.get_available_memory() * be.device_count()
     if args.memory is None:
         args.memory = int(args.memory_scaling * available_memory)
@@ -976,6 +966,8 @@ def main():
         target_dim=target.metadata.get("batch_dimension", None),
         template_dim=template.metadata.get("batch_dimension", None),
     )
+    args.batch_dims = tuple(int(x) for x in np.where(matching_data._batch_mask)[0])
+
     splits, schedule = compute_schedule(args, matching_data, callback_class)
 
     n_splits = np.prod(list(splits.values()))
@@ -1004,7 +996,6 @@ def main():
     compute_options = {
         "Backend": be._BACKEND_REGISTRY[be._backend_name],
         "Compute Devices": f"CPU [{args.cores}], GPU [{gpus_used}]",
-        "Use Mixed Precision": args.use_mixed_precision,
         "Assigned Memory [MB]": f"{args.memory // 1e6} [out of {available_memory//1e6}]",
         "Temporary Directory": args.temp_directory,
         "Target Splits": f"{target_split} [N={n_splits}]",
@@ -1025,7 +1016,6 @@ def main():
         "Tilt Angles": args.tilt_angles,
         "Tilt Weighting": args.tilt_weighting,
         "Reconstruction Filter": args.reconstruction_filter,
-        "Extend Filter Grid": args.pad_filter,
     }
     if args.ctf_file is not None or args.defocus is not None:
         filter_args["CTF File"] = args.ctf_file
@@ -1081,7 +1071,6 @@ def main():
         callback_class_args=analyzer_args,
         target_splits=splits,
         pad_target_edges=args.pad_edges,
-        pad_template_filter=args.pad_filter,
         interpolation_order=args.interpolation_order,
     )
 

{pytme-0.3b0.post1.data → pytme-0.3.1.dev20250731.data}/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"
 

{pytme-0.3b0.post1.data → pytme-0.3.1.dev20250731.data}/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)