并行快速傅里叶变换 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, ax