mpi4py_fft package¶

Subpackages¶

Submodules¶

mpi4py_fft.libfft module¶

class mpi4py_fft.libfft.FFT(shape, axes=None, dtype=<class 'float'>, padding=False, use_pyfftw=False, transforms=None, **kw)[source]¶

Bases: mpi4py_fft.libfft.FFTBase

Class for serial FFT transforms

Parameters:

shape (list or tuple of ints) – shape of input array planned for
axes (None, int or tuple of ints, optional) – axes to transform over. If None transform over all axes
dtype (np.dtype, optional) – Type of input array
padding (bool, number or list of numbers) – If False, then no padding. If number, then apply this number as padding factor for all axes. If list of numbers, then each number gives the padding for each axis. Must be same length as axes.
kw (dict) – Parameters passed to serial transform object

forward(input_array=None, output_array=None, **options)¶

Generic serial forward transform

Parameters:	input_array (array, optional) output_array (array, optional) options (dict) – parameters to serial transforms
Returns:	output_array
Return type:	array

backward(input_array=None, output_array=None, **options)¶

Generic serial backward transform

Parameters:	input_array (array, optional) output_array (array, optional) options (dict) – parameters to serial transforms
Returns:	output_array
Return type:	array

class mpi4py_fft.libfft.FFTBase(shape, axes=None, dtype=<class 'float'>, padding=False)[source]¶

Bases: object

Base class for serial FFT transforms

Parameters:

shape (list or tuple of ints) – shape of input array planned for
axes (None, int or tuple of ints, optional) – axes to transform over. If None transform over all axes
dtype (np.dtype, optional) – Type of input array
padding (bool, number or list of numbers) – If False, then no padding. If number, then apply this number as padding factor for all axes. If list of numbers, then each number gives the padding for each axis. Must be same length as axes.

class mpi4py_fft.libfft.FFTNumPy(shape, axes=None, dtype=<class 'float'>, padding=False, **kw)[source]¶

Bases: mpi4py_fft.libfft.FFTBase

Class for serial FFT transforms using Numpy FFT

Parameters:

shape (list or tuple of ints) – shape of input array planned for
axes (None, int or tuple of ints, optional) – axes to transform over. If None transform over all axes
dtype (np.dtype, optional) – Type of input array
padding (bool, number or list of numbers) – If False, then no padding. If number, then apply this number as padding factor for all axes. If list of numbers, then each number gives the padding for each axis. Must be same length as axes.
kw (dict) – Parameters passed to serial transform object
forward(input_array=None, output_array=None, **options) – Generic serial forward transform
Parameters:
input_array (array, optional)

output_array (array, optional)

options (dict) – parameters to serial transforms
returns: output_array

rtype: array
backward(input_array=None, output_array=None, **options) – Generic serial backward transform
Parameters:
input_array (array, optional)

output_array (array, optional)

options (dict) – parameters to serial transforms
returns: output_array

rtype: array

mpi4py_fft.mpifft module¶

class mpi4py_fft.mpifft.PFFT(comm, shape=None, axes=None, dtype=<class 'float'>, slab=False, padding=False, collapse=False, use_pyfftw=False, transforms=None, darray=None, **kw)[source]¶

Bases: object

Base class for parallel FFT transforms

Parameters:

comm (MPI communicator)
shape (sequence of ints, optional) – shape of input array planned for
axes (None, int, sequence of ints or sequence of sequence of ints, optional) – axes to transform over.
- None -> All axes are transformed
- int -> Just one axis to transform over
- sequence of ints -> e.g., (0, 1, 2) or (0, 2, 1)
- sequence of sequence of ints -> e.g., ((0,), (1,)) or ((0,), (1, 2)) For seq. of seq. of ints only the last inner sequence may be longer than 1. This corresponds to collapsing axes, where serial FFTs are performed for all collapsed axes in one single call
dtype (np.dtype, optional) – Type of input array
slab (bool, optional) – If True then distribute only one index of the global array
padding (bool, number or sequence of numbers, optional) – If False, then no padding. If number, then apply this number as padding factor for all axes. If sequence of numbers, then each number gives the padding for each axis. Must be same length as axes.
collapse (bool, optional) – If True try to collapse several serial transforms into one
use_pyfftw (bool, optional) – Use pyfftw for serial transforms instead of local wrappers
transforms (None or dict, optional) – Dictionary of axes to serial transforms (forward and backward) along those axes. For example:
```
{(0, 1): (dctn, idctn), (2, 3): (dstn, idstn)}
```
If missing the default is to use rfftn/irfftn for real input arrays and fftn/ifftn for complex input arrays. Real-to-real transforms can be configured using this dictionary and real-to-real transforms from the fftw.xfftn module. See Examples.

