pytme 0.2.1__cp311-cp311-macosx_14_0_arm64.whl → 0.2.2__cp311-cp311-macosx_14_0_arm64.whl

tme/memory.py

@@ -0,0 +1,377 @@
+""" Compute memory consumption of template matching components.
+
+    Copyright (c) 2023 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+
+from abc import ABC, abstractmethod
+from typing import Tuple
+
+import numpy as np
+from pyfftw import next_fast_len
+
+
+class MatchingMemoryUsage(ABC):
+    """
+    Class specification for estimating the memory requirements of template matching.
+
+    Parameters
+    ----------
+    fast_shape : tuple of int
+        Shape of the real array.
+    ft_shape : tuple of int
+        Shape of the complex array.
+    float_nbytes : int
+        Number of bytes of the used float, e.g. 4 for float32.
+    complex_nbytes : int
+        Number of bytes of the used complex, e.g. 8 for complex64.
+    integer_nbytes : int
+        Number of bytes of the used integer, e.g. 4 for int32.
+
+    Attributes
+    ----------
+    real_array_size : int
+        Number of elements in real array.
+    complex_array_size : int
+        Number of elements in complex array.
+    float_nbytes : int
+        Number of bytes of the used float, e.g. 4 for float32.
+    complex_nbytes : int
+        Number of bytes of the used complex, e.g. 8 for complex64.
+    integer_nbytes : int
+        Number of bytes of the used integer, e.g. 4 for int32.
+
+    Methods
+    -------
+    base_usage():
+        Returns the base memory usage in bytes.
+    per_fork():
+        Returns the memory usage in bytes per fork.
+    """
+
+    def __init__(
+        self,
+        fast_shape: Tuple[int],
+        ft_shape: Tuple[int],
+        float_nbytes: int,
+        complex_nbytes: int,
+        integer_nbytes: int,
+    ):
+        self.real_array_size = np.prod(fast_shape)
+        self.complex_array_size = np.prod(ft_shape)
+        self.float_nbytes = float_nbytes
+        self.complex_nbytes = complex_nbytes
+        self.integer_nbytes = integer_nbytes
+
+    @abstractmethod
+    def base_usage(self) -> int:
+        """Return the base memory usage in bytes."""
+
+    @abstractmethod
+    def per_fork(self) -> int:
+        """Return the memory usage per fork in bytes."""
+
+
+class CCMemoryUsage(MatchingMemoryUsage):
+    """
+    Memory usage estimation for CC scoring.
+
+    See Also
+    --------
+    :py:meth:`tme.matching_exhaustive.cc_setup`.
+    """
+
+    def base_usage(self) -> int:
+        float_arrays = self.real_array_size * self.float_nbytes
+        complex_arrays = self.complex_array_size * self.complex_nbytes
+        return float_arrays + complex_arrays
+
+    def per_fork(self) -> int:
+        float_arrays = self.real_array_size * self.float_nbytes
+        complex_arrays = self.complex_array_size * self.complex_nbytes
+        return float_arrays + complex_arrays
+
+
+class LCCMemoryUsage(CCMemoryUsage):
+    """
+    Memory usage estimation for LCC scoring.
+    See Also
+    --------
+    :py:meth:`tme.matching_exhaustive.lcc_setup`.
+    """
+
+
+class CORRMemoryUsage(MatchingMemoryUsage):
+    """
+    Memory usage estimation for CORR scoring.
+
+    See Also
+    --------
+    :py:meth:`tme.matching_exhaustive.corr_setup`.
+    """
+
+    def base_usage(self) -> int:
+        float_arrays = self.real_array_size * self.float_nbytes * 4
+        complex_arrays = self.complex_array_size * self.complex_nbytes
+        return float_arrays + complex_arrays
+
+    def per_fork(self) -> int:
+        float_arrays = self.real_array_size * self.float_nbytes
+        complex_arrays = self.complex_array_size * self.complex_nbytes
+        return float_arrays + complex_arrays
+
+
+class CAMMemoryUsage(CORRMemoryUsage):
+    """
+    Memory usage estimation for CAM scoring.
+
+    See Also
+    --------
+    :py:meth:`tme.matching_exhaustive.cam_setup`.
+    """
+
+
+class FLCSphericalMaskMemoryUsage(CORRMemoryUsage):
+    """
+    Memory usage estimation for FLCMSphericalMask scoring.
+
+    See Also
+    --------
+    :py:meth:`tme.matching_exhaustive.flcSphericalMask_setup`.
+    """
+
+
+class FLCMemoryUsage(MatchingMemoryUsage):
+    """
+    Memory usage estimation for FLC scoring.
+
+    See Also
+    --------
+    :py:meth:`tme.matching_exhaustive.flc_setup`.
+    """
+
+    def base_usage(self) -> int:
+        float_arrays = self.real_array_size * self.float_nbytes * 2
+        complex_arrays = self.complex_array_size * self.complex_nbytes * 2
+        return float_arrays + complex_arrays
+
+    def per_fork(self) -> int:
+        float_arrays = self.real_array_size * self.float_nbytes * 3
+        complex_arrays = self.complex_array_size * self.complex_nbytes * 2
+        return float_arrays + complex_arrays
+
+
+class MCCMemoryUsage(MatchingMemoryUsage):
+    """
+    Memory usage estimation for MCC scoring.
+
+    See Also
+    --------
+    :py:meth:`tme.matching_exhaustive.mcc_setup`.
+    """
+
+    def base_usage(self) -> int:
+        float_arrays = self.real_array_size * self.float_nbytes * 2
+        complex_arrays = self.complex_array_size * self.complex_nbytes * 3
+        return float_arrays + complex_arrays
+
+    def per_fork(self) -> int:
+        float_arrays = self.real_array_size * self.float_nbytes * 6
+        complex_arrays = self.complex_array_size * self.complex_nbytes
+        return float_arrays + complex_arrays
+
+
+class MaxScoreOverRotationsMemoryUsage(MatchingMemoryUsage):
+    """
+    Memory usage estimation MaxScoreOverRotations Analyzer.
+
+    See Also
+    --------
+    :py:class:`tme.analyzer.MaxScoreOverRotations`.
+    """
+
+    def base_usage(self) -> int:
+        float_arrays = self.real_array_size * self.float_nbytes * 2
+        return float_arrays
+
+    def per_fork(self) -> int:
+        return 0
+
+
+class PeakCallerMaximumFilterMemoryUsage(MatchingMemoryUsage):
+    """
+    Memory usage estimation MaxScoreOverRotations Analyzer.
+
+    See Also
+    --------
+    :py:class:`tme.analyzer.PeakCallerMaximumFilter`.
+    """
+
+    def base_usage(self) -> int:
+        float_arrays = self.real_array_size * self.float_nbytes
+        return float_arrays
+
+    def per_fork(self) -> int:
+        float_arrays = self.real_array_size * self.float_nbytes
+        return float_arrays
+
+
+class CupyBackendMemoryUsage(MatchingMemoryUsage):
+    """
+    Memory usage estimation for CupyBackend.
+
+    See Also
+    --------
+    :py:class:`tme.backends.CupyBackend`.
+    """
+
+    def base_usage(self) -> int:
+        # FFT plans, overhead from assigning FFT result, rotation interpolation
+        complex_arrays = self.real_array_size * self.complex_nbytes * 3
+        float_arrays = self.complex_array_size * self.float_nbytes * 2
+        return float_arrays + complex_arrays
+
+    def per_fork(self) -> int:
+        return 0
+
+
+def _compute_convolution_shapes(
+    arr1_shape: Tuple[int], arr2_shape: Tuple[int]
+) -> Tuple[Tuple[int], Tuple[int], Tuple[int]]:
+    """
+    Computes regular, optimized and fourier convolution shape.
+
+    Parameters
+    ----------
+    arr1_shape : tuple
+        Tuple of integers corresponding to array1 shape.
+    arr2_shape : tuple
+        Tuple of integers corresponding to array2 shape.
+
+    Returns
+    -------
+    tuple
+        Tuple with regular convolution shape, convolution shape optimized for faster
+        fourier transform, shape of the forward fourier transform
+        (see :py:meth:`build_fft`).
+    """
+    convolution_shape = np.add(arr1_shape, arr2_shape) - 1
+    fast_shape = [next_fast_len(x) for x in convolution_shape]
+    fast_ft_shape = list(fast_shape[:-1]) + [fast_shape[-1] // 2 + 1]
+
+    return convolution_shape, fast_shape, fast_ft_shape
+
+
+MATCHING_MEMORY_REGISTRY = {
+    "CC": CCMemoryUsage,
+    "LCC": LCCMemoryUsage,
+    "CORR": CORRMemoryUsage,
+    "CAM": CAMMemoryUsage,
+    "MCC": MCCMemoryUsage,
+    "FLCSphericalMask": FLCSphericalMaskMemoryUsage,
+    "FLC": FLCMemoryUsage,
+    "MaxScoreOverRotations": MaxScoreOverRotationsMemoryUsage,
+    "PeakCallerMaximumFilter": PeakCallerMaximumFilterMemoryUsage,
+    "cupy": CupyBackendMemoryUsage,
+    "pytorch": CupyBackendMemoryUsage,
+}
+
+
+def estimate_ram_usage(
+    shape1: Tuple[int],
+    shape2: Tuple[int],
+    matching_method: str,
+    ncores: int,
+    analyzer_method: str = None,
+    backend: str = None,
+    float_nbytes: int = 4,
+    complex_nbytes: int = 8,
+    integer_nbytes: int = 4,
+) -> int:
+    """
+    Estimate the RAM usage for a given convolution operation based on input shapes,
+    matching_method, and number of cores.
+
+    Parameters
+    ----------
+    shape1 : tuple
+        The shape of the input target.
+    shape2 : tuple
+        The shape of the input template.
+    matching_method : str
+        The method used for the operation.
+    is_gpu : bool, optional
+        Whether the computation is performed on GPU. This factors in FFT
+        plan caching.
+    analyzer_method : str, optional
+        The method used for score analysis.
+    backend : str, optional
+        Backend used for computation.
+    ncores : int
+        The number of CPU cores used for the operation.
+    float_nbytes : int
+        Number of bytes of the used float, e.g. 4 for float32.
+    complex_nbytes : int
+        Number of bytes of the used complex, e.g. 8 for complex64.
+    integer_nbytes : int
+        Number of bytes of the used integer, e.g. 4 for int32.
+
+    Returns
+    -------
+    int
+        The estimated RAM usage for the operation in bytes.
+
+    Notes
+    -----
+        Residual memory from other objects that may remain allocated during
+        template matching, e.g. the full sized target when using splitting,
+        are not considered by this function.
+
+    Raises
+    ------
+    ValueError
+        If an unsupported matching_methode is provided.
+    """
+    if matching_method not in MATCHING_MEMORY_REGISTRY:
+        raise ValueError(
+            f"Supported options are {','.join(MATCHING_MEMORY_REGISTRY.keys())}"
+        )
+
+    convolution_shape, fast_shape, ft_shape = _compute_convolution_shapes(
+        shape1, shape2
+    )
+
+    memory_instance = MATCHING_MEMORY_REGISTRY[matching_method](
+        fast_shape=fast_shape,
+        ft_shape=ft_shape,
+        float_nbytes=float_nbytes,
+        complex_nbytes=complex_nbytes,
+        integer_nbytes=integer_nbytes,
+    )
+
+    nbytes = memory_instance.base_usage() + memory_instance.per_fork() * ncores
+
+    analyzer_instance = MATCHING_MEMORY_REGISTRY.get(analyzer_method, None)
+    if analyzer_instance is not None:
+        analyzer_instance = analyzer_instance(
+            fast_shape=fast_shape,
+            ft_shape=ft_shape,
+            float_nbytes=float_nbytes,
+            complex_nbytes=complex_nbytes,
+            integer_nbytes=integer_nbytes,
+        )
+        nbytes += analyzer_instance.base_usage() + analyzer_instance.per_fork() * ncores
+
+    backend_instance = MATCHING_MEMORY_REGISTRY.get(backend, None)
+    if backend_instance is not None:
+        backend_instance = backend_instance(
+            fast_shape=fast_shape,
+            ft_shape=ft_shape,
+            float_nbytes=float_nbytes,
+            complex_nbytes=complex_nbytes,
+            integer_nbytes=integer_nbytes,
+        )
+        nbytes += backend_instance.base_usage() + backend_instance.per_fork() * ncores
+
+    return nbytes

tme/orientations.py

@@ -62,16 +62,16 @@ class Orientations:
         Array with additional orientation details (n, ).
     """
 
