並行快速傅里葉變換 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