Other Parameters:

darray (DistArray object, optional) – Create PFFT using information contained in darray, neglecting most optional Parameters above

forward(input_array=None, output_array=None, **kw)¶

Parallel forward transform. The method is an instance of the Transform class. See Transform.__call__()

Parameters:	input_array (array, optional) output_array (array, optional) kw (dict) – parameters to serial transforms
Returns:	output_array
Return type:	array

backward(input_array=None, output_array=None, **kw)¶

Parallel backward transform. The method is an instance of the Transform class. See Transform.__call__()

Parameters:	input_array (array, optional) output_array (array, optional) kw (dict) – parameters to serial transforms
Returns:	output_array
Return type:	array

Examples

>>> import numpy as np
>>> from mpi4py import MPI
>>> from mpi4py_fft import PFFT, newDistArray
>>> N = np.array([12, 14, 15], dtype=int)
>>> fft = PFFT(MPI.COMM_WORLD, N, axes=(0, 1, 2))
>>> u = newDistArray(fft, False)
>>> u[:] = np.random.random(u.shape).astype(u.dtype)
>>> u_hat = fft.forward(u)
>>> uj = np.zeros_like(u)
>>> uj = fft.backward(u_hat, uj)
>>> assert np.allclose(uj, u)

Now configure with real-to-real discrete cosine transform type 3

>>> from mpi4py_fft.fftw import rfftn, irfftn, dctn, idctn
>>> import functools
>>> dct = functools.partial(dctn, type=3)
>>> idct = functools.partial(idctn, type=3)
>>> transforms = {(1, 2): (dct, idct)}
>>> r2c = PFFT(MPI.COMM_WORLD, N, axes=((0,), (1, 2)), transforms=transforms)
>>> u = newDistArray(r2c, False)
>>> u[:] = np.random.random(u.shape).astype(u.dtype)
>>> u_hat = r2c.forward(u)
>>> uj = np.zeros_like(u)
>>> uj = r2c.backward(u_hat, uj)
>>> assert np.allclose(uj, u)

destroy()[source]¶

dimensions()[source]¶: The number of dimensions for transformed arrays

dtype(forward_output=False)[source]¶

The type of transformed arrays

Parameters:	forward_output (bool, optional) – If True then return dtype of an array that is the result of a forward transform. Otherwise, return the dtype of an array that is input to a forward transform.

local_shape(forward_output=True)[source]¶

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)

local_slice(forward_output=True)[source]¶

The local view into the global data

Parameters:	forward_output (bool, optional) – Return local slices of output array (spectral space) if True, else return local slices of input array (physical space)

shape(forward_output=False)[source]¶

Return shape of tensor for space

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 space, i.e., the input to a forward transfer.

class mpi4py_fft.mpifft.Transform(xfftn, transfer, pencil)[source]¶

Bases: object

Class for performing any parallel transform, forward or backward

Parameters:	xfftn (list of serial transform objects) transfer (list of global redistribution objects) pencil (list of two pencil objects) – The two pencils represent the input and final output configuration of the distributed global arrays

__call__(input_array=None, output_array=None, **kw)[source]¶

Compute transform

Parameters:	input_array (array, optional) output_array (array, optional) kw (dict) – parameters to serial transforms

Note

If input_array/output_array are not given, then use predefined arrays as planned with serial transform object _xfftn.