-    #: Return a numpy array with translations of each orientation (n x d).
+    #: Array with translations of each orientation (n, d).
     translations: np.ndarray
 
-    #: Return a numpy array with euler angles of each orientation in zxy format (n x d).
+    #: Array with zyx euler angles of each orientation (n, d).
     rotations: np.ndarray
 
-    #: Return a numpy array with the score of each orientation (n, ).
+    #: Array with scores of each orientation (n, ).
     scores: np.ndarray
 
-    #: Return a numpy array with additional orientation details (n, ).
+    #: Array with additional details of each orientation(n, ).
     details: np.ndarray
 
     def __post_init__(self):
@@ -130,9 +130,21 @@ class Orientations:
             "scores",
             "details",
         )
-        kwargs = {attr: getattr(self, attr)[indices] for attr in attributes}
+        kwargs = {attr: getattr(self, attr)[indices].copy() for attr in attributes}
         return self.__class__(**kwargs)
 
+    def copy(self) -> "Orientations":
+        """
+        Create a copy of the current class instance.
+
+        Returns
+        -------
+        :py:class:`Orientations`
+            Copy of the class instance.
+        """
+        indices = np.arange(self.scores.size)
+        return self[indices]
+
     def to_file(self, filename: str, file_format: type = None, **kwargs) -> None:
         """
         Save the current class instance to a file in the specified format.
@@ -146,7 +158,7 @@ class Orientations:
             the file_format from the typical extension. Supported formats are
 
             +---------------+----------------------------------------------------+
-            | text          | pyTME's standard tab-separated orientations file   |
+            | text          | pytme's standard tab-separated orientations file   |
             +---------------+----------------------------------------------------+
             | relion        | Creates a STAR file of orientations                |
             +---------------+----------------------------------------------------+
@@ -207,11 +219,11 @@ class Orientations:
         with open(filename, mode="w", encoding="utf-8") as ofile:
             _ = ofile.write(f"{header}\n")
             for translation, angles, score, detail in self:
-                translation_string = "\t".join([str(x) for x in translation])
-                angle_string = "\t".join([str(x) for x in angles])
-                _ = ofile.write(
-                    f"{translation_string}\t{angle_string}\t{score}\t{detail}\n"
+                out_string = (
+                    "\t".join([str(x) for x in (*translation, *angles, score, detail)])
+                    + "\n"
                 )
+                _ = ofile.write(out_string)
         return None
 
     def _to_dynamo_tbl(
@@ -465,8 +477,10 @@ class Orientations:
 
         Notes
         -----
-        The text file is expected to have a header and data in columns corresponding to
-        z, y, x, euler_z, euler_y, euler_x, score, detail.
+        The text file is expected to have a header and data in columns. Colums containing
+        the name euler are considered to specify rotations. The second last and last
+        column correspond to score and detail. Its possible to only specify translations,
+        in this case the remaining columns will be filled with trivial values.
         """
         with open(filename, mode="r", encoding="utf-8") as infile:
             data = [x.strip().split("\t") for x in infile.read().split("\n")]
