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

tme/filters/reconstruction.py

@@ -1,8 +1,9 @@
-""" Defines filters on tomographic tilt series.
+"""
+Defines filters on tomographic tilt series.
 
-    Copyright (c) 2024 European Molecular Biology Laboratory
+Copyright (c) 2024 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 from typing import Tuple
@@ -15,27 +16,26 @@ from ..backends import backend as be
 
 from .compose import ComposableFilter
 from ..rotations import euler_to_rotationmatrix
-from ._utils import (
-    crop_real_fourier,
-    shift_fourier,
-    create_reconstruction_filter,
-)
+from ._utils import crop_real_fourier, shift_fourier, create_reconstruction_filter
 
 __all__ = ["ReconstructFromTilt"]
 
 
 @dataclass
 class ReconstructFromTilt(ComposableFilter):
-    """Reconstruct a volume from a tilt series."""
+    """
+    Reconstruct a d+1 array from a d-dimensional input projection using weighted
+    backprojection (WBP).
+    """
 
     #: Shape of the reconstruction.
     shape: Tuple[int] = None
     #: Angle of each individual tilt.
     angles: Tuple[float] = None
-    #: The axis around which the volume is opened.
-    opening_axis: int = 0
-    #: Axis the plane is tilted over.
-    tilt_axis: int = 2
+    #: Projection axis, defaults to 2 (z).
+    opening_axis: int = 2
+    #: Tilt axis, defaults to 0 (x).
+    tilt_axis: int = 0
     #: Whether to return a share compliant with rfftn.
     return_real_fourier: bool = True
     #: Interpolation order used for rotation
@@ -43,19 +43,57 @@ class ReconstructFromTilt(ComposableFilter):
     #: Filter window applied during reconstruction.
     reconstruction_filter: str = None
 
-    def __call__(self, **kwargs):
+    def __call__(self, return_real_fourier: bool = False, **kwargs):
+        """
+        Reconstruct a  d+1 array from a d-dimensional input using WBP.
+
+        Parameters
+        ----------
+        shape : tuple of int
+            The shape of the reconstruction volume.
+        data : BackendArray
+            D-dimensional image stack with shape (n, ...)
+        angles : tuple of float
+            Angle of each individual tilt.
+        return_real_fourier : bool, optional
+            Return a shape compliant
+        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). Defaults
+            to False.
+        reconstruction_filter : bool, optional
+           Filter window applied during reconstruction.
+           See :py:meth:`create_reconstruction_filter` for available options.
+        tilt_axis : int
+            Axis the plane is tilted over, defaults to 0 (x).
+        opening_axis : int
+            The projection axis, defaults to 2 (z).
+
+        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.
+        """
+
         func_args = vars(self).copy()
         func_args.update(kwargs)
 
         ret = self.reconstruct(**func_args)
 
+        if return_real_fourier:
+            ret = crop_real_fourier(ret)
+
         return {
             "data": ret,
-            "shape": ret.shape,
-            "shape_is_real_fourier": func_args["return_real_fourier"],
-            "angles": func_args["angles"],
-            "tilt_axis": func_args["tilt_axis"],
-            "opening_axis": func_args["opening_axis"],
+            "shape": func_args["shape"],
+            "shape_is_real_fourier": return_real_fourier,
             "is_multiplicative_filter": False,
         }
 
@@ -67,7 +105,6 @@ class ReconstructFromTilt(ComposableFilter):
         opening_axis: int,
         tilt_axis: int,
         interpolation_order: int = 1,
-        return_real_fourier: bool = True,
         reconstruction_filter: str = None,
         **kwargs,
     ):
@@ -88,8 +125,6 @@ class ReconstructFromTilt(ComposableFilter):
             Axis the plane is tilted over.
         interpolation_order : int, optional
             Interpolation order used for rotation, defaults to 1.
-        return_real_fourier : bool, optional
-           Whether to return a shape compliant with rfftn, defaults to True.
         reconstruction_filter : bool, optional
            Filter window applied during reconstruction.
            See :py:meth:`create_reconstruction_filter` for available options.