input_array¶: Return input array of Transform

input_pencil¶: Return input pencil of Transform

output_array¶: Return output array of Transform

output_pencil¶: Return output pencil of Transform

mpi4py_fft.pencil module¶

class mpi4py_fft.pencil.Pencil(subcomm, shape, axis=-1)[source]¶

Bases: object

Class to represent a distributed array (pencil)

Parameters:	subcomm (MPI communicator) shape (sequence of ints) – Shape of global array axis (int, optional) – Pencil is aligned in this direction

Examples

Create two pencils for a 4-dimensional array of shape (8, 8, 8, 8) using 4 processors in total. The input pencil will be distributed in the first two axes, whereas the output pencil will be distributed in axes 1 and 2. Note that the Subcomm instance below may distribute any axis where an entry 0 is found, whereas an entry of 1 means that this axis should not be distributed.

>>> import subprocess
>>> fx = open('pencil_script.py', 'w')
>>> h = fx.write('''
... import numpy as np
... from mpi4py import MPI
... from mpi4py_fft.pencil import Subcomm, Pencil
... comm = MPI.COMM_WORLD
... N = (8, 8, 8, 8)
... subcomms = Subcomm(comm, [0, 0, 1, 0])
... axis = 2
... p0 = Pencil(subcomms, N, axis)
... p1 = p0.pencil(0)
... shape0 = comm.gather(p0.subshape)
... shape1 = comm.gather(p1.subshape)
... if comm.Get_rank() == 0:
...     print('Subshapes all 4 processors pencil p0:')
...     print(np.array(shape0))
...     print('Subshapes all 4 processors pencil p1:')
...     print(np.array(shape1))''')
>>> fx.close()
>>> print(subprocess.getoutput('mpirun -np 4 python pencil_script.py'))
Subshapes all 4 processors pencil p0:
[[4 4 8 8]
 [4 4 8 8]
 [4 4 8 8]
 [4 4 8 8]]
Subshapes all 4 processors pencil p1:
[[8 4 4 8]
 [8 4 4 8]
 [8 4 4 8]
 [8 4 4 8]]

Two index sets of the global data of shape (8, 8, 8, 8) are distributed. This means that the current distribution is using two groups of processors, with 2 processors in each group (4 in total). One group shares axis 0 and the other axis 1 on the input arrays. On the output, one group shares axis 1, whereas the other shares axis 2. Note that the call p1 = p0.pencil(0) creates a new pencil (p1) that is non-distributed in axes 0. It is, in other words, aligned in axis 0. Hence the first 8 in the lists with [8 4 4 8] above. The alignment is configurable, and p1 = p0.pencil(1) would lead to an output pencil aligned in axis 1.

pencil(axis)[source]¶

Return a Pencil aligned with axis

Parameters:	axis (int) – The axis along which the pencil is aligned

transfer(pencil, dtype)[source]¶

Return an appropriate instance of the Transfer class

The returned Transfer class is used for global redistribution from this pencil’s instance to the pencil instance provided.

Parameters:	pencil (`Pencil`) – The receiving pencil of a forward transform dtype (dtype) – The type of the sending pencil

class mpi4py_fft.pencil.Subcomm[source]¶

Bases: tuple

Class returning a tuple of subcommunicators of any dimensionality

Parameters:	comm (A communicator or group of communicators) dims (None, int or sequence of ints) – dims = [0, 0, 1] will give communicators distributed in the two first indices, whereas the third will not be distributed

Examples

>>> import subprocess
>>> fx = open('subcomm_script.py', 'w')
>>> h = fx.write('''
... from mpi4py import MPI
... comm = MPI.COMM_WORLD
... from mpi4py_fft.pencil import Subcomm
... subcomms = Subcomm(comm, [0, 0, 1])
... if comm.Get_rank() == 0:
...     for subcomm in subcomms:
...         print(subcomm.Get_size())''')
>>> fx.close()
>>> print(subprocess.getoutput('mpirun -np 4 python subcomm_script.py'))
2
2
1
>>> print(subprocess.getoutput('mpirun -np 6 python subcomm_script.py'))
3
2
1

