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

tme/backends/jax_backend.py

@@ -1,17 +1,38 @@
 """ Backend using jax for template matching.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+    Copyright (c) 2023-2024 European Molecular Biology Laboratory
 
     Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
-from typing import Tuple, Callable
+from functools import wraps
+from typing import Tuple, List, Callable
+
+from ..types import BackendArray
+from .npfftw_backend import NumpyFFTWBackend, shm_type
+
+
+def emulate_out(func):
+    """
+    Adds an out argument to write output of ``func`` to.
+    """
+
+    @wraps(func)
+    def inner(*args, out=None, **kwargs):
+        ret = func(*args, **kwargs)
+        if out is not None:
+            out = out.at[:].set(ret)
+            return out
+        return ret
+
+    return inner
 
-from .npfftw_backend import NumpyFFTWBackend
 
 class JaxBackend(NumpyFFTWBackend):
-    def __init__(
-        self, float_dtype=None, complex_dtype=None, int_dtype=None, **kwargs
-    ):
+    """
+    A jax-based matching backend.
+    """
+
+    def __init__(self, float_dtype=None, complex_dtype=None, int_dtype=None, **kwargs):
         import jax.scipy as jsp
         import jax.numpy as jnp
 
@@ -19,24 +40,33 @@ class JaxBackend(NumpyFFTWBackend):
         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,
+            overflow_safe_dtype=float_dtype,
         )
+        self.scipy = jsp
+        self._create_ufuncs()
+        try:
+            from ._jax_utils import scan as _
+
+            self.scan = self._scan
+        except Exception:
+            pass
 
-    def to_backend_array(self, arr):
-        return self._array_backend.asarray(arr)
+    def from_sharedarr(self, arr: BackendArray) -> BackendArray:
+        return arr
 
-    def preallocate_array(self, shape: Tuple[int], dtype: type):
-        arr = self._array_backend.zeros(shape, dtype=dtype)
+    @staticmethod
+    def to_sharedarr(arr: BackendArray, shared_memory_handler: type = None) -> shm_type:
         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)
+    def topleft_pad(
+        self, arr: BackendArray, shape: Tuple[int], padval: int = 0
+    ) -> BackendArray:
+        b = self.full(shape=shape, dtype=arr.dtype, fill_value=padval)
         aind = [slice(None, None)] * arr.ndim
         bind = [slice(None, None)] * arr.ndim
         for i in range(arr.ndim):
@@ -47,43 +77,29 @@ class JaxBackend(NumpyFFTWBackend):
         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 _create_ufuncs(self):
+        ufuncs = [
+            "add",
+            "subtract",
+            "multiply",
+            "divide",
+            "square",
+            "sqrt",
+            "maximum",
+        ]
+        for ufunc in ufuncs:
+            backend_method = emulate_out(getattr(self._array_backend, ufunc))
+            setattr(self, ufunc, staticmethod(backend_method))
+
+        ufuncs = ["zeros", "full"]
+        for ufunc in ufuncs:
+            backend_method = getattr(self._array_backend, ufunc)
+            setattr(self, ufunc, staticmethod(backend_method))
+
+    def fill(self, arr: BackendArray, value: float) -> BackendArray:
+        return self._array_backend.full(
+            shape=arr.shape, dtype=arr.dtype, fill_value=value
+        )
 
     def build_fft(
         self,
@@ -92,127 +108,189 @@ class JaxBackend(NumpyFFTWBackend):
         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 rfftn(arr, out, shape=fast_shape):
+            return 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))
+        def irfftn(arr, out, shape=fast_shape):
+            return 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(
+    def rigid_transform(
         self,
-        arr,
-        rotation_matrix,
-        arr_mask = None,
-        translation = None,
-        use_geometric_center: bool = False,
-        out = None,
-        out_mask = None,
-        order: int = 3,
-    ) -> None:
+        arr: BackendArray,
+        rotation_matrix: BackendArray,
+        out: BackendArray = None,
+        out_mask: BackendArray = None,
+        translation: BackendArray = None,
+        arr_mask: BackendArray = None,
+        order: int = 1,
+        **kwargs,
+    ) -> Tuple[BackendArray, BackendArray]:
         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]
+        center = self.divide(self.to_backend_array(arr.shape) - 1, 2)[:, None]
+        indices = self._array_backend.indices(arr.shape, dtype=self._float_dtype)
+        indices = indices.reshape((arr.ndim, -1))
         indices = indices.at[:].add(-center)
-        rotation_matrix = self._array_backend.linalg.inv(rotation_matrix)
-        indices = self._array_backend.matmul(rotation_matrix, indices)
+        indices = self._array_backend.matmul(rotation_matrix.T, indices)
         indices = indices.at[:].add(center)
+        if translation is not None:
+            indices = indices.at[:].add(translation)
 
-        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)
+        out = self.scipy.ndimage.map_coordinates(arr, indices, order=order).reshape(
+            arr.shape
         )
 
+        out_mask = arr_mask
         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)
-            )
+            out_mask = self.scipy.ndimage.map_coordinates(
+                arr_mask, indices, order=order
+            ).reshape(arr_mask.shape)
 
-        match return_type:
-            case 0:
-                return None
-            case 1:
-                return out
-            case 2:
-                return out_mask
-            case 3:
-                return out, out_mask
+        return out, out_mask
 
     def max_score_over_rotations(
         self,
-        score_space,
-        internal_scores,
-        internal_rotations,
+        scores: BackendArray,
+        max_scores: BackendArray,
+        rotations: BackendArray,
         rotation_index: int,
-    ):
+    ) -> Tuple[BackendArray, BackendArray]:
+        update = self.greater(max_scores, scores)
+        max_scores = max_scores.at[:].set(self.where(update, max_scores, scores))
+        rotations = rotations.at[:].set(self.where(update, rotations, rotation_index))
+        return max_scores, rotations
+
+    def _scan(
+        self,
+        matching_data: type,
+        splits: Tuple[Tuple[slice, slice]],
+        n_jobs: int,
+        callback_class,
+        rotate_mask: bool = False,
+        **kwargs,
+    ) -> List:
         """
