pytme 0.2.9__cp311-cp311-macosx_15_0_arm64.whl → 0.3.0__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
@@ -10,32 +11,31 @@ from dataclasses import dataclass
 
 import numpy as np
 
-from ..types import NDArray
 from ..backends import backend as be
+from ..types import NDArray, BackendArray
 
 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"]
+__all__ = ["ReconstructFromTilt", "ShiftFourier"]
 
 
 @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,60 @@ 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, ...). The data is assumed to be
+            a Fourier transform of the stack you are trying to reconstruct with
+            DC component in the center.
+        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)
 
+        ret = shift_fourier(data=ret, shape_is_real_fourier=False)
+        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"],
+            "return_real_fourier": return_real_fourier,
             "is_multiplicative_filter": False,
         }
 
@@ -67,7 +108,6 @@ class ReconstructFromTilt(ComposableFilter):
         opening_axis: int,
         tilt_axis: int,
         interpolation_order: int = 1,
-        return_real_fourier: bool = True,
         reconstruction_filter: str = None,
         **kwargs,
     ):
@@ -77,7 +117,7 @@ class ReconstructFromTilt(ComposableFilter):
         Parameters
         ----------
         data : NDArray
-            The tilt series data.
+            The Fourier transform of tilt series data.
         shape : tuple of int
             Shape of the reconstruction.
         angles : tuple of float
@@ -88,8 +128,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.
@@ -103,9 +141,9 @@ class ReconstructFromTilt(ComposableFilter):
             return data
 
         data = be.to_backend_array(data)
