pytme 0.2.9.post1__cp311-cp311-macosx_15_0_arm64.whl → 0.3.0__cp311-cp311-macosx_15_0_arm64.whl

tme/analyzer/base.py

@@ -0,0 +1,127 @@
+"""
+Implements abstract base class for template matching analyzers.
+
+Copyright (c) 2025 European Molecular Biology Laboratory
+
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+
+from typing import Tuple, List
+from abc import ABC, abstractmethod
+
+__all__ = ["AbstractAnalyzer"]
+
+
+class AbstractAnalyzer(ABC):
+    """
+    Abstract base class for template matching analyzers.
+    """
+
+    @property
+    def shareable(self):
+        """
+        Indicate whether the analyzer can be shared across processes.
+
+        Returns
+        -------
+        bool
+            True if the analyzer supports shared memory operations
+            and can be safely used across multiple processes, False
+            if it should only be used within a single process.
+        """
+        return False
+
+    @abstractmethod
+    def init_state(self, *args, **kwargs) -> Tuple:
+        """
+        Initialize the analyzer state.
+
+        Returns
+        -------
+        state : tuple
+            Initial state tuple of the analyzer instance. The exact structure
+            depends on the specific implementation.
+
+        Notes
+        -----
+        This method creates the initial state that will be passed to
+        :py:meth:`AbstractAnalyzer.__call__` and finally to
+        :py:meth:`AbstractAnalyzer.result`. The state should contain all necessary
+        data structures for accumulating analysis results.
+        """
+
+    @abstractmethod
+    def __call__(self, state, scores, rotation_matrix, **kwargs) -> Tuple:
+        """
+        Update the analyzer state with new scoring data.
+
+        Parameters
+        ----------
+        state : tuple
+            Current analyzer state as returned :py:meth:`AbstractAnalyzer.init_state`
+            or previous invocations of :py:meth:`AbstractAnalyzer.__call__`.
+        scores : BackendArray
+            Array of new scores with dimensionality d.
+        rotation_matrix : BackendArray
+            Rotation matrix used to generate scores with shape (d,d).
+        **kwargs : dict
+            Keyword arguments used by specific implementations.
+
+        Returns
+        -------
+        tuple
+            Updated analyzer state incorporating the new data.
+        """
+
+    @abstractmethod
+    def result(self, state: Tuple, **kwargs) -> Tuple:
+        """
+        Finalize the analysis by performing potential post processing.
+
+        Parameters
+        ----------
+        state : tuple
+            Analyzer state containing accumulated data.
+        **kwargs : dict
+            Keyword arguments used by specific implementations.
+
+        Returns
+        -------
+        result
+            Final analysis result. The exact struccture depends on the
+            analyzer implementation.
+
+        Notes
+        -----
+        This method converts the internal analyzer state into the
+        final output format expected by the template matching pipeline.
+        It may apply postprocessing operations like convolution mode
+        correction or coordinate transformations.
+        """
+
+    @classmethod
+    @abstractmethod
+    def merge(cls, results: List[Tuple], **kwargs) -> Tuple:
+        """
+        Merge multiple analyzer results.
+
+        Parameters
+        ----------
+        results : list of tuple
+            List of tuple objects returned by :py:meth:`AbstractAnalyzer.result`
+            from different instances of the same analyzer class.
+        **kwargs : dict
+            Keyword arguments used by specific implementations.
+
+        Returns
+        -------
+        tuple
+            Single result object combining all input results.
+
+        Notes
+        -----
+        This method enables parallel processing by allowing results
+        from different processes or splits to be combined into a
+        unified result. The merge operation should handle overlapping
+        data appropriately and maintain consistency.
+        """

tme/analyzer/peaks.py

@@ -1,23 +1,25 @@
-""" 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 functools import wraps
-from abc import ABC, abstractmethod
-from typing import Tuple, List, Dict, Generator
+from abc import abstractmethod
+from typing import Tuple, List, Dict
 
 import numpy as np
 from skimage.feature import peak_local_max
 from skimage.registration._phase_cross_correlation import _upsampled_dft
 
+from .base import AbstractAnalyzer
 from ._utils import score_to_cart
 from ..backends import backend as be
-from ..matching_utils import split_shape
 from ..types import BackendArray, NDArray
 from ..rotations import euler_to_rotationmatrix
