pytme 0.3b0.post1__cp311-cp311-macosx_15_0_arm64.whl → 0.3.1__cp311-cp311-macosx_15_0_arm64.whl

tme/backends/cupy_backend.py

@@ -6,9 +6,9 @@ Copyright (c) 2023 European Molecular Biology Laboratory
 Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
+from typing import Tuple, List
 from importlib.util import find_spec
 from contextlib import contextmanager
-from typing import Tuple, Callable, List
 
 from .npfftw_backend import NumpyFFTWBackend
 from ..types import CupyArray, NDArray, shm_type
@@ -111,53 +111,14 @@ class CupyBackend(NumpyFFTWBackend):
     def unravel_index(self, indices, shape):
         return self._array_backend.unravel_index(indices=indices, dims=shape)
 
-    def build_fft(
-        self,
-        fwd_shape: Tuple[int],
-        inv_shape: Tuple[int],
-        inv_output_shape: Tuple[int] = None,
-        fwd_axes: Tuple[int] = None,
-        inv_axes: Tuple[int] = None,
-        **kwargs,
-    ) -> Tuple[Callable, Callable]:
-        cache = self._array_backend.fft.config.get_plan_cache()
-        current_device = self._array_backend.cuda.device.get_device_id()
-
-        previous_transform = [fwd_shape, inv_shape]
-        if current_device in PLAN_CACHE:
-            previous_transform = PLAN_CACHE[current_device]
-
-        real_diff, cmplx_diff = True, True
-        if len(fwd_shape) == len(previous_transform[0]):
-            real_diff = fwd_shape == previous_transform[0]
-        if len(inv_shape) == len(previous_transform[1]):
-            cmplx_diff = inv_shape == previous_transform[1]
-
-        if real_diff or cmplx_diff:
-            cache.clear()
-
-        rfft_shape = self._format_fft_shape(fwd_shape, fwd_axes)
-        irfft_shape = fwd_shape if inv_output_shape is None else inv_output_shape
-        irfft_shape = self._format_fft_shape(irfft_shape, inv_axes)
-
-        def rfftn(
-            arr: CupyArray, out: CupyArray = None, s=rfft_shape, axes=fwd_axes
-        ) -> CupyArray:
-            return self.rfftn(arr, s=s, axes=fwd_axes, overwrite_x=True)
-
-        def irfftn(
-            arr: CupyArray, out: CupyArray = None, s=irfft_shape, axes=inv_axes
-        ) -> CupyArray:
-            return self.irfftn(arr, s=s, axes=inv_axes, overwrite_x=True)
-
-        PLAN_CACHE[current_device] = [fwd_shape, inv_shape]
-        return rfftn, irfftn
+    def free_cache(self):
+        self._array_backend.fft.config.get_plan_cache().clear()
 
     def rfftn(self, arr: CupyArray, out: CupyArray = None, **kwargs) -> CupyArray:
         return self._cufft.rfftn(arr, **kwargs)
 
     def irfftn(self, arr: CupyArray, out: CupyArray = None, **kwargs) -> CupyArray:
-        return self._cufft.irfftn(arr, **kwargs)
+        return self._cufft.irfftn(arr, **kwargs).astype(self._float_dtype)
 
     def compute_convolution_shapes(
         self, arr1_shape: Tuple[int], arr2_shape: Tuple[int]
@@ -235,7 +196,7 @@ class CupyBackend(NumpyFFTWBackend):
             )
             return None
 
-        if data.ndim == 3 and cache and self.texture_available:
+        if data.ndim == 3 and cache and self.texture_available and not batched:
             # Device memory pool (should) come to rescue performance
             temp = self.zeros(data.shape, data.dtype)
             texture = self._get_texture(data, order=order, prefilter=prefilter)

tme/backends/jax_backend.py

@@ -7,7 +7,9 @@ Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 from functools import wraps
-from typing import Tuple, List, Callable, Dict
+from typing import Tuple, List, Dict, Any
+
+import numpy as np
 
 from ..types import BackendArray
 from .npfftw_backend import NumpyFFTWBackend, shm_type
