pytme 0.1.9__cp311-cp311-macosx_14_0_arm64.whl → 0.2.0b0__cp311-cp311-macosx_14_0_arm64.whl

scripts/postprocess.py

@@ -5,19 +5,18 @@
 
     Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
-from os import getcwd
-from os.path import join
 import argparse
 from sys import exit
-from typing import List, Tuple
+from os import getcwd
+from os.path import join, abspath
+from typing import List
 from os.path import splitext
-from dataclasses import dataclass
 
 import numpy as np
-from scipy.spatial.transform import Rotation
 from numpy.typing import NDArray
+from scipy.special import erfcinv
 
-from tme import Density, Structure
+from tme import Density, Structure, Orientations
 from tme.analyzer import (
     PeakCallerSort,
     PeakCallerMaximumFilter,
@@ -29,7 +28,6 @@ from tme.matching_utils import (
     load_pickle,
     euler_to_rotationmatrix,
     euler_from_rotationmatrix,
-    centered_mask,
 )
 
 PEAK_CALLERS = {
@@ -45,420 +43,325 @@ def parse_args():
     parser = argparse.ArgumentParser(
         description="Peak Calling for Template Matching Outputs"
     )
-    parser.add_argument(
+
+    input_group = parser.add_argument_group("Input")
+    output_group = parser.add_argument_group("Output")
+    peak_group = parser.add_argument_group("Peak Calling")
+    additional_group = parser.add_argument_group("Additional Parameters")
+
+    input_group.add_argument(
         "--input_file",
         required=True,
+        nargs="+",
         help="Path to the output of match_template.py.",
     )
-    parser.add_argument(
+    input_group.add_argument(
+        "--target_mask",
+        required=False,
+        type=str,
+        help="Path to an optional mask applied to template matching scores.",
+    )
+    input_group.add_argument(
+        "--orientations",
+        required=False,
+        type=str,
+        help="Path to file generated using output_format orientations. Can be filtered "
+        "to exclude false-positive peaks. If this file is provided, peak calling "
+        "is skipped and corresponding parameters ignored.",
+    )
+
+    output_group.add_argument(
         "--output_prefix",
         required=True,
-        help="Prefix for the output file name. Extension depends on output_format.",
+        help="Output filename, extension will be added based on output_format.",
     )
-    parser.add_argument(
-        "--number_of_peaks",
-        type=int,
-        default=1000,
-        help="Number of peaks to consider. Note, this is the number of called peaks "
-        ", subject to min_distance and min_boundary_distance filtering. Therefore, the "
-        "returned number of peaks will be at most equal to number_of_peaks. "
-        "Ignored when --orientations is provided.",
+    output_group.add_argument(
+        "--output_format",
+        choices=[
+            "orientations",
+            "alignment",
+            "extraction",
+            "relion",
+            "backmapping",
+            "average",
+        ],
+        default="orientations",
+        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), "
+        "relion (perform extraction step and generate corresponding star files), "
+        "backmapping (map template to target using identified peaks),"
+        "average (extract matched regions from target and average them).",
     )
-    parser.add_argument(
+
+    peak_group.add_argument(
+        "--peak_caller",
+        choices=list(PEAK_CALLERS.keys()),
+        default="PeakCallerScipy",
+        help="Peak caller for local maxima identification.",
+    )
+    peak_group.add_argument(
+        "--minimum_score",
+        type=float,
+        default=None,
+        help="Minimum score from which peaks will be considered.",
+    )
+    peak_group.add_argument(
+        "--maximum_score",
+        type=float,
+        default=None,
+        help="Maximum score until which peaks will be considered.",
+    )
+    peak_group.add_argument(
         "--min_distance",
         type=int,
         default=5,
-        help="Minimum distance between peaks. Ignored when --orientations is provided.",
+        help="Minimum distance between peaks.",
     )
-    parser.add_argument(
+    peak_group.add_argument(
         "--min_boundary_distance",
         type=int,
         default=0,
-        help="Minimum distance from target boundaries. Ignored when --orientations "
-        "is provided.",
+        help="Minimum distance of peaks to target edges.",
     )
-    parser.add_argument(
+    peak_group.add_argument(
         "--mask_edges",
         action="store_true",
         default=False,
-        help="Whether to mask edges of the input score array according to the template shape."
-        "Uses twice the value of --min_boundary_distance if boht are provided.",
+        help="Whether candidates should not be identified from scores that were "
+        "computed from padded densities. Superseded by min_boundary_distance.",
     )
-    parser.add_argument(
-        "--wedge_mask",
-        type=str,
+    peak_group.add_argument(
+        "--number_of_peaks",
+        type=int,
         default=None,
-        help="Path to Fourier space mask. Only considered if output_format is relion.",
+        required=False,
+        help="Upper limit of peaks to call, subject to filtering parameters. Default 1000. "
+        "If minimum_score is provided all peaks scoring higher will be reported.",
     )
-    parser.add_argument(
-        "--peak_caller",
-        choices=list(PEAK_CALLERS.keys()),
-        default="PeakCallerScipy",
-        help="Peak caller to use for analysis. Ignored if input_file contains peaks or when "
-        "--orientations is provided.",
+    peak_group.add_argument(
+        "--peak_oversampling",
+        type=int,
+        default=1,
+        help="1 / factor equals voxel precision, e.g. 2 detects half voxel "
+        "translations. Useful for matching structures to electron density maps.",
     )
-    parser.add_argument(
-        "--orientations",
+
+    additional_group.add_argument(
+        "--subtomogram_box_size",
+        type=int,
         default=None,
-        help="Path to orientations file to overwrite orientations computed from"
-        " match_template.py output.",
+        help="Subtomogram box size, by default equal to the centered template. Will be "
+        "padded to even values if output_format is relion.",
     )
-    parser.add_argument(
-        "--output_format",
-        choices=["orientations", "alignment", "extraction", "relion"],
-        default="orientations",
-        help="Choose the output format. Available formats are: "
-        "orientations (translation, rotation, and score), "
-        "alignment (aligned template to target based on orientations), "
-        "extraction (extract regions around peaks from targets, i.e. subtomograms). "
-        "relion (perform extraction step and generate corresponding star files).",
+    additional_group.add_argument(
+        "--mask_subtomograms",
+        action="store_true",
+        default=False,
+        help="Whether to mask subtomograms using the template mask. The mask will be "
+        "rotated according to determined angles.",
+    )
+    additional_group.add_argument(
+        "--invert_target_contrast",
+        action="store_true",
+        default=False,
+        help="Whether to invert the target contrast.",
+    )
+    additional_group.add_argument(
+        "--wedge_mask",
+        type=str,
+        default=None,
+        help="Path to file used as ctf_mask for output_format relion.",
+    )
+    additional_group.add_argument(
+        "--n_false_positives",
+        type=int,
+        default=None,
+        required=False,
+        help="Number of accepted false-positives picks to determine minimum score.",
     )
-    args = parser.parse_args()
-
-    return args
-
 
-@dataclass
-class Orientations:
-    #: Return a numpy array with translations of each orientation (n x d).
-    translations: np.ndarray
-
-    #: Return a numpy array with euler angles of each orientation in zxy format (n x d).
-    rotations: np.ndarray
-
-    #: Return a numpy array with the score of each orientation (n, ).
-    scores: np.ndarray
-
-    #: Return a numpy array with additional orientation details (n, ).
-    details: np.ndarray
-
-    def __iter__(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
-        """
-        Iterate over the current class instance. Each iteration returns a orientation
-        defined by its translation, rotation, score and additional detail.
-
-        Yields
-        ------
-        Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]
-            A tuple of arrays defining the given orientation.
-        """
-        yield from zip(self.translations, self.rotations, self.scores, self.details)
-
-    def __getitem__(self, indices: List[int]) -> "Orientations":
-        """
-        Retrieve a subset of orientations based on the provided indices.
-
-        Parameters
-        ----------
-        indices : List[int]
-            A list of indices specifying the orientations to be retrieved.
-
-        Returns
-        -------
-        :py:class:`Orientations`
-            A new :py:class:`Orientations`instance containing only the selected orientations.
-        """
-        indices = np.asarray(indices)
-        attributes = (
-            "translations",
-            "rotations",
-            "scores",
-            "details",
-        )
-        kwargs = {attr: getattr(self, attr)[indices] for attr in attributes}
-        return self.__class__(**kwargs)
-
-    def to_file(self, filename: str, file_format: type, **kwargs) -> None:
-        """
-        Save the current class instance to a file in the specified format.
-
-        Parameters
-        ----------
-        filename : str
-            The name of the file where the orientations will be saved.
-        file_format : type
-            The format in which to save the orientations. Supported formats are 'text' and 'relion'.
-        **kwargs : dict
-            Additional keyword arguments specific to the file format.
-
-        Raises
-        ------
-        ValueError
-            If an unsupported file format is specified.
-        """
-        mapping = {
-            "text": self._to_text,
-            "relion": self._to_relion_star,
-        }
-
-        func = mapping.get(file_format, None)
-        if func is None:
-            raise ValueError(
-                f"{file_format} not implemented. Supported are {','.join(mapping.keys())}."
-            )
+    args = parser.parse_args()
 
-        return func(filename=filename, **kwargs)
+    if args.wedge_mask is not None:
+        args.wedge_mask = abspath(args.wedge_mask)
 
-    def _to_text(self, filename: str) -> None:
-        """
-        Save orientations in a text file format.
+    if args.output_format == "relion" and args.subtomogram_box_size is not None:
+        args.subtomogram_box_size += args.subtomogram_box_size % 2
 
-        Parameters
-        ----------
-        filename : str
-            The name of the file to save the orientations.
+    if args.orientations is not None:
+        args.orientations = Orientations.from_file(filename=args.orientations)
 
-        Notes
-        -----
-        The file is saved with a header specifying each column: z, y, x, euler_z,
-        euler_y, euler_x, score, detail. Each row in the file corresponds to an orientation.
-        """
-        header = "\t".join(
-            ["z", "y", "x", "euler_z", "euler_y", "euler_x", "score", "detail"]
-        )
-        with open(filename, mode="w", encoding="utf-8") as ofile:
-            _ = ofile.write(f"{header}\n")
-            for translation, angles, score, detail in self:
-                translation_string = "\t".join([str(x) for x in translation])
-                angle_string = "\t".join([str(x) for x in angles])
-                _ = ofile.write(
-                    f"{translation_string}\t{angle_string}\t{score}\t{detail}\n"
-                )
-        return None
-
-    def _to_relion_star(
-        self,
-        filename: str,
-        name_prefix: str = None,
-        ctf_image: str = None,
-        sampling_rate: float = 1.0,
-        subtomogram_size: int = 0,
-    ) -> None:
-        """
-        Save orientations in RELION's STAR file format.
-
-        Parameters
-        ----------
-        filename : str
-            The name of the file to save the orientations.
-        name_prefix : str, optional
-            A prefix to add to the image names in the STAR file.
-        ctf_image : str, optional
-            Path to CTF or wedge mask RELION.
-        sampling_rate : float, optional
-            Subtomogram sampling rate in angstrom per voxel
-        subtomogram_size : int, optional
-            Size of the square shaped subtomogram.
-
-        Notes
-        -----
-        The file is saved with a standard header used in RELION STAR files.
-        Each row in the file corresponds to an orientation.
-        """
-        optics_header = [
-            "# version 30001",
-            "data_optics",
-            "",
-            "loop_",
-            "_rlnOpticsGroup",
-            "_rlnOpticsGroupName",
-            "_rlnSphericalAberration",
-            "_rlnVoltage",
-            "_rlnImageSize",
-            "_rlnImageDimensionality",
-            "_rlnImagePixelSize",
-        ]
-        optics_data = [
-            "1",
-            "opticsGroup1",
-            "2.700000",
-            "300.000000",
-            str(int(subtomogram_size)),
-            "3",
-            str(float(sampling_rate)),
-        ]
-        optics_header = "\n".join(optics_header)
-        optics_data = "\t".join(optics_data)
-
-        header = [
-            "data_particles",
-            "",
-            "loop_",
-            "_rlnCoordinateX",
-            "_rlnCoordinateY",
-            "_rlnCoordinateZ",
-            "_rlnImageName",
-            "_rlnAngleRot",
-            "_rlnAngleTilt",
-            "_rlnAnglePsi",
-            "_rlnOpticsGroup",
-        ]
-        if ctf_image is not None:
-            header.append("_rlnCtfImage")
-
-        ctf_image = "" if ctf_image is None else f"\t{ctf_image}"
-
-        header = "\n".join(header)
-        name_prefix = "" if name_prefix is None else name_prefix
-
-        with open(filename, mode="w", encoding="utf-8") as ofile:
-            _ = ofile.write(f"{optics_header}\n")
-            _ = ofile.write(f"{optics_data}\n")
-
-            _ = ofile.write("\n# version 30001\n")
-            _ = ofile.write(f"{header}\n")
-
-            # pyTME uses a zyx data layout
-            for index, (translation, rotation, score, detail) in enumerate(self):
-                rotation = Rotation.from_euler("zyx", rotation, degrees=True)
-                rotation = rotation.as_euler(seq="xyx", degrees=True)
-
-                translation_string = "\t".join([str(x) for x in translation][::-1])
-                angle_string = "\t".join([str(x) for x in rotation])
-                name = f"{name_prefix}_{index}.mrc"
-                _ = ofile.write(
-                    f"{translation_string}\t{name}\t{angle_string}\t1{ctf_image}\n"
-                )
+    if args.minimum_score is not None or args.n_false_positives is not None:
+        args.number_of_peaks = np.iinfo(np.int64).max
+    else:
+        args.number_of_peaks = 1000
 
-        return None
-
-    @classmethod
-    def from_file(cls, filename: str, file_format: type, **kwargs) -> "Orientations":
-        """
-        Create an instance of :py:class:`Orientations` from a file.
-
-        Parameters
-        ----------
-        filename : str
-            The name of the file from which to read the orientations.
-        file_format : type
-            The format of the file. Currently, only 'text' format is supported.
-        **kwargs : dict
-            Additional keyword arguments specific to the file format.
-
-        Returns
-        -------
-        :py:class:`Orientations`
-            An instance of :py:class:`Orientations` populated with data from the file.
-
-        Raises
-        ------
-        ValueError
-            If an unsupported file format is specified.
-        """
-        mapping = {
-            "text": cls._from_text,
-        }
-
-        func = mapping.get(file_format, None)
-        if func is None:
-            raise ValueError(
-                f"{file_format} not implemented. Supported are {','.join(mapping.keys())}."
-            )
+    return args
 
-        translations, rotations, scores, details, *_ = func(filename=filename, **kwargs)
-        return cls(
-            translations=translations,
-            rotations=rotations,
-            scores=scores,
-            details=details,
-        )
 
-    @staticmethod
-    def _from_text(
-        filename: str,
-    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
-        """
-        Read orientations from a text file.
-
-        Parameters
-        ----------
-        filename : str
-            The name of the file from which to read the orientations.
-
-        Returns
-        -------
-        Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]
-            A tuple containing numpy arrays for translations, rotations, scores,
-            and details.
-
-        Notes
-        -----
-        The text file is expected to have a header and data in columns corresponding to
-        z, y, x, euler_z, euler_y, euler_x, score, detail.
-        """
-        with open(filename, mode="r", encoding="utf-8") as infile:
-            data = [x.strip().split("\t") for x in infile.read().split("\n")]
-            _ = data.pop(0)
-
-        translation, rotation, score, detail = [], [], [], []
-        for candidate in data:
-            if len(candidate) <= 1:
-                continue
-            if len(candidate) != 8:
-                candidate.append(-1)
-
-            candidate = [float(x) for x in candidate]
-            translation.append((candidate[0], candidate[1], candidate[2]))
-            rotation.append((candidate[3], candidate[4], candidate[5]))
-            score.append(candidate[6])
-            detail.append(candidate[7])
-
-        translation = np.vstack(translation).astype(int)
-        rotation = np.vstack(rotation).astype(float)
-        score = np.array(score).astype(float)
-        detail = np.array(detail).astype(float)
-
-        return translation, rotation, score, detail
-
-
-def load_template(filepath: str, sampling_rate: NDArray) -> "Density":
+def load_template(filepath: str, sampling_rate: NDArray, center: bool = True):
     try:
         template = Density.from_file(filepath)
-        template, _ = template.centered(0)
         center_of_mass = template.center_of_mass(template.data)
+        template_is_density = True
     except ValueError:
         template = Structure.from_file(filepath)
         center_of_mass = template.center_of_mass()[::-1]
         template = Density.from_structure(template, sampling_rate=sampling_rate)
+        template_is_density = False
+
+    translation = np.zeros_like(center_of_mass)
+    if center:
+        template, translation = template.centered(0)
+
+    return template, center_of_mass, translation, template_is_density
+
 
-    return template, center_of_mass
+def merge_outputs(data, filepaths: List[str], args):
+    if len(filepaths) == 0:
+        return data, 1
+
+    if data[0].ndim != data[2].ndim:
+        return data, 1
+
+    from tme.matching_exhaustive import _normalize_under_mask
+
+    def _norm_scores(data, args):
+        target_origin, _, sampling_rate, cli_args = data[-1]
+
+        _, template_extension = splitext(cli_args.template)
+        ret = load_template(
+            filepath=cli_args.template,
+            sampling_rate=sampling_rate,
+            center=not cli_args.no_centering,
+        )
+        template, center_of_mass, translation, template_is_density = ret
+
+        if args.mask_edges and args.min_boundary_distance == 0:
+            max_shape = np.max(template.shape)
+            args.min_boundary_distance = np.ceil(np.divide(max_shape, 2))
+
+        target_mask = 1
+        if args.target_mask is not None:
+            target_mask = Density.from_file(args.target_mask).data
+        elif cli_args.target_mask is not None:
+            target_mask = Density.from_file(args.target_mask).data
+
+        mask = np.ones_like(data[0])
+        np.multiply(mask, target_mask, out=mask)
+
+        cropped_shape = np.subtract(
+            mask.shape, np.multiply(args.min_boundary_distance, 2)
+        ).astype(int)
+        mask[cropped_shape] = 0
+        _normalize_under_mask(template=data[0], mask=mask, mask_intensity=mask.sum())
+        return data[0]
+
+    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)
+        indices = new_scores > data[0]
+        entities[indices] = index + 1
+        data[0][indices] = new_scores[indices]
+
+    return data, entities
 
 
 def main():
     args = parse_args()
-    data = load_pickle(args.input_file)
+    data = load_pickle(args.input_file[0])
 
-    meta = data[-1]
-    target_origin, _, sampling_rate, cli_args = meta
+    target_origin, _, sampling_rate, cli_args = data[-1]
 
-    if args.orientations is not None:
-        orientations = Orientations.from_file(
-            filename=args.orientations, file_format="text"
+    _, template_extension = splitext(cli_args.template)
+    ret = load_template(
+        filepath=cli_args.template,
+        sampling_rate=sampling_rate,
+        center=not cli_args.no_centering,
+    )
+    template, center_of_mass, translation, template_is_density = ret
+
+    if args.output_format == "relion" and args.subtomogram_box_size is None:
+        new_shape = np.add(template.shape, np.mod(template.shape, 2))
+        new_shape = np.repeat(new_shape.max(), new_shape.size).astype(int)
+        print(f"Padding template from {template.shape} to {new_shape} for RELION.")
+        template.pad(new_shape)
+
+    template_mask = template.empty
+    template_mask.data[:] = 1
+    if cli_args.template_mask is not None:
+        template_mask = Density.from_file(cli_args.template_mask)
+        template_mask.pad(template.shape, center=False)
+        origin_translation = np.divide(
+            np.subtract(template.origin, template_mask.origin), template.sampling_rate
         )
+        translation = np.add(translation, origin_translation)
 
-    else:
+        template_mask = template_mask.rigid_transform(
+            rotation_matrix=np.eye(template_mask.data.ndim),
+            translation=-translation,
+            order=1,
+        )
+
+    if args.mask_edges and args.min_boundary_distance == 0:
+        max_shape = np.max(template.shape)
+        args.min_boundary_distance = np.ceil(np.divide(max_shape, 2))
+
+    # data, entities = merge_outputs(data=data, filepaths=args.input_file[1:], args=args)
+
+    orientations = args.orientations
+    if orientations is None:
         translations, rotations, scores, details = [], [], [], []
         # Output is MaxScoreOverRotations
         if data[0].ndim == data[2].ndim:
             scores, offset, rotation_array, rotation_mapping, meta = data
-            if args.mask_edges:
-                template, center_of_mass = load_template(
-                    cli_args.template, sampling_rate=sampling_rate
+
+            if args.target_mask is not None:
+                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 = tuple(
+                    slice(
+                        int(args.min_boundary_distance),
+                        int(x - args.min_boundary_distance),
+                    )
+                    for x in scores.shape
                 )
-                if not cli_args.no_centering:
-                    template, *_ = template.centered(0)
-                mask_size = template.shape
-                if args.min_boundary_distance > 0:
-                    mask_size = 2 * args.min_boundary_distance
-                scores = centered_mask(scores, np.subtract(scores.shape, mask_size) + 1)
+                # Rickgauer et al. 2017
+                n_correlations = np.size(scores[cropped_shape]) * len(rotation_mapping)
+                minimum_score = np.multiply(
+                    erfcinv(2 * args.n_false_positives / n_correlations),
+                    np.sqrt(2) * np.std(scores[cropped_shape]),
+                )
+                print(f"Determined minimum score cutoff: {minimum_score}.")
+                minimum_score = max(minimum_score, 0)
+                args.minimum_score = minimum_score
 
             peak_caller = PEAK_CALLERS[args.peak_caller](
                 number_of_peaks=args.number_of_peaks,
                 min_distance=args.min_distance,
                 min_boundary_distance=args.min_boundary_distance,
             )
-            peak_caller(scores, rotation_matrix=np.eye(3))
+            if args.minimum_score is not None:
+                args.number_of_peaks = np.inf
+
+            peak_caller(
+                scores,
+                rotation_matrix=np.eye(3),
+                mask=template.data,
+                rotation_mapping=rotation_mapping,
+                rotation_array=rotation_array,
+                minimum_score=args.minimum_score,
+            )
             candidates = peak_caller.merge(
                 candidates=[tuple(peak_caller)],
                 number_of_peaks=args.number_of_peaks,
@@ -466,16 +369,15 @@ def main():
                 min_boundary_distance=args.min_boundary_distance,
             )
             if len(candidates) == 0:
-                exit(
-                    "Found no peaks. Try reducing min_distance or min_boundary_distance."
-                )
+                print("Found no peaks. Consider changing peak calling parameters.")
+                exit(-1)
 
             for translation, _, score, detail in zip(*candidates):
                 rotations.append(rotation_mapping[rotation_array[tuple(translation)]])
 
         else:
             candidates = data
-            translation, rotation, score, detail, *_ = data
+            translation, rotation, *_ = data
             for i in range(translation.shape[0]):
                 rotations.append(euler_from_rotationmatrix(rotation[i]))
 
@@ -488,25 +390,35 @@ def main():
             details=details,
         )
 
+    if args.minimum_score is not None:
+        keep = orientations.scores >= args.minimum_score
+        orientations = orientations[keep]
+
+    if args.maximum_score is not None:
+        keep = orientations.scores <= args.maximum_score
+        orientations = orientations[keep]
+
     if args.output_format == "orientations":
         orientations.to_file(filename=f"{args.output_prefix}.tsv", file_format="text")
         exit(0)
 
-    _, template_extension = splitext(cli_args.template)
-    template, center_of_mass = load_template(
-        filepath=cli_args.template, sampling_rate=sampling_rate
-    )
-    template_is_density, index = isinstance(template, Density), 0
-
-    if args.output_format == "relion":
-        new_shape = np.add(template.shape, np.mod(template.shape, 2))
-        new_shape = np.repeat(new_shape.max(), new_shape.size).astype(int)
-        print(f"Padding template from {template.shape} to {new_shape} for RELION.")
-        template.pad(new_shape)
+    target = Density.from_file(cli_args.target)
+    if args.invert_target_contrast:
+        if args.output_format == "relion":
+            target.data = target.data * -1
+            target.data = np.divide(
+                np.subtract(target.data, target.data.mean()), target.data.std()
+            )
+        else:
+            target.data = (
+                -np.divide(
+                    np.subtract(target.data, target.data.min()),
+                    np.subtract(target.data.max(), target.data.min()),
+                )
+                + 1
+            )
 
     if args.output_format in ("extraction", "relion"):
-        target = Density.from_file(cli_args.target)
-
         if not np.all(np.divide(target.shape, template.shape) > 2):
             print(
                 "Target might be too small relative to template to extract"
@@ -514,26 +426,19 @@ def main():
                 f" Target : {target.shape}, template : {template.shape}."
             )
 
-        peaks = orientations.translations.astype(int)
-        max_shape = np.max(template.shape).astype(int)
-        half_shape = max_shape // 2
-
-        left_pad = half_shape
-        right_pad = np.add(half_shape, max_shape % 2)
-        starts = np.subtract(peaks, left_pad)
-        stops = np.add(peaks, right_pad)
-
-        candidate_starts = np.maximum(starts, 0).astype(int)
-        candidate_stops = np.minimum(stops, target.shape).astype(int)
-        keep_peaks = (
-            np.sum(
-                np.multiply(starts == candidate_starts, stops == candidate_stops),
-                axis=1,
+        extraction_shape = template.shape
+        if args.subtomogram_box_size is not None:
+            extraction_shape = np.repeat(
+                args.subtomogram_box_size, len(extraction_shape)
             )
-            == peaks.shape[1]
+
+        orientations, cand_slices, obs_slices = orientations.get_extraction_slices(
+            target_shape=target.shape,
+            extraction_shape=extraction_shape,
+            drop_out_of_box=True,
+            return_orientations=True,
         )
 
-        orientations = orientations[keep_peaks]
         working_directory = getcwd()
         if args.output_format == "relion":
             orientations.to_file(
@@ -542,62 +447,101 @@ def main():
                 name_prefix=join(working_directory, args.output_prefix),
                 ctf_image=args.wedge_mask,
                 sampling_rate=target.sampling_rate.max(),
-                subtomogram_size=template.shape[0],
+                subtomogram_size=extraction_shape[0],
             )
 
-        peaks = peaks[keep_peaks,]
-        starts = starts[keep_peaks,]
-        stops = stops[keep_peaks,]
-        candidate_starts = candidate_starts[keep_peaks,]
-        candidate_stops = candidate_stops[keep_peaks,]
-
-        if not len(peaks):
-            print(
-                "No peak remaining after filtering. Started with"
-                f" {orientations.translations.shape[0]} filtered to {peaks.shape[0]}."
-                " Consider reducing min_distance, increase num_peaks or use"
-                " a different peak caller."
-            )
-            exit(-1)
-
-        observation_starts = np.subtract(candidate_starts, starts).astype(int)
-        observation_stops = np.subtract(np.add(max_shape, candidate_stops), stops)
-        observation_stops = observation_stops.astype(int)
-
-        candidate_slices = [
-            tuple(slice(s, e) for s, e in zip(start_row, stop_row))
-            for start_row, stop_row in zip(candidate_starts, candidate_stops)
-        ]
-
-        observation_slices = [
-            tuple(slice(s, e) for s, e in zip(start_row, stop_row))
-            for start_row, stop_row in zip(observation_starts, observation_stops)
-        ]
-        observations = np.zeros(
-            (len(candidate_slices), max_shape, max_shape, max_shape)
-        )
-
-        slices = zip(candidate_slices, observation_slices)
+        observations = np.zeros((len(cand_slices), *extraction_shape))
+        slices = zip(cand_slices, obs_slices)
         for idx, (cand_slice, obs_slice) in enumerate(slices):
-            observations[idx][:] = np.mean(target.data[cand_slice])
-            observations[idx][obs_slice] = target.data[cand_slice]
+            observations[idx][:] = np.mean(target.data[obs_slice])
+            observations[idx][cand_slice] = target.data[obs_slice]
 
         for index in range(observations.shape[0]):
+            cand_start = [x.start for x in cand_slices[index]]
             out_density = Density(
                 data=observations[index],
                 sampling_rate=sampling_rate,
-                origin=candidate_starts[index] * sampling_rate,
+                origin=np.multiply(cand_start, sampling_rate),
             )
-            # out_density.data = out_density.data * template_mask.data
+            if args.mask_subtomograms:
+                rotation_matrix = euler_to_rotationmatrix(orientations.rotations[index])
+                mask_transfomed = template_mask.rigid_transform(
+                    rotation_matrix=rotation_matrix, order=1
+                )
+                out_density.data = out_density.data * mask_transfomed.data
             out_density.to_file(
                 join(working_directory, f"{args.output_prefix}_{index}.mrc")
             )
 
         exit(0)
 
-    for translation, angles, *_ in orientations:
-        rotation_matrix = euler_to_rotationmatrix(angles)
+    if args.output_format == "backmapping":
+        orientations, cand_slices, obs_slices = orientations.get_extraction_slices(
+            target_shape=target.shape,
+            extraction_shape=template.shape,
+            drop_out_of_box=True,
+            return_orientations=True,
+        )
+        ret, template_sum = target.empty, template.data.sum()
+        for index in range(len(cand_slices)):
+            rotation_matrix = euler_to_rotationmatrix(orientations.rotations[index])
 
+            transformed_template = template.rigid_transform(
+                rotation_matrix=rotation_matrix
+            )
+            transformed_template.data = np.multiply(
+                transformed_template.data,
+                np.divide(template_sum, transformed_template.data.sum()),
+            )
+            cand_slice, obs_slice = cand_slices[index], obs_slices[index]
+            ret.data[obs_slice] += transformed_template.data[cand_slice]
+        ret.to_file(f"{args.output_prefix}_backmapped.mrc")
+        exit(0)
+
+    if args.output_format == "average":
+        orientations, cand_slices, obs_slices = orientations.get_extraction_slices(
+            target_shape=target.shape,
+            extraction_shape=np.multiply(template.shape, 2),
+            drop_out_of_box=True,
+            return_orientations=True,
+        )
+        out = np.zeros_like(template.data)
+        out = np.zeros(np.multiply(template.shape, 2).astype(int))
+        for index in range(len(cand_slices)):
+            from scipy.spatial.transform import Rotation
+
+            rotation = Rotation.from_euler(
+                angles=orientations.rotations[index], seq="zyx", degrees=True
+            )
+            rotation_matrix = rotation.inv().as_matrix()
+
+            # rotation_matrix = euler_to_rotationmatrix(orientations.rotations[index])
+            subset = Density(target.data[obs_slices[index]])
+            subset = subset.rigid_transform(rotation_matrix=rotation_matrix, order=1)
+
+            np.add(out, subset.data, out=out)
+        out /= len(cand_slices)
+        ret = Density(out, sampling_rate=template.sampling_rate, origin=0)
+        ret.pad(template.shape, center=True)
+        ret.to_file(f"{args.output_prefix}_average.mrc")
+        exit(0)
+
+    if args.peak_oversampling > 1:
+        peak_caller = peak_caller = PEAK_CALLERS[args.peak_caller]()
+        if data[0].ndim != data[2].ndim:
+            print(
+                "Input pickle does not contain template matching scores."
+                " Cannot oversample peaks."
+            )
+            exit(-1)
+        orientations.translations = peak_caller.oversample_peaks(
+            score_space=data[0],
+            translations=orientations.translations,
+            oversampling_factor=args.oversampling_factor,
+        )
+
+    for index, (translation, angles, *_) in enumerate(orientations):
+        rotation_matrix = euler_to_rotationmatrix(angles)
         if template_is_density:
             translation = np.subtract(translation, center_of_mass)
             transformed_template = template.rigid_transform(
@@ -606,6 +550,7 @@ def main():
             new_origin = np.add(target_origin / sampling_rate, translation)
             transformed_template.origin = np.multiply(new_origin, sampling_rate)
         else:
+            template = Structure.from_file(cli_args.template)
             new_center_of_mass = np.add(
                 np.multiply(translation, sampling_rate), target_origin
             )
@@ -614,7 +559,7 @@ def main():
                 translation=translation[::-1],
                 rotation_matrix=rotation_matrix[::-1, ::-1],
             )
-        # template_extension should contain the extension '.'
+        # template_extension should contain '.'
         transformed_template.to_file(
             f"{args.output_prefix}_{index}{template_extension}"
         )