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

tme/preprocessing/frequency_filters.py

@@ -4,13 +4,15 @@
 
     Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
+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
+from ._utils import fftfreqn, crop_real_fourier, shift_fourier
 from ..backends import backend
 
 
@@ -87,6 +89,9 @@ class BandPassFilter:
         NDArray
             The bandpass filter in Fourier space.
         """
+        if shape_is_real_fourier:
+            return_real_fourier = False
+
         grid = fftfreqn(
             shape=shape,
             sampling_rate=0.5,
@@ -103,15 +108,9 @@ class BandPassFilter:
         lowcut = np.max(2 * sampling_rate / highpass)
 
         bandpass_filter = ((grid <= highcut) & (grid >= lowcut)) * 1.0
-        shift = backend.add(
-            backend.astype(backend.divide(bandpass_filter.shape, 2), int),
-            backend.mod(bandpass_filter.shape, 2),
-        )
-        if shape_is_real_fourier:
-            shift[-1] = 0
 
-        bandpass_filter = backend.roll(
-            bandpass_filter, shift, tuple(i for i in range(len(shift)))
+        bandpass_filter = shift_fourier(
+            data=bandpass_filter, shape_is_real_fourier=shape_is_real_fourier
         )
 
         if return_real_fourier:
@@ -166,15 +165,15 @@ class BandPassFilter:
         grid = -backend.square(grid)
 
         lowpass_filter, highpass_filter = 1, 1
-        norm = float(backend.sqrt(2 * backend.log(2)))
+        norm = float(sqrt(2 * log(2)))
         upper_sampling = float(backend.max(backend.multiply(2, sampling_rate)))
 
         if lowpass is not None:
             lowpass = float(lowpass)
-            lowpass = backend.maximum(lowpass, backend.eps(lowpass))
+            lowpass = backend.maximum(lowpass, backend.eps(backend._float_dtype))
         if highpass is not None:
             highpass = float(highpass)
-            highpass = backend.maximum(highpass, backend.eps(highpass))
+            highpass = backend.maximum(highpass, backend.eps(backend._float_dtype))
 
         if lowpass is not None:
             lowpass = upper_sampling / (lowpass * norm)
@@ -185,22 +184,15 @@ class BandPassFilter:
             highpass = backend.multiply(2, backend.square(highpass))
             highpass_filter = 1 - backend.exp(backend.divide(grid, highpass))
 
-        lowpass_filter = backend.multiply(lowpass_filter, highpass_filter)
-        shift = backend.add(
-            backend.astype(backend.divide(lowpass_filter.shape, 2), int),
-            backend.mod(lowpass_filter.shape, 2),
-        )
-        if shape_is_real_fourier:
-            shift[-1] = 0
-
-        lowpass_filter = backend.roll(
-            lowpass_filter, shift, tuple(i for i in range(len(shift)))
+        bandpass_filter = backend.multiply(lowpass_filter, highpass_filter)
+        bandpass_filter = shift_fourier(
+            data=bandpass_filter, shape_is_real_fourier=shape_is_real_fourier
         )
 
         if return_real_fourier:
-            lowpass_filter = crop_real_fourier(lowpass_filter)
+            bandpass_filter = crop_real_fourier(bandpass_filter)
 
-        return lowpass_filter
+        return bandpass_filter
 
     def __call__(self, **kwargs):
         func_args = vars(self)
@@ -228,6 +220,16 @@ class LinearWhiteningFilter:
     -----------
     **kwargs : Dict, optional
         Additional keyword arguments.
+
+
+    References
+    ----------
+    .. [1] de Teresa-Trueba, I.; Goetz, S. K.; Mattausch, A.; Stojanovska, F.; Zimmerli, C. E.;
+        Toro-Nahuelpan, M.; Cheng, D. W. C.; Tollervey, F.; Pape, C.; Beck, M.; Diz-Munoz,
+        A.; Kreshuk, A.; Mahamid, J.; Zaugg, J. B. Nat. Methods 2023, 20, 284–294.
+    .. [2]  M. L. Chaillet, G. van der Schot, I. Gubins, S. Roet,
+        R. C. Veltkamp, and F. Förster, Int. J. Mol. Sci. 24,
+        13375 (2023)
     """
 
     def __init__(self, **kwargs):
@@ -235,7 +237,7 @@ class LinearWhiteningFilter:
 
     @staticmethod
     def _compute_spectrum(
-        data_rfft: NDArray, n_bins: int = None
+        data_rfft: NDArray, n_bins: int = None, batch_dimension: int = None
     ) -> Tuple[NDArray, NDArray]:
         """
         Compute the spectrum of the input data.
@@ -246,6 +248,8 @@ class LinearWhiteningFilter:
             The Fourier transform of the input data.
         n_bins : int, optional
             The number of bins for computing the spectrum, defaults to None.
+        batch_dimension : int, optional
+            Batch dimension to average over.
 
         Returns:
         --------
@@ -254,35 +258,70 @@ class LinearWhiteningFilter:
         radial_averages : NDArray
             Array containing the radial averages of the spectrum.
         """
