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

tme/backends/cupy_backend.py

@@ -25,29 +25,37 @@ class CupyBackend(NumpyFFTWBackend):
     """
 
     def __init__(
-        self, default_dtype=None, complex_dtype=None, default_dtype_int=None, **kwargs
+        self,
+        float_dtype=None,
+        complex_dtype=None,
+        int_dtype=None,
+        overflow_safe_dtype=None,
+        **kwargs,
     ):
         import cupy as cp
         from cupyx.scipy.fft import get_fft_plan
         from cupyx.scipy.ndimage import affine_transform
         from cupyx.scipy.ndimage import maximum_filter
 
-        default_dtype = cp.float32 if default_dtype is None else default_dtype
+        float_dtype = cp.float32 if float_dtype is None else float_dtype
         complex_dtype = cp.complex64 if complex_dtype is None else complex_dtype
-        default_dtype_int = cp.int32 if default_dtype_int is None else default_dtype_int
+        int_dtype = cp.int32 if int_dtype is None else int_dtype
+        if overflow_safe_dtype is None:
+            overflow_safe_dtype = cp.float32
 
         super().__init__(
             array_backend=cp,
-            default_dtype=default_dtype,
+            float_dtype=float_dtype,
             complex_dtype=complex_dtype,
-            default_dtype_int=default_dtype_int,
+            int_dtype=int_dtype,
+            overflow_safe_dtype=overflow_safe_dtype,
         )
         self.get_fft_plan = get_fft_plan
         self.affine_transform = affine_transform
         self.maximum_filter = maximum_filter
 
-        floating = f"float{self.datatype_bytes(default_dtype) * 8}"
-        integer = f"int{self.datatype_bytes(default_dtype_int) * 8}"
+        floating = f"float{self.datatype_bytes(float_dtype) * 8}"
+        integer = f"int{self.datatype_bytes(int_dtype) * 8}"
         self._max_score_over_rotations = self._array_backend.ElementwiseKernel(
             f"{floating} internal_scores, {floating} scores, {integer} rot_index",
             f"{floating} out1, {integer} rotations",
@@ -119,12 +127,11 @@ class CupyBackend(NumpyFFTWBackend):
         fast_ft_shape: Tuple[int],
         real_dtype: type,
         complex_dtype: type,
-        fftargs: Dict = {},
         inverse_fast_shape: Tuple[int] = None,
         **kwargs,
     ) -> Tuple[Callable, Callable]:
         """
-        Build pyFFTW builder functions.
+        Build rfftn and irfftn functions.
 
         Parameters
         ----------
@@ -140,8 +147,6 @@ class CupyBackend(NumpyFFTWBackend):
             Numpy dtype of the fourier transform.
         inverse_fast_shape : tuple, optional
             Output shape of the inverse Fourier transform. By default fast_shape.
-        fftargs : dict, optional
-            Dictionary passed to pyFFTW builders.
         **kwargs: dict, optional
             Unused keyword arguments.
 
@@ -260,11 +265,13 @@ class CupyBackend(NumpyFFTWBackend):
         return_type = (out is None) + 2 * rotate_mask * (out_mask is None)
         translation = self.zeros(arr.ndim) if translation is None else translation
 
-        center = self.divide(self.to_backend_array(arr.shape), 2)
+        center = self.divide(arr.shape, 2)
         if not use_geometric_center:
             center = self.center_of_mass(arr, cutoff=0)
 
-        rotation_matrix_inverted = self.linalg.inv(rotation_matrix)
+        rotation_matrix_inverted = self.linalg.inv(
+            rotation_matrix.astype(self._overflow_safe_dtype)
+        ).astype(self._float_dtype)
         transformed_center = rotation_matrix_inverted @ center.reshape(-1, 1)
         transformed_center = transformed_center.reshape(-1)
         base_offset = self.subtract(center, transformed_center)

tme/backends/jax_backend.py

