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

tme/preprocessing/_utils.py

@@ -5,12 +5,13 @@
     Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
-from typing import Tuple
+from typing import Tuple, List
 
 import numpy as np
-from numpy.typing import NDArray
 
-from ..backends import backend
+from ..backends import backend as be
+from ..backends import NumpyFFTWBackend
+from ..types import BackendArray, NDArray
 from ..matching_utils import euler_to_rotationmatrix
 
 
@@ -93,18 +94,27 @@ def frequency_grid_at_angle(
     tilt_shape = compute_tilt_shape(
         shape=shape, opening_axis=opening_axis, reduce_dim=False
     )
-    index_grid = centered_grid(shape=tilt_shape)
+
+    if angle == 0:
+        index_grid = fftfreqn(
+            tuple(x for x in tilt_shape if x != 1),
+            sampling_rate=1,
+            compute_euclidean_norm=True,
+        )
+
     if angle != 0:
         angles = np.zeros(len(shape))
         angles[tilt_axis] = angle
         rotation_matrix = euler_to_rotationmatrix(np.roll(angles, opening_axis - 1))
+
+        index_grid = fftfreqn(tilt_shape, sampling_rate=None)
         index_grid = np.einsum("ij,j...->i...", rotation_matrix, index_grid)
+        norm = np.multiply(sampling_rate, shape).astype(int)
 
-    norm = np.divide(1, 2 * sampling_rate * np.divide(shape, 2).astype(int))
+        index_grid = np.divide(index_grid.T, norm).T
+        index_grid = np.squeeze(index_grid)
+        index_grid = np.linalg.norm(index_grid, axis=(0))
 
-    index_grid = np.multiply(index_grid.T, norm).T
-    index_grid = np.squeeze(index_grid)
-    index_grid = np.linalg.norm(index_grid, axis=(0))
     return index_grid
 
 
@@ -113,9 +123,10 @@ def fftfreqn(
     sampling_rate: Tuple[float],
     compute_euclidean_norm: bool = False,
     shape_is_real_fourier: bool = False,
+    return_sparse_grid: bool = False,
 ) -> NDArray:
     """
-    Generate the n-dimensional discrete Fourier Transform sample frequencies.
+    Generate the n-dimensional discrete Fourier transform sample frequencies.
 
     Parameters:
     -----------
@@ -133,56 +144,74 @@ def fftfreqn(
     NDArray
         The sample frequencies.
     """
-    center = backend.astype(backend.divide(shape, 2), backend._int_dtype)
-
-    norm = np.ones(len(shape))
+    # There is no real need to have these operations on GPU right now
+    temp_backend = NumpyFFTWBackend()
+    norm = temp_backend.full(len(shape), fill_value=1)
+    center = temp_backend.astype(temp_backend.divide(shape, 2), temp_backend._int_dtype)
     if sampling_rate is not None:
-        norm = backend.astype(backend.multiply(shape, sampling_rate), int)
+        norm = temp_backend.astype(temp_backend.multiply(shape, sampling_rate), int)
 
     if shape_is_real_fourier:
-        center[-1] = 0
-        norm[-1] = 1
+        center[-1], norm[-1] = 0, 1
         if sampling_rate is not None:
             norm[-1] = (shape[-1] - 1) * 2 * sampling_rate
 
-    indices = backend.transpose(backend.indices(shape))
-    indices -= center
-    indices = backend.divide(indices, norm)
-    indices = backend.transpose(indices)
+    grids = []
+    for i, x in enumerate(shape):
+        baseline_dims = tuple(1 if i != t else x for t in range(len(shape)))
+        grid = (temp_backend.arange(x) - center[i]) / norm[i]
+        grids.append(temp_backend.reshape(grid, baseline_dims))
 
     if compute_euclidean_norm:
-        indices = backend.square(indices)
-        indices = backend.sum(indices, axis=0)
-        backend.sqrt(indices, out=indices)
+        grids = sum(temp_backend.square(x) for x in grids)
+        grids = temp_backend.sqrt(grids, out=grids)
+        return grids
+
+    if return_sparse_grid:
+        return grids
 
-    return indices
+    grid_flesh = temp_backend.full(shape, fill_value=1)
+    grids = temp_backend.stack(tuple(grid * grid_flesh for grid in grids))
 
+    return grids
 
-def crop_real_fourier(data: NDArray) -> NDArray:
+
+def crop_real_fourier(data: BackendArray) -> BackendArray:
     """
     Crop the real part of a Fourier transform.
 
     Parameters:
     -----------
-    data : NDArray
+    data : BackendArray
         The Fourier transformed data.
 
     Returns:
     --------
-    NDArray
+    BackendArray
         The cropped data.
     """
     stop = 1 + (data.shape[-1] // 2)
     return data[..., :stop]
 
 
-def shift_fourier(data: NDArray, shape_is_real_fourier: bool = False):
-    shift = backend.add(
-        backend.astype(backend.divide(data.shape, 2), int),
-        backend.mod(data.shape, 2),
-    )
+def compute_fourier_shape(
+    shape: Tuple[int], shape_is_real_fourier: bool = False
+) -> List[int]:
+    if shape_is_real_fourier:
+        return shape
+    shape = [int(x) for x in shape]
+    shape[-1] = 1 + shape[-1] // 2
+    return shape
+
+
+def shift_fourier(
+    data: BackendArray, shape_is_real_fourier: bool = False
+) -> BackendArray:
+    shape = be.to_backend_array(data.shape)
+    shift = be.add(be.divide(shape, 2), be.mod(shape, 2))
+    shift = [int(x) for x in shift]
     if shape_is_real_fourier:
         shift[-1] = 0
 
-    data = backend.roll(data, shift, tuple(i for i in range(len(shift))))
+    data = be.roll(data, shift, tuple(i for i in range(len(shift))))
     return data

tme/preprocessing/composable_filter.py

@@ -17,15 +17,18 @@ class ComposableFilter(ABC):
     @abstractmethod
     def __call__(self, *args, **kwargs) -> Dict:
         """
-        Parameters:
-        -----------
+
+        Parameters
+        ----------
+
         *args : tuple
             Variable length argument list.
         **kwargs : dict
             Arbitrary keyword arguments.
 
-        Returns:
-        --------
+        Returns
+        -------
+
         Dict
             A dictionary representing the result of the filtering operation.
         """

tme/preprocessing/compose.py

@@ -7,7 +7,7 @@
 
 from typing import Tuple, Dict
 
-from tme.backends import backend
+from tme.backends import backend as be
 
 
 class Compose:
@@ -42,9 +42,13 @@ class Compose:
             kwargs.update(meta)
             ret = transform(**kwargs)
 
+            if "data" not in ret:
+                continue
+
             if ret.get("is_multiplicative_filter", False):
-                backend.multiply(ret["data"], meta["data"], out=ret["data"])
-                ret["merge"] = None
+                prev_data = meta.pop("data")
+                ret["data"] = be.multiply(ret["data"], prev_data, out=ret["data"])
+                ret["merge"], prev_data = None, None
 
             meta = ret
 

tme/preprocessing/frequency_filters.py

@@ -8,12 +8,12 @@ from math import log, sqrt
 from typing import Tuple, Dict
 
 import numpy as np
-from numpy.typing import NDArray
 from scipy.ndimage import mean as ndimean
 from scipy.ndimage import map_coordinates
 
-from ._utils import fftfreqn, crop_real_fourier, shift_fourier
-from ..backends import backend
+from ..types import BackendArray
+from ..backends import backend as be
+from ._utils import fftfreqn, crop_real_fourier, shift_fourier, compute_fourier_shape
 
 
 class BandPassFilter:
@@ -29,7 +29,7 @@ class BandPassFilter:
     highpass : float, optional
         The highpass cutoff, defaults to None.
     sampling_rate : Tuple[float], optional
-        The sampling rate in Fourier space, defaults to 1.
+        The sampling r_position_to_molmapate in Fourier space, defaults to 1.
     use_gaussian : bool, optional
         Whether to use Gaussian bandpass filter, defaults to True.
     return_real_fourier : bool, optional
@@ -63,7 +63,7 @@ class BandPassFilter:
         return_real_fourier: bool = False,
         shape_is_real_fourier: bool = False,
         **kwargs,
-    ) -> NDArray:
+    ) -> BackendArray:
         """
         Generate a bandpass filter using discrete frequency cutoffs.
 
@@ -86,7 +86,7 @@ class BandPassFilter:
 
         Returns:
         --------
-        NDArray
+        BackendArray
             The bandpass filter in Fourier space.
         """
         if shape_is_real_fourier:
