pytme 0.3b0__cp311-cp311-macosx_15_0_arm64.whl → 0.3b0.post1__cp311-cp311-macosx_15_0_arm64.whl

tme/filters/wedge.py

@@ -7,8 +7,8 @@ Copyright (c) 2024 European Molecular Biology Laboratory
 Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
-import warnings
 from typing import Tuple, Dict
+from dataclasses import dataclass
 
 import numpy as np
 
@@ -16,8 +16,8 @@ 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 ..parser import XMLParser, StarParser, MDOCParser
 from ._utils import (
     centered_grid,
     frequency_grid_at_angle,
@@ -31,56 +31,28 @@ 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, defaults to 0 (x).
-    opening_axis : int
-        The projection axis, defaults to 2 (z).
-    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
-        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],
-        angles: Tuple[float],
-        weights: Tuple[float],
-        tilt_axis: int = 0,
-        opening_axis: int = 2,
-        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":
@@ -93,6 +65,8 @@ class Wedge(ComposableFilter):
             +-------+---------------------------------------------------------+
             | .xml  | WARP/M XML file                                         |
             +-------+---------------------------------------------------------+
+            | .mdoc | SerialEM file                                           |
+            +-------+---------------------------------------------------------+
             | .*    | Tab-separated file with optional column names           |
             +-------+---------------------------------------------------------+
 
@@ -111,6 +85,8 @@ class Wedge(ComposableFilter):
             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)
@@ -132,6 +108,9 @@ class Wedge(ComposableFilter):
         )
 
     def __call__(self, **kwargs: Dict) -> NDArray:
+        """
+        Returns a Wedge stack of chosen parameters with DC component in the center.
+        """
         func_args = vars(self).copy()
         func_args.update(kwargs)
 
@@ -170,6 +149,7 @@ class Wedge(ComposableFilter):
         return {
             "data": ret,
             "shape": func_args["shape"],
+            "return_real_fourier": func_args.get("return_real_fourier", False),
             "is_multiplicative_filter": True,
         }
 
@@ -196,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
@@ -211,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(
@@ -219,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
@@ -239,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:
         """
@@ -264,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"):
@@ -283,55 +268,35 @@ class Wedge(ComposableFilter):
         return wedges
 
 
+@dataclass
 class WedgeReconstructed:
     """
     Initialize :py:class:`WedgeReconstructed`.
-
-    Parameters
-    ----------
-    angles :tuple of float, optional
-        The tilt angles, defaults to None.
-    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
-        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 = 2,
-        tilt_axis: int = 0,
-        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
+    #: 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
+        self, shape: Tuple[int], return_real_fourier: bool = False, **kwargs
     ) -> Dict:
         """
         Generate the reconstructed wedge.
@@ -341,10 +306,8 @@ 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
+            Return a shape compliant with rfftn. Defaults to False.
+        **kwargs : dict
             Additional keyword arguments.
 
         Returns
@@ -373,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(
@@ -565,7 +527,31 @@ def _from_xml(filename: str, **kwargs) -> Dict:
 
 def _from_star(filename: str, **kwargs) -> Dict:
     """
-    Read tilt data from a tomostar STAR file.
+    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
     ----------
@@ -577,8 +563,9 @@ def _from_star(filename: str, **kwargs) -> Dict:
     Dict
         A dictionary with one key for each column.
     """
-    data = StarParser(filename, delimiter=None)["data_"]
-    return {"angles": data["_wrpAxisAngle"], "weights": data["_wrpDose"]}
+    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:
@@ -604,10 +591,6 @@ def _from_text(filename: str, **kwargs) -> Dict:
     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 "

tme/filters/whitening.py

@@ -24,11 +24,6 @@ class LinearWhiteningFilter(ComposableFilter):
     """
     Compute Fourier power spectrums and perform whitening.
 
-    Parameters
-    ----------
-    **kwargs : Dict, optional
-        Additional keyword arguments.
-
     References
     ----------
     .. [1] de Teresa-Trueba, I.; Goetz, S. K.; Mattausch, A.; Stojanovska, F.; Zimmerli, C. E.;
