在上一篇中我們介紹了一個使用 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

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