@@ -152,9 +187,4 @@ class ReconstructFromTilt(ComposableFilter):
             )
             volume = be.add(volume, volume_temp_rotated, out=volume)
 
-        volume = shift_fourier(data=volume, shape_is_real_fourier=False)
-
-        if return_real_fourier:
-            volume = crop_real_fourier(volume)
-
-        return volume
+        return shift_fourier(data=volume, shape_is_real_fourier=False)

tme/filters/wedge.py

@@ -1,11 +1,13 @@
-""" Implements class Wedge and WedgeReconstructed to create Fourier
-    filter representations.
+"""
+Implements class Wedge and WedgeReconstructed to create Fourier
+filter representations.
 
-    Copyright (c) 2024 European Molecular Biology Laboratory
+Copyright (c) 2024 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
+import warnings
 from typing import Tuple, Dict
 
 import numpy as np
@@ -14,6 +16,7 @@ from ..types import NDArray
 from ..backends import backend as be
 from .compose import ComposableFilter
 from ..matching_utils import centered
+from ..parser import XMLParser, StarParser
 from ..rotations import euler_to_rotationmatrix
 from ._utils import (
     centered_grid,
@@ -37,9 +40,9 @@ class Wedge(ComposableFilter):
     shape : tuple of int
         The shape of the reconstruction volume.
     tilt_axis : int
-        Axis the plane is tilted over.
+        Axis the plane is tilted over, defaults to 0 (x).
     opening_axis : int
-        The axis around which the volume is opened.
+        The projection axis, defaults to 2 (z).
     angles : tuple of float
         The tilt angles.
     weights : tuple of float
@@ -51,17 +54,24 @@ class Wedge(ComposableFilter):
 
     Returns
     -------
-    Dict
-        A dictionary containing weighted wedges and related information.
+    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.
     """
 
     def __init__(
         self,
         shape: Tuple[int],
-        tilt_axis: int,
-        opening_axis: int,
         angles: Tuple[float],
         weights: Tuple[float],
+        tilt_axis: int = 0,
+        opening_axis: int = 2,
         weight_type: str = None,
         frequency_cutoff: float = 0.5,
     ):
@@ -75,8 +85,16 @@ class Wedge(ComposableFilter):
     @classmethod
     def from_file(cls, filename: str) -> "Wedge":
         """
-        Generate a :py:class:`Wedge` instance by reading tilt angles and weights
-        from a tab-separated text file.
+        Generate a :py:class:`Wedge` instance by reading tilt angles and weights.
+        Supported extensions are:
+
+            +-------+---------------------------------------------------------+
+            | .star | Tomostar STAR file                                      |
+            +-------+---------------------------------------------------------+
+            | .xml  | WARP/M XML file                                         |
+            +-------+---------------------------------------------------------+
+            | .*    | Tab-separated file with optional column names           |
+            +-------+---------------------------------------------------------+
 
         Parameters
         ----------
@@ -88,8 +106,13 @@ class Wedge(ComposableFilter):
         :py:class:`Wedge`
            Class instance instance initialized with angles and weights from the file.
         """
-        data = cls._from_text(filename)
+        func = _from_text
+        if filename.lower().endswith("xml"):
+            func = _from_xml
+        elif filename.lower().endswith("star"):
+            func = _from_star
 
+        data = func(filename)
         angles, weights = data.get("angles", None), data.get("weights", None)
         if angles is None:
             raise ValueError(f"Could not find colum angles in {filename}")
@@ -108,32 +131,6 @@ class Wedge(ComposableFilter):
             weights=np.array(weights, dtype=np.float32),
         )
 
-    @staticmethod
-    def _from_text(filename: str, delimiter="\t") -> Dict:
-        """
-        Read column data from a text file.
-
-        Parameters
-        ----------
-        filename : str
-            The path to the text file.
-        delimiter : str, optional
-            The delimiter used in the file, defaults to '\t'.
-
-        Returns
-        -------
-        Dict
-            A dictionary with one key for each column.
-        """
-        with open(filename, mode="r", encoding="utf-8") as infile:
-            data = [x.strip() for x in infile.read().split("\n")]
-            data = [x.split("\t") for x in data if len(x)]
-
-        headers = data.pop(0)
-        ret = {header: list(column) for header, column in zip(headers, zip(*data))}
-
-        return ret
-
     def __call__(self, **kwargs: Dict) -> NDArray:
         func_args = vars(self).copy()
         func_args.update(kwargs)
@@ -172,10 +169,7 @@ class Wedge(ComposableFilter):
 
         return {
             "data": ret,
-            "angles": func_args["angles"],
-            "tilt_axis": func_args["tilt_axis"],
-            "opening_axis": func_args["opening_axis"],
-            "sampling_rate": func_args.get("sampling_rate", 1),
+            "shape": func_args["shape"],
             "is_multiplicative_filter": True,
         }
 
@@ -297,10 +291,10 @@ class WedgeReconstructed:
     ----------
     angles :tuple of float, optional
         The tilt angles, defaults to None.
-    opening_axis : int, optional
-        The axis around which the wedge is opened.
-    tilt_axis : int, optional
-        The axis along which the tilt is applied.
+    tilt_axis : int
+        Axis the plane is tilted over, defaults to 0 (x).
+    opening_axis : int
+        The projection axis, defaults to 2 (z).
     weights : tuple of float, optional
         Weights to assign to individual wedge components.
     weight_wedge : bool, optional
@@ -317,8 +311,8 @@ class WedgeReconstructed:
 
     def __init__(
         self,
-        opening_axis: int,
-        tilt_axis: int,
+        opening_axis: int = 2,
+        tilt_axis: int = 0,
         angles: Tuple[float] = None,
         weights: Tuple[float] = None,
         weight_wedge: bool = False,
@@ -336,7 +330,9 @@ class WedgeReconstructed:
         self.create_continuous_wedge = create_continuous_wedge
         self.frequency_cutoff = frequency_cutoff
 
-    def __call__(self, shape: Tuple[int], **kwargs: Dict) -> Dict:
+    def __call__(
+        self, shape: Tuple[int], return_real_fourier: bool = False, **kwargs: Dict
+    ) -> Dict:
         """
         Generate the reconstructed wedge.
 
@@ -344,20 +340,28 @@ class WedgeReconstructed:
         ----------
         shape : tuple of int
             The shape of the reconstruction volume.
+        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). Defaults
+            to False.
         **kwargs : Dict
             Additional keyword arguments.
 
         Returns
         -------
