pytme 0.3.1.post2__cp311-cp311-macosx_15_0_arm64.whl → 0.3.2.dev0__cp311-cp311-macosx_15_0_arm64.whl

tme/backends/npfftw_backend.py

@@ -6,7 +6,6 @@ Copyright (c) 2023 European Molecular Biology Laboratory
 Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
-import os
 from psutil import virtual_memory
 from contextlib import contextmanager
 from typing import Tuple, List, Type
@@ -26,12 +25,6 @@ from ..types import NDArray, BackendArray, shm_type
 from .matching_backend import MatchingBackend, _create_metafunction
 
 
-os.environ["MKL_NUM_THREADS"] = "1"
-os.environ["OMP_NUM_THREADS"] = "1"
-os.environ["PYFFTW_NUM_THREADS"] = "1"
-os.environ["OPENBLAS_NUM_THREADS"] = "1"
-
-
 def create_ufuncs(obj):
     ufuncs = [
         "add",
@@ -247,12 +240,13 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
 
         shm = shared_memory_handler.SharedMemory(size=arr.nbytes)
         np_array = self.ndarray(arr.shape, dtype=arr.dtype, buffer=shm.buf)
-        np_array[:] = arr[:].copy()
+        np_array[:] = arr[:]
         return shm, arr.shape, arr.dtype
 
-    def topleft_pad(self, arr: NDArray, shape: Tuple[int], padval: int = 0) -> NDArray:
-        b = self.zeros(shape, arr.dtype)
-        self.add(b, padval, out=b)
+    def topleft_pad(
+        self, arr: NDArray, shape: Tuple[int], padval: float = 0
+    ) -> NDArray:
+        b = self.full(shape, fill_value=padval, dtype=arr.dtype)
         aind = [slice(None, None)] * arr.ndim
         bind = [slice(None, None)] * arr.ndim
         for i in range(arr.ndim):
@@ -311,15 +305,6 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
             **kwargs,
         )
 
-    def extract_center(self, arr: NDArray, newshape: Tuple[int]) -> NDArray:
-        new_shape = self.to_backend_array(newshape)
-        current_shape = self.to_backend_array(arr.shape)
-        starts = self.subtract(current_shape, new_shape)
-        starts = self.astype(self.divide(starts, 2), self._int_dtype)
-        stops = self.astype(self.add(starts, new_shape), self._int_dtype)
-        box = tuple(slice(start, stop) for start, stop in zip(starts, stops))
-        return arr[box]
-
     def compute_convolution_shapes(
         self, arr1_shape: Tuple[int], arr2_shape: Tuple[int]
     ) -> Tuple[List[int], List[int], List[int]]:
@@ -329,62 +314,65 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
 
         return convolution_shape, fast_shape, fast_ft_shape
 
-    def _rigid_transform_matrix(
+    def _build_transform_matrix(
         self,
-        rotation_matrix: NDArray,
-        translation: NDArray = None,
-        center: NDArray = None,
-    ) -> NDArray:
+        rotation_matrix: BackendArray,
+        translation: BackendArray = None,
+        center: BackendArray = None,
+        **kwargs,
+    ) -> BackendArray:
         ndim = rotation_matrix.shape[0]
-        matrix = self.identity(ndim + 1, dtype=self._float_dtype)
 
+        spatial_slice = slice(0, ndim)
+        matrix = self.eye(ndim + 1, dtype=self._float_dtype)
+
+        rotation_matrix = self.astype(rotation_matrix, self._float_dtype)
+        matrix = self.at(matrix, (spatial_slice, spatial_slice), rotation_matrix)
+
+        total_translation = self.zeros(ndim, dtype=self._float_dtype)
         if translation is not None:
-            translation_matrix = self.identity(ndim + 1, dtype=self._float_dtype)
-            translation_matrix[:ndim, ndim] = -translation
-            self.dot(matrix, translation_matrix, out=matrix)
+            translation = self.astype(translation, self._float_dtype)
+            total_translation = self.subtract(total_translation, translation)
 
         if center is not None:
-            center_matrix = self.identity(ndim + 1, dtype=self._float_dtype)
-            center_matrix[:ndim, ndim] = center
-            self.dot(matrix, center_matrix, out=matrix)
+            total_translation = self.add(total_translation, center)
+            rotated_center = self.matmul(rotation_matrix, center)
+            total_translation = self.subtract(total_translation, rotated_center)
 
-        if rotation_matrix is not None:
-            rmat = self.identity(ndim + 1, dtype=self._float_dtype)
-            rmat[:ndim, :ndim] = self._array_backend.linalg.inv(rotation_matrix)
-            self.dot(matrix, rmat, out=matrix)
+        matrix = self.at(matrix, (spatial_slice, ndim), total_translation)
+        return self.to_backend_array(matrix)
 
