pytme 0.2.1__cp311-cp311-macosx_14_0_arm64.whl → 0.2.2__cp311-cp311-macosx_14_0_arm64.whl

scripts/match_template.py

@@ -8,7 +8,6 @@
 import os
 import argparse
 import warnings
-import importlib.util
 from sys import exit
 from time import time
 from typing import Tuple
@@ -22,7 +21,6 @@ from tme.matching_utils import (
     get_rotation_matrices,
     get_rotations_around_vector,
     compute_parallelization_schedule,
-    euler_from_rotationmatrix,
     scramble_phases,
     generate_tempfile_name,
     write_pickle,
@@ -33,7 +31,7 @@ from tme.analyzer import (
     MaxScoreOverRotations,
     PeakCallerMaximumFilter,
 )
-from tme.backends import backend
+from tme.backends import backend as be
 from tme.preprocessing import Compose
 
 
@@ -52,7 +50,7 @@ def print_block(name: str, data: dict, label_width=20) -> None:
 
 def print_entry() -> None:
     width = 80
-    text = f" pyTME v{__version__} "
+    text = f" pytme v{__version__} "
     padding_total = width - len(text) - 2
     padding_left = padding_total // 2
     padding_right = padding_total - padding_left
@@ -273,7 +271,7 @@ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Com
                 return_real_fourier=True,
             )
         ctf.sampling_rate = template.sampling_rate
-        ctf.flip_phase = not args.no_flip_phase
+        ctf.flip_phase = args.no_flip_phase
         ctf.amplitude_contrast = args.amplitude_contrast
         ctf.spherical_aberration = args.spherical_aberration
         ctf.acceleration_voltage = args.acceleration_voltage * 1e3
@@ -306,6 +304,12 @@ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Com
             if highpass is not None:
                 highpass = np.max(np.divide(template.sampling_rate, highpass))
 
+        try:
+            if args.lowpass >= args.highpass:
+                warnings.warn("--lowpass should be smaller than --highpass.")
+        except Exception:
+            pass
+
         bandpass = BandPassFilter(
             use_gaussian=args.no_pass_smooth,
             lowpass=lowpass,
@@ -313,7 +317,9 @@ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Com
             sampling_rate=template.sampling_rate,
         )
         template_filter.append(bandpass)
-        target_filter.append(bandpass)
+
+        if not args.no_filter_target:
+            target_filter.append(bandpass)
 
     if args.whiten_spectrum:
         whitening_filter = LinearWhiteningFilter()