@@ -493,6 +507,32 @@ class Orientations:
         score = np.array(score)
         detail = np.array(detail)
 
+        if translation.shape[1] == len(header):
+            rotation = np.zeros(translation.shape, dtype=np.float32)
+            score = np.zeros(translation.shape[0], dtype=np.float32)
+            detail = np.zeros(translation.shape[0], dtype=np.float32) - 1
+
+        if rotation.size == 0 and translation.shape[0] != 0:
+            rotation = np.zeros(translation.shape, dtype=np.float32)
+
+        header_order = tuple(x for x in header if x in ascii_lowercase)
+        header_order = zip(header_order, range(len(header_order)))
+        sort_order = tuple(
+            x[1] for x in sorted(header_order, key=lambda x: x[0], reverse=True)
+        )
+        translation = translation[..., sort_order]
+
+        header_order = tuple(
+            x
+            for x in header
+            if "euler" in x and x.replace("euler_", "") in ascii_lowercase
+        )
+        header_order = zip(header_order, range(len(header_order)))
+        sort_order = tuple(
+            x[1] for x in sorted(header_order, key=lambda x: x[0], reverse=True)
+        )
+        rotation = rotation[..., sort_order]
+
         return translation, rotation, score, detail
 
     @staticmethod

tme/parser.py

