Source code for fluidfft

"""Efficient and easy Fast Fourier Transform for Python
=======================================================

The fft and related `operators` classes are in the two subpackages

.. autosummary::
   :toctree:

   fft2d
   fft3d

The two commands ``fluidfft-bench`` and ``fluidfft-bench-analysis`` can be used to
benchmark the classes on particular cases and computers. These commands are
implemented in the following modules

.. autosummary::
   :toctree:

   bench
   bench_analysis

This root module provides two helper functions to import fft classes and create
fft objects:

.. autofunction:: get_plugins

.. autofunction:: get_methods

.. autofunction:: import_fft_class

.. autofunction:: create_fft_object

"""

import importlib
import os
import subprocess
import sys
import logging

if sys.version_info < (3, 10):
    from importlib_metadata import entry_points, EntryPoint
else:
    from importlib.metadata import entry_points, EntryPoint

from fluiddyn.util import mpi

from fluidfft._version import __version__

try:
    from pyfftw import empty_aligned, byte_align
except ImportError:
    import numpy as np

    empty_aligned = np.empty

    def byte_align(values, *args):
        """False byte_align function used when pyfftw can not be imported"""
        return values


__citation__ = r"""
@article{fluiddyn,
doi = {10.5334/jors.237},
year = {2019},
publisher = {Ubiquity Press,  Ltd.},
volume = {7},
author = {Pierre Augier and Ashwin Vishnu Mohanan and Cyrille Bonamy},
title = {{FluidDyn}: A Python Open-Source Framework for Research and Teaching in Fluid Dynamics
    by Simulations,  Experiments and Data Processing},
journal = {Journal of Open Research Software}
}

@article{fluidfft,
doi = {10.5334/jors.238},
year = {2019},
publisher = {Ubiquity Press,  Ltd.},
volume = {7},
author = {Ashwin Vishnu Mohanan and Cyrille Bonamy and Pierre Augier},
title = {{FluidFFT}: Common {API} (C$\mathplus\mathplus$ and Python)
    for Fast Fourier Transform {HPC} Libraries},
journal = {Journal of Open Research Software}
}
"""


__all__ = [
    "__citation__",
    "__version__",
    "byte_align",
    "create_fft_object",
    "empty_aligned",
    "get_module_fullname_from_method",
    "get_plugins",
    "get_methods",
    "import_fft_class",
]


_plugins = None


[docs] def get_plugins(reload=False, ndim=None, sequential=None): """Discover the fluidfft plugins installed""" global _plugins if _plugins is None or reload: _plugins = entry_points(group="fluidfft.plugins") if not _plugins: raise RuntimeError("No Fluidfft plugins were found.") if ndim is None and sequential is None: return _plugins if ndim is None: index = 6 prefix = "" elif ndim in (2, 3): index = 0 prefix = f"fft{ndim}d." else: raise ValueError(f"Unsupported value for {ndim = }") if sequential is not None and not sequential: prefix += "mpi_" elif sequential: prefix += "with_" return tuple( plugin for plugin in _plugins if plugin.name[index:].startswith(prefix) )
[docs] def get_methods(ndim=None, sequential=None): """Get available methods""" plugins = get_plugins(ndim=ndim, sequential=sequential) return set(plug.name for plug in plugins)
def get_module_fullname_from_method(method): """Get the module name from a method string Parameters ---------- method : str Name of module or string characterizing a method. """ plugins = get_plugins() selected_plugins = plugins.select(name=method) if len(selected_plugins) == 0: raise ValueError( f"Cannot find a fluidfft plugin for {method = }. {plugins}" ) elif len(selected_plugins) > 1: logging.warning( f"{len(selected_plugins)} plugins were found for {method = }" ) return selected_plugins[method].value def _normalize_method_name(method): """Normalize a method name""" if method == "sequential": method = "fft2d.with_fftw2d" elif method.startswith("fluidfft:"): method = method.removeprefix("fluidfft:") return method def _check_failure(method): """Check if a tiny fft maker can be created""" if not any(method.endswith(postfix) for postfix in ("pfft", "p3dfft")): return False # for few methods, try before real import because importing can lead to # a fatal error (Illegal instruction) if mpi.rank == 0: if mpi.nb_proc > 1: # We need to filter out the MPI environment variables. # Fragile because it is specific to MPI implementations env = { key: value for key, value in os.environ.items() if not ("MPI" in key or key.startswith("PMI_")) } else: env = os.environ try: # TODO: capture stdout and stderr and include last line in case of failure subprocess.check_call( [ sys.executable, "-c", f"from fluidfft import create_fft_object as c; c('{method}', 2, 2, 2, check=False)", ], env=env, shell=False, ) failure = False except subprocess.CalledProcessError: failure = True else: failure = None if mpi.nb_proc > 1: failure = mpi.comm.bcast(failure, root=0) return failure
[docs] def import_fft_class(method, raise_import_error=True, check=True): """Import a fft class. Parameters ---------- method : str Name of module or string characterizing a method. It has to correspond to a module of fluidfft. The first part "fluidfft." of the module "path" can be omitted. raise_import_error : {True}, False If raise_import_error == False and if there is an import error, the function handles the error and returns None. Returns ------- The corresponding FFT class. """ if isinstance(method, EntryPoint): module_fullname = method.value method = method.name else: method = _normalize_method_name(method) module_fullname = get_module_fullname_from_method(method) if check: failure = _check_failure(method) if failure: if not raise_import_error: mpi.printby0("ImportError during check:", module_fullname) return None else: raise ImportError(module_fullname) try: mod = importlib.import_module(module_fullname) except ImportError: if raise_import_error: raise mpi.printby0("ImportError:", module_fullname) return None return mod.FFTclass
def _get_classes(ndim, sequential): plugins = get_plugins(ndim=ndim, sequential=sequential) return { plugin.name: import_fft_class(plugin, raise_import_error=False) for plugin in plugins }
[docs] def create_fft_object(method, n0, n1, n2=None, check=True): """Helper for creating fft objects. Parameters ---------- method : str Name of module or string characterizing a method. It has to correspond to a module of fluidfft. The first part "fluidfft." of the module "path" can be omitted. n0, n1, n2 : int Dimensions of the real space array (in sequential). Returns ------- The corresponding FFT object. """ cls = import_fft_class(method, check=check) str_module = cls.__module__ if n2 is None and "fft3d" in str_module: raise ValueError("Arguments incompatible") if n2 is not None and "fft2d" in str_module: n2 = None if n2 is None: return cls(n0, n1) else: return cls(n0, n1, n2)