@@ -335,7 +341,10 @@ def setup_filter(args, template: Density, target: Density) -> Tuple[Compose, Com
 
 
 def parse_args():
-    parser = argparse.ArgumentParser(description="Perform template matching.")
+    parser = argparse.ArgumentParser(
+        description="Perform template matching.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
 
     io_group = parser.add_argument_group("Input / Output")
     io_group.add_argument(
@@ -405,13 +414,6 @@ def parse_args():
         choices=list(MATCHING_EXHAUSTIVE_REGISTER.keys()),
         help="Template matching scoring function.",
     )
-    scoring_group.add_argument(
-        "-p",
-        dest="peak_calling",
-        action="store_true",
-        default=False,
-        help="Perform peak calling instead of score aggregation.",
-    )
 
     angular_group = parser.add_argument_group("Angular Sampling")
     angular_exclusive = angular_group.add_mutually_exclusive_group(required=True)
@@ -445,7 +447,7 @@ def parse_args():
         type=check_positive,
         default=360.0,
         required=False,
-        help="Sampling angle along the z-axis of the cone. Defaults to 360.",
+        help="Sampling angle along the z-axis of the cone.",
     )
     angular_group.add_argument(
         "--axis_sampling",
@@ -513,8 +515,7 @@ def parse_args():
         required=False,
         type=float,
         default=0.85,
-        help="Fraction of available memory that can be used. Defaults to 0.85 and is "
-        "ignored if --ram is set",
+        help="Fraction of available memory to be used. Ignored if --ram is set.",
     )
     computation_group.add_argument(
         "--temp_directory",
@@ -522,7 +523,13 @@ def parse_args():
         default=None,
         help="Directory for temporary objects. Faster I/O improves runtime.",
     )
-
+    computation_group.add_argument(
+        "--backend",
+        dest="backend",
+        default=None,
+        choices=be.available_backends(),
+        help="[Expert] Overwrite default computation backend.",
+    )
     filter_group = parser.add_argument_group("Filters")
     filter_group.add_argument(
         "--lowpass",
@@ -552,9 +559,9 @@ def parse_args():
         dest="pass_format",
         type=str,
         required=False,
+        default="sampling_rate",
         choices=["sampling_rate", "voxel", "frequency"],
-        help="How values passed to --lowpass and --highpass should be interpreted. "
-        "By default, they are assumed to be in units of sampling rate, e.g. Ångstrom.",
+        help="How values passed to --lowpass and --highpass should be interpreted. ",
     )
     filter_group.add_argument(
         "--whiten_spectrum",
@@ -613,6 +620,13 @@ def parse_args():
         required=False,
         help="Analogous to --interpolation_order but for reconstruction.",
     )
+    filter_group.add_argument(
+        "--no_filter_target",
+        dest="no_filter_target",
+        action="store_true",
+        default=False,
+        help="Whether to not apply potential filters to the target.",
+    )
 
     ctf_group = parser.add_argument_group("Contrast Transfer Function")
     ctf_group.add_argument(
@@ -647,7 +661,7 @@ def parse_args():
         type=float,
         required=False,
         default=300,
-        help="Acceleration voltage in kV, defaults to 300.",
+        help="Acceleration voltage in kV.",
     )
     ctf_group.add_argument(
         "--spherical_aberration",
@@ -663,14 +677,14 @@ def parse_args():
         type=float,
         required=False,
         default=0.07,
-        help="Amplitude contrast, defaults to 0.07.",
+        help="Amplitude contrast.",
     )
     ctf_group.add_argument(
         "--no_flip_phase",
         dest="no_flip_phase",
         action="store_false",
         required=False,
-        help="Whether the phase of the computed CTF should not be flipped.",
+        help="Do not perform phase-flipping CTF correction.",
     )
     ctf_group.add_argument(
         "--correct_defocus_gradient",
@@ -721,14 +735,22 @@ def parse_args():
         "for numerical stability. When working with very large targets, e.g. tomograms, "
         "it is safe to use this flag and benefit from the performance gain.",
     )
+    performance_group.add_argument(
+        "--no_filter_padding",
+        dest="no_filter_padding",
+        action="store_true",
+        default=False,
+        help="Omits padding of optional template filters. Particularly effective when "
+        "the target is much larger than the template. However, for fast osciliating "
+        "filters setting this flag can introduce aliasing effects.",
+    )
     performance_group.add_argument(
         "--interpolation_order",
         dest="interpolation_order",
         required=False,
         type=int,
         default=3,
-        help="Spline interpolation used for template rotations. If less than zero "
-        "no interpolation is performed.",
+        help="Spline interpolation used for rotations.",
     )
     performance_group.add_argument(
         "--use_mixed_precision",
@@ -755,7 +777,20 @@ def parse_args():
         default=0,
         help="Minimum template matching scores to consider for analysis.",
     )
-
+    analyzer_group.add_argument(
+        "-p",
+        dest="peak_calling",
+        action="store_true",
+        default=False,
+        help="Perform peak calling instead of score aggregation.",
+    )
+    analyzer_group.add_argument(
+        "--number_of_peaks",
+        dest="number_of_peaks",
+        action="store_true",
+        default=1000,
+        help="Number of peaks to call, 1000 by default.",
+    )
     args = parser.parse_args()
     args.version = __version__
 
@@ -924,44 +959,56 @@ def main():
             template.data, noise_proportion=1.0, normalize_power=True
         )
 
-    available_memory = backend.get_available_memory()
+    # Determine suitable backend for the selected operation
+    available_backends = be.available_backends()
+    if args.backend is not None:
+        req_backend = args.backend
+        if req_backend not in available_backends:
+            raise ValueError("Requested backend is not available.")
+        available_backends = [req_backend]
+
+    be_selection = ("numpyfftw", "pytorch", "jax", "mlx")
     if args.use_gpu:
         args.cores = len(args.gpu_indices)
-        has_torch = importlib.util.find_spec("torch") is not None
-        has_cupy = importlib.util.find_spec("cupy") is not None
-
-        if not has_torch and not has_cupy:
-            raise ValueError(
-                "Found neither CuPy nor PyTorch installation. You need to install"
-                " either to enable GPU support."
-            )
+        be_selection = ("pytorch", "cupy", "jax")
+    if args.use_mixed_precision:
+        be_selection = tuple(x for x in be_selection if x in ("cupy", "numpyfftw"))
 