@@ -137,8 +137,7 @@ class Parser(ABC):
 
 class PDBParser(Parser):
     """
-    A Parser subclass for converting PDB file data into a dictionary representation.
-    This class is specifically designed to work with PDB file format.
+    Convert PDB file data into a dictionary representation [1]_.
 
     References
     ----------
@@ -228,8 +227,8 @@ class PDBParser(Parser):
 
 class MMCIFParser(Parser):
     """
-    A Parser subclass for converting MMCIF file data into a dictionary representation.
-    This implementation heavily relies on the atomium library [1]_.
+    Convert MMCIF file data into a dictionary representation. This implementation
+    heavily relies on the atomium library [1]_.
 
     References
     ----------

tme/preprocessing/_utils.py

@@ -5,12 +5,13 @@
     Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
 """
 
-from typing import Tuple
+from typing import Tuple, List
 
 import numpy as np
-from numpy.typing import NDArray
 
-from ..backends import backend
+from ..backends import backend as be
+from ..backends import NumpyFFTWBackend
+from ..types import BackendArray, NDArray
 from ..matching_utils import euler_to_rotationmatrix
 
 
@@ -93,18 +94,27 @@ def frequency_grid_at_angle(
     tilt_shape = compute_tilt_shape(
         shape=shape, opening_axis=opening_axis, reduce_dim=False
     )
-    index_grid = centered_grid(shape=tilt_shape)
+
+    if angle == 0:
+        index_grid = fftfreqn(
+            tuple(x for x in tilt_shape if x != 1),
+            sampling_rate=1,
+            compute_euclidean_norm=True,
+        )
+
     if angle != 0:
         angles = np.zeros(len(shape))
         angles[tilt_axis] = angle
         rotation_matrix = euler_to_rotationmatrix(np.roll(angles, opening_axis - 1))
+
+        index_grid = fftfreqn(tilt_shape, sampling_rate=None)
         index_grid = np.einsum("ij,j...->i...", rotation_matrix, index_grid)
+        norm = np.multiply(sampling_rate, shape).astype(int)
 
-    norm = np.divide(1, 2 * sampling_rate * np.divide(shape, 2).astype(int))
+        index_grid = np.divide(index_grid.T, norm).T
+        index_grid = np.squeeze(index_grid)
+        index_grid = np.linalg.norm(index_grid, axis=(0))
 
-    index_grid = np.multiply(index_grid.T, norm).T
-    index_grid = np.squeeze(index_grid)
-    index_grid = np.linalg.norm(index_grid, axis=(0))
     return index_grid
 
 
@@ -113,9 +123,10 @@ def fftfreqn(
     sampling_rate: Tuple[float],
     compute_euclidean_norm: bool = False,
     shape_is_real_fourier: bool = False,
+    return_sparse_grid: bool = False,
 ) -> NDArray:
     """