-        Dict
-            A dictionary containing the reconstructed wedge and related information.
+        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.
         """
         func_args = vars(self).copy()
         func_args.update(kwargs)
 
-        if kwargs.get("is_fourier_shape", False):
-            print("Cannot create continuous wedge mask based on real fourier shape.")
-
         func = self.step_wedge
         if func_args.get("create_continuous_wedge", False):
             func = self.continuous_wedge
@@ -386,17 +390,15 @@ class WedgeReconstructed:
         ret = be.astype(be.to_backend_array(ret), be._float_dtype)
 
         ret = shift_fourier(data=ret, shape_is_real_fourier=False)
-        if func_args.get("return_real_fourier", False):
+
+        if return_real_fourier:
             ret = crop_real_fourier(ret)
 
         return {
             "data": ret,
-            "shape_is_real_fourier": func_args["return_real_fourier"],
-            "shape": ret.shape,
-            "tilt_axis": func_args["tilt_axis"],
-            "opening_axis": func_args["opening_axis"],
+            "shape": shape,
+            "return_real_fourier": return_real_fourier,
             "is_multiplicative_filter": True,
-            "angles": func_args["angles"],
         }
 
     @staticmethod
@@ -426,6 +428,7 @@ class WedgeReconstructed:
         NDArray
             Wedge mask.
         """
+        angles = np.abs(np.asarray(angles))
         aspect_ratio = shape[opening_axis] / shape[tilt_axis]
         angles = np.degrees(np.arctan(np.tan(np.radians(angles)) * aspect_ratio))
 
