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

tme/analyzer.py

@@ -43,14 +43,29 @@ def filter_points_indices_bucket(
 
 
 def filter_points_indices(
-    coordinates: NDArray, min_distance: float, bucket_cutoff: int = 1e4
+    coordinates: NDArray,
+    min_distance: float,
+    bucket_cutoff: int = 1e4,
+    batch_dims: Tuple[int] = None,
 ) -> NDArray:
     if min_distance <= 0:
         return backend.arange(coordinates.shape[0])
+    if coordinates.shape[0] == 0:
+        return ()
+
+    if batch_dims is not None:
+        coordinates_new = backend.zeros(coordinates.shape, coordinates.dtype)
+        coordinates_new[:] = coordinates
+        coordinates_new[..., batch_dims] = backend.astype(
+            coordinates[..., batch_dims] * (2 * min_distance), coordinates_new.dtype
+        )
+        coordinates = coordinates_new
 
     if isinstance(coordinates, np.ndarray):
         return find_candidate_indices(coordinates, min_distance)
-    elif coordinates.shape[0] > bucket_cutoff:
+    elif coordinates.shape[0] > bucket_cutoff or not isinstance(
+        coordinates, np.ndarray
+    ):
         return filter_points_indices_bucket(coordinates, min_distance)
     distances = np.linalg.norm(coordinates[:, None] - coordinates, axis=-1)
     distances = np.tril(distances)
@@ -59,8 +74,10 @@ def filter_points_indices(
     return indices[keep == indices]
 
 
-def filter_points(coordinates: NDArray, min_distance: Tuple[int]) -> NDArray:
-    unique_indices = filter_points_indices(coordinates, min_distance)
+def filter_points(
+    coordinates: NDArray, min_distance: Tuple[int], batch_dims: Tuple[int] = None
+) -> NDArray:
+    unique_indices = filter_points_indices(coordinates, min_distance, batch_dims)
     coordinates = coordinates[unique_indices]
     return coordinates
 
@@ -77,6 +94,8 @@ class PeakCaller(ABC):
         Minimum distance between peaks.
     min_boundary_distance : int, optional
         Minimum distance to array boundaries.
+    batch_dims : int, optional
+        Peak calling batch dimensions.
     **kwargs
         Additional keyword arguments.
 
@@ -92,6 +111,7 @@ class PeakCaller(ABC):
         number_of_peaks: int = 1000,
         min_distance: int = 1,
         min_boundary_distance: int = 0,
+        batch_dims: Tuple[int] = None,
         **kwargs,
     ):
         number_of_peaks = int(number_of_peaks)
@@ -114,6 +134,10 @@ class PeakCaller(ABC):
         self.min_boundary_distance = min_boundary_distance
         self.number_of_peaks = number_of_peaks
 
+        self.batch_dims = batch_dims
+        if batch_dims is not None:
+            self.batch_dims = tuple(int(x) for x in self.batch_dims)
+
         # Postprocesing arguments
         self.fourier_shift = kwargs.get("fourier_shift", None)
         self.convolution_mode = kwargs.get("convolution_mode", None)
@@ -128,10 +152,41 @@ class PeakCaller(ABC):
         self.peak_list = [backend.to_cpu_array(arr) for arr in self.peak_list]
         yield from self.peak_list
 
+    @staticmethod
+    def _batchify(shape: Tuple[int], batch_dims: Tuple[int] = None) -> List:
+        if batch_dims is None:
+            yield (tuple(slice(None) for _ in shape), tuple(0 for _ in shape))
+            return None
+
+        batch_ranges = [range(shape[dim]) for dim in batch_dims]
+
+        def _generate_slices_recursive(current_dim, current_indices):
+            if current_dim == len(batch_dims):
+                slice_list, offset_list, batch_index = [], [], 0
+                for i in range(len(shape)):
+                    if i in batch_dims:
+                        index = current_indices[batch_index]
+                        slice_list.append(slice(index, index + 1))
+                        offset_list.append(index)
+                        batch_index += 1
+                    else:
+                        slice_list.append(slice(None))
+                        offset_list.append(0)
+                yield (tuple(slice_list), tuple(offset_list))
+            else:
+                for index in batch_ranges[current_dim]:
+                    yield from _generate_slices_recursive(
+                        current_dim + 1, current_indices + (index,)
+                    )
+
+        yield from _generate_slices_recursive(0, ())
+
     def __call__(
         self,
         score_space: NDArray,
         rotation_matrix: NDArray,
+        minimum_score: float = None,
+        maximum_score: float = None,
         **kwargs,
     ) -> None:
         """
@@ -143,59 +198,91 @@ class PeakCaller(ABC):
             Array containing the score space.
         rotation_matrix : NDArray
             Rotation matrix used to obtain the score array.
+        minimum_score : float
+            Minimum score from which to consider peaks. If provided, superseeds limits
+            presented by :py:attr:`PeakCaller.number_of_peaks`.
+        maximum_score : float
+            Maximum score upon which to consider peaks,
         **kwargs
             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
-        )
+        for subset, offset in self._batchify(score_space.shape, self.batch_dims):
+            peak_positions, peak_details = self.call_peaks(
+                score_space=score_space[subset],
+                rotation_matrix=rotation_matrix,
+                minimum_score=minimum_score,
+                maximum_score=maximum_score,
+                **kwargs,
+            )
 
-        if peak_positions is None:
-            return None
+            if peak_positions is None:
+                continue
+            if peak_positions.shape[0] == 0:
+                continue
 
-        peak_positions = backend.astype(peak_positions, int)
-        if peak_positions.shape[0] == 0:
-            return None
+            if peak_details is None:
+                peak_details = backend.full((peak_positions.shape[0],), fill_value=-1)
 
-        if peak_details is None:
-            peak_details = backend.to_backend_array([-1] * peak_positions.shape[0])
+            backend.add(peak_positions, offset, out=peak_positions)
+            peak_positions = backend.astype(peak_positions, int)
+            if self.min_boundary_distance > 0:
+                upper_limit = backend.subtract(
+                    score_space.shape, self.min_boundary_distance
+                )
+                valid_peaks = backend.multiply(
+                    peak_positions < upper_limit,
+                    peak_positions >= self.min_boundary_distance,
+                )
+                if self.batch_dims is not None:
+                    valid_peaks[..., self.batch_dims] = True
 
-        if self.min_boundary_distance > 0:
-            upper_limit = backend.subtract(
-                score_space.shape, self.min_boundary_distance
-            )
-            valid_peaks = (
-                backend.sum(
-                    backend.multiply(
-                        peak_positions < upper_limit,
-                        peak_positions >= self.min_boundary_distance,
-                    ),
-                    axis=1,
+                valid_peaks = (
+                    backend.sum(valid_peaks, axis=1) == peak_positions.shape[1]
                 )
-                == peak_positions.shape[1]
-            )
-            if backend.sum(valid_peaks) == 0:
-                return None
 
-            peak_positions, peak_details = (
-                peak_positions[valid_peaks],
-                peak_details[valid_peaks],
+                if backend.sum(valid_peaks) == 0:
+                    continue
+
+                peak_positions, peak_details = (
+                    peak_positions[valid_peaks],
+                    peak_details[valid_peaks],
+                )
+
+            peak_scores = score_space[tuple(peak_positions.T)]
+            if minimum_score is not None:
+                valid_peaks = peak_scores >= minimum_score
+                peak_positions, peak_details, peak_scores = (
+                    peak_positions[valid_peaks],
+                    peak_details[valid_peaks],
+                    peak_scores[valid_peaks],
+                )
+            if maximum_score is not None:
+                valid_peaks = peak_scores <= maximum_score
+                peak_positions, peak_details, peak_scores = (
+                    peak_positions[valid_peaks],
+                    peak_details[valid_peaks],
+                    peak_scores[valid_peaks],
+                )
+
+            if peak_positions.shape[0] == 0:
+                continue
+
+            rotations = backend.repeat(
+                rotation_matrix.reshape(1, *rotation_matrix.shape),
+                peak_positions.shape[0],
+                axis=0,
             )
 
-        rotations = backend.repeat(
-            rotation_matrix.reshape(1, *rotation_matrix.shape),
-            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=peak_scores,
+                rotations=rotations,
+                batch_offset=offset,
+                **kwargs,
+            )
 
-        self._update(
-            peak_positions=peak_positions,
-            peak_details=peak_details,
-            peak_scores=peak_scores,
-            rotations=rotations,
-            **kwargs,
-        )
+        return None
 
     @abstractmethod
     def call_peaks(
@@ -212,10 +299,8 @@ class PeakCaller(ABC):
         ----------
         score_space : NDArray
             Data array of scores.
-        minimum_score : float
-            Minimum score value to consider.
-        min_distance : float
-            Minimum distance between maxima.
+        **kwargs : Dict, optional
+            Keyword arguments passed to __call__.
 
         Returns
         -------
@@ -364,25 +449,52 @@ class PeakCaller(ABC):
         backend.add(peak_positions, translation_offset, out=peak_positions)
         if not len(self.peak_list):
             self.peak_list = [peak_positions, rotations, peak_scores, peak_details]
-            dim = peak_positions.shape[1]
-            peak_scores = backend.zeros((0,), peak_scores.dtype)
-            peak_details = backend.zeros((0,), peak_details.dtype)
-            rotations = backend.zeros((0, dim, dim), rotations.dtype)
-            peak_positions = backend.zeros((0, dim), peak_positions.dtype)
 
-        peaks = backend.concatenate((self.peak_list[0], peak_positions))
+        peak_positions = backend.concatenate((self.peak_list[0], peak_positions))
         rotations = backend.concatenate((self.peak_list[1], rotations))
         peak_scores = backend.concatenate((self.peak_list[2], peak_scores))
         peak_details = backend.concatenate((self.peak_list[3], peak_details))
 
-        top_n = min(backend.size(peak_scores), self.number_of_peaks)
-        top_scores, *_ = backend.topk_indices(peak_scores, top_n)
+        if self.batch_dims is None:
+            top_n = min(backend.size(peak_scores), self.number_of_peaks)
+            top_scores, *_ = backend.topk_indices(peak_scores, top_n)
+        else:
+            # Not very performant but fairly robust
+            batch_indices = peak_positions[..., self.batch_dims]
+            backend.subtract(
+                batch_indices, backend.min(batch_indices, axis=0), out=batch_indices
+            )
+            multiplier = backend.power(
+                backend.max(batch_indices, axis=0) + 1,
+                backend.arange(batch_indices.shape[1]),
+            )
+            backend.multiply(batch_indices, multiplier, out=batch_indices)
+            batch_indices = backend.sum(batch_indices, axis=1)
+            unique_indices, batch_counts = backend.unique(
+                batch_indices, return_counts=True
+            )
+            total_indices = backend.arange(peak_scores.shape[0])
+            batch_indices = [total_indices[batch_indices == x] for x in unique_indices]
+            top_scores = backend.concatenate(
+                [
+                    total_indices[indices][
+                        backend.topk_indices(
+                            peak_scores[indices], min(y, self.number_of_peaks)
+                        )
+                    ]
+                    for indices, y in zip(batch_indices, batch_counts)
+                ]
+            )
 
         final_order = top_scores[
-            filter_points_indices(peaks[top_scores], self.min_distance)
+            filter_points_indices(
+                coordinates=peak_positions[top_scores],
+                min_distance=self.min_distance,
+                batch_dims=self.batch_dims,
+            )
         ]
 
-        self.peak_list[0] = peaks[final_order,]
+        self.peak_list[0] = peak_positions[final_order,]
         self.peak_list[1] = rotations[final_order,]
         self.peak_list[2] = peak_scores[final_order]
         self.peak_list[3] = peak_details[final_order]
@@ -390,6 +502,9 @@ class PeakCaller(ABC):
     def _postprocess(
         self, fourier_shift, convolution_mode, targetshape, templateshape, **kwargs
     ):
+        if not len(self.peak_list):
+            return self
+
         peak_positions = self.peak_list[0]
         if not len(peak_positions):
             return self
@@ -402,12 +517,13 @@ class PeakCaller(ABC):
 
         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),
+                    backend.astype(
+                        backend.divide(peak_positions, score_space_shape), int
+                    ),
                     score_space_shape,
                 ),
                 out=peak_positions,
@@ -446,14 +562,10 @@ class PeakCaller(ABC):
 class PeakCallerSort(PeakCaller):
     """
     A :py:class:`PeakCaller` subclass that first selects ``number_of_peaks``
-    highest scores and subsequently filters local maxima to suffice a distance
-    from one another of ``min_distance``.
-
+    highest scores.
     """
 
-    def call_peaks(
-        self, score_space: NDArray, minimum_score: float = None, **kwargs
-    ) -> Tuple[NDArray, NDArray]:
+    def call_peaks(self, score_space: NDArray, **kwargs) -> Tuple[NDArray, NDArray]:
         """
         Call peaks in the score space.
 
@@ -461,9 +573,6 @@ class PeakCallerSort(PeakCaller):
         ----------
         score_space : NDArray
             Data array of scores.
-        minimum_score : float
-            Minimum score value to consider. If provided, superseeds limit given
-            by :py:attr:`PeakCaller.number_of_peaks`.
 
         Returns
         -------
@@ -473,16 +582,12 @@ class PeakCallerSort(PeakCaller):
         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)
         coordinates = backend.transpose(backend.stack(coordinates))
 