-        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.
+        Emulates output of :py:meth:`tme.matching_exhaustive.scan` using
+        :py:class:`tme.analyzer.MaxScoreOverRotations`.
         """
-        indices = score_space > internal_scores
-        internal_scores.at[indices].set(score_space[indices])
-        internal_rotations.at[indices].set(rotation_index)
+        from ._jax_utils import scan as scan_inner
+
+        pad_target = True if len(splits) > 1 else False
+        convolution_mode = "valid" if pad_target else "same"
+        target_pad = matching_data.target_padding(pad_target=pad_target)
+
+        target_shape = tuple(
+            (x.stop - x.start + p) for x, p in zip(splits[0][0], target_pad)
+        )
+        conv_shape, fast_shape, fast_ft_shape, shift = matching_data._fourier_padding(
+            target_shape=self.to_numpy_array(target_shape),
+            template_shape=self.to_numpy_array(matching_data._template.shape),
+            pad_fourier=False,
+        )
+
+        analyzer_args = {
+            "convolution_mode": convolution_mode,
+            "fourier_shift": shift,
+            "targetshape": target_shape,
+            "templateshape": matching_data.template.shape,
+            "convolution_shape": conv_shape,
+        }
+
+        create_target_filter = matching_data.target_filter is not None
+        create_template_filter = matching_data.template_filter is not None
+        create_filter = create_target_filter or create_template_filter
+
+        # Applying the filter leads to more FFTs
+        fastt_shape = matching_data._template.shape
+        if create_template_filter:
+            _, fastt_shape, _, tshift = matching_data._fourier_padding(
+                target_shape=self.to_numpy_array(matching_data._template.shape),
+                template_shape=self.to_numpy_array(
+                    [1 for _ in matching_data._template.shape]
+                ),
+                pad_fourier=False,
+            )
+
+        ret, template_filter, target_filter = [], 1, 1
+        rotation_mapping = {
+            self.tobytes(matching_data.rotations[i]): i
+            for i in range(matching_data.rotations.shape[0])
+        }
+        for split_start in range(0, len(splits), n_jobs):
+            split_subset = splits[split_start : (split_start + n_jobs)]
+            if not len(split_subset):
+                continue
+
+            targets, translation_offsets = [], []
+            for target_split, template_split in split_subset:
+                base = matching_data.subset_by_slice(
+                    target_slice=target_split,
+                    target_pad=target_pad,
+                    template_slice=template_split,
+                )
+                translation_offsets.append(base._translation_offset)
+                targets.append(self.topleft_pad(base._target, fast_shape))
+
+            if create_filter:
+                filter_args = {
+                    "data_rfft": self.fft.rfftn(targets[0]),
+                    "return_real_fourier": True,
+                    "shape_is_real_fourier": False,
+                }
+
+            if create_template_filter:
+                template_filter = matching_data.template_filter(
+                    shape=fastt_shape, **filter_args
+                )["data"]
+                template_filter = template_filter.at[(0,) * template_filter.ndim].set(0)
+
+            if create_target_filter:
+                target_filter = matching_data.target_filter(
+                    shape=fast_shape, **filter_args
+                )["data"]
+                target_filter = target_filter.at[(0,) * target_filter.ndim].set(0)
+
+            create_filter, create_template_filter, create_target_filter = (False,) * 3
+            base, targets = None, self._array_backend.stack(targets)
+            scores, rotations = scan_inner(
+                targets,
+                self.topleft_pad(matching_data.template, fastt_shape),
+                self.topleft_pad(matching_data.template_mask, fastt_shape),
+                matching_data.rotations,
+                template_filter,
+                target_filter,
+                fast_shape,
+                rotate_mask,
+            )
+
+            for index in range(scores.shape[0]):
+                temp = callback_class(
+                    scores=scores[index],
+                    rotations=rotations[index],
+                    thread_safe=False,
+                    offset=translation_offsets[index],
+                )
+                temp.rotation_mapping = rotation_mapping
+                ret.append(tuple(temp._postprocess(**analyzer_args)))
+
+        return ret
+
+    def get_available_memory(self) -> int:
+        import jax
+
+        _memory = {"cpu": 0, "gpu": 0}
+        for device in jax.devices():
+            if device.platform == "cpu":
+                _memory["cpu"] = super().get_available_memory()
+            else:
+                mem_stats = device.memory_stats()
+                _memory["gpu"] += mem_stats.get("bytes_limit", 0)
+
+        if _memory["gpu"] > 0:
+            return _memory["gpu"]
+        return _memory["cpu"]