pytme 0.3.1.post2__cp311-cp311-macosx_15_0_arm64.whl → 0.3.2.dev0__cp311-cp311-macosx_15_0_arm64.whl

tme/filters/whitening.py

@@ -1,5 +1,5 @@
 """
-Implements class BandPassFilter to create Fourier filter representations.
+Implements class LinearWhiteningFilter
 
 Copyright (c) 2024 European Molecular Biology Laboratory
 
@@ -7,22 +7,26 @@ Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 from typing import Tuple, Dict
+from dataclasses import dataclass
 
 import numpy as np
 from scipy.ndimage import mean as ndimean
 from scipy.ndimage import map_coordinates
 
+from ._utils import fftfreqn
 from ..types import BackendArray
+from ..analyzer.peaks import batchify
 from ..backends import backend as be
 from .compose import ComposableFilter
-from ._utils import fftfreqn, compute_fourier_shape, shift_fourier
+
 
 __all__ = ["LinearWhiteningFilter"]
 
 
+@dataclass
 class LinearWhiteningFilter(ComposableFilter):
     """
-    Compute Fourier power spectrums and perform whitening.
+    Generate Fourier whitening filters.
 
     References
     ----------
@@ -34,12 +38,9 @@ class LinearWhiteningFilter(ComposableFilter):
         13375 (2023)
     """
 
-    def __init__(self, *args, **kwargs):
-        pass
-
     @staticmethod
     def _compute_spectrum(
-        data_rfft: BackendArray, n_bins: int = None, batch_dimension: int = None
+        data_rfft: BackendArray, n_bins: int = None
     ) -> Tuple[BackendArray, BackendArray]:
         """
         Compute the power spectrum of the input data.
@@ -50,8 +51,6 @@ class LinearWhiteningFilter(ComposableFilter):
             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
         -------
@@ -60,7 +59,7 @@ class LinearWhiteningFilter(ComposableFilter):
         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)
+        shape = data_rfft.shape
 
         max_bins = max(max(shape[:-1]) // 2 + 1, shape[-1])
         n_bins = max_bins if n_bins is None else n_bins
@@ -71,25 +70,22 @@ class LinearWhiteningFilter(ComposableFilter):
             sampling_rate=0.5,
             shape_is_real_fourier=True,
             compute_euclidean_norm=True,
+            fftshift=False,
         )
         bins = be.to_numpy_array(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)
+
+        fourier_spectrum = np.abs(data_rfft)
+        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)
-
+        radial_averages = np.sqrt(radial_averages, out=radial_averages)
+        radial_averages = np.where(radial_averages != 0, 1 / radial_averages, 0)
+        norm_factor = radial_averages.max()
+        if norm_factor != 0:
+            radial_averages = np.divide(radial_averages, norm_factor)
         return bins, radial_averages
 
     @staticmethod
@@ -104,21 +100,19 @@ class LinearWhiteningFilter(ComposableFilter):
             sampling_rate=0.5,
             shape_is_real_fourier=shape_is_real_fourier,
             compute_euclidean_norm=True,
+            fftshift=False,
         )
         grid = be.to_numpy_array(grid)
-        np.multiply(grid, (spectrum.shape[0] - 1), out=grid) + 0.5
+        grid = np.floor(np.multiply(grid, spectrum.shape[0] - 1) + 0.5)
         spectrum = map_coordinates(spectrum, grid.reshape(1, -1), order=order)
         return spectrum.reshape(grid.shape)
 
