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

tme/filters/__init__.py

@@ -1,6 +1,6 @@
 from .ctf import CTF, CTFReconstructed
 from .compose import Compose, ComposableFilter
-from .bandpass import BandPassFilter
+from .bandpass import BandPass, BandPassReconstructed
 from .whitening import LinearWhiteningFilter
 from .wedge import Wedge, WedgeReconstructed
-from .reconstruction import ReconstructFromTilt
+from .reconstruction import ReconstructFromTilt, ShiftFourier

tme/filters/_utils.py

@@ -15,6 +15,18 @@ from ..backends import NumpyFFTWBackend
 from ..types import BackendArray, NDArray
 from ..rotations import euler_to_rotationmatrix
 
+__all__ = [
+    "compute_tilt_shape",
+    "centered_grid",
+    "frequency_grid_at_angle",
+    "fftfreqn",
+    "crop_real_fourier",
+    "compute_fourier_shape",
+    "shift_fourier",
+    "create_reconstruction_filter",
+    "pad_to_length",
+]
+
 
 def compute_tilt_shape(shape: Tuple[int], opening_axis: int, reduce_dim: bool = False):
     """
@@ -71,21 +83,27 @@ def frequency_grid_at_angle(
     """
     Generate a frequency grid from 0 to 1/(2 * sampling_rate) in each axis.
 
+    Conceptually, this function generates accurate frequency grid of tilted
+    projections. Given a non-cubical shape, it no longer accurate to compute
+    frequences as Euclidean distances from a centered index grid. This function
+    solves this issue, and makes it possible to create complex filters on
+    non-cubical input shapes.
+
     Parameters
     ----------
-    shape : Tuple[int]
+    shape : tuple of int
         The shape of the grid.
     angle : float
         The angle at which to generate the grid.
-    sampling_rate : Tuple[float]
+    sampling_rate : tuple of float
         The sampling rate for each dimension.
     opening_axis : int, optional
-        The axis to be opened, defaults to None.
+        The projection axis, defaults to None.
     tilt_axis : int, optional
         The axis along which the grid is tilted, defaults to None.
 
-    Returns:
-    --------
+    Returns
+    -------
     NDArray
         The frequency grid.
     """
@@ -231,7 +249,9 @@ def shift_fourier(
 def create_reconstruction_filter(
     filter_shape: Tuple[int], filter_type: str, **kwargs: Dict
 ):
-    """Create a reconstruction filter of given filter_type.
+    """
+    Create a reconstruction filter of given filter_type. The DC component of
+    the filter will be located in the array center.
 
     Parameters
     ----------
@@ -299,7 +319,7 @@ def create_reconstruction_filter(
         ret = fftfreqn((size,), sampling_rate=1, compute_euclidean_norm=True)
         min_increment = np.radians(np.min(np.abs(np.diff(np.sort(tilt_angles)))))
         ret *= min_increment * size
-        np.fmin(ret, 1, out=ret)
+        ret = np.fmin(ret, 1, out=ret)
     elif filter_type == "shepp-logan":
         ret = freq * np.sinc(freq / 2)
     elif filter_type == "cosine":
@@ -310,3 +330,8 @@ def create_reconstruction_filter(
         raise ValueError("Unsupported filter type")
 
     return ret
+
+
+def pad_to_length(arr, length: int):
+    ret = np.atleast_1d(arr)
+    return np.repeat(ret, length // ret.size)

tme/filters/bandpass.py

@@ -8,225 +8,264 @@ Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 
 from typing import Tuple
 from math import log, sqrt
+from dataclasses import dataclass
+
+import numpy as np
 
 from ..types import BackendArray
 from ..backends import backend as be
 from .compose import ComposableFilter
-from ._utils import fftfreqn, crop_real_fourier, shift_fourier
+from ._utils import (
+    crop_real_fourier,
+    shift_fourier,
+    pad_to_length,
+    frequency_grid_at_angle,
+    fftfreqn,
+)
 
-__all__ = ["BandPassFilter"]
+__all__ = ["BandPass", "BandPassReconstructed"]
 
 
-class BandPassFilter(ComposableFilter):
+@dataclass
+class BandPass(ComposableFilter):
     """