@@ -0,0 +1,218 @@
+""" Backend using jax for template matching.
+
+    Copyright (c) 2023 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+from typing import Tuple, Callable
+
+from .npfftw_backend import NumpyFFTWBackend
+
+class JaxBackend(NumpyFFTWBackend):
+    def __init__(
+        self, float_dtype=None, complex_dtype=None, int_dtype=None, **kwargs
+    ):
+        import jax.scipy as jsp
+        import jax.numpy as jnp
+
+        float_dtype = jnp.float32 if float_dtype is None else float_dtype
+        complex_dtype = jnp.complex64 if complex_dtype is None else complex_dtype
+        int_dtype = jnp.int32 if int_dtype is None else int_dtype
+
+        self.scipy = jsp
+        super().__init__(
+            array_backend=jnp,
+            float_dtype=float_dtype,
+            complex_dtype=complex_dtype,
+            int_dtype=int_dtype,
+        )
+
+    def to_backend_array(self, arr):
+        return self._array_backend.asarray(arr)
+
+    def preallocate_array(self, shape: Tuple[int], dtype: type):
+        arr = self._array_backend.zeros(shape, dtype=dtype)
+        return arr
+
+    def topleft_pad(self, arr, shape: Tuple[int], padval: int = 0):
+        b = self.preallocate_array(shape, arr.dtype)
+        self.add(b, padval, out=b)
+        aind = [slice(None, None)] * arr.ndim
+        bind = [slice(None, None)] * arr.ndim
+        for i in range(arr.ndim):
+            if arr.shape[i] > shape[i]:
+                aind[i] = slice(0, shape[i])
+            elif arr.shape[i] < shape[i]:
+                bind[i] = slice(0, arr.shape[i])
+        b = b.at[tuple(bind)].set(arr[tuple(aind)])
+        return b
+
+
+    def add(self, x1, x2, out = None, *args, **kwargs):
+        x1 = self.to_backend_array(x1)
+        x2 = self.to_backend_array(x2)
+        ret = self._array_backend.add(x1, x2, *args, **kwargs)
+
+        if out is not None:
+            out = out.at[:].set(ret)
+        return ret
+
+    def subtract(self, x1, x2, out = None, *args, **kwargs):
+        x1 = self.to_backend_array(x1)
+        x2 = self.to_backend_array(x2)
+        ret = self._array_backend.subtract(x1, x2, *args, **kwargs)
+        if out is not None:
+            out = out.at[:].set(ret)
+        return ret
+
+    def multiply(self, x1, x2, out = None, *args, **kwargs):
+        x1 = self.to_backend_array(x1)
+        x2 = self.to_backend_array(x2)
+        ret = self._array_backend.multiply(x1, x2, *args, **kwargs)
+        if out is not None:
+            out = out.at[:].set(ret)
+        return ret
+
+    def divide(self, x1, x2, out = None, *args, **kwargs):
+        x1 = self.to_backend_array(x1)
+        x2 = self.to_backend_array(x2)
+        ret = self._array_backend.divide(x1, x2, *args, **kwargs)
+        if out is not None:
+            out = out.at[:].set(ret)
+        return ret
+
+    def fill(self, arr, value: float) -> None:
+        arr.at[:].set(value)
+
+
+    def build_fft(
+        self,
+        fast_shape: Tuple[int],
+        fast_ft_shape: Tuple[int],
+        inverse_fast_shape: Tuple[int] = None,
+        **kwargs,
+    ) -> Tuple[Callable, Callable]:
+        """
+        Build fft builder functions.
+
+        Parameters
+        ----------
+        fast_shape : tuple
+            Tuple of integers corresponding to fast convolution shape
+            (see :py:meth:`PytorchBackend.compute_convolution_shapes`).
+        fast_ft_shape : tuple
+            Tuple of integers corresponding to the shape of the Fourier
+            transform array (see :py:meth:`PytorchBackend.compute_convolution_shapes`).
+        inverse_fast_shape : tuple, optional
+            Output shape of the inverse Fourier transform. By default fast_shape.
+        **kwargs : dict, optional
+            Unused keyword arguments.
+
+        Returns
+        -------
+        tuple
+            Tupple containing callable rfft and irfft object.
+        """
+        if inverse_fast_shape is None:
+            inverse_fast_shape = fast_shape
+
+        def rfftn(
+            arr, out, shape: Tuple[int] = fast_shape
+        ) -> None:
+            out = out.at[:].set(self._array_backend.fft.rfftn(arr, s=shape))
+
+        def irfftn(
+            arr, out, shape: Tuple[int] = inverse_fast_shape
+        ) -> None:
+            out = out.at[:].set(self._array_backend.fft.irfftn(arr, s=shape))
+
+        return rfftn, irfftn
+
+    def sharedarr_to_arr(self, shm, shape: Tuple[int], dtype: str):
+        return shm
+
+    @staticmethod
+    def arr_to_sharedarr(arr, shared_memory_handler: type = None):
+        return arr
+
+    def rotate_array(
+        self,
+        arr,
+        rotation_matrix,
+        arr_mask = None,
+        translation = None,
+        use_geometric_center: bool = False,
+        out = None,
+        out_mask = None,
+        order: int = 3,
+    ) -> None:
+        rotate_mask = arr_mask is not None
+        return_type = (out is None) + 2 * rotate_mask * (out_mask is None)
+        translation = self.zeros(arr.ndim) if translation is None else translation
+
+        indices = self._array_backend.indices(arr.shape).reshape(
+            (len(arr.shape), -1)
+        ).astype(self._float_dtype)
+
+        center = self.divide(arr.shape, 2)
+        if not use_geometric_center:
+            center = self.center_of_mass(arr, cutoff=0)
+        center = center[:, None]
+        indices = indices.at[:].add(-center)
+        rotation_matrix = self._array_backend.linalg.inv(rotation_matrix)
+        indices = self._array_backend.matmul(rotation_matrix, indices)
+        indices = indices.at[:].add(center)
+
+        out = self.zeros_like(arr) if out is None else out
+        out_slice = tuple(slice(0, stop) for stop in arr.shape)
+
+        out = out.at[out_slice].set(
+            self.scipy.ndimage.map_coordinates(
+                arr, indices, order=order
+            ).reshape(arr.shape)
+        )
+
+        if rotate_mask:
+            out_mask = self.zeros_like(arr_mask) if out_mask is None else out_mask
+            out_mask_slice = tuple(slice(0, stop) for stop in arr_mask.shape)
+            out_mask = out_mask.at[out_mask_slice].set(
+                self.scipy.ndimage.map_coordinates(
+                    arr_mask, indices, order=order
+                ).reshape(arr.shape)
+            )
+
+        match return_type:
+            case 0:
+                return None
+            case 1:
+                return out
+            case 2:
+                return out_mask
+            case 3:
+                return out, out_mask
+
+    def max_score_over_rotations(
+        self,
+        score_space,
+        internal_scores,
+        internal_rotations,
+        rotation_index: int,
+    ):
+        """
+        Modify internal_scores and internal_rotations inplace with scores and rotation
+        index respectively, wherever score_sapce is larger than internal scores.
+
+        Parameters
+        ----------
+        score_space : CupyArray
+            The score space to compare against internal_scores.
+        internal_scores : CupyArray
+            The internal scores to update with maximum scores.
+        internal_rotations : CupyArray
+            The internal rotations corresponding to the maximum scores.
+        rotation_index : int
+            The index representing the current rotation.
+        """
+        indices = score_space > internal_scores
+        internal_scores.at[indices].set(score_space[indices])
+        internal_rotations.at[indices].set(rotation_index)

tme/backends/matching_backend.py

@@ -29,32 +29,39 @@ class MatchingBackend(ABC):
     ----------
     array_backend : object
         The backend object providing array functionalities.
-    default_dtype : type
+    float_dtype : type
         Data type of real array instances, e.g. np.float32.
     complex_dtype : type
         Data type of complex array instances, e.g. np.complex64.
-    default_dtype_int : type
+    int_dtype : type
         Data type of integer array instances, e.g. np.int32.
+    overflow_safe_dtype : type
+        Data type than can be used for reduction operations to avoid overflows.
 
     Attributes
     ----------
     _array_backend : object
         The backend object used to delegate method and attribute calls.
-    _default_dtype : type
+    _float_dtype : type
         Data type of real array instances, e.g. np.float32.
     _complex_dtype : type
         Data type of complex array instances, e.g. np.complex64.
-    _default_dtype_int : type
+    _int_dtype : type
         Data type of integer array instances, e.g. np.int32.
+    _overflow_safe_dtype : type
+        Data type than can be used for reduction operations to avoid overflows.
+    _fundamental_dtypes : Dict
+        Mapping between fundamental int, float and complex python types to
+        array backend specific data types.
 
     Examples
     --------
     >>> import numpy as np
     >>> backend = MatchingBackend(
         array_backend = np,
-        default_dtype = np.float32,
+        float_dtype = np.float32,
         complex_dtype = np.complex64,
-        default_dtype_int = np.int32
+        int_dtype = np.int32
     )
     >>> arr = backend.array([1, 2, 3])
     >>> print(arr)
@@ -69,14 +76,22 @@ class MatchingBackend(ABC):
     def __init__(
         self,
         array_backend,
-        default_dtype: type,
+        float_dtype: type,
         complex_dtype: type,
-        default_dtype_int: type,
+        int_dtype: type,
+        overflow_safe_dtype: type,
     ):
         self._array_backend = array_backend
-        self._default_dtype = default_dtype
+        self._float_dtype = float_dtype
         self._complex_dtype = complex_dtype
-        self._default_dtype_int = default_dtype_int
+        self._int_dtype = int_dtype
+        self._overflow_safe_dtype = overflow_safe_dtype
+
+        self._fundamental_dtypes = {
+            int: self._int_dtype,
+            float: self._float_dtype,
+            complex: self._complex_dtype,
+        }
 
     def __getattr__(self, name: str):
         """

