pytme 0.1.6__cp311-cp311-macosx_14_0_arm64.whl → 0.1.8__cp311-cp311-macosx_14_0_arm64.whl

tme/backends/mlx_backend.py

@@ -0,0 +1,269 @@
+""" Backend Apple's MLX library for template matching.
+
+    Copyright (c) 2024 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+from typing import Tuple, List, Callable
+
+import numpy as np
+
+from .npfftw_backend import NumpyFFTWBackend
+from ..types import NDArray, MlxArray, Scalar
+
+
+class MLXBackend(NumpyFFTWBackend):
+    """
+    A MLX based backend for template matching.
+    """
+
+    def __init__(
+        self,
+        device="cpu",
+        default_dtype=None,
+        complex_dtype=None,
+        default_dtype_int=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
+        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
+
+        super().__init__(
+            array_backend=mx,
+            default_dtype=default_dtype,
+            complex_dtype=complex_dtype,
+            default_dtype_int=default_dtype_int,
+        )
+
+        self.device = device
+
+    def to_backend_array(self, arr: NDArray) -> MlxArray:
+        return self._array_backend.array(arr)
+
+    def to_numpy_array(self, arr: MlxArray) -> NDArray:
+        return np.array(arr)
+
+    def to_cpu_array(self, arr: MlxArray) -> NDArray:
+        return arr
+
+    def free_cache(self):
+        pass
+
+    def mod(self, arr1: MlxArray, arr2: MlxArray, out: MlxArray = None) -> MlxArray:
+        if out is not None:
+            out[:] = arr1 % arr2
+            return None
+        return arr1 % arr2
+
+    def add(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.add(x1, x2, **kwargs)
+            return None
+        return self._array_backend.add(x1, x2, **kwargs)
+
+    def std(self, arr: MlxArray, axis) -> Scalar:
+        return self._array_backend.sqrt(arr.var(axis=axis))
+
+    def unique(self, *args, **kwargs):
+        ret = np.unique(*args, **kwargs)
+        if isinstance(ret, tuple):
+            ret = [self.to_backend_array(x) for x in ret]
+        return ret
+
+    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:
+        arr[:] = value
+
+    def zeros(self, shape: Tuple[int], dtype: type = None) -> MlxArray:
+        return self._array_backend.zeros(shape=shape, dtype=dtype)
+
+    def roll(self, a: MlxArray, shift, axis, **kwargs):
+        a = self.to_numpy_array(a)
+        ret = NumpyFFTWBackend().roll(
+            a,
+            shift=shift,
+            axis=axis,
+            **kwargs,
+        )
+        return self.to_backend_array(ret)
+
+    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), self._default_dtype_int)
+        stops = self.astype(self.add(starts, newshape), self._default_dtype_int)
+        starts, stops = starts.tolist(), stops.tolist()
+        box = tuple(slice(start, stop) for start, stop in zip(starts, stops))
+        return arr[box]
+
+    def build_fft(
+        self, fast_shape: Tuple[int], fast_ft_shape: Tuple[int], **kwargs
+    ) -> Tuple[Callable, Callable]:
+        """
+        Build fft 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`).
+        **kwargs : dict, optional
+            Additional parameters that are not used for now.
+
+        Returns
+        -------
+        tuple
+            Tupple containing callable rfft and irfft object.
+        """
+
+        # Runs on mlx.core.cpu until Metal support is available
+        def rfftn(arr: MlxArray, out: MlxArray, shape: Tuple[int] = fast_shape) -> None:
+            out[:] = self._array_backend.fft.rfftn(
+                arr, s=shape, stream=self._array_backend.cpu
+            )
+
+        def irfftn(
+            arr: MlxArray, out: MlxArray, shape: Tuple[int] = fast_shape
+        ) -> None:
+            out[:] = self._array_backend.fft.irfftn(
+                arr, s=shape, stream=self._array_backend.cpu
+            )
+
+        return rfftn, irfftn
+
+    def sharedarr_to_arr(
+        self, shape: Tuple[int], dtype: str, shm: MlxArray
+    ) -> MlxArray:
+        return shm
+
+    @staticmethod
+    def arr_to_sharedarr(arr: MlxArray, shared_memory_handler: type = None) -> MlxArray:
+        return arr
+
+    def topk_indices(self, arr: NDArray, k: int):
+        arr = self.to_numpy_array(arr)
+        ret = NumpyFFTWBackend().topk_indices(arr=arr, k=k)
+        ret = [self.to_backend_array(x) for x in ret]
+        return ret
+
+    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:
+        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)
+
+        if arr_mask is not None:
+            arr_mask = self.to_numpy_array(arr_mask)
+
+        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)
+
+        ret = NumpyFFTWBackend().rotate_array(
+            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
+
+        if out is not None:
+            out[:] = self.to_backend_array(out_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
+
+    def indices(self, arr: List) -> MlxArray:
+        ret = NumpyFFTWBackend().indices(arr)
+        return self.to_backend_array(ret)

tme/backends/npfftw_backend.py

@@ -38,7 +38,7 @@ class NumpyFFTWBackend(MatchingBackend):
             array_backend=array_backend,
             default_dtype=default_dtype,
             complex_dtype=complex_dtype,
-            default_dtype_int=np.int32,
+            default_dtype_int=default_dtype_int,
         )
         self.affine_transform = affine_transform
 
@@ -424,8 +424,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), int)
-        stops = self.add(starts, newshape)
+        starts = self.astype(self.divide(starts, 2), self._default_dtype_int)
+        stops = self.astype(self.add(starts, newshape), self._default_dtype_int)
         box = tuple(slice(start, stop) for start, stop in zip(starts, stops))
         return arr[box]
 
@@ -493,7 +493,9 @@ class NumpyFFTWBackend(MatchingBackend):
         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.
+            Spline interpolation order. Has to be in the range 0-5. Non-zero
+            elements will be converted into a point-cloud and rotated according
+            to ``rotation_matrix`` if order is None.
         """
 
         if order is None:
@@ -611,13 +613,6 @@ class NumpyFFTWBackend(MatchingBackend):
         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)

tme/backends/pytorch_backend.py

@@ -6,13 +6,12 @@
     Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
-from typing import Tuple, Dict, Callable
+from typing import Tuple, Callable
 from contextlib import contextmanager
 from multiprocessing import shared_memory
 from multiprocessing.managers import SharedMemoryManager
 
-from numpy.typing import NDArray
-
+import numpy as np
 from .npfftw_backend import NumpyFFTWBackend
 from ..types import NDArray, TorchTensor
 
@@ -47,7 +46,6 @@ class PytorchBackend(NumpyFFTWBackend):
         )
         self.device = device
         self.F = F
-        self._default_dtype_int = torch.int32
 
     def to_backend_array(self, arr: NDArray) -> TorchTensor:
         if isinstance(arr, self._array_backend.Tensor):
@@ -57,6 +55,8 @@ class PytorchBackend(NumpyFFTWBackend):
         return self.tensor(arr, device=self.device)
 
     def to_numpy_array(self, arr: TorchTensor) -> NDArray:
+        if isinstance(arr, np.ndarray):
+            return arr
         return arr.cpu().numpy()
 
     def to_cpu_array(self, arr: TorchTensor) -> NDArray:
@@ -131,7 +131,7 @@ class PytorchBackend(NumpyFFTWBackend):
 
     def full(self, shape, fill_value, dtype=None):
         return self._array_backend.full(
-            size = shape, dtype=dtype, fill_value=fill_value, device=self.device
+            size=shape, dtype=dtype, fill_value=fill_value, device=self.device
         )
 
     def datatype_bytes(self, dtype: type) -> int:

tme/density.py

@@ -186,8 +186,8 @@ class Density:
         Notes
         -----
         If ``filename`` ends with ".em" or ".em.gz" the method will parse it as EM file.
-        Otherwise it defaults to the CCP4/MRC format and on failure, defaults to
-        skimage.io.imread regardless of the extension. Currently, the later does not
+        Otherwise it defaults to the CCP4/MRC format and on failure, switches to
+        :obj:`skimage.io.imread` regardless of the extension. Currently, the later does not
         extract origin or sampling_rate information from the file.
 
         See Also
@@ -199,16 +199,16 @@ class Density:
             func = cls._load_mrc
             if filename.endswith(".em") or filename.endswith(".em.gz"):
                 func = cls._load_em
-            data, origin, sampling_rate = func(
+            data, origin, sampling_rate, meta = func(
                 filename=filename, subset=subset, use_memmap=use_memmap
             )
         except ValueError:
-            data, origin, sampling_rate = cls._load_skio(filename=filename)
+            data, origin, sampling_rate, meta = cls._load_skio(filename=filename)
             if subset is not None:
                 cls._validate_slices(slices=subset, shape=data.shape)
                 data = data[subset].copy()
 
-        return cls(data=data, origin=origin, sampling_rate=sampling_rate)
+        return cls(data=data, origin=origin, sampling_rate=sampling_rate, metadata=meta)
 
     @classmethod
     def _load_mrc(
@@ -305,6 +305,15 @@ class Density:
             ) and not np.all(start == 0):
                 origin = np.multiply(start, sampling_rate)
 
+            extended_header = mrc.header.nsymbt
+
+            metadata = {
+                "min": float(mrc.header.dmin),
+                "max": float(mrc.header.dmax),
+                "mean": float(mrc.header.dmean),
+                "std": float(mrc.header.rms),
+            }
+
         if is_gzipped(filename):
             if use_memmap:
                 warnings.warn(
@@ -325,9 +334,9 @@ class Density:
                 slices=subset,
                 data_shape=data_shape,
                 dtype=data_type,
-                header_size=1024,
+                header_size=1024 + extended_header,
             )
-            return data, origin, sampling_rate
+            return data, origin, sampling_rate, metadata
 
         if not use_memmap:
             with mrcfile.open(filename, header_only=False) as mrc:
@@ -341,7 +350,7 @@ class Density:
             data = np.transpose(data, crs_index)
             start = np.take(start, crs_index)
 
-        return data, origin, sampling_rate
+        return data, origin, sampling_rate, metadata
 
     @classmethod
     def _load_em(
@@ -446,7 +455,7 @@ class Density:
             pixel_size = 1
         sampling_rate = np.repeat(pixel_size, data.ndim).astype(data.dtype)
 
-        return data, origin, sampling_rate
+        return data, origin, sampling_rate, {}
 
     @staticmethod
     def _validate_slices(slices: Tuple[slice], shape: Tuple[int]):
@@ -553,12 +562,12 @@ class Density:
     @staticmethod
     def _load_skio(filename: str) -> Tuple[NDArray]:
         """
-        Uses skimage.io.imread to extract data from filename.
+        Uses :obj:`skimage.io.imread` to extract data from filename [1]_.
 
         Parameters
         ----------
         filename : str
-            Path to a file whose format is supported by skimage.io.imread.
+            Path to a file whose format is supported by :obj:`skimage.io.imread`.
 
         Returns
         -------
@@ -590,7 +599,7 @@ class Density:
         warnings.warn(
             "origin and sampling_rate are not yet extracted from non CCP4/MRC files."
         )
-        return data, np.zeros(data.ndim), np.ones(data.ndim)
+        return data, np.zeros(data.ndim), np.ones(data.ndim), {}
 
     @classmethod
     def from_structure(
@@ -720,7 +729,7 @@ class Density:
         :py:meth:`tme.structure.Structure.to_volume`
         """
         structure = filename_or_structure
-        if type(filename_or_structure) == str:
+        if isinstance(filename_or_structure, str):
             structure = Structure.from_file(
                 filename=filename_or_structure,
                 filter_by_elements=filter_by_elements,
@@ -790,8 +799,8 @@ class Density:
         Notes
         -----
         If ``filename`` ends with ".em" or ".em.gz", the method will create an EM file.
-        Otherwise, it defaults to the CCP4/MRC format, and on failure, it falls back
-        to `skimage.io.imsave`.
+        Otherwise, it defaults to the CCP4/MRC format, and on failure, falls back
+        to :obj:`skimage.io.imsave`.
 
         See Also
         --------
@@ -879,12 +888,12 @@ class Density:
 
     def _save_skio(self, filename: str, gzip: bool) -> None:
         """
-        Uses skimage.io.imsave to write data to filename.
+        Uses :obj:`skimage.io.imsave` to write data to filename [1]_.
 
         Parameters
         ----------
         filename : str
-            Path to write to with a format supported by skimage.io.imsave.
+            Path to write to with a format supported by :obj:`skimage.io.imsave`.
         gzip : bool, optional
             If True, the output will be gzip compressed.
 
@@ -907,7 +916,8 @@ class Density:
         Returns a copy of the current class instance with all elements in
         :py:attr:`Density.data` set to zero. :py:attr:`Density.origin` and
         :py:attr:`Density.sampling_rate` will be copied, while
-        :py:attr:`Density.metadata` will be initialized to an empty dictionary.
+        :py:attr:`Density.metadata` will be initialized to contain min, max,
+        mean and standard deviation of :py:attr:`Density.data`.
 
         Examples
         --------
@@ -922,6 +932,7 @@ class Density:
             data=np.zeros_like(self.data),
             origin=deepcopy(self.origin),
             sampling_rate=deepcopy(self.sampling_rate),
+            metadata={"min": 0, "max": 0, "mean": 0, "std": 0},
         )
 
     def copy(self) -> "Density":
@@ -1407,9 +1418,9 @@ class Density:
 
     def centered(self, cutoff: float = 0) -> Tuple["Density", NDArray]:
         """
-        Shifts the data center of mass to the center of the data array. The box size
-        of the return Density object is at least equal to the box size of the class
-        instance.
+        Shifts the data center of mass to the center of the data array using linear
+        interpolation. The box size of the returned :py:class:`Density` object is at
+        least equal to the box size of the class instance.
 
         Parameters
         ----------
@@ -1442,8 +1453,8 @@ class Density:
         --------
         :py:meth:`Density.centered` returns a tuple containing a centered version
         of the current :py:class:`Density` instance, as well as an array with
-        translations. The translation corresponds to the shift that was used to
-        center the current :py:class:`Density` instance.
+        translations. The translation corresponds to the shift that between the
+        center of mass and the center of the internal :py:attr:`Density.data` attribute.
 
         >>> import numpy as np
         >>> from tme import Density
@@ -1463,12 +1474,13 @@ class Density:
         internal :py:attr:`Density.data` attribute:
 
         >>> centered_dens.data
-        array([[0., 0., 0., 0., 0., 0.],
-               [0., 1., 1., 1., 1., 1.],
-               [0., 1., 1., 1., 1., 1.],
-               [0., 1., 1., 1., 1., 1.],
-               [0., 1., 1., 1., 1., 1.],
-               [0., 1., 1., 1., 1., 1.]])
+        array([[0., 0., 0., 0., 0., 0., 0.],
+               [0., 1., 1., 1., 1., 1., 0.],
+               [0., 1., 1., 1., 1., 1., 0.],
+               [0., 1., 1., 1., 1., 1., 0.],
+               [0., 1., 1., 1., 1., 1., 0.],
+               [0., 1., 1., 1., 1., 1., 0.],
+               [0., 0., 0., 0., 0., 0., 0.]])
 
         `centered_dens` is sufficiently large to represent all rotations that
         could be applied to the :py:attr:`Density.data` attribute. Lets look
@@ -1494,14 +1506,15 @@ class Density:
         ret.pad(new_shape)
 
         center = self.center_of_mass(ret.data, cutoff)
-        shift = np.subtract(np.divide(ret.shape, 2), center).astype(int)
+        shift = np.subtract(np.divide(ret.shape, 2), center)
 
         ret = ret.rigid_transform(
             translation=shift,
             rotation_matrix=np.eye(ret.data.ndim),
             use_geometric_center=False,
+            order=1,
         )
-        offset = np.subtract(center, self.center_of_mass(ret.data))
+        offset = np.subtract(center, self.center_of_mass(ret.data, cutoff))
 
         return ret, offset