使用 FFT 和 Python 高效计算 3D 拉普拉斯算子答案

【问题标题】：Efficiently computing the 3D Laplacian using FFT and Python使用 FFT 和 Python 高效计算 3D 拉普拉斯算子
【发布时间】：2014-03-29 10:37:30
【问题描述】：

为了求解 PDE（薛定谔方程），我需要计算三个维度的拉普拉斯算子。我目前的解决方案是这样的（到目前为止需要最多时间的部分代码）：

for n in range(Ntstep): # loop         

 for i in range(self.Nixyz[0]): # internal levels of wavefunction
    wf.psi[i,:,:,:]=self.expu * wf.psi[i,:,:,:] # potential      

    if n < Ntstep - 1: # compute laplacian in 3d
        wf.psi[i,:,:,:]=\
            sf.ifft(self.expkx*sf.fft(wf.psi[i,:,:,:],
                axis=0,**fft_args),axis=0,**fft_args)
        wf.psi[i,:,:,:]=\
            sf.ifft(self.expky*sf.fft(wf.psi[i,:,:,:],
                axis=1,**fft_args),axis=1,**fft_args)
        wf.psi[i,:,:,:]=\
            sf.ifft(self.expkz*sf.fft(wf.psi[i,:,:,:],
                axis=2,**fft_args),axis=2,**fft_args)

为了获得更好的性能，我尝试/做了/考虑了以下几点：

不要直接执行 3D FFT。拉普拉斯算子是可分离的，因此可以分成三个 1D FFT，这应该会将复杂度从 n^3 降低到 3n。（在上面的代码中完成。）
我针对 MKL 编译了 numpy 和 scipy，希望获得一些性能，特别是希望能够启用多线程计算。对于某些操作，使用了多个线程（矩阵向量乘法），但 numpy.fft 和 scipy.fftpack 都没有使用多个内核。
我编译了 libfftw 和 pyfftw 并将其用作 np/sp 的替代品。我有一个 Intel Core i7-3770K，即四核八线程。当使用带有 fftw 的两个或四个线程时，与 np/sp 相比，我的性能大约是 np/sp 的两倍。出于某种原因，一个线程或四个以上的线程会更慢。

所以，我现在的主要问题基本上是：

FFT(W) 是否可并行化，性能是否随可用内核/线程的数量而扩展？如果是，我需要考虑什么？目前，两到四个线程似乎是我的最佳选择。更多（或更少）会更慢，尽管我的 CPU 上有 8 个线程可用。
我应该尝试并行化我的 Python 代码吗？例如。将三个一维 FFT 放在三个不同的内核上。当然我必须确保我不会同时在不同线程中读取和写入同一个变量，所以我需要在上面的代码中添加一些额外的“临时”变量，例如：
- 线程 1：TempA = FFT(psi..., axis=0)
- 线程 2：TempB = FFT(psi..., axis=1)
- 线程 3：TempC = FFT(psi..., axis=1)
- 最后一步：psi = TempA + TempB + TempC
axis=0 的 FFT 所需的时间是其他轴的两倍（！）。是否有可能摆脱这种差异并使所有 FFT 都同样快？
（新）究竟是 FFT 方法是最佳选择，还是用户 Rory 的有限差分方法总是更好，至少在性能方面？

我认为有效计算拉普拉斯算子是一个已被广泛研究的主题，因此即使是论文、书籍等的一些链接或提示也可能会有所帮助。

【问题讨论】：

标签： python 3d fft pde

【解决方案1】：

我真的没有答案，但我的 fft laplacian 看起来比你的更简单：

def laplacian3d(field, KX, KY, KZ):
    return ifft(-KX**2*fft(field, axis = 0), axis = 0) + 
        ifft(-KY**2*fft(field, axis = 1), axis = 1) + 
        ifft(-KZ**2*fft(field, axis = 2), axis = 2)

其中 KX、KY 和 KZ 是由以下各项组成的 3D 数组： KX, KY, KZ = meshgrid(kx, ky, kz, indexing='ij')，而 feild 是 3D 实空间场（波函数），kx = 2*pi*fftfreq(len(x), (x[1]-x[0]))，（其中 x 是包含均匀间隔位置的实空间 1D 数组）