-        if center is not None:
-            center_matrix[:ndim, ndim] = -center_matrix[:ndim, ndim]
-            self.dot(matrix, center_matrix, out=matrix)
+    def _batch_transform_matrix(self, matrix: NDArray) -> NDArray:
+        ndim = matrix.shape[0] + 1
 
-        matrix /= matrix[ndim, ndim]
-        return matrix
+        ret = self.zeros((ndim, ndim), dtype=matrix.dtype)
+        ret = self.at(ret, (0, 0), 1)
 
-    def _rigid_transform(
+        spatial_slice = slice(1, ndim)
+        ret = self.at(ret, (spatial_slice, spatial_slice), matrix)
+        return ret
+
+    def _compute_transform_center(
+        self, arr: NDArray, use_geometric_center: bool, batched: bool = False
+    ) -> NDArray:
+        center = self.divide(self.to_backend_array(arr.shape) - 1, 2)
+        if not use_geometric_center:
+            center = self.center_of_mass(arr, cutoff=0)
+        if batched:
+            return center[1:]
+        return center
+
+    def _transform(
         self,
         data: NDArray,
         matrix: NDArray,
         output: NDArray,
         prefilter: bool,
         order: int,
-        cache: bool = False,
-        batched=False,
-    ) -> None:
-        if batched:
-            for i in range(data.shape[0]):
-                self._rigid_transform(
-                    data=data[i],
-                    matrix=matrix,
-                    output=output[i],
-                    prefilter=prefilter,
-                    order=order,
-                    cache=cache,
-                    batched=False,
-                )
-            return None
-
+        **kwargs,
+    ) -> NDArray:
         out_slice = tuple(slice(0, stop) for stop in data.shape)
