Start working on IO for DistArrays

Author: mikaem

mpi4py_fft/distarray.py

@@ -322,10 +322,10 @@ def redistribute(self, axis=None, out=None):
         p1, transfer = self.get_pencil_and_transfer(axis)
         if out is None:
             out = DistArray(self.global_shape,
-                               subcomm=p1.subcomm,
-                               dtype=self.dtype,
-                               alignment=axis,
-                               rank=self.rank)
+                            subcomm=p1.subcomm,
+                            dtype=self.dtype,
+                            alignment=axis,
+                            rank=self.rank)
 
         if self.rank == 0:
             transfer.forward(self, out)

mpi4py_fft/io/file_base.py

@@ -8,8 +8,9 @@ class FileBase(object):
 
     Parameters
     ----------
-    T : PFFT
-        Instance of a :class:`.PFFT` class.
+    global
+    u : Distributed array, optional
+        Instance of :class:`.DistArray`
     domain : sequence, optional
         The spatial domain. Sequence of either
 
@@ -18,11 +19,15 @@ class FileBase(object):
             - Arrays of coordinates, e.g., np.linspace(0, 2*pi, N). One
               array per dimension.
     """
-    def __init__(self, T, domain=None, **kw):
+    def __init__(self, global_shape=None, u=None, rank=None, domain=None, **kw):
         self.f = None
         self.filename = None
-        self.T = T
-        self.domain = domain if domain is not None else ((0, 2*np.pi),)*T.dimensions()
+        if u is not None:
+            assert isinstance(u, DistArray)
+        self.global_shape = global_shape if global_shape is not None else u.global_shape
+        self.rank = rank if rank is not None else u.rank
+        self.dimensions = len(self.global_shape[self.rank:])
+        self.domain = domain if domain is not None else ((0, 2*np.pi),)*self.dimensions
 
     def write(self, step, fields, **kw):
         """Write snapshot ``step`` of ``fields`` to file

mpi4py_fft/io/h5py_file.py

@@ -14,8 +14,7 @@ class HDF5File(FileBase):
     h5name : str
         Name of hdf5 file to be created.
     T : PFFT
-        Instance of a :class:`PFFT` class. Must be the same as the space
-        used for storing with :class:`HDF5Writer.write`.
+        Instance of a :class:`PFFT` class.
     domain : Sequence, optional
         The spatial domain. Sequence of either
 

mpi4py_fft/io/nc_file.py

@@ -17,8 +17,7 @@ class NCFile(FileBase):
     ncname : str
         Name of netcdf file to be created
     T : PFFT
-        Instance of a :class:`.PFFT` class. Must be the same as the space
-        used for storing with :class:`NCWriter.write`
+        Instance of a :class:`.PFFT` class.
     domain : Sequence
         The spatial domain. Sequence of either
 

setup.py

@@ -90,7 +90,7 @@ def version():
     generate_extensions(fftw_libs)
     setup(name="mpi4py-fft",
           version=version(),
-          description="mpi4py-fft -- FFT with MPI",
+          description="mpi4py-fft -- Parallel Fast Fourier Transforms (FFTs) using MPI for Python",
           long_description=long_description,
           author="Lisandro Dalcin and Mikael Mortensen",
           url='https://bitbucket.org/mpi4py/mpi4py-fft',
@@ -112,5 +112,6 @@ def version():
               ],
           ext_modules=get_extensions(fftw_libs),
           install_requires=["mpi4py", "numpy"],
-          setup_requires=["setuptools>=18.0", "cython>=0.25"]
+          setup_requires=["setuptools>=18.0", "cython>=0.25"],
+          keywords=['Python', 'FFTW', 'FFT', 'DCT', 'DST', 'MPI']
           )