并行快速傅里叶变换 mpi4py-fft

本文从本人简书博客同步过来

在上一篇中我们介绍了一个使用 mpi4py 实现的并行日志工具 —— python-mpi-logger，下面将介绍一个并行的快速傅里叶变换库 —— mpi4py-fft。

快速傅里叶变换（FFT）简介

快速傅里叶变换（Fast Fourier Transform, FFT），是快速计算序列的离散傅里叶变换（DFT）或其逆变换的方法。傅里叶分析将信号从原始域（通常是时间或空间）转换到频域的表示或者逆过来转换。FFT 会通过把 DFT 矩阵分解为稀疏（大多为零）因子之积来快速计算此类变换。因此，它能够将计算 DFT 的复杂度从只用 DFT 定义计算需要的 O(n²)降低到 O(n log n)，其中 n 为数据大小。

快速傅里叶变换广泛地应用于工程、科学和数学领域。1994年美国数学家吉尔伯特·斯特朗把 FFT 描述为“我们一生中最重要的数值算法”，它还被 IEEE 科学与工程计算期刊列入 20 世纪十大算法。

numpy.fft 模块提供了进行快速傅里叶变换的函数，可以对 1 维，2 维或更高维的数值进行快速傅里叶变换，不过 numpy.fft 只能对完整的数据集或者分布式数据集的某些完整的数据轴进行快速傅里叶变换操作，无法直接对一个分布在不同进程上的数据轴进行快速傅里叶变换。

下面是 numpy.fft 模块用于快速傅里叶变换的主要函数：

fft(a[, n, axis, norm])

计算数组 a 沿轴 axis 的 1 维快速傅里叶变换。

ifft(a[, n, axis, norm])

计算数组 a 沿轴 axis 的 1 维快速逆傅里叶变换。

fft2(a[, s, axes, norm])

计算数组 a 沿轴 axes 的 2 维快速傅里叶变换。

ifft2(a[, s, axes, norm])

计算数组 a 沿轴 axes 的 2 维快速逆傅里叶变换。

fftn(a[, s, axes, norm])

计算数组 a 沿轴 axes 的 n 维快速傅里叶变换。

ifftn(a[, s, axes, norm])

计算数组 a 沿轴 axes 的 n 维快速逆傅里叶变换。

关于 numpy.fft 模块的更多快速傅里叶变换相关的方法及其参数介绍请参加其文档。

下面我们介绍一个使用 mpi4py 对（大的）多维数组（可以分布在多个 MPI 进程上）进行快速傅里叶变换的工具 —— mpi4py-fft。

mpi4py-fft

mpi4py-fft 是一个用于并行快速傅里叶变换计算的 Python 包。它可以分布及重新分布以多种方式分布在多个 MPI 进程上的多维数组，并对其沿任意指定轴（可以是多个轴，包括数据分布轴）进行快速傅里叶变换。mpi4py-fft 支持 FFTW 库（快速傅里叶变换的最快的自由软件实现）中的各种变换方法及它们的各种组合，它提供了 FFTW 库中各种变换方法的 Python 接口。另外，mpi4py-fft 也可以仅仅只对大的数组进行分布及重新分布（借助 mpi4py）而不做任何傅里叶变换操作。

mpi4py-fft 包含如下 4 个主要模块：

• mpifft
• pencil
• libfft
• fftw

mpifft.PFFT 类是做并行快速傅里叶变换的主要入口，这是一个功能强大而丰富的类，它能够自动完成沿各个轴的各种变换所需的（大）数组分布等任务。

pencil 具体完成数组的全局分布（通过 mpi4py）。一般情况下并不直接使用该模块，除非你仅仅只需进行数组的分布操作而不做任何傅里叶变换。mpifft.PFFT 在底层大量使用 pencil 模块的相关功能。

libfft 模块提供了一个（非并行）调用 FFTW 库中相关变换函数的通用的接口。

fftw 模块是对 FFTW 库中相应变换函数的具体包装。通过此模块，用户基本上可以直接在 Python 中使用 FFTW 库中的几乎任何功能。

下载和安装 mpi4py-fft

用下面的命令下载和安装 mpi4py-fft；

git clone https://bitbucket.org/mpi4py/mpi4py-fft.git
cd mpi4py-fft
export FFTW_ROOT=/path/to/your/FFTW
python setup.py install [--user] 

数组的全局分布

在高性能计算中，大的多维数组一般会分布在多个处理器上（这些处理器可以处于不同的物理机器上），我们称之为数组的全局分布。

