pytme 0.3b0__cp311-cp311-macosx_15_0_arm64.whl → 0.3.1__cp311-cp311-macosx_15_0_arm64.whl

tests/test_matching_utils.py

@@ -10,9 +10,6 @@ from tme.backends import backend as be
 from tme.memory import MATCHING_MEMORY_REGISTRY
 from tme.matching_utils import (
     compute_parallelization_schedule,
-    elliptical_mask,
-    box_mask,
-    tube_mask,
     create_mask,
     scramble_phases,
     apply_convolution_mode,
@@ -50,73 +47,26 @@ class TestMatchingUtils:
             max_splits=256,
         )
 
-    def test_create_mask(self):
+    @pytest.mark.parametrize("mask_type", ["ellipse", "box", "tube", "membrane"])
+    def test_create_mask(self, mask_type: str):
         create_mask(
-            mask_type="ellipse",
+            mask_type=mask_type,
             shape=self.density.shape,
             radius=5,
             center=np.divide(self.density.shape, 2),
+            height=np.max(self.density.shape) // 2,
+            size=np.divide(self.density.shape, 2).astype(int),
+            thickness=2,
+            separation=2,
+            symmetry_axis=1,
+            inner_radius=5,
+            outer_radius=10,
         )
 
     def test_create_mask_error(self):
         with pytest.raises(ValueError):
             create_mask(mask_type=None)
 
-    def test_elliptical_mask(self):
-        elliptical_mask(
-            shape=self.density.shape,
-            radius=5,
-            center=np.divide(self.density.shape, 2),
-        )
-
-    def test_box_mask(self):
-        box_mask(
-            shape=self.density.shape,
-            height=[5, 10, 20],
-            center=np.divide(self.density.shape, 2),
-        )
-
-    def test_tube_mask(self):
-        tube_mask(
-            shape=self.density.shape,
-            outer_radius=10,
-            inner_radius=5,
-            height=5,
-            base_center=np.divide(self.density.shape, 2),
-            symmetry_axis=1,
-        )
-
-    def test_tube_mask_error(self):
-        with pytest.raises(ValueError):
-            tube_mask(
-                shape=self.density.shape,
-                outer_radius=5,
-                inner_radius=10,
-                height=5,
-                base_center=np.divide(self.density.shape, 2),
-                symmetry_axis=1,
-            )
-
-        with pytest.raises(ValueError):
-            tube_mask(
-                shape=self.density.shape,
-                outer_radius=5,
-                inner_radius=10,
-                height=10 * np.max(self.density.shape),
-                base_center=np.divide(self.density.shape, 2),
-                symmetry_axis=1,
-            )
-
-        with pytest.raises(ValueError):
-            tube_mask(
-                shape=self.density.shape,
-                outer_radius=5,
-                inner_radius=10,
-                height=10 * np.max(self.density.shape),
-                base_center=np.divide(self.density.shape, 2),
-                symmetry_axis=len(self.density.shape) + 1,
-            )
-
     def test_scramble_phases(self):
         scramble_phases(arr=self.density.data, noise_proportion=0.5)
 
@@ -139,7 +89,7 @@ class TestMatchingUtils:
             expected_size = np.subtract(
                 self.density.shape, self.structure_density.shape
             )
-            expected_size += np.mod(self.structure_density.shape, 2)
+            expected_size += 1
         assert np.allclose(ret.shape, expected_size)
 
     def test_apply_convolution_mode_error(self):

tests/test_rotations.py

@@ -8,6 +8,7 @@ from tme import Density
 from scipy.spatial.transform import Rotation
 from scipy.signal import correlate
 
