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

tme/backends/_cupy_utils.py

@@ -1,34 +1,35 @@
-""" Utility functions for cupy backend.
+"""
+Utility functions for cupy backend.
 
-    The functions spline_filter, _prepad_for_spline_filter, _filter_input,
-    _get_coord_affine_batched and affine_transform are largely copied from
-    cupyx.scipy.ndimage which operates under the following license
+The functions spline_filter, _prepad_for_spline_filter, _filter_input,
+_get_coord_affine_batched and affine_transform are largely copied from
+cupyx.scipy.ndimage which operates under the following license
 
-    Copyright (c) 2015 Preferred Infrastructure, Inc.
-    Copyright (c) 2015 Preferred Networks, Inc.
+Copyright (c) 2015 Preferred Infrastructure, Inc.
+Copyright (c) 2015 Preferred Networks, Inc.
 
-    Permission is hereby granted, free of charge, to any person obtaining a copy
-    of this software and associated documentation files (the "Software"), to deal
-    in the Software without restriction, including without limitation the rights
-    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-    copies of the Software, and to permit persons to whom the Software is
-    furnished to do so, subject to the following conditions:
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
 
-    The above copyright notice and this permission notice shall be included in
-    all copies or substantial portions of the Software.
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
 
-    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
-    THE SOFTWARE.
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
 
-    I have since extended the functionality of the cupyx.scipy.ndimage functions
-    in question to support batched inputs.
+I have since extended the functionality of the cupyx.scipy.ndimage functions
+in question to support batched inputs.
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 import numpy

tme/backends/_jax_utils.py

@@ -1,8 +1,9 @@
-""" Utility functions for jax backend.
+"""
+Utility functions for jax backend.
 
-    Copyright (c) 2023-2024 European Molecular Biology Laboratory
+Copyright (c) 2023-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

tme/backends/cupy_backend.py

@@ -1,8 +1,9 @@
-""" Backend using cupy for template matching.
+"""
+Backend using cupy for template matching.
 
-    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
@@ -113,16 +114,6 @@ class CupyBackend(NumpyFFTWBackend):
     def unravel_index(self, indices, shape):
         return self._array_backend.unravel_index(indices=indices, dims=shape)
 
-    def unique(self, ar, axis=None, *args, **kwargs):
-        if axis is None:
-            return self._array_backend.unique(ar=ar, axis=axis, *args, **kwargs)
-
-        warnings.warn("Axis argument not yet supported in CupY, falling back to NumPy.")
-        ret = np.unique(ar=self.to_numpy_array(ar), axis=axis, *args, **kwargs)
-        if not isinstance(ret, tuple):
-            return self.to_backend_array(ret)
-        return tuple(self.to_backend_array(k) for k in ret)
-
     def build_fft(
         self,
         fwd_shape: Tuple[int],

tme/backends/jax_backend.py

@@ -1,8 +1,9 @@
-""" Backend using jax for template matching.
+"""
+Backend using jax for template matching.
 
-    Copyright (c) 2023-2024 European Molecular Biology Laboratory
+Copyright (c) 2023-2024 European Molecular Biology Laboratory
 
-    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
 from functools import wraps
@@ -288,13 +289,10 @@ class JaxBackend(NumpyFFTWBackend):
             for index in range(scores.shape[0]):
                 temp = callback_class(
                     shape=scores.shape,
-                    scores=scores[index],
-                    rotations=rotations[index],
-                    thread_safe=False,
                     offset=translation_offsets[index],
                 )
-                temp.rotation_mapping = rotation_mapping
-                ret.append(tuple(temp._postprocess(**analyzer_args)))
+                state = (scores, rotations, rotation_mapping)
+                ret.append(temp.result(state, **analyzer_args))
 
         return ret
 

tme/backends/matching_backend.py

@@ -1,8 +1,9 @@
-""" Strategy pattern to allow for flexible array / FFT backends.
+"""
+Strategy pattern to allow for flexible array / FFT backends.
 
-    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>
 """
 
 from abc import ABC, abstractmethod

tme/backends/mlx_backend.py

@@ -1,8 +1,9 @@
-""" Backend using Apple's MLX library for template matching.
+"""
+Backend using Apple's MLX library for template matching.
 
-    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, Callable

tme/backends/npfftw_backend.py

@@ -1,8 +1,9 @@
-""" Backend using numpy and pyFFTW for template matching.
+"""
+Backend using numpy and pyFFTW for template matching.
 
-    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 os
@@ -159,8 +160,9 @@ class NumpyFFTWBackend(_NumpyWrapper, MatchingBackend):
         temp = self._array_backend.zeros(1, dtype=dtype)
         return temp.nbytes
 
-    @staticmethod
-    def astype(arr, dtype: Type) -> NDArray:
+    def astype(self, arr, dtype: Type) -> NDArray:
+        if self._array_backend.iscomplexobj(arr):
+            arr = arr.real
         return arr.astype(dtype)
 
     @staticmethod

tme/backends/pytorch_backend.py

@@ -1,9 +1,10 @@
-""" Backend using pytorch and optionally GPU acceleration for
-    template matching.
+"""
+Backend using pytorch and optionally GPU acceleration for
+template matching.
 
-    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>
 """
 
 from typing import Tuple, Callable
@@ -134,6 +135,15 @@ class PytorchBackend(NumpyFFTWBackend):
     def astype(self, arr: TorchTensor, dtype: type) -> TorchTensor:
         return arr.to(dtype)
 
+    @staticmethod
+    def at(arr, idx, value) -> NDArray:
+        arr[idx] = value
+        return arr
+
+    @staticmethod
+    def addat(arr, indices, *args, **kwargs) -> NDArray:
+        return arr.index_put_(indices, *args, accumulate=True, **kwargs)
+
     def flip(self, a, axis, **kwargs):
         return self._array_backend.flip(input=a, dims=axis, **kwargs)
 

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

tme/filters/__init__.py

@@ -1,4 +1,4 @@
-from .ctf import CTF
+from .ctf import CTF, CTFReconstructed
 from .compose import Compose, ComposableFilter
 from .bandpass import BandPassFilter
 from .whitening import LinearWhiteningFilter

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

tme/filters/bandpass.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
@@ -225,6 +226,7 @@ class BandPassFilter(ComposableFilter):
 
         return {
             "data": be.to_backend_array(mask),
-            "sampling_rate": func_args.get("sampling_rate", 1),
+            "shape": func_args["shape"],
+            "return_real_fourier": func_args.get("return_real_fourier"),
             "is_multiplicative_filter": True,
         }

tme/filters/compose.py

@@ -1,8 +1,9 @@
-""" Combine filters using an interface analogous to pytorch's Compose.
+"""
+Combine filters using an interface analogous to pytorch's Compose.
 
-    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
@@ -60,7 +61,7 @@ class Compose:
 
 class ComposableFilter(ABC):
     """
-    Strategy class for composable filters.
+    Base class for composable filters.
     """
 
     @abstractmethod