mpi4py-fft 的 pencil 模块提供了 Pencil、 Subcomm 和 Transfer 三个类用来进行数组的全局分布。

下面是这三个类的定义及主要方法接口：

Pencil 类

class Pencil(object)

Pencil 类，表示一个分布式的数组。

def __init__(self, subcomm, shape, axis=-1)

初始化方法，subcomm 为 MPI 通信子或者一系列的 MPI 通信子。shape 为数组的 shape，axis 为数组对齐的轴（即不分布的轴）。

def pencil(self, axis)

返回一个在 axis 轴上对齐的 Pencil 对象。

def transfer(self, pencil, dtype)

返回一个 Transfer 类对象，该对象能进行当前的 Pencil 对象和参数 pencil 所表示的全局分布之间的相互转换。datatyp 为所使用的数据类型。

Subcomm 类

class Subcomm(tuple)

Subcomm 类，为表示多维数组分布的一系列通信子对象所组成的 tuple。

def __new__(cls, comm, dims=None, reorder=True)

构造方法，comm 为一个或一系列 MPI 通信子，dim 可为 None，一个整数或者一系列整数，指定分布在各维的进程数，对大于 0 的值会保持该值指定的进程数，对等于 0 的值所对应的维会根据总的进程数进行分配。所用应该使用 0 指定自动进程分配的分布轴，用 1 指定在该轴上不做分布，其它大于 0 的数值指定在该轴上分布的进程数。

Transfer 类

class Transfer(object)

Transfer 类，执行数组的全局分布操作。

def __init__(self, comm, shape, dtype, subshapeA, axisA, subshapeB, axisB)

初始化方法，comm 为一个 MPI 通信子，shape 为输入数组的 shape，dtype 为输入数组的 dtype，subshpaeA 为输入 pencil 的 shape，axisA 为输入数组对齐的轴，subshpaeA 为输出 pencil 的 shape，axisA 为输出数组对齐的轴。

def forward(self, arrayA, arrayB)

执行从 arrayA 到 arrayB 的全局分布。

def backward(self, arrayB, arrayA)

执行从 arrayB 到 arrayA 的全局分布。

下面这个例子展示这几个类的使用。

# pencil.py

import numpy as np
from mpi4py_fft import pencil
from mpi4py import MPI

comm = MPI.COMM_WORLD # MPI communicator

N = (8, 8) # shape of the global array
subcomm = pencil.Subcomm(comm, [0, 1]) # distribute on axis 0, not axis 1
p0 = pencil.Pencil(subcomm, N, axis=1) # create a pencil that aligned in axis 1
p1 = p0.pencil(0) # return a new pencil aligned in axis 0
a0 = np.zeros(p0.subshape) # local array with p0's distribution
a0[:] = comm.Get_rank()
print 'p0.subshape: ', a0.shape
transfer = p0.transfer(p1, np.float) # return a Transfer object from p0 to p1
a1 = np.zeros(p1.subshape) # local array with p1's distribution
print 'p1.subshape: ', a1.shape
transfer.forward(a0, a1) # global distribute from a0 to a1
transfer.backward(a1, a0) # global distribute from a1 to a0

以 4 个进程执行以上程序，结果为：

$ mpirun -np 4 python pencils.py
p0.subshape: (2, 8)
p0.subshape: (2, 8)
p0.subshape: (2, 8)
p0.subshape: (2, 8)
p1.subshape: (8, 2)
p1.subshape: (8, 2)
p1.subshape: (8, 2)
p1.subshape: (8, 2)

下图展示了该 8 × 8 数组从 p0 pencil 到 p1 pencil 的全局分布过程。

fftw 模块

fftw 模块提供了 FFTW 库中各种变换函数的包装接口，在 fftw.xfftn 子模块中有下列变换函数：

• fftn() - 复到复快速傅里叶变换
• ifftn() - 复到复快速逆傅里叶变换
• rfftn() - 实到实快速傅里叶变换
• irfftn() - 实到实快速逆傅里叶变换
• dctn() - 实到实离散余弦变换 (DCT)
• idctn() - 实到实离散逆余弦变换
• dstn() - 实到实离散正弦变换 (DST)
• idstn() - 实到实离散逆正弦变换
• hfftn() - 厄密对称序列的复到实快速傅里叶变换
• ihfftn() - 厄密对称序列的实到复快速逆傅里叶变换

并行快速傅里叶变换

