pytme 0.1.8__cp311-cp311-macosx_14_0_arm64.whl → 0.2.0__cp311-cp311-macosx_14_0_arm64.whl

tme/analyzer.py

@@ -1,4 +1,4 @@
-""" Implements classes to analyze score spaces from systematic fitting.
+""" Implements classes to analyze outputs from exhaustive template matching.
 
     Copyright (c) 2023 European Molecular Biology Laboratory
 
@@ -15,20 +15,22 @@ from numpy.typing import NDArray
 from scipy.stats import entropy
 from sklearn.cluster import DBSCAN
 from skimage.feature import peak_local_max
-
-from .extensions import max_index_by_label, online_statistics
+from skimage.registration._phase_cross_correlation import _upsampled_dft
+from .extensions import max_index_by_label, online_statistics, find_candidate_indices
 from .matching_utils import (
     split_numpy_array_slices,
     array_to_memmap,
     generate_tempfile_name,
+    euler_to_rotationmatrix,
+    apply_convolution_mode,
 )
-
 from .backends import backend
 
 
-def filter_points_indices(coordinates: NDArray, min_distance: Tuple[int]):
-    if min_distance <= 0:
-        return backend.arange(coordinates.shape[0])
+def filter_points_indices_bucket(
+    coordinates: NDArray, min_distance: Tuple[int]
+) -> NDArray:
+    coordinates = backend.subtract(coordinates, backend.min(coordinates, axis=0))
     bucket_indices = backend.astype(backend.divide(coordinates, min_distance), int)
     multiplier = backend.power(
         backend.max(bucket_indices, axis=0) + 1, backend.arange(bucket_indices.shape[1])
@@ -40,7 +42,24 @@ def filter_points_indices(coordinates: NDArray, min_distance: Tuple[int]):
     return unique_indices
 
 
-def filter_points(coordinates: NDArray, min_distance: Tuple[int]):
+def filter_points_indices(
+    coordinates: NDArray, min_distance: float, bucket_cutoff: int = 1e4
+) -> NDArray:
+    if min_distance <= 0:
+        return backend.arange(coordinates.shape[0])
+
+    if isinstance(coordinates, np.ndarray):
+        return find_candidate_indices(coordinates, min_distance)
+    elif coordinates.shape[0] > bucket_cutoff:
+        return filter_points_indices_bucket(coordinates, min_distance)
+    distances = np.linalg.norm(coordinates[:, None] - coordinates, axis=-1)
+    distances = np.tril(distances)
+    keep = np.sum(distances > min_distance, axis=1)
+    indices = np.arange(coordinates.shape[0])
+    return indices[keep == indices]
+
+
+def filter_points(coordinates: NDArray, min_distance: Tuple[int]) -> NDArray:
     unique_indices = filter_points_indices(coordinates, min_distance)
     coordinates = coordinates[unique_indices]
     return coordinates
@@ -95,6 +114,12 @@ class PeakCaller(ABC):
         self.min_boundary_distance = min_boundary_distance
         self.number_of_peaks = number_of_peaks
 
+        # Postprocesing arguments
+        self.fourier_shift = kwargs.get("fourier_shift", None)
+        self.convolution_mode = kwargs.get("convolution_mode", None)
+        self.targetshape = kwargs.get("targetshape", None)
+        self.templateshape = kwargs.get("templateshape", None)
+
     def __iter__(self):
         """
         Returns a generator to list objects containing translation,
@@ -104,7 +129,10 @@ class PeakCaller(ABC):
         yield from self.peak_list
 
     def __call__(
-        self, score_space: NDArray, rotation_matrix: NDArray, **kwargs
+        self,
+        score_space: NDArray,
+        rotation_matrix: NDArray,
+        **kwargs,
     ) -> None:
         """
         Update the internal parameter store based on input array.
@@ -116,7 +144,7 @@ class PeakCaller(ABC):
         rotation_matrix : NDArray
             Rotation matrix used to obtain the score array.
         **kwargs
-            Additional keyword arguments.
+            Optional keyword arguments passed to :py:meth:`PeakCaller.call_peak`.
         """
         peak_positions, peak_details = self.call_peaks(
             score_space=score_space, rotation_matrix=rotation_matrix, **kwargs
@@ -159,11 +187,12 @@ class PeakCaller(ABC):
             peak_positions.shape[0],
             axis=0,
         )
+        peak_scores = score_space[tuple(peak_positions.T)]
 
         self._update(
             peak_positions=peak_positions,
             peak_details=peak_details,
-            peak_scores=score_space[tuple(peak_positions.T)],
+            peak_scores=peak_scores,
             rotations=rotations,
             **kwargs,
         )
@@ -190,10 +219,8 @@ class PeakCaller(ABC):
 
         Returns
         -------
-        NDArray
-            Array of peak coordiantes.
-        NDArray
-            Array of peak details.
+        Tuple[NDArray, NDArray]
+            Array of peak coordinates and peak details.
         """
 
     @classmethod
@@ -231,6 +258,79 @@ class PeakCaller(ABC):
             )
         return tuple(base)
 