tme/backends/mlx_backend.py

@@ -1,4 +1,4 @@
-""" Backend Apple's MLX library for template matching.
+""" Backend using Apple's MLX library for template matching.
 
     Copyright (c) 2024 European Molecular Biology Laboratory
 
@@ -20,23 +20,27 @@ class MLXBackend(NumpyFFTWBackend):
     def __init__(
         self,
         device="cpu",
-        default_dtype=None,
+        float_dtype=None,
         complex_dtype=None,
-        default_dtype_int=None,
+        int_dtype=None,
+        overflow_safe_dtype=None,
         **kwargs,
     ):
         import mlx.core as mx
 
         device = mx.cpu if device == "cpu" else mx.gpu
-        default_dtype = mx.float32 if default_dtype is None else default_dtype
+        float_dtype = mx.float32 if float_dtype is None else float_dtype
         complex_dtype = mx.complex64 if complex_dtype is None else complex_dtype
-        default_dtype_int = mx.int32 if default_dtype_int is None else default_dtype_int
+        int_dtype = mx.int32 if int_dtype is None else int_dtype
+        if overflow_safe_dtype is None:
+            overflow_safe_dtype = mx.float32
 
         super().__init__(
             array_backend=mx,
-            default_dtype=default_dtype,
+            float_dtype=float_dtype,
             complex_dtype=complex_dtype,
-            default_dtype_int=default_dtype_int,
+            int_dtype=int_dtype,
+            overflow_safe_dtype=overflow_safe_dtype,
         )
 
         self.device = device