并行快速傅里叶变换可以通过数组的全局分布和串行的快速傅里叶变换组合实现，mpi4py-fft 的子模块 mpifft 中的 PFFT 类实现并行快速傅里叶变换的功能（在底层自动完成数据的全局分布和再分布，无需使用者关心）。

下面是 PFFT 类及一个辅助的 Function 类的定义及主要方法：

PFFT 类

class PFFT(object)

PFFT 类，实现并行快速傅里叶变换。

def __init__(self, comm, shape, axes=None, dtype=float, slab=False, padding=False, collapse=False, use_pyfftw=False, transforms=None, **kw)

初始化方法，comm 为 MPI 通信子，shape 为要进行变换的数组的 shape， axes 指定变换轴，可以为 None、一个整数或者一系列整数，dtype 为输入数组的 dtype，slab 如果为 True，则数组只会分布到一个轴上，padding 可以为 True/False、一个数值或一系列数值，如果为 False 则变换的数据不做任何填充。如果为一个数值，则所用轴都会以该值作为填充因子进行填充，如果为一系列数值，则每个轴以该系列中对应的数值作为填充因子进行填充，在这种情况下，padding 的长度必须同 axes 的长度，collapse 指明是否将多个串行的变换合并为一个一起操作，pyfftw 指明是否使用 pyfftw 而不是本地的 FFTW 包装接口，transforms 可为 None 或一个字典，在字典中指定一系列变换轴及需执行的变换，如 {(0, 1): (dctn, idctn), (2, 3): (dstn, idstn)}，如果为 None，则对实的输入数组默认使用 rfftn/irfftn，对复的输入数组使用 fftn/ifftn。

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

并行的正向变换，input_array 为输入数组，output_array 为输出数组。

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

并行的反向变换，input_array 为输入数组，output_array 为输出数组。

Function 类

class Function(np.ndarray)

Function 类，继承自 numpy.ndarray，本质上是一个分布式的 numpy 数组，其 shape 同创建此类对象的 PFFT 实例的输入变换或输出变换的 shape。

def __new__(cls, pfft, forward_output=True, val=0, tensor=None)

构造方法，pfft 为一个 PFFT 类对象，forward_output 如果为 True，则创建的 Function 数组的shape 同 pfft 变换的输出 shape，否则同 pfft 变换的输入 shape，val 为创建的数组的初始填充值。

以上简单介绍了 mpi4py-fft 工具及其中主要的模块和类方法，更多内容可以参见其文档。

例程

下面给出使用例程：

# transform.py

"""
Demonstrates parallel FFT operations with mpi4py-fft package.

Run this with 2 processes like:
$ mpiexec -n 2 python transform.py
"""

import functools
import numpy as np
from mpi4py import MPI
from mpi4py_fft.mpifft import PFFT, Function
from mpi4py_fft.fftw import dctn, idctn


comm = MPI.COMM_WORLD

# global size of global array
N = [8, 8, 8]

dct = functools.partial(dctn, type=3) # the Discrete Cosine Transform
idct = functools.partial(idctn, type=3) # the inverse Discrete Cosine Transform

# transform: dct on axis 1, idct on axis 2
transforms = {(1, 2): (dct, idct)}

# FFT without padding
fft = PFFT(MPI.COMM_WORLD, N, axes=None, collapse=True, slab=True, transforms=transforms)
# FFT with padding, padding factor = [1.5, 1.0, 1.0] for each axis
pfft = PFFT(MPI.COMM_WORLD, N, axes=((0,), (1, 2)), slab=True, padding=[1.5, 1.0, 1.0], transforms=transforms)

# create and initialize the input array
u = Function(fft, False)
u[:] = np.random.random(u.shape).astype(u.dtype)

u_hat = Function(fft) # the output array
u_hat = fft.forward(u, u_hat) # the forward transform

# check the transform
uj = np.zeros_like(u)
uj = fft.backward(u_hat, uj)
print np.allclose(uj, u)

# transform with padding
u_padded = Function(pfft, False)
uc = u_hat.copy()
u_padded = pfft.backward(u_hat, u_padded)
u_hat = pfft.forward(u_padded, u_hat)
print np.allclose(u_hat, uc)

# complex transform, default use fftn/ifftn for the forward/backward transform
cfft = PFFT(MPI.COMM_WORLD, N, dtype=complex)

uc = np.random.random(cfft.backward.input_array.shape).astype(complex)
u2 = cfft.backward(uc)
u3 = uc.copy()
u3 = cfft.forward(u2, u3)

print np.allclose(uc, u3)

执行结果为：

$ mpiexec -n 2 python transform.py
True
True
True
True
True
True