pytme 0.2.3__cp311-cp311-macosx_14_0_arm64.whl → 0.2.4__cp311-cp311-macosx_14_0_arm64.whl

tme/matching_memory.py

@@ -1,383 +0,0 @@
-""" 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):
-    """
-    Base class for estimating the memory usage of template matching.
-
-    This class provides a template for estimating memory usage for
-    different matching methods. Users should subclass it and implement the
-    `base_usage` and `per_fork` methods to specify custom memory usage
-    estimates.
-
-    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 the CC fitter.
-
-    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,
-    "FLC2": 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