destroy()[source]¶

class mpi4py_fft.pencil.Transfer(comm, shape, dtype, subshapeA, axisA, subshapeB, axisB)[source]¶

Bases: object

Class for performing global redistributions

Parameters:

comm (MPI communicator)
shape (sequence of ints) – shape of input array planned for
dtype (np.dtype, optional) – Type of input array
subshapeA (sequence of ints) – Shape of input pencil
axisA (int) – Input array aligned in this direction
subshapeB (sequence of ints) – Shape of output pencil
axisB (int) – Output array aligned in this direction

Examples

Create two pencils for a 4-dimensional array of shape (8, 8, 8, 8) using 4 processors in total. The input pencil will be distributed in the first two axes, whereas the output pencil will be distributed in axes 1 and 2. Create a random array of shape according to the input pencil and transfer its values to an array of the output shape.

>>> import subprocess
>>> fx = open('transfer_script.py', 'w')
>>> h = fx.write('''
... import numpy as np
... from mpi4py import MPI
... from mpi4py_fft.pencil import Subcomm, Pencil
... comm = MPI.COMM_WORLD
... N = (8, 8, 8, 8)
... subcomms = Subcomm(comm, [0, 0, 1, 0])
... axis = 2
... p0 = Pencil(subcomms, N, axis)
... p1 = p0.pencil(0)
... transfer = p0.transfer(p1, np.float)
... a0 = np.zeros(p0.subshape, dtype=np.float)
... a1 = np.zeros(p1.subshape)
... a0[:] = np.random.random(a0.shape)
... transfer.forward(a0, a1)
... s0 = comm.reduce(np.sum(a0**2))
... s1 = comm.reduce(np.sum(a1**2))
... if comm.Get_rank() == 0:
...     assert np.allclose(s0, s1)''')
>>> fx.close()
>>> h=subprocess.getoutput('mpirun -np 4 python transfer_script.py')

backward(arrayB, arrayA)[source]¶

Global redistribution from arrayB to arrayA

Parameters:	arrayB (array) – Array of shape subshapeB, containing data to be redistributed arrayA (array) – Array of shape subshapeA, for receiving data

destroy()[source]¶

forward(arrayA, arrayB)[source]¶

Global redistribution from arrayA to arrayB

Parameters:	arrayA (array) – Array of shape subshapeA, containing data to be redistributed arrayB (array) – Array of shape subshapeB, for receiving data

mpi4py_fft.distarray module¶

class mpi4py_fft.distarray.DistArray[source]¶

Bases: numpy.ndarray

Distributed Numpy array

This Numpy array is part of a larger global array. Information about the 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) – Describes how to distribute the array
val (int or None, optional) – Initialize array with this int if buffer is not given
dtype (np.dtype, optional) – Type of array
buffer (np.ndarray, optional) – Array of correct shape
alignment (None or int, optional) – Make sure array is aligned in this direction. Note that alignment does not take rank into consideration.
rank (int, optional) – Rank of tensor (scalar is zero, vector one, matrix two)

For more information, see numpy.ndarray

Note

Tensors of rank higher than 0 are not distributed in the first rank indices. For example,

>>> from mpi4py_fft import DistArray
>>> a = DistArray((3, 8, 8, 8), rank=1)
>>> print(a.pencil.shape)
(8, 8, 8)

The array a cannot be distributed in the first axis of length 3 since rank is 1 and this first index represent the vector component. The pencil attribute of a thus only considers the last three axes.

Also note that the alignment keyword does not take rank into consideration. Setting alignment=2 for the array above means that the last axis will be aligned, also when rank>0.

alignment¶: Return alignment of local self array

commsizes¶: Return number of processors along each axis of self

get_global_slice(gslice)[source]¶

Return global slice of self

Parameters:	gslice (sequence of slice(None) and ints) – The slice of the global array.
Returns:	The slice of the global array is returned on rank 0, whereas the remaining ranks return None
Return type:	Numpy array

Example