-        volume_temp = be.zeros(shape, dtype=be._float_dtype)
-        volume_temp_rotated = be.zeros(shape, dtype=be._float_dtype)
-        volume = be.zeros(shape, dtype=be._float_dtype)
+        volume_temp = be.zeros(shape, dtype=data.dtype)
+        volume_temp_rotated = be.zeros(shape, dtype=data.dtype)
+        volume = be.zeros(shape, dtype=data.dtype)
 
         slices = tuple(slice(a // 2, (a // 2) + 1) for a in shape)
         subset = tuple(
@@ -152,9 +190,30 @@ class ReconstructFromTilt(ComposableFilter):
             )
             volume = be.add(volume, volume_temp_rotated, out=volume)
 
-        volume = shift_fourier(data=volume, shape_is_real_fourier=False)
+        return volume
 
-        if return_real_fourier:
-            volume = crop_real_fourier(volume)
 
-        return volume
+class ShiftFourier(ComposableFilter):
+    def __call__(
+        self,
+        data: BackendArray,
+        shape_is_real_fourier: bool = False,
+        return_real_fourier: bool = True,
+        **kwargs,
+    ):
+        ret = []
+        for index in range(data.shape[0]):
+            mask = be.to_numpy_array(data[index])
+
+            mask = shift_fourier(data=mask, shape_is_real_fourier=shape_is_real_fourier)
+            if return_real_fourier:
+                mask = crop_real_fourier(mask)
+            ret.append(mask[None])
+        ret = np.concatenate(ret, axis=0)
+
+        return {
+            "data": ret,
+            "shape": kwargs.get("shape"),
+            "return_real_fourier": return_real_fourier,
+            "is_multiplicative_filter": False,
+        }

tme/filters/wedge.py

@@ -1,12 +1,14 @@
-""" 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>
 """
 
 from typing import Tuple, Dict
+from dataclasses import dataclass
 
 import numpy as np
 
@@ -15,6 +17,7 @@ from ..backends import backend as be
 from .compose import ComposableFilter
 from ..matching_utils import centered
 from ..rotations import euler_to_rotationmatrix
+from ..parser import XMLParser, StarParser, MDOCParser
 from ._utils import (
     centered_grid,
     frequency_grid_at_angle,
@@ -28,55 +31,44 @@ from ._utils import (
 __all__ = ["Wedge", "WedgeReconstructed"]
 
 
+@dataclass
 class Wedge(ComposableFilter):
     """
     Generate wedge mask for tomographic data.
-
-    Parameters
-    ----------
-    shape : tuple of int
-        The shape of the reconstruction volume.
-    tilt_axis : int
-        Axis the plane is tilted over.
-    opening_axis : int
-        The axis around which the volume is opened.
-    angles : tuple of float
-        The tilt angles.
-    weights : tuple of float
-        The weights corresponding to each tilt angle.
-    weight_type : str, optional
-        The type of weighting to apply, defaults to None.
-    frequency_cutoff : float, optional
-        Frequency cutoff for created mask. Nyquist 0.5 by default.
-
-    Returns
-    -------
-    Dict
-        A dictionary containing weighted wedges and related information.
     """
 
-    def __init__(
-        self,
-        shape: Tuple[int],
-        tilt_axis: int,
-        opening_axis: int,
-        angles: Tuple[float],
-        weights: Tuple[float],
-        weight_type: str = None,
-        frequency_cutoff: float = 0.5,
-    ):
-        self.shape = shape
-        self.tilt_axis = tilt_axis
-        self.opening_axis = opening_axis
-        self.angles = angles
-        self.weights = weights
-        self.frequency_cutoff = frequency_cutoff
+    #: The shape of the reconstruction volume.
+    shape: Tuple[int] = None
+    #: The tilt angles.
+    angles: Tuple[float] = None
+    #: The weights corresponding to each tilt angle.
+    weights: 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 type of weighting to apply, defaults to None.
+    weight_type: str = None
+    #: Frequency cutoff for created mask. Nyquist 0.5 by default.
+    frequency_cutoff: float = 0.5
+    #: The sampling rate, defaults to 1 Ångstrom / voxel.
+    sampling_rate: Tuple[float] = 1
 
     @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                                         |
+            +-------+---------------------------------------------------------+
+            | .mdoc | SerialEM file                                           |
+            +-------+---------------------------------------------------------+
+            | .*    | Tab-separated file with optional column names           |
+            +-------+---------------------------------------------------------+
 
         Parameters
         ----------
@@ -88,8 +80,15 @@ 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
+        elif filename.lower().endswith("mdoc"):
+            func = _from_mdoc
+
+        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,33 +107,10 @@ class Wedge(ComposableFilter):
             weights=np.array(weights, dtype=np.float32),
         )
 
-    @staticmethod
-    def _from_text(filename: str, delimiter="\t") -> Dict:
+    def __call__(self, **kwargs: Dict) -> NDArray:
         """
-        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.
+        Returns a Wedge stack of chosen parameters with DC component in the center.
         """
-        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 +148,8 @@ 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"],
+            "return_real_fourier": func_args.get("return_real_fourier", False),
             "is_multiplicative_filter": True,
         }
 
@@ -202,7 +176,12 @@ class Wedge(ComposableFilter):
         return wedges
 
     def weight_relion(
-        self, shape: Tuple[int], opening_axis: int, tilt_axis: int, **kwargs
+        self,
+        shape: Tuple[int],
+        opening_axis: int,
+        tilt_axis: int,
+        sampling_rate: float = 1.0,
+        **kwargs,
     ) -> NDArray:
         """
         Generate weighted wedges based on the RELION 1.4 formalism, weighting each
@@ -217,7 +196,6 @@ class Wedge(ComposableFilter):
         tilt_shape = compute_tilt_shape(
             shape=shape, opening_axis=opening_axis, reduce_dim=True
         )
-
         wedges = np.zeros((len(self.angles), *tilt_shape))
         for index, angle in enumerate(self.angles):
             frequency_grid = frequency_grid_at_angle(
@@ -225,7 +203,7 @@ class Wedge(ComposableFilter):
                 opening_axis=opening_axis,
                 tilt_axis=tilt_axis,
                 angle=angle,
-                sampling_rate=1,
+                sampling_rate=sampling_rate,
             )
             sigma = np.sqrt(self.weights[index] * 4 / (8 * np.pi**2))
             sigma = -2 * np.pi**2 * sigma**2
@@ -245,6 +223,7 @@ class Wedge(ComposableFilter):
         amplitude: float = 0.245,
         power: float = -1.665,
         offset: float = 2.81,
+        sampling_rate: float = 1.0,
         **kwargs,
     ) -> NDArray:
         """
@@ -270,7 +249,7 @@ class Wedge(ComposableFilter):
                 opening_axis=opening_axis,
                 tilt_axis=tilt_axis,
                 angle=angle,
-                sampling_rate=1,
+                sampling_rate=sampling_rate,
             )
 
             with np.errstate(divide="ignore"):
@@ -289,54 +268,36 @@ class Wedge(ComposableFilter):
         return wedges
 
 
+@dataclass
 class WedgeReconstructed:
     """
     Initialize :py:class:`WedgeReconstructed`.
-
-    Parameters
-    ----------
-    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.
-    weights : tuple of float, optional
-        Weights to assign to individual wedge components.
-    weight_wedge : bool, optional
-        Whether individual wedge components should be weighted. If True and weights
-        is None, uses the cosine of the angle otherwise weights.
-    create_continuous_wedge: bool, optional
-        Whether to create a continous wedge or a per-component wedge. Weights are only
-        considered for non-continuous wedges.
-    frequency_cutoff : float, optional
-        Filter window applied during reconstruction.
-    **kwargs : Dict
-        Additional keyword arguments.
     """
 
-    def __init__(
-        self,
-        opening_axis: int,
-        tilt_axis: int,
-        angles: Tuple[float] = None,
-        weights: Tuple[float] = None,
-        weight_wedge: bool = False,
-        create_continuous_wedge: bool = False,
-        frequency_cutoff: float = 0.5,
-        reconstruction_filter: str = None,
-        **kwargs: Dict,
-    ):
-        self.angles = angles
-        self.opening_axis = opening_axis
-        self.tilt_axis = tilt_axis
-        self.weights = weights
-        self.weight_wedge = weight_wedge
-        self.reconstruction_filter = reconstruction_filter
-        self.create_continuous_wedge = create_continuous_wedge
-        self.frequency_cutoff = frequency_cutoff
-
-    def __call__(self, shape: Tuple[int], **kwargs: Dict) -> Dict:
+    #: The tilt angles, defaults to None.
+    angles: Tuple[float] = None
+    #: Weights to assign to individual wedge components. Not considered for continuous wedge
+    weights: Tuple[float] = None
+    #: Whether individual wedge components should be weighted.
+    weight_wedge: bool = False
+    #: Whether to create a continous wedge or a per-component wedge.
+    create_continuous_wedge: bool = False
+    #: Frequency cutoff of filter
+    frequency_cutoff: float = 0.5
+    #: 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
+    #: Filter window applied during reconstruction.
+    reconstruction_filter: str = None
+
+    def __post_init__(self):
+        if self.create_continuous_wedge:
+            self.angles = (min(self.angles), max(self.angles))
+
+    def __call__(
+        self, shape: Tuple[int], return_real_fourier: bool = False, **kwargs
+    ) -> Dict:
         """
         Generate the reconstructed wedge.
 
@@ -344,20 +305,26 @@ class WedgeReconstructed:
         ----------
         shape : tuple of int
             The shape of the reconstruction volume.
-        **kwargs : Dict
+        return_real_fourier : tuple of int
+            Return a shape compliant with rfftn. 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
@@ -369,7 +336,6 @@ class WedgeReconstructed:
             )
 
         ret = func(shape=shape, **func_args)
-
         frequency_cutoff = func_args.get("frequency_cutoff", None)
         if frequency_cutoff is not None:
             frequency_mask = fftfreqn(
@@ -386,17 +352,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 +390,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 +505,99 @@ 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 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)
+    if "data_stopgap_wedgelist" in data:
+        angles = data["data_stopgap_wedgelist"]["_tilt_angle"]
+        weights = data["data_stopgap_wedgelist"]["_exposure"]
+    else:
+        angles = data["data_"]["_wrpAxisAngle"]
+        weights = data["data_"]["_wrpDose"]
+    return {"angles": angles, "weights": weights}
+
+
+def _from_mdoc(filename: str, **kwargs) -> Dict:
+    """
+    Read tilt data from a SerialEM MDOC file.
+
+    Parameters
+    ----------
+    filename : str
+        The path to the text file.
+
+    Returns
+    -------
+    Dict
+        A dictionary with one key for each column.
+    """
+    data = MDOCParser(filename)
+    cumulative_exposure = np.multiply(np.add(1, data["ZValue"]), data["ExposureDose"])
+    return {"angles": data["TiltAngle"], "weights": cumulative_exposure}
+
+
+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:
+        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