+from tme.mask import elliptical_mask
 from tme.rotations import (
     euler_from_rotationmatrix,
     euler_to_rotationmatrix,
@@ -16,7 +17,6 @@ from tme.rotations import (
     get_rotation_matrices,
 )
 from tme.matching_utils import (
-    elliptical_mask,
     split_shape,
     compute_full_convolution_index,
 )

tme/__version__.py

@@ -1 +1 @@
-__version__ = "0.3.b0"
+__version__ = "0.3.1"

tme/analyzer/__init__.py

@@ -1,4 +1,4 @@
 from .peaks import *
 from .aggregation import *
 from .proxy import *
-from .base import *
+from .base import *

tme/analyzer/_utils.py

@@ -47,10 +47,7 @@ def _convmode_to_shape(
     if convolution_mode == "same":
         output_shape = targetshape
     elif convolution_mode == "valid":
-        output_shape = be.add(
-            be.subtract(targetshape, templateshape),
-            be.mod(templateshape, 2),
-        )
+        output_shape = be.subtract(targetshape, templateshape) + 1
     return be.to_backend_array(output_shape)
 
 
@@ -96,22 +93,22 @@ def cart_to_score(
     templateshape = be.to_backend_array(templateshape)
     convolution_shape = be.to_backend_array(convolution_shape)
 
-    # Compute removed padding
     output_shape = _convmode_to_shape(
         convolution_mode=convolution_mode,
         targetshape=targetshape,
         templateshape=templateshape,
         convolution_shape=convolution_shape,
     )
-    valid_positions = be.multiply(positions >= 0, positions < output_shape)
-    valid_positions = be.sum(valid_positions, axis=1) == positions.shape[1]
 
+    # Offset from padding the target
     starts = be.astype(
         be.divide(be.subtract(convolution_shape, output_shape), 2),
         be._int_dtype,
     )
-
     positions = be.add(positions, starts)
+
+    valid_positions = be.multiply(positions >= 0, positions < fast_shape)
+    valid_positions = be.sum(valid_positions, axis=1) == positions.shape[1]
     if fourier_shift is not None:
         fourier_shift = be.to_backend_array(fourier_shift)
         positions = be.subtract(positions, fourier_shift)

tme/analyzer/aggregation.py

@@ -57,22 +57,23 @@ class MaxScoreOverRotations(AbstractAnalyzer):
     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()
+    >>> state = analyzer.init_state()
     >>> for rotation_number in range(10):
     >>>     scores = np.random.rand(50,50)
     >>>     rotation = np.random.rand(scores.ndim, scores.ndim)
-    >>>     state, analyzer(state, scores = scores, rotation_matrix = rotation)
+    >>>     state = analyzer(state, scores=scores, rotation_matrix=rotation)
 
     The aggregated scores can be extracted by invoking the result method of
     ``analyzer``
 
-    >>> results = analyzer.result()
+    >>> 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
@@ -100,6 +101,7 @@ class MaxScoreOverRotations(AbstractAnalyzer):
         shm_handler: object = None,
         use_memmap: bool = False,
         inversion_mapping: bool = False,
+        jax_mode: bool = False,
         **kwargs,
     ):
         self._use_memmap = use_memmap
@@ -107,6 +109,10 @@ class MaxScoreOverRotations(AbstractAnalyzer):
         self._shape = tuple(int(x) for x in shape)
         self._inversion_mapping = inversion_mapping
 
+        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)
@@ -138,6 +144,7 @@ class MaxScoreOverRotations(AbstractAnalyzer):
         state: Tuple,
         scores: BackendArray,
         rotation_matrix: BackendArray,
+        **kwargs,
     ) -> Tuple:
         """
         Update the parameter store.
@@ -167,6 +174,8 @@ class MaxScoreOverRotations(AbstractAnalyzer):
         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)
@@ -536,13 +545,19 @@ class MaxScoreOverRotationsConstrained(MaxScoreOverRotations):
                 )
 
     def __call__(
-        self, state: Tuple, scores: BackendArray, rotation_matrix: BackendArray
+        self,
+        state: Tuple,
+        scores: BackendArray,
+        rotation_matrix: BackendArray,
+        **kwargs,
     ) -> 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)
+        return super().__call__(
+            state, scores=scores, rotation_matrix=rotation_matrix, **kwargs
+        )
 
     def _get_constraint(self, rotation_matrix: BackendArray) -> BackendArray:
         """
@@ -627,7 +642,11 @@ class MaxScoreOverTranslations(MaxScoreOverRotations):
         return scores, rotations, {}
 
     def __call__(
-        self, state, scores: BackendArray, rotation_matrix: BackendArray
+        self,
+        state,
+        scores: BackendArray,
+        rotation_matrix: BackendArray,
+        **kwargs,
     ) -> Tuple:
         prev_scores, rotations, rotation_mapping = state
 
@@ -738,5 +757,5 @@ class MaxScoreOverTranslations(MaxScoreOverRotations):
             cls._invert_rmap(master_rotation_mapping),
         )
 
-    def _postprocess(self, **kwargs):
-        return self
+    def result(self, state: Tuple, **kwargs) -> Tuple:
+        return state

tme/analyzer/base.py

@@ -38,16 +38,16 @@ class AbstractAnalyzer(ABC):
 
         Returns
         -------
-        state
-            Initial state tuple containing the analyzer's internal data
-            structures. The exact structure depends on the specific
-            implementation.
+        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
-        subsequent calls to __call__. The state should contain all
-        necessary data structures for accumulating analysis results.
+        :py:meth:`AbstractAnalyzer.__call__` and finally to
+        :py:meth:`AbstractAnalyzer.result`. The state should contain all necessary
+        data structures for accumulating analysis results.
         """
 
     @abstractmethod
@@ -57,49 +57,39 @@ class AbstractAnalyzer(ABC):
 
         Parameters
         ----------
-        state : object
-            Current analyzer state as returned by init_state() or
-            previous calls to __call__.
+        state : tuple
+            Current analyzer state as returned :py:meth:`AbstractAnalyzer.init_state`
+            or previous invocations of :py:meth:`AbstractAnalyzer.__call__`.
         scores : BackendArray
-            Array of scores computed for the current rotation.
+            Array of new scores with dimensionality d.
         rotation_matrix : BackendArray
-            Rotation matrix used to generate the scores.
+            Rotation matrix used to generate scores with shape (d,d).
         **kwargs : dict
-            Additional keyword arguments specific to the analyzer
-            implementation.
+            Keyword arguments used by specific implementations.
 
         Returns
         -------
-        state
-            Updated analyzer state with the new scoring data incorporated.
-
-        Notes
-        -----
-        This method should be pure functional - it should not modify
-        the input state but return a new state with the updates applied.
-        The exact signature may vary between implementations.
+        tuple
+            Updated analyzer state incorporating the new data.
         """
-        pass
 
     @abstractmethod
     def result(self, state: Tuple, **kwargs) -> Tuple:
         """
-        Finalize the analysis and produce the final result.
+        Finalize the analysis by performing potential post processing.
 
         Parameters
         ----------
         state : tuple
-            Final analyzer state containing all accumulated data.
+            Analyzer state containing accumulated data.
         **kwargs : dict
-            Additional keyword arguments for result processing,
-            such as postprocessing parameters.
+            Keyword arguments used by specific implementations.
 
         Returns
         -------
         result
-            Final analysis result. The exact format depends on the
-            analyzer implementation but typically includes processed
-            scores, rotation information, and metadata.
+            Final analysis result. The exact struccture depends on the
+            analyzer implementation.
 
         Notes
         -----
@@ -108,25 +98,24 @@ class AbstractAnalyzer(ABC):
         It may apply postprocessing operations like convolution mode
         correction or coordinate transformations.
         """
-        pass
 
     @classmethod
     @abstractmethod
     def merge(cls, results: List[Tuple], **kwargs) -> Tuple:
         """
-        Merge results from multiple analyzer instances.
+        Merge multiple analyzer results.
 
         Parameters
         ----------
-        results : list
-            List of result objects as returned by the result() method
-            from multiple analyzer instances.
+        results : list of tuple
+            List of tuple objects returned by :py:meth:`AbstractAnalyzer.result`
+            from different instances of the same analyzer class.
         **kwargs : dict
-            Additional keyword arguments for merge configuration.
+            Keyword arguments used by specific implementations.
 
         Returns
         -------
-        merged_result
+        tuple
             Single result object combining all input results.
 
         Notes

tme/analyzer/peaks.py

@@ -17,9 +17,9 @@ 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",
@@ -228,10 +228,15 @@ class PeakCaller(AbstractAnalyzer):
         translations = be.full(
             (self.num_peaks, ndim), fill_value=-1, dtype=be._int_dtype
         )
+
+        rdim = len(self.shape)
+        if self.batch_dims:
+            rdim -= len(self.batch_dims)
+
         rotations = be.full(
-            (self.num_peaks, ndim, ndim), fill_value=0, dtype=be._float_dtype
+            (self.num_peaks, rdim, rdim), fill_value=0, dtype=be._float_dtype
         )
-        for i in range(ndim):
+        for i in range(rdim):
             rotations[:, i, i] = 1.0
 
         scores = be.full((self.num_peaks,), fill_value=-1, dtype=be._float_dtype)
@@ -750,8 +755,7 @@ class PeakCallerRecursiveMasking(PeakCaller):
             Dictionary mapping values in rotations to Euler angles.
             By default None
         min_score : float
-            Minimum score value to consider. If provided, superseeds limit given
-            by :py:attr:`PeakCaller.num_peaks`.
+            Minimum score value to consider.
 
         Returns
         -------
@@ -765,54 +769,57 @@ 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
+        peaks = []
+        box = tuple(self.min_distance for _ in range(scores.ndim))
 
-        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)
+        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)
 
-        rotated_template = be.zeros(mask.shape, dtype=mask.dtype)
-
-        peak_limit = self.num_peaks
-        if min_score is not None:
-            peak_limit = be.size(scores)
-        else:
+        if min_score is None:
             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) >= self.num_peaks:
                 break
 
-        peaks = be.to_backend_array(coordinates)
-        return peaks, None
+        return be.to_backend_array(peaks), None
 
     @staticmethod
     def _get_rotation_matrix(
@@ -845,93 +852,13 @@ class PeakCallerRecursiveMasking(PeakCaller):
 
         rotation = rotation_mapping[rotation_space[tuple(peak)]]
 
-        # TODO: Newer versions of rotation mapping contain rotation matrices not angles
+        # Old versions of rotation mapping contained Euler angles
         if rotation.ndim != 2:
             rotation = be.to_backend_array(
                 euler_to_rotationmatrix(be.to_numpy_array(rotation))
             )
         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):
     """

tme/analyzer/proxy.py

@@ -18,6 +18,7 @@ from ..backends import backend as be
 
 __all__ = ["StatelessSharedAnalyzerProxy", "SharedAnalyzerProxy"]
 
+
 class StatelessSharedAnalyzerProxy:
     """
     Proxy that wraps functional analyzers for concurrent access via shared memory.