@@ -127,7 +127,7 @@ class BandPassFilter:
         return_real_fourier: bool = False,
         shape_is_real_fourier: bool = False,
         **kwargs,
-    ) -> NDArray:
+    ) -> BackendArray:
         """
         Generate a bandpass filter using Gaussian functions.
 
@@ -150,7 +150,7 @@ class BandPassFilter:
 
         Returns:
         --------
-        NDArray
+        BackendArray
             The bandpass filter in Fourier space.
         """
         if shape_is_real_fourier:
@@ -162,29 +162,32 @@ class BandPassFilter:
             shape_is_real_fourier=shape_is_real_fourier,
             compute_euclidean_norm=True,
         )
-        grid = -backend.square(grid)
+        grid = be.to_backend_array(grid)
+        grid = -be.square(grid)
 
         lowpass_filter, highpass_filter = 1, 1
         norm = float(sqrt(2 * log(2)))
-        upper_sampling = float(backend.max(backend.multiply(2, sampling_rate)))
+        upper_sampling = float(
+            be.max(be.multiply(2, be.to_backend_array(sampling_rate)))
+        )
 
         if lowpass is not None:
             lowpass = float(lowpass)
-            lowpass = backend.maximum(lowpass, backend.eps(backend._float_dtype))
+            lowpass = be.maximum(lowpass, be.eps(be._float_dtype))
         if highpass is not None:
             highpass = float(highpass)