-        self.affine_transform(
+        return self.affine_transform(
             input=data,
             matrix=matrix,
             mode="constant",
@@ -393,72 +381,87 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
             prefilter=prefilter,
         )
 
-    def rigid_transform(
+    def _rigid_transform(
         self,
         arr: NDArray,
-        rotation_matrix: NDArray,
+        matrix: NDArray,
         arr_mask: NDArray = None,
-        translation: NDArray = None,
-        use_geometric_center: bool = False,
         out: NDArray = None,
         out_mask: NDArray = None,
         order: int = 3,
         cache: bool = False,
-        batched: bool = False,
+        **kwargs,
     ) -> Tuple[NDArray, NDArray]:
         if out is None:
             out = self.zeros_like(arr)
 
-        # Check whether rotation_matrix is already a rigid transform matrix
-        matrix = rotation_matrix
-        if matrix.shape[-1] == (arr.ndim - int(batched)):
-            center = self.divide(self.to_backend_array(arr.shape) - 1, 2)
-            if not use_geometric_center:
-                center = self.center_of_mass(arr, cutoff=0)
-
-            offset = int(arr.ndim - rotation_matrix.shape[0])
-            center = center[offset:]
-            translation = (
-                self.zeros(center.size) if translation is None else translation
-            )
-            matrix = self._rigid_transform_matrix(
-                rotation_matrix=rotation_matrix,
-                translation=translation,
-                center=center,
-            )
-
-        self._rigid_transform(
+        out = self._transform(
             data=arr,
             matrix=matrix,
             output=out,
             order=order,
             prefilter=True,
             cache=cache,
-            batched=batched,
         )
 
-        # Applying the prefilter leads to artifacts in the mask.
         if arr_mask is not None:
             if out_mask is None:
-                out_mask = self.zeros_like(arr_mask)
+                out_mask = self.zeros_like(arr)
 
-            self._rigid_transform(
+            # Applying the prefilter leads to artifacts in the mask.
+            out_mask = self._transform(
                 data=arr_mask,
                 matrix=matrix,
                 output=out_mask,
                 order=order,
                 prefilter=False,
                 cache=cache,
-                batched=batched,
             )
 
         return out, out_mask
 
+    def rigid_transform(
+        self,
+        arr: NDArray,
+        rotation_matrix: NDArray,
+        arr_mask: NDArray = None,
+        translation: NDArray = None,
+        use_geometric_center: bool = False,
+        out: NDArray = None,
+        out_mask: NDArray = None,
+        order: int = 3,
+        cache: bool = False,
+        batched: bool = False,
+    ) -> Tuple[NDArray, NDArray]:
+        matrix = rotation_matrix
+
+        # Build transformation matrix from rotation matrix
+        if matrix.shape[-1] == (arr.ndim - int(batched)):
+            center = self._compute_transform_center(arr, use_geometric_center, batched)
+            matrix = self._build_transform_matrix(
+                rotation_matrix=rotation_matrix,
+                translation=translation,
+                center=self.astype(center, self._float_dtype),
+                shape=arr.shape[1:] if batched else arr.shape,
+            )
+
+        if batched:
+            matrix = self._batch_transform_matrix(matrix)
+
+        return self._rigid_transform(
+            arr=arr,
+            arr_mask=arr_mask,
+            out=out,
+            out_mask=out_mask,
+            matrix=matrix,
+            cache=cache,
+            order=order,
+            batched=batched,
+        )
+
     def center_of_mass(self, arr: BackendArray, cutoff: float = None) -> BackendArray:
         """
-        Computes the center of mass of a numpy ndarray instance using all available
-        elements. For template matching it typically makes sense to only input
-        positive densities.
+        Computes the center of mass of an array larger than cutoff.
 
         Parameters
         ----------
@@ -466,19 +469,19 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
             Array to compute the center of mass of.
         cutoff : float, optional
             Densities less than or equal to cutoff are nullified for center
-            of mass computation. By default considers all values.
+            of mass computation. Defaults to None.
 
         Returns
         -------
         BackendArray
             Center of mass with shape (arr.ndim).
         """
-        cutoff = self.min(arr) - 1 if cutoff is None else cutoff
-
-        arr = self.where(arr > cutoff, arr, 0)
-        denominator = self.sum(arr)
+        arr = self.abs(arr)
+        if cutoff is not None:
+            arr = self.where(arr > cutoff, arr, 0)
 
         grids = []
+        denominator = self.sum(arr)
         for i, x in enumerate(arr.shape):
             baseline_dims = tuple(1 if i != t else x for t in range(len(arr.shape)))
             grids.append(
@@ -587,7 +590,7 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
         sq_exp = self.sqrt(sq_exp, out=sq_exp)
 
         # Assume that low stdev regions also have low scores
-        # See :py:meth:`tme.matching_exhaustive.flcSphericalMask_setup` for correct norm
+        # See :py:meth:`tme.matching_scores.flcSphericalMask_setup` for correct norm
         sq_exp[sq_exp < eps] = 1
         sq_exp = self.multiply(sq_exp, n_obs, out=sq_exp)
         return self.divide(arr, sq_exp, out=out)

tme/backends/pytorch_backend.py

@@ -281,22 +281,45 @@ class PytorchBackend(NumpyFFTWBackend):
         kwargs["dim"] = kwargs.pop("axes", None)
         return self._array_backend.fft.irfftn(arr, **kwargs)
 
-    def _rigid_transform_matrix(self, rotation_matrix, *args, **kwargs):
-        return rotation_matrix
+    def _build_transform_matrix(
+        self,
+        shape: Tuple[int],
+        rotation_matrix: TorchTensor,
+        translation: TorchTensor = None,
+        center: TorchTensor = None,
+        **kwargs,
+    ) -> TorchTensor:
+        """
+        Express the transform matrix in normalized coordinates.
+        """
+        shape = self.to_backend_array(shape) - 1
+
+        scale_factors = 2.0 / shape
+        if center is not None:
+            center = center - shape / 2
+            center = center * scale_factors
+
+        if translation is not None:
+            translation = translation * scale_factors
+
+        return super()._build_transform_matrix(
+            rotation_matrix=self.flip(rotation_matrix, [0, 1]),
+            translation=translation,
+            center=center,
+        )
 
-    def rigid_transform(
+    def _rigid_transform(
         self,
         arr: TorchTensor,
-        rotation_matrix: TorchTensor,
+        matrix: TorchTensor,
         arr_mask: TorchTensor = None,
-        translation: TorchTensor = None,
-        use_geometric_center: bool = False,
         out: TorchTensor = None,
         out_mask: TorchTensor = None,
         order: int = 1,
-        cache: bool = False,
+        batched: bool = False,
         **kwargs,
-    ):
+    ) -> Tuple[TorchTensor, TorchTensor]:
+        """Apply rigid transformation using homogeneous transformation matrix."""
         _mode_mapping = {0: "nearest", 1: "bilinear", 3: "bicubic"}
         mode = _mode_mapping.get(order, None)
         if mode is None:
@@ -305,90 +328,54 @@ class PytorchBackend(NumpyFFTWBackend):
                 f"Got {order} but supported interpolation orders are: {modes}."
             )
 
-        out = self.zeros_like(arr) if out is None else out
-
-        if translation is None:
-            translation = self._array_backend.zeros(arr.ndim, device=arr.device)
-
-        normalized_translation = self.divide(
-            -2.0 * translation, self.tensor(arr.shape, device=arr.device)
-        )
-        rotation_matrix_pull = self.linalg.inv(self.flip(rotation_matrix, [0, 1]))
-
-        out_slice = tuple(slice(0, x) for x in arr.shape)
-        subset = tuple(slice(None) for _ in range(arr.ndim))
-        offset = max(int(arr.ndim - rotation_matrix.shape[0]) - 1, 0)
-        if offset > 0:
-            normalized_translation = normalized_translation[offset:]
-            subset = tuple(0 if i < offset else slice(None) for i in range(arr.ndim))
-            out_slice = tuple(
-                slice(0, 1) if i < offset else slice(0, x)
-                for i, x in enumerate(arr.shape)
-            )
-
-        out[out_slice] = self._affine_transform(
-            arr=arr[subset],
-            rotation_matrix=rotation_matrix_pull,
-            translation=normalized_translation,
-            mode=mode,
-        )
-
-        if arr_mask is not None:
-            out_mask_slice = tuple(slice(0, x) for x in arr_mask.shape)
-            if out_mask is None:
-                out_mask = self._array_backend.zeros_like(arr_mask)
-            out_mask[out_mask_slice] = self._affine_transform(
-                arr=arr_mask[subset],
-                rotation_matrix=rotation_matrix_pull,
-                translation=normalized_translation,
-                mode=mode,
-            )
-
-        return out, out_mask
-
-    def _affine_transform(
-        self,
-        arr: TorchTensor,
-        rotation_matrix: TorchTensor,
-        translation: TorchTensor,
-        mode,
-    ) -> TorchTensor:
-        batched = arr.ndim != rotation_matrix.shape[0]
-
         batch_size, spatial_dims = 1, arr.shape
+        out_slice = tuple(slice(0, x) for x in arr.shape)
         if batched:
-            translation = translation[1:]
+            matrix = matrix[1:, 1:]
             batch_size, *spatial_dims = arr.shape
 
-        n_dims = len(spatial_dims)
-        transformation_matrix = self._array_backend.zeros(
-            n_dims, n_dims + 1, device=arr.device, dtype=arr.dtype
-        )
-
-        transformation_matrix[:, :n_dims] = rotation_matrix
-        transformation_matrix[:, n_dims] = translation
-        transformation_matrix = transformation_matrix.unsqueeze(0).expand(
-            batch_size, -1, -1
-        )
+        # Remove homogeneous row and expand for batch processing
+        matrix = matrix[:-1, :].to(arr.dtype)
+        matrix = matrix.unsqueeze(0).expand(batch_size, -1, -1)
 
-        if not batched:
-            arr = arr.unsqueeze(0)
-
-        size = self.Size([batch_size, 1, *spatial_dims])
         grid = self.F.affine_grid(
-            theta=transformation_matrix, size=size, align_corners=False
+            theta=matrix.to(arr.dtype),
+            size=self.Size([batch_size, 1, *spatial_dims]),
+            align_corners=False,
         )
-        output = self.F.grid_sample(
+
+        arr = arr.unsqueeze(0) if not batched else arr
+        ret = self.F.grid_sample(
             input=arr.unsqueeze(1),
             grid=grid,
             mode=mode,
             align_corners=False,
-        )
+        ).squeeze(1)
+
+        ret_mask = None
+        if arr_mask is not None:
+            arr_mask = arr_mask.unsqueeze(0) if not batched else arr_mask
+            ret_mask = self.F.grid_sample(
+                input=arr_mask.unsqueeze(1),
+                grid=grid,
+                mode=mode,
+                align_corners=False,
+            ).squeeze(1)
 
         if not batched:
-            output = output.squeeze(0)
+            ret = ret.squeeze(0)
+            ret_mask = ret_mask.squeeze(0) if arr_mask is not None else None
+
+        if out is not None:
+            out[out_slice] = ret
+        else:
+            out = ret
 
-        return output.squeeze(1)
+        if out_mask is not None:
+            out_mask[out_slice] = ret_mask
+        else:
+            out_mask = ret_mask
+        return out, out_mask
 
     def get_available_memory(self) -> int:
         if self.device == "cpu":

tme/cli.py

@@ -52,7 +52,7 @@ def match_template(
     """
     from .matching_data import MatchingData
     from .analyzer import MaxScoreOverRotations
-    from .matching_exhaustive import scan_subsets, MATCHING_EXHAUSTIVE_REGISTER
+    from .matching_exhaustive import match_exhaustive, MATCHING_EXHAUSTIVE_REGISTER
 
     if rotations is None:
         rotations = np.eye(target.ndim).reshape(1, target.ndim, target.ndim)
@@ -73,7 +73,7 @@ def match_template(
     matching_setup, matching_score = MATCHING_EXHAUSTIVE_REGISTER[score]
 
     candidates = list(
-        scan_subsets(
+        match_exhaustive(
             matching_data=matching_data,
             matching_score=matching_score,
             matching_setup=matching_setup,