pytme 0.2.9.post1__cp311-cp311-macosx_15_0_arm64.whl → 0.3b0.post1__cp311-cp311-macosx_15_0_arm64.whl

tme/analyzer/aggregation.py

@@ -1,34 +1,34 @@
-""" Implements classes to analyze outputs from exhaustive template matching.
+"""
+Implements classes to analyze outputs from exhaustive template matching.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
-from contextlib import nullcontext
-from multiprocessing import Manager
-from typing import Tuple, List, Dict, Generator
+from typing import Tuple, List, Dict
 
 import numpy as np
 
+from .base import AbstractAnalyzer
 from ..types import BackendArray
 from ._utils import cart_to_score
 from ..backends import backend as be
 from ..matching_utils import (
     create_mask,
     array_to_memmap,
-    generate_tempfile_name,
     apply_convolution_mode,
+    generate_tempfile_name,
 )
 
-
 __all__ = [
     "MaxScoreOverRotations",
+    "MaxScoreOverRotationsConstrained",
     "MaxScoreOverTranslations",
 ]
 
 
-class MaxScoreOverRotations:
+class MaxScoreOverRotations(AbstractAnalyzer):
     """
     Determine the rotation maximizing the score over all possible translations.
 
@@ -36,10 +36,6 @@ class MaxScoreOverRotations:
     ----------
     shape : tuple of int
         Shape of array passed to :py:meth:`MaxScoreOverRotations.__call__`.
-    scores : BackendArray, optional
-        Array mapping translations to scores.
-    rotations : BackendArray, optional
-        Array mapping translations to rotation indices.
     offset : BackendArray, optional
         Coordinate origin considered during merging, zero by default.
     score_threshold : float, optional
@@ -50,58 +46,47 @@ class MaxScoreOverRotations:
         Memmap internal arrays, False by default.
     thread_safe: bool, optional
         Allow class to be modified by multiple processes, True by default.
-    only_unique_rotations : bool, optional
-        Whether each rotation will be shown only once, False by default.
-
-    Attributes
-    ----------
-    scores : BackendArray
-        Mapping of translations to scores.
-    rotations : BackendArray
-        Mmapping of translations to rotation indices.
-    rotation_mapping : Dict
-        Mapping of rotations to rotation indices.
-    offset : BackendArray, optional
-        Coordinate origin considered during merging, zero by default
+    inversion_mapping : bool, optional
+        Do not use rotation matrix bytestrings for intermediate data handling.
+        This is useful for GPU backend where analyzers are not shared across
+        devices and every rotation is only observed once. It is generally
+        safe to deactivate inversion mapping, but at a cost of performance.
 
     Examples
     --------
     The following achieves the minimal definition of a :py:class:`MaxScoreOverRotations`
     instance
 
+    >>> import numpy as np
     >>> from tme.analyzer import MaxScoreOverRotations
-    >>> analyzer = MaxScoreOverRotations(shape = (50, 50))
+    >>> analyzer = MaxScoreOverRotations(shape=(50, 50))
 
     The following simulates a template matching run by creating random data for a range
     of rotations and sending it to ``analyzer`` via its __call__ method
 
+    >>> state = analyzer.init_state()
     >>> for rotation_number in range(10):
     >>>     scores = np.random.rand(50,50)
     >>>     rotation = np.random.rand(scores.ndim, scores.ndim)
-    >>>     analyzer(scores = scores, rotation_matrix = rotation)
+    >>>     state = analyzer(state, scores=scores, rotation_matrix=rotation)
 
-    The aggregated scores can be extracted by invoking the __iter__ method of
+    The aggregated scores can be extracted by invoking the result method of
     ``analyzer``
 
-    >>> results = tuple(analyzer)
+    >>> results = analyzer.result(state)
 
     The ``results`` tuple contains (1) the maximum scores for each translation,
     (2) an offset which is relevant when merging results from split template matching
     using :py:meth:`MaxScoreOverRotations.merge`, (3) the rotation used to obtain a
-    score for a given translation, (4) a dictionary mapping rotation matrices to the
-    indices used in (2).
+    score for a given translation, (4) a dictionary mapping indices used in (2) to
+    rotation matrices (2).
 
     We can extract the ``optimal_score``, ``optimal_translation`` and ``optimal_rotation``
     as follows
 
     >>> optimal_score = results[0].max()
     >>> optimal_translation = np.where(results[0] == results[0].max())
-    >>> optimal_rotation_index = results[2][optimal_translation]
-    >>> for key, value in results[3].items():
-    >>>     if value != optimal_rotation_index:
-    >>>         continue
-    >>>     optimal_rotation = np.frombuffer(key, rotation.dtype)
-    >>>     optimal_rotation = optimal_rotation.reshape(scores.ndim, scores.ndim)
+    >>> optimal_rotation = results[2][optimal_translation]
 
     The outlined procedure is a trivial method to identify high scoring peaks.
     Alternatively, :py:class:`PeakCaller` offers a range of more elaborate approaches
@@ -111,156 +96,221 @@ class MaxScoreOverRotations:
     def __init__(
         self,
         shape: Tuple[int],
-        scores: BackendArray = None,
-        rotations: BackendArray = None,
         offset: BackendArray = None,
         score_threshold: float = 0,
         shm_handler: object = None,
         use_memmap: bool = False,
-        thread_safe: bool = True,
-        only_unique_rotations: bool = False,
+        inversion_mapping: bool = False,
+        jax_mode: bool = False,
         **kwargs,
     ):
+        self._use_memmap = use_memmap
+        self._score_threshold = score_threshold
         self._shape = tuple(int(x) for x in shape)
+        self._inversion_mapping = inversion_mapping
 
-        self.scores = scores
-        if self.scores is None:
-            self.scores = be.full(
-                shape=self._shape, dtype=be._float_dtype, fill_value=score_threshold
-            )
-        self.rotations = rotations
-        if self.rotations is None:
-            self.rotations = be.full(self._shape, dtype=be._int_dtype, fill_value=-1)
-
-        self.scores = be.to_sharedarr(self.scores, shm_handler)
-        self.rotations = be.to_sharedarr(self.rotations, shm_handler)
+        self._jax_mode = jax_mode
+        if self._jax_mode:
+            self._inversion_mapping = False
 
         if offset is None:
             offset = be.zeros(len(self._shape), be._int_dtype)
-        self.offset = be.astype(be.to_backend_array(offset), int)
+        self._offset = be.astype(be.to_backend_array(offset), int)
 
-        self._use_memmap = use_memmap
-        self._lock = Manager().Lock() if thread_safe else nullcontext()
-        self._lock_is_nullcontext = isinstance(self.scores, type(be.zeros((1))))
-        self._inversion_mapping = self._lock_is_nullcontext and only_unique_rotations
-        self.rotation_mapping = Manager().dict() if thread_safe else {}
+    @property
+    def shareable(self):
+        return True
+
+    def init_state(self):
+        """
+        Initialize the analysis state.
+
+        Returns
+        -------
+        tuple
+            Initial state tuple containing (scores, rotations, rotation_mapping) where:
+            - scores : BackendArray of shape `self._shape` filled with `score_threshold`.
+            - rotations : BackendArray of shape `self._shape` filled with -1.
+            - rotation_mapping : dict, empty mapping from rotation bytes to indices.
+        """
+        scores = be.full(
+            shape=self._shape, dtype=be._float_dtype, fill_value=self._score_threshold
+        )
+        rotations = be.full(self._shape, dtype=be._int_dtype, fill_value=-1)
+        return scores, rotations, {}
+
+    def __call__(
+        self,
+        state: Tuple,
+        scores: BackendArray,
+        rotation_matrix: BackendArray,
+        **kwargs,
+    ) -> Tuple:
+        """
+        Update the parameter store.
+
+        Parameters
+        ----------
+        state : tuple
+            Current state tuple (scores, rotations, rotation_mapping) where:
+            - scores : BackendArray, current maximum scores.
+            - rotations : BackendArray, current rotation indices.
+            - rotation_mapping : dict, mapping from rotation bytes to indices.
+        scores : BackendArray
+            Array of new scores to update analyzer with.
+        rotation_matrix : BackendArray
+            Square matrix used to obtain the current rotation.
+        Returns
+        -------
+        tuple
+            Updated state tuple (scores, rotations, rotation_mapping).
+        """
+        # be.tobytes behaviour caused overhead for certain GPU/CUDA combinations
+        # If the analyzer is not shared and each rotation is unique, we can
+        # use index to rotation mapping and invert prior to merging.
+        prev_scores, rotations, rotation_mapping = state
+
+        rotation_index = len(rotation_mapping)
+        rotation_matrix = be.astype(rotation_matrix, be._float_dtype)
+        if self._inversion_mapping:
+            rotation_mapping[rotation_index] = rotation_matrix
+        elif self._jax_mode:
+            rotation_index = kwargs.get("rotation_index", 0)
+        else:
+            rotation = be.tobytes(rotation_matrix)
+            rotation_index = rotation_mapping.setdefault(rotation, rotation_index)
+
+        scores, rotations = be.max_score_over_rotations(
+            scores=scores,
+            max_scores=prev_scores,
+            rotations=rotations,
+            rotation_index=rotation_index,
+        )
+        return scores, rotations, rotation_mapping
 
-    def _postprocess(
+    @staticmethod
+    def _invert_rmap(rotation_mapping: dict) -> dict:
+        """
+        Invert dictionary from rotation matrix bytestrings mapping to rotation
+        indices ro rotation indices mapping to rotation matrices.
+        """
+        new_map, ndim = {}, None
+        for k, v in rotation_mapping.items():
+            nbytes = be.datatype_bytes(be._float_dtype)
+            dtype = np.float32 if nbytes == 4 else np.float16
+            rmat = np.frombuffer(k, dtype=dtype)
+            if ndim is None:
+                ndim = int(np.sqrt(rmat.size))
+            new_map[v] = rmat.reshape(ndim, ndim)
+        return new_map
+
+    def result(
         self,
-        targetshape: Tuple[int],
-        templateshape: Tuple[int],
-        convolution_shape: Tuple[int],
+        state,
+        targetshape: Tuple[int] = None,
+        templateshape: Tuple[int] = None,
+        convolution_shape: Tuple[int] = None,
         fourier_shift: Tuple[int] = None,
         convolution_mode: str = None,
-        shm_handler=None,
         **kwargs,
-    ) -> "MaxScoreOverRotations":
-        """Correct padding to Fourier shape and convolution mode."""
-        scores = be.from_sharedarr(self.scores)
-        rotations = be.from_sharedarr(self.rotations)
+    ) -> Tuple:
+        """
+        Finalize the analysis result with optional postprocessing.
+
+        Parameters
+        ----------
+        state : tuple
+            Current state tuple (scores, rotations, rotation_mapping) where:
+            - scores : BackendArray, current maximum scores.
+            - rotations : BackendArray, current rotation indices.
+            - rotation_mapping : dict, mapping from rotation indices to matrices.
+        targetshape : Tuple[int], optional
+            Shape of the target for convolution mode correction.
+        templateshape : Tuple[int], optional
+            Shape of the template for convolution mode correction.
+        convolution_shape : Tuple[int], optional
+            Shape used for convolution.
+        fourier_shift : Tuple[int], optional.
+            Shift to apply for Fourier correction.
+        convolution_mode : str, optional
+            Convolution mode for padding correction.
+        **kwargs
+            Additional keyword arguments.
+
+        Returns
+        -------
+        tuple
+            Final result tuple (scores, offset, rotations, rotation_mapping).
+        """
+        scores, rotations, rotation_mapping = state
+
+        # Apply postprocessing if parameters are provided
         if fourier_shift is not None:
             axis = tuple(i for i in range(len(fourier_shift)))
             scores = be.roll(scores, shift=fourier_shift, axis=axis)
             rotations = be.roll(rotations, shift=fourier_shift, axis=axis)
 
-        convargs = {
-            "s1": targetshape,
-            "s2": templateshape,
-            "convolution_mode": convolution_mode,
-            "convolution_shape": convolution_shape,
-        }
         if convolution_mode is not None:
+            convargs = {
+                "s1": targetshape,
+                "s2": templateshape,
+                "convolution_mode": convolution_mode,
+                "convolution_shape": convolution_shape,
+            }
             scores = apply_convolution_mode(scores, **convargs)
             rotations = apply_convolution_mode(rotations, **convargs)
 
-        self._shape, self.scores, self.rotations = scores.shape, scores, rotations
-        if shm_handler is not None:
-            self.scores = be.to_sharedarr(scores, shm_handler)
-            self.rotations = be.to_sharedarr(rotations, shm_handler)
-        return self
-
-    def __iter__(self) -> Generator:
-        scores = be.from_sharedarr(self.scores)
-        rotations = be.from_sharedarr(self.rotations)
-
         scores = be.to_numpy_array(scores)
         rotations = be.to_numpy_array(rotations)
         if self._use_memmap:
             scores = array_to_memmap(scores)
             rotations = array_to_memmap(rotations)
-        else:
-            if type(self.scores) is not type(scores):
-                # Copy to avoid invalidation by shared memory handler
-                scores, rotations = scores.copy(), rotations.copy()
 
         if self._inversion_mapping:
-            self.rotation_mapping = {
-                be.tobytes(v): k for k, v in self.rotation_mapping.items()
-            }
+            rotation_mapping = {be.tobytes(v): k for k, v in rotation_mapping.items()}
 
-        param_store = (
+        return (
             scores,
-            be.to_numpy_array(self.offset),
+            be.to_numpy_array(self._offset),
             rotations,
-            dict(self.rotation_mapping),
+            self._invert_rmap(rotation_mapping),
         )
-        yield from param_store
 
-    def __call__(self, scores: BackendArray, rotation_matrix: BackendArray):
+    def _harmonize_states(states: List[Tuple]):
         """