-            highpass = backend.maximum(highpass, backend.eps(backend._float_dtype))
+            highpass = be.maximum(highpass, be.eps(be._float_dtype))
 
         if lowpass is not None:
             lowpass = upper_sampling / (lowpass * norm)
-            lowpass = backend.multiply(2, backend.square(lowpass))
-            lowpass_filter = backend.exp(backend.divide(grid, lowpass))
+            lowpass = be.multiply(2, be.square(lowpass))
+            lowpass_filter = be.exp(be.divide(grid, lowpass))
         if highpass is not None:
             highpass = upper_sampling / (highpass * norm)
-            highpass = backend.multiply(2, backend.square(highpass))
-            highpass_filter = 1 - backend.exp(backend.divide(grid, highpass))
+            highpass = be.multiply(2, be.square(highpass))
+            highpass_filter = 1 - be.exp(be.divide(grid, highpass))
 
-        bandpass_filter = backend.multiply(lowpass_filter, highpass_filter)
+        bandpass_filter = be.multiply(lowpass_filter, highpass_filter)
         bandpass_filter = shift_fourier(
             data=bandpass_filter, shape_is_real_fourier=shape_is_real_fourier
         )
@@ -205,7 +208,7 @@ class BandPassFilter:
         mask = func(**func_args)
 
         return {
-            "data": backend.to_backend_array(mask),
+            "data": be.to_backend_array(mask),
             "sampling_rate": func_args.get("sampling_rate", 1),
             "is_multiplicative_filter": True,
         }
@@ -237,14 +240,14 @@ class LinearWhiteningFilter:
 
     @staticmethod
     def _compute_spectrum(
-        data_rfft: NDArray, n_bins: int = None, batch_dimension: int = None
-    ) -> Tuple[NDArray, NDArray]:
+        data_rfft: BackendArray, n_bins: int = None, batch_dimension: int = None
+    ) -> Tuple[BackendArray, BackendArray]:
         """
         Compute the spectrum of the input data.
 
         Parameters:
         -----------
-        data_rfft : NDArray
+        data_rfft : BackendArray
             The Fourier transform of the input data.
         n_bins : int, optional
             The number of bins for computing the spectrum, defaults to None.
@@ -253,9 +256,9 @@ class LinearWhiteningFilter:
 
         Returns:
         --------
-        bins : NDArray
+        bins : BackendArray
             Array containing the bin indices for the spectrum.
-        radial_averages : NDArray
+        radial_averages : BackendArray
             Array containing the radial averages of the spectrum.
         """
         shape = tuple(x for i, x in enumerate(data_rfft.shape) if i != batch_dimension)