-    Generate bandpass filters in Fourier space.
-
-    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 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
-        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.
+    Generate per-slice Fourier Bandpass filter
     """
 
-    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,
-    ) -> BackendArray:
+    #: The tilt angles.
+    angles: Tuple[float]
+    #: The lowpass cutoffs. Either one or one per angle, defaults to None.
+    lowpass: Tuple[float] = None
+    #: The highpass cutoffs. Either one or one per angle, defaults to None.
+    highpass: Tuple[float] = None
+    #: The shape of the to-be created mask.
+    shape: Tuple[int] = None
+    #: Axis the plane is tilted over, defaults to 0 (x).
+    tilt_axis: int = 0
+    #: The projection axis, defaults to 2 (z).
+    opening_axis: int = 2
+    #: The sampling rate, defaults to 1 Ångstrom / voxel.
+    sampling_rate: Tuple[float] = 1
+    #: Whether to use Gaussian bandpass filter, defaults to True.
+    use_gaussian: bool = True
+    #: Whether to return a mask for rfft
+    return_real_fourier: bool = False
+
+    def __call__(self, **kwargs):
         """
-        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
-        -------
-        BackendArray
-            The bandpass filter in Fourier space.
+        Returns a Bandpass stack of chosen parameters with DC component in the center.
         """
+        func_args = vars(self).copy()
+        func_args.update(kwargs)
+
+        func = discrete_bandpass
+        if func_args.get("use_gaussian"):
+            func = gaussian_bandpass
+
+        return_real_fourier = kwargs.get("return_real_fourier", True)
+        shape_is_real_fourier = kwargs.get("shape_is_real_fourier", False)
         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 = be.astype(be.to_backend_array(grid), be._float_dtype)
-        sampling_rate = be.to_backend_array(sampling_rate)
+        angles = np.atleast_1d(func_args["angles"])
+        _lowpass = pad_to_length(func_args["lowpass"], angles.size)
+        _highpass = pad_to_length(func_args["highpass"], angles.size)
+
+        masks = []
+        for index, angle in enumerate(angles):
+            frequency_grid = frequency_grid_at_angle(
+                shape=func_args["shape"],
+                tilt_axis=func_args["tilt_axis"],
+                opening_axis=func_args["opening_axis"],
+                angle=angle,
+                sampling_rate=1,
+            )
+            func_args["lowpass"] = _lowpass[index]
+            func_args["highpass"] = _highpass[index]
+            mask = func(grid=frequency_grid, **func_args)
 
-        highcut = grid.max()
-        if lowpass is not None:
-            highcut = be.max(2 * sampling_rate / lowpass)
+            mask = shift_fourier(data=mask, shape_is_real_fourier=shape_is_real_fourier)
+            if return_real_fourier:
+                mask = crop_real_fourier(mask)
+            masks.append(mask[None])
 
-        lowcut = 0
-        if highpass is not None:
-            lowcut = be.max(2 * sampling_rate / highpass)
+        masks = be.concatenate(masks, axis=0)
+        return {
+            "data": be.to_backend_array(masks),
+            "shape": func_args["shape"],
+            "return_real_fourier": return_real_fourier,
+            "is_multiplicative_filter": True,
+        }
 
-        bandpass_filter = ((grid <= highcut) & (grid >= lowcut)) * 1.0
-        bandpass_filter = shift_fourier(
-            data=bandpass_filter, shape_is_real_fourier=shape_is_real_fourier
-        )
 
-        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,
-    ) -> BackendArray:
-        """
-        Generate a bandpass filter using Gaussians.
-
-        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
-        -------
-        BackendArray
-            The bandpass filter in Fourier space.
-        """
+@dataclass
+class BandPassReconstructed(ComposableFilter):
+    """
+    Generate reconstructed bandpass filters in Fourier space.
+    """
+
+    #: The lowpass cutoff, defaults to None.
+    lowpass: float = None
+    #: The highpass cutoff, defaults to None.
+    highpass: float = None
+    #: The shape of the to-be created mask.
+    shape: Tuple[int] = None
+    #: Axis the plane is tilted over, defaults to 0 (x).
+    tilt_axis: int = 0
+    #: The projection axis, defaults to 2 (z).
+    opening_axis: int = 2
+    #: The sampling rate, defaults to 1 Ångstrom / voxel.
+    sampling_rate: Tuple[float] = 1
+    #: Whether to use Gaussian bandpass filter, defaults to True.
+    use_gaussian: bool = True
+    #: Whether to return a mask for rfft
+    return_real_fourier: bool = False
+
+    def __call__(self, **kwargs):
+        func_args = vars(self).copy()
+        func_args.update(kwargs)
+
+        func = discrete_bandpass
+        if func_args.get("use_gaussian"):
+            func = gaussian_bandpass
+
+        return_real_fourier = func_args.get("return_real_fourier", True)
+        shape_is_real_fourier = func_args.get("shape_is_real_fourier", False)
         if shape_is_real_fourier:
             return_real_fourier = False
 
         grid = fftfreqn(
-            shape=shape,
+            shape=func_args["shape"],
             sampling_rate=0.5,
             shape_is_real_fourier=shape_is_real_fourier,
             compute_euclidean_norm=True,
         )
-        grid = be.astype(be.to_backend_array(grid), be._float_dtype)
-        grid = -be.square(grid, out=grid)
+        mask = func(grid=grid, **func_args)
 
-        has_lowpass, has_highpass = False, False
-        norm = float(sqrt(2 * log(2)))
-        upper_sampling = float(
-            be.max(be.multiply(2, be.to_backend_array(sampling_rate)))
-        )
+        mask = shift_fourier(data=mask, shape_is_real_fourier=shape_is_real_fourier)
+        if return_real_fourier:
+            mask = crop_real_fourier(mask)
 
-        if lowpass is not None:
-            lowpass, has_lowpass = float(lowpass), True
-            lowpass = be.maximum(lowpass, be.eps(be._float_dtype))
-        if highpass is not None:
-            highpass, has_highpass = float(highpass), True
-            highpass = be.maximum(highpass, be.eps(be._float_dtype))
-
-        if has_lowpass:
-            lowpass = upper_sampling / (lowpass * norm)
-            lowpass = be.multiply(2, be.square(lowpass))
-            if not has_highpass:
-                lowpass_filter = be.divide(grid, lowpass, out=grid)
-            else:
-                lowpass_filter = be.divide(grid, lowpass)
-            lowpass_filter = be.exp(lowpass_filter, out=lowpass_filter)
-
-        if has_highpass:
-            highpass = upper_sampling / (highpass * norm)
-            highpass = be.multiply(2, be.square(highpass))
-            highpass_filter = be.divide(grid, highpass, out=grid)
-            highpass_filter = be.exp(highpass_filter, out=highpass_filter)
-            highpass_filter = be.subtract(1, highpass_filter, out=highpass_filter)
-
-        if has_lowpass and not has_highpass:
-            bandpass_filter = lowpass_filter
-        elif not has_lowpass and has_highpass:
-            bandpass_filter = highpass_filter
-        elif has_lowpass and has_highpass:
-            bandpass_filter = be.multiply(
-                lowpass_filter, highpass_filter, out=lowpass_filter
-            )
-        else:
-            bandpass_filter = be.full(shape, fill_value=1, dtype=be._float_dtype)
+        return {
+            "data": be.to_backend_array(mask),
+            "shape": func_args["shape"],
+            "return_real_fourier": return_real_fourier,
+            "is_multiplicative_filter": True,
+        }
 
-        bandpass_filter = shift_fourier(
-            data=bandpass_filter, shape_is_real_fourier=shape_is_real_fourier
-        )
 
-        if return_real_fourier:
-            bandpass_filter = crop_real_fourier(bandpass_filter)
+def discrete_bandpass(
+    grid: BackendArray,
+    lowpass: float,
+    highpass: float,
+    sampling_rate: Tuple[float],
+    **kwargs,
+) -> BackendArray:
+    """
+    Generate a bandpass filter using discrete frequency cutoffs.
 
