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

{pytme-0.2.1.data → pytme-0.2.2.data}/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.")
 
 

{pytme-0.2.1.data → pytme-0.2.2.data}/scripts/postprocess.py

@@ -8,9 +8,8 @@
 import argparse
 from sys import exit
 from os import getcwd
-from os.path import join, abspath
 from typing import List, Tuple
-from os.path import splitext
+from os.path import join, abspath, splitext
 
 import numpy as np
 from numpy.typing import NDArray
@@ -26,6 +25,7 @@ from tme.analyzer import (
 )
 from tme.matching_utils import (
     load_pickle,
+    centered_mask,
     euler_to_rotationmatrix,
     euler_from_rotationmatrix,
 )
@@ -41,9 +41,7 @@ PEAK_CALLERS = {
 
 
 def parse_args():
-    parser = argparse.ArgumentParser(
-        description="Peak Calling for Template Matching Outputs"
-    )
+    parser = argparse.ArgumentParser(description="Analyze Template Matching Outputs")
 
     input_group = parser.add_argument_group("Input")
     output_group = parser.add_argument_group("Output")
@@ -56,6 +54,13 @@ def parse_args():
         nargs="+",
         help="Path to the output of match_template.py.",
     )
+    input_group.add_argument(
+        "--background_file",
+        required=False,
+        nargs="+",
+        help="Path to an output of match_template.py used for normalization. "
+        "For instance from --scramble_phases or a different template.",
+    )
     input_group.add_argument(
         "--target_mask",
         required=False,
@@ -87,7 +92,7 @@ def parse_args():
             "average",
         ],
         default="orientations",
-        help="Available output formats:"
+        help="Available output formats: "
         "orientations (translation, rotation, and score), "
         "alignment (aligned template to target based on orientations), "
         "extraction (extract regions around peaks from targets, i.e. subtomograms), "
@@ -206,6 +211,15 @@ def parse_args():
     elif args.number_of_peaks is None:
         args.number_of_peaks = 1000
 
+    if args.background_file is None:
+        args.background_file = [None]
+    if len(args.background_file) == 1:
+        args.background_file = args.background_file * len(args.input_file)
+    elif len(args.background_file) not in (0, len(args.input_file)):
+        raise ValueError(
+            "--background_file needs to be specified once or for each --input_file."
+        )
+
     return args
 
 