-        Update the parameter store.
-
-        Parameters
-        ----------
-        scores : BackendArray
-            Array of scores.
-        rotation_matrix : BackendArray
-            Square matrix describing the current rotation.
+        Create consistent reference frame for merging different analyzer
+        instances, w.r.t. to rotations and output shape from different
+        splits of the target.
         """
-        # be.tobytes behaviour caused overhead for certain GPU/CUDA combinations
-        # If the analyzer is not shared and each rotation is unique, we can
-        # use index to rotation mapping and invert prior to merging.
-        if self._lock_is_nullcontext:
-            rotation_index = len(self.rotation_mapping)
-            if self._inversion_mapping:
-                self.rotation_mapping[rotation_index] = rotation_matrix
-            else:
-                rotation = be.tobytes(rotation_matrix)
-                rotation_index = self.rotation_mapping.setdefault(
-                    rotation, rotation_index
-                )
-            self.scores, self.rotations = be.max_score_over_rotations(
-                scores=scores,
-                max_scores=self.scores,
-                rotations=self.rotations,
-                rotation_index=rotation_index,
-            )
-            return None
+        new_rotation_mapping, out_shape = {}, None
+        for i in range(len(states)):
+            if states[i] is None:
+                continue
 
-        rotation = be.tobytes(rotation_matrix)
-        with self._lock:
-            rotation_index = self.rotation_mapping.setdefault(
-                rotation, len(self.rotation_mapping)
-            )
-            internal_scores = be.from_sharedarr(self.scores)
-            internal_rotations = be.from_sharedarr(self.rotations)
-            internal_sores, internal_rotations = be.max_score_over_rotations(
-                scores=scores,
-                max_scores=internal_scores,
-                rotations=internal_rotations,
-                rotation_index=rotation_index,
-            )
-            return None
+            scores, offset, rotations, rotation_mapping = states[i]
+            if out_shape is None:
+                out_shape = np.zeros(scores.ndim, int)
+            out_shape = np.maximum(out_shape, np.add(offset, scores.shape))
+
+            new_param = {}
+            for key, value in rotation_mapping.items():
+                rotation_bytes = be.tobytes(value)
+                new_param[rotation_bytes] = key
+                if rotation_bytes not in new_rotation_mapping:
+                    new_rotation_mapping[rotation_bytes] = len(new_rotation_mapping)
+            states[i] = (scores, offset, rotations, new_param)
+        out_shape = tuple(int(x) for x in out_shape)
+        return new_rotation_mapping, out_shape, states
 
     @classmethod
-    def merge(cls, param_stores: List[Tuple], **kwargs) -> Tuple:
+    def merge(cls, results: List[Tuple], **kwargs) -> Tuple:
         """
         Merge multiple instances of the current class.
 
         Parameters
         ----------