-        return bandpass_filter
+    Parameters
+    ----------
+    grid : BackendArray
+        Frequencies in Fourier space.
+    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.
+    **kwargs : dict
+        Additional keyword arguments.
+
+    Returns
+    -------
+    BackendArray
+        The bandpass filter in Fourier space.
+    """
+    grid = be.astype(be.to_backend_array(grid), be._float_dtype)
+    sampling_rate = be.to_backend_array(sampling_rate)
 
-    def __call__(self, **kwargs):
-        func_args = vars(self)
-        func_args.update(kwargs)
+    highcut = grid.max()
+    if lowpass is not None:
+        highcut = be.max(2 * sampling_rate / lowpass)
 
-        func = self.discrete_bandpass
-        if func_args.get("use_gaussian"):
-            func = self.gaussian_bandpass
+    lowcut = 0
+    if highpass is not None:
+        lowcut = be.max(2 * sampling_rate / highpass)
 
-        mask = func(**func_args)
+    bandpass_filter = ((grid <= highcut) & (grid >= lowcut)) * 1.0
+    return bandpass_filter
 
-        return {
-            "data": be.to_backend_array(mask),
-            "shape": func_args["shape"],
-            "return_real_fourier": func_args.get("return_real_fourier"),
-            "is_multiplicative_filter": True,
-        }
+
+def gaussian_bandpass(
+    grid: BackendArray,
+    lowpass: float = None,
+    highpass: float = None,
+    sampling_rate: float = 1,
+    **kwargs,
+) -> BackendArray:
+    """
+    Generate a bandpass filter using Gaussians.
+
+    Parameters
+    ----------
+    grid : BackendArray
+        Frequency grid in Fourier space.
+    lowpass : float, optional
+        The lowpass cutoff in units of sampling rate, defaults to None.
+    highpass : float, optional
+        The highpass cutoff in units of sampling rate, defaults to None.
+    sampling_rate : float, optional
+        The sampling rate in Fourier space, defaults to one.
+    **kwargs : dict
+        Additional keyword arguments.
+
+    Returns
+    -------
+    BackendArray
+        The bandpass filter in Fourier space.
+    """
+    grid = be.astype(be.to_backend_array(grid), be._float_dtype)
+    grid = -be.square(grid, out=grid)
+
+    has_lowpass, has_highpass = False, False
+    norm = float(sqrt(2 * log(2)))
+    upper_sampling = float(be.max(be.multiply(2, be.to_backend_array(sampling_rate))))
+
+    if lowpass is not None:
+        lowpass, has_lowpass = float(lowpass), True
+        lowpass = be.maximum(lowpass, be.eps(be._float_dtype))
+    if highpass is not None:
+        highpass, has_highpass = float(highpass), True
+        highpass = be.maximum(highpass, be.eps(be._float_dtype))
+
+    if has_lowpass:
+        lowpass = upper_sampling / (lowpass * norm)
+        lowpass = be.multiply(2, be.square(lowpass))
+        if not has_highpass:
+            lowpass_filter = be.divide(grid, lowpass, out=grid)
+        else:
+            lowpass_filter = be.divide(grid, lowpass)
+        lowpass_filter = be.exp(lowpass_filter, out=lowpass_filter)
+
+    if has_highpass:
+        highpass = upper_sampling / (highpass * norm)
+        highpass = be.multiply(2, be.square(highpass))
+        highpass_filter = be.divide(grid, highpass, out=grid)
+        highpass_filter = be.exp(highpass_filter, out=highpass_filter)
+        highpass_filter = be.subtract(1, highpass_filter, out=highpass_filter)
+
+    if has_lowpass and not has_highpass:
+        bandpass_filter = lowpass_filter
+    elif not has_lowpass and has_highpass:
+        bandpass_filter = highpass_filter
+    elif has_lowpass and has_highpass:
+        bandpass_filter = be.multiply(
+            lowpass_filter, highpass_filter, out=lowpass_filter
+        )
+    else:
+        bandpass_filter = be.full(grid.shape, fill_value=1, dtype=be._float_dtype)
+
+    return bandpass_filter