+from ..matching_utils import split_shape, compute_extraction_box
 
 __all__ = [
     "PeakCaller",
@@ -141,7 +143,7 @@ def batchify(shape: Tuple[int], batch_dims: Tuple[int] = None) -> List:
     yield from _generate_slices_recursive(0, ())
 
 
-class PeakCaller(ABC):
+class PeakCaller(AbstractAnalyzer):
     """
     Base class for peak calling algorithms.
 
@@ -190,19 +192,7 @@ class PeakCaller(ABC):
         if min_boundary_distance < 0:
             raise ValueError("min_boundary_distance has to be non-negative.")
 
-        ndim = len(shape)
-        self.translations = be.full(
-            (num_peaks, ndim), fill_value=-1, dtype=be._int_dtype
-        )
-        self.rotations = be.full(
-            (num_peaks, ndim, ndim), fill_value=0, dtype=be._float_dtype
-        )
-        for i in range(ndim):
-            self.rotations[:, i, i] = 1.0
-
-        self.scores = be.full((num_peaks,), fill_value=0, dtype=be._float_dtype)
-        self.details = be.full((num_peaks,), fill_value=0, dtype=be._float_dtype)
-
+        self.shape = shape
         self.num_peaks = int(num_peaks)
         self.min_distance = int(min_distance)
         self.min_boundary_distance = int(min_boundary_distance)
@@ -213,31 +203,47 @@ class PeakCaller(ABC):
 
         self.min_score, self.max_score = min_score, max_score
 
-        # Postprocessing 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) -> Generator:
+    @abstractmethod
+    def call_peaks(self, scores: BackendArray, **kwargs) -> PeakType:
         """
-        Returns a generator to list objects containing translation,
-        rotation, score and details of a given candidate.
+        Call peaks in the score space.
+
+        Parameters
+        ----------
+        scores : BackendArray
+            Score array to update analyzer with.
+        **kwargs : dict
+            Optional keyword arguments passed to underlying implementations.
+
+        Returns
+        -------
+        BackendArray
+            Peak positions (n, d).
+        BackendArray
+            Peak details (n, d).
         """
-        self.peak_list = [
-            be.to_cpu_array(self.translations),
-            be.to_cpu_array(self.rotations),
-            be.to_cpu_array(self.scores),
-            be.to_cpu_array(self.details),
-        ]
-        yield from self.peak_list
+
+    def init_state(self):
+        ndim = len(self.shape)
+        translations = be.full(
+            (self.num_peaks, ndim), fill_value=-1, dtype=be._int_dtype
+        )
+        rotations = be.full(
+            (self.num_peaks, ndim, ndim), fill_value=0, dtype=be._float_dtype
+        )
+        for i in range(ndim):
+            rotations[:, i, i] = 1.0
+
+        scores = be.full((self.num_peaks,), fill_value=-1, dtype=be._float_dtype)
+        details = be.full((self.num_peaks,), fill_value=-1, dtype=be._float_dtype)
+        return translations, rotations, scores, details
 
     def _get_peak_mask(self, peaks: BackendArray, scores: BackendArray) -> BackendArray:
         if not len(peaks):
             return None
 
         valid_peaks = be.full((peaks.shape[0],), fill_value=1) == 1
-        if self.min_boundary_distance > 0:
+        if self.min_boundary_distance >= 0:
             upper_limit = be.subtract(
                 be.to_backend_array(scores.shape), self.min_boundary_distance
             )
@@ -318,20 +324,34 @@ class PeakCaller(ABC):
         peak_positions = be.astype(peak_positions, int)
         return peak_positions, peak_details
 
-    def __call__(self, scores: BackendArray, rotation_matrix: BackendArray, **kwargs):
+    def __call__(
+        self,
+        state: Tuple,
+        scores: BackendArray,
+        rotation_matrix: BackendArray,
+        **kwargs,
+    ) -> Tuple:
         """
         Update the internal parameter store based on input array.
 
         Parameters
         ----------
+        state : tuple
+            Current state tuple where:
+            - positions : BackendArray, (n, d) of peak positions
+            - rotations : BackendArray, (n, d, d) of correponding rotations
+            - scores : BackendArray, (n, ) of peak scores
+            - details : BackendArray, (n, ) of peak details
         scores : BackendArray
-            Score space data.
+            Array of new scores to update analyzer with.
         rotation_matrix : BackendArray
             Rotation matrix used to obtain the score array.
         **kwargs
             Optional keyword aguments passed to :py:meth:`PeakCaller.call_peaks`.
         """
-        for ret in self._call_peaks(scores=scores, rotation_matrix=rotation_matrix):
+        for ret in self._call_peaks(
+            scores=scores, rotation_matrix=rotation_matrix, **kwargs
+        ):
             peak_positions, peak_details = ret
             if peak_positions is None:
                 continue
@@ -344,7 +364,6 @@ class PeakCaller(ABC):
             peak_scores = scores[tuple(peak_positions.T)]
             if peak_details is not None:
                 peak_details = peak_details[valid_peaks]
-                # peak_details, peak_scores = peak_scores, -peak_details
             else:
                 peak_details = be.full(peak_scores.shape, fill_value=-1)
 
@@ -354,66 +373,57 @@ class PeakCaller(ABC):
                 axis=0,
             )
 
-            self._update(
+            state = self._update(
+                state,
                 peak_positions=peak_positions,
                 peak_details=peak_details,
                 peak_scores=peak_scores,
-                rotations=rotations,
+                peak_rotations=rotations,
             )
 
-        return None
-
-    @abstractmethod
-    def call_peaks(self, scores: BackendArray, **kwargs) -> PeakType:
-        """
-        Call peaks in the score space.
-
-        Parameters
-        ----------
-        scores : BackendArray
-            Score array.
-        **kwargs : dict
-            Optional keyword arguments passed to underlying implementations.
-
-        Returns
-        -------
-        Tuple[BackendArray, BackendArray]
-            Array of peak coordinates and peak details.
-        """
+        return state
 
     @classmethod
-    def merge(cls, candidates=List[List], **kwargs) -> Tuple:
+    def merge(cls, results=List[Tuple], **kwargs) -> Tuple:
         """
         Merge multiple instances of :py:class:`PeakCaller`.
 
         Parameters
         ----------
-        candidates : list of lists
-            Obtained by invoking list on the generator returned by __iter__.
+        results : list of tuple
+            List of instance results created by applying `result`.
         **kwargs
             Optional keyword arguments.
 
         Returns
         -------
-        Tuple
-            Tuple of translation, rotation, score and details of candidates.
+        NDArray
+            Peak positions (n, d).
+        NDArray
+            Peak rotation matrices (n, d, d).
+        NDArray
+            Peak scores (n, ).
+        NDArray
+            Peak details (n,).
         """
         if "shape" not in kwargs:
-            kwargs["shape"] = tuple(1 for _ in range(candidates[0][0].shape[1]))
+            kwargs["shape"] = tuple(1 for _ in range(results[0][0].shape[1]))
 
         base = cls(**kwargs)
-        for candidate in candidates:
-            if len(candidate) == 0:
+        base_state = base.init_state()
+        for result in results:
+            if len(result) == 0:
                 continue
-            peak_positions, rotations, peak_scores, peak_details = candidate
-            base._update(
+            peak_positions, rotations, peak_scores, peak_details = result
+            base_state = base._update(
+                base_state,
                 peak_positions=be.to_backend_array(peak_positions),
                 peak_details=be.to_backend_array(peak_details),
                 peak_scores=be.to_backend_array(peak_scores),
-                rotations=be.to_backend_array(rotations),
+                peak_rotations=be.to_backend_array(rotations),
                 offset=kwargs.get("offset", None),
             )
-        return tuple(base)
+        return base_state
 
     @staticmethod
     def oversample_peaks(
@@ -520,10 +530,11 @@ class PeakCaller(ABC):
 
     def _update(
         self,
+        state,
         peak_positions: BackendArray,
         peak_details: BackendArray,
         peak_scores: BackendArray,
-        rotations: BackendArray,
+        peak_rotations: BackendArray,
         offset: BackendArray = None,
     ):
         """
@@ -542,14 +553,15 @@ class PeakCaller(ABC):
         offset : BackendArray, optional
             Translation offset, e.g. from splitting, (d, ).
         """
+        translations, rotations, scores, details = state
         if offset is not None:
             offset = be.astype(be.to_backend_array(offset), peak_positions.dtype)
             peak_positions = be.add(peak_positions, offset, out=peak_positions)
 
-        positions = be.concatenate((self.translations, peak_positions))
-        rotations = be.concatenate((self.rotations, rotations))
-        scores = be.concatenate((self.scores, peak_scores))
-        details = be.concatenate((self.details, peak_details))
+        positions = be.concatenate((translations, peak_positions))
+        rotations = be.concatenate((rotations, peak_rotations))
+        scores = be.concatenate((scores, peak_scores))
+        details = be.concatenate((details, peak_details))
 
         # topk filtering after distances yields more distributed peak calls
         distance_order = filter_points_indices(
@@ -564,22 +576,69 @@ class PeakCaller(ABC):
         )
         final_order = distance_order[top_scores]
 
-        self.translations = positions[final_order, :]
-        self.rotations = rotations[final_order, :]
-        self.scores = scores[final_order]
-        self.details = details[final_order]
+        translations = positions[final_order, :]
+        rotations = rotations[final_order, :]
+        scores = scores[final_order]
+        details = details[final_order]
+        return translations, rotations, scores, details
 
-    def _postprocess(self, **kwargs):
-        if not len(self.translations):
-            return self
+    def result(
+        self,
+        state,
+        fast_shape: Tuple[int] = None,
+        targetshape: Tuple[int] = None,
+        templateshape: Tuple[int] = None,
+        convolution_shape: Tuple[int] = None,
+        fourier_shift: Tuple[int] = None,
+        convolution_mode: str = None,
+        **kwargs,
+    ):
+        """
+        Finalize the analysis result with optional postprocessing.
 
-        positions, valid_peaks = score_to_cart(self.translations, **kwargs)
+        Parameters
+        ----------
+        state : tuple
+            Current state tuple where:
+            - positions : BackendArray, (n, d) of peak positions
+            - rotations : BackendArray, (n, d, d) of correponding rotations
+            - scores : BackendArray, (n, ) of peak scores
+            - details : BackendArray, (n, ) of peak details
+        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.
 
-        self.translations = positions[valid_peaks]
-        self.rotations = self.rotations[valid_peaks]
-        self.scores = self.scores[valid_peaks]
-        self.details = self.details[valid_peaks]
-        return self
+        Returns
+        -------
+        tuple
+            Final result tuple (positions, rotations, scores, details).
+        """
+        translations, rotations, scores, details = state
+
+        positions, valid_peaks = score_to_cart(
+            positions=translations,
+            fast_shape=fast_shape,
+            targetshape=targetshape,
+            templateshape=templateshape,
+            convolution_shape=convolution_shape,
+            fourier_shift=fourier_shift,
+            convolution_mode=convolution_mode,
+            **kwargs,
+        )
+        translations = be.to_cpu_array(positions[valid_peaks])
+        rotations = be.to_cpu_array(rotations[valid_peaks])
+        scores = be.to_cpu_array(scores[valid_peaks])
+        details = be.to_cpu_array(details[valid_peaks])
+        return translations, rotations, scores, details
 
 
 class PeakCallerSort(PeakCaller):
@@ -706,14 +765,14 @@ class PeakCallerRecursiveMasking(PeakCaller):
         values. If rotations and rotation_mapping is provided, the respective
         rotation will be applied to the mask, otherwise rotation_matrix is used.
         """
-        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(scores.ndim))
-            mask = be.zeros(shape, dtype=be._float_dtype)
+        peaks = []
+        box = tuple(self.min_distance for _ in range(scores.ndim))
 
