Fixing import export wisdom for MPI

Author: mikaem

mpi4py_fft/fftw/factory.py

@@ -1,5 +1,6 @@
 import six
 import numpy as np
+from mpi4py import MPI
 from . import fftw_xfftn
 try:
     from . import fftwf_xfftn
@@ -16,6 +17,8 @@
     if v is not None:
         fftlib[k] = v
 
+comm = MPI.COMM_WORLD
+
 def get_planned_FFT(input_array, output_array, axes=(-1,), kind=FFTW_FORWARD,
                     threads=1, flags=(FFTW_MEASURE,), normalize=1):
     """Return instance of transform class
@@ -84,17 +87,20 @@ def export_wisdom(filename):
     Note
     ----
     Wisdom is stored for all precisions available: float, double and long
-    double, using, respectively, prefix ``F_``, ``D_`` and ``G_``. Wisdom is
-    imported using :func:`.import_wisdom`.
+    double, using, respectively, prefix ``Fn_``, ``Dn_`` and ``Gn_``, where
+    n is the rank of the processor.
+    Wisdom is imported using :func:`.import_wisdom`, which must be called
+    with the same MPI configuration as used with :func:`.export_wisdom`.
 
     See also
     --------
     :func:`.import_wisdom`
 
     """
+    rank = str(comm.Get_rank())
     e = []
     for key, lib in six.iteritems(fftlib):
-        e.append(lib.export_wisdom(bytearray(key+'_'+filename, 'utf-8')))
+        e.append(lib.export_wisdom(bytearray(key+rank+'_'+filename, 'utf-8')))
     assert np.all(np.array(e) == 1), "Not able to export wisdom {}".format(filename)
 
 def import_wisdom(filename):
@@ -108,17 +114,22 @@ def import_wisdom(filename):
     Note
     ----
     Wisdom is imported for all available precisions: float, double and long