@@ -141,8 +145,8 @@ class MLXBackend(NumpyFFTWBackend):
         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._default_dtype_int)
-        stops = self.astype(self.add(starts, newshape), self._default_dtype_int)
+        starts = self.astype(self.divide(starts, 2), self._int_dtype)
+        stops = self.astype(self.add(starts, newshape), self._int_dtype)
         starts, stops = starts.tolist(), stops.tolist()
         box = tuple(slice(start, stop) for start, stop in zip(starts, stops))
         return arr[box]

tme/backends/npfftw_backend.py

@@ -29,16 +29,18 @@ class NumpyFFTWBackend(MatchingBackend):
     def __init__(
         self,
         array_backend=np,
-        default_dtype=np.float32,
+        float_dtype=np.float32,
         complex_dtype=np.complex64,
-        default_dtype_int=np.int32,
+        int_dtype=np.int32,
+        overflow_safe_dtype=np.float32,
         **kwargs,
     ):
         super().__init__(
             array_backend=array_backend,
-            default_dtype=default_dtype,
+            float_dtype=float_dtype,
             complex_dtype=complex_dtype,
-            default_dtype_int=default_dtype_int,
+            int_dtype=int_dtype,
+            overflow_safe_dtype=overflow_safe_dtype,
         )
         self.affine_transform = affine_transform
 
