For consistency with DistArray local_shape -> shape and shape -> global_shape

Author: mikaem

examples/darray.py

@@ -70,12 +70,12 @@
 s1 = MPI.COMM_WORLD.reduce(np.linalg.norm(z2)**2)
 if MPI.COMM_WORLD.Get_rank() == 0:
     assert abs(s0-s1) < 1e-12
-print('hei')
+
 N = (3, 3, 6, 6, 6)
 z2 = DistArray(N, dtype=float, val=1, alignment=2, rank=2)
 z2[:] = MPI.COMM_WORLD.Get_rank()
-#z1 = z2.redistribute(1)
-#z0 = z1.redistribute(0)
+z1 = z2.redistribute(1)
+z0 = z1.redistribute(0)
 
 s0 = MPI.COMM_WORLD.reduce(np.linalg.norm(z2)**2)
 s1 = MPI.COMM_WORLD.reduce(np.linalg.norm(z0)**2)

examples/spectral_dns_solver.py

@@ -28,32 +28,32 @@
 FFT_pad = FFT
 
 # Declare variables needed to solve Navier-Stokes
-U = newDistArray(FFT, False, rank=1)       # Velocity
-U_hat = newDistArray(FFT, rank=1)          # Velocity transformed
-P = newDistArray(FFT, False)               # Pressure (scalar)
-P_hat = newDistArray(FFT)                  # Pressure transformed
-U_hat0 = newDistArray(FFT, rank=1)         # Runge-Kutta work array
-U_hat1 = newDistArray(FFT, rank=1)         # Runge-Kutta work array
+U = newDistArray(FFT, False, rank=1, view=True)       # Velocity
+U_hat = newDistArray(FFT, rank=1, view=True)          # Velocity transformed
+P = newDistArray(FFT, False, view=True)               # Pressure (scalar)
+P_hat = newDistArray(FFT, view=True)                  # Pressure transformed
+U_hat0 = newDistArray(FFT, rank=1, view=True)         # Runge-Kutta work array
+U_hat1 = newDistArray(FFT, rank=1, view=True)         # Runge-Kutta work array
 a = [1./6., 1./3., 1./3., 1./6.]        # Runge-Kutta parameter
 b = [0.5, 0.5, 1.]                      # Runge-Kutta parameter
-dU = newDistArray(FFT, rank=1)             # Right hand side of ODEs
-curl = newDistArray(FFT, False, rank=1)
-U_pad = newDistArray(FFT_pad, False, rank=1)
-curl_pad = newDistArray(FFT_pad, False, rank=1)
+dU = newDistArray(FFT, rank=1, view=True)             # Right hand side of ODEs
+curl = newDistArray(FFT, False, rank=1, view=True)
+U_pad = newDistArray(FFT_pad, False, rank=1, view=True)
+curl_pad = newDistArray(FFT_pad, False, rank=1, view=True)
 
 def get_local_mesh(FFT, L):
     """Returns local mesh."""
     X = np.ogrid[FFT.local_slice(False)]
-    N = FFT.shape()
+    N = FFT.global_shape()
     for i in range(len(N)):
         X[i] = (X[i]*L[i]/N[i])
-    X = [np.broadcast_to(x, FFT.local_shape(False)) for x in X]
+    X = [np.broadcast_to(x, FFT.shape(False)) for x in X]
     return X
 
 def get_local_wavenumbermesh(FFT, L):
     """Returns local wavenumber mesh."""
     s = FFT.local_slice()
-    N = FFT.shape()
+    N = FFT.global_shape()
     # Set wavenumbers in grid
     k = [np.fft.fftfreq(n, 1./n).astype(int) for n in N[:-1]]
     k.append(np.fft.rfftfreq(N[-1], 1./N[-1]).astype(int))
@@ -62,7 +62,7 @@ def get_local_wavenumbermesh(FFT, L):
     Lp = 2*np.pi/L
     for i in range(3):
         Ks[i] = (Ks[i]*Lp[i]).astype(float)
-    return [np.broadcast_to(k, FFT.local_shape(True)) for k in Ks]
+    return [np.broadcast_to(k, FFT.shape(True)) for k in Ks]
 
 X = get_local_mesh(FFT, L)
 K = get_local_wavenumbermesh(FFT, L)

mpi4py_fft/distarray.py

