Fixing generate_xdmf and FFTNumpy

Author: mikaem

mpi4py_fft/io/generate_xdmf.py

@@ -40,10 +40,10 @@ def get_geometry(kind=0, dim=2):
 
         return """<Geometry Type="VXVY">
           <DataItem Format="HDF" NumberType="Float" Precision="{0}" Dimensions="{1}">
-            {3}:/mesh/{4}
+            {3}:{6}/mesh/{4}
           </DataItem>
           <DataItem Format="HDF" NumberType="Float" Precision="{0}" Dimensions="{2}">
-            {3}:/mesh/{5}
+            {3}:{6}/mesh/{5}
           </DataItem>
         </Geometry>"""
 
@@ -60,13 +60,13 @@ def get_geometry(kind=0, dim=2):
 
         return """<Geometry Type="VXVYVZ">
           <DataItem Format="HDF" NumberType="Float" Precision="{0}" Dimensions="{3}">
-            {4}:/mesh/{5}
+            {4}:{8}/mesh/{5}
           </DataItem>
           <DataItem Format="HDF" NumberType="Float" Precision="{0}" Dimensions="{2}">
-            {4}:/mesh/{6}
+            {4}:{8}/mesh/{6}
           </DataItem>
           <DataItem Format="HDF" NumberType="Float" Precision="{0}" Dimensions="{1}">
-            {4}:/mesh/{7}
+            {4}:{8}/mesh/{7}
           </DataItem>
         </Geometry>"""
 
@@ -233,17 +233,17 @@ def generate_xdmf(h5filename, periodic=True, order='paraview'):
 
                     if ndim == 2 and ('slice' not in slices or len(f[group].attrs.get('shape')) > 3):
                         if order.lower() == 'paraview':
-                            sig = (prec, N[0], N[1], h5filename, cc[0], cc[1])
+                            sig = (prec, N[0], N[1], h5filename, cc[0], cc[1], group)
                         else:
-                            sig = (prec, N[1], N[0], h5filename, cc[1], cc[0])
+                            sig = (prec, N[1], N[0], h5filename, cc[1], cc[0], group)
                     else:
                         if ndim == 2: # 2D slice in 3D domain
                             pos = f[group+"/mesh/x{}".format(kk)][sl]
                             z = re.findall(r'<DataItem(.*?)</DataItem>', geo, re.DOTALL)
                             geo = geo.replace(z[2-kk], ' Format="XML" NumberType="Float" Precision="{0}" Dimensions="{%d}">\n           {%d}\n          '%(1+kk, 7-kk))
                             cc = list(cc)
                             cc.insert(kk, pos)
-                        sig = (prec, N[0], N[1], N[2], h5filename, cc[2], cc[1], cc[0])
+                        sig = (prec, N[0], N[1], N[2], h5filename, cc[2], cc[1], cc[0], group)
                     geometry[slices] = geo.format(*sig)
                     topology[slices] = get_topology(N, kind=1)
                 grid[slices] = ''

mpi4py_fft/libfft.py

@@ -1,3 +1,4 @@
+import functools
 import numpy as np
 from . import fftw
 