@@ -270,7 +273,7 @@ class LinearWhiteningFilter:
             shape_is_real_fourier=True,
             compute_euclidean_norm=True,
         )
-        bins = backend.to_numpy_array(bins)
+        bins = be.to_numpy_array(bins)
 
         # Implicit lowpass to nyquist
         bins = np.floor(bins * (n_bins - 1) + 0.5).astype(int)
@@ -292,11 +295,11 @@ class LinearWhiteningFilter:
 
     @staticmethod
     def _interpolate_spectrum(
-        spectrum: NDArray,
+        spectrum: BackendArray,
         shape: Tuple[int],
         shape_is_real_fourier: bool = True,
         order: int = 1,
-    ) -> NDArray:
+    ) -> BackendArray:
         """
         References
         ----------
@@ -306,19 +309,19 @@ class LinearWhiteningFilter:
         """
         grid = fftfreqn(
             shape=shape,
-            sampling_rate=.5,
+            sampling_rate=0.5,
             shape_is_real_fourier=shape_is_real_fourier,
             compute_euclidean_norm=True,
         )
-        grid = backend.to_numpy_array(grid)
-        np.multiply(grid, (spectrum.shape[0] - 1), out = grid) + 0.5
+        grid = be.to_numpy_array(grid)
+        np.multiply(grid, (spectrum.shape[0] - 1), out=grid) + 0.5
         spectrum = map_coordinates(spectrum, grid.reshape(1, -1), order=order)
         return spectrum.reshape(grid.shape)
 
     def __call__(
         self,
-        data: NDArray = None,
-        data_rfft: NDArray = None,
+        data: BackendArray = None,
+        data_rfft: BackendArray = None,
         n_bins: int = None,
         batch_dimension: int = None,
         order: int = 1,
@@ -329,9 +332,9 @@ class LinearWhiteningFilter:
 
         Parameters:
         -----------
-        data : NDArray, optional
+        data : BackendArray, optional
             The input data, defaults to None.
-        data_rfft : NDArray, optional
+        data_rfft : BackendArray, optional
             The Fourier transform of the input data, defaults to None.
         n_bins : int, optional
             The number of bins for computing the spectrum, defaults to None.
@@ -348,9 +351,9 @@ class LinearWhiteningFilter:
            Filter data and associated parameters.
         """
         if data_rfft is None:
-            data_rfft = np.fft.rfftn(backend.to_numpy_array(data))
+            data_rfft = np.fft.rfftn(be.to_numpy_array(data))
 
-        data_rfft = backend.to_numpy_array(data_rfft)
+        data_rfft = be.to_numpy_array(data_rfft)
 
         bins, radial_averages = self._compute_spectrum(
             data_rfft, n_bins, batch_dimension
@@ -358,21 +361,28 @@ class LinearWhiteningFilter:
 
         if order is None:
             cutoff = bins < radial_averages.size
-            filter_mask = np.zeros(data_rfft.shape, radial_averages.dtype)
+            filter_mask = np.zeros(bins.shape, radial_averages.dtype)
             filter_mask[cutoff] = radial_averages[bins[cutoff]]
         else:
+            shape = bins.shape
+            if kwargs.get("shape", False):
+                shape = compute_fourier_shape(
+                    shape=kwargs.get("shape"),
+                    shape_is_real_fourier=kwargs.get("shape_is_real_fourier", False),
+                )
+
             filter_mask = self._interpolate_spectrum(
                 spectrum=radial_averages,
-                shape=data_rfft.shape,
+                shape=shape,
                 shape_is_real_fourier=True,
             )
 
         filter_mask = np.fft.ifftshift(
             filter_mask,
-            axes=tuple(i for i in range(data_rfft.ndim - 1) if i != batch_dimension)
+            axes=tuple(i for i in range(data_rfft.ndim - 1) if i != batch_dimension),
         )
 
         return {
-            "data": backend.to_backend_array(filter_mask),
+            "data": be.to_backend_array(filter_mask),
             "is_multiplicative_filter": True,
         }