@@ -10,19 +10,19 @@ class DistArray(np.ndarray):
     """Distributed Numpy array
 
     This Numpy array is part of a larger global array. Information about the
-    distribution is contained in the attributes
+    distribution is contained in the attributes.
 
     Parameters
     ----------
     global_shape : sequence of ints
         Shape of non-distributed global array
-    subcomm : None, Subcomm instance or sequence of ints, optional
+    subcomm : None, :class:`.Subcomm` object or sequence of ints, optional
         Describes how to distribute the array
     val : Number or None, optional
         Initialize array with this number if buffer is not given
     dtype : np.dtype, optional
         Type of array
-    buffer : np.ndarray, optional
+    buffer : Numpy array, optional
         Array of correct shape
     alignment : None or int, optional
         Make sure array is aligned in this direction. Note that alignment does
@@ -100,7 +100,13 @@ def __array_finalize__(self, obj):
 
     @property
     def alignment(self):
-        """Return alignment of local ``self`` array"""
+        """Return alignment of local ``self`` array
+
+        Note
+        ----
+        For tensors of rank > 0 the array is actually aligned along
+        ``alignment+rank``
+        """
         return self._p0.axis
 
     @property
@@ -130,7 +136,7 @@ def pencil(self):
 
     @property
     def rank(self):
-        """Return rank of ``self``"""
+        """Return tensor rank of ``self``"""
         return self._rank
 
     def __getitem__(self, i):
@@ -191,9 +197,9 @@ def get_global_slice(self, gslice):
         s = self.local_slice()
         sp = np.nonzero([isinstance(x, slice) for x in gslice])[0]
         sf = tuple(np.take(s, sp))
-        N = self.global_shape
-        f.require_dataset('0', shape=tuple(np.take(N, sp)), dtype=self.dtype)
+        f.require_dataset('data', shape=tuple(np.take(self.global_shape, sp)), dtype=self.dtype)
         gslice = list(gslice)
+        # We are required to check if the indices in si are on this processor
         si = np.nonzero([isinstance(x, int) and not z == slice(None) for x, z in zip(gslice, s)])[0]
         on_this_proc = True
         for i in si:
@@ -202,12 +208,12 @@ def get_global_slice(self, gslice):
             else:
                 on_this_proc = False
         if on_this_proc:
-            f["0"][sf] = self[tuple(gslice)]
+            f["data"][sf] = self[tuple(gslice)]
         f.close()
         c = None
         if comm.Get_rank() == 0:
             h = h5py.File('tmp.h5', 'r')
-            c = h['0'].__array__()
+            c = h['data'].__array__()
             h.close()
             os.remove('tmp.h5')
         return c
@@ -279,10 +285,10 @@ def redistribute(self, axis=None, darray=None):
 
         Returns
         -------
-        :class:`.DistArray` : darray
+        DistArray : darray
             The ``self`` array globally redistributed. If keyword ``darray`` is
             None then a new DistArray (aligned along ``axis``) is created
-            and returned
+            and returned. Otherwise the provided darray is returned.
         """
         if axis is None:
             assert isinstance(darray, np.ndarray)
@@ -308,7 +314,7 @@ def redistribute(self, axis=None, darray=None):
         return darray
 
 def newDistArray(pfft, forward_output=True, val=0, rank=0, view=False):
-    """Return a :class:`.DistArray` for provided :class:`.PFFT` object
+    """Return a new :class:`.DistArray` object for provided :class:`.PFFT` object
 
     Parameters
     ----------
@@ -317,15 +323,21 @@ def newDistArray(pfft, forward_output=True, val=0, rank=0, view=False):
         If False then create newDistArray of shape/type for input to
         forward transform, otherwise create newDistArray of shape/type for
         output from forward transform.
-    val : int or float
+    val : int or float, optional
         Value used to initialize array.
-    rank: int
+    rank: int, optional
         Scalar has rank 0, vector 1 and matrix 2
-    view : bool
+    view : bool, optional
         If True return view of the underlying Numpy array, i.e., return
         cls.view(np.ndarray). Note that the DistArray still will
         be accessible through the base attribute of the view.
 
+    Returns
+    -------
+    Distarray
+        A new :class:`.DistArray` object. Return the ``ndarray`` view if
+        keyword ``view`` is True.
+
     Examples
     --------
     >>> from mpi4py import MPI
