Minor change to paper

Author: mikaem

docs/paper/paper.md

@@ -27,7 +27,7 @@ bibliography: paper.bib
 # Summary
 
 The fast Fourier transform (FFT) is an algorithm that efficiently
-computes the discrete Fourier transform. The FFT is a ubiquitous
+computes the discrete Fourier transform. The FFT is a celebrated
 algorithm utilized throughout science and engineering. Since the dawn
 of our digital society, the FFT has permeated to the heart of everyday
 life applications involving audio, image, and video processing. The
@@ -84,7 +84,7 @@ unknowns.
 which allows for further reuse in applications beyond the FFT. In
 fact, the distributed array interface can be used for boosting
 performance through MPI-based parallelism in any algorithm able to
-operate on local arrays by processing undivided axes one at the time.
+operate on local arrays by processing undivided axes.
 
 # Acknowledgements
 

examples/darray.py

@@ -24,9 +24,6 @@
 s0, s1 = np.linalg.norm(z2), np.linalg.norm(z2c)
 assert abs(s0-s1) < 1e-12, s0-s1
 
-print(z3.get((5, 4, 5)))
-print(z3.local_slice(), z3.substart, z3.commsizes)
-
 v0 = newDistArray(fft, forward_output=False, rank=1)
 #v0 = Function(fft, forward_output=False, rank=1)
 v0[:] = np.random.random(v0.shape)
@@ -40,8 +37,6 @@
 s0, s1 = np.linalg.norm(v0c), np.linalg.norm(v0)
 assert abs(s0-s1) < 1e-12
 
-print(v0.substart, v0.commsizes)
-
 nfft = PFFT(MPI.COMM_WORLD, darray=v0[0], axes=(0, 2, 1))
 for i in range(3):
     v1[i] = nfft.forward(v0[i], v1[i])
@@ -53,17 +48,11 @@
 N = (6, 6, 6)
 z = DistArray(N, dtype=float, alignment=0)
 z[:] = MPI.COMM_WORLD.Get_rank()
-g = z.get((0, slice(None), 0))
-if MPI.COMM_WORLD.Get_rank() == 0:
-    print(g)
-
-z2 = DistArray(N, dtype=float, alignment=2)
-z.redistribute(darray=z2)
-
-g = z2.get((0, slice(None), 0))
-if MPI.COMM_WORLD.Get_rank() == 0:
-    print(g)
-
+g0 = z.get((0, slice(None), 0))
+z2 = z.redistribute(2)
+z = z2.redistribute(darray=z)
+g1 = z.get((0, slice(None), 0))
+assert np.all(g0 == g1)
 s0 = MPI.COMM_WORLD.reduce(np.linalg.norm(z)**2)
 s1 = MPI.COMM_WORLD.reduce(np.linalg.norm(z2)**2)
 if MPI.COMM_WORLD.Get_rank() == 0:
@@ -79,3 +68,16 @@
 s1 = MPI.COMM_WORLD.reduce(np.linalg.norm(z0)**2)
 if MPI.COMM_WORLD.Get_rank() == 0:
     assert abs(s0-s1) < 1e-12
+
+z1 = z0.redistribute(darray=z1)
+z0 = z1.redistribute(darray=z0)
+
+N = (6, 6, 6, 6, 6)
+m0 = DistArray(N, dtype=float, alignment=2)
+m0[:] = MPI.COMM_WORLD.Get_rank()
+m1 = m0.redistribute(4)
+m0 = m1.redistribute(darray=m0)
+s0 = MPI.COMM_WORLD.reduce(np.linalg.norm(m0)**2)
+s1 = MPI.COMM_WORLD.reduce(np.linalg.norm(m1)**2)
+if MPI.COMM_WORLD.Get_rank() == 0:
+    assert abs(s0-s1) < 1e-12

mpi4py_fft/distarray.py

@@ -2,7 +2,7 @@
 from numbers import Number
 import numpy as np
 from mpi4py import MPI
-from .pencil import Pencil, Subcomm
+from .pencil import Pencil, Subcomm, Transfer
 
 comm = MPI.COMM_WORLD
 
@@ -276,6 +276,10 @@ def get_pencil_and_transfer(self, axis):
     def redistribute(self, axis=None, darray=None):
         """Global redistribution of local ``self`` array
 
+        Note
+        ----
+        Use either ``axis`` or ``darray``, not both.
+
         Parameters
         ----------
         axis : int, optional
@@ -290,17 +294,40 @@ def redistribute(self, axis=None, darray=None):
             None then a new DistArray (aligned along ``axis``) is created
             and returned. Otherwise the provided darray is returned.
         """
-        if axis is None:
+        # Take care of some trivial cases first
+        if axis == self.alignment:
+            return self
+
+        # Check if self is already aligned along axis. In that case just switch
+        # axis of pencil (both axes are undivided) and return
+        if axis is not None:
+            if self.commsizes[self.rank+axis] == 1:
+                self._p0.axis = axis
+                return self
+
+        if axis is None: # darray interface
             assert isinstance(darray, np.ndarray)
             assert self.global_shape == darray.global_shape
             axis = darray.alignment
+            if self.commsizes == darray.commsizes:
+                # Just a copy required. Should probably not be here
+                darray[:] = self
+                return darray
+
+            # Check that arrays are compatible
+            for i in range(len(self._p0.shape)):
+                if i != self._p0.axis and i != darray._p0.axis:
+                    assert self._p0.subcomm[i] == darray._p0.subcomm[i]
+                    assert self._p0.subshape[i] == darray._p0.subshape[i]
+
         p1, transfer = self.get_pencil_and_transfer(axis)
         if darray is None:
             darray = DistArray(self.global_shape,
                                subcomm=p1.subcomm,
                                dtype=self.dtype,
                                alignment=axis,
                                rank=self.rank)
+
         if self.rank == 0:
             transfer.forward(self, darray)
         elif self.rank == 1:

tests/test_darray.py

@@ -45,6 +45,8 @@ def test_2Darray():
             s1 = MPI.COMM_WORLD.reduce(np.linalg.norm(b)**2)
             if MPI.COMM_WORLD.Get_rank() == 0:
                 assert abs(s0-s1) < 1e-1
+            c = a.redistribute(a.alignment)
+            assert c is a
 
 def test_3Darray():
     N = (8, 8, 8)