Calcolo in modo efficiente il laplaciano 3D utilizzando FFT e Python

Question

Per risolvere un PDE (Equazione Schrödinger), ho bisogno di calcolare l'operatore di Laplace in tre dimensioni. La mia soluzione attuale è questa (parte del codice che richiede di gran lunga il più tempo):

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)

.

Per ottenere prestazioni migliori, ho provato / ha fatto / ho pensato a quanto segue:

.

• Non fare direttamente il FFT 3D. La laplaciana è separabile e quindi può essere divisa in tre fft 1D, che dovrebbe abbattere la complessità da n^3 a 3n. (Fatto nel codice sopra.)

• I Compilato Numpy e Scipy contro il MKL con la speranza di ottenere alcune prestazioni, sperando soprattutto di abilitare calcoli multi-thread. Per alcune operazioni vengono utilizzati thread multipli (moltiplicazione di matrice vettoriali), ma né Numpy.ftt né Scipy.FFTPtack utilizza più core.

• ho compilato libfftw und pyfftw e l'ho usato come sostituzione a drop-in per NP / SP. Ho un Intel Core I7-3770K, I.e. Quattro nuclei e otto fili. Ottieni circa il doppio delle prestazioni rispetto a NP / SP quando usi due o quattro fili con FFTW. Un filo o più di quattro fili è più lento, per qualche motivo.

Allora, le mie domande principali sono fondamentalmente ora:

.

1. è il fft (w) parallelizzabile in un modo le scale di prestazione con il numero di core / fili disponibili? Se sì, cosa devo considerare? Attualmente, due o quattro fili sembrano essere il punto dolce per me. Più (o meno) è più lento, anche se ho otto fili disponibili sulla mia CPU.

2. Dovrei provare a parallelizzare il mio codice Python? Per esempio. Metti i tre fft 1D su tre diversi core. Naturalmente devo assicurarmi di non leggere e scrivere dalla stessa variabile in fili diversi contemporaneamente, quindi ho bisogno di alcune variabili aggiuntive "temp" nel codice sopra, ad esempio.:

   .
   • Thread 1: Tempa= FFT (PSI ..., Axis= 0)
   • Discussione 2: tempb= FFT (PSI ..., Axis= 1)
   • Discussione 3: TEMPC= FFT (PSI ..., AXIS= 1)
   • Fata finale: PSI= Tempa + tempb + TEMPC

3. L'FFT di axis=0 prende due volte (!) Finché per gli altri assi. È possibile sbarazzarsi di questa differenza e rendere tutte le FFTs ugualmente veloci?

4. (nuovo) è l'approccio FFT la scelta migliore dopo tutto, o è l'approccio delle differenze finite per User Rory sempre migliore, almeno in termini di prestazioni?

   < / li >.

   Penso che calcolando in modo efficiente il laplaciano è un argomento che è stato estento ricercato, quindi anche alcuni collegamenti o suggerimenti su documenti, libri ecc. Possono essere utili.

Answer 1

Non ho davvero una risposta, ma il mio laplaciano FFT sembra più semplice del tuo:

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)

.

DOVE KX, KY e KZ sono realizzati con array 3D: KX, KY, KZ = meshgrid(kx, ky, kz, indexing='ij') e Feild è il campo 3D Real Space Space (The WaveFunction) e kx = 2*pi*fftfreq(len(x), (x[1]-x[0])), (con X è il vero array dello spazio 1D contenente posizioni uniformemente distanziate)

In pratica, ho trovato la differenza finita Laplaciani implementata in cython per essere più veloce di circa 10 volte:

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

.

Il codice sopra può essere copiato in un file denominato "cython_finite_diff.pyx" e compilato utilizzando un file setup.py come questo:

#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),
)

.

Scusa per la formattazione, sono un noob per pubblicare su stack overflow.Si noti inoltre che le laplaciane della differenza finita danno condizioni di confine riflettenti.Tu e renderli periodici impostando il ciclo per includere la prima fila di punti sul confine opposto.