-        max_bins = max(max(data_rfft.shape[:-1]) // 2 + 1, data_rfft.shape[-1])
+        shape = tuple(x for i, x in enumerate(data_rfft.shape) if i != batch_dimension)
+
+        max_bins = max(max(shape[:-1]) // 2 + 1, shape[-1])
         n_bins = max_bins if n_bins is None else n_bins
         n_bins = int(min(n_bins, max_bins))
 
-        grid = fftfreqn(
-            shape=data_rfft.shape,
-            sampling_rate=None,
+        bins = fftfreqn(
+            shape=shape,
+            sampling_rate=0.5,
             shape_is_real_fourier=True,
             compute_euclidean_norm=True,
         )
-        _, bin_edges = np.histogram(grid, bins=n_bins - 1)
-        bins = np.digitize(grid, bins=bin_edges, right=True)
+        bins = backend.to_numpy_array(bins)
 
-        fft_shift_axes = tuple(range(data_rfft.ndim - 1))
-        fourier_transform = np.fft.fftshift(data_rfft, axes=fft_shift_axes)
-        fourier_spectrum = np.square(np.abs(fourier_transform))
-        radial_averages = ndimean(fourier_spectrum, labels=bins, index=np.unique(bins))
+        # Implicit lowpass to nyquist
+        bins = np.floor(bins * (n_bins - 1) + 0.5).astype(int)
+        fft_shift_axes = tuple(
+            i for i in range(data_rfft.ndim - 1) if i != batch_dimension
+        )
+        fourier_spectrum = np.fft.fftshift(data_rfft, axes=fft_shift_axes)
+        fourier_spectrum = np.abs(fourier_spectrum)
+        np.square(fourier_spectrum, out=fourier_spectrum)
 
+        radial_averages = ndimean(
+            fourier_spectrum, labels=bins, index=np.arange(n_bins)
+        )
         np.sqrt(radial_averages, out=radial_averages)
         np.reciprocal(radial_averages, out=radial_averages)
         np.divide(radial_averages, radial_averages.max(), out=radial_averages)
 
         return bins, radial_averages
 
+    @staticmethod
+    def _interpolate_spectrum(
+        spectrum: NDArray,
+        shape: Tuple[int],
+        shape_is_real_fourier: bool = True,
+        order: int = 1,
+    ) -> NDArray:
+        """
+        References
+        ----------
+        .. [1]  M. L. Chaillet, G. van der Schot, I. Gubins, S. Roet,
+            R. C. Veltkamp, and F. Förster, Int. J. Mol. Sci. 24,
+            13375 (2023)
+        """
+        grid = fftfreqn(
+            shape=shape,
+            sampling_rate=.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
+        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,
         n_bins: int = None,
+        batch_dimension: int = None,
+        order: int = 1,
         **kwargs: Dict,
     ) -> Dict:
         """
@@ -296,27 +335,44 @@ class LinearWhiteningFilter:
             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.
+        batch_dimension : int, optional
+            Batch dimension to average over.
+        order : int, optional
+            Interpolation order to use.
         **kwargs : Dict
             Additional keyword arguments.
 
         Returns:
         --------
         Dict
-            A dictionary containing the whitened data and information
-            about the filter being a multiplicative filter.
+           Filter data and associated parameters.
         """
         if data_rfft is None:
             data_rfft = np.fft.rfftn(backend.to_numpy_array(data))
 
         data_rfft = backend.to_numpy_array(data_rfft)
 
-        bins, radial_averages = self._compute_spectrum(data_rfft, n_bins)
+        bins, radial_averages = self._compute_spectrum(
+            data_rfft, n_bins, batch_dimension
+        )
 
-        radial_averages = np.fft.ifftshift(
-            radial_averages[bins], axes=tuple(range(data_rfft.ndim - 1))
+        if order is None:
+            cutoff = bins < radial_averages.size
+            filter_mask = np.zeros(data_rfft.shape, radial_averages.dtype)
+            filter_mask[cutoff] = radial_averages[bins[cutoff]]
+        else:
+            filter_mask = self._interpolate_spectrum(
+                spectrum=radial_averages,
+                shape=data_rfft.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)
         )
 
         return {
-            "data": backend.to_backend_array(radial_averages),
+            "data": backend.to_backend_array(filter_mask),
             "is_multiplicative_filter": True,
         }