在实践中，我发现在 cython 中实现的有限差分拉普拉斯算子大约快 10 倍：

cimport numpy as np
cimport cython
import numpy as np

#3D laplacian of a complex function
@cython.boundscheck(False) # turn of bounds-checking for entire function
def laplacianFD3dcomplex(np.ndarray[double complex, ndim=3] f, double complex dx, double complex dy, double complex dz):
    cdef unsigned int i, j, k, ni, nj, nk
    cdef double complex ifactor, jfactor, kfactor, ijkfactor
    ni = f.shape[0]
    nj = f.shape[1]
    nk = f.shape[2]
    cdef np.ndarray[double complex, ndim=3] lapf = np.zeros((ni,nj,nk)) +0.0J

    ifactor = 1/dx**2
    jfactor = 1/dy**2
    kfactor = 1/dz**2
    ijkfactor = 2.0*(ifactor + jfactor + kfactor)

    for i in xrange(1,ni-1):
        for j in xrange(1, nj-1):
            for k in xrange(1, nk-1):
                lapf[i, j, k] = (f[i, j, k-1] + f[i, j, k+1])*kfactor + (f[i, j-1, k] + f[i, j+1, k])*jfactor + (f[i-1, j, k] + f[i+1, j, k])*ifactor - f[i,j,k]*ijkfactor
    return lapf

#3D laplacian of a real function
@cython.boundscheck(False) # turn of bounds-checking for entire function
def laplacianFD3dreal(np.ndarray[double, ndim=3] f, double dx, double dy, double dz):
    cdef unsigned int i, j, k, ni, nj, nk
    cdef double ifactor, jfactor, kfactor, ijkfactor
    ni = f.shape[0]
    nj = f.shape[1]
    nk = f.shape[2]
    cdef np.ndarray[double, ndim=3] lapf = np.zeros((ni,nj,nk))

    ifactor = 1/dx**2
    jfactor = 1/dy**2
    kfactor = 1/dz**2
    ijkfactor = 2.0*(ifactor + jfactor + kfactor)

    for i in xrange(1,ni-1):
        for j in xrange(1, nj-1):
            for k in xrange(1, nk-1):
                lapf[i, j, k] = (f[i, j, k-1] + f[i, j, k+1])*kfactor + (f[i, j-1, k] + f[i, j+1, k])*jfactor + (f[i-1, j, k] + f[i+1, j, k])*ifactor - f[i,j,k]*ijkfactor
    return lapf

上面的代码可以复制到一个名为“cython_finite_diff.pyx”的文件中，并使用如下的 setup.py 文件进行编译：

#To build the cython code in the .pyx file, type in the terminal:
#"python setup.py build_ext --inplace"
from distutils.core import setup
from distutils.extension import Extension
from Cython.Build import cythonize
import numpy

extensions = [
    Extension("cython_finite_diff", ["cython_finite_diff.pyx"],
                include_dirs = [numpy.get_include()]),
]

setup(
    name = "my_cython_fd",
    ext_modules = cythonize(extensions, annotate=True),
)

对格式感到抱歉，我是在堆栈溢出上发帖的菜鸟。另请注意，有限差分拉普拉斯算子给出反射边界条件。您可以通过将循环设置为包括相对边界上的第一行点来使它们周期性。

【讨论】：

感谢您的回答并让我注意到 FD 方法。现在，我想知道 FFT 方法是否可以达到同等或更高的性能。你的方法中的多线程呢？
这两种方法都是可并行的，但我不知道如何实现。使用 FD 方法，您可以将整个域切割成立方体，并在每个子域上并行计算拉普拉斯算子。使用 FFT 方法，您实际上只是计算 3N^2 1D FFT，因此您可以（几乎）尽可能多地并行化。不过，与往常一样，您需要注意启动新进程等的开销不会超过您首先通过并行化获得的收益。
至于哪种方法更快，FFT方法需要O(N^3logN)次计算（假设N是2的幂），FD方法需要O(N^3)次计算（没有N）的限制），所以我认为FD方法总是会获胜（在速度上），但差异会缓慢增长。