@@ -64,6 +66,10 @@ class JaxBackend(NumpyFFTWBackend):
         arr = arr.at[idx].set(value)
         return arr
 
+    def addat(self, arr, indices, values):
+        arr = arr.at[indices].add(values)
+        return arr
+
     def topleft_pad(
         self, arr: BackendArray, shape: Tuple[int], padval: int = 0
     ) -> BackendArray:
@@ -88,6 +94,7 @@ class JaxBackend(NumpyFFTWBackend):
             "sqrt",
             "maximum",
             "exp",
+            "mod",
         ]
         for ufunc in ufuncs:
             backend_method = emulate_out(getattr(self._array_backend, ufunc))
@@ -103,27 +110,6 @@ class JaxBackend(NumpyFFTWBackend):
             shape=arr.shape, dtype=arr.dtype, fill_value=value
         )
 
-    def build_fft(
-        self,
-        fwd_shape: Tuple[int],
-        inv_shape: Tuple[int] = None,
-        inv_output_shape: Tuple[int] = None,
-        fwd_axes: Tuple[int] = None,
-        inv_axes: Tuple[int] = None,
-        **kwargs,
-    ) -> Tuple[Callable, Callable]:
-        rfft_shape = self._format_fft_shape(fwd_shape, fwd_axes)
-        irfft_shape = fwd_shape if inv_output_shape is None else inv_output_shape
-        irfft_shape = self._format_fft_shape(irfft_shape, inv_axes)
-
-        def rfftn(arr, out=None, s=rfft_shape, axes=fwd_axes):
-            return self._array_backend.fft.rfftn(arr, s=s, axes=axes)
-
-        def irfftn(arr, out=None, s=irfft_shape, axes=inv_axes):
-            return self._array_backend.fft.irfftn(arr, s=s, axes=axes)
-
-        return rfftn, irfftn
-
     def rfftn(self, arr: BackendArray, *args, **kwargs) -> BackendArray:
         return self._array_backend.fft.rfftn(arr, **kwargs)
 
@@ -194,12 +180,30 @@ class JaxBackend(NumpyFFTWBackend):
 
         return convolution_shape, fast_shape, fast_ft_shape
 
+    def _to_hashable(self, obj: Any) -> Tuple[str, Tuple]:
+        if isinstance(obj, np.ndarray):
+            return ("array", (tuple(obj.flatten().tolist()), obj.shape))
+        return ("other", obj)
+
+    def _from_hashable(self, type_info: str, data: Any) -> Any:
+        if type_info == "array":
+            data, shape = data
+            return self.array(data).reshape(shape)
+        return data
+
+    def _dict_to_tuple(self, data: Dict) -> Tuple:
+        return tuple((k, self._to_hashable(v)) for k, v in data.items())
+
+    def _tuple_to_dict(self, data: Tuple) -> Dict:
+        return {x[0]: self._from_hashable(*x[1]) for x in data}
+
     def scan(
         self,
         matching_data: type,
         splits: Tuple[Tuple[slice, slice]],
         n_jobs: int,
-        callback_class,
+        callback_class: object,
+        callback_class_args: Dict,
         rotate_mask: bool = False,
         **kwargs,
     ) -> List:
@@ -220,16 +224,20 @@ class JaxBackend(NumpyFFTWBackend):
             target_shape=self.to_numpy_array(target_shape),
             template_shape=self.to_numpy_array(matching_data._template.shape),
             batch_mask=self.to_numpy_array(matching_data._batch_mask),
-            pad_target=pad_target,
         )
         analyzer_args = {
-            "convolution_mode": convolution_mode,
+            "shape": fast_shape,
             "fourier_shift": shift,
+            "fast_shape": fast_shape,
             "targetshape": target_shape,
             "templateshape": matching_data.template.shape,
             "convolution_shape": conv_shape,
+            "convolution_mode": convolution_mode,
+            "thread_safe": False,
+            "aggregate_axis": matching_data._batch_axis(matching_data._batch_mask),
+            "n_rotations": matching_data.rotations.shape[0],
+            "jax_mode": True,
         }