-        peaks = filter_points(coordinates, self.min_distance)
-        return peaks, None
+        return coordinates, None
 
 
 class PeakCallerMaximumFilter(PeakCaller):
@@ -492,9 +597,7 @@ class PeakCallerMaximumFilter(PeakCaller):
     skimage.feature.peak_local_max.
     """
 
-    def call_peaks(
-        self, score_space: NDArray, minimum_score: float = None, **kwargs
-    ) -> Tuple[NDArray, NDArray]:
+    def call_peaks(self, score_space: NDArray, **kwargs) -> Tuple[NDArray, NDArray]:
         """
         Call peaks in the score space.
 
@@ -502,9 +605,8 @@ class PeakCallerMaximumFilter(PeakCaller):
         ----------
         score_space : NDArray
             Data array of scores.
-        minimum_score : float
-            Minimum score value to consider. If provided, superseeds limit given
-            by :py:attr:`PeakCaller.number_of_peaks`.
+        kwargs: Dict, optional
+            Optional keyword arguments.
 
         Returns
         -------
@@ -513,17 +615,6 @@ class PeakCallerMaximumFilter(PeakCaller):
         """
         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
-        )
-        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
 
 
@@ -536,9 +627,7 @@ class PeakCallerFast(PeakCaller):
 
     """
 
