pytme 0.3.1.post1__cp311-cp311-macosx_15_0_arm64.whl → 0.3.2__cp311-cp311-macosx_15_0_arm64.whl

tme/filters/bandpass.py

@@ -1,157 +1,157 @@
 """
-Implements class BandPassFilter to create Fourier filter representations.
+Implements class BandPass and BandPassReconstructed.
 
 Copyright (c) 2024 European Molecular Biology Laboratory
 
 Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
-from typing import Tuple
 from math import log, sqrt
-from dataclasses import dataclass
+from typing import Tuple, Union, Optional
+from pydantic.dataclasses import dataclass
 
 import numpy as np
 
 from ..types import BackendArray
 from ..backends import backend as be
 from .compose import ComposableFilter
-from ._utils import (
-    crop_real_fourier,
-    shift_fourier,
-    pad_to_length,
-    frequency_grid_at_angle,
-    fftfreqn,
-)
+from ._utils import pad_to_length, frequency_grid_at_angle, fftfreqn
 
 __all__ = ["BandPass", "BandPassReconstructed"]
 
 
-@dataclass
+@dataclass(config=dict(extra="allow"))
 class BandPass(ComposableFilter):
     """
-    Generate per-slice Fourier Bandpass filter
+    Generate per-tilt Bandpass filter.
+
+    Examples
+    --------
+
+    The following creates an instance of :py:class:`BandPass`
+
+    >>> from tme.filters import BandPass
+    >>> bpf_instance = BandPass(
+    >>>     angles=(-70, 0, 30),
+    >>>     lowpass=30,
+    >>>     sampling_rate=10
+    >>> )
+
+    Differently from :py:class:`tme.filters.BandPassReconstructed`, the filter
+    masks are intended to be used in subsequent reconstruction using
+    :py:class:`tme.filters.ReconstructFromTilt`.
+
+    The ``opening_axis``, ``tilt_axis`` and ``angles`` parameter are used
+    to determine the correct frequencies for non-cubical input shapes. The
+    ``shape`` argument contains the shape of the reconstruction.
+
+    >>> ret = bpf_instance(shape=(50,50,25))
+    >>> mask = ret["data"]
+    >>> mask.shape # 3, 50, 50
+
+    Note that different from its reconstructed counterpart, the DC
+    component is at the center of the array.
+
+    >>> import matplotlib.pyplot as plt
+    >>> fix, ax = plt.subplots(nrows=1, ncols=3)
+    >>> _ = [ax[i].imshow(mask[i]) for i in range(mask.shape[0])]
+    >>> plt.show()
+
     """
 
-    #: 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
+    #: The lowpass cutoff, defaults to None.
+    lowpass: Optional[float] = None
+    #: The highpass cutoff, defaults to None.
+    highpass: Optional[float] = None
+    #: The sampling rate, defaults to 1 Ångstrom / voxel.
+    sampling_rate: Union[Tuple[float, ...], float] = 1
+    #: Whether to use Gaussian bandpass filter, defaults to True.
+    use_gaussian: bool = True
+    #: The tilt angles in degrees.
+    angles: Tuple[float, ...] = 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):
+    def _evaluate(self, shape: Tuple[int, ...], **kwargs):
         """
         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"):
+        if kwargs.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
-
-        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)
+        angles = np.atleast_1d(kwargs["angles"])
+        _lowpass = pad_to_length(kwargs["lowpass"], angles.size)
+        _highpass = pad_to_length(kwargs["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"],
+                shape=shape,
+                tilt_axis=kwargs["tilt_axis"],
+                opening_axis=kwargs["opening_axis"],
                 angle=angle,
                 sampling_rate=1,
             )
-            func_args["lowpass"] = _lowpass[index]
-            func_args["highpass"] = _highpass[index]
-            mask = func(grid=frequency_grid, **func_args)
-
-            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])
-
-        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,
-        }
+            kwargs["lowpass"] = _lowpass[index]
+            kwargs["highpass"] = _highpass[index]
+            mask = func(grid=frequency_grid, **kwargs)
+            masks.append(be.to_backend_array(mask[None]))
+        return {"data": be.concatenate(masks, axis=0), "shape": shape}
 
 
-@dataclass
+@dataclass(config=dict(extra="allow"))
 class BandPassReconstructed(ComposableFilter):
     """
-    Generate reconstructed bandpass filters in Fourier space.
+    Generate Bandpass filter for reconstructions.
+
+    Examples
+    --------
+    The following creates an instance of :py:class:`BandPassReconstructed`
+
+    >>> from tme.filters import BandPassReconstructed
+    >>> bpf_instance = BandPassReconstructed(
+    >>>     lowpass=30,
+    >>>     sampling_rate=10
+    >>> )
+
+    We can use its call method to create filters of given shape
+
+    >>> import matplotlib.pyplot as plt
+    >>> ret = bpf_instance(shape=(50,50))
+
+    The ``data`` key of the returned dictionary contains the corresponding
+    Fourier filter mask. The DC component is located at the origin.
+
+    >>> plt.imshow(ret["data"])
+    >>> plt.show()
     """
 
     #: The lowpass cutoff, defaults to None.