@@ -335,17 +347,13 @@ def newDistArray(pfft, forward_output=True, val=0, rank=0, view=False):
     >>> u_hat = newDistArray(FFT, True, rank=1)
 
     """
+    global_shape = pfft.global_shape(forward_output)
+    p0 = pfft.pencil[forward_output]
     if forward_output is True:
-        shape = pfft.forward.output_array.shape
         dtype = pfft.forward.output_array.dtype
-        p0 = pfft.pencil[1]
     else:
-        shape = pfft.forward.input_array.shape
         dtype = pfft.forward.input_array.dtype
-        p0 = pfft.pencil[0]
-    commsizes = [s.Get_size() for s in p0.subcomm]
-    global_shape = tuple([s*p for s, p in zip(shape, commsizes)])
-    global_shape = (len(shape),)*rank + global_shape
+    global_shape = (len(global_shape),)*rank + global_shape
     z = DistArray(global_shape, subcomm=p0.subcomm, val=val, dtype=dtype,
                   rank=rank)
     return z.v if view else z

mpi4py_fft/mpifft.py

@@ -327,7 +327,7 @@ def destroy(self):
         for trans in self.transfer:
             trans.destroy()
 
-    def local_shape(self, forward_output=True):
+    def shape(self, forward_output=True):
         """The local (to each processor) shape of data
 
         Parameters
@@ -360,14 +360,27 @@ def local_slice(self, forward_output=True):
                                                                    ip.subshape)]
         return tuple(s)
 
-    def shape(self, forward_output=False):
-        """Return shape of tensor for space
+    def local_shape(self, forward_output=False):
+        """The local (to each processor) shape of data
+
+        Parameters
+        ----------
+        forward_output : bool, optional
+            Return shape of output array (spectral space) if True, else return
+            shape of input array (physical space)
+        """
+        import warnings
+        warnings.warn("local_shape() is deprecated; use shape().", FutureWarning)
+        return self.shape(forward_output)
+
+    def global_shape(self, forward_output=False):
+        """Return global shape of associated tensors
 
         Parameters
         ----------
         forward_output : bool, optional
-            If True then return shape of spectral space, i.e., the input to
-            a backward transfer. If False then return shape of physical
+            If True then return global shape of spectral space, i.e., the input
+            to a backward transfer. If False then return shape of physical
             space, i.e., the input to a forward transfer.
         """
         if forward_output:

mpi4py_fft/utilities/h5py_file.py

@@ -140,7 +140,7 @@ def _write_slice_step(self, name, step, slices, field, **kw):
         group = "/".join((name, "{}D".format(ndims), slname))
         if group not in self.f:
             self.f.create_group(group)
-        N = self.T.shape(forward_output)
+        N = self.T.global_shape(forward_output)
         self.f[group].require_dataset(str(step), shape=tuple(np.take(N, sp)), dtype=field.dtype)
         if inside == 1:
             self.f["/".join((group, str(step)))][sf] = field[sl]
@@ -151,5 +151,5 @@ def _write_group(self, name, u, step, **kw):
         group = "/".join((name, "{}D".format(self.T.dimensions())))
         if group not in self.f:
             self.f.create_group(group)
-        self.f[group].require_dataset(str(step), shape=self.T.shape(forward_output), dtype=u.dtype)
+        self.f[group].require_dataset(str(step), shape=self.T.global_shape(forward_output), dtype=u.dtype)
         self.f["/".join((group, str(step)))][s] = u

mpi4py_fft/utilities/nc_file.py

@@ -42,7 +42,7 @@ def __init__(self, ncname, T, domain=None, clobber=True, mode='r', **kw):
         from netCDF4 import Dataset
         self.filename = ncname
         self.f = f = Dataset(ncname, mode=mode, clobber=clobber, parallel=True, comm=comm, **kw)
-        self.N = N = T.shape(False)[-T.dimensions():]
+        self.N = N = T.global_shape(False)[-T.dimensions():]
         dtype = self.T.dtype(False)
         assert dtype.char in 'fdg'
         self._dtype = dtype
@@ -66,7 +66,7 @@ def __init__(self, ncname, T, domain=None, clobber=True, mode='r', **kw):
                 self.dims.append(xyz)
                 nc_xyz = f.createVariable(xyz, self._dtype, (xyz))
                 nc_xyz[:] = d[i]
-            f.setncatts({"ndim": T.dimensions(), "shape": T.shape(False)})
+            f.setncatts({"ndim": T.dimensions(), "shape": T.global_shape(False)})
             f.sync()
         self.close()