pytme 0.3b0__cp311-cp311-macosx_15_0_arm64.whl → 0.3.1__cp311-cp311-macosx_15_0_arm64.whl

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
+from typing import Tuple, List, Dict, Any
+
+import numpy as np
 
 from ..types import BackendArray
 from .npfftw_backend import NumpyFFTWBackend, shm_type
@@ -51,12 +53,6 @@ class JaxBackend(NumpyFFTWBackend):
         )
         self.scipy = jsp
         self._create_ufuncs()
-        try:
-            from ._jax_utils import scan as _
-
-            self.scan = self._scan
-        except Exception:
-            pass
 
     def from_sharedarr(self, arr: BackendArray) -> BackendArray:
         return arr
@@ -70,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:
@@ -94,6 +94,7 @@ class JaxBackend(NumpyFFTWBackend):
             "sqrt",
             "maximum",
             "exp",
+            "mod",
         ]
         for ufunc in ufuncs:
             backend_method = emulate_out(getattr(self._array_backend, ufunc))
@@ -109,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)
 
@@ -189,12 +169,41 @@ class JaxBackend(NumpyFFTWBackend):
         rotations = rotations.at[:].set(self.where(update, rotations, rotation_index))
         return max_scores, rotations
 
-    def _scan(
+    def compute_convolution_shapes(
+        self, arr1_shape: Tuple[int], arr2_shape: Tuple[int]
+    ) -> Tuple[List[int], List[int], List[int]]:
+        from scipy.fft import next_fast_len
+
+        convolution_shape = [int(x + y - 1) for x, y in zip(arr1_shape, arr2_shape)]
+        fast_shape = [next_fast_len(x, real=True) 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 _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:
@@ -214,17 +223,21 @@ class JaxBackend(NumpyFFTWBackend):
         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,
+            batch_mask=self.to_numpy_array(matching_data._batch_mask),
         )
-
         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
@@ -240,25 +253,34 @@ 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
 
             targets, translation_offsets = [], []
             for target_split, template_split in split_subset:
-                base = matching_data.subset_by_slice(
+                base, translation_offset = 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))
+                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(_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:
@@ -275,25 +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.shape,
-                    offset=translation_offsets[index],
-                )
-                state = (scores, rotations, 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)
@@ -398,33 +401,33 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
         out_mask: NDArray = None,
         order: int = 3,
         cache: bool = False,
+        batched: bool = False,
     ) -> Tuple[NDArray, NDArray]:
-        out = self.zeros_like(arr) if out is None else out
-        batched = arr.ndim != rotation_matrix.shape[0]
-
-        center = self.divide(self.to_backend_array(arr.shape) - 1, 2)
-        if not use_geometric_center:
-            center = self.center_of_mass(arr, cutoff=0)
-
-        offset = int(arr.ndim - rotation_matrix.shape[0])
-        center = center[offset:]
-        translation = self.zeros(center.size) if translation is None else translation
-        matrix = self._rigid_transform_matrix(
-            rotation_matrix=rotation_matrix,
-            translation=translation,
-            center=center,
-        )
-
-        subset = tuple(slice(None) for _ in range(arr.ndim))
-        if offset > 1:
-            subset = tuple(
-                0 if i < (offset - 1) else slice(None) for i in range(arr.ndim)
+        if out is None:
+            out = self.zeros_like(arr)
+
+        # Check whether rotation_matrix is already a rigid transform matrix
+        matrix = rotation_matrix
+        if matrix.shape[-1] == (arr.ndim - int(batched)):
+            center = self.divide(self.to_backend_array(arr.shape) - 1, 2)
+            if not use_geometric_center:
+                center = self.center_of_mass(arr, cutoff=0)
+
+            offset = int(arr.ndim - rotation_matrix.shape[0])
+            center = center[offset:]
+            translation = (
+                self.zeros(center.size) if translation is None else translation
+            )
+            matrix = self._rigid_transform_matrix(
+                rotation_matrix=rotation_matrix,
+                translation=translation,
+                center=center,
             )
 
         self._rigid_transform(
-            data=arr[subset],
+            data=arr,
             matrix=matrix,
-            output=out[subset],
+            output=out,
             order=order,
             prefilter=True,
             cache=cache,
@@ -433,11 +436,13 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
 
         # Applying the prefilter leads to artifacts in the mask.
         if arr_mask is not None:
-            out_mask = self.zeros_like(arr_mask) if out_mask is None else out_mask
+            if out_mask is None:
+                out_mask = self.zeros_like(arr_mask)
+
             self._rigid_transform(
-                data=arr_mask[subset],
+                data=arr_mask,
                 matrix=matrix,
-                output=out_mask[subset],
+                output=out_mask,
                 order=order,
                 prefilter=False,
                 cache=cache,

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)
@@ -306,6 +281,9 @@ class PytorchBackend(NumpyFFTWBackend):
         kwargs["dim"] = kwargs.pop("axes", None)
         return self._array_backend.fft.irfftn(arr, **kwargs)
 
+    def _rigid_transform_matrix(self, rotation_matrix, *args, **kwargs):
+        return rotation_matrix
+
     def rigid_transform(
         self,
         arr: TorchTensor,
@@ -317,6 +295,7 @@ class PytorchBackend(NumpyFFTWBackend):
         out_mask: TorchTensor = None,
         order: int = 1,
         cache: bool = False,
+        **kwargs,
     ):
         _mode_mapping = {0: "nearest", 1: "bilinear", 3: "bicubic"}
         mode = _mode_mapping.get(order, None)

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
 
@@ -1763,12 +1765,13 @@ class Density:
                 axis=axis,
             )
 
-            arr_ft = np.fft.fftn(self.data)
+            mask, mask_ret = np.where(mask), np.where(mask_ret)
+
+            arr_ft = np.fft.fftn(self.data)[mask]
             arr_ft *= np.prod(ret_shape) / np.prod(self.shape)
             ret_ft = np.zeros(ret_shape, dtype=arr_ft.dtype)
-            ret_ft[mask_ret] = arr_ft[mask]
-            ret.data = np.real(np.fft.ifftn(ret_ft))
-
+            np.add.at(ret_ft, mask_ret, arr_ft)
+            ret.data = np.real(np.fft.ifftn(ret_ft)).astype(self.data.dtype)
         ret.sampling_rate = new_sampling_rate
         return ret
 
@@ -2256,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"