@@ -233,8 +247,8 @@ def load_template(
     return template, center, translation, template_is_density
 
 
-def merge_outputs(data, filepaths: List[str], args):
-    if len(filepaths) == 0:
+def merge_outputs(data, foreground_paths: List[str], background_paths: List[str], args):
+    if len(foreground_paths) == 0:
         return data, 1
 
     if data[0].ndim != data[2].ndim:
@@ -275,8 +289,11 @@ def merge_outputs(data, filepaths: List[str], args):
 
     entities = np.zeros_like(data[0])
     data[0] = _norm_scores(data=data, args=args)
-    for index, filepath in enumerate(filepaths):
-        new_scores = _norm_scores(data=load_pickle(filepath), args=args)
+    for index, filepath in enumerate(foreground_paths):
+        new_scores = _norm_scores(
+            data=load_match_template_output(filepath, background_paths[index]),
+            args=args,
+        )
         indices = new_scores > data[0]
         entities[indices] = index + 1
         data[0][indices] = new_scores[indices]
@@ -284,9 +301,18 @@ def merge_outputs(data, filepaths: List[str], args):
     return data, entities
 
 
+def load_match_template_output(foreground_path, background_path):
+    data = load_pickle(foreground_path)
+    if background_path is not None:
+        data_background = load_pickle(background_path)
+        data[0] = (data[0] - data_background[0]) / (1 - data_background[0])
+        np.fmax(data[0], 0, out=data[0])
+    return data
+
+
 def main():
     args = parse_args()
-    data = load_pickle(args.input_file[0])
+    data = load_match_template_output(args.input_file[0], args.background_file[0])
 
     target_origin, _, sampling_rate, cli_args = data[-1]
 
@@ -326,7 +352,12 @@ def main():
 
     entities = None
     if len(args.input_file) > 1:
-        data, entities = merge_outputs(data=data, filepaths=args.input_file, args=args)
+        data, entities = merge_outputs(
+            data=data,
+            foreground_paths=args.input_file,
+            background_paths=args.background_file,
+            args=args,
+        )
 
     orientations = args.orientations
     if orientations is None:
@@ -339,24 +370,27 @@ def main():
                 target_mask = Density.from_file(args.target_mask)
                 scores = scores * target_mask.data
 
-            if args.n_false_positives is not None:
-                args.n_false_positives = max(args.n_false_positives, 1)
-                cropped_shape = np.subtract(
-                    scores.shape, np.multiply(args.min_boundary_distance, 2)
-                ).astype(int)
+            cropped_shape = np.subtract(
+                scores.shape, np.multiply(args.min_boundary_distance, 2)
+            ).astype(int)
 
-                cropped_shape = tuple(
+            if args.min_boundary_distance > 0:
+                scores = centered_mask(scores, new_shape=cropped_shape)
+
+            if args.n_false_positives is not None:
+                # Rickgauer et al. 2017
+                cropped_slice = tuple(
                     slice(
                         int(args.min_boundary_distance),
                         int(x - args.min_boundary_distance),
                     )
                     for x in scores.shape
                 )
-                # Rickgauer et al. 2017
-                n_correlations = np.size(scores[cropped_shape]) * len(rotation_mapping)
+                args.n_false_positives = max(args.n_false_positives, 1)
+                n_correlations = np.size(scores[cropped_slice]) * len(rotation_mapping)
                 minimum_score = np.multiply(
                     erfcinv(2 * args.n_false_positives / n_correlations),
-                    np.sqrt(2) * np.std(scores[cropped_shape]),
+                    np.sqrt(2) * np.std(scores[cropped_slice]),
                 )
                 print(f"Determined minimum score cutoff: {minimum_score}.")
                 minimum_score = max(minimum_score, 0)
@@ -371,6 +405,8 @@ def main():
                 "min_distance": args.min_distance,
                 "min_boundary_distance": args.min_boundary_distance,
                 "batch_dims": args.batch_dims,
+                "minimum_score": args.minimum_score,
+                "maximum_score": args.maximum_score,
             }
 
             peak_caller = PEAK_CALLERS[args.peak_caller](**peak_caller_kwargs)
@@ -380,7 +416,6 @@ def main():
                 mask=template.data,
                 rotation_mapping=rotation_mapping,
                 rotation_array=rotation_array,
-                minimum_score=args.minimum_score,
             )
             candidates = peak_caller.merge(
                 candidates=[tuple(peak_caller)], **peak_caller_kwargs
@@ -388,10 +423,16 @@ def main():
             if len(candidates) == 0:
                 candidates = [[], [], [], []]
                 print("Found no peaks, consider changing peak calling parameters.")
-                exit(0)
+                exit(-1)
 
             for translation, _, score, detail in zip(*candidates):
-                rotations.append(rotation_mapping[rotation_array[tuple(translation)]])
+                rotation_index = rotation_array[tuple(translation)]
+                rotation = rotation_mapping.get(
+                    rotation_index, np.zeros(template.data.ndim, int)
+                )
+                if rotation.ndim == 2:
+                    rotation = euler_from_rotationmatrix(rotation)
+                rotations.append(rotation)
 
         else:
             candidates = data
@@ -430,7 +471,7 @@ def main():
             )
             exit(-1)
         orientations.translations = peak_caller.oversample_peaks(
-            score_space=data[0],
+            scores=data[0],
             peak_positions=orientations.translations,
             oversampling_factor=args.peak_oversampling,
         )
@@ -570,7 +611,7 @@ def main():
             return_orientations=True,
         )
         out = np.zeros_like(template.data)
-        out = np.zeros(np.multiply(template.shape, 2).astype(int))
+        # out = np.zeros(np.multiply(template.shape, 2).astype(int))
         for index in range(len(cand_slices)):
             from scipy.spatial.transform import Rotation