@@ -540,3 +543,78 @@ class WedgeReconstructed:
         wedge_volume = np.tile(wedge_volume, tile_dimensions)
 
         return wedge_volume
+
+
+def _from_xml(filename: str, **kwargs) -> Dict:
+    """
+    Read tilt data from a WARP/M XML file.
+
+    Parameters
+    ----------
+    filename : str
+        The path to the text file.
+
+    Returns
+    -------
+    Dict
+        A dictionary with one key for each column.
+    """
+    data = XMLParser(filename)
+    return {"angles": data["Angles"], "weights": data["Dose"]}
+
+
+def _from_star(filename: str, **kwargs) -> Dict:
+    """
+    Read tilt data from a tomostar STAR file.
+
+    Parameters
+    ----------
+    filename : str
+        The path to the text file.
+
+    Returns
+    -------
+    Dict
+        A dictionary with one key for each column.
+    """
+    data = StarParser(filename, delimiter=None)["data_"]
+    return {"angles": data["_wrpAxisAngle"], "weights": data["_wrpDose"]}
+
+
+def _from_text(filename: str, **kwargs) -> Dict:
+    """
+    Read column data from a text file.
+
+    Parameters
+    ----------
+    filename : str
+        The path to the text file.
+    delimiter : str, optional
+        The delimiter used in the file, defaults to '\t'.
+
+    Returns
+    -------
+    Dict
+        A dictionary with one key for each column.
+    """
+    with open(filename, mode="r", encoding="utf-8") as infile:
+        data = [x.strip() for x in infile.read().split("\n")]
+        data = [x.split("\t") for x in data if len(x)]
+
+    if "angles" in data[0]:
+        headers = data.pop(0)
+    else:
+        warnings.warn(
+            f"Did not find a column named 'angles' in {filename}. Assuming "
+            "first column specifies angles."
+        )
+        if len(data[0]) != 1:
+            raise ValueError(
+                "Found more than one column without column names. Please add "
+                "column names to your file. If you only want to specify tilt "
+                "angles without column names, use a single column file."
+            )
+        headers = ("angles",)
+    ret = {header: list(column) for header, column in zip(headers, zip(*data))}
+
+    return ret

tme/filters/whitening.py

@@ -1,8 +1,9 @@
-""" Implements class BandPassFilter to create Fourier filter representations.
+"""
+Implements class BandPassFilter to create Fourier filter representations.
 
-    Copyright (c) 2024 European Molecular Biology Laboratory
+Copyright (c) 2024 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 from typing import Tuple, Dict
@@ -14,7 +15,7 @@ from scipy.ndimage import map_coordinates
 from ..types import BackendArray
 from ..backends import backend as be
 from .compose import ComposableFilter
-from ._utils import fftfreqn, compute_fourier_shape
+from ._utils import fftfreqn, compute_fourier_shape, shift_fourier
 
 __all__ = ["LinearWhiteningFilter"]
 
@@ -103,13 +104,6 @@ class LinearWhiteningFilter(ComposableFilter):
         shape_is_real_fourier: bool = True,
         order: int = 1,
     ) -> BackendArray:
-        """
-        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=0.5,
@@ -123,11 +117,13 @@ class LinearWhiteningFilter(ComposableFilter):
 
     def __call__(
         self,
+        shape: Tuple[int],
         data: BackendArray = None,
         data_rfft: BackendArray = None,
         n_bins: int = None,
         batch_dimension: int = None,
         order: int = 1,
+        return_real_fourier: bool = True,
         **kwargs: Dict,
     ) -> Dict:
         """
@@ -135,6 +131,8 @@ class LinearWhiteningFilter(ComposableFilter):
 
         Parameters
         ----------
+        shape : tuple of ints
+            Shape of the returned whitening filter.
         data : BackendArray, optional
             The input data, defaults to None.
         data_rfft : BackendArray, optional
@@ -143,49 +141,51 @@ class LinearWhiteningFilter(ComposableFilter):
             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.
+        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)
         **kwargs : Dict
             Additional keyword arguments.
 
         Returns
         -------