-
         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
@@ -245,6 +253,9 @@ class JaxBackend(NumpyFFTWBackend):
             for i in range(matching_data.rotations.shape[0])
         }
         for split_start in range(0, len(splits), n_jobs):
+
+            analyzer_kwargs = []
+
             split_subset = splits[split_start : (split_start + n_jobs)]
             if not len(split_subset):
                 continue
@@ -256,8 +267,15 @@ class JaxBackend(NumpyFFTWBackend):
                     target_pad=target_pad,
                     template_slice=template_split,
                 )
+                cur_args = analyzer_args.copy()
+                cur_args["offset"] = translation_offset
+                cur_args.update(callback_class_args)
+
+                analyzer_kwargs.append(self._dict_to_tuple(cur_args))
+
+                _target = self.astype(base._target, self._float_dtype)
                 translation_offsets.append(translation_offset)
-                targets.append(self.topleft_pad(base._target, fast_shape))
+                targets.append(self.topleft_pad(_target, fast_shape))
 
             if create_filter:
                 filter_args = {
@@ -279,24 +297,27 @@ class JaxBackend(NumpyFFTWBackend):
 
             create_filter, create_template_filter, create_target_filter = (False,) * 3
             base, targets = None, self._array_backend.stack(targets)
-            scores, rotations = scan_inner(
+
+            analyzer_kwargs = tuple(analyzer_kwargs)
+            states = scan_inner(
                 self.astype(targets, self._float_dtype),
-                matching_data.template,
-                matching_data.template_mask,
+                self.astype(matching_data.template, self._float_dtype),
+                self.astype(matching_data.template_mask, self._float_dtype),
                 matching_data.rotations,
                 template_filter,
                 target_filter,
                 fast_shape,
                 rotate_mask,
+                callback_class,
+                analyzer_kwargs,
             )
 
-            for index in range(scores.shape[0]):
-                temp = callback_class(
-                    shape=scores[index].shape,
-                    offset=translation_offsets[index],
-                )
-                state = (scores[index], rotations[index], rotation_mapping)
-                ret.append(temp.result(state, **analyzer_args))
+            for index in range(targets.shape[0]):
+                kwargs = self._tuple_to_dict(analyzer_kwargs[index])
+                analyzer = callback_class(**kwargs)
+
+                state = (states[0][index], states[1][index], rotation_mapping)
+                ret.append(analyzer.result(state, **kwargs))
         return ret
 
     def get_available_memory(self) -> int:

tme/backends/matching_backend.py

@@ -8,7 +8,7 @@ Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 
 from abc import ABC, abstractmethod
 from multiprocessing import shared_memory
-from typing import Tuple, Callable, List, Any, Union, Optional, Generator, Dict
+from typing import Tuple, Callable, List, Any, Union, Optional, Generator
 
 from ..types import BackendArray, NDArray, Scalar, shm_type
 
@@ -1087,57 +1087,12 @@ class MatchingBackend(ABC):
         """
 
     @abstractmethod
-    def build_fft(
-        self,
-        fwd_shape: Tuple[int],
-        inv_shape: Tuple[int],
-        real_dtype: type,
-        cmpl_dtype: type,
-        inv_output_shape: Tuple[int] = None,
-        temp_fwd: NDArray = None,
-        temp_inv: NDArray = None,
-        fwd_axes: Tuple[int] = None,
-        inv_axes: Tuple[int] = None,
-        fftargs: Dict = {},
-    ) -> Tuple[Callable, Callable]:
-        """
-        Build forward and inverse real fourier transform functions. The returned
-        callables have two parameters ``arr`` and ``out`` which correspond to the
-        input and output of the Fourier transform. The methods return the output
-        of the respective function call, regardless of ``out`` being provided or not,
-        analogous to most numpy functions.
+    def rfftn(self, **kwargs):
+        """Perform an n-D real FFT."""
 
-        Parameters
-        ----------
-        fwd_shape : tuple
-            Input shape for the forward Fourier transform.
-            (see `compute_convolution_shapes`).
-        inv_shape : tuple
-            Input shape for the inverse Fourier transform.
-        real_dtype : dtype
-            Data type of the forward Fourier transform.
-        complex_dtype : dtype
-            Data type of the inverse Fourier transform.
-        inv_output_shape : tuple, optional
-            Output shape of the inverse Fourier transform. By default fast_shape.
-        fftargs : dict, optional
-            Dictionary passed to pyFFTW builders.
-        temp_fwd : NDArray, optional
-            Temporary array to build the forward transform. Superseeds shape defined by
-            fwd_shape if provided.
-        temp_inv : NDArray, optional
-            Temporary array to build the inverse transform. Superseeds shape defined by
-            inv_shape if provided.
-        fwd_axes : tuple of int
-            Axes to perform the forward Fourier transform over.
-        inv_axes : tuple of int
-            Axes to perform the inverse Fourier transform over.
-
-        Returns
-        -------
-        tuple
-            Tuple of callables for forward and inverse real Fourier transform.
-        """
+    @abstractmethod
+    def irfftn(self, **kwargs):
+        """Perform an n-D real inverse FFT."""
 
     def extract_center(self, arr: BackendArray, newshape: Tuple[int]) -> BackendArray:
         """

tme/backends/mlx_backend.py

@@ -6,7 +6,7 @@ Copyright (c) 2024 European Molecular Biology Laboratory
 Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
-from typing import Tuple, List, Callable
+from typing import Tuple, List
 
 import numpy as np
 
@@ -144,32 +144,6 @@ class MLXBackend(NumpyFFTWBackend):
         box = tuple(slice(start, stop) for start, stop in zip(starts, stops))
         return arr[box]
 
-    def build_fft(
-        self,
-        fwd_shape: Tuple[int],
-        inv_shape: Tuple[int] = None,
-        inv_output_shape: Tuple[int] = None,
-        fwd_axes: Tuple[int] = None,
-        inv_axes: Tuple[int] = None,
-        **kwargs,
-    ) -> Tuple[Callable, Callable]:
-        # Runs on mlx.core.cpu until Metal support is available
-        rfft_shape = self._format_fft_shape(fwd_shape, fwd_axes)
-        irfft_shape = fwd_shape if inv_output_shape is None else inv_output_shape
-        irfft_shape = self._format_fft_shape(irfft_shape, inv_axes)
-
-        def rfftn(arr: MlxArray, out: MlxArray = None, s=rfft_shape, axes=fwd_axes):
-            out[:] = self._array_backend.fft.rfftn(
-                arr, s=s, axes=axes, stream=self._array_backend.cpu
-            )
-
-        def irfftn(arr: MlxArray, out: MlxArray = None, s=irfft_shape, axes=inv_axes):
-            out[:] = self._array_backend.fft.irfftn(
-                arr, s=s, axes=axes, stream=self._array_backend.cpu
-            )
-
-        return rfftn, irfftn
-
     def rfftn(self, arr, *args, **kwargs):
         return self.fft.rfftn(arr, stream=self._array_backend.cpu, **kwargs)
 

tme/backends/npfftw_backend.py

@@ -9,17 +9,23 @@ Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 import os
 from psutil import virtual_memory
 from contextlib import contextmanager
-from typing import Tuple, Dict, List, Type
+from typing import Tuple, List, Type
 
 import scipy
 import numpy as np
 from scipy.ndimage import maximum_filter, affine_transform
-from pyfftw.builders import rfftn as rfftn_builder, irfftn as irfftn_builder
-from pyfftw import zeros_aligned, simd_alignment, FFTW, next_fast_len, interfaces
+from pyfftw import (
+    zeros_aligned,
+    simd_alignment,
+    next_fast_len,
+    interfaces,
+    config,
+)
 
 from ..types import NDArray, BackendArray, shm_type
 from .matching_backend import MatchingBackend, _create_metafunction
 
+
 os.environ["MKL_NUM_THREADS"] = "1"
 os.environ["OMP_NUM_THREADS"] = "1"
 os.environ["PYFFTW_NUM_THREADS"] = "1"
@@ -103,6 +109,20 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
         self.solve_triangular = self._solve_triangular
         self.linalg.solve_triangular = scipy.linalg.solve_triangular
 
+        try:
+            from ._numpyfftw_utils import rfftn as rfftn_cache
+            from ._numpyfftw_utils import irfftn as irfftn_cache
+
+            self._rfftn = rfftn_cache
+            self._irfftn = irfftn_cache
+        except Exception as e:
+            print(e)
+
+        config.NUM_THREADS = 1
+        config.PLANNER_EFFORT = "FFTW_MEASURE"
+        interfaces.cache.enable()
+        interfaces.cache.set_keepalive_time(360)
+
     def _linalg_cholesky(self, arr, lower=False, *args, **kwargs):
         # Upper argument is not supported until numpy 2.0
         ret = self._array_backend.linalg.cholesky(arr, *args, **kwargs)
@@ -138,7 +158,7 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
         return float
 
     def free_cache(self):
-        pass
+        interfaces.cache.disable()
 
     def transpose(self, arr: NDArray, *args, **kwargs) -> NDArray:
         return self._array_backend.transpose(arr, *args, **kwargs)
@@ -240,70 +260,53 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
         b[tuple(bind)] = arr[tuple(aind)]
         return b
 
-    def build_fft(
-        self,
-        fwd_shape: Tuple[int],
-        inv_shape: Tuple[int],
-        real_dtype: type,
-        cmpl_dtype: type,
-        fftargs: Dict = {},
-        inv_output_shape: Tuple[int] = None,
-        temp_fwd: NDArray = None,
-        temp_inv: NDArray = None,
-        fwd_axes: Tuple[int] = None,
-        inv_axes: Tuple[int] = None,
-    ) -> Tuple[FFTW, FFTW]:
-        if temp_fwd is None:
-            temp_fwd = (
-                self.zeros(fwd_shape, real_dtype) if temp_fwd is None else temp_fwd
-            )
-        if temp_inv is None:
-            temp_inv = (
-                self.zeros(inv_shape, cmpl_dtype) if temp_inv is None else temp_inv
-            )
-
-        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]
-
-        rfft_shape = self._format_fft_shape(temp_fwd.shape, fwd_axes)
-        _rfftn = rfftn_builder(temp_fwd, s=rfft_shape, axes=fwd_axes, **fftargs)
-        overwrite_input = fftargs.pop("overwrite_input", None)
-
-        irfft_shape = fwd_shape if inv_output_shape is None else inv_output_shape
-        irfft_shape = self._format_fft_shape(irfft_shape, inv_axes)
-        _irfftn = irfftn_builder(temp_inv, s=irfft_shape, axes=inv_axes, **fftargs)
-
-        def _rfftn_wrapper(arr, out, *args, **kwargs):
-            return _rfftn(arr, out)
-
-        def _irfftn_wrapper(arr, out, *args, **kwargs):
-            return _irfftn(arr, out)
-
-        fftargs["overwrite_input"] = overwrite_input
-        return _rfftn_wrapper, _irfftn_wrapper
+    def _rfftn(self, arr, out=None, **kwargs):
+        ret = interfaces.numpy_fft.rfftn(arr, **kwargs)
+        if out is not None:
+            out[:] = ret
+            return out
+        return ret
 
-    @staticmethod
-    def _format_fft_shape(shape: Tuple[int], axes: Tuple[int] = None):
-        if axes is None:
-            return shape
-        axes = tuple(sorted(range(len(shape))[i] for i in axes))
-        return tuple(shape[i] for i in axes)
+    def _irfftn(self, arr, out=None, **kwargs):
+        ret = interfaces.numpy_fft.irfftn(arr, **kwargs)
+        if out is not None:
+            out[:] = ret
+            return out
+        return ret
 
-    def rfftn(self, arr: NDArray, *args, **kwargs) -> NDArray:
-        return interfaces.numpy_fft.rfftn(arr, **kwargs)
+    def rfftn(
+        self,
+        arr: NDArray,
+        out=None,
+        auto_align_input: bool = False,
+        auto_contiguous: bool = False,
+        overwrite_input: bool = True,
+        **kwargs,
+    ) -> NDArray:
+        return self._rfftn(
+            arr,
+            auto_align_input=auto_align_input,
+            auto_contiguous=auto_contiguous,
+            overwrite_input=overwrite_input,
+            **kwargs,
+        )
 
-    def irfftn(self, arr: NDArray, *args, **kwargs) -> NDArray:
-        return interfaces.numpy_fft.irfftn(arr, **kwargs)
+    def irfftn(
+        self,
+        arr: NDArray,
+        out=None,
+        auto_align_input: bool = False,
+        auto_contiguous: bool = False,
+        overwrite_input: bool = True,
+        **kwargs,
+    ) -> NDArray:
+        return self._irfftn(
+            arr,
+            auto_align_input=auto_align_input,
+            auto_contiguous=auto_contiguous,
+            overwrite_input=overwrite_input,
+            **kwargs,
+        )
 
     def extract_center(self, arr: NDArray, newshape: Tuple[int]) -> NDArray:
         new_shape = self.to_backend_array(newshape)

tme/backends/pytorch_backend.py

@@ -7,7 +7,7 @@ Copyright (c) 2023 European Molecular Biology Laboratory
 Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
-from typing import Tuple, Callable
+from typing import Tuple
 from contextlib import contextmanager
 from multiprocessing import shared_memory
 from multiprocessing.managers import SharedMemoryManager
@@ -273,31 +273,6 @@ class PytorchBackend(NumpyFFTWBackend):
             kwargs["device"] = self.device
         return self._array_backend.eye(*args, **kwargs)
 
-    def build_fft(
-        self,
-        fwd_shape: Tuple[int],
-        inv_shape: Tuple[int],
-        inv_output_shape: Tuple[int] = None,
-        fwd_axes: Tuple[int] = None,
-        inv_axes: Tuple[int] = None,
-        **kwargs,
-    ) -> Tuple[Callable, Callable]:
-        rfft_shape = self._format_fft_shape(fwd_shape, fwd_axes)
-        irfft_shape = fwd_shape if inv_output_shape is None else inv_output_shape
-        irfft_shape = self._format_fft_shape(irfft_shape, inv_axes)
-
-        def rfftn(
-            arr: TorchTensor, out: TorchTensor, s=rfft_shape, axes=fwd_axes
-        ) -> TorchTensor:
-            return self._array_backend.fft.rfftn(arr, s=s, out=out, dim=axes)
-
-        def irfftn(
-            arr: TorchTensor, out: TorchTensor = None, s=irfft_shape, axes=inv_axes
-        ) -> TorchTensor:
-            return self._array_backend.fft.irfftn(arr, s=s, out=out, dim=axes)
-
-        return rfftn, irfftn
-
     def rfftn(self, arr: NDArray, *args, **kwargs) -> NDArray:
         kwargs["dim"] = kwargs.pop("axes", None)
         return self._array_backend.fft.rfftn(arr, **kwargs)

tme/density.py

@@ -36,6 +36,7 @@ from .matching_utils import (
     array_to_memmap,
     memmap_to_array,
     minimum_enclosing_box,
+    is_gzipped,
 )
 
 __all__ = ["Density"]
@@ -331,6 +332,7 @@ class Density:
         if non_standard_crs:
             data = np.transpose(data, crs_index)
             origin = np.take(origin, crs_index)
+            sampling_rate = np.take(sampling_rate, crs_index)
 
         return data.T, origin[::-1], sampling_rate[::-1], metadata
 
@@ -2257,9 +2259,3 @@ class Density:
         coordinates = np.array(np.where(data > 0))
         weights = self.data[tuple(coordinates)]
         return align_to_axis(coordinates.T, weights=weights, axis=axis, flip=flip)
-
-
-def is_gzipped(filename: str) -> bool:
-    """Check if a file is a gzip file by reading its magic number."""
-    with open(filename, "rb") as f:
-        return f.read(2) == b"\x1f\x8b"