pytme 0.2.0b0__cp311-cp311-macosx_14_0_arm64.whl → 0.2.2__cp311-cp311-macosx_14_0_arm64.whl

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
 
@@ -9,34 +9,38 @@ from typing import Tuple, List, Callable
 import numpy as np
 
 from .npfftw_backend import NumpyFFTWBackend
-from ..types import NDArray, MlxArray, Scalar
+from ..types import NDArray, MlxArray, Scalar, shm_type
 
 
 class MLXBackend(NumpyFFTWBackend):
     """
-    A MLX based backend for template matching.
+    A mlx-based matching backend.
     """
 
     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
@@ -68,6 +72,15 @@ class MLXBackend(NumpyFFTWBackend):
             return None
         return self._array_backend.add(x1, x2, **kwargs)
 
+    def multiply(self, x1, x2, out: MlxArray = None, **kwargs) -> MlxArray:
+        x1 = self.to_backend_array(x1)
+        x2 = self.to_backend_array(x2)
+
+        if out is not None:
+            out[:] = self._array_backend.multiply(x1, x2, **kwargs)
+            return None
+        return self._array_backend.multiply(x1, x2, **kwargs)
+
     def std(self, arr: MlxArray, axis) -> Scalar:
         return self._array_backend.sqrt(arr.var(axis=axis))
 
@@ -80,30 +93,12 @@ class MLXBackend(NumpyFFTWBackend):
     def tobytes(self, arr):
         return self.to_numpy_array(arr).tobytes()
 
-    def preallocate_array(self, shape: Tuple[int], dtype: type = None) -> NDArray:
-        """
-        Returns a byte-aligned array of zeros with specified shape and dtype.
-
-        Parameters
-        ----------
-        shape : Tuple[int]
-            Desired shape for the array.
-        dtype : type, optional
-            Desired data type for the array.
-
-        Returns
-        -------
-        NDArray
-            Byte-aligned array of zeros with specified shape and dtype.
-        """
-        arr = self._array_backend.zeros(shape, dtype=dtype)
-        return arr
-
     def full(self, shape, fill_value, dtype=None):
         return self._array_backend.full(shape=shape, dtype=dtype, vals=fill_value)
 
-    def fill(self, arr: MlxArray, value: Scalar) -> None:
+    def fill(self, arr: MlxArray, value: Scalar) -> MlxArray:
         arr[:] = value
+        return arr
 
     def zeros(self, shape: Tuple[int], dtype: type = None) -> MlxArray:
         return self._array_backend.zeros(shape=shape, dtype=dtype)
@@ -141,8 +136,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]
@@ -185,13 +180,11 @@ class MLXBackend(NumpyFFTWBackend):
 
         return rfftn, irfftn
 
-    def sharedarr_to_arr(
-        self, shape: Tuple[int], dtype: str, shm: MlxArray
-    ) -> MlxArray:
-        return shm
+    def from_sharedarr(self, arr: MlxArray) -> MlxArray:
+        return arr
 
     @staticmethod
-    def arr_to_sharedarr(arr: MlxArray, shared_memory_handler: type = None) -> MlxArray:
+    def to_sharedarr(arr: MlxArray, shared_memory_handler: type = None) -> shm_type:
         return arr
 
     def topk_indices(self, arr: NDArray, k: int):
@@ -200,7 +193,7 @@ class MLXBackend(NumpyFFTWBackend):
         ret = [self.to_backend_array(x) for x in ret]
         return ret
 
-    def rotate_array(
+    def rigid_transform(
         self,
         arr: NDArray,
         rotation_matrix: NDArray,
@@ -210,10 +203,8 @@ class MLXBackend(NumpyFFTWBackend):
         out: NDArray = None,
         out_mask: NDArray = None,
         order: int = 3,
+        **kwargs,
     ) -> None:
-        rotate_mask = arr_mask is not None
-        return_type = (out is None) + 2 * rotate_mask * (out_mask is None)
-
         arr = self.to_numpy_array(arr)
         rotation_matrix = self.to_numpy_array(rotation_matrix)
 
@@ -223,46 +214,32 @@ class MLXBackend(NumpyFFTWBackend):
         if translation is not None:
             translation = self.to_numpy_array(translation)
 
-        out_pass, out_mask_pass = None, None
-        if out is not None:
-            out_pass = self.to_numpy_array(out)
-        if out_mask is not None:
-            out_mask_pass = self.to_numpy_array(out_mask)
+        if out is None:
+            out = self.zeros(arr.shape)
+        if out_mask is None and arr_mask is not None:
+            out_mask_pass = self.zeros(arr_mask.shape)
 
-        ret = NumpyFFTWBackend().rotate_array(
+        ret = NumpyFFTWBackend().rigid_transform(
             arr=arr,
             rotation_matrix=rotation_matrix,
             arr_mask=arr_mask,
             translation=translation,
             use_geometric_center=use_geometric_center,
-            out=out_pass,
-            out_mask=out_mask_pass,
             order=order,
         )
 
-        if ret is not None:
-            if len(ret) == 1 and out is None:
-                out_pass = ret
-            elif len(ret) == 1 and out_mask is None:
-                out_mask_pass = ret
-            else:
-                out_pass, out_mask_pass = ret
+        out_pass, out_mask_pass = ret
+        out[:] = self.to_backend_array(out_pass)
 
-        if out is not None:
-            out[:] = self.to_backend_array(out_pass)
+        if out_mask_pass is not None:
+            out_mask_pass = self.to_backend_array(out_mask_pass)
 
         if out_mask is not None:
-            out_mask[:] = self.to_backend_array(out_mask_pass)
-
-        match return_type:
-            case 0:
-                return None
-            case 1:
-                return out
-            case 2:
-                return out_mask
-            case 3:
-                return out, out_mask
+            out_mask[:] = out_mask_pass
+        else:
+            out_mask = out_mask_pass
+
+        return out, out_mask
 
     def indices(self, arr: List) -> MlxArray:
         ret = NumpyFFTWBackend().indices(arr)