-    double, using, respectively, prefix ``F_``, ``D_`` and ``G_``. Wisdom is
-    exported using :func:`.export_wisdom`.
+    double, using, respectively, prefix ``Fn_``, ``Dn_`` and ``Gn_``, where
+    n is the rank of the processor.
+    Wisdom is exported using :func:`.export_wisdom`.
+    Note that importing wisdom only works when using the same MPI configuration
+    as used with :func:`.export_wisdom`.
+
 
     See also
     --------
     :func:`.export_wisdom`
 
     """
+    rank = str(comm.Get_rank())
     e = []
     for key, lib in six.iteritems(fftlib):
-        e.append(lib.import_wisdom(bytearray(key+'_'+filename, 'utf-8')))
+        e.append(lib.import_wisdom(bytearray(key+rank+'_'+filename, 'utf-8')))
     assert np.all(np.array(e) == 1), "Not able to import wisdom {}".format(filename)
 
 def forget_wisdom():

tests/test_fftw.py

@@ -1,11 +1,18 @@
 from __future__ import print_function
-import six
 from time import time
+import six
 import numpy as np
-import scipy
-import pyfftw
+from scipy.fftpack import dctn as scipy_dctn
+from scipy.fftpack import dstn as scipy_dstn
+import scipy.fftpack
 from mpi4py_fft import fftw
 
+has_pyfftw = True
+try:
+    import pyfftw
+except ImportError:
+    has_pyfftw = False
+
 abstol = dict(f=5e-4, d=1e-12, g=1e-14)
 
 kinds = {'dst4': fftw.FFTW_RODFT11, # no scipy to compare with
@@ -46,26 +53,27 @@ def test_fftw():
                             outshape = list(shape)
                             outshape[axes[-1]] = shape[axes[-1]]//2+1
                             output_array = fftw.aligned(outshape, dtype=typecode.upper())
-                            oa = output_array if typecode=='d' else None # Test for both types of signature
+                            oa = output_array if typecode == 'd' else None # Test for both types of signature
                             rfftn = fftw.rfftn(input_array, None, axes, threads, fflags, output_array=oa)
                             A = np.random.random(shape).astype(typecode)
                             input_array[:] = A
                             B = rfftn()
                             assert id(B) == id(rfftn.output_array)
-                            B2 = pyfftw.interfaces.numpy_fft.rfftn(input_array, axes=axes)
-                            assert allclose(B, B2), np.linalg.norm(B-B2)
-                            ia = input_array if typecode=='d' else None
+                            if has_pyfftw:
+                                B2 = pyfftw.interfaces.numpy_fft.rfftn(input_array, axes=axes)
+                                assert allclose(B, B2), np.linalg.norm(B-B2)
+                            ia = input_array if typecode == 'd' else None
                             sa = np.take(input_array.shape, axes) if shape[axes[-1]] % 2 == 1 else None
                             irfftn = fftw.irfftn(output_array, sa, axes, threads, iflags, output_array=ia)
-                            irfftn.input_array[...] = B2
+                            irfftn.input_array[...] = B
                             A2 = irfftn(normalize=True)
                             assert allclose(A, A2), np.linalg.norm(A-A2)
                             hfftn = fftw.hfftn(output_array, sa, axes, threads, fflags, output_array=ia)
-                            hfftn.input_array[...] = B2
+                            hfftn.input_array[...] = B
                             AC = hfftn().copy()
                             ihfftn = fftw.ihfftn(input_array, None, axes, threads, iflags, output_array=oa)
                             A2 = ihfftn(AC, implicit=False, normalize=True)
-                            assert allclose(A2, B2), print(np.linalg.norm(A2-B2))
+                            assert allclose(A2, B), print(np.linalg.norm(A2-B))
 
                             # c2c
                             input_array = fftw.aligned(shape, dtype=typecode.upper())
@@ -79,8 +87,9 @@ def test_fftw():
                             ifftn.input_array[...] = D
                             C2 = ifftn(normalize=True)
                             assert allclose(C, C2), np.linalg.norm(C-C2)
-                            D2 = pyfftw.interfaces.numpy_fft.fftn(C, axes=axes)
-                            assert allclose(D, D2), np.linalg.norm(D-D2)
+                            if has_pyfftw:
+                                D2 = pyfftw.interfaces.numpy_fft.fftn(C, axes=axes)
+                                assert allclose(D, D2), np.linalg.norm(D-D2)
 
                             # r2r
                             input_array = fftw.aligned(shape, dtype=typecode)
@@ -93,7 +102,7 @@ def test_fftw():
                                 A2 = idct(B, implicit=True, normalize=True)
                                 assert allclose(A, A2), np.linalg.norm(A-A2)
                                 if typecode is not 'g' and not type is 4:
-                                    B2 = scipy.fftpack.dctn(A, axes=axes, type=type)
+                                    B2 = scipy_dctn(A, axes=axes, type=type)
                                     assert allclose(B, B2), np.linalg.norm(B-B2)
 
                                 dst = fftw.dstn(input_array, None, axes, type, threads, fflags, output_array=oa)
@@ -102,7 +111,7 @@ def test_fftw():
                                 A2 = idst(B, implicit=True, normalize=True)
                                 assert allclose(A, A2), np.linalg.norm(A-A2)
                                 if typecode is not 'g' and not type is 4:
-                                    B2 = scipy.fftpack.dstn(A, axes=axes, type=type)
+                                    B2 = scipy_dstn(A, axes=axes, type=type)
                                     assert allclose(B, B2), np.linalg.norm(B-B2)
 
                             # Different r2r transforms along all axes. Just pick
@@ -128,8 +137,8 @@ def test_fftw():
 
 def test_wisdom():
     # Test a simple export/import call
-    fftw.export_wisdom('wisdom.dat')
-    fftw.import_wisdom('wisdom.dat')
+    fftw.export_wisdom('newwisdom.dat')
+    fftw.import_wisdom('newwisdom.dat')
     fftw.forget_wisdom()
 
 def test_timelimit():
@@ -151,4 +160,3 @@ def test_timelimit():
     test_fftw()
     test_wisdom()
     test_timelimit()
-

tests/test_libfft.py

@@ -3,6 +3,12 @@
 import numpy as np
 from mpi4py_fft.libfft import FFT
 
+has_pyfftw = True
+try:
+    import pyfftw
+except ImportError:
+    has_pyfftw = False
+
 abstol = dict(f=5e-5, d=1e-14, g=1e-15)
 
 def allclose(a, b):
@@ -17,6 +23,8 @@ def test_libfft():
     types = 'fdgFDG'
 
     for use_pyfftw in (False, True):
+        if has_pyfftw is False and use_pyfftw is True:
+            continue
         t0 = 0
         for typecode in types:
             for dim in dims:
@@ -49,6 +57,8 @@ def test_libfft():
     # difficult to initialize. We solve this problem by making one extra
     # transform
     for use_pyfftw in (True, False):
+        if has_pyfftw is False and use_pyfftw is True:
+            continue
         for padding in (1.5, 2.0):
             for typecode in types:
                 for dim in dims:

tests/test_mpifft.py

@@ -1,12 +1,17 @@
 from __future__ import print_function
 import functools
 import numpy as np
-import pyfftw
 from mpi4py import MPI
 from mpi4py_fft.mpifft import PFFT
 from mpi4py_fft.pencil import Subcomm
 from mpi4py_fft import fftw, Function
 
+has_pyfftw = True
+try:
+    import pyfftw
+except ImportError:
+    has_pyfftw = False
+
 abstol = dict(f=0.1, d=2e-10, g=1e-10)
 
 def allclose(a, b):
@@ -64,6 +69,8 @@ def test_mpifft():
                     padding = False
                     for collapse in (True, False):
                         for use_pyfftw in (False, True):
+                            if has_pyfftw is False and use_pyfftw is True:
+                                continue
                             transforms = None
                             if dim < 3:
                                 allaxes = [None, (-1,), (-2,),
@@ -159,6 +166,8 @@ def test_mpifft():
 
                     padding = [1.5]*len(shape)
                     for use_pyfftw in (True, False):
+                        if has_pyfftw is False and use_pyfftw is True:
+                            continue
                         if dim < 3:
                             allaxes = [None, (-1,), (-2,),
                                        (-1, -2,), (-2, -1),