-        rotated_template = be.zeros(mask.shape, dtype=mask.dtype)
+        scores = be.to_backend_array(scores)
+        if mask is not None:
+            box = mask.shape
+            mask = be.to_backend_array(mask)
+            mask_buffer = be.zeros(mask.shape, dtype=mask.dtype)
 
         peak_limit = self.num_peaks
         if min_score is not None:
@@ -721,39 +780,45 @@ class PeakCallerRecursiveMasking(PeakCaller):
         else:
             min_score = be.min(scores) - 1
 
-        scores_copy = be.zeros(scores.shape, dtype=scores.dtype)
-        scores_copy[:] = scores
-
+        _scores = be.zeros(scores.shape, dtype=scores.dtype)
+        _scores[:] = scores[:]
         while True:
-            be.argmax(scores_copy)
-            peak = be.unravel_index(
-                indices=be.argmax(scores_copy), shape=scores_copy.shape
-            )
-            if scores_copy[tuple(peak)] < min_score:
+            peak = be.unravel_index(indices=be.argmax(_scores), shape=_scores.shape)
+            if _scores[tuple(peak)] < min_score:
                 break
+            peaks.append(peak)
 
-            coordinates.append(peak)
-
-            current_rotation_matrix = self._get_rotation_matrix(
-                peak=peak,
-                rotation_space=rotations,
-                rotation_mapping=rotation_mapping,
-                rotation_matrix=rotation_matrix,
+            score_beg, score_end, tmpl_beg, tmpl_end, _ = compute_extraction_box(
+                centers=be.to_backend_array(peak)[None],
+                extraction_shape=box,
+                original_shape=scores.shape,
             )
-
-            masking_function(
-                scores=scores_copy,
-                rotation_matrix=current_rotation_matrix,
-                peak=peak,
-                mask=mask,
-                rotated_template=rotated_template,
+            score_slice = tuple(
+                slice(int(x), int(y)) for x, y in zip(score_beg[0], score_end[0])
+            )
+            tmpl_slice = tuple(
+                slice(int(x), int(y)) for x, y in zip(tmpl_beg[0], tmpl_end[0])
             )
 
-            if len(coordinates) >= peak_limit:
+            score_mask = 0
+            if mask is not None:
+                mask_buffer.fill(0)
+                rmat = self._get_rotation_matrix(
+                    peak=peak,
+                    rotation_space=rotations,
+                    rotation_mapping=rotation_mapping,
+                    rotation_matrix=rotation_matrix,
+                )
+                be.rigid_transform(
+                    arr=mask, rotation_matrix=rmat, order=1, out=mask_buffer
+                )
+                score_mask = mask_buffer[tmpl_slice] <= 0.1
+
+            _scores[score_slice] = be.multiply(_scores[score_slice], score_mask)
+            if len(peaks) >= peak_limit:
                 break
 
-        peaks = be.to_backend_array(coordinates)
-        return peaks, None
+        return be.to_backend_array(peaks), None
 
     @staticmethod
     def _get_rotation_matrix(
@@ -793,86 +858,6 @@ class PeakCallerRecursiveMasking(PeakCaller):
             )
         return rotation
 
-    @staticmethod
-    def _mask_scores_box(
-        scores: BackendArray, peak: BackendArray, mask: BackendArray, **kwargs: Dict
-    ) -> None:
-        """
-        Mask scores in a box around a peak.
-
-        Parameters
-        ----------
-        scores : BackendArray
-            Data array of scores.
-        peak : BackendArray
-            Peak coordinates.
-        mask : BackendArray
-            Mask array.
-        """
-        start = be.maximum(be.subtract(peak, mask.shape), 0)
-        stop = be.minimum(be.add(peak, mask.shape), scores.shape)
-        start, stop = be.astype(start, int), be.astype(stop, int)
-        coords = tuple(slice(*pos) for pos in zip(start, stop))
-        scores[coords] = 0
-        return None
-
-    @staticmethod
-    def _mask_scores_rotate(
-        scores: BackendArray,
-        peak: BackendArray,
-        mask: BackendArray,
-        rotated_template: BackendArray,
-        rotation_matrix: BackendArray,
-        **kwargs: Dict,
-    ) -> None:
-        """
-        Mask scores using mask rotation around a peak.
-
-        Parameters
-        ----------
-        scores : BackendArray
-            Data array of scores.
-        peak : BackendArray
-            Peak coordinates.
-        mask : BackendArray
-            Mask array.
-        rotated_template : BackendArray
-            Empty array to write mask rotations to.
-        rotation_matrix : BackendArray
-            Rotation matrix.
-        """
-        left_pad = be.divide(mask.shape, 2).astype(int)
-        right_pad = be.add(left_pad, be.mod(mask.shape, 2).astype(int))
-
-        score_start = be.subtract(peak, left_pad)
-        score_stop = be.add(peak, right_pad)
-
-        template_start = be.subtract(be.maximum(score_start, 0), score_start)
-        template_stop = be.subtract(score_stop, be.minimum(score_stop, scores.shape))
-        template_stop = be.subtract(mask.shape, template_stop)
-
-        score_start = be.maximum(score_start, 0)
-        score_stop = be.minimum(score_stop, scores.shape)
-        score_start = be.astype(score_start, int)
-        score_stop = be.astype(score_stop, int)
-
-        template_start = be.astype(template_start, int)
-        template_stop = be.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)
-        be.rigid_transform(
-            arr=mask, rotation_matrix=rotation_matrix, order=1, out=rotated_template
-        )
-
-        scores[coords_score] = be.multiply(
-            scores[coords_score], (rotated_template[coords_template] <= 0.1)
-        )
-        return None
-
 
 class PeakCallerScipy(PeakCaller):
     """