@@ -39,7 +34,7 @@ class LinearWhiteningFilter(ComposableFilter):
         13375 (2023)
     """
 
-    def __init__(self, **kwargs):
+    def __init__(self, *args, **kwargs):
         pass
 
     @staticmethod

tme/matching_data.py

@@ -175,7 +175,7 @@ class MatchingData:
         target_pad: NDArray = None,
         template_pad: NDArray = None,
         invert_target: bool = False,
-    ) -> "MatchingData":
+    ) -> Tuple["MatchingData", Tuple]:
         """
         Subset class instance based on slices.
 
@@ -194,6 +194,8 @@ class MatchingData:
         -------
         :py:class:`MatchingData`
             Newly allocated subset of class instance.
+        Tuple
+            Translation offset to merge analyzers.
 
         Examples
         --------
@@ -251,8 +253,9 @@ class MatchingData:
         target_offset[mask] = [x.start for x in target_slice]
         mask = np.subtract(1, self._target_batch).astype(bool)
         template_offset = np.zeros(len(self._output_template_shape), dtype=int)
-        template_offset[mask] = [x.start for x in template_slice]
-        ret._translation_offset = tuple(x for x in target_offset)
+        template_offset[mask] = [x.start for x, b in zip(template_slice, mask) if b]
+
+        translation_offset = tuple(x for x in target_offset)
 
         ret.target_filter = self.target_filter
         ret.template_filter = self.template_filter
@@ -262,7 +265,7 @@ class MatchingData:
             template_dim=getattr(self, "_template_dim", None),
         )
 
-        return ret
+        return ret, translation_offset
 
     def to_backend(self):
         """
@@ -323,11 +326,6 @@ class MatchingData:
 
         target_ndim -= len(target_dims)
         template_ndim -= len(template_dims)
-
-        if target_ndim != template_ndim:
-            raise ValueError(
-                f"Dimension mismatch: Target ({target_ndim}) Template ({template_ndim})."
-            )
         self._set_matching_dimension(
             target_dims=target_dims, template_dims=template_dims
         )
@@ -492,29 +490,26 @@ class MatchingData:
     def _fourier_padding(
         target_shape: Tuple[int],
         template_shape: Tuple[int],
-        pad_fourier: bool,
+        pad_target: bool = False,
         batch_mask: Tuple[int] = None,
     ) -> Tuple[Tuple, Tuple, Tuple, Tuple]:
-        fourier_pad = template_shape
-        fourier_shift = np.zeros_like(template_shape)
-
         if batch_mask is None:
             batch_mask = np.zeros_like(template_shape)
         batch_mask = np.asarray(batch_mask)
 
-        if not pad_fourier:
-            fourier_pad = np.ones(len(fourier_pad), dtype=int)
+        fourier_pad = np.ones(len(template_shape), dtype=int)
         fourier_pad = np.multiply(fourier_pad, 1 - batch_mask)
         fourier_pad = np.add(fourier_pad, batch_mask)
 
+        # Avoid padding batch dimensions
         pad_shape = np.maximum(target_shape, template_shape)
+        pad_shape = np.maximum(target_shape, np.multiply(1 - batch_mask, pad_shape))
         ret = be.compute_convolution_shapes(pad_shape, fourier_pad)
         conv_shape, fast_shape, fast_ft_shape = ret
 
         template_mod = np.mod(template_shape, 2)
-        if not pad_fourier:
-            fourier_shift = 1 - np.divide(template_shape, 2).astype(int)
-            fourier_shift = np.subtract(fourier_shift, template_mod)
+        fourier_shift = 1 - np.divide(template_shape, 2).astype(int)
+        fourier_shift = np.subtract(fourier_shift, template_mod)
 
         shape_diff = np.multiply(
             np.subtract(target_shape, template_shape), 1 - batch_mask
@@ -523,34 +518,27 @@ class MatchingData:
         if np.sum(shape_mask):
             shape_shift = np.divide(shape_diff, 2)
             offset = np.mod(shape_diff, 2)
-            if pad_fourier:
-                offset = -np.subtract(
-                    offset,
-                    np.logical_and(np.mod(target_shape, 2) == 0, template_mod == 1),
-                )
-            else:
-                warnings.warn(
-                    "Template is larger than target and padding is turned off. Consider "
-                    "swapping them or activate padding. Correcting the shift for now."
-                )
+            warnings.warn(
+                "Template is larger than target and padding is turned off. Consider "
+                "swapping them or activate padding. Correcting the shift for now."
+            )
             shape_shift = np.multiply(np.add(shape_shift, offset), shape_mask)
             fourier_shift = np.subtract(fourier_shift, shape_shift).astype(int)
 
-        fourier_shift = tuple(np.multiply(fourier_shift, 1 - batch_mask).astype(int))
+        if pad_target:
+            fourier_shift = np.subtract(fourier_shift, np.subtract(1, template_mod))
 
+        fourier_shift = tuple(np.multiply(fourier_shift, 1 - batch_mask).astype(int))
         return tuple(conv_shape), tuple(fast_shape), tuple(fast_ft_shape), fourier_shift
 
-    def fourier_padding(
-        self, pad_fourier: bool = False
-    ) -> Tuple[Tuple, Tuple, Tuple, Tuple]:
+    def fourier_padding(self, pad_target: bool = False) -> Tuple:
         """
         Computes efficient shape four Fourier transforms and potential associated shifts.
 
         Parameters
         ----------