-        if args.peak_calling:
-            preferred_backend = "pytorch"
-            if not has_torch:
-                preferred_backend = "cupy"
-            backend.change_backend(backend_name=preferred_backend, device="cuda")
-        else:
-            preferred_backend = "cupy"
-            if not has_cupy:
-                preferred_backend = "pytorch"
-            backend.change_backend(backend_name=preferred_backend, device="cuda")
-            if args.use_mixed_precision and preferred_backend == "pytorch":
+    available_backends = [x for x in available_backends if x in be_selection]
+    if args.peak_calling:
+        if "jax" in available_backends:
+            available_backends.remove("jax")
+        if args.use_gpu and "pytorch" in available_backends:
+            available_backends = ("pytorch",)
+            if args.interpolation_order == 3:
                 raise NotImplementedError(
-                    "pytorch backend does not yet support mixed precision."
-                    " Consider installing CuPy to enable this feature."
+                    "Pytorch does not support --interpolation_order 3, 1 is supported."
                 )
-            elif args.use_mixed_precision:
-                backend.change_backend(
-                    backend_name="cupy",
-                    default_dtype=backend._array_backend.float16,
-                    complex_dtype=backend._array_backend.complex64,
-                    default_dtype_int=backend._array_backend.int16,
-                )
-        available_memory = backend.get_available_memory() * args.cores
-        if preferred_backend == "pytorch" and args.interpolation_order == 3:
-            args.interpolation_order = 1
+    # dim_match = len(template.shape) == len(target.shape) <= 3
+    # if dim_match and args.use_gpu and "jax" in available_backends:
+    #     args.interpolation_order = 1
+    #     available_backends = ["jax"]
 
+    backend_preference = ("numpyfftw", "pytorch", "jax", "mlx")
+    if args.use_gpu:
+        backend_preference = ("cupy", "pytorch", "jax")
+    for pref in backend_preference:
+        if pref not in available_backends:
+            continue
+        be.change_backend(pref)
+        if pref == "pytorch":
+            be.change_backend(pref, device="cuda" if args.use_gpu else "cpu")
+
+        if args.use_mixed_precision:
+            be.change_backend(
+                backend_name=pref,
+                default_dtype=be._array_backend.float16,
+                complex_dtype=be._array_backend.complex64,
+                default_dtype_int=be._array_backend.int16,
+            )
+        break
+
+    available_memory = be.get_available_memory() * be.device_count()
     if args.memory is None:
         args.memory = int(args.memory_scaling * available_memory)
 
@@ -978,17 +1025,15 @@ def main():
         rotations=parse_rotation_logic(args=args, ndim=template.data.ndim),
     )
 
-    template_filter, target_filter = setup_filter(args, template, target)
-    matching_data.template_filter = template_filter
-    matching_data.target_filter = target_filter
+    matching_data.template_filter, matching_data.target_filter = setup_filter(
+        args, template, target
+    )
 
     template_box = matching_data._output_template_shape
     if not args.pad_fourier:
-        template_box = np.ones(len(template_box), dtype=int)
+        template_box = tuple(0 for _ in range(len(template_box)))
 
-    target_padding = np.zeros(
-        (backend.size(matching_data._output_template_shape)), dtype=int
-    )
+    target_padding = tuple(0 for _ in range(len(template_box)))
     if args.pad_target_edges:
         target_padding = matching_data._output_template_shape
 
@@ -1001,10 +1046,10 @@ def main():
         split_only_outer=args.use_gpu,
         matching_method=args.score,
         analyzer_method=callback_class.__name__,
-        backend=backend._backend_name,
-        float_nbytes=backend.datatype_bytes(backend._float_dtype),
-        complex_nbytes=backend.datatype_bytes(backend._complex_dtype),
-        integer_nbytes=backend.datatype_bytes(backend._int_dtype),
+        backend=be._backend_name,
+        float_nbytes=be.datatype_bytes(be._float_dtype),
+        complex_nbytes=be.datatype_bytes(be._complex_dtype),
+        integer_nbytes=be.datatype_bytes(be._int_dtype),
     )
 
     if splits is None:
@@ -1021,27 +1066,36 @@ def main():
     )
     gpus_used = 0 if args.gpu_indices is None else len(args.gpu_indices)
     options = {
-        "CPU Cores": args.cores,
-        "Run on GPU": f"{args.use_gpu} [N={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,
+        "Angular Sampling": f"{args.angular_sampling}"
+        f" [{matching_data.rotations.shape[0]} rotations]",
+        "Center Template": not args.no_centering,
+        "Scramble Template": args.scramble_phases,
+        "Invert Contrast": args.invert_target_contrast,
         "Extend Fourier Grid": not args.no_fourier_padding,
         "Extend Target Edges": not args.no_edge_padding,
         "Interpolation Order": args.interpolation_order,
-        "Score": f"{args.score}",
         "Setup Function": f"{get_func_fullname(matching_setup)}",
         "Scoring Function": f"{get_func_fullname(matching_score)}",
-        "Angular Sampling": f"{args.angular_sampling}"
-        f" [{matching_data.rotations.shape[0]} rotations]",
-        "Scramble Template": args.scramble_phases,
-        "Target Splits": f"{target_split} [N={n_splits}]",
     }
 
     print_block(
-        name="Template Matching Options",
+        name="Template Matching",
         data=options,
-        label_width=max(len(key) for key in options.keys()) + 2,
+        label_width=max(len(key) for key in options.keys()) + 3,
+    )
+
+    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}]",
+    }
+    print_block(
+        name="Computation",
+        data=compute_options,
+        label_width=max(len(key) for key in options.keys()) + 3,
     )
 
     filter_args = {
@@ -1059,7 +1113,7 @@ def main():
         filter_args["CTF File"] = args.ctf_file
         filter_args["Defocus"] = args.defocus
         filter_args["Phase Shift"] = args.phase_shift
-        filter_args["No Flip Phase"] = args.no_flip_phase
+        filter_args["Flip Phase"] = args.no_flip_phase
         filter_args["Acceleration Voltage"] = args.acceleration_voltage
         filter_args["Spherical Aberration"] = args.spherical_aberration
         filter_args["Amplitude Contrast"] = args.amplitude_contrast
@@ -1070,20 +1124,20 @@ def main():
         print_block(
             name="Filters",
             data=filter_args,
-            label_width=max(len(key) for key in options.keys()) + 2,
+            label_width=max(len(key) for key in options.keys()) + 3,
         )
 
     analyzer_args = {
         "score_threshold": args.score_threshold,
-        "number_of_peaks": 1000,
-        "convolution_mode": "valid",
+        "number_of_peaks": args.number_of_peaks,
+        "min_distance": max(template.shape) // 2,
+        "min_boundary_distance": max(template.shape) // 2,
         "use_memmap": args.use_memmap,
     }
-    analyzer_args = {"Analyzer": callback_class, **analyzer_args}
     print_block(
-        name="Score Analysis Options",
-        data=analyzer_args,
-        label_width=max(len(key) for key in options.keys()) + 2,
+        name="Analyzer",
+        data={"Analyzer": callback_class, **analyzer_args},
+        label_width=max(len(key) for key in options.keys()) + 3,
     )
     print("\n" + "-" * 80)
 
@@ -1104,6 +1158,7 @@ def main():
         target_splits=splits,
         pad_target_edges=args.pad_target_edges,
         pad_fourier=args.pad_fourier,
+        pad_template_filter=not args.no_filter_padding,
         interpolation_order=args.interpolation_order,
     )
 
@@ -1113,19 +1168,18 @@ def main():
             candidates[0] *= target_mask.data
         with warnings.catch_warnings():
             warnings.simplefilter("ignore", category=UserWarning)
-            nbytes = backend.datatype_bytes(backend._float_dtype)
+            nbytes = be.datatype_bytes(be._float_dtype)
             dtype = np.float32 if nbytes == 4 else np.float16
             rot_dim = matching_data.rotations.shape[1]
             candidates[3] = {
-                x: euler_from_rotationmatrix(
-                    np.frombuffer(i, dtype=dtype).reshape(rot_dim, rot_dim)
-                )
+                x: np.frombuffer(i, dtype=dtype).reshape(rot_dim, rot_dim)
                 for i, x in candidates[3].items()
             }
     candidates.append((target.origin, template.origin, template.sampling_rate, args))
     write_pickle(data=candidates, filename=args.output)
 
     runtime = time() - start
+    print("\n" + "-" * 80)
     print(f"\nRuntime real: {runtime:.3f}s user: {(runtime * args.cores):.3f}s.")