-        param_stores : list of tuple
-            List of instance's internal state created by applying `tuple(instance)`.
+        results : list of tuple
+            List of instance's internal state created by applying `result`.
         **kwargs : dict, optional
             Optional keyword arguments.
 
@@ -276,8 +326,8 @@ class MaxScoreOverRotations:
             Mapping between rotations and rotation indices.
         """
         use_memmap = kwargs.get("use_memmap", False)
-        if len(param_stores) == 1:
-            ret = param_stores[0]
+        if len(results) == 1:
+            ret = results[0]
             if use_memmap:
                 scores, offset, rotations, rotation_mapping = ret
                 scores = array_to_memmap(scores)
@@ -287,25 +337,12 @@ class MaxScoreOverRotations:
             return ret
 
         # Determine output array shape and create consistent rotation map
-        new_rotation_mapping, out_shape = {}, None
-        for i in range(len(param_stores)):
-            if param_stores[i] is None:
-                continue
-
-            scores, offset, rotations, rotation_mapping = param_stores[i]
-            if out_shape is None:
-                out_shape = np.zeros(scores.ndim, int)
-                scores_dtype, rotations_dtype = scores.dtype, rotations.dtype
-            out_shape = np.maximum(out_shape, np.add(offset, scores.shape))
-
-            for key, value in rotation_mapping.items():
-                if key not in new_rotation_mapping:
-                    new_rotation_mapping[key] = len(new_rotation_mapping)
-
+        master_rotation_mapping, out_shape, results = cls._harmonize_states(results)
         if out_shape is None:
             return None
 
-        out_shape = tuple(int(x) for x in out_shape)
+        scores_dtype = results[0][0].dtype
+        rotations_dtype = results[0][2].dtype
         if use_memmap:
             scores_out_filename = generate_tempfile_name()
             rotations_out_filename = generate_tempfile_name()
@@ -331,8 +368,8 @@ class MaxScoreOverRotations:
             )
             rotations_out = np.full(out_shape, fill_value=-1, dtype=rotations_dtype)
 
-        for i in range(len(param_stores)):
-            if param_stores[i] is None:
+        for i in range(len(results)):
+            if results[i] is None:
                 continue
 
             if use_memmap:
@@ -348,7 +385,7 @@ class MaxScoreOverRotations:
                     shape=out_shape,
                     dtype=rotations_dtype,
                 )
-            scores, offset, rotations, rotation_mapping = param_stores[i]
+            scores, offset, rotations, rotation_mapping = results[i]
             stops = np.add(offset, scores.shape).astype(int)
             indices = tuple(slice(*pos) for pos in zip(offset, stops))
 
@@ -359,7 +396,7 @@ class MaxScoreOverRotations:
                 len(rotation_mapping) + 1, dtype=rotations_out.dtype
             )
             for key, value in rotation_mapping.items():
-                lookup_table[value] = new_rotation_mapping[key]
+                lookup_table[value] = master_rotation_mapping[key]
 
             updated_rotations = rotations[indices_update]
             if len(updated_rotations):
@@ -372,7 +409,7 @@ class MaxScoreOverRotations:
                 rotations_out.flush()
                 scores_out, rotations_out = None, None
 
-            param_stores[i] = None
+            results[i] = None
             scores, rotations = None, None
 
         if use_memmap:
@@ -390,13 +427,165 @@ class MaxScoreOverRotations:
             scores_out,
             np.zeros(scores_out.ndim, dtype=int),
             rotations_out,
-            new_rotation_mapping,
+            cls._invert_rmap(master_rotation_mapping),
         )
 
-    @property
-    def is_shareable(self) -> bool:
-        """Boolean indicating whether class instance can be shared across processes."""
-        return True
+
+class MaxScoreOverRotationsConstrained(MaxScoreOverRotations):
+    """
+    Implements constrained template matching using rejection sampling.
+
+    Parameters
+    ----------
+    cone_angle : float
+        Maximum accepted rotational deviation in degrees.
+    positions : BackendArray
+        Array of shape (n, d) with n seed point translations.
+    positions : BackendArray
+        Array of shape (n, d, d) with n seed point rotation matrices.
+    reference : BackendArray
+        Reference orientation of the template, wlog defaults to (0,0,1).
+    acceptance_radius : int or tuple of ints
+        Translational acceptance radius around seed point in voxels.
+    **kwargs : dict, optional
+        Keyword aguments passed to the constructor of :py:class:`MaxScoreOverRotations`.
+    """
+
+    def __init__(
+        self,
+        cone_angle: float,
+        positions: BackendArray,
+        rotations: BackendArray,
+        reference: BackendArray = (0, 0, 1),
+        acceptance_radius: int = 10,
+        **kwargs,
+    ):
+        MaxScoreOverRotations.__init__(self, **kwargs)
+
+        if not isinstance(acceptance_radius, (int, Tuple)):
+            raise ValueError("acceptance_radius needs to be of type int or tuple.")
+
+        if isinstance(acceptance_radius, int):
+            acceptance_radius = (
+                acceptance_radius,
+                acceptance_radius,
+                acceptance_radius,
+            )
+        acceptance_radius = tuple(int(x) for x in acceptance_radius)
+
+        self._cone_angle = float(np.radians(cone_angle))
+        self._cone_cutoff = float(np.tan(self._cone_angle))
+        self._reference = be.astype(
+            be.reshape(be.to_backend_array(reference), (-1,)), be._float_dtype
+        )
+        positions = be.astype(be.to_backend_array(positions), be._int_dtype)
+
+        ndim = positions.shape[1]
+        rotate_mask = len(set(acceptance_radius)) != 1
+        extend = max(acceptance_radius)
+        mask = create_mask(
+            mask_type="ellipse",
+            radius=acceptance_radius,
+            shape=tuple(2 * extend + 1 for _ in range(ndim)),
+            center=tuple(extend for _ in range(ndim)),
+        )
+        self._score_mask = be.astype(be.to_backend_array(mask), be._float_dtype)
+
+        # Map position from real space to shifted score space
+        lower_limit = be.to_backend_array(self._offset)
+        positions = be.subtract(positions, lower_limit)
+        positions, valid_positions = cart_to_score(
+            positions=positions,
+            fast_shape=kwargs.get("fast_shape", None),
+            targetshape=kwargs.get("targetshape", None),
+            templateshape=kwargs.get("templateshape", None),
+            fourier_shift=kwargs.get("fourier_shift", None),
+            convolution_mode=kwargs.get("convolution_mode", None),
+            convolution_shape=kwargs.get("convolution_shape", None),
+        )
+
+        self._positions = positions[valid_positions]
+        rotations = be.to_backend_array(rotations)[valid_positions]
+        ex = be.astype(be.to_backend_array((1, 0, 0)), be._float_dtype)
+        ey = be.astype(be.to_backend_array((0, 1, 0)), be._float_dtype)
+        ez = be.astype(be.to_backend_array((0, 0, 1)), be._float_dtype)
+
+        self._normals_x = (rotations @ ex[..., None])[..., 0]
+        self._normals_y = (rotations @ ey[..., None])[..., 0]
+        self._normals_z = (rotations @ ez[..., None])[..., 0]
+
+        # Periodic wrapping could be avoided by padding the target
+        shape = be.to_backend_array(self._shape)
+        starts = be.subtract(self._positions, extend)
+        ret, (n, d), mshape = [], self._positions.shape, self._score_mask.shape
+        if starts.shape[0] > 0:
+            for i in range(d):
+                indices = starts[:, slice(i, i + 1)] + be.arange(mshape[i])[None]
+                indices = be.mod(indices, shape[i], out=indices)
+                indices_shape = (n, *tuple(1 if k != i else -1 for k in range(d)))
+                ret.append(be.reshape(indices, indices_shape))
+
+        self._index_grid = tuple(ret)
+        self._mask_shape = tuple(1 if i != 0 else -1 for i in range(1 + ndim))
+
+        if rotate_mask:
+            self._score_mask = be.zeros(
+                (rotations.shape[0], *self._score_mask.shape), dtype=be._float_dtype
+            )
+            for i in range(rotations.shape[0]):
+                mask = create_mask(
+                    mask_type="ellipse",
+                    radius=acceptance_radius,
+                    shape=tuple(2 * extend + 1 for _ in range(ndim)),
+                    center=tuple(extend for _ in range(ndim)),
+                    orientation=be.to_numpy_array(rotations[i]),
+                )
+                self._score_mask[i] = be.astype(
+                    be.to_backend_array(mask), be._float_dtype
+                )
+
+    def __call__(
+        self, state: Tuple, scores: BackendArray, rotation_matrix: BackendArray
+    ) -> Tuple:
+        mask = self._get_constraint(rotation_matrix)
+        mask = self._get_score_mask(mask=mask, scores=scores)
+
+        scores = be.multiply(scores, mask, out=scores)
+        return super().__call__(state, scores=scores, rotation_matrix=rotation_matrix)
+
+    def _get_constraint(self, rotation_matrix: BackendArray) -> BackendArray:
+        """
+        Determine whether the angle between projection of reference w.r.t to
+        a given rotation matrix and a set of rotations fall within the set
+        cone_angle cutoff.
+
+        Parameters
+        ----------
+        rotation_matrix : BackendArray
+            Rotation matrix with shape (d,d).
+
+        Returns
+        -------
+        BackerndArray
+            Boolean mask of shape (n, )
+        """
+        template_rot = rotation_matrix @ self._reference
+
+        x = be.sum(be.multiply(self._normals_x, template_rot), axis=1)
+        y = be.sum(be.multiply(self._normals_y, template_rot), axis=1)
+        z = be.sum(be.multiply(self._normals_z, template_rot), axis=1)
+
+        return be.sqrt(x**2 + y**2) <= (z * self._cone_cutoff)
+
+    def _get_score_mask(self, mask: BackendArray, scores: BackendArray, **kwargs):
+        score_mask = be.zeros(scores.shape, scores.dtype)
+
+        if be.sum(mask) == 0:
+            return score_mask
+        mask = be.reshape(mask, self._mask_shape)
+
+        score_mask = be.addat(score_mask, self._index_grid, self._score_mask * mask)
+        return score_mask > 0
 
 
 class MaxScoreOverTranslations(MaxScoreOverRotations):
@@ -437,43 +626,40 @@ class MaxScoreOverTranslations(MaxScoreOverRotations):
         super().__init__(
             shape=shape_reduced, shm_handler=shm_handler, offset=offset, **kwargs
         )
-
-        self.rotations = be.full(1, dtype=be._int_dtype, fill_value=-1)
-        self.rotations = be.to_sharedarr(self.rotations, shm_handler)
         self._aggregate_axis = aggregate_axis
 
-    def __call__(self, scores: BackendArray, rotation_matrix: BackendArray):
-        if self._lock_is_nullcontext:
-            rotation_index = len(self.rotation_mapping)
-            if self._inversion_mapping:
-                self.rotation_mapping[rotation_index] = rotation_matrix
-            else:
-                rotation = be.tobytes(rotation_matrix)
-                rotation_index = self.rotation_mapping.setdefault(
-                    rotation, rotation_index
-                )
-            max_score = be.max(scores, axis=self._aggregate_axis)
-            self.scores[rotation_index] = max_score
-            return None
+    def init_state(self):
+        scores = be.full(
+            shape=self._shape, dtype=be._float_dtype, fill_value=self._score_threshold
+        )
+        rotations = be.full(1, dtype=be._int_dtype, fill_value=-1)
+        return scores, rotations, {}
 
-        rotation = be.tobytes(rotation_matrix)
-        with self._lock:
-            rotation_index = self.rotation_mapping.setdefault(
-                rotation, len(self.rotation_mapping)
-            )
-            internal_scores = be.from_sharedarr(self.scores)
-            max_score = be.max(scores, axis=self._aggregate_axis)
-            internal_scores[rotation_index] = max_score
-            return None
+    def __call__(
+        self, state, scores: BackendArray, rotation_matrix: BackendArray
+    ) -> Tuple:
+        prev_scores, rotations, rotation_mapping = state
+
+        rotation_index = len(rotation_mapping)
+        if self._inversion_mapping:
+            rotation_mapping[rotation_index] = rotation_matrix
+        else:
+            rotation = be.tobytes(rotation_matrix)
+            rotation_index = rotation_mapping.setdefault(rotation, rotation_index)
+        max_score = be.max(scores, axis=self._aggregate_axis)
+
+        update = prev_scores[rotation_index]
+        update = be.maximum(max_score, update, out=update)
+        return prev_scores, rotations, rotation_mapping
 
     @classmethod
-    def merge(cls, param_stores: List[Tuple], **kwargs) -> Tuple:
+    def merge(cls, states: List[Tuple], **kwargs) -> Tuple:
         """
         Merge multiple instances of the current class.
 
         Parameters
         ----------
