Python Bindings

Fast-Fourier Transform

The core of µFFT is the Fast-Fourier-Transform (FFT) abstraction layer, implemented as a class which is specialized for a given FFT engine. Currently supported are pocketfft, FFTW, MPIFFTW and PFFT. In our experience, PFFT has the best scaling properties for large-scale parallel FFTs because it uses pencil decomposition. Only pocketfft is enabled by default.

The FFTs operate on µGrid fields. The engines creates fields with the right memory layout for the respective FFT engine. Please read the µGrid documentation before starting here.

The FFT class

Instantiating an FFT object is a simple as

from muFFT import FFT
fft = FFT((nx, ny, nz), engine='pocketfft')

where [nx, ny, nz] is the shape of the grid and the optional engine is the FFT engine to use. The FFT class takes another optional comm parameter that can be either a mpi4py communicator or an native muFFT Communicator. If no communicator is provided, the FFT class will use the default communicator MPI_COMM_SELF.

µGrid interface

The FFT class provides methods to obtain real-valued real-space fields and complex-valued Fourier-space fields. The transform operates between these fields. An example of a simple transform is shown here:

import numpy as np
from muFFT import FFT

# Instantiate a FFT object with the PocketFFT engine
nb_grid_pts = (32, 32)
fft = FFT(nb_grid_pts, engine='pocketfft')

# Obtain a real and a Fourie-space field
rfield = fft.real_space_field('rfield')
ffield = fft.fourier_space_field('ffield')

# Fill field with random numbers
rfield.p = np.random.rand(*nb_grid_pts)

# Perform a forward FFT
fft.fft(rfield, ffield)

# Perform an inverse FFT
r2field = fft.real_space_field('second rfield')
fft.ifft(ffield, r2field)

# Check that the original and the reconstructed field are the same
np.testing.assert_allclose(rfield.p, r2field.p * fft.normalisation)

Convenience numpy interface

The FFT class also provides a convenience interface that allows usage of numpy arrays directly. Internally, the numpy arrays are converted to µGrid fields. The following shows an example of the convenience interface:

import numpy as np
from muFFT import FFT

# Instantiate a FFT object with the PocketFFT engine
nb_grid_pts = (32, 32)
fft = FFT(nb_grid_pts, engine='pocketfft',
          allow_temporary_buffer=True)

# Create a random real-space field as a numpy array
rarr = np.random.rand(*nb_grid_pts)

# The convenience interface can work with numpy array directly, but this
# implies intermediate temporary copies of the numpy arrays
farr = fft.fft(rarr)

# Convert back
r2arr = fft.ifft(farr)

# Check that the original and the reconstructed field are the same
np.testing.assert_allclose(rarr, r2arr * fft.normalisation)

The downside of the convenience interface is that temporary copies are typically created. The reason for this is that most FFT engines have strict requirements on the memory layout. The µGrid interface allows to create fields with the right memory layout directly, avoiding copies.

Temporary copies can be disabled entirely by setting the allow_temporary_buffer option to false. The above example then still runs for the pocketfft engine, but will fail with the error:

RuntimeError: Incompatible memory layout for the real-space field and no temporary
copies are allowed.

for fftw which requires a padded memory layout.

Normalization

All µFFT transforms are not normalized. A round-trip forward-inverse transform will pick up a global factor, that can be compensated by multiplying with the value of the normlization property of the FFT object as in the examples above.

The reason for this is that there is a choice of putting the factor into the forward or inverse transform, or putting the square-root of the factor into both. µFFT leaves the choice of normalization to the user.

More information can be found for example in the section on normalization of the numpy documentation.

Coordinates and wavevectors

Wavevectors are obtained via the fftfreq property of the FFT object. The result is identical to the numpy.fft.fftfreq function. µFFT adds a second convenience method ifftfreq, which returns an array of integers that indicate the index of the wavevector in the Fourier-space field. This is the same as the output of numpy.fft.fftfreq with d=1/n. The following example shows how to obtain a gradient field using a Fourier derivative that uses the ifftfreq property:

import numpy as np
from muFFT import FFT

# Instantiate a FFT object with the PocketFFT engine
nb_grid_pts = (54, 17)
physical_sizes = (1.4, 2.3)  # Sizes of the domain (in arbitrary units)
fft = FFT(nb_grid_pts)

# Compute wavevectors (2 * pi * k / L for all k and in all directions)
wavevectors = (2 * np.pi * fft.ifftfreq.T / np.array(physical_sizes)).T