-    def __call__(
+    def _evaluate(
         self,
-        shape: Tuple[int],
-        data: BackendArray = None,
-        data_rfft: BackendArray = None,
-        n_bins: int = None,
-        batch_dimension: int = None,
+        shape: Tuple[int, ...],
+        data_rfft: BackendArray,
+        axes: Tuple[int] = (),
         order: int = 1,
-        return_real_fourier: bool = True,
         **kwargs: Dict,
     ) -> Dict:
         """
@@ -128,59 +122,29 @@ class LinearWhiteningFilter(ComposableFilter):
         ----------
         shape : tuple of ints
             Shape of the returned whitening filter.
-        data : BackendArray, optional
-            The input data, defaults to None.
         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.
-        batch_dimension : int, optional
-            Batch dimension to average over.
-        return_real_fourier : tuple of int
-            Return a shape compliant with rfft, i.e., omit the negative frequencies
-            terms resulting in a return shape (*shape[:-1], shape[-1]//2+1)
+        axes : tuple of ints, optional
+            Axes to compute spectrum for independently.
         **kwargs : Dict
             Additional keyword arguments.
-
-        Returns
-        -------
-        dict
-            data: BackendArray
-                The filter mask.
-            shape: tuple of ints
-                The requested filter shape
-            return_real_fourier: bool
-                Whether data is compliant with rfftn.
-            is_multiplicative_filter: bool
-                Whether the filter is multiplicative in Fourier space.
         """
-        if data_rfft is None:
-            data_rfft = be.rfftn(data)
+        if isinstance(axes, int):
+            axes = (axes,)
 
+        stack = []
         data_rfft = be.to_numpy_array(data_rfft)
-        bins, radial_averages = self._compute_spectrum(
-            data_rfft, n_bins, batch_dimension
-        )
-        shape = tuple(int(x) for i, x in enumerate(shape) if i != batch_dimension)
-
-        shape_filter = shape
-        if return_real_fourier:
-            shape_filter = compute_fourier_shape(
+        for subset, _ in batchify(data_rfft.shape, axes):
+            _, radial_avg = self._compute_spectrum(np.squeeze(data_rfft[subset]))
+            ret = self._interpolate_spectrum(
+                spectrum=radial_avg,
                 shape=shape,
                 shape_is_real_fourier=False,
+                order=order,
             )
+            stack.append(ret)
 
-        ret = self._interpolate_spectrum(
-            spectrum=radial_averages,
-            shape=shape_filter,
-            shape_is_real_fourier=return_real_fourier,
-        )
-
-        ret = shift_fourier(data=ret, shape_is_real_fourier=return_real_fourier)
-
-        return {
-            "data": be.to_backend_array(ret),
-            "shape": shape,
-            "return_real_fourier": return_real_fourier,
-            "is_multiplicative_filter": True,
-        }
+        ret = np.array(stack)
+        if not len(axes):
+            ret = np.squeeze(ret)
+        return {"data": be.to_backend_array(ret), "shape": shape}

tme/mask.py

@@ -11,7 +11,7 @@ from typing import Tuple, Optional
 
 from .types import NDArray
 from scipy.ndimage import gaussian_filter
-from .matching_utils import rigid_transform
+from .matching_utils import _rigid_transform
 
 __all__ = ["elliptical_mask", "tube_mask", "box_mask", "membrane_mask"]
 
@@ -76,7 +76,7 @@ def elliptical_mask(
     if orientation is not None:
         return_shape = indices.shape
         indices = indices.reshape(n, -1)
-        rigid_transform(
+        _rigid_transform(
             coordinates=indices,
             rotation_matrix=np.asarray(orientation),
             out=indices,

tme/matching_data.py

@@ -15,7 +15,7 @@ from . import Density
 from .filters import Compose
 from .backends import backend as be
 from .types import BackendArray, NDArray
-from .matching_utils import compute_parallelization_schedule
+from .matching_utils import compute_parallelization_schedule, copy_docstring
 
 __all__ = ["MatchingData"]
 
@@ -249,8 +249,8 @@ class MatchingData:
         target_offset = np.zeros(len(self._output_target_shape), dtype=int)
         target_offset[mask] = [x.start for x in target_slice]
         mask = np.subtract(1, self._target_batch).astype(bool)
-        template_offset = np.zeros(len(self._output_template_shape), dtype=int)
-        template_offset[mask] = [x.start for x in template_slice]
+        # template_offset = np.zeros(len(self._output_template_shape), dtype=int)
+        # template_offset[mask] = [x.start for x in template_slice]
 
         translation_offset = tuple(x for x in target_offset)
 
@@ -485,14 +485,18 @@ class MatchingData:
 
     @staticmethod
     def _fourier_padding(
-        target_shape: Tuple[int],
-        template_shape: Tuple[int],
-        batch_mask: Tuple[int] = None,
+        target_shape: NDArray,
+        template_shape: NDArray,
+        target_batch: NDArray = None,
+        template_batch: NDArray = None,
         **kwargs,
     ) -> Tuple[Tuple, Tuple, Tuple, Tuple]:
-        if batch_mask is None:
-            batch_mask = np.zeros_like(template_shape)
-        batch_mask = np.asarray(batch_mask)
+        if target_batch is None:
+            target_batch = np.zeros_like(target_shape)
+        if template_batch is None:
+            template_batch = np.zeros_like(target_shape)
+
+        batch_mask = np.logical_or(target_batch, template_batch)
 
         fourier_pad = np.ones(len(template_shape), dtype=int)
         fourier_pad = np.multiply(fourier_pad, 1 - batch_mask)
@@ -500,7 +504,9 @@ class MatchingData:
 
         # Avoid padding batch dimensions
         pad_shape = np.maximum(target_shape, template_shape)
-        pad_shape = np.maximum(pad_shape, np.multiply(1 - batch_mask, pad_shape))
+        pad_shape = np.where(target_batch, target_shape, pad_shape)
+        pad_shape = np.where(template_batch, template_shape, pad_shape)
+
         ret = be.compute_convolution_shapes(pad_shape, fourier_pad)
         conv_shape, fast_shape, fast_ft_shape = ret
 
@@ -538,10 +544,15 @@ class MatchingData:
         --------
         >>> conv, fwd, inv, shift = matching_data.fourier_padding(pad_fourier=True)
         """
+        target_shape = kwargs.get("target_shape", self._output_target_shape)
+        template_shape = kwargs.get("template_shape", self._output_template_shape)
+        target_batch = kwargs.get("target_batch", self._target_batch)
+        template_batch = kwargs.get("template_batch", self._template_batch)
         return self._fourier_padding(
-            target_shape=be.to_numpy_array(self._output_target_shape),
-            template_shape=be.to_numpy_array(self._output_template_shape),
-            batch_mask=be.to_numpy_array(self._batch_mask),
+            target_shape=be.to_numpy_array(target_shape),
+            template_shape=be.to_numpy_array(template_shape),
+            target_batch=be.to_numpy_array(target_batch),
+            template_batch=be.to_numpy_array(template_batch),
         )
 
     def _score_mask(self, fast_shape: Tuple[int], shift: Tuple[int]) -> BackendArray:
@@ -568,6 +579,68 @@ class MatchingData:
         )
         return be.to_backend_array(score_mask)
 
+    def _transform_data(
+        self, method: str, data: BackendArray, batch_mask: Tuple[int], **kwargs
+    ) -> BackendArray:
+        """
+        Transform data using the specified method.
+
+        Parameters
+        ----------
+        method : str, optional
+            Transformation method, default "phase_randomization".
+            - "phase_randomization": Scrambles phase while preserving amplitude spectrum
+            - "standardize": Standardize to zero mean and unit variance
+            - "laplace": Applies Laplacian edge detection filter
+        **kwargs : dict
+            Method-specific arguments (e.g., mode="wrap" for laplace).
+
+        Returns
+        -------
+        BackendArray
+            Transformed data.
+        """
+        from scipy.ndimage import laplace
+        from .matching_utils import scramble_phases, standardize
+
+        def _standardize(arr: NDArray, **kwargs) -> NDArray:
+            return standardize(arr, 1, arr.size)
+
+        _supported_methods = {
+            "phase_randomization": scramble_phases,
+            "laplace": laplace,
+            "standardize": _standardize,
+        }
+        func = _supported_methods.get(method)
+        if func is None:
+            _supported = ",".join([str(x) for x in _supported_methods])
+            raise ValueError(f"Only methods {_supported} are supported.")
+
+        data = be.to_numpy_array(data)
+
+        ret = np.zeros_like(data)
+        for subset in self._batch_iter(data.shape, batch_mask):
+            ret[subset] = func(data[subset], **kwargs)
+        return be.to_backend_array(ret)
+
+    @copy_docstring(_transform_data)
+    def transform_target(self, method: str = "phase_randomization", **kwargs):
+        return self._transform_data(method, self.target, self._target_batch, **kwargs)
+
+    @copy_docstring(_transform_data)
+    def transform_template(
+        self, method: str = "phase_randomization", reverse: bool = False, **kwargs
+    ):
+        """
+        Notes
+        -----
+        The returned template is in the original not reversed orientation.
+        """
+        template = self._get_data(
+            self._template, self._output_template_shape, reverse, self._template_dim
+        )
+        return self._transform_data(method, template, self._template_batch, **kwargs)
+
     def computation_schedule(
         self,
         matching_method: str = "FLCSphericalMask",
@@ -723,7 +796,6 @@ class MatchingData:
         _output_shape = self._output_template_shape
         if np.prod([int(i) for i in template_mask.shape]) != np.prod(_output_shape):
             _output_shape = self._batch_shape(_output_shape, self._template_batch, True)
-
         return self._get_data(template_mask, _output_shape, True, self._template_dim)
 
     @target.setter
@@ -855,7 +927,7 @@ class MatchingData:
             rot_list.append(self.rotations[init_rot:end_rot])
         return rot_list
 
-    def _free_data(self):
+    def free(self):
         """
         Dereference data arrays owned by the class instance.
         """