pytme 0.2.9__cp311-cp311-macosx_15_0_arm64.whl

tme/backends/__init__.py

@@ -0,0 +1,171 @@
+""" pyTME backend manager.
+
+    Copyright (c) 2023 European Molecular Biology Laboratory
+
+    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
+"""
+
+from typing import Dict, List
+from importlib.util import find_spec
+
+from .matching_backend import MatchingBackend
+from .npfftw_backend import NumpyFFTWBackend
+from .pytorch_backend import PytorchBackend
+from .cupy_backend import CupyBackend
+from .mlx_backend import MLXBackend
+from .jax_backend import JaxBackend
+
+
+class BackendManager:
+    """
+    Manager for template matching backends.
+
+    This class serves as an interface to various computational backends (e.g.,
+    CPU, GPU). It allows users to seamlessly swap between different backend
+    implementations while preserving the consistency and functionality of the API.
+    Direct attribute and method calls to the manager are delegated to the current
+    active backend.
+
+    Examples
+    --------
+    >>> from tme.backends import backend
+    >>> backend.multiply(arr1, arr2)
+    # This will use the default NumpyFFTWBackend's multiply method
+
+    >>> backend.change_backend("pytorch")
+    >>> backend.multiply(arr1, arr2)
+    # This will use the pytorchs multiply method
+
+    >>> backend.available_backends()
+    # Backends available on your system
+
+    Notes
+    -----
+    The backend has to be reinitialzed when using fork-based parallelism.
+    """
+
+    def __init__(self):
+        self._BACKEND_REGISTRY = {
+            "numpyfftw": NumpyFFTWBackend,
+            "pytorch": PytorchBackend,
+            "cupy": CupyBackend,
+            "mlx": MLXBackend,
+            "jax": JaxBackend,
+        }
+        self._backend = NumpyFFTWBackend()
+        self._backend_name = "numpyfftw"
+        self._backend_args = {}
+
+    def __repr__(self):
+        return f"<BackendManager: using {self._backend_name}>"
+
+    def __getattr__(self, name):
+        return getattr(self._backend, name)
+
+    def __dir__(self) -> List:
+        """
+        Return a list of attributes available in this object,
+        including those from the backend.
+
+        Returns
+        -------
+        list
+            Sorted list of attributes.
+        """
+        base_attributes = []
+        base_attributes.extend(dir(self.__class__))
+        base_attributes.extend(self.__dict__.keys())
+        base_attributes.extend(dir(self._backend))
+        return sorted(base_attributes)
+
+    def add_backend(self, backend_name: str, backend_class: type):
+        """
+        Adds a custom backend to the registry.
+
+        Parameters
+        ----------
+        backend_name : str
+            Name by which the backend can be referenced.
+        backend_class : :py:class:`MatchingBackend`
+            An instance of the backend to be added.
+
+        Examples
+        --------
+        >>> from tme.backends import backend as be
+        >>> from tme.backends import NumpyFFTWBackend
+        >>> be.add_backend("numpy2", NumpyFFTWBackend)
+
+        Raises
+        ------
+        ValueError
+            If the provided backend_instance does not inherit from
+            :py:class:`MatchingBackend`.
+        """
+        if not issubclass(backend_class, MatchingBackend):
+            raise ValueError("backend_class needs to inherit from MatchingBackend.")
+        self._BACKEND_REGISTRY[backend_name] = backend_class
+
+    def change_backend(self, backend_name: str, **backend_kwargs: Dict) -> None:
+        """
+        Change the current computation backend.
+
+        Parameters
+        ----------
+        backend_name : str
+            Name of the new backend that should be used.
+        **backend_kwargs : dict, optional
+            Parameters passed to __init__ method of backend.
+
+        Examples
+        --------
+        >>> from tme.backends import backend as be
+        >>> be.change_backend("cupy")
+
+        Raises
+        ------
+        NotImplementedError
+            If no backend is found with the provided name.
+        """
+        if backend_name not in self._BACKEND_REGISTRY:
+            available_backends = ", ".join(self.available_backends())
+            raise NotImplementedError(
+                f"Available backends are {available_backends} - not {backend_name}."
+            )
+        self._backend = self._BACKEND_REGISTRY[backend_name](**backend_kwargs)
+        self._backend_name = backend_name
+        self._backend_args = backend_kwargs
+
+    def available_backends(self) -> List[str]:
+        """
+        Returns available computation backends.
+
+        Returns
+        -------
+        list of str
+            Available computation backends.
+
+        Examples
+        --------
+        >>> from tme.backends import backend as be
+        >>> be.available_backends()
+        """
+        # This is an approximation but avoids runtime polution
+        _dependencies = {
+            "numpyfftw": "numpy",
+            "cupy": "cupy",
+            "pytorch": "torch",
+            "mlx": "mlx",
+            "jax": "jax",
+        }
+        available_backends = []
+        for name, backend in self._BACKEND_REGISTRY.items():
+            if name not in _dependencies:
+                continue
+
+            if find_spec(_dependencies[name]) is not None:
+                available_backends.append(name)
+
+        return available_backends
+
+
+backend = BackendManager()