-        pad_fourier : bool, optional
-            If true, returns the shape of the full-convolution defined as sum of target
-            shape and template shape minus one, False by default.
+        pad_target : bool, optional
+            Whether the target has been padded to the full convolution shape.
 
         Returns
         -------
@@ -565,7 +553,7 @@ class MatchingData:
             target_shape=be.to_numpy_array(self._output_target_shape),
             template_shape=be.to_numpy_array(self._output_template_shape),
             batch_mask=be.to_numpy_array(self._batch_mask),
-            pad_fourier=pad_fourier,
+            pad_target=pad_target,
         )
 
     def computation_schedule(

tme/matching_exhaustive.py

@@ -149,7 +149,7 @@ def scan(
     n_jobs: int = 4,
     callback_class: CallbackClass = None,
     callback_class_args: Dict = {},
-    pad_fourier: bool = True,
+    pad_target: bool = True,
     pad_template_filter: bool = True,
     interpolation_order: int = 3,
     jobs_per_callback_class: int = 8,
@@ -176,8 +176,8 @@ def scan(
         Analyzer class pointer to operate on computed scores.
     callback_class_args : dict, optional
         Arguments passed to the callback_class. Default is an empty dictionary.
-    pad_fourier: bool, optional
-        Whether to pad target and template to the full convolution shape.
+    pad_target: bool, optional
+        Whether to pad target to the full convolution shape.
     pad_template_filter: bool, optional
         Whether to pad potential template filters to the full convolution shape.
     interpolation_order : int, optional
@@ -210,17 +210,17 @@ def scan(
     >>> )
 
     """
-    matching_data = matching_data.subset_by_slice(
+    matching_data, translation_offset = matching_data.subset_by_slice(
         target_slice=target_slice,
         template_slice=template_slice,
-        target_pad=matching_data.target_padding(pad_target=pad_fourier),
+        target_pad=matching_data.target_padding(pad_target=pad_target),
     )
 
     matching_data.to_backend()
     template_shape = matching_data._batch_shape(
         matching_data.template.shape, matching_data._template_batch
     )
-    conv, fwd, inv, shift = matching_data.fourier_padding(pad_fourier=False)
+    conv, fwd, inv, shift = matching_data.fourier_padding(pad_target=pad_target)
 
     template_filter = _setup_template_filter_apply_target_filter(
         matching_data=matching_data,
@@ -231,14 +231,14 @@ def scan(
 
     default_callback_args = {
         "shape": fwd,
-        "offset": matching_data._translation_offset,
+        "offset": translation_offset,
         "fourier_shift": shift,
         "fast_shape": fwd,
         "targetshape": matching_data._output_shape,
         "templateshape": template_shape,
         "convolution_shape": conv,
         "thread_safe": n_jobs > 1,
-        "convolution_mode": "valid" if pad_fourier else "same",
+        "convolution_mode": "valid" if pad_target else "same",
         "shm_handler": shm_handler,
         "only_unique_rotations": True,
         "aggregate_axis": matching_data._batch_axis(matching_data._batch_mask),
@@ -425,14 +425,17 @@ def scan_subsets(
     splits = tuple(product(target_splits, template_splits))
 
     outer_jobs, inner_jobs = job_schedule
-    if hasattr(be, "scan"):
+    if be._backend_name == "jax":
+        func = be.scan
+
         corr_scoring = MATCHING_EXHAUSTIVE_REGISTER.get("CORR", (None, None))[1]
-        results = be.scan(
+        results = func(
             matching_data=matching_data,
             splits=splits,
             n_jobs=outer_jobs,
             rotate_mask=matching_score != corr_scoring,
             callback_class=callback_class,
+            callback_class_args=callback_class_args,
         )
     else:
         results = Parallel(n_jobs=outer_jobs, verbose=verbose)(
@@ -447,7 +450,7 @@ def scan_subsets(
                     callback_class=callback_class,
                     callback_class_args=callback_class_args,
                     interpolation_order=interpolation_order,
-                    pad_fourier=pad_target_edges,
+                    pad_target=pad_target_edges,
                     gpu_index=index % outer_jobs,
                     pad_template_filter=pad_template_filter,
                     target_slice=target_split,

tme/matching_scores.py

@@ -593,20 +593,27 @@ def corr_scoring(
         **_fftargs,
     )
 
+    center = be.divide(be.to_backend_array(template.shape) - 1, 2)
     unpadded_slice = tuple(slice(0, stop) for stop in template.shape)
+
+    template_rot = be.zeros(template.shape, be._float_dtype)
     for index in range(rotations.shape[0]):
+        # d+1, d+1 rigid transform matrix from d,d rotation matrix
         rotation = rotations[index]
-        arr = be.fill(arr, 0)
-        arr, _ = be.rigid_transform(
+        matrix = be._rigid_transform_matrix(rotation_matrix=rotation, center=center)
+        template_rot, _ = be.rigid_transform(
             arr=template,
-            rotation_matrix=rotation,
-            out=arr,
-            use_geometric_center=True,
+            rotation_matrix=matrix,
+            out=template_rot,
             order=interpolation_order,
-            cache=False,
+            cache=True,
         )
-        arr = template_filter_func(arr, ft_temp, template_filter)
-        norm_template(arr[unpadded_slice], template_mask, mask_sum)
+
+        template_rot = template_filter_func(template_rot, ft_temp, template_filter)
+        norm_template(template_rot, template_mask, mask_sum)
+
+        arr = be.fill(arr, 0)
+        arr[unpadded_slice] = template_rot
 
         ft_temp = rfftn(arr, ft_temp)
         ft_temp = be.multiply(ft_target, ft_temp, out=ft_temp)
@@ -729,7 +736,7 @@ def flc_scoring(
             out_mask=temp,
             use_geometric_center=True,
             order=interpolation_order,
-            cache=False,
+            cache=True,
         )
 
         n_obs = be.sum(temp)
@@ -875,7 +882,7 @@ def mcc_scoring(
             out_mask=temp,
             use_geometric_center=True,
             order=interpolation_order,
-            cache=False,
+            cache=True,
         )
 
         template_filter_func(template_rot, temp_ft, template_filter)
@@ -1035,7 +1042,8 @@ def flc_scoring2(
             out_mask=tmp_sqz,
             use_geometric_center=True,
             order=interpolation_order,
-            cache=False,
+            cache=True,
+            batched=True,
         )
         n_obs = be.sum(tmp_sqz, axis=data_axes, keepdims=True)
         arr_norm = template_filter_func(arr_sqz, ft_temp, template_filter)
@@ -1155,7 +1163,8 @@ def corr_scoring2(
             out=arr_sqz,
             use_geometric_center=True,
             order=interpolation_order,
-            cache=False,
+            cache=True,
+            batched=True,
         )
         arr_norm = template_filter_func(arr_sqz, ft_sqz, template_filter)
         norm_template(arr_norm[unpadded_slice], template_mask, mask_sum, axis=data_axes)