# Obtain a real field and fill it
rfield = fft.real_space_field('scalar field')
x, y = fft.coords
rfield.p = np.sin(2 * np.pi * x)  # Just a sine

# Compute Fourier transform
ffield = fft.fourier_space_field('scalar field')
fft.fft(rfield, ffield)

# Compute Fourier gradient by multiplying with wavevector
fgrad = fft.fourier_space_field('gradient field', (2,))
fgrad.p = 1j * wavevectors * ffield.p

# Inverse transform to get gradient in real space
rgrad = fft.real_space_field('gradient field', (2,))
fft.ifft(fgrad, rgrad)

# Normalize gradient (µFFT does not normalize the transform)
gradx, grady = rgrad.p * fft.normalisation

# Gradient in x is cosine
lx, ly = physical_sizes
np.testing.assert_allclose(gradx, 2 * np.pi * np.cos(2 * np.pi * x) / lx, atol=1e-12)
# Gradient in y is zero
np.testing.assert_allclose(grady, 0, atol=1e-12)

The result of fftfreq can be regarded as a fractional Fourier-space coordinate, the result of ifftfreq and an integer Fourier-space coordinate. There are also convenience properties coords and icoords that yield the fractional and integer coordinates in real-space. These properties are useful in particular when running MPI-parallel with domain decomposition. All properties then return just the coordinates of the local domain.

Derivatives

While it is possible to manually compute the derivative as in the above example, µFFT provides utility classes that compute the derivative in Fourier space. The following example shows how to compute the gradient of a field using the FourierDerivative class:

import numpy as np
from muFFT import FFT, FourierDerivative

# Instantiate a FFT object with the PocketFFT engine
nb_grid_pts = (54, 17)
physical_sizes = (1.4, 2.3)  # Sizes of the domain (in arbitrary units)
grid_spacing_x, grid_spacing_y = np.array(physical_sizes) / np.array(nb_grid_pts)
fft = FFT(nb_grid_pts)

# Construct Fourier derivative
fourier_x = FourierDerivative(2, 0)  # Derivative in x direction
fourier_y = FourierDerivative(2, 1)  # Derivative in y direction

# Obtain a real field and fill it
rfield = fft.real_space_field('scalar field')
x, y = fft.coords
rfield.p = np.sin(2 * np.pi * x)  # Just a sine

# Compute Fourier transform
ffield = fft.fourier_space_field('scalar field')
fft.fft(rfield, ffield)

# Compute Fourier gradient by multiplying with wavevector
fgrad = fft.fourier_space_field('gradient field', (2,))
fgrad.p[0] = fourier_x.fourier(fft.fftfreq) * ffield.p / grid_spacing_x
fgrad.p[1] = fourier_y.fourier(fft.fftfreq) * ffield.p / grid_spacing_y

# Inverse transform to get gradient in real space
rgrad = fft.real_space_field('gradient field', (2,))
fft.ifft(fgrad, rgrad)

# Normalize gradient (µFFT does not normalize the transform)
gradx, grady = rgrad.p * fft.normalisation

# Gradient in x is cosine
lx, ly = physical_sizes
np.testing.assert_allclose(gradx, 2 * np.pi * np.cos(2 * np.pi * x) / lx, atol=1e-12)
# Gradient in y is zero
np.testing.assert_allclose(grady, 0, atol=1e-12)

In a similar fashion, µFFT provides utility classes for discrete derivatives. The DiscreteDerivative class accepts a stencil and computes the derivative in Fourier space. The following example shows how to compute the gradient of a field using the DiscreteDerivative class:

import numpy as np
from muFFT import FFT
from muFFT.Stencils2D import upwind_x, upwind_y

# Instantiate a FFT object with the PocketFFT engine
nb_grid_pts = (54, 17)
physical_sizes = (1.4, 2.3)  # Sizes of the domain (in arbitrary units)
grid_spacing_x, grid_spacing_y = np.array(physical_sizes) / np.array(nb_grid_pts)
fft = FFT(nb_grid_pts)

# Obtain a real field and fill it
rfield = fft.real_space_field('scalar field')
x, y = fft.coords
rfield.p = np.sin(2 * np.pi * x)  # Just a sine

# Compute Fourier transform
ffield = fft.fourier_space_field('scalar field')
fft.fft(rfield, ffield)

