Эффективное вычисление трехмерного лапласиана с использованием БПФ и Python

Для решения 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 БПФ напрямую. Лапласиан разделим и, следовательно, может быть разделен на три одномерных БПФ, что должно снизить сложность с n^3 до 3n. (Выполнено в коде выше.)

  • Я скомпилировал numpy и scipy для MKL в надежде получить некоторую производительность, особенно надеясь включить многопоточные вычисления. Для некоторых операций используется несколько потоков (умножение матрицы на вектор), но ни numpy.fft, ни scipy.fftpack не используют несколько ядер.

  • Я скомпилировал libfftw и pyfftw и использовал его как замену np / sp. У меня Intel Core i7-3770K, то есть четыре ядра и восемь потоков. Я получаю примерно вдвое большую производительность по сравнению с np / sp при использовании двух или четырех потоков с fftw. Один или более четырех потоков по какой-то причине работают медленнее.

Итак, мои основные вопросы в основном сейчас:

  1. Можно ли распараллелить БПФ (W) таким образом, чтобы производительность масштабировалась с количеством доступных ядер / потоков? Если да, что мне нужно учесть? В настоящее время мне кажется, что от двух до четырех потоков лучше всего. Более (или менее) медленнее, хотя на моем процессоре доступно восемь потоков.

  2. Стоит ли мне попытаться распараллелить свой код Python? Например. поместите три одномерных БПФ на три разных ядра. Конечно, я должен убедиться, что я не читаю и не записываю одну и ту же переменную в разных потоках одновременно, поэтому мне нужны дополнительные временные переменные в приведенном выше коде, например:

    • Thread 1: TempA = FFT(psi..., axis=0)
    • Поток 2: TempB = FFT (psi ..., ось = 1)
    • Резьба 3: TempC = FFT (psi ..., ось = 1)
    • Последний шаг: psi = TempA + TempB + TempC

  3. БПФ для axis=0 занимает вдвое (!) Больше времени, чем для других осей. Можно ли избавиться от этой разницы и сделать все БПФ одинаково быстрыми?

  4. (Новое) Является ли подход БПФ лучшим выбором, или подход пользователя Рори, основанный на конечных различиях, всегда лучше, по крайней мере, с точки зрения производительности?

Я думаю, что эффективное вычисление лапласиана - это тема, которая широко исследована, поэтому даже некоторые ссылки или подсказки к статьям, книгам и т. Д. Могут быть полезны.


python 3d pde fft
person Micha    schedule 26.02.2014    source источник

Ответы (1)


arrow_upward
2
arrow_downward

На самом деле у меня нет ответа, но мой лапласианский fft выглядит проще, чем ваш:

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 - это трехмерное поле реального пространства (волновая функция) и kx = 2*pi*fftfreq(len(x), (x[1]-x[0])) (где x - это одномерный массив реального пространства, содержащий равномерно распределенные позиции)

На практике я обнаружил, что лапласианы конечных разностей, реализованные в 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),
)

Извините за форматирование, я новичок в публикации при переполнении стека. Также обратите внимание, что конечные разностные лапласианы задают отражающие граничные условия. Вы и сделаете их периодическими, установив цикл так, чтобы он включал первый ряд точек на противоположной границе.

person Rory    schedule 26.02.2014
comment
Спасибо за ваш ответ и за то, что привлекли мое внимание к методу FD. Теперь мне интересно, может ли подход БПФ быть приведен к такой же или большей производительности. А как насчет многопоточности с вашим подходом? - person Micha; 27.02.2014
comment
Оба метода можно распараллелить, но я не знаю, как это сделать. С помощью метода FD вы можете разрезать весь домен на кубы и параллельно вычислять лапласиан по каждому подобластю. Используя метод БПФ, вы фактически просто вычисляете 3N ^ 2 1D БПФ, поэтому вы можете распараллеливать (почти) столько, сколько захотите. Однако, как всегда, вы должны быть осторожны, чтобы накладные расходы на запуск новых процессов и т. Д. Не перевешивали выгоды, которые вы получаете, в первую очередь, за счет распараллеливания. - person Rory; 07.03.2014
comment
Что касается того, какой метод быстрее, метод FFT требует вычислений O (N ^ 3logN) (при условии, что N является степенью 2), метод FD требует вычислений O (N ^ 3) (без ограничений на N), поэтому я думаю метод FD всегда будет выигрывать (по скорости), но разница будет медленно расти. - person Rory; 07.03.2014