+    @staticmethod
+    def oversample_peaks(
+        score_space: NDArray, peak_positions: NDArray, oversampling_factor: int = 8
+    ):
+        """
+        Refines peaks positions in the corresponding score space.
+
+        Parameters
+        ----------
+        score_space : NDArray
+            The d-dimensional array representing the score space.
+        peak_positions : NDArray
+            An array of shape (n, d) containing the peak coordinates
+            to be refined, where n is the number of peaks and d is the
+            dimensionality of the score space.
+        oversampling_factor : int, optional
+            The oversampling factor for Fourier transforms. Defaults to 8.
+
+        Returns
+        -------
+        NDArray
+            An array of shape (n, d) containing the refined subpixel
+            coordinates of the peaks.
+
+        Notes
+        -----
+        Floating point peak positions are determined by oversampling the
+        score_space around peak_positions. The accuracy
+        of refinement scales with 1 / oversampling_factor.
+
+        References
+        ----------
+        .. [1]  https://scikit-image.org/docs/stable/api/skimage.registration.html
+        .. [2]  Manuel Guizar-Sicairos, Samuel T. Thurman, and
+                James R. Fienup, “Efficient subpixel image registration
+                algorithms,” Optics Letters 33, 156-158 (2008).
+                DOI:10.1364/OL.33.000156
+
+        """
+        score_space = backend.to_numpy_array(score_space)
+        peak_positions = backend.to_numpy_array(peak_positions)
+
+        peak_positions = np.round(
+            np.divide(
+                np.multiply(peak_positions, oversampling_factor), oversampling_factor
+            )
+        )
+        upsampled_region_size = np.ceil(np.multiply(oversampling_factor, 1.5))
+        dftshift = np.round(np.divide(upsampled_region_size, 2.0))
+        sample_region_offset = np.subtract(
+            dftshift, np.multiply(peak_positions, oversampling_factor)
+        )
+
+        score_space_ft = np.fft.fftn(score_space).conj()
+        for index in range(sample_region_offset.shape[0]):
+            cross_correlation_upsampled = _upsampled_dft(
+                data=score_space_ft,
+                upsampled_region_size=upsampled_region_size,
+                upsample_factor=oversampling_factor,
+                axis_offsets=sample_region_offset[index],
+            ).conj()
+
+            maxima = np.unravel_index(
+                np.argmax(np.abs(cross_correlation_upsampled)),
+                cross_correlation_upsampled.shape,
+            )
+            maxima = np.divide(np.subtract(maxima, dftshift), oversampling_factor)
+            peak_positions[index] = np.add(peak_positions[index], maxima)
+
+        peak_positions = backend.to_backend_array(peak_positions)
+
+        return peak_positions
+
     def _update(
         self,
         peak_positions: NDArray,
@@ -287,6 +387,61 @@ class PeakCaller(ABC):
         self.peak_list[2] = peak_scores[final_order]
         self.peak_list[3] = peak_details[final_order]
 
+    def _postprocess(
+        self, fourier_shift, convolution_mode, targetshape, templateshape, **kwargs
+    ):
+        peak_positions = self.peak_list[0]
+        if not len(peak_positions):
+            return self
+
+        if targetshape is None or templateshape is None:
+            return self
+
+        # Remove padding to next fast fourier length
+        score_space_shape = backend.add(targetshape, templateshape) - 1
+
+        if fourier_shift is not None:
+            peak_positions = backend.add(peak_positions, fourier_shift)
+            backend.divide(peak_positions, score_space_shape).astype(int)
+
+            backend.subtract(
+                peak_positions,
+                backend.multiply(
+                    backend.divide(peak_positions, score_space_shape).astype(int),
+                    score_space_shape,
+                ),
+                out=peak_positions,
+            )
+
+        if convolution_mode is None:
+            return None
+
+        if convolution_mode == "full":
+            output_shape = score_space_shape
+        elif convolution_mode == "same":
+            output_shape = targetshape
+        elif convolution_mode == "valid":
+            output_shape = backend.add(
+                backend.subtract(targetshape, templateshape),
+                backend.mod(templateshape, 2),
+            )
+
+        output_shape = backend.to_backend_array(output_shape)
+        starts = backend.divide(backend.subtract(score_space_shape, output_shape), 2)
+        starts = backend.astype(starts, int)
+        stops = backend.add(starts, output_shape)
+
+        valid_peaks = (
+            backend.sum(
+                backend.multiply(peak_positions > starts, peak_positions <= stops),
+                axis=1,
+            )
+            == peak_positions.shape[1]
+        )
+        self.peak_list[0] = backend.subtract(peak_positions, starts)
+        self.peak_list = [x[valid_peaks] for x in self.peak_list]
+        return self
+
 
 class PeakCallerSort(PeakCaller):
     """
@@ -297,7 +452,7 @@ class PeakCallerSort(PeakCaller):
     """
 
     def call_peaks(
-        self, score_space: NDArray, rotation_matrix: NDArray, **kwargs
+        self, score_space: NDArray, minimum_score: float = None, **kwargs
     ) -> Tuple[NDArray, NDArray]:
         """
         Call peaks in the score space.
@@ -307,20 +462,20 @@ class PeakCallerSort(PeakCaller):
         score_space : NDArray
             Data array of scores.
         minimum_score : float
-            Minimum score value to consider.
-        min_distance : float
-            Minimum distance between maxima.
+            Minimum score value to consider. If provided, superseeds limit given
+            by :py:attr:`PeakCaller.number_of_peaks`.
 
         Returns
         -------
-        NDArray
-            Array of peak coordiantes.
-        NDArray
-            Array of peak details.
+        Tuple[NDArray, NDArray]
+            Array of peak coordinates and peak details.
         """
         flat_score_space = score_space.reshape(-1)
         k = min(self.number_of_peaks, backend.size(flat_score_space))
 
+        if minimum_score is not None:
+            k = backend.sum(score_space >= minimum_score)
+
         top_k_indices, *_ = backend.topk_indices(flat_score_space, k)
 
         coordinates = backend.unravel_index(top_k_indices, score_space.shape)
@@ -333,12 +488,12 @@ class PeakCallerSort(PeakCaller):
 class PeakCallerMaximumFilter(PeakCaller):
     """
     Find local maxima by applying a maximum filter and enforcing a distance
-    constraint subseqquently. This is similar to the strategy implemented in
+    constraint subsequently. This is similar to the strategy implemented in
     skimage.feature.peak_local_max.
     """
 
     def call_peaks(
-        self, score_space: NDArray, rotation_matrix: NDArray, **kwargs
+        self, score_space: NDArray, minimum_score: float = None, **kwargs
     ) -> Tuple[NDArray, NDArray]:
         """
         Call peaks in the score space.
@@ -348,25 +503,25 @@ class PeakCallerMaximumFilter(PeakCaller):
         score_space : NDArray
             Data array of scores.
         minimum_score : float
-            Minimum score value to consider.
-        min_distance : float
-            Minimum distance between maxima.
+            Minimum score value to consider. If provided, superseeds limit given
+            by :py:attr:`PeakCaller.number_of_peaks`.
 
         Returns
         -------
-        NDArray
-            Array of peak coordiantes.
-        NDArray
-            Array of peak details.
+        Tuple[NDArray, NDArray]
+            Array of peak coordinates and peak details.
         """
         peaks = backend.max_filter_coordinates(score_space, self.min_distance)
 
+        scores = score_space[tuple(peaks.T)]
+
         input_candidates = min(
             self.number_of_peaks, peaks.shape[0] - 1, backend.size(score_space) - 1
         )
-        top_indices = backend.topk_indices(
-            score_space[tuple(peaks.T)], input_candidates
-        )
+        if minimum_score is not None:
+            input_candidates = backend.sum(scores >= minimum_score)
+
+        top_indices = backend.topk_indices(scores, input_candidates)
         peaks = peaks[top_indices]
 
         return peaks, None
@@ -382,7 +537,7 @@ class PeakCallerFast(PeakCaller):
     """
 
     def call_peaks(
-        self, score_space: NDArray, rotation_matrix: NDArray, **kwargs
+        self, score_space: NDArray, minimum_score: float = None, **kwargs
     ) -> Tuple[NDArray, NDArray]:
         """
         Call peaks in the score space.
@@ -392,16 +547,13 @@ class PeakCallerFast(PeakCaller):
         score_space : NDArray
             Data array of scores.
         minimum_score : float
-            Minimum score value to consider.
-        min_distance : float
-            Minimum distance between maxima.
+            Minimum score value to consider. If provided, superseeds limit given
+            by :py:attr:`PeakCaller.number_of_peaks`.
 
         Returns
         -------
-        NDArray
-            Array of peak coordiantes.
-        NDArray
-            Array of peak details.
+        Tuple[NDArray, NDArray]
+            Array of peak coordinates and peak details.
         """
         splits = {
             axis: score_space.shape[axis] // self.min_distance
@@ -457,7 +609,14 @@ class PeakCallerRecursiveMasking(PeakCaller):
     """
 
     def call_peaks(
-        self, score_space: NDArray, rotation_matrix: NDArray, **kwargs
+        self,
+        score_space: NDArray,
+        rotation_matrix: NDArray,
+        mask: NDArray = None,
+        minimum_score: float = None,
+        rotation_space: NDArray = None,
+        rotation_mapping: Dict = None,
+        **kwargs,
     ) -> Tuple[NDArray, NDArray]:
         """
         Call peaks in the score space.
@@ -466,36 +625,195 @@ class PeakCallerRecursiveMasking(PeakCaller):
         ----------
         score_space : NDArray
             Data array of scores.
+        rotation_matrix : NDArray
+            Rotation matrix.
+        mask : NDArray, optional
+            Mask array, by default None.
+        rotation_space : NDArray, optional
+            Rotation space array, by default None.
+        rotation_mapping : Dict optional
+            Dictionary mapping values in rotation_space to Euler angles.
+            By default None
         minimum_score : float
-            Minimum score value to consider.
-        min_distance : float
-            Minimum distance between maxima.
+            Minimum score value to consider. If provided, superseeds limit given
+            by :py:attr:`PeakCaller.number_of_peaks`.
 
         Returns
         -------
-        NDArray
-            Array of peak coordiantes.
-        NDArray
-            Array of peak details.
+        Tuple[NDArray, NDArray]
+            Array of peak coordinates and peak details.
+
+        Notes
+        -----
+        By default, scores are masked using a box with edge length self.min_distance.
+        If mask is provided, elements around each peak will be multiplied by the mask
+        values. If rotation_space and rotation_mapping is provided, the respective
+        rotation will be applied to the mask, otherwise rotation_matrix is used.
         """
-        score_box = tuple(self.min_distance for _ in range(score_space.ndim))
-        coordinates = []
+        coordinates, masking_function = [], self._mask_scores_rotate
+
+        if mask is None:
+            masking_function = self._mask_scores_box
+            shape = tuple(self.min_distance for _ in range(score_space.ndim))
+            mask = backend.zeros(shape, dtype=backend._default_dtype)
+
+        rotated_template = backend.zeros(mask.shape, dtype=mask.dtype)
+
+        peak_limit = self.number_of_peaks
+        if minimum_score is not None:
+            peak_limit = backend.size(score_space)
+        else:
+            minimum_score = backend.min(score_space) - 1
+
         while True:
             backend.argmax(score_space)
-            max_coord = backend.unravel_index(
+            peak = backend.unravel_index(
                 indices=backend.argmax(score_space), shape=score_space.shape
             )
-            coordinates.append(max_coord)
-            start = backend.maximum(backend.subtract(max_coord, score_box), 0)
-            stop = backend.minimum(backend.add(max_coord, score_box), score_space.shape)
-            start, stop = backend.astype(start, int), backend.astype(stop, int)
-            coords = tuple(slice(*pos) for pos in zip(start, stop))
-            score_space[coords] = 0
-            if len(coordinates) >= self.number_of_peaks:
+            if score_space[tuple(peak)] < minimum_score:
+                break
+
+            coordinates.append(peak)
+
+            current_rotation_matrix = self._get_rotation_matrix(
+                peak=peak,
+                rotation_space=rotation_space,
+                rotation_mapping=rotation_mapping,
+                rotation_matrix=rotation_matrix,
+            )
+
+            masking_function(
+                score_space=score_space,
+                rotation_matrix=current_rotation_matrix,
+                peak=peak,
+                mask=mask,
+                rotated_template=rotated_template,
+            )
+
+            if len(coordinates) >= peak_limit:
                 break
+
         peaks = backend.to_backend_array(coordinates)
         return peaks, None
 
+    @staticmethod
+    def _get_rotation_matrix(
+        peak: NDArray,
+        rotation_space: NDArray,
+        rotation_mapping: NDArray,
+        rotation_matrix: NDArray,
+    ) -> NDArray:
+        """
+        Get rotation matrix based on peak and rotation data.
+
+        Parameters
+        ----------
+        peak : NDArray
+            Peak coordinates.
+        rotation_space : NDArray
+            Rotation space array.
+        rotation_mapping : Dict
+            Dictionary mapping values in rotation_space to Euler angles.
+        rotation_matrix : NDArray
+            Current rotation matrix.
+
+        Returns
+        -------
+        NDArray
+            Rotation matrix.
+        """
+        if rotation_space is None or rotation_mapping is None:
+            return rotation_matrix
+
+        rotation = rotation_mapping[rotation_space[tuple(peak)]]
+
+        rotation_matrix = backend.to_backend_array(
+            euler_to_rotationmatrix(backend.to_numpy_array(rotation))
+        )
+        return rotation_matrix
+
+    @staticmethod
+    def _mask_scores_box(
+        score_space: NDArray, peak: NDArray, mask: NDArray, **kwargs: Dict
+    ) -> None:
+        """
+        Mask scores in a box around a peak.
+
+        Parameters
+        ----------
+        score_space : NDArray
+            Data array of scores.
+        peak : NDArray
+            Peak coordinates.
+        mask : NDArray
+            Mask array.
+        """
+        start = backend.maximum(backend.subtract(peak, mask.shape), 0)
+        stop = backend.minimum(backend.add(peak, mask.shape), score_space.shape)
+        start, stop = backend.astype(start, int), backend.astype(stop, int)
+        coords = tuple(slice(*pos) for pos in zip(start, stop))
+        score_space[coords] = 0
+        return None
+
+    @staticmethod
+    def _mask_scores_rotate(
+        score_space: NDArray,
+        peak: NDArray,
+        mask: NDArray,
+        rotated_template: NDArray,
+        rotation_matrix: NDArray,
+        **kwargs: Dict,
+    ) -> None:
+        """
+        Mask score_space using mask rotation around a peak.
+
+        Parameters
+        ----------
+        score_space : NDArray
+            Data array of scores.
+        peak : NDArray
+            Peak coordinates.
+        mask : NDArray
+            Mask array.
+        rotated_template : NDArray
+            Empty array to write mask rotations to.
+        rotation_matrix : NDArray
+            Rotation matrix.
+        """
+        left_pad = backend.divide(mask.shape, 2).astype(int)
+        right_pad = backend.add(left_pad, backend.mod(mask.shape, 2).astype(int))
+
+        score_start = backend.subtract(peak, left_pad)
+        score_stop = backend.add(peak, right_pad)
+
+        template_start = backend.subtract(backend.maximum(score_start, 0), score_start)
+        template_stop = backend.subtract(
+            score_stop, backend.minimum(score_stop, score_space.shape)
+        )
+        template_stop = backend.subtract(mask.shape, template_stop)
+
+        score_start = backend.maximum(score_start, 0)
+        score_stop = backend.minimum(score_stop, score_space.shape)
+        score_start = backend.astype(score_start, int)
+        score_stop = backend.astype(score_stop, int)
+
+        template_start = backend.astype(template_start, int)
+        template_stop = backend.astype(template_stop, int)
+        coords_score = tuple(slice(*pos) for pos in zip(score_start, score_stop))
+        coords_template = tuple(
+            slice(*pos) for pos in zip(template_start, template_stop)
+        )
+
+        rotated_template.fill(0)
+        backend.rotate_array(
+            arr=mask, rotation_matrix=rotation_matrix, order=1, out=rotated_template
+        )
+
+        score_space[coords_score] = backend.multiply(
+            score_space[coords_score], (rotated_template[coords_template] <= 0.1)
+        )
+        return None
+
 
 class PeakCallerScipy(PeakCaller):
     """
@@ -503,7 +821,7 @@ class PeakCallerScipy(PeakCaller):
     """
 
     def call_peaks(
-        self, score_space: NDArray, rotation_matrix: NDArray, **kwargs
+        self, score_space: NDArray, minimum_score: float = None, **kwargs
     ) -> Tuple[NDArray, NDArray]:
         """
         Call peaks in the score space.
@@ -513,21 +831,25 @@ class PeakCallerScipy(PeakCaller):
         score_space : NDArray
             Data array of scores.
         minimum_score : float
-            Minimum score value to consider.
-        min_distance : float
-            Minimum distance between maxima.
+            Minimum score value to consider. If provided, superseeds limit given
+            by :py:attr:`PeakCaller.number_of_peaks`.
 
         Returns
         -------
-        NDArray
-            Array of peak coordiantes.
-        NDArray
-            Array of peak details.
+        Tuple[NDArray, NDArray]
+            Array of peak coordinates and peak details.
         """
+
+        score_space = backend.to_numpy_array(score_space)
+        num_peaks = self.number_of_peaks
+        if minimum_score is not None:
+            num_peaks = np.inf
+
         peaks = peak_local_max(
             score_space,
-            num_peaks=self.number_of_peaks,
+            num_peaks=num_peaks,
             min_distance=self.min_distance,
+            threshold_abs=minimum_score,
         )
         return peaks, None
 
@@ -879,8 +1201,63 @@ class MaxScoreOverRotations:
 
         self.use_memmap = use_memmap
         self.lock = Manager().Lock() if thread_safe else nullcontext()
+        self.lock_is_nullcontext = isinstance(
+            self.score_space, type(backend.zeros((1)))
+        )
         self.observed_rotations = Manager().dict() if thread_safe else {}
 
+    def _postprocess(
+        self,
+        fourier_shift,
+        convolution_mode,
+        targetshape,
+        templateshape,
+        shared_memory_handler=None,
+        **kwargs,
+    ):
+        internal_scores = backend.sharedarr_to_arr(
+            shape=self.score_space_shape,
+            dtype=self.score_space_dtype,
+            shm=self.score_space,
+        )
+        internal_rotations = backend.sharedarr_to_arr(
+            shape=self.score_space_shape,
+            dtype=self.rotation_space_dtype,
+            shm=self.rotations,
+        )
+
+        if fourier_shift is not None:
+            axis = tuple(i for i in range(len(fourier_shift)))
+            internal_scores = backend.roll(
+                internal_scores, shift=fourier_shift, axis=axis
+            )
+            internal_rotations = backend.roll(
+                internal_rotations, shift=fourier_shift, axis=axis
+            )
+
+        if convolution_mode is not None:
+            internal_scores = apply_convolution_mode(
+                internal_scores,
+                convolution_mode=convolution_mode,
+                s1=targetshape,
+                s2=templateshape,
+            )
+            internal_rotations = apply_convolution_mode(
+                internal_rotations,
+                convolution_mode=convolution_mode,
+                s1=targetshape,
+                s2=templateshape,
+            )
+
+        self.score_space_shape = internal_scores.shape
+        self.score_space = backend.arr_to_sharedarr(
+            internal_scores, shared_memory_handler
+        )
+        self.rotations = backend.arr_to_sharedarr(
+            internal_rotations, shared_memory_handler
+        )
+        return self
+
     def __iter__(self):
         internal_scores = backend.sharedarr_to_arr(
             shape=self.score_space_shape,
@@ -938,11 +1315,21 @@ class MaxScoreOverRotations:
         **kwargs
             Arbitrary keyword arguments.
         """
+        rotation = backend.tobytes(rotation_matrix)
+        rotation_index = self.observed_rotations.setdefault(
+            rotation, len(self.observed_rotations)
+        )
+
+        if self.lock_is_nullcontext:
+            backend.max_score_over_rotations(
+                score_space=score_space,
+                internal_scores=self.score_space,
+                internal_rotations=self.rotations,
+                rotation_index=rotation_index,
+            )
+            return None
+
         with self.lock:
-            rotation = backend.tobytes(rotation_matrix)
-            if rotation not in self.observed_rotations:
-                self.observed_rotations[rotation] = len(self.observed_rotations)
-            rotation_index = self.observed_rotations[rotation]
             internal_scores = backend.sharedarr_to_arr(
                 shape=self.score_space_shape,
                 dtype=self.score_space_dtype,
@@ -953,9 +1340,14 @@ class MaxScoreOverRotations:
                 dtype=self.rotation_space_dtype,
                 shm=self.rotations,
             )
-            indices = score_space > internal_scores
-            internal_scores[indices] = score_space[indices]
-            internal_rotations[indices] = rotation_index
+
+            backend.max_score_over_rotations(
+                score_space=score_space,
+                internal_scores=internal_scores,
+                internal_rotations=internal_rotations,
+                rotation_index=rotation_index,
+            )
+            return None
 
     @classmethod
     def merge(cls, param_stores=List[Tuple], **kwargs) -> Tuple[NDArray]:
@@ -1007,14 +1399,22 @@ class MaxScoreOverRotations:
             scores_out = np.memmap(
                 scores_out_filename, mode="w+", shape=base_max, dtype=scores_out_dtype
             )
+            scores_out.fill(kwargs.get("score_threshold", 0))
+            scores_out.flush()
             rotations_out = np.memmap(
                 rotations_out_filename,
                 mode="w+",
                 shape=base_max,
                 dtype=rotations_out_dtype,
             )
+            rotations_out.fill(-1)
+            rotations_out.flush()
         else:
-            scores_out = np.zeros(base_max, dtype=scores_out_dtype)
+            scores_out = np.full(
+                base_max,
+                fill_value=kwargs.get("score_threshold", 0),
+                dtype=scores_out_dtype,
+            )
             rotations_out = np.full(base_max, fill_value=-1, dtype=rotations_out_dtype)
 
         for i in range(len(param_stores)):
@@ -1079,6 +1479,61 @@ class MaxScoreOverRotations:
         )
 
 
+class _MaxScoreOverTranslations(MaxScoreOverRotations):
+    """
+    Obtain the maximum translation score over various rotations.
+
+    Attributes
+    ----------
+    score_space : NDArray
+        The score space for the observed rotations.
+    rotations : NDArray
+        The rotation identifiers for each score.
+    translation_offset : NDArray, optional
+        The offset applied during translation.
+    observed_rotations : int
+        Count of observed rotations.
+    use_memmap : bool, optional
+        Whether to offload internal data arrays to disk
+    thread_safe: bool, optional
+        Whether access to internal data arrays should be thread safe
+    """
+
+    def __call__(
+        self, score_space: NDArray, rotation_matrix: NDArray, **kwargs
+    ) -> None:
+        """
+        Update internal parameter store based on `score_space`.
+
+        Parameters
+        ----------
+        score_space : ndarray
+            Numpy array containing the score space.
+        rotation_matrix : ndarray
+            Square matrix describing the current rotation.
+        **kwargs
+            Arbitrary keyword arguments.
+        """
+        from tme.matching_utils import centered_mask
+
+        with self.lock:
+            rotation = backend.tobytes(rotation_matrix)
+            if rotation not in self.observed_rotations:
+                self.observed_rotations[rotation] = len(self.observed_rotations)
+            score_space = centered_mask(score_space, kwargs["template_shape"])
+            rotation_index = self.observed_rotations[rotation]
+            internal_scores = backend.sharedarr_to_arr(
+                shape=self.score_space_shape,
+                dtype=self.score_space_dtype,
+                shm=self.score_space,
+            )
+            max_score = score_space.max(axis=(1, 2, 3))
+            mean_score = score_space.mean(axis=(1, 2, 3))
+            std_score = score_space.std(axis=(1, 2, 3))
+            z_score = (max_score - mean_score) / std_score
+            internal_scores[rotation_index] = z_score
+
+
 class MemmapHandler:
     """
     Create numpy memmap objects to write score spaces to.