Working on DistributedArray, adding getDarray instead of Function

Author: mikaem

docs/Makefile

@@ -14,6 +14,9 @@ help:
 
 .PHONY: help Makefile
 
+doctest:
+	@$(SPHINXBUILD) -b doctest "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
 # Catch-all target: route all unknown targets to Sphinx using the new
 # "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
 %: Makefile

docs/source/mpi4py_fft.rst

@@ -38,6 +38,14 @@ mpi4py\_fft.pencil module
     :undoc-members:
     :show-inheritance:
 
+mpi4py\_fft.distributedarray module
+-----------------------------------
+
+.. automodule:: mpi4py_fft.distributedarray
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
 Module contents
 ---------------
 

examples/darray.py

@@ -1,25 +1,49 @@
 import numpy as np
 from mpi4py import MPI
 from mpi4py_fft.pencil import Subcomm
-from mpi4py_fft.distributedarray import DistributedArray
-from mpi4py_fft.mpifft import PFFT, Function
+from mpi4py_fft.distributedarray import DistributedArray, getDarray
+from mpi4py_fft.mpifft import PFFT
 
 # Test DistributedArray. Start with alignment in axis 0, then tranfer to 1 and
 # finally to 2
 N = (16, 14, 12)
 z0 = DistributedArray(N, dtype=np.float, alignment=0)
 z0[:] = np.random.randint(0, 10, z0.shape)
 s0 = MPI.COMM_WORLD.allreduce(np.sum(z0))
-print(MPI.COMM_WORLD.Get_rank(), z0.shape)
-z1 = z0.redistribute(1)
+z1 = z0.redistribute(2)
 s1 = MPI.COMM_WORLD.allreduce(np.sum(z1))
-print(MPI.COMM_WORLD.Get_rank(), z1.shape)
-z2 = z1.redistribute(2)
-print(MPI.COMM_WORLD.Get_rank(), z2.shape)
-
+z2 = z1.redistribute(1)
 s2 = MPI.COMM_WORLD.allreduce(np.sum(z2))
 assert s0 == s1 == s2
 
-fft = PFFT(Subcomm(MPI.COMM_WORLD, [s.Get_size() for s in z2.p0.subcomm]), N, dtype=z2.dtype)
-z3 = Function(fft, True)
-fft.forward(z2, z3)
+fft = PFFT(MPI.COMM_WORLD, darray=z2, axes=(0, 2, 1))
+z3 = getDarray(fft, forward_output=True)
+z2c = z2.copy()
+fft.forward(z2, z3)
+fft.backward(z3, z2)
+s0, s1 = np.linalg.norm(z2), np.linalg.norm(z2c)
+assert abs(s0-s1) < 1e-12, s0-s1
+
+print(z3.local_slice(), z3.substart, z3.commsizes)
+
+v0 = getDarray(fft, forward_output=False, rank=1)
+v0[:] = np.random.random(v0.shape)
+v0c = v0.copy()
+v1 = getDarray(fft, forward_output=True, rank=1)
+
+for i in range(3):
+    v1[i] = fft.forward(v0[i], v1[i])
+for i in range(3):
+    v0[i] = fft.backward(v1[i], v0[i])
+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])
+for i in range(3):
+    v0[i] = nfft.backward(v1[i], v0[i])
+s0, s1 = np.linalg.norm(v0c), np.linalg.norm(v0)
+assert abs(s0-s1) < 1e-12

examples/spectral_dns_solver.py

@@ -10,7 +10,7 @@
 from time import time
 import numpy as np
 from mpi4py import MPI
-from mpi4py_fft.mpifft import PFFT, Function
+from mpi4py_fft import PFFT, getDarray
 
 # Set viscosity, end time and time step
 nu = 0.000625
@@ -28,18 +28,18 @@
 FFT_pad = FFT
 
 # Declare variables needed to solve Navier-Stokes
-U = Function(FFT, False, tensor=3)       # Velocity
-U_hat = Function(FFT, tensor=3)          # Velocity transformed
-P = Function(FFT, False)                 # Pressure (scalar)
-P_hat = Function(FFT)                    # Pressure transformed
-U_hat0 = Function(FFT, tensor=3)         # Runge-Kutta work array
-U_hat1 = Function(FFT, tensor=3)         # 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 = Function(FFT, tensor=3)             # Right hand side of ODEs
-curl = Function(FFT, False, tensor=3)
-U_pad = Function(FFT_pad, False, tensor=3)
-curl_pad = Function(FFT_pad, False, tensor=3)
+U = getDarray(FFT, False, rank=1)       # Velocity
+U_hat = getDarray(FFT, rank=1)          # Velocity transformed
+P = getDarray(FFT, False)               # Pressure (scalar)
+P_hat = getDarray(FFT)                  # Pressure transformed
+U_hat0 = getDarray(FFT, rank=1)         # Runge-Kutta work array
+U_hat1 = getDarray(FFT, rank=1)         # 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 = getDarray(FFT, rank=1)             # Right hand side of ODEs
+curl = getDarray(FFT, False, rank=1)
+U_pad = getDarray(FFT_pad, False, rank=1)
+curl_pad = getDarray(FFT_pad, False, rank=1)
 
 def get_local_mesh(FFT, L):
     """Returns local mesh."""
@@ -52,10 +52,8 @@ def get_local_mesh(FFT, L):
 
 def get_local_wavenumbermesh(FFT, L):
     """Returns local wavenumber mesh."""
-
     s = FFT.local_slice()
     N = FFT.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))

examples/transforms.py

@@ -1,7 +1,7 @@
 import functools
 import numpy as np
 from mpi4py import MPI
-from mpi4py_fft.mpifft import PFFT, Function
+from mpi4py_fft import PFFT, DistributedArray
 from mpi4py_fft.fftw import dctn, idctn
 
 # Set global size of the computational box
@@ -17,16 +17,16 @@
 
 assert fft.axes == pfft.axes
 
-u = Function(fft, False)
+u = DistributedArray(pfft=fft, forward_output=False)
 u[:] = np.random.random(u.shape).astype(u.dtype)
 
-u_hat = Function(fft)
+u_hat = DistributedArray(pfft=fft, forward_output=True)
 u_hat = fft.forward(u, u_hat)
 uj = np.zeros_like(u)
 uj = fft.backward(u_hat, uj)
 assert np.allclose(uj, u)
 
-u_padded = Function(pfft, False)
+u_padded = DistributedArray(pfft=pfft, forward_output=False)
 uc = u_hat.copy()
 u_padded = pfft.backward(u_hat, u_padded)
 u_hat = pfft.forward(u_padded, u_hat)

mpi4py_fft/__init__.py

@@ -19,7 +19,7 @@
 __version__ = '1.1.2'
 __author__ = 'Lisandro Dalcin and Mikael Mortensen'
 
-from .distributedarray import DistributedArray
-from .mpifft import PFFT, Function
+from .distributedarray import DistributedArray, getDarray
+from .mpifft import PFFT
 from . import fftw
 from .utilities import HDF5File, NCFile, generate_xdmf