使用FFT和Python高效计算3d拉普拉斯

StackOverflow https://stackoverflow.com//questions/22044126

  •  21-12-2019
  •  | 
  •  

题

为了求解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。拉普拉斯是可分离的，因此可以分裂成三个一维Fft，这应该降低复杂性。 n^3 到 3n.（在上面的代码中完成。)

  • 我针对MKL编译了numpy和scipy，希望获得一些性能，特别是希望能够实现多线程计算。对于某些操作，使用多个线程（矩阵向量乘法），但都不是numpy。fft也不是scipy。fftpack利用多个内核。

  • 我编译了libfftw und pyfftw，并将其用作np/sp的下拉替换。我有一个英特尔酷睿i7-3770k，即四个内核和八个线程。当使用两个或四个带有fftw的线程时，与np/sp相比，我获得了大约两倍的性能。由于某种原因，一个线程或四个以上的线程较慢。

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

  1. FFT（W）在某种程度上是可并行的，性能会随着可用的内核/线程数量而扩展吗？如果是，我需要考虑什么？目前，两到四个线程似乎是我的甜蜜点。更多（或更少）更慢，尽管我的CPU上有八个线程可用。

  2. 我应该尝试并行化我的Python代码吗？例如。将三个1D Fft放在三个不同的内核上。当然，我必须确保我不会在同一时间在不同的线程中读取和写入相同的变量，所以我需要在上面的代码中添加一些额外的"temp"变量，例如:

    • 线程1:TempA=FFT（psi。..，轴=0）
    • 线程2:TempB=FFT（psi。..，轴=1）
    • 螺纹3:TempC=FFT（psi。..，轴=1）
    • 最后一步:psi=TempA+TempB+TempC

  3. 的FFT axis=0 需要两次（！）只要对于其他轴。是否有可能摆脱这种差异，使所有Fft同样快？

  4. (新) FFT方法毕竟是最好的选择，还是用户Rory的有限差异方法总是更好，至少在性能方面？

我认为有效地计算拉普拉斯是一个被广泛研究的主题，所以甚至一些链接或提示论文，书籍等。可能会有帮助。

有帮助吗？

解决方案

我真的没有答案，但我的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是由:KX, KY, KZ = meshgrid(kx, ky, kz, indexing='ij'), ，而feild是3D实空间场（波函数）和 kx = 2*pi*fftfreq(len(x), (x[1]-x[0])), ，（x是包含均匀间隔位置的实空间1D数组）

在实践中，我发现在cython中实现的有限差分laplacians大约快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),
)

抱歉的格式,我是一个noob张贴在堆栈溢出.另请注意，有限差分laplacians给出了反射边界条件。您并通过将循环设置为包含相反边界上的第一行点来使它们具有周期性。

许可以下： CC-BY-SA 和 归因
不隶属于 StackOverflow
scroll top