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

tme/preprocessing/_utils.py

@@ -0,0 +1,176 @@
+""" Utilities for the generation of frequency grids.
+
+    Copyright (c) 2024 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+
+from typing import Tuple
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ..backends import backend
+from ..matching_utils import euler_to_rotationmatrix
+
+
+def compute_tilt_shape(shape: Tuple[int], opening_axis: int, reduce_dim: bool = False):
+    """
+    Given an opening_axis, computes the shape of the remaining dimensions.
+
+    Parameters:
+    -----------
+    shape : Tuple[int]
+        The shape of the input array.
+    opening_axis : int
+        The axis along which the array will be tilted.
+    reduce_dim : bool, optional (default=False)
+        Whether to reduce the dimensionality after tilting.
+
+    Returns:
+    --------
+    Tuple[int]
+        The shape of the array after tilting.
+    """
+    tilt_shape = tuple(x if i != opening_axis else 1 for i, x in enumerate(shape))
+    if reduce_dim:
+        tilt_shape = tuple(x for i, x in enumerate(shape) if i != opening_axis)
+
+    return tilt_shape
+
+
+def centered_grid(shape: Tuple[int]) -> NDArray:
+    """
+    Generate an integer valued grid centered around size // 2
+
+    Parameters:
+    -----------
+    shape : Tuple[int]
+        The shape of the grid.
+
+    Returns:
+    --------
+    NDArray
+        The centered grid.
+    """
+    index_grid = np.array(
+        np.meshgrid(*[np.arange(size) - size // 2 for size in shape], indexing="ij")
+    )
+    return index_grid
+
+
+def frequency_grid_at_angle(
+    shape: Tuple[int],
+    angle: float,
+    sampling_rate: Tuple[float],
+    opening_axis: int = None,
+    tilt_axis: int = None,
+) -> NDArray:
+    """
+    Generate a frequency grid from 0 to 1/(2 * sampling_rate) in each axis.
+
+    Parameters:
+    -----------
+    shape : Tuple[int]
+        The shape of the grid.
+    angle : float
+        The angle at which to generate the grid.
+    sampling_rate : Tuple[float]
+        The sampling rate for each dimension.
+    opening_axis : int, optional
+        The axis to be opened, defaults to None.
+    tilt_axis : int, optional
+        The axis along which the grid is tilted, defaults to None.
+
+    Returns:
+    --------
+    NDArray
+        The frequency grid.
+    """
+    sampling_rate = np.array(sampling_rate)
+    sampling_rate = np.repeat(sampling_rate, len(shape) // sampling_rate.size)
+
+    tilt_shape = compute_tilt_shape(
+        shape=shape, opening_axis=opening_axis, reduce_dim=False
+    )
+    index_grid = centered_grid(shape=tilt_shape)
+    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 = np.einsum("ij,j...->i...", rotation_matrix, index_grid)
+
+    norm = np.divide(1, 2 * sampling_rate * np.divide(shape, 2).astype(int))
+
+    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
+
+
+def fftfreqn(
+    shape: Tuple[int],
+    sampling_rate: Tuple[float],
+    compute_euclidean_norm: bool = False,
+    shape_is_real_fourier: bool = False,
+) -> NDArray:
+    """
+    Generate the n-dimensional discrete Fourier Transform sample frequencies.
+
+    Parameters:
+    -----------
+    shape : Tuple[int]
+        The shape of the data.
+    sampling_rate : float or Tuple[float]
+        The sampling rate.
+    compute_euclidean_norm : bool, optional
+        Whether to compute the Euclidean norm, defaults to False.
+    shape_is_real_fourier : bool, optional
+        Whether the shape corresponds to a real Fourier transform, defaults to False.
+
+    Returns:
+    --------
+    NDArray
+        The sample frequencies.
+    """
+    center = backend.astype(backend.divide(shape, 2), backend._default_dtype_int)
+
+    norm = np.ones(3)
+    if sampling_rate is not None:
+        norm = backend.multiply(shape, sampling_rate).astype(int)
+
+    if shape_is_real_fourier:
+        center[-1] = 0
+        norm[-1] = 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)
+
+    if compute_euclidean_norm:
+        backend.square(indices, indices)
+        indices = backend.sum(indices, axis=0)
+        indices = backend.sqrt(indices)
+
+    return indices
+
+
+def crop_real_fourier(data: NDArray) -> NDArray:
+    """
+    Crop the real part of a Fourier transform.
+
+    Parameters:
+    -----------
+    data : NDArray
+        The Fourier transformed data.
+
+    Returns:
+    --------
+    NDArray
+        The cropped data.
+    """
+    stop = 1 + (data.shape[-1] // 2)
+    return data[..., :stop]

tme/preprocessing/composable_filter.py

@@ -0,0 +1,30 @@
+""" Defines a specification for filters that can be used with
+    :py:class:`tme.preprocessing.compose.Compose`.
+
+    Copyright (c) 2024 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+from typing import Dict
+from abc import ABC, abstractmethod
+
+class ComposableFilter(ABC):
+    """
+    Strategy class for composable filters.
+    """
+
+    @abstractmethod
+    def __call__(self, *args, **kwargs) -> Dict:
+        """
+        Parameters:
+        -----------
+        *args : tuple
+            Variable length argument list.
+        **kwargs : dict
+            Arbitrary keyword arguments.
+
+        Returns:
+        --------
+        Dict
+            A dictionary representing the result of the filtering operation.
+        """