@@ -53,6 +55,16 @@ class NumpyFFTWBackend(MatchingBackend):
     def to_cpu_array(self, arr: NDArray) -> NDArray:
         return arr
 
+    def get_fundamental_dtype(self, arr):
+        dt = arr.dtype
+        if self._array_backend.issubdtype(dt, self._array_backend.integer):
+            return int
+        elif self._array_backend.issubdtype(dt, self._array_backend.floating):
+            return float
+        elif self._array_backend.issubdtype(dt, self._array_backend.complexfloating):
+            return complex
+        return float
+
     def free_cache(self):
         pass
 
@@ -429,8 +441,8 @@ class NumpyFFTWBackend(MatchingBackend):
         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._default_dtype_int)
-        stops = self.astype(self.add(starts, newshape), self._default_dtype_int)
+        starts = self.astype(self.divide(starts, 2), self._int_dtype)
+        stops = self.astype(self.add(starts, newshape), self._int_dtype)
         box = tuple(slice(start, stop) for start, stop in zip(starts, stops))
         return arr[box]
 
@@ -722,11 +734,11 @@ class NumpyFFTWBackend(MatchingBackend):
         arr = self._array_backend.where(arr > cutoff, arr, 0)
         denominator = self.sum(arr)
         grids = self._array_backend.ogrid[tuple(slice(0, i) for i in arr.shape)]
-        grids = [grid.astype(self._default_dtype) for grid in grids]
+        grids = [grid.astype(self._float_dtype) for grid in grids]
 
         center_of_mass = self.array(
             [
-                self.sum(self.multiply(arr, grids[dim])) / denominator
+                self.sum(self.multiply(arr, grids[dim]) / denominator)
                 for dim in range(arr.ndim)
             ]
         )

tme/backends/pytorch_backend.py

@@ -24,25 +24,27 @@ class PytorchBackend(NumpyFFTWBackend):
     def __init__(
         self,
         device="cuda",
-        default_dtype=None,
+        float_dtype=None,
         complex_dtype=None,
-        default_dtype_int=None,
+        int_dtype=None,
+        overflow_safe_dtype=None,
         **kwargs,
     ):
         import torch
         import torch.nn.functional as F
 
-        default_dtype = torch.float32 if default_dtype is None else default_dtype
+        float_dtype = torch.float32 if float_dtype is None else float_dtype
         complex_dtype = torch.complex64 if complex_dtype is None else complex_dtype
-        default_dtype_int = (
-            torch.int32 if default_dtype_int is None else default_dtype_int
-        )
+        int_dtype = torch.int32 if int_dtype is None else int_dtype
+        if overflow_safe_dtype is None:
+            overflow_safe_dtype = torch.float32
 
         super().__init__(
             array_backend=torch,
-            default_dtype=default_dtype,
+            float_dtype=float_dtype,
             complex_dtype=complex_dtype,
-            default_dtype_int=default_dtype_int,
+            int_dtype=int_dtype,
+            overflow_safe_dtype=overflow_safe_dtype,
         )
         self.device = device
         self.F = F
@@ -57,11 +59,20 @@ class PytorchBackend(NumpyFFTWBackend):
     def to_numpy_array(self, arr: TorchTensor) -> NDArray:
         if isinstance(arr, np.ndarray):
             return arr
-        return arr.cpu().numpy()
+        elif isinstance(arr, self._array_backend.Tensor):
+            return arr.cpu().numpy()
+        return np.array(arr)
 
     def to_cpu_array(self, arr: TorchTensor) -> NDArray:
         return arr.cpu()
 
+    def get_fundamental_dtype(self, arr):
+        if self._array_backend.is_floating_point(arr):
+            return float
+        elif self._array_backend.is_complex(arr):
+            return complex
+        return int
+
     def free_cache(self):
         self._array_backend.cuda.empty_cache()