-        param_stores : list of tuple
+        states : list of tuple
             List of instance's internal state created by applying `tuple(instance)`.
         **kwargs : dict, optional
             Optional keyword arguments.
@@ -489,31 +675,18 @@ class MaxScoreOverTranslations(MaxScoreOverRotations):
         Dict
             Mapping between rotations and rotation indices.
         """
-        if len(param_stores) == 1:
-            return param_stores[0]
+        if len(states) == 1:
+            return states[0]
 
         # Determine output array shape and create consistent rotation map
-        new_rotation_mapping, out_shape = {}, None
-        for i in range(len(param_stores)):
-            if param_stores[i] is None:
-                continue
-
-            scores, offset, rotations, rotation_mapping = param_stores[i]
-            if out_shape is None:
-                out_shape = np.zeros(scores.ndim, int)
-                scores_dtype, rotations_out = scores.dtype, rotations
-            out_shape = np.maximum(out_shape, np.add(offset, scores.shape))
-
-            for key, value in rotation_mapping.items():
-                if key not in new_rotation_mapping:
-                    new_rotation_mapping[key] = len(new_rotation_mapping)
-
+        states, master_rotation_mapping, out_shape = cls._harmonize_states(states)
         if out_shape is None:
             return None
 
-        out_shape[0] = len(new_rotation_mapping)
+        out_shape[0] = len(master_rotation_mapping)
         out_shape = tuple(int(x) for x in out_shape)
 
+        scores_dtype = states[0][0].dtype
         use_memmap = kwargs.get("use_memmap", False)
         if use_memmap:
             scores_out_filename = generate_tempfile_name()
@@ -529,8 +702,8 @@ class MaxScoreOverTranslations(MaxScoreOverRotations):
                 dtype=scores_dtype,
             )
 
-        for i in range(len(param_stores)):
-            if param_stores[i] is None:
+        for i in range(len(states)):
+            if states[i] is None:
                 continue
 
             if use_memmap:
@@ -540,11 +713,11 @@ class MaxScoreOverTranslations(MaxScoreOverRotations):
                     shape=out_shape,
                     dtype=scores_dtype,
                 )
-            scores, offset, rotations, rotation_mapping = param_stores[i]
+            scores, offset, rotations, rotation_mapping = states[i]
 
             outer_table = np.arange(len(rotation_mapping), dtype=int)
             lookup_table = np.array(
-                [new_rotation_mapping[key] for key in rotation_mapping.keys()],
+                [master_rotation_mapping[key] for key in rotation_mapping.keys()],
                 dtype=int,
             )
 
@@ -560,7 +733,7 @@ class MaxScoreOverTranslations(MaxScoreOverRotations):
                 scores_out.flush()
                 scores_out = None
 
-            param_stores[i], scores = None, None
+            states[i], scores = None, None
 
         if use_memmap:
             scores_out = np.memmap(
@@ -570,9 +743,9 @@ class MaxScoreOverTranslations(MaxScoreOverRotations):
         return (
             scores_out,
             np.zeros(scores_out.ndim, dtype=int),
-            rotations_out,
-            new_rotation_mapping,
+            states[2],
+            cls._invert_rmap(master_rotation_mapping),
         )
 
-    def _postprocess(self, **kwargs):
-        return self
+    def result(self, state: Tuple, **kwargs) -> Tuple:
+        return state