# Compute gradient by multiplying with wavevector
fgrad = fft.fourier_space_field('gradient field', (2,))
fgrad.p[0] = upwind_x.fourier(fft.fftfreq) * ffield.p / grid_spacing_x
fgrad.p[1] = upwind_y.fourier(fft.fftfreq) * ffield.p / grid_spacing_y

# Inverse transform to get gradient in real space
rgrad = fft.real_space_field('gradient field', (2,))
fft.ifft(fgrad, rgrad)

# Normalize gradient (µFFT does not normalize the transform)
gradx, grady = rgrad.p * fft.normalisation

# Gradient in x is the finite difference of the sine
lx, ly = physical_sizes
np.testing.assert_allclose(gradx, (np.roll(rfield.p, -1, 0) - rfield.p) / grid_spacing_x, atol=1e-12)
# Gradient in y is zero
np.testing.assert_allclose(grady, 0, atol=1e-12)

The modules muFFT.Stencils1D, muFFT.Stencils2D and muFFT.Stencils3D provide a number of predefined stencils that are implemented via the DiscreteDerivative class.

Parallelization

The previous example can be parallelized by initializing the FFT engine with an MPI communicator:

import numpy as np
from mpi4py import MPI
from muGrid import FileIONetCDF, OpenMode, Communicator
from muFFT import FFT

# Instantiate a FFT object with the PocketFFT engine
nb_grid_pts = (32, 32, 32)
physical_sizes = (2, 2, 2)  # Sizes of the domain (in arbitrary units)
fft = FFT(nb_grid_pts, engine='mpi', communicator=MPI.COMM_WORLD)

if MPI.COMM_WORLD.rank == 0:
    print('  Rank   Size          Domain       Subdomain        Location')
    print('  ----   ----          ------       ---------        --------')
MPI.COMM_WORLD.Barrier()  # Barrier so header is printed first

print(f'{MPI.COMM_WORLD.rank:6} {MPI.COMM_WORLD.size:6} {str(fft.nb_domain_grid_pts):>15} '
      f'{str(fft.nb_subdomain_grid_pts):>15} {str(fft.subdomain_locations):>15}')

# Compute wavevectors (2 * pi * k / L for all k and in all directions)
wavevectors = (2 * np.pi * fft.ifftfreq.T / np.array(physical_sizes)).T

# Obtain a real field and fill it
rfield = fft.real_space_field('scalar-field')
x, y, z = fft.coords
rfield.p = np.sin(2 * np.pi * x + 4 * np.pi * y)  # Just a sine

# Compute Fourier transform
ffield = fft.fourier_space_field('scalar-field')
fft.fft(rfield, ffield)

# Compute Fourier gradient by multiplying with wavevector
fgrad = fft.fourier_space_field('gradient-field', (3,))
fgrad.p = 1j * wavevectors * ffield.p

# Inverse transform to get gradient in real space
rgrad = fft.real_space_field('gradient-field', (3,))
fft.ifft(fgrad, rgrad)

# Normalize gradient (µFFT does not normalize the transform)
gradx, grady, gradz = rgrad.p * fft.normalisation

# Gradient in x is cosine
lx, ly, lz = physical_sizes
np.testing.assert_allclose(
    gradx, 2 * np.pi * np.cos(2 * np.pi * x + 4 * np.pi * y) / lx, atol=1e-12)
# Gradient in y is also cosine
np.testing.assert_allclose(
    grady, 4 * np.pi * np.cos(2 * np.pi * x + 4 * np.pi * y) / ly, atol=1e-12)
# Gradient in z is zero
np.testing.assert_allclose(gradz, 0, atol=1e-12)

# I/O example
file = FileIONetCDF('example.nc', open_mode=OpenMode.Overwrite,
                    communicator=Communicator(MPI.COMM_WORLD))
file.register_field_collection(fft.real_field_collection)
file.append_frame().write()

The engine that is selected must support MPI parallelization. Currently, only the fftwmpi engine and pfft engine support MPI parallelization. Using mpi as in the example autoselects an engine.

The parallelization employs domain decomposition. The domain is split into stripe-shaped subdomains for fftwmpi and pencil-shaped subdomains for pfft. pfft scales better to large numbers of MPI processes because of this pencil decomposition. The number of grid points on the local domain is returned by nb_subdomain_grid_pts and the location of the domain is given by subdomain_locations. Note that in a typical code uses coords, icoord, fftfreq and ifftfreq as above and does not need to care about the actual details of the decomposition, as those properties return the domain-local coordinates and wavevectors.

The above example also illustrates how to write the global field to a file.