-    lowpass: float = None
+    lowpass: Optional[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
+    highpass: Optional[float] = None
     #: The sampling rate, defaults to 1 Ångstrom / voxel.
-    sampling_rate: Tuple[float] = 1
+    sampling_rate: Union[Tuple[float, ...], 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)
 
+    def _evaluate(self, shape: Tuple[int, ...], **kwargs):
         func = discrete_bandpass
-        if func_args.get("use_gaussian"):
+        if kwargs.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=func_args["shape"],
+            shape=shape,
             sampling_rate=0.5,
-            shape_is_real_fourier=shape_is_real_fourier,
+            shape_is_real_fourier=False,
             compute_euclidean_norm=True,
+            fftshift=False,
         )
-        mask = func(grid=grid, **func_args)
-
-        mask = shift_fourier(data=mask, shape_is_real_fourier=shape_is_real_fourier)
-        if return_real_fourier:
-            mask = crop_real_fourier(mask)
-
-        return {
-            "data": be.to_backend_array(mask),
-            "shape": func_args["shape"],
-            "return_real_fourier": return_real_fourier,
-            "is_multiplicative_filter": True,
-        }
+        ret = be.to_backend_array(func(grid=grid, **kwargs))
+        return {"data": ret, "shape": shape}
 
 
 def discrete_bandpass(
@@ -169,11 +169,9 @@ def discrete_bandpass(
     grid : BackendArray
         Frequencies in Fourier space.
     lowpass : float
-        The lowpass cutoff in units of sampling rate.
+        The lowpass cutoff in spatial 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.
+        The highpass cutoff in spatial units of sampling rate.
     sampling_rate : float
         The sampling rate in Fourier space.
     **kwargs : dict
@@ -244,16 +242,13 @@ def gaussian_bandpass(
     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.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.divide(grid, highpass)
         highpass_filter = be.exp(highpass_filter, out=highpass_filter)
         highpass_filter = be.subtract(1, highpass_filter, out=highpass_filter)
 
@@ -267,5 +262,4 @@ def gaussian_bandpass(
         )
     else:
         bandpass_filter = be.full(grid.shape, fill_value=1, dtype=be._float_dtype)
-
     return bandpass_filter

tme/filters/compose.py

@@ -9,74 +9,168 @@ Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 from typing import Tuple, Dict
 from abc import ABC, abstractmethod
 
-from tme.backends import backend as be
+from ._utils import crop_real_fourier
+from ..backends import backend as be
 
 __all__ = ["Compose", "ComposableFilter"]
 
 
-class Compose:
+class ComposableFilter(ABC):
     """
-    Compose a series of transformations.
+    Base class for composable filters.
 
-    This class allows composing multiple transformations together. Each transformation
-    is expected to be a callable that accepts keyword arguments and returns metadata.
+    This class provides a standard interface for filters used in template matching
+    and reconstruction. It automatically handles:
 
-    Parameters
-    ----------
-    transforms : Tuple[object]
-        A tuple containing transformation objects.
+    - Parameter merging between instance attributes and runtime arguments
+    - Fourier space shifting when needed
+    - Real Fourier transform cropping for efficiency
+    - Standardized result dictionary formatting
 
-    Returns
-    -------
-    Dict
-        Metadata resulting from the composed transformations.
+    Subclasses need to implement :py:meth:`ComposableFilter._evaluate` which
+    contains the core filter computation logic.
 
+    By default, all filters are assumed to be multiplicative in Fourier space,
+    which covers the vast majority of use cases (bandpass, CTF, wedge, whitening, etc.).
+    Only explicitly specify non-multiplicative behavior when needed.
     """
 
-    def __init__(self, transforms: Tuple[object]):
-        self.transforms = transforms
+    @abstractmethod
+    def _evaluate(self, **kwargs) -> Dict:
+        """
+        Compute the actual filter given a set of keyword parameters.
 
-    def __call__(self, **kwargs: Dict) -> Dict:
-        meta = {}
-        if not len(self.transforms):
-            return meta
+        Parameters
+        ----------
+        **kwargs : dict
+            Merged parameters from instance attributes and runtime arguments
+            passed to :py:meth:`__call__`. This includes both the filter's
+            configuration parameters and any runtime overrides.
 
-        meta = self.transforms[0](**kwargs)
-        for transform in self.transforms[1:]:
-            kwargs.update(meta)
-            ret = transform(**kwargs)
+        Returns
+        -------
+        Dict
+            Dictionary containing the filter result and metadata. Required keys:
 
-            if "data" not in ret:
-                continue
+            - data : BackendArray or array-like
+                The computed filter data
+            - shape : tuple of int
+                Input shape the filter was built for.
 
-            if ret.get("is_multiplicative_filter", False):
-                prev_data = meta.pop("data")
-                ret["data"] = be.multiply(ret["data"], prev_data)
-                ret["merge"], prev_data = None, None
+            Optional keys:
+            - is_multiplicative_filter : bool
+                Whether the filter is multiplicative in Fourier space (default True)
+        """
 
-            meta = ret
+    def __call__(self, return_real_fourier: bool = False, **kwargs) -> Dict:
+        """
+        This method provides the standard interface for creating of composable
+        filter masks. It merges instance attributes with runtime parameters,
+        and ensures Fourier conventions are consistent across filters.
 
-        return meta
+        Parameters
+        ----------
+        return_real_fourier : bool, optional
+            Whether to crop the filter mask for compatibility with real input
+            FFTs (i.e., :py:func:`numpy.fft.rfft`). When True, only the
+            positive frequency components are returned, reducing memory usage
+            and computation time for real-valued inputs. Default is False.
+        **kwargs : dict
+            Additional keyword arguments passed to :py:meth:`_evaluate`.
+            These will override any matching instance attributes during
+            parameter merging.
+
+        Returns
+        -------
+        Dict
+            - data : BackendArray
+                The processed filter data, converted to the appropriate backend
+                array type and with fourier operations applied as needed
+            - shape : tuple of int or None
+                Shape for which the filter was created
+            - return_real_fourier : bool
+                The value of the return_real_fourier parameter
+            - is_multiplicative_filter : bool
+                Whether the filter is multiplicative in Fourier space
+            - Additional metadata from the filter implementation
+        """
+        ret = self._evaluate(**(vars(self) | kwargs))
 
+        # This parameter is only here to allow for using Composable filters outside
+        # the context of a Compose operation. Internally, we require return_real_fourier
+        # to be False, e.g., for filters that require reconstruction.
+        if return_real_fourier:
+            ret["data"] = crop_real_fourier(ret["data"])
 
-class ComposableFilter(ABC):
+        ret["data"] = be.to_backend_array(ret["data"])
+        ret["return_real_fourier"] = return_real_fourier
+        return ret
+
+
+class Compose:
     """
-    Base class for composable filters.
+    Compose a series of filters.
+
+    Parameters
+    ----------
+    transforms : tuple of :py:class:`ComposableFilter`.
+        Tuple of filter instances.
     """
 
-    @abstractmethod
-    def __call__(self, *args, **kwargs) -> Dict:
+    def __init__(self, transforms: Tuple[ComposableFilter, ...]):
+        for transform in transforms:
+            if not isinstance(transform, ComposableFilter):
+                raise ValueError(f"{transform} is not a child of {ComposableFilter}.")
+
+        self.transforms = transforms
+
+    def __call__(self, return_real_fourier: bool = False, **kwargs) -> Dict:
         """
+        Apply the sequence of filters in order, chaining their outputs.
 
         Parameters
         ----------
-        *args : tuple
-            Variable length argument list.
+        return_real_fourier : bool, optional
+            Whether to crop the filter mask for compatibility with real input
+            FFTs (i.e., :py:func:`numpy.fft.rfft`). When True, only the
+            positive frequency components are returned, reducing memory usage
+            and computation time for real-valued inputs. Default is False.
         **kwargs : dict
-            Arbitrary keyword arguments.
+            Keyword arguments passed to the first filter and propagated through
+            the pipeline.
 
         Returns
         -------
         Dict
-            A dictionary representing the result of the filtering operation.
+            Result dictionary from the final filter in the composition, containing:
+
+            - data : BackendArray
+                The final composite filter data. For multiplicative filters, this is
+                the element-wise product of all individual filter outputs.
+            - shape : tuple of int
+                Shape of the filter data
+            - return_real_fourier : bool
+                Whether the output is compatible with real FFTs
+            - Additional metadata from the filter pipeline
         """
+        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 "data" not in ret:
+                continue
+
+            if ret.get("is_multiplicative_filter", True):
+                prev_data = meta.pop("data")
+                ret["data"] = be.multiply(ret["data"], prev_data)
+                ret["merge"], prev_data = None, None
+            meta = ret
+
+        if return_real_fourier:
+            meta["data"] = crop_real_fourier(meta["data"])
+        return meta