pytme 0.2.9.post1__cp311-cp311-macosx_15_0_arm64.whl → 0.3.0__cp311-cp311-macosx_15_0_arm64.whl

tme/cli.py

@@ -0,0 +1,126 @@
+#!python3
+"""
+CLI utility functions.
+
+Copyright (c) 2025 European Molecular Biology Laboratory
+
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+
+import argparse
+
+import numpy as np
+from . import __version__
+from .types import BackendArray
+
+
+def match_template(
+    target: BackendArray,
+    template: BackendArray,
+    template_mask: BackendArray = None,
+    score="FLCSphericalMask",
+    rotations=None,
+):
+    """
+    Simple template matching run.
+
+    Parameters
+    ----------
+    target : BackendArray
+        Target array.
+    template : BackendArray
+        Template to be matched against target.
+    template_mask : BackendArray, optional
+        Template mask for normalization, defaults to None.
+    score : str, optional
+        Scoring method to use, defaults to 'FLCSphericalMask'.
+    rotations: BackendArray, optional
+        Rotation matrices with shape (n, d, d), where d is the dimension
+        of the target. Defaults to the identity rotation matrix.
+
+    Returns
+    -------
+    tuple
+        scores : BackendArray
+            Computed cross-correlation scores.
+        offset : BackendArray
+            Offset in target, defaults to 0.
+        rotations : BackendArray
+            Map between translations and rotation indices
+        rotation_mapping : dict
+            Map between rotation indices and rotation matrices
+    """
+    from .matching_data import MatchingData
+    from .analyzer import MaxScoreOverRotations
+    from .matching_exhaustive import scan_subsets, MATCHING_EXHAUSTIVE_REGISTER
+
+    if rotations is None:
+        rotations = np.eye(target.ndim).reshape(1, target.ndim, target.ndim)
+
+    if rotations.shape[-1] != target.ndim:
+        print(
+            f"Dimension of rotation matrix {rotations.shape[-1]} does not "
+            "match target dimension."
+        )
+
+    matching_data = MatchingData(
+        target=target,
+        template=template,
+        template_mask=template_mask,
+        rotations=rotations,
+    )
+    matching_data.template_mask = template_mask
+    matching_setup, matching_score = MATCHING_EXHAUSTIVE_REGISTER[score]
+
+    candidates = list(
+        scan_subsets(
+            matching_data=matching_data,
+            matching_score=matching_score,
+            matching_setup=matching_setup,
+            callback_class=MaxScoreOverRotations,
+            callback_class_args={
+                "score_threshold": -1,
+            },
+            pad_target_edges=True,
+            job_schedule=(1, 1),
+        )
+    )
+    return candidates
+
+
+def sanitize_name(name: str):
+    return name.title().replace("_", " ").replace("-", " ")
+
+
+def print_entry() -> None:
+    width = 80
+    text = f" pytme v{__version__} "
+    padding_total = width - len(text) - 2
+    padding_left = padding_total // 2
+    padding_right = padding_total - padding_left
+
+    print("*" * width)
+    print(f"*{ ' ' * padding_left }{text}{ ' ' * padding_right }*")
+    print("*" * width)
+
+
+def get_func_fullname(func) -> str:
+    """Returns the full name of the given function, including its module."""
+    return f"<function '{func.__module__}.{func.__name__}'>"
+
+
+def print_block(name: str, data: dict, label_width=20) -> None:
+    """Prints a formatted block of information."""
+    print(f"\n> {name}")
+    for key, value in data.items():
+        if isinstance(value, np.ndarray):
+            value = value.shape
+        formatted_value = str(value)
+        print(f"  - {str(key) + ':':<{label_width}} {formatted_value}")
+
+
+def check_positive(value):
+    ivalue = float(value)
+    if ivalue <= 0:
+        raise argparse.ArgumentTypeError("%s is an invalid positive float." % value)
+    return ivalue

tme/density.py

@@ -1,8 +1,9 @@
-""" Representation of N-dimensional densities
+"""
+Representation of N-dimensional densities
 
-    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
@@ -1762,12 +1763,13 @@ class Density:
                 axis=axis,
             )
 
-            arr_ft = np.fft.fftn(self.data)
+            mask, mask_ret = np.where(mask), np.where(mask_ret)
+
+            arr_ft = np.fft.fftn(self.data)[mask]
             arr_ft *= np.prod(ret_shape) / np.prod(self.shape)
             ret_ft = np.zeros(ret_shape, dtype=arr_ft.dtype)
-            ret_ft[mask_ret] = arr_ft[mask]
-            ret.data = np.real(np.fft.ifftn(ret_ft))
-
+            np.add.at(ret_ft, mask_ret, arr_ft)
+            ret.data = np.real(np.fft.ifftn(ret_ft)).astype(self.data.dtype)
         ret.sampling_rate = new_sampling_rate
         return ret
 

tme/filters/__init__.py

@@ -1,6 +1,6 @@
-from .ctf import CTF
+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

@@ -1,8 +1,9 @@
-""" Utilities for the generation of frequency grids.
+"""
+Utilities for the generation of frequency grids.
 
-    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, List, Dict
@@ -14,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):
     """
@@ -70,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.
     """
@@ -230,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
     ----------
@@ -298,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":
@@ -309,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)