-        Dict
-           Filter data and associated parameters.
+        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 = np.fft.rfftn(be.to_numpy_array(data))
+            data_rfft = be.rfftn(data)
 
         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)
 
-        if order is None:
-            cutoff = bins < radial_averages.size
-            filter_mask = np.zeros(bins.shape, radial_averages.dtype)
-            filter_mask[cutoff] = radial_averages[bins[cutoff]]
-        else:
-            shape = bins.shape
-            if kwargs.get("shape", False):
-                shape = compute_fourier_shape(
-                    shape=kwargs.get("shape"),
-                    shape_is_real_fourier=kwargs.get("shape_is_real_fourier", False),
-                )
-
-            filter_mask = self._interpolate_spectrum(
-                spectrum=radial_averages,
+        shape_filter = shape
+        if return_real_fourier:
+            shape_filter = compute_fourier_shape(
                 shape=shape,
-                shape_is_real_fourier=True,
+                shape_is_real_fourier=False,
             )
 
-        filter_mask = np.fft.ifftshift(
-            filter_mask,
-            axes=tuple(i for i in range(data_rfft.ndim - 1) if i != batch_dimension),
+        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(filter_mask),
+            "data": be.to_backend_array(ret),
+            "shape": shape,
+            "return_real_fourier": return_real_fourier,
             "is_multiplicative_filter": True,
         }

tme/matching_data.py

@@ -1,8 +1,9 @@
-""" Class representation of template matching data.
+"""
+Class representation of template matching data.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 import warnings

tme/matching_exhaustive.py

@@ -1,8 +1,9 @@
-""" Implements cross-correlation based template matching using different metrics.
+"""
+Implements cross-correlation based template matching using different metrics.
 
-    Copyright (c) 2023 European Molecular Biology Laboratory
+Copyright (c) 2023 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 import sys
@@ -19,6 +20,7 @@ from .filters import Compose
 from .backends import backend as be
 from .matching_utils import split_shape
 from .types import CallbackClass, MatchingData
+from .analyzer.proxy import SharedAnalyzerProxy
 from .matching_scores import MATCHING_EXHAUSTIVE_REGISTER
 from .memory import MatchingMemoryUsage, MATCHING_MEMORY_REGISTRY
 
@@ -242,6 +244,7 @@ def scan(
         "aggregate_axis": matching_data._batch_axis(matching_data._batch_mask),
         "n_rotations": matching_data.rotations.shape[0],
     }
+    callback_class_args["inversion_mapping"] = n_jobs == 1
     default_callback_args.update(callback_class_args)
 
     setup = matching_setup(
@@ -257,13 +260,17 @@ def scan(
     matching_data._free_data()
     be.free_cache()
 
-    # Some analyzers cannot be shared across processes
-    if not getattr(callback_class, "is_shareable", False):
-        jobs_per_callback_class = 1
-
     n_callback_classes = max(n_jobs // jobs_per_callback_class, 1)
     callback_classes = [
-        callback_class(**default_callback_args) if callback_class else None
+        (
+            SharedAnalyzerProxy(
+                callback_class,
+                default_callback_args,
+                shm_handler=shm_handler if n_jobs > 1 else None,
+            )
+            if callback_class
+            else None
+        )
         for _ in range(n_callback_classes)
     ]
     ret = Parallel(n_jobs=n_jobs)(
@@ -276,14 +283,9 @@ def scan(
         )
         for index, rotation in enumerate(matching_data._split_rotations_on_jobs(n_jobs))
     )
-
-    # TODO: Make sure peak callers are thread safe to begin with
-    if not getattr(callback_class, "is_shareable", False):
-        callback_classes = ret
-
     callbacks = [
-        tuple(callback._postprocess(**default_callback_args))
-        for callback in callback_classes
+        callback.result(**default_callback_args)
+        for callback in ret[:n_callback_classes]
         if callback
     ]
     be.free_cache()
@@ -454,7 +456,6 @@ def scan_subsets(
                 for index, (target_split, template_split) in enumerate(splits)
             ]
         )
-
     matching_data._free_data()
     if callback_class is not None:
         return callback_class.merge(results, **callback_class_args)