Add FFT methods to utils.interface

Author: pulkkins

pysteps/utils/fft.py

@@ -6,27 +6,28 @@
     # TODO: Caching and multithreading are currently disabled because they give 
     # a segfault with dask.
     pyfftw.interfaces.cache.disable()
-    pyfftw_kwargs = {"threads":1, "planner_effort":"FFTW_ESTIMATE"}
     pyfftw_imported = True
 except ImportError:
     pyfftw_imported = False
-try:
-    import scipy.fftpack as scipy_fft
-    scipy_fft_kwargs = {}
-except ImportError:
-    scipy_imported = False
-try:
-    import numpy.fft as numpy_fft
-    numpy_fft_kwargs = {}
-except ImportError:
-    numpy_imported = False
+import scipy.fftpack as scipy_fft
+import numpy.fft as numpy_fft
+
+from pysteps.exceptions import MissingOptionalDependency
+
+# use numpy implementation of rfft2/irfft2 because they have not been
+# implemented in scipy.fftpack
+scipy_fft.rfft2  = numpy_fft.rfft2
+scipy_fft.irfft2 = numpy_fft.irfft2
 
 def get_method(name):
+    """Return a callable function for the FFT method corresponding to the given name."""
     if name == "numpy":
-        return numpy_fft
+        return numpy_fft,{}
     elif name == "scipy":
-        return scipy_fft
+        return scipy_fft,{}
     elif name == "pyfftw":
-        return pyfftw_fft
+        if not pyfftw_imported:
+            raise MissingOptionalDependency("pyfftw is required but it is not installed")
+        return pyfftw_fft,{"threads":1, "planner_effort":"FFTW_ESTIMATE"}
     else:
         raise ValueError("unknown method %s, the available methods are 'numpy', 'scipy' and 'pyfftw'" % name)

pysteps/utils/interface.py

@@ -2,10 +2,12 @@
 from . import conversion
 from . import transformation
 from . import dimension
+from . import fft
 
 def get_method(name):
     """Return a callable function for the utility method corresponding to the
-    given name.\n\
+    given name. For the FFT methods, the return value is a two-element tuple
+    containing the function and a dictionary of keyword arguments.\n\
 
     Conversion methods:
 
@@ -50,6 +52,18 @@ def get_method(name):
     |  upscale          | upscale the field                                      |
     +-------------------+--------------------------------------------------------+
 
+    FFT methods (wrappers to different implementations):
+
+    +-------------------+--------------------------------------------------------+
+    |     Name          |              Description                               |
+    +===================+========================================================+
+    |  numpy_fft        | numpy.fft                                              |
+    +-------------------+--------------------------------------------------------+
+    |  scipy_fft        | scipy.fftpack                                          |
+    +-------------------+--------------------------------------------------------+
+    |  pyfftw_fft       | pyfftw.interfaces.numpy_fft                            |
+    +-------------------+--------------------------------------------------------+
+
     """
 
     if name is None:
@@ -82,6 +96,10 @@ def donothing(R, metadata, *args, **kwargs):
     methods_objects["clip"]         = dimension.clip_domain
     methods_objects["square"]       = dimension.square_domain
     methods_objects["upscale"]      = dimension.aggregate_fields_space
+    # FFT methods
+    methods_objects["numpy_fft"]    = fft.get_method("numpy")
+    methods_objects["scipy_fft"]    = fft.get_method("scipy")
+    methods_objects["pyfftw_fft"]   = fft.get_method("pyfftw")
 
     try:
         return methods_objects[name]