@@ -371,28 +372,37 @@ class FFTNumPy(FFTBase): #pragma: no cover
 
     """
 
-    def __init__(self, shape, axes=None, dtype=float, padding=False, **kw):
+    def __init__(self, shape, axes=None, dtype=float, padding=False,
+                 transforms=None, **kw):
         FFTBase.__init__(self, shape, axes, dtype, padding)
         typecode = self.dtype.char
 
         self.sizes = list(np.take(self.shape, self.axes))
         arrayA = np.zeros(self.shape, self.dtype)
-        if self.real_transform:
-            axis = self.axes[-1]
-            self.shape[axis] = self.shape[axis]//2 + 1
-            arrayB = np.zeros(self.shape, typecode.upper())
-            fwd = np.fft.rfftn
-            bck = np.fft.irfftn
+        transforms = {} if transforms is None else transforms
+        if tuple(self.axes) in transforms:
+            fwd, bck = transforms[tuple(self.axes)]
+            arrayB = fwd(arrayA, axes=self.axes).astype(typecode)
+            self.fwd = functools.partial(fwd, shape=self.sizes)
+            self.bck = functools.partial(bck, shape=self.sizes)
+
         else:
-            arrayB = np.zeros(self.shape, typecode)
-            fwd = np.fft.fftn
-            bck = np.fft.ifftn
+            if self.real_transform:
+                fwd = np.fft.rfftn
+                bck = np.fft.irfftn
+                arrayB = fwd(arrayA, s=self.sizes, axes=self.axes).astype(typecode.upper())
+                self.shape = arrayB.shape
+            else:
+                fwd = np.fft.fftn
+                bck = np.fft.ifftn
+                arrayB = fwd(arrayA, s=self.sizes, axes=self.axes).astype(typecode)
+            self.fwd = functools.partial(fwd, s=self.sizes)
+            self.bck = functools.partial(bck, s=self.sizes)
 
-        fwd.input_array = arrayA
-        fwd.output_array = arrayB
-        bck.input_array = arrayB
-        bck.output_array = arrayA
-        self.fwd, self.bck = fwd, bck
+        self.fwd_input_array = arrayA
+        self.fwd_output_array = arrayB
+        self.bck_input_array = arrayB
+        self.bck_output_array = arrayA
 
         self.padding_factor = 1
         if padding is not False:
@@ -407,14 +417,14 @@ def __init__(self, shape, axes=None, dtype=float, padding=False, **kw):
             self.backward = _Xfftn_wrap(self._backward, arrayB, arrayA)
 
     def _forward(self, **kw):
-        self.fwd.output_array[:] = self.fwd(self.fwd.input_array, s=self.sizes,
+        self.fwd_output_array[:] = self.fwd(self.fwd_input_array,
                                             axes=self.axes, **kw)
-        self._truncation_forward(self.fwd.output_array, self.forward.output_array)
+        self._truncation_forward(self.fwd_output_array, self.forward.output_array)
         return self.forward.output_array
 
     def _backward(self, **kw):
-        self._padding_backward(self.backward.input_array, self.bck.input_array)
-        self.backward.output_array[:] = self.bck(self.bck.input_array, s=self.sizes,
+        self._padding_backward(self.backward.input_array, self.bck_input_array)
+        self.backward.output_array[:] = self.bck(self.bck_input_array,
                                                  axes=self.axes, **kw)
         return self.backward.output_array
 

mpi4py_fft/mpifft.py

@@ -288,8 +288,8 @@ def __init__(self, comm, shape=None, axes=None, dtype=float, slab=False,
 
         axes = self.axes[-1]
         pencil = Pencil(self.subcomm, shape, axes[-1])
-        xfftn = FFT(pencil.subshape, axes, dtype, padding, use_pyfftw,
-                    transforms, **kw)
+        xfftn = FFT(pencil.subshape, axes, dtype, padding, use_pyfftw=use_pyfftw,
+                    transforms=transforms, **kw)
         self.xfftn.append(xfftn)
         self.pencil[0] = pencilA = pencil
         if not shape[axes[-1]] == xfftn.forward.output_array.shape[axes[-1]]:
@@ -300,8 +300,8 @@ def __init__(self, comm, shape=None, axes=None, dtype=float, slab=False,
         for axes in reversed(self.axes[:-1]):
             pencilB = pencilA.pencil(axes[-1])
             transAB = pencilA.transfer(pencilB, dtype)
-            xfftn = FFT(pencilB.subshape, axes, dtype, padding, use_pyfftw,
-                        transforms, **kw)
+            xfftn = FFT(pencilB.subshape, axes, dtype, padding, use_pyfftw=use_pyfftw,
+                        transforms=transforms, **kw)
             self.xfftn.append(xfftn)
             self.transfer.append(transAB)
             pencilA = pencilB

tests/test_libfft.py

@@ -10,7 +10,7 @@
 except ImportError:
     has_pyfftw = False
 
-abstol = dict(f=5e-5, d=1e-14, g=1e-15)
+abstol = dict(f=5e-5, d=1e-14, g=1e-14)
 
 def allclose(a, b):
     atol = abstol[a.dtype.char.lower()]
@@ -91,7 +91,7 @@ def test_libfft():
 
                             B.fill(0)
                             B = fft.forward(A, B)
-                            assert allclose(B, X)
+                            assert allclose(B, X), np.linalg.norm(B-X)
 
 
 if __name__ == '__main__':