-    Generate the n-dimensional discrete Fourier Transform sample frequencies.
+    Generate the n-dimensional discrete Fourier transform sample frequencies.
 
     Parameters:
     -----------
@@ -133,56 +144,74 @@ def fftfreqn(
     NDArray
         The sample frequencies.
     """
-    center = backend.astype(backend.divide(shape, 2), backend._int_dtype)
-
-    norm = np.ones(len(shape))
+    # There is no real need to have these operations on GPU right now
+    temp_backend = NumpyFFTWBackend()
+    norm = temp_backend.full(len(shape), fill_value=1)
+    center = temp_backend.astype(temp_backend.divide(shape, 2), temp_backend._int_dtype)
     if sampling_rate is not None:
-        norm = backend.astype(backend.multiply(shape, sampling_rate), int)
+        norm = temp_backend.astype(temp_backend.multiply(shape, sampling_rate), int)
 
     if shape_is_real_fourier:
-        center[-1] = 0
-        norm[-1] = 1
+        center[-1], norm[-1] = 0, 1
         if sampling_rate is not None:
             norm[-1] = (shape[-1] - 1) * 2 * sampling_rate
 
-    indices = backend.transpose(backend.indices(shape))
-    indices -= center
-    indices = backend.divide(indices, norm)
-    indices = backend.transpose(indices)
+    grids = []
+    for i, x in enumerate(shape):
+        baseline_dims = tuple(1 if i != t else x for t in range(len(shape)))
+        grid = (temp_backend.arange(x) - center[i]) / norm[i]
+        grids.append(temp_backend.reshape(grid, baseline_dims))
 
     if compute_euclidean_norm:
-        indices = backend.square(indices)
-        indices = backend.sum(indices, axis=0)
-        backend.sqrt(indices, out=indices)
+        grids = sum(temp_backend.square(x) for x in grids)
+        grids = temp_backend.sqrt(grids, out=grids)
+        return grids
+
+    if return_sparse_grid:
+        return grids
 
-    return indices
+    grid_flesh = temp_backend.full(shape, fill_value=1)
+    grids = temp_backend.stack(tuple(grid * grid_flesh for grid in grids))
 
+    return grids
 
-def crop_real_fourier(data: NDArray) -> NDArray:
+
+def crop_real_fourier(data: BackendArray) -> BackendArray:
     """
     Crop the real part of a Fourier transform.
 
     Parameters:
     -----------
-    data : NDArray
+    data : BackendArray
         The Fourier transformed data.
 
     Returns:
     --------
-    NDArray
+    BackendArray
         The cropped data.
     """
     stop = 1 + (data.shape[-1] // 2)
     return data[..., :stop]
 
 
-def shift_fourier(data: NDArray, shape_is_real_fourier: bool = False):
-    shift = backend.add(
-        backend.astype(backend.divide(data.shape, 2), int),
-        backend.mod(data.shape, 2),
-    )
+def compute_fourier_shape(
+    shape: Tuple[int], shape_is_real_fourier: bool = False
+) -> List[int]:
+    if shape_is_real_fourier:
+        return shape
+    shape = [int(x) for x in shape]
+    shape[-1] = 1 + shape[-1] // 2
+    return shape
+
+
+def shift_fourier(
+    data: BackendArray, shape_is_real_fourier: bool = False
+) -> BackendArray:
+    shape = be.to_backend_array(data.shape)
+    shift = be.add(be.divide(shape, 2), be.mod(shape, 2))
+    shift = [int(x) for x in shift]
     if shape_is_real_fourier:
         shift[-1] = 0
 
-    data = backend.roll(data, shift, tuple(i for i in range(len(shift))))
+    data = be.roll(data, shift, tuple(i for i in range(len(shift))))
     return data

tme/preprocessing/compose.py

@@ -7,7 +7,7 @@
 
 from typing import Tuple, Dict
 
-from tme.backends import backend
+from tme.backends import backend as be
 
 
 class Compose:
@@ -42,9 +42,13 @@ class Compose:
             kwargs.update(meta)
             ret = transform(**kwargs)
 
+            if "data" not in ret:
+                continue
+
             if ret.get("is_multiplicative_filter", False):
-                backend.multiply(ret["data"], meta["data"], out=ret["data"])
-                ret["merge"] = None
+                prev_data = meta.pop("data")
+                ret["data"] = be.multiply(ret["data"], prev_data, out=ret["data"])
+                ret["merge"], prev_data = None, None
 
             meta = ret