pytme 0.1.5__cp311-cp311-macosx_14_0_arm64.whl

tme/backends/npfftw_backend.py

@@ -0,0 +1,763 @@
+""" Backend using numpy and pyFFTW for template matching.
+
+    Copyright (c) 2023 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+
+from typing import Tuple, Dict, List
+from multiprocessing import shared_memory
+from multiprocessing.managers import SharedMemoryManager
+from contextlib import contextmanager
+
+import numpy as np
+from psutil import virtual_memory
+from numpy.typing import NDArray
+from pyfftw import zeros_aligned, simd_alignment, FFTW, next_fast_len
+from pyfftw.builders import rfftn as rfftn_builder, irfftn as irfftn_builder
+from scipy.ndimage import maximum_filter, affine_transform
+
+from .matching_backend import MatchingBackend
+from ..matching_utils import rigid_transform
+
+
+class NumpyFFTWBackend(MatchingBackend):
+    """
+    A numpy and pyfftw based backend for template matching.
+    """
+
+    def __init__(
+        self,
+        array_backend=np,
+        default_dtype=np.float32,
+        complex_dtype=np.complex64,
+        default_dtype_int=np.int32,
+        **kwargs,
+    ):
+        super().__init__(
+            array_backend=array_backend,
+            default_dtype=default_dtype,
+            complex_dtype=complex_dtype,
+            default_dtype_int=np.int32,
+        )
+        self.affine_transform = affine_transform
+
+    def to_backend_array(self, arr: NDArray) -> NDArray:
+        if isinstance(arr, self._array_backend.ndarray):
+            return arr
+        return self._array_backend.asarray(arr)
+
+    def to_numpy_array(self, arr: NDArray) -> NDArray:
+        return arr
+
+    def to_cpu_array(self, arr: NDArray) -> NDArray:
+        return arr
+
+    def free_cache(self):
+        pass
+
+    def add(self, x1, x2, *args, **kwargs) -> NDArray:
+        x1 = self.to_backend_array(x1)
+        x2 = self.to_backend_array(x2)
+        return self._array_backend.add(x1, x2, *args, **kwargs)
+
+    def subtract(self, x1, x2, *args, **kwargs) -> NDArray:
+        x1 = self.to_backend_array(x1)
+        x2 = self.to_backend_array(x2)
+        return self._array_backend.subtract(x1, x2, *args, **kwargs)
+
+    def multiply(self, x1, x2, *args, **kwargs) -> NDArray:
+        x1 = self.to_backend_array(x1)
+        x2 = self.to_backend_array(x2)
+        return self._array_backend.multiply(x1, x2, *args, **kwargs)
+
+    def divide(self, x1, x2, *args, **kwargs) -> NDArray:
+        x1 = self.to_backend_array(x1)
+        x2 = self.to_backend_array(x2)
+        return self._array_backend.divide(x1, x2, *args, **kwargs)
+
+    def mod(self, x1, x2, *args, **kwargs):
+        x1 = self.to_backend_array(x1)
+        x2 = self.to_backend_array(x2)
+        return self._array_backend.mod(x1, x2, *args, **kwargs)
+
+    def sum(self, *args, **kwargs) -> NDArray:
+        return self._array_backend.sum(*args, **kwargs)
+
+    def einsum(self, *args, **kwargs) -> NDArray:
+        return self._array_backend.einsum(*args, **kwargs)
+
+    def mean(self, *args, **kwargs) -> NDArray:
+        return self._array_backend.mean(*args, **kwargs)
+
+    def std(self, *args, **kwargs) -> NDArray:
+        return self._array_backend.std(*args, **kwargs)
+
+    def max(self, *args, **kwargs) -> NDArray:
+        return self._array_backend.max(*args, **kwargs)
+
+    def min(self, *args, **kwargs) -> NDArray:
+        return self._array_backend.min(*args, **kwargs)
+
+    def maximum(self, x1, x2, *args, **kwargs) -> NDArray:
+        x1 = self.to_backend_array(x1)
+        x2 = self.to_backend_array(x2)
+        return self._array_backend.maximum(x1, x2, *args, **kwargs)
+
+    def minimum(self, x1, x2, *args, **kwargs) -> NDArray:
+        x1 = self.to_backend_array(x1)
+        x2 = self.to_backend_array(x2)
+        return self._array_backend.minimum(x1, x2, *args, **kwargs)
+
+    def sqrt(self, *args, **kwargs) -> NDArray:
+        return self._array_backend.sqrt(*args, **kwargs)
+
+    def square(self, *args, **kwargs) -> NDArray:
+        return self._array_backend.square(*args, **kwargs)
+
+    def abs(self, *args, **kwargs) -> NDArray:
+        return self._array_backend.abs(*args, **kwargs)
+
+    def transpose(self, arr):
+        return arr.T
+
+    def power(self, *args, **kwargs):
+        return self._array_backend.power(*args, **kwargs)
+
+    def tobytes(self, arr):
+        return arr.tobytes()
+
+    def size(self, arr):
+        return arr.size
+
+    def fill(self, arr: NDArray, value: float) -> None:
+        arr.fill(value)
+
+    def zeros(self, shape, dtype=np.float64) -> NDArray:
+        return self._array_backend.zeros(shape=shape, dtype=dtype)
+
+    def full(self, shape, fill_value, dtype=None, **kwargs) -> NDArray:
+        return self._array_backend.full(
+            shape, dtype=dtype, fill_value=fill_value, **kwargs
+        )
+
+    def eps(self, dtype: type) -> NDArray:
+        """
+        Returns the eps defined as diffeerence between 1.0 and the next
+        representable floating point value larger than 1.0.
+
+        Parameters
+        ----------
+        dtype : type
+            Data type for which eps should be returned.
+
+        Returns
+        -------
+        Scalar
+            The eps for the given data type
+        """
+        return self._array_backend.finfo(dtype).eps
+
+    def datatype_bytes(self, dtype: type) -> NDArray:
+        """
+        Return the number of bytes occupied by a given datatype.
+
+        Parameters
+        ----------
+        dtype : type
+            Datatype for which the number of bytes is to be determined.
+
+        Returns
+        -------
+        int
+            Number of bytes occupied by the datatype.
+        """
+        temp = self._array_backend.zeros(1, dtype=dtype)
+        return temp.nbytes
+
+    def clip(self, *args, **kwargs) -> NDArray:
+        return self._array_backend.clip(*args, **kwargs)
+
+    def flip(self, a, axis, **kwargs):
+        return self._array_backend.flip(a, axis, **kwargs)
+
+    @staticmethod
+    def astype(arr, dtype):
+        return arr.astype(dtype)
+
+    def arange(self, *args, **kwargs):
+        return self._array_backend.arange(*args, **kwargs)
+
+    def stack(self, *args, **kwargs):
+        return self._array_backend.stack(*args, **kwargs)
+
+    def concatenate(self, *args, **kwargs):
+        return self._array_backend.concatenate(*args, **kwargs)
+
+    def repeat(self, *args, **kwargs):
+        return self._array_backend.repeat(*args, **kwargs)
+
+    def topk_indices(self, arr: NDArray, k: int):
+        temp = arr.reshape(-1)
+        indices = self._array_backend.argpartition(temp, -k)[-k:][:k]
+        sorted_indices = indices[self._array_backend.argsort(temp[indices])][::-1]
+        sorted_indices = self.unravel_index(indices=sorted_indices, shape=arr.shape)
+        return sorted_indices
+
+    def indices(self, *args, **kwargs) -> NDArray:
+        return self._array_backend.indices(*args, **kwargs)
+
+    def roll(self, a, shift, axis, **kwargs):
+        return self._array_backend.roll(
+            a,
+            shift=shift,
+            axis=axis,
+            **kwargs,
+        )
+
+    def unique(self, *args, **kwargs):
+        return self._array_backend.unique(*args, **kwargs)
+
+    def argsort(self, *args, **kwargs):
+        return self._array_backend.argsort(*args, **kwargs)
+
+    def unravel_index(self, indices, shape):
+        return self._array_backend.unravel_index(indices=indices, shape=shape)
+
+    def tril_indices(self, *args, **kwargs):
+        return self._array_backend.tril_indices(*args, **kwargs)
+
+    def max_filter_coordinates(self, score_space, min_distance: Tuple[int]):
+        score_box = tuple(min_distance for _ in range(score_space.ndim))
+        max_filter = maximum_filter(score_space, size=score_box, mode="constant")
+        max_filter = max_filter == score_space
+
+        peaks = np.array(np.nonzero(max_filter)).T
+        return peaks
+
+    @staticmethod
+    def preallocate_array(shape: Tuple[int], dtype: type) -> NDArray:
+        """
+        Returns a byte-aligned array of zeros with specified shape and dtype.
+
+        Parameters
+        ----------
+        shape : Tuple[int]
+            Desired shape for the array.
+        dtype : type
+            Desired data type for the array.
+
+        Returns
+        -------
+        NDArray
+            Byte-aligned array of zeros with specified shape and dtype.
+        """
+        arr = zeros_aligned(shape, dtype=dtype, n=simd_alignment)
+        return arr
+
+    def sharedarr_to_arr(
+        self, shape: Tuple[int], dtype: str, shm: shared_memory.SharedMemory
+    ) -> NDArray:
+        """
+        Returns an array of given shape and dtype from shared memory location.
+
+        Parameters
+        ----------
+        shape : tuple
+            Tuple of integers specifying the shape of the array.
+        dtype : str
+            String specifying the dtype of the array.
+        shm : shared_memory.SharedMemory
+            Shared memory object where the array is stored.
+
+        Returns
+        -------
+        NDArray
+            Array of the specified shape and dtype from the shared memory location.
+        """
+        return self.ndarray(shape, dtype, shm.buf)
+
+    def arr_to_sharedarr(
+        self, arr: NDArray, shared_memory_handler: type = None
+    ) -> shared_memory.SharedMemory:
+        """
+        Converts a numpy array to an object shared in memory.
+
+        Parameters
+        ----------
+        arr : NDArray
+            Numpy array to convert.
+        shared_memory_handler : type, optional
+            The type of shared memory handler. Default is None.
+
+        Returns
+        -------
+        shared_memory.SharedMemory
+            The shared memory object containing the numpy array.
+        """
+        if type(shared_memory_handler) == SharedMemoryManager:
+            shm = shared_memory_handler.SharedMemory(size=arr.nbytes)
+        else:
+            shm = shared_memory.SharedMemory(create=True, size=arr.nbytes)
+        np_array = self.ndarray(arr.shape, dtype=arr.dtype, buffer=shm.buf)
+        np_array[:] = arr[:].copy()
+        return shm
+
+    def topleft_pad(self, arr: NDArray, shape: Tuple[int], padval: int = 0) -> NDArray:
+        """
+        Returns an array that has been padded to a specified shape with a padding
+        value at the top-left corner.
+
+        Parameters
+        ----------
+        arr : NDArray
+            Input array to be padded.
+        shape : Tuple[int]
+            Desired shape for the output array.
+        padval : int, optional
+            Value to use for padding, default is 0.
+
+        Returns
+        -------
+        NDArray
+            Array that has been padded to the specified shape.
+        """
+        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[tuple(bind)] = arr[tuple(aind)]
+        return b
+
+    def build_fft(
+        self,
+        fast_shape: Tuple[int],
+        fast_ft_shape: Tuple[int],
+        real_dtype: type,
+        complex_dtype: type,
+        fftargs: Dict = {},
+        temp_real: NDArray = None,
+        temp_fft: NDArray = None,
+    ) -> Tuple[FFTW, FFTW]:
+        """
+        Build pyFFTW builder functions.
+
+        Parameters
+        ----------
+        fast_shape : tuple
+            Tuple of integers corresponding to fast convolution shape
+            (see `compute_convolution_shapes`).
+        fast_ft_shape : tuple
+            Tuple of integers corresponding to the shape of the fourier
+            transform array (see `compute_convolution_shapes`).
+        real_dtype : dtype
+            Numpy dtype of the inverse fourier transform.
+        complex_dtype : dtype
+            Numpy dtype of the fourier transform.
+        fftargs : dict, optional
+            Dictionary passed to pyFFTW builders.
+        temp_real : NDArray, optional
+            Temporary real numpy array, by default None.
+        temp_fft : NDArray, optional
+            Temporary fft numpy array, by default None.
+
+        Returns
+        -------
+        tuple
+            Tuple containing callable pyFFTW objects for forward and inverse
+            fourier transform.
+        """
+        if temp_real is None:
+            temp_real = self.preallocate_array(fast_shape, real_dtype)
+        if temp_fft is None:
+            temp_fft = self.preallocate_array(fast_ft_shape, complex_dtype)
+
+        default_values = {
+            "planner_effort": "FFTW_MEASURE",
+            "auto_align_input": False,
+            "auto_contiguous": False,
+            "avoid_copy": True,
+            "overwrite_input": True,
+            "threads": 1,
+        }
+        for key in default_values:
+            if key in fftargs:
+                continue
+            fftargs[key] = default_values[key]
+
+        rfftn = rfftn_builder(temp_real, s=fast_shape, **fftargs)
+
+        overwrite_input = None
+        if "overwrite_input" in fftargs:
+            overwrite_input = fftargs.pop("overwrite_input")
+        irfftn = irfftn_builder(temp_fft, s=fast_shape, **fftargs)
+
+        if overwrite_input is not None:
+            fftargs["overwrite_input"] = overwrite_input
+        return rfftn, irfftn
+
+    def extract_center(self, arr: NDArray, newshape: Tuple[int]) -> NDArray:
+        """
+        Extract the centered portion of an array based on a new shape.
+
+        Parameters
+        ----------
+        arr : NDArray
+            Input array.
+        newshape : tuple
+            Desired shape for the central portion.
+
+        Returns
+        -------
+        NDArray
+            Central portion of the array with shape `newshape`.
+
+        References
+        ----------
+        .. [1] https://github.com/scipy/scipy/blob/v1.11.2/scipy/signal/_signaltools.py
+        """
+        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), int)
+        stops = self.add(starts, newshape)
+        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]]:
+        """
+        Computes regular, optimized and fourier convolution shape.
+
+        Parameters
+        ----------
+        arr1_shape : tuple
+            Tuple of integers corresponding to array1 shape.
+        arr2_shape : tuple
+            Tuple of integers corresponding to array2 shape.
+
+        Returns
+        -------
+        tuple
+            Tuple with regular convolution shape, convolution shape optimized for faster
+            fourier transform, shape of the forward fourier transform
+            (see :py:meth:`build_fft`).
+        """
+        convolution_shape = [
+            int(x) + int(y) - 1 for x, y in zip(arr1_shape, arr2_shape)
+        ]
+        fast_shape = [next_fast_len(x) for x in convolution_shape]
+        fast_ft_shape = list(fast_shape[:-1]) + [fast_shape[-1] // 2 + 1]
+
+        return convolution_shape, fast_shape, fast_ft_shape
+
+    def rotate_array(
+        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,
+    ) -> None:
+        """
+        Rotates coordinates of arr according to rotation_matrix.
+
+        If no output array is provided, this method will compute an array with
+        sufficient space to hold all elements. If both `arr` and `arr_mask`
+        are provided, `arr_mask` will be centered according to arr.
+
+        Parameters
+        ----------
+        arr : NDArray
+            The input array to be rotated.
+        arr_mask : NDArray, optional
+            The mask of `arr` that will be equivalently rotated.
+        rotation_matrix : NDArray
+            The rotation matrix to apply [d x d].
+        translation : NDArray
+            The translation to apply [d].
+        use_geometric_center : bool, optional
+            Whether the rotation should be centered around the geometric
+            or mass center. Default is mass center.
+        out : NDArray, optional
+            The output array to write the rotation of `arr` to.
+        out_mask : NDArray, optional
+            The output array to write the rotation of `arr_mask` to.
+        order : int, optional
+            Spline interpolation order. Has to be in the range 0-5.
+        """
+
+        if order is None:
+            mask_coordinates = None
+            if arr_mask is not None:
+                mask_coordinates = np.array(np.where(arr_mask > 0))
+            return self.rotate_array_coordinates(
+                arr=arr,
+                arr_mask=arr_mask,
+                coordinates=np.array(np.where(arr > 0)),
+                mask_coordinates=mask_coordinates,
+                out=out,
+                out_mask=out_mask,
+                rotation_matrix=rotation_matrix,
+                translation=translation,
+                use_geometric_center=use_geometric_center,
+            )
+
+        rotate_mask = arr_mask is not None
+        return_type = (out is None) + 2 * rotate_mask * (out_mask is None)
+        translation = np.zeros(arr.ndim) if translation is None else translation
+
+        center = np.divide(arr.shape, 2)
+        if not use_geometric_center:
+            center = self.center_of_mass(arr, cutoff=0)
+
+        rotation_matrix_inverted = np.linalg.inv(rotation_matrix)
+        transformed_center = rotation_matrix_inverted @ center.reshape(-1, 1)
+        transformed_center = transformed_center.reshape(-1)
+        base_offset = np.subtract(center, transformed_center)
+        offset = np.subtract(base_offset, translation)
+
+        out = np.zeros_like(arr) if out is None else out
+        out_slice = tuple(slice(0, stop) for stop in arr.shape)
+
+        # Applying the prefilter can cause artifacts in the mask
+        affine_transform(
+            input=arr,
+            matrix=rotation_matrix_inverted,
+            offset=offset,
+            mode="constant",
+            output=out[out_slice],
+            order=order,
+            prefilter=True,
+        )
+
+        if rotate_mask:
+            out_mask = np.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)
+            affine_transform(
+                input=arr_mask,
+                matrix=rotation_matrix_inverted,
+                offset=offset,
+                mode="constant",
+                output=out_mask[out_mask_slice],
+                order=order,
+                prefilter=False,
+            )
+
+        match return_type:
+            case 0:
+                return None
+            case 1:
+                return out
+            case 2:
+                return out_mask
+            case 3:
+                return out, out_mask
+
+    @staticmethod
+    def rotate_array_coordinates(
+        arr: NDArray,
+        coordinates: NDArray,
+        rotation_matrix: NDArray,
+        translation: NDArray = None,
+        out: NDArray = None,
+        use_geometric_center: bool = True,
+        arr_mask: NDArray = None,
+        mask_coordinates: NDArray = None,
+        out_mask: NDArray = None,
+    ) -> None:
+        """
+        Rotates coordinates of arr according to rotation_matrix.
+
+        If no output array is provided, this method will compute an array with
+        sufficient space to hold all elements. If both `arr` and `arr_mask`
+        are provided, `arr_mask` will be centered according to arr.
+
+        No centering will be performed if the rotation matrix is the identity matrix.
+
+        Parameters
+        ----------
+        arr : NDArray
+            The input array to be rotated.
+        coordinates : NDArray
+            The pointcloud [d x N] containing elements of `arr` that should be rotated.
+            See :py:meth:`Density.to_pointcloud` on how to obtain the coordinates.
+        rotation_matrix : NDArray
+            The rotation matrix to apply [d x d].
+        rotation_matrix : NDArray
+            The translation to apply [d].
+        out : NDArray, optional
+            The output array to write the rotation of `arr` to.
+        use_geometric_center : bool, optional
+            Whether the rotation should be centered around the geometric
+            or mass center.
+        arr_mask : NDArray, optional
+            The mask of `arr` that will be equivalently rotated.
+        mask_coordinates : NDArray, optional
+            Equivalent to `coordinates`, but containing elements of `arr_mask`
+            that should be rotated.
+        out_mask : NDArray, optional
+            The output array to write the rotation of `arr_mask` to.
+        """
+        rotate_mask = arr_mask is not None and mask_coordinates is not None
+        return_type = (out is None) + 2 * rotate_mask * (out_mask is None)
+
+        # Otherwise array might be slightly shifted by centering
+        if np.allclose(
+            rotation_matrix,
+            np.eye(rotation_matrix.shape[0], dtype=rotation_matrix.dtype),
+        ):
+            center_rotation = False
+
+        coordinates_rotated = np.empty(coordinates.shape, dtype=rotation_matrix.dtype)
+        mask_rotated = (
+            np.empty(mask_coordinates.shape, dtype=rotation_matrix.dtype)
+            if rotate_mask
+            else None
+        )
+
+        center = np.array(arr.shape) // 2 if use_geometric_center else None
+        if translation is None:
+            translation = np.zeros(coordinates_rotated.shape[0])
+
+        rigid_transform(
+            coordinates=coordinates,
+            coordinates_mask=mask_coordinates,
+            out=coordinates_rotated,
+            out_mask=mask_rotated,
+            rotation_matrix=rotation_matrix,
+            translation=translation,
+            use_geometric_center=use_geometric_center,
+            center=center,
+        )
+
+        coordinates_rotated = coordinates_rotated.astype(int)
+        offset = coordinates_rotated.min(axis=1)
+        np.multiply(offset, offset < 0, out=offset)
+        coordinates_rotated -= offset[:, None]
+
+        out_offset = np.zeros(
+            coordinates_rotated.shape[0], dtype=coordinates_rotated.dtype
+        )
+        if out is None:
+            out_offset = coordinates_rotated.min(axis=1)
+            coordinates_rotated -= out_offset[:, None]
+            out = np.zeros(coordinates_rotated.max(axis=1) + 1, dtype=arr.dtype)
+
+        if rotate_mask:
+            mask_rotated = mask_rotated.astype(int)
+            if out_mask is None:
+                mask_rotated -= out_offset[:, None]
+                out_mask = np.zeros(
+                    coordinates_rotated.max(axis=1) + 1, dtype=arr.dtype
+                )
+
+            in_box = np.logical_and(
+                mask_rotated < np.array(out_mask.shape)[:, None],
+                mask_rotated >= 0,
+            ).min(axis=0)
+            out_of_box = np.invert(in_box).sum()
+            if out_of_box != 0:
+                print(
+                    f"{out_of_box} elements out of bounds. Perhaps increase"
+                    " *arr_mask* size."
+                )
+
+            mask_coordinates = tuple(mask_coordinates[:, in_box])
+            mask_rotated = tuple(mask_rotated[:, in_box])
+            np.add.at(out_mask, mask_rotated, arr_mask[mask_coordinates])
+
+        # Negative coordinates would be (mis)interpreted as reverse index
+        in_box = np.logical_and(
+            coordinates_rotated < np.array(out.shape)[:, None], coordinates_rotated >= 0
+        ).min(axis=0)
+        out_of_box = np.invert(in_box).sum()
+        if out_of_box != 0:
+            print(f"{out_of_box} elements out of bounds. Perhaps increase *out* size.")
+
+        coordinates = coordinates[:, in_box]
+        coordinates_rotated = coordinates_rotated[:, in_box]
+
+        coordinates = tuple(coordinates)
+        coordinates_rotated = tuple(coordinates_rotated)
+        np.add.at(out, coordinates_rotated, arr[coordinates])
+
+        match return_type:
+            case 0:
+                return None
+            case 1:
+                return out
+            case 2:
+                return out_mask
+            case 3:
+                return out, out_mask
+
+    def center_of_mass(self, arr: NDArray, cutoff: float = None) -> NDArray:
+        """
+        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.
+
+        Parameters
+        ----------
+        arr : NDArray
+            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.
+
+        Returns
+        -------
+        NDArray
+            Center of mass with shape (arr.ndim).
+        """
+        cutoff = arr.min() - 1 if cutoff is None else cutoff
+        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]
+
+        center_of_mass = self.array(
+            [
+                self.sum(self.multiply(arr, grids[dim])) / denominator
+                for dim in range(arr.ndim)
+            ]
+        )
+
+        return center_of_mass
+
+    def get_available_memory(self) -> int:
+        return virtual_memory().available
+
+    @contextmanager
+    def set_device(self, device_index: int):
+        yield None
+
+    def device_count(self) -> int:
+        return 1
+
+    @staticmethod
+    def reverse(arr: NDArray) -> NDArray:
+        """
+        Reverse the order of elements in an array along all its axes.
+
+        Parameters
+        ----------
+        arr : NDArray
+            Input array.
+
+        Returns
+        -------
+        NDArray
+            Reversed array.
+        """
+        return arr[(slice(None, None, -1),) * arr.ndim]