tme/preprocessing/compose.py

@@ -0,0 +1,52 @@
+""" Combine filters using an interface analogous to pytorch's Compose.
+
+    Copyright (c) 2024 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+
+from typing import Tuple, Dict
+
+from tme.backends import backend
+
+
+class Compose:
+    """
+    Compose a series of transformations.
+
+    This class allows composing multiple transformations together. Each transformation
+    is expected to be a callable that accepts keyword arguments and returns metadata.
+
+    Parameters:
+    -----------
+    transforms : Tuple[object]
+        A tuple containing transformation objects.
+
+    Returns:
+    --------
+    Dict
+        Metadata resulting from the composed transformations.
+
+    """
+
+    def __init__(self, transforms: Tuple[object]):
+        self.transforms = transforms
+
+    def __call__(self, **kwargs: Dict) -> Dict:
+        meta = {}
+        if not len(self.transforms):
+            return meta
+
+        meta = self.transforms[0](**kwargs)
+        for transform in self.transforms[1:]:
+
+            kwargs.update(meta)
+            ret = transform(**kwargs)
+
+            if ret.get("is_multiplicative_filter", False):
+                backend.multiply(ret["data"], meta["data"], ret["data"])
+                ret["merge"] = None
+
+            meta = ret
+
+        return meta

tme/preprocessing/frequency_filters.py

@@ -0,0 +1,322 @@
+""" Defines Fourier frequency filters.
+
+    Copyright (c) 2024 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+from typing import Tuple, Dict
+
+import numpy as np
+from numpy.typing import NDArray
+from scipy.ndimage import mean as ndimean
+
+from ._utils import fftfreqn, crop_real_fourier
+from ..backends import backend
+
+
+class BandPassFilter:
+    """
+    This class provides methods to generate bandpass filters in Fourier space,
+    either by directly specifying the frequency cutoffs (discrete_bandpass) or
+    by using Gaussian functions (gaussian_bandpass).
+
+    Parameters:
+    -----------
+    lowpass : float, optional
+        The lowpass cutoff, defaults to None.
+    highpass : float, optional
+        The highpass cutoff, defaults to None.
+    sampling_rate : Tuple[float], optional
+        The sampling rate in Fourier space, defaults to 1.
+    use_gaussian : bool, optional
+        Whether to use Gaussian bandpass filter, defaults to True.
+    return_real_fourier : bool, optional
+        Whether to return only the real Fourier space, defaults to False.
+    shape_is_real_fourier : bool, optional
+        Whether the shape represents the real Fourier space, defaults to False.
+    """
+
+    def __init__(
+        self,
+        lowpass: float = None,
+        highpass: float = None,
+        sampling_rate: Tuple[float] = 1,
+        use_gaussian: bool = True,
+        return_real_fourier: bool = False,
+        shape_is_real_fourier: bool = False,
+    ):
+        self.lowpass = lowpass
+        self.highpass = highpass
+        self.use_gaussian = use_gaussian
+        self.return_real_fourier = return_real_fourier
+        self.shape_is_real_fourier = shape_is_real_fourier
+        self.sampling_rate = sampling_rate
+
+    @staticmethod
+    def discrete_bandpass(
+        shape: Tuple[int],
+        lowpass: float,
+        highpass: float,
+        sampling_rate: Tuple[float],
+        return_real_fourier: bool = False,
+        shape_is_real_fourier: bool = False,
+        **kwargs,
+    ) -> NDArray:
+        """
+        Generate a bandpass filter using discrete frequency cutoffs.
+
+        Parameters:
+        -----------
+        shape : tuple of int
+            The shape of the bandpass filter.
+        lowpass : float
+            The lowpass cutoff in units of sampling rate.
+        highpass : float
+            The highpass cutoff in units of sampling rate.
+        return_real_fourier : bool, optional
+            Whether to return only the real Fourier space, defaults to False.
+        sampling_rate : float
+            The sampling rate in Fourier space.
+        shape_is_real_fourier : bool, optional
+            Whether the shape represents the real Fourier space, defaults to False.
+        **kwargs : dict
+            Additional keyword arguments.
+
+        Returns:
+        --------
+        NDArray
+            The bandpass filter in Fourier space.
+        """
+        grid = fftfreqn(
+            shape=shape,
+            sampling_rate=0.5,
+            shape_is_real_fourier=shape_is_real_fourier,
+            compute_euclidean_norm=True,
+        )
+
+        lowpass = 0 if lowpass is None else lowpass
+        highpass = 1e10 if highpass is None else highpass
+
+        highcut = grid.max()
+        if lowpass > 0:
+            highcut = np.max(2 * sampling_rate / lowpass)
+        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)))
+        )
+
+        if return_real_fourier:
+            bandpass_filter = crop_real_fourier(bandpass_filter)
+
+        return bandpass_filter
+
+    @staticmethod
+    def gaussian_bandpass(
+        shape: Tuple[int],
+        lowpass: float,
+        highpass: float,
+        sampling_rate: float,
+        return_real_fourier: bool = False,
+        shape_is_real_fourier: bool = False,
+        **kwargs,
+    ) -> NDArray:
+        """
+        Generate a bandpass filter using Gaussian functions.
+
+        Parameters:
+        -----------
+        shape : tuple of int
+            The shape of the bandpass filter.
+        lowpass : float
+            The lowpass cutoff in units of sampling rate.
+        highpass : float
+            The highpass cutoff in units of sampling rate.
+        sampling_rate : float
+            The sampling rate in Fourier space.
+        return_real_fourier : bool, optional
+            Whether to return only the real Fourier space, defaults to False.
+        shape_is_real_fourier : bool, optional
+            Whether the shape represents the real Fourier space, defaults to False.
+        **kwargs : dict
+            Additional keyword arguments.
+
+        Returns:
+        --------
+        NDArray
+            The bandpass filter in Fourier space.
+        """
+        if shape_is_real_fourier:
+            return_real_fourier = False
+
+        grid = fftfreqn(
+            shape=shape,
+            sampling_rate=0.5,
+            shape_is_real_fourier=shape_is_real_fourier,
+            compute_euclidean_norm=True,
+        )
+        grid = -backend.square(grid)
+
+        lowpass_filter, highpass_filter = 1, 1
+        norm = float(backend.sqrt(2 * backend.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))
+        if highpass is not None:
+            highpass = float(highpass)
+            highpass = backend.maximum(highpass, backend.eps(highpass))
+
+        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))
+        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))
+
+        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)))
+        )
+
+        if return_real_fourier:
+            lowpass_filter = crop_real_fourier(lowpass_filter)
+
+        return lowpass_filter
+
+    def __call__(self, **kwargs):
+        func_args = vars(self)
+        func_args.update(kwargs)
+
+        func = self.discrete_bandpass
+        if func_args.get("use_gaussian"):
+            func = self.gaussian_bandpass
+
+        mask = func(**func_args)
+
+        return {
+            "data": backend.to_backend_array(mask),
+            "sampling_rate": func_args.get("sampling_rate", 1),
+            "is_multiplicative_filter": True,
+        }
+
+
+class LinearWhiteningFilter:
+    """
+    This class provides methods to compute the spectrum of the input data and
+    apply linear whitening to the Fourier coefficients.
+
+    Parameters:
+    -----------
+    **kwargs : Dict, optional
+        Additional keyword arguments.
+    """
+
+    def __init__(self, **kwargs):
+        pass
+
+    @staticmethod
+    def _compute_spectrum(
+        data_rfft: NDArray, n_bins: int = None
+    ) -> Tuple[NDArray, NDArray]:
+        """
+        Compute the spectrum of the input data.
+
+        Parameters:
+        -----------
+        data_rfft : NDArray
+            The Fourier transform of the input data.
+        n_bins : int, optional
+            The number of bins for computing the spectrum, defaults to None.
+
+        Returns:
+        --------
+        bins : NDArray
+            Array containing the bin indices for the spectrum.
+        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])
+        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,
+            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)
+
+        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))
+
+        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
+
+    def __call__(
+        self,
+        data: NDArray = None,
+        data_rfft: NDArray = None,
+        n_bins: int = None,
+        **kwargs: Dict,
+    ) -> Dict:
+        """
+        Apply linear whitening to the data and return the result.
+
+        Parameters:
+        -----------
+        data : NDArray, optional
+            The input data, defaults to None.
+        data_rfft : NDArray, 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.
+        **kwargs : Dict
+            Additional keyword arguments.
+
+        Returns:
+        --------
+        Dict
+            A dictionary containing the whitened data and information
+            about the filter being a multiplicative filter.
+        """
+        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)
+
+        radial_averages = np.fft.ifftshift(
+            radial_averages[bins], axes=tuple(range(data_rfft.ndim - 1))
+        )
+
+        return {
+            "data": backend.to_backend_array(radial_averages),
+            "is_multiplicative_filter": True,
+        }