>>> import subprocess
>>> fx = open('gs_script.py', 'w')
>>> h = fx.write('''
... from mpi4py import MPI
... from mpi4py_fft.distarray import DistArray
... comm = MPI.COMM_WORLD
... N = (6, 6, 6)
... z = DistArray(N, dtype=float, alignment=0)
... z[:] = comm.Get_rank()
... g = z.get_global_slice((0, slice(None), 0))
... if comm.Get_rank() == 0:
...     print(g)''')
>>> fx.close()
>>> print(subprocess.getoutput('mpirun -np 4 python gs_script.py'))
[0. 0. 0. 2. 2. 2.]

get_pencil_and_transfer(axis)[source]¶

Return pencil and transfer objects for alignment along axis

Parameters:	axis (int) – The new axis to align data with
Returns:	2-tuple where first item is a `Pencil` aligned in `axis`. Second item is a `Transfer` object for executing the redistribution of data
Return type:	2-tuple

global_shape¶: Return global shape of self

local_slice()[source]¶

Return local view into global self array

Returns:	Each item of the returned list is the slice along that axis, describing the view of the `self` array into the global array.
Return type:	List of slices

Example

Print local_slice of a global array of shape (16, 14, 12) using 4 processors.

>>> import subprocess
>>> fx = open('ls_script.py', 'w')
>>> h = fx.write('''
... from mpi4py import MPI
... from mpi4py_fft.distarray import DistArray
... comm = MPI.COMM_WORLD
... N = (16, 14, 12)
... z = DistArray(N, dtype=float, alignment=0)
... ls = comm.gather(z.local_slice())
... if comm.Get_rank() == 0:
...     for l in ls:
...         print(l)''')
>>> fx.close()
>>> print(subprocess.getoutput('mpirun -np 4 python ls_script.py'))
[slice(0, 16, None), slice(0, 7, None), slice(0, 6, None)]
[slice(0, 16, None), slice(0, 7, None), slice(6, 12, None)]
[slice(0, 16, None), slice(7, 14, None), slice(0, 6, None)]
[slice(0, 16, None), slice(7, 14, None), slice(6, 12, None)]

pencil¶: Return pencil describing distribution of self

rank¶: Return rank of self

redistribute(axis=None, darray=None)[source]¶

Global redistribution of local self array

Parameters:	axis (int, optional) – Align local `self` array along this axis darray (`DistArray`, optional) – Copy data to this array of possibly different alignment
Returns:	:class:`.DistArray` – The `self` array globally redistributed. If keyword `darray` is None then a new DistArray (aligned along `axis`) is created and returned
Return type:	darray

subcomm¶: Return tuple of subcommunicators for all axes of self

substart¶: Return starting indices of local self array

v¶: Return local self array as an ndarray object

mpi4py_fft.distarray.Function(*args, **kwargs)[source]¶

mpi4py_fft.distarray.newDistArray(pfft, forward_output=True, val=0, rank=0, view=False)[source]¶

Return a DistArray for provided PFFT object

Parameters:

pfft (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 output from forward transform.
val (int or float) – Value used to initialize array.
rank (int) – Scalar has rank 0, vector 1 and matrix 2
view (bool) – 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.

Examples

>>> from mpi4py import MPI
>>> from mpi4py_fft import PFFT, newDistArray
>>> FFT = PFFT(MPI.COMM_WORLD, [64, 64, 64])
>>> u = newDistArray(FFT, False, rank=1)
>>> u_hat = newDistArray(FFT, True, rank=1)

Module contents¶

This is the mpi4py-fft package

What is mpi4py-fft?¶

The Python package mpi4py-fft is a tool primarily for working with Fast Fourier Transforms (FFTs) of (large) multidimensional arrays. There is really no limit as to how large the arrays can be, just as long as there is sufficient computing powers available. Also, there are no limits as to how transforms can be configured. Just about any combination of transforms from the FFTW library is supported. Furthermore, mpi4py-fft can also be used simply to perform global redistributions (distribute and communicate) of large arrays with MPI, without any transforms at all.

For more information, see documentation.