Working on IO with rank of distarrays

Author: mikaem

docs/source/io.rst

@@ -4,7 +4,7 @@ Storing datafiles
 mpi4py-fft works with regular Numpy arrays. However, since arrays in parallel
 can become very large, and the arrays live on multiple processors, we require
 parallel IO capabilities that goes beyond Numpys regular methods.
-In the :mod:`mpi4py_fft.utilities` module there are two helper classes for dumping
+In the :mod:`mpi4py_fft.io` module there are two helper classes for dumping
 dataarrays to either `HDF5 <https://www.hdf5.org>`_ or
 `NetCDF <https://www.unidata.ucar.edu/software/netcdf/>`_ format:
 
@@ -17,46 +17,50 @@ reads data in parallel. A simple example of usage is::
     from mpi4py import MPI
     import numpy as np
     from mpi4py_fft import PFFT, HDF5File, NCFile, newDistArray
-
     N = (128, 256, 512)
     T = PFFT(MPI.COMM_WORLD, N)
     u = newDistArray(T, forward_output=False)
     v = newDistArray(T, forward_output=False, val=2)
-    u[:] = np.random.random(N)
-
+    u[:] = np.random.random(u.shape)
+    # Store by first creating output files
     fields = {'u': [u], 'v': [v]}
-    f0 = HDF5File('h5test.h5', T)
-    f1 = NCFile('nctest.nc', T)
+    f0 = HDF5File('h5test.h5', global_shape=N, mode='w')
+    f1 = NCFile('nctest.nc', global_shape=N, mode='w')
     f0.write(0, fields)
     f1.write(0, fields)
     v[:] = 3
     f0.write(1, fields)
     f1.write(1, fields)
+    # Alternatively, just use write method of each distributed array
+    u.write('h5test.h5', 'u', step=2)
+    v.write('h5test.h5', 'v', step=2)
+    u.write('nctest.nc', 'u', step=2)
+    v.write('nctest.nc', 'v', step=2)
 
-Note that we are creating two datafiles ``h5test.h5`` and ``nctest.nc``,
+Note that we are here creating two datafiles ``h5test.h5`` and ``nctest.nc``,
 for storing in HDF5 or NetCDF4 formats respectively. Normally, one would be
 satisfied using only one format, so this is only for illustration. We store
-the fields ``u`` and ``v`` using method ``write`` on two different occasions,
-so the datafiles will contain two snapshots of each field ``u`` and ``v``.
+the fields ``u`` and ``v`` on three different occasions,
+so the datafiles will contain three snapshots of each field ``u`` and ``v``.
 
 The stored dataarrays can be retrieved later on::
 
-    f0 = HDF5File('h5test.h5', T, mode='r')
-    f1 = NCFile('nctest.nc', T, mode='r')
     u0 = newDistArray(T, forward_output=False)
     u1 = newDistArray(T, forward_output=False)
-    f0.read(u0, 'u', 0)
-    f0.read(u1, 'u', 1)
-    f1.read(u0, 'u', 0)
-    f1.read(u1, 'u', 1)
+    u0.read('h5test.h5', 'u', 0)
+    u1.read('h5test.h5', 'u', 1)
+    # or alternatively for netcdf
+    #u0.read('nctest.nc', 'u', 0)
+    #u1.read('nctest.nc', 'u', 1)
+
 
 Note that one does not have to use the same number of processors when
 retrieving the data as when they were stored.
 
 It is also possible to store only parts of the, potentially large, arrays.
 Any chosen slice may be stored, using a *global* view of the arrays::
 
-    f2 = HDF5File('variousfields.h5', T, mode='w')
+    f2 = HDF5File('variousfields.h5', global_shape=N, mode='w')
     fields = {'u': [u,
                     (u, [slice(None), slice(None), 4]),
                     (u, [5, 5, slice(None)])],
@@ -65,6 +69,8 @@ Any chosen slice may be stored, using a *global* view of the arrays::
     f2.write(0, fields)
     f2.write(1, fields)
     f2.write(2, fields)
+    # or, using write method of field, e.g.
+    #u.write('variousfields.h5', 'u', 0, [slice(None), slice(None), 4])
 
 This will lead to an hdf5-file with groups::
 
@@ -107,14 +113,14 @@ two different ways when creating the datafiles:
        originates from the origin, with lengths :math:`\pi, 2\pi, 3\pi`, can be
        given as::
 
-        f0 = HDF5File('filename.h5', T, domain=((0, pi), (0, 2*np.pi), (0, 3*np.pi)))
+        f0 = HDF5File('filename.h5', global_shape=N, domain=((0, pi), (0, 2*np.pi), (0, 3*np.pi)))
 
     2) A sequence of arrays giving the coordinates for each dimension. For example::
 
         d = (np.arange(N[0], dtype=np.float)*1*np.pi/N[0],
              np.arange(N[1], dtype=np.float)*2*np.pi/N[1],
              np.arange(N[2], dtype=np.float)*2*np.pi/N[2])
-        f0 = HDF5File('filename.h5', T, domain=d)
+        f0 = HDF5File('filename.h5', global_shape=N, domain=d)
 
 With NetCDF4 the layout is somewhat different. For ``variousfields`` above,
 if we were using :class:`.NCFile` instead of :class:`.HDF5File`,
@@ -147,9 +153,9 @@ opened with `Visit <https://www.visitusers.org>`_.
 
 To view the HDF5-files we first need to generate some light-weight *xdmf*-files that can
 be understood by both Paraview and Visit. To generate such files, simply throw the
-module :mod:`.utilities.generate_xdmf` on the HDF5-files::
+module :mod:`.io.generate_xdmf` on the HDF5-files::
 
-    from mpi4py_fft.utilities import generate_xdmf
+    from mpi4py_fft.io import generate_xdmf
     generate_xdmf('variousfields.h5')
 
 This will create a number of xdmf-files, one for each group that contains 2D

docs/source/mpi4py_fft.utilities.rst renamed to docs/source/mpi4py_fft.io.rst

@@ -1,45 +1,45 @@
-mpi4py_fft.utilities package
+mpi4py_fft.io package
 =============================
 
 Submodules
 ----------
 
-mpi4py_fft.utilities.generate_xdmf module
+mpi4py_fft.io.generate_xdmf module
 -----------------------------------------
 
-.. automodule:: mpi4py_fft.utilities.generate_xdmf
+.. automodule:: mpi4py_fft.io.generate_xdmf
     :members:
     :undoc-members:
     :show-inheritance:
 
-mpi4py_fft.utilities.h5py_file module
+mpi4py_fft.io.h5py_file module
 -------------------------------------
 
-.. automodule:: mpi4py_fft.utilities.h5py_file
+.. automodule:: mpi4py_fft.io.h5py_file
     :members:
     :undoc-members:
     :show-inheritance:
 
-mpi4py_fft.utilities.nc_file module
+mpi4py_fft.io.nc_file module
 -----------------------------------
 
-.. automodule:: mpi4py_fft.utilities.nc_file
+.. automodule:: mpi4py_fft.io.nc_file
     :members:
     :undoc-members:
     :show-inheritance:
 
-mpi4py_fft.utilities.file_base module
+mpi4py_fft.io.file_base module
 -------------------------------------
 
-.. automodule:: mpi4py_fft.utilities.file_base
+.. automodule:: mpi4py_fft.io.file_base
     :members:
     :undoc-members:
     :show-inheritance:
 
 Module contents
 ---------------
 
-.. automodule:: mpi4py_fft.utilities
+.. automodule:: mpi4py_fft.io
     :members:
     :undoc-members:
     :show-inheritance:

docs/source/mpi4py_fft.rst

@@ -7,7 +7,7 @@ Subpackages
 .. toctree::
 
     mpi4py_fft.fftw
-    mpi4py_fft.utilities
+    mpi4py_fft.io
 
 
 Submodules

mpi4py_fft/distarray.py

@@ -3,6 +3,7 @@
 import numpy as np
 from mpi4py import MPI
 from .pencil import Pencil, Subcomm
+from .io import HDF5File, NCFile, FileBase
 
 comm = MPI.COMM_WORLD
 
@@ -54,8 +55,12 @@ class DistArray(np.ndarray):
     """
     def __new__(cls, global_shape, subcomm=None, val=None, dtype=np.float,
                 buffer=None, alignment=None, rank=0):
-        if rank > 0:
-            assert global_shape[:rank] == (len(global_shape[rank:]),)*rank
+        if len(global_shape) < 2:
+            obj = np.ndarray.__new__(cls, global_shape, dtype=dtype, buffer=buffer)
+            if buffer is None and isinstance(val, Number):
+                obj.fill(val)
+            obj._rank = rank
+            return obj
 
         if isinstance(subcomm, Subcomm):
             pass
@@ -139,13 +144,24 @@ def rank(self):
         """Return tensor rank of ``self``"""
         return self._rank
 
+    @property
+    def dimensions(self):
+        """Return dimensions of array not including rank"""
+        return len(self._p0.shape)
+
     def __getitem__(self, i):
         # Return DistArray if the result is a component of a tensor
         # Otherwise return ndarray view
-        if isinstance(i, int) and self.rank > 0:
+        if self.ndim == 1:
+            return np.ndarray.__getitem__(self, i)
+
+        if isinstance(i, (int, slice)) and self.rank > 0:
             v0 = np.ndarray.__getitem__(self, i)
-            v0._rank -= 1
+            v0._rank = self.rank - (self.ndim - v0.ndim)
+            #if v0.ndim < self.ndim:
+            #    v0._rank -= 1
             return v0
+
         if isinstance(i, tuple) and len(i) == 2 and self.rank == 2:
             v0 = np.ndarray.__getitem__(self, i)
             v0._rank = 0
@@ -246,14 +262,14 @@ def local_slice(self):
         ...         print(l)''')
         >>> fx.close()
         >>> print(subprocess.getoutput('mpirun -np 4 python ls_script.py'))
-        [slice(0, 16, None), slice(0, 7, None), slice(0, 6, None)]
-        [slice(0, 16, None), slice(0, 7, None), slice(6, 12, None)]
-        [slice(0, 16, None), slice(7, 14, None), slice(0, 6, None)]
-        [slice(0, 16, None), slice(7, 14, None), slice(6, 12, None)]
+        (slice(0, 16, None), slice(0, 7, None), slice(0, 6, None))
+        (slice(0, 16, None), slice(0, 7, None), slice(6, 12, None))
+        (slice(0, 16, None), slice(7, 14, None), slice(0, 6, None))
+        (slice(0, 16, None), slice(7, 14, None), slice(6, 12, None))
         """
         v = [slice(start, start+shape) for start, shape in zip(self._p0.substart,
                                                                self._p0.subshape)]
-        return [slice(0, s) for s in self.shape[:self.rank]] + v
+        return tuple([slice(0, s) for s in self.shape[:self.rank]] + v)
 
     def get_pencil_and_transfer(self, axis):
         """Return pencil and transfer objects for alignment along ``axis``
@@ -339,6 +355,74 @@ def redistribute(self, axis=None, out=None):
 
         return out
 
+    def write(self, filename, name='darray', step=0, global_slice=None,
+              as_scalar=False):
+        """Write snapshot ``step`` of ``self`` to file ``filename``
+
+        Parameters
+        ----------
+        filename : str or instance of :class:`.FileBase`
+            The name of the file (or the file itself) that is used to store the
+            requested data in ``self``
+        name : str, optional
+            Name used for storing snapshot in file.
+        step : int, optional
+            Index used for snapshot in file.
+        global_slice : sequence of slices or integers, optional
+            Store only this global slice of ``self``
+        as_scalar : boolean, optional
+            Whether to store rank > 0 arrays as scalars. Default is False.
+
+        Example
+        -------
+        >>> from mpi4py_fft import DistArray
+        >>> u = DistArray((8, 8), val=1)
+        >>> u.write('h5file.h5', 'u', 0)
+        >>> u.write('h5file.h5', 'u', (slice(None), 4))
+        """
+        if isinstance(filename, str):
+            writer = HDF5File if filename.endswith('.h5') else NCFile
+            f = writer(filename, u=self, mode='a')
+        elif isinstance(filename, FileBase):
+            f = filename
+        field = [self] if global_slice is None else [(self, global_slice)]
+        f.write(step, {name: field}, as_scalar=as_scalar)
+
+    def read(self, filename, name='darray', step=0):
+        """Read from file ``filename`` into array ``self``
+
+        Note
+        ----
+        Only whole arrays can be read from file, not slices.
+
+        Parameters
+        ----------
+        filename : str or instance of :class:`.FileBase`
+            The name of the file (or the file itself) holding the data that is
+            loaded into ``self``.
+        name : str, optional
+            Internal name in file of snapshot to be read.
+        step : int, optional
+            Index of field to be read. Default is 0.
+
+        Example
+        -------
+        >>> from mpi4py_fft import DistArray
+        >>> u = DistArray((8, 8), val=1)
+        >>> u.write('h5file.h5', 'u', 0)
+        >>> v = DistArray((8, 8))
+        >>> v.read('h5file.h5', 'u', 0)
+        >>> assert np.allclose(u, v)
+
+        """
+        if isinstance(filename, str):
+            writer = HDF5File if filename.endswith('.h5') else NCFile
+            f = writer(filename, u=self, mode='r')
+        elif isinstance(filename, FileBase):
+            f = filename
+        f.read(self, name, step=step)
+
+
 def newDistArray(pfft, forward_output=True, val=0, rank=0, view=False):
     """Return a new :class:`.DistArray` object for provided :class:`.PFFT` object
 

mpi4py_fft/io/__init__.py

@@ -1,3 +1,4 @@
 from .h5py_file import *
 from .nc_file import *
+from .file_base import *
 from .generate_xdmf import generate_xdmf