Hatree-Fork全向量化優化以及雙粒子格林函數FFT計算思路

心得體會，搜索官方文檔。

充分利用numpy廣播特性，優化了Hatree-Fork自洽計算過程（Hatree-Fork過程不是計算瓶頸)

# Hatree_Fork 自洽計算過程
def Hatree_Fock(ks, Nup_avg,Ndown_avg, U, mu,T, N, ncc):
    m = 2*N
    nk = ks.size
    for i in range(ncc):
        Nup_avg_tmp = np.zeros((m,m),dtype='double')
        Ndown_avg_tmp = np.zeros((m,m),dtype='double')
        for i in range(nk):
            Hk0 = Hamiltonian_H0(ks[i],N)
            Hk_up = Hk0 + U*(Ndown_avg -0.5*np.eye(m))
            Hk_down = Hk0 + U*(Nup_avg - 0.5*np.eye(m))
            Ek_up, Ak_up = eigh(Hk_up)
            Ek_down, Ak_down = eigh(Hk_down)
            fermi_up = fermi(Ek_up,mu,T).reshape(1,m)
            fermi_down = fermi(Ek_down,mu,T).reshape(1,m)
            Nk_up = (Ak_up*Ak_up.conj()*fermi_up).sum(axis=1,dtype='double')
            Nk_down = (Ak_down*Ak_down.conj()*fermi_down).sum(axis=1,dtype='double')
            Nup_avg_tmp += np.diag(Nk_up)
            Ndown_avg_tmp += np.diag(Nk_down)
        Nup_avg = Nup_avg_tmp/nk
        Ndown_avg = Ndown_avg_tmp/nk
        if i%5==0:
            print(f"已經完成{i}次迭代")
        
    return Nup_avg, Ndown_avg

利用numpy多維數組操作，對多維數組指定二維求trace，得到譜函數，這個計算也不是瓶頸。O(n^2)

nw = 1024
delta = 0.002
nk = 256
N = 4
ks = np.linspace(0,2*pi,nk)
omega =   np.linspace(-6,6,nw) 
Gk_up_r =   np.zeros((nw,nk,m,m),dtype='complex')
Gk_down_r = np.zeros((nw,nk,m,m),dtype='complex')
Gk_up_a = np.zeros((nw,nk,m,m),dtype='complex')
Gk_down_a = np.zeros((nw,nk,m,m),dtype='complex')
for i in range(nw):
    for j in range(nk):
        Hk0 = Hamiltonian_H0(ks[j],N)
        Hk_up = Hk0 + U*(Ndown_avg -0.5*np.eye(m))
        Hk_down = Hk0 + U*(Nup_avg - 0.5*np.eye(m))
        Gk_up_r[i,j,:,:] = inv((omega[i]+1j*delta)*np.eye(m)-Hk_up)
        Gk_down_r[i,j,:,:] = inv((omega[i]+1j*delta)*np.eye(m)-Hk_down)
        Gk_up_a[i,j,:,:] =   inv((omega[i]-1j*delta)*np.eye(m)-Hk_up)
        Gk_down_a[i,j,:,:] = inv((omega[i]-1j*delta)*np.eye(m)-Hk_down)

Ak_up = -Gk_up_r.imag.trace(axis1=2,axis2=3)/pi
Ak_down = -Gk_down_r.imag.trace(axis1=2,axis2=3)/pi
X,Y = np.meshgrid(ks,omega)
plt.pcolormesh(X, Y, Ak_up)
plt.colorbar()

這個計算費米函數相當於一個自相關過程，可以用快速傅里葉變換技巧來優化。最後的實現過程如下：

f = fermi(omega,mu,T).reshape(nw,1,1,1)

Gk_up_r_mult_f = f*Gk_up_r.imag
Gk_down_r_mult_f = f*Gk_down_r.imag

Gk_up_r_mult_f_transpose = np.transpose(Gk_up_r_mult_f,(0,1,3,2))
Gk_up_a_transpose = np.transpose(Gk_up_a,(0,1,3,2))

Gk_up_r_mult_f_transpose_fft = np.fft.fftn(Gk_up_r_mult_f_transpose,axes=(0,1))
Gk_down_r_ifft = np.fft.ifftn(Gk_down_r,axes=(0,1))
Gk_down_r_mult_f_fft = np.fft.fftn(Gk_down_r_mult_f,axes=(0,1))
Gk_up_a_fft = np.fft.ifftn(Gk_up_a,axes=(0,1))

 

chi0_f = Gk_up_r_mult_f_transpose_fft*Gk_down_r_ifft#+Gk_down_r_mult_f_fft*Gk_up_a_fft
chi0_f += Gk_down_r_mult_f_fft*Gk_up_a_fft_transpose

chi0 = np.fft.ifftn(chi0_f)

 

plt.plot(ks,-chi0[0,:,0,0].real/pi)

 

plt.pcolormesh(X, Y, chi0[:,:,0,0].imag)
plt.ylim((0,1))
plt.colorbar()