Mostly docstrings of IO

Author: mikaem

mpi4py_fft/distarray.py

@@ -397,7 +397,8 @@ def write(self, filename, name='darray', step=0, global_slice=None,
         f.write(step, {name: field}, as_scalar=as_scalar)
 
     def read(self, filename, name='darray', step=0):
-        """Read from file ``filename`` into array ``self``
+        """Read data ``name`` at index ``step``from file ``filename`` into
+        ``self``
 
         Note
         ----
@@ -425,7 +426,7 @@ def read(self, filename, name='darray', step=0):
         """
         if isinstance(filename, str):
             writer = HDF5File if filename.endswith('.h5') else NCFile
-            f = writer(filename, u=self, mode='r')
+            f = writer(filename, mode='r')
         elif isinstance(filename, FileBase):
             f = filename
         f.read(self, name, step=step)
@@ -438,8 +439,8 @@ def newDistArray(pfft, forward_output=True, val=0, rank=0, view=False):
     ----------
     pfft : :class:`.PFFT` object
     forward_output: boolean, optional
-        If False then create newDistArray of shape/type for input to
-        forward transform, otherwise create newDistArray of shape/type for
+        If False then create DistArray of shape/type for input to
+        forward transform, otherwise create DistArray of shape/type for
         output from forward transform.
     val : int or float, optional
         Value used to initialize array.

mpi4py_fft/io/file_base.py

@@ -10,6 +10,8 @@ class FileBase(object):
 
     Parameters
     ----------
+    filename : str, optional
+        Name of backend file used to store data
     domain : sequence, optional
         An optional spatial mesh or domain to go with the data.
         Sequence of either
@@ -18,14 +20,15 @@ class FileBase(object):
               of each dimension, e.g., (0, 2*pi).
             - Arrays of coordinates, e.g., np.linspace(0, 2*pi, N). One
               array per dimension.
+
     """
-    def __init__(self, domain=None, **kw):
+    def __init__(self, filename=None, domain=None):
         self.f = None
-        self.filename = None
+        self.filename = filename
         self.domain = domain
 
     def _check_domain(self, group, field):
-        """Write domain to file"""
+        """Check dimensions and store (if missing) self.domain"""
         raise NotImplementedError
 
     def write(self, step, fields, **kw):
@@ -75,25 +78,35 @@ def _write(group, u, sl, step, kw, k=None):
                                 _write(g, u[k, l], sl, step, kw)
 
     def read(self, u, name, **kw):
-        """Read into array ``u``
+        """Read field ``name`` into distributed array ``u``
 
         Parameters
         ----------
         u : array
-            The array to read into.
+            The :class:`.DistArray` to read into.
         name : str
-            Name of array to be read.
+            Name of field to be read.
+        step : int, optional
+            Index of field to be read. Default is 0.
         """
         raise NotImplementedError
 
     def close(self):
         self.f.close()
 
-    def open(self):
+    def open(self, mode='r+'):
+        """Open the self.filename file for reading or writing
+
+        Parameters
+        ----------
+        mode : str
+           Open file in this mode. Default is 'r+'.
+        """
         raise NotImplementedError
 
     @staticmethod
     def backend():
+        """Return which backend is used to store data"""
         raise NotImplementedError
 
     def _write_slice_step(self, name, step, slices, field, **kwargs):

mpi4py_fft/io/h5py_file.py

@@ -23,16 +23,17 @@ class HDF5File(FileBase):
               array per dimension.
     mode : str, optional
         ``r``, ``w`` or ``a`` for read, write or append. Default is ``a``.
+    kw : dict, optional
+        Optional additional keyword arguments used when creating the file
+        used to store data.
     """
     def __init__(self, h5name, domain=None, mode='a', **kw):
-        FileBase.__init__(self, domain=domain, **kw)
+        FileBase.__init__(self, h5name, domain=domain)
         import h5py
-        self.filename = h5name
-        self.f = h5py.File(h5name, mode, driver="mpio", comm=comm)
+        self.f = h5py.File(h5name, mode, driver="mpio", comm=comm, **kw)
         self.close()
 
     def _check_domain(self, group, field):
-        """Check dimensions of domain and write to file"""
         if self.domain is None:
             self.domain = ((0, 2*np.pi),)*field.dimensions
         assert len(self.domain) == field.dimensions
@@ -66,9 +67,9 @@ def _check_domain(self, group, field):
     def backend():
         return 'hdf5'
 
-    def open(self):
+    def open(self, mode='r+'):
         import h5py
-        self.f = h5py.File(self.filename, 'r+', driver="mpio", comm=comm)
+        self.f = h5py.File(self.filename, mode, driver="mpio", comm=comm)
 
     def write(self, step, fields, **kw):
         """Write snapshot ``step`` of ``fields`` to HDF5 file
@@ -82,6 +83,8 @@ def write(self, step, fields, **kw):
             and either arrays or 2-tuples, respectively. The arrays are complete
             arrays to be stored, whereas 2-tuples are arrays with associated
             *global* slices.
+        as_scalar : boolean, optional
+            Whether to store rank > 0 arrays as scalars. Default is False.
 
         Example
         -------
@@ -116,17 +119,6 @@ def write(self, step, fields, **kw):
         self.close()
 
     def read(self, u, name, **kw):
-        """Read from file ``self`` into array ``u``
-
-        Parameters
-        ----------
-        u : array
-            The array to read into.
-        name : str
-            Name of array to be read.
-        step : int, optional
-            Index of field to be read. Default is 0.
-        """
         step = kw.get('step', 0)
         self.open()
         s = u.local_slice()

mpi4py_fft/io/nc_file.py

@@ -29,19 +29,24 @@ class NCFile(FileBase):
     mode : str
         ``r``, ``w`` or ``a`` for read, write or append. Default is ``a``.
     clobber : bool, optional
+        If True (default), opening a file with mode='w' will clobber an
+        existing file with the same name. If False, an exception will be
+        raised if a file with the same name already exists.
+    kw : dict, optional
+        Optional additional keyword arguments used when creating the backend
+        file.
 
     Note
     ----
     Each class instance creates one unique NetCDF4-file, with one step-counter.
     It is possible to store multiple fields in each file, but all snapshots of
     the fields must be taken at the same time. If you want one field stored
     every 10th timestep and another every 20th timestep, then use two different
-    class instances and as such two NetCDF4-files.
+    class instances with two different filenames ``ncname``.
     """
     def __init__(self, ncname, domain=None, mode='a', clobber=True, **kw):
-        FileBase.__init__(self, domain=domain, **kw)
+        FileBase.__init__(self, ncname, domain=domain)
         from netCDF4 import Dataset
-        self.filename = ncname
         # netCDF4 does not seem to handle 'a' if the file does not already exist
         if mode == 'a' and not os.path.exists(ncname):
             mode = 'w'
@@ -51,11 +56,9 @@ def __init__(self, ncname, domain=None, mode='a', clobber=True, **kw):
         if not 'time' in self.f.variables:
             self.f.createDimension('time', None)
             self.f.createVariable('time', np.float, ('time'))
-
         self.close()
 
-    def _check_domain(self, write_domain, field):
-        """Check dimensions of domain and write to file if missing"""
+    def _check_domain(self, group, field):
         N = field.global_shape[field.rank:]
         if self.domain is None:
             self.domain = []
@@ -92,22 +95,59 @@ def _check_domain(self, write_domain, field):
     def backend():
         return 'netcdf4'
 
-    def open(self):
+    def open(self, mode='r+'):
         from netCDF4 import Dataset
-        self.f = Dataset(self.filename, mode='r+', parallel=True, comm=comm)
+        self.f = Dataset(self.filename, mode=mode, parallel=True, comm=comm)
 
     def write(self, step, fields, **kw):
-        """Write snapshot step of ``fields`` to NetCDF4 file
+        """Write snapshot ``step`` of ``fields`` to NetCDF4 file
 
         Parameters
         ----------
         step : int
-            Index of snapshot
+            Index of snapshot.
         fields : dict
             The fields to be dumped to file. (key, value) pairs are group name
             and either arrays or 2-tuples, respectively. The arrays are complete
             arrays to be stored, whereas 2-tuples are arrays with associated
             *global* slices.
+        as_scalar : boolean, optional
+            Whether to store rank > 0 arrays as scalars. Default is False.
+
+        Example
+        -------
+        >>> from mpi4py import MPI
+        >>> from mpi4py_fft import PFFT, NCFile, newDistArray
+        >>> comm = MPI.COMM_WORLD
+        >>> T = PFFT(comm, (15, 16, 17))
+        >>> u = newDistArray(T, forward_output=False, val=1)
+        >>> v = newDistArray(T, forward_output=False, val=2)
+        >>> f = NCFile('ncfilename.nc', mode='w')
+        >>> f.write(0, {'u': [u, (u, [slice(None), 4, slice(None)])],
+        ...             'v': [v, (v, [slice(None), 5, 5])]})
+        >>> f.write(1, {'u': [u, (u, [slice(None), 4, slice(None)])],
+        ...             'v': [v, (v, [slice(None), 5, 5])]})
+        >>> f.close()
+
+        This stores the following datasets to the file ``ncfilename.nc``.
+        Using in a terminal 'ncdump -h ncfilename.nc', one gets::
+
+            netcdf ncfilename {
+            dimensions:
+                    time = UNLIMITED ; // (2 currently)
+                    x = 15 ;
+                    y = 16 ;
+                    z = 17 ;
+            variables:
+                    double time(time) ;
+                    double x(x) ;
+                    double y(y) ;
+                    double z(z) ;
+                    double u(time, x, y, z) ;
+                    double u_slice_4_slice(time, x, z) ;
+                    double v(time, x, y, z) ;
+                    double v_slice_5_5(time, x) ;
+            }
 
         """
         self.open()
@@ -122,17 +162,6 @@ def write(self, step, fields, **kw):
         self.close()
 
     def read(self, u, name, **kw):
-        """Read into array ``u``
-
-        Parameters
-        ----------
-        u : array
-            The array to read into.
-        name : str
-            Name of array to be read.
-        step : int, optional
-            Index of field to be read. Default is 0.
-        """
         step = kw.get('step', 0)
         self.open()
         s = u.local_slice()
@@ -141,10 +170,9 @@ def read(self, u, name, **kw):
         self.close()
 
     def _write_slice_step(self, name, step, slices, field, **kw):
-        assert name not in self.dims
+        assert name not in self.dims # Crashes if user tries to name fields x, y, z, .
         rank = field.rank
-        slices = (slice(None),)*rank + tuple(slices)
-        slices = list(slices)
+        slices = list((slice(None),)*rank + tuple(slices))
         slname = self._get_slice_name(slices[rank:])
         s = field.local_slice()
         slices, inside = self._get_local_slices(slices, s)
@@ -168,7 +196,7 @@ def _write_slice_step(self, name, step, slices, field, **kw):
         self.f.sync()
 
     def _write_group(self, name, u, step, **kw):
-        assert name not in self.dims
+        assert name not in self.dims # Crashes if user tries to name fields x, y, z, .
         s = u.local_slice()
         if name not in self.f.variables:
             h = self.f.createVariable(name, u.dtype, self.dims)

tests/test_io.py

@@ -51,7 +51,7 @@ def test_2D(backend, forward_output):
                 generate_xdmf(filename, order='visit')
 
             u0 = newDistArray(T, forward_output=forward_output, rank=rank)
-            read = reader[backend](filename, u=u0)
+            read = reader[backend](filename)
             read.read(u0, 'u', step=0)
             u0.read(filename, 'u', 2)
             u0.read(read, 'u', 2)
@@ -117,7 +117,7 @@ def test_3D(backend, forward_output):
                 generate_xdmf('v'+filename, order='visit')
 
             u0 = newDistArray(T, forward_output=forward_output, rank=rank)
-            read = reader[backend]('uv'+filename, u=u0)
+            read = reader[backend]('uv'+filename)
             read.read(u0, 'u', step=0)
             assert np.allclose(u0, u)
             read.read(u0, 'v', step=0)
@@ -153,7 +153,7 @@ def test_4D(backend, forward_output):
                 generate_xdmf('uv'+filename)
 
             u0 = newDistArray(T, forward_output=forward_output, rank=rank)
-            read = reader[backend]('uv'+filename, u=u0)
+            read = reader[backend]('uv'+filename)
             read.read(u0, 'u', step=0)
             assert np.allclose(u0, u)
             read.read(u0, 'v', step=0)