-    def call_peaks(
-        self, score_space: NDArray, minimum_score: float = None, **kwargs
-    ) -> Tuple[NDArray, NDArray]:
+    def call_peaks(self, score_space: NDArray, **kwargs) -> Tuple[NDArray, NDArray]:
         """
         Call peaks in the score space.
 
@@ -546,9 +635,6 @@ class PeakCallerFast(PeakCaller):
         ----------
         score_space : NDArray
             Data array of scores.
-        minimum_score : float
-            Minimum score value to consider. If provided, superseeds limit given
-            by :py:attr:`PeakCaller.number_of_peaks`.
 
         Returns
         -------
@@ -580,26 +666,24 @@ class PeakCallerFast(PeakCaller):
         if coordinates.shape[0] == 0:
             return None
 
-        peaks = filter_points(coordinates, self.min_distance)
-
-        starts = backend.maximum(peaks - self.min_distance, 0)
-        stops = backend.minimum(peaks + self.min_distance, score_space.shape)
+        starts = backend.maximum(coordinates - self.min_distance, 0)
+        stops = backend.minimum(coordinates + self.min_distance, score_space.shape)
         slices_list = [
             tuple(slice(*coord) for coord in zip(start_row, stop_row))
             for start_row, stop_row in zip(starts, stops)
         ]
 
-        scores = score_space[tuple(peaks.T)]
+        scores = score_space[tuple(coordinates.T)]
         keep = [
             score >= backend.max(score_space[subvol])
             for subvol, score in zip(slices_list, scores)
         ]
-        peaks = peaks[keep,]
+        coordinates = coordinates[keep,]
 
-        if len(peaks) == 0:
-            return peaks, None
+        if len(coordinates) == 0:
+            return coordinates, None
 
-        return peaks, None
+        return coordinates, None
 
 
 class PeakCallerRecursiveMasking(PeakCaller):
@@ -655,7 +739,7 @@ class PeakCallerRecursiveMasking(PeakCaller):
         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)
+            mask = backend.zeros(shape, dtype=backend._float_dtype)
 
         rotated_template = backend.zeros(mask.shape, dtype=mask.dtype)
 
@@ -665,12 +749,15 @@ class PeakCallerRecursiveMasking(PeakCaller):
         else:
             minimum_score = backend.min(score_space) - 1
 
+        scores = backend.zeros(score_space.shape, dtype=score_space.dtype)
+        scores[:] = score_space
+
         while True:
-            backend.argmax(score_space)
+            backend.argmax(scores)
             peak = backend.unravel_index(
-                indices=backend.argmax(score_space), shape=score_space.shape
+                indices=backend.argmax(scores), shape=scores.shape
             )
-            if score_space[tuple(peak)] < minimum_score:
+            if scores[tuple(peak)] < minimum_score:
                 break
 
             coordinates.append(peak)
@@ -683,7 +770,7 @@ class PeakCallerRecursiveMasking(PeakCaller):
             )
 
             masking_function(
-                score_space=score_space,
+                score_space=scores,
                 rotation_matrix=current_rotation_matrix,
                 peak=peak,
                 mask=mask,
@@ -839,18 +926,23 @@ class PeakCallerScipy(PeakCaller):
         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
 
+        non_squeezable_dims = tuple(
+            i for i, x in enumerate(score_space.shape) if x != 1
+        )
         peaks = peak_local_max(
-            score_space,
+            np.squeeze(score_space),
             num_peaks=num_peaks,
             min_distance=self.min_distance,
             threshold_abs=minimum_score,
         )
+        peaks_full = np.zeros((peaks.shape[0], score_space.ndim), peaks.dtype)
+        peaks_full[..., non_squeezable_dims] = peaks[:]
+        peaks = backend.to_backend_array(peaks_full)
         return peaks, None
 
 
@@ -1164,6 +1256,53 @@ class MaxScoreOverRotations:
         Whether to offload internal data arrays to disk
     thread_safe: bool, optional
         Whether access to internal data arrays should be thread safe
+
+    Examples
+    --------
+    The following achieves the minimal definition of a :py:class:`MaxScoreOverRotations`
+    instance
+
+    >>> from tme.analyzer import MaxScoreOverRotations
+    >>> analyzer = MaxScoreOverRotations(
+    >>>    score_space_shape = (50, 50),
+    >>>    score_space_dtype = np.float32,
+    >>>    rotation_space_dtype = np.int32,
+    >>> )
+
+    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
+
+    >>> for rotation_number in range(10):
+    >>>     scores = np.random.rand(50,50)
+    >>>     rotation = np.random.rand(scores.ndim, scores.ndim)
+    >>>     analyzer(score_space = scores, rotation_matrix = rotation)
+
+    The aggregated scores can be exctracted by invoking the __iter__ method of
+    ``analyzer``
+
+    >>> results = tuple(analyzer)
+
+    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).
+
+    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)
+
+    The outlined procedure is a trivial method to identify high scoring peaks.
+    Alternatively, :py:class:`PeakCaller` offers a range of more elaborate approaches
+    that can be used.
     """
 
     def __init__(
@@ -1316,11 +1455,11 @@ class MaxScoreOverRotations:
             Arbitrary keyword arguments.
         """
         rotation = backend.tobytes(rotation_matrix)
-        rotation_index = self.observed_rotations.setdefault(
-            rotation, len(self.observed_rotations)
-        )
 
         if self.lock_is_nullcontext:
+            rotation_index = self.observed_rotations.setdefault(
+                rotation, len(self.observed_rotations)
+            )
             backend.max_score_over_rotations(
                 score_space=score_space,
                 internal_scores=self.score_space,
@@ -1330,6 +1469,9 @@ class MaxScoreOverRotations:
             return None
 
         with self.lock:
+            rotation_index = self.observed_rotations.setdefault(
+                rotation, len(self.observed_rotations)
+            )
             internal_scores = backend.sharedarr_to_arr(
                 shape=self.score_space_shape,
                 dtype=self.score_space_dtype,

tme/backends/__init__.py

@@ -57,6 +57,7 @@ class BackendManager:
             "cpu_backend": NumpyFFTWBackend,
             "pytorch": PytorchBackend,
             "cupy": CupyBackend,
+            "mlx": MLXBackend,
         }
         self._backend = NumpyFFTWBackend()
         self._backend_name = "cpu_backend"