"""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)