python求解二维拉普拉斯方程（密度泛函系列）

gaohaikuo1 · 发表于 2023-6-27 07:22:05

主方程

import numpy as np
import time
from dft2function import *
time_begin = time.time()
N =30 # 网格点数
dx = dy =1 / (N +1) # 步长
x = np.linspace(dx,1 - dx, N)
y = np.linspace(dy,1 - dy, N)
X, Y = np.meshgrid(x, y)
num_electron = 5
M_laplace=get_laplace(N);
M_T = -1./2.*M_laplace;
max_iter=100
energy_tolerance=1e-5
log={"energy":[float("inf")], "energy_diff":[float("inf")]}
nxy=np.zeros([N,N])
def print_log(i,log):
print(f"step: {i:<5} energy: {log['energy'][-1]:<10.4f} energy_diff: {log['energy_diff'][-1]:.10f}")
for i in range(max_iter):
ex_energy, ex_potential=get_exchange(nxy,dx,dy)
ha_potential=get_hatree(nxy,x,y)
ext = get_ext(x,y)
# Hamiltonian
H=M_T+np.diagflat((ex_potential+ha_potential+ext).flatten())
energy, psi= np.linalg.eigh(H)
# log
log["energy"].append(energy[0])
energy_diff=energy[0]-log["energy"][-2]
log["energy_diff"].append(energy_diff)
print_log(i,log)
# convergence
if abs(energy_diff) < energy_tolerance:
print("converged!")
break
# update density
nxy=get_nx(num_electron, psi, dx,dy)
nxy.resize([N,N])
else:
print("not converged")
time_end = time.time()
print("求解时间: {:.3f}".format(time_end-time_begin))
import matplotlib.pyplot as plt
fig, axs = plt.subplots(10, 10)
for i in range(10):
for j in range(10):
axs[i, j].set_title("EigV: {:.3f}".format(energy[10*i+j]))
axs[i, j].imshow(psi[:, 10*i+j].reshape(N, N))
fig.tight_layout()
plt.show()

复制代码

函数方程

import numpy as np
# 构造系数矩阵
def get_laplace(N):
T = np.zeros((N**2, N**2))
for i in range(N):
for j in range(N):
k = i * N + j # 当前点索引
T[k,k] = -4
if i>0 : T[k,k-N] = 1
if i<N-1: T[k,k+N] = 1
if j>0 : T[k,k-1] = 1
if j<N-1: T[k,k+1] = 1
return T
def integral(dx,dy,z,axis=0):
return np.sum(z*dx*dy, axis=axis)
def get_nx(num_electron, psi, dx,dy):
# normalization
I=integral(dx,dy,psi**2)
normed_psi=psi/np.sqrt(I)[None, :]
fn=[2 for _ in range(num_electron//2)]
if num_electron % 2:
fn.append(1)
# density
res=np.zeros_like(normed_psi[:,0])
for ne, psi in zip(fn,normed_psi.T):
res += ne*(psi**2)
return res
def get_exchange(nxy,dx,dy):
energy=-3./4.*(3./np.pi)**(1./3.)*integral(dx,dy,nxy**(4./3.))
potential=-(3./np.pi)**(1./3.)*nxy**(1./3.)
return energy, potential
def get_hatree(nxy,x,y,eps=1e-1):
dx =x[1]-x[0]
dy =y[1]-y[0]
x1 = x[:,None,None,None]
x2= x[None,:,None,None]
y1 = y[None,None,:,None]
y2= y[None,None,None,:]
nx2y2 = nxy[None,:None,:]
potential=np.sum(nx2y2*dx*dy/np.sqrt((x1-x2)**2+(y1-y2)**2+eps),axis=(-1,-3))
return potential
def get_ext(x,y):
x2y2 = x[:,None]**2*y[None,:]**2
return x2y2

复制代码

gaohaikuo1 · 发表于 2023-6-27 07:23:47

二维拉普拉斯本征方程可以表示为：

$$\nabla^2 u(x,y) = -\lambda u(x,y)$$
其中，$
\nabla^2$ 表示拉普拉斯算子，$u(x,y)$ 是未知的函数，$\lambda$ 是本征值。

我们可以使用有限差分法来求解这个方程。具体来说，我们可以将二维区域 $(0,1) \times (0,1)$ 划分为 $N \times N$个网格点，每个网格点的坐标为 $(i \Delta x, j \Delta y)$，其中 $\Delta x = \Delta y = \frac{1}{N+1}$。我们可以用 $u_{i,j}$ 表示在 $(i \Delta x, j \Delta y)$ 处的函数值，用 $A$ 表示系数矩阵（也就是拉普拉斯算子的离散形式），用 $\lambda_k$ 和 $v_k$ 表示第 $k$个本征值和本征向量，则方程可以表示为：

$$A \mathbf{u} = -\lambda \mathbf{u}$$其中，$\mathbf{u}$ 是 $N^2$ 维的列向量，$\lambda$ 是标量。

我们可以通过求解上面的特征值问题来得到本征值和本征向量。具体来说，我们可以使用 numpy 中的 linalg.eig 函数来求解。
求解过程中，虽然最终的波函数是二维的，但是eig函数只能求解一维的，因此需要把二维降为一维的。、
具体操作时
1.把最终的波函数矩阵按行拼接后转置变为列向量（这只是分解习惯，也可以按列拼接），假设图像N×N，那么这个列向量就有N×N个元素；
2.特征矩阵A的每行都相当于一个算子，对这个图像进行变换（第k行对第k个点进行运算,得到一个新的第k个点，k= 1,2,3...N×N)，具体到拉不拉斯算子，就是对每个点求二阶导，更具体点，就是让“4个相邻点的和”减去4倍当前点"
3.根据”第k行对第k个点进行运算“的原则，假设要处理的点为i行j列的点，对应的列向量中第k = (i-1)*N+(j-1)个点(k从0开始算），前面的点就是k-1,后面的点就是k+1，上面的点就是上一行的点k-N，下面的点就是下一行的点k+N。

代码如下：

import numpy as np
import time
time_begin = time.time()
N =10 # 网格点数
dx = dy =1 / (N +1) # 步长
x = np.linspace(dx,1 - dx, N)
y = np.linspace(dy,1 - dy, N)
X, Y = np.meshgrid(x, y)
# 构造系数矩阵
A = np.zeros((N**2, N**2))
for i in range(N):
for j in range(N):
k = i * N + j # 当前点索引
A[k,k] = -4
if i>0 : A[k,k-N] = 1
if i<N-1: A[k,k+N] = 1
if j>0 : A[k,k-1] = 1
if j<N-1: A[k,k+1] = 1
#A[k,k] += (i-N/2)**2+(j-N/2)**2
eig, psi=np.linalg.eigh(A)
time_end = time.time()
print("求解时间: {:.3f}".format(time_end-time_begin))
import matplotlib.pyplot as plt
fig, axs = plt.subplots(10, 10)
for i in range(10):
for j in range(10):
axs[i, j].set_title("EigV: {:.3f}".format(eig[10*i+j]))
axs[i, j].imshow(psi[:, 10*i+j].reshape(N, N))
fig.tight_layout()
plt.show()

复制代码

这里面有几个需要着重理解的地方

运行结果如下图所示：
可以看到，我们成功地求解了二维拉普拉斯本征方程，并得到了前十个本征值和本征向量。

gaohaikuo1 · 发表于 2023-6-27 08:32:14

https://www.qyyshop.com/info/722458.html

gaohaikuo1 · 发表于 2023-6-27 18:07:56

尝试加快运算速度

import numpy as np
from numpy.linalg import eigh
from scipy.sparse.linalg import eigs
import time
time_begin = time.time()
N =100 # 网格点数
dx = dy =1 / (N +1) # 步长
x = np.linspace(dx,1 - dx, N)
y = np.linspace(dy,1 - dy, N)
X, Y = np.meshgrid(x, y)
# 构造系数矩阵
A = np.zeros((N**2, N**2))
for i in range(N):
for j in range(N):
k = i * N + j # 当前点索引
A[k,k] = -4
if i>0 : A[k,k-N] = 1
if i<N-1: A[k,k+N] = 1
if j>0 : A[k,k-1] = 1
if j<N-1: A[k,k+1] = 1
#A[k,k] += (i-N/2)**2+(j-N/2)**2
#eig, psi=eigh(A)
eig, psi=eigs(A,10)
time_end = time.time()
print("求解时间: {:.3f}".format(time_end-time_begin))
import matplotlib.pyplot as plt
fig, axs = plt.subplots(2, 5)
for i in range(2):
for j in range(5):
axs[i, j].set_title("EigV: {:.3f}".format(eig[i]))
axs[i, j].imshow(psi[:, 2*i+j].reshape(N, N).real)
fig.tight_layout()
plt.show()

复制代码

gaohaikuo1 · 发表于 2023-6-27 23:15:43

稀疏矩阵求解

import numpy as np
from scipy.sparse import diags
from scipy.sparse.linalg import eigsh
# 定义求解区域和边界条件
L =1.0 # 矩形的长
W =1.0 # 矩形的宽
N =50 # 矩形的行数
M =50 # 矩形的列数
# 离散化步长
dx = L / (N -1)
dy = W / (M -1)
# 定义系数矩阵和右端向量
A = diags([-4.0,1.0,1.0,1.0,1.0], [0, -1,1, -N, N], shape=(N*M, N*M))
A = -A / dx**2 #+ diags([1.0], [0], shape=(N*M, N*M))
b = np.zeros(N*M)
# 设置边界条件
for i in range(N):
b[i] =0.0
b[N*(M-1)+i] =0.0
for j in range(M):
b[j*N] =0.0
b[(j+1)*N-1] =0.0
# 求解本征值和本征向量
n_eigs =100 # 求解前5个本征值和本征向量
lambdas, us = eigsh(A, k=n_eigs, which='SM')
idx = np.argsort(lambdas)[:100]
lambdas = lambdas[idx]
us = us[:, idx]
#print('可视化本征函数中。。。。')
# 可视化本征函数
import matplotlib.pyplot as plt
fig, axs = plt.subplots(10,10, figsize=(10,4))
for i in range(10):
for j in range(10):
k = i *10 + j
u = us[:, k].reshape((N, M))
print(u)
axs[i, j].imshow(u.real, cmap='coolwarm')#, extent=[0,1,0,1])
axs[i, j].set_title(f'λk+1= {lambdas[k]:.3f}')
plt.tight_layout()
plt.show()

复制代码

如果M和N都是1000，时间会增长为22668s(折合7小时）

gaohaikuo1 · 发表于 2023-6-29 08:24:41

接下来需要求Hatree和exchange，这两个都需要用到电子密度。
电子密度为所有归一化波函数的平方和
所以积分过程变为：

def integral(dx,dy,z,axis=(0,1)):
return np.sum(z*dx*dy, axis=axis)
def get_nx(num_electron, psi, x):
# normalization
I=integral(dx,dy,psi**2)
normed_psi=psi/np.sqrt(I)[None, :]
# occupation num
fn=[2 for _ in range(num_electron//2)]
if num_electron % 2:
fn.append(1)
# density
res=np.zeros_like(normed_psi[:,0])
for ne, psi in zip(fn,normed_psi.T):
res += ne*(psi**2)
return res

复制代码

gaohaikuo1 · 发表于 2023-7-1 21:53:08

LDA项也就是交换关联项
$$\mathrm{v_X^{LDA}[n]=\frac{\partial E_X^{LDA}}{\partial n}=-(\frac3\pi)^{1/3}n^{1/3}}$$
势函数可以写为

def get_exchange(nxy,dx,dy):
energy=-3./4.*(3./np.pi)**(1./3.)*integral(dx,dy,nxy**(4./3.))
potential=-(3./np.pi)**(1./3.)*nxy**(1./3.)
return energy, potential

复制代码

gaohaikuo1 · 发表于 2023-7-1 22:01:22

最后一项是Hatree项
$$\hat{V}_H =\frac{e^2}{4\pi\epsilon_0}\int\frac{\rho(\vec r')}{|\vec r-\vec r'|}dr'$$

def get_hatree(nxy,x,y,eps=1e-1):
dx  =x[1]-x[0]
dy  =y[1]-y[0]
x1 = x[:,None,None,None]
x2=  x[None,:,None,None]
y1 = y[None,None,:,None]
y2=  y[None,None,None,:]
nx2y2 = nxy[None,:None,:]
potential=np.sum(nx2y2*dx*dy/np.sqrt((x1-x2)**2+(y1-y2)**2+eps),axis=(-1,-3))
return  potential

gaohaikuo1 · 发表于 2023-7-1 22:05:32

def get_hatree(nxy,x,y,eps=1e-1):
x1 = x[:,None,None,None]
x2= x[None,:,None,None]
y1 = y[None,None,:,None]
y2= y[None,None,None,:]
nx2y2 = nxy[None,:None,:]
potential=np.sum(nx2y2*h/np.sqrt((x1-x2)**2+(y1-y2)**2+eps),axis=(-1,-3))
return energy, potential

gaohaikuo1 · 发表于 2023-7-1 22:22:23

最后进入主循环
for i in range(max_iter):
ex_energy, ex_potential=get_exchange(nx,x)
ha_potential=get_hatree(nx,x)

# Hamiltonian
H=-D2/2+np.diagflat(ex_potential+ha_potential+x*x)

energy, psi= np.linalg.eigh(H)

# log
log["energy"].append(energy[0])
energy_diff=energy[0]-log["energy"][-2]
log["energy_diff"].append(energy_diff)
print_log(i,log)

# convergence
if abs(energy_diff) < energy_tolerance:
print("converged!")
break

# update density
nx=get_nx(num_electron,psi,x)
else:
print("not converged")

plt.plot(nx)
plt.plot(get_nx(num_electron,psi_harm,x), label="no-interaction")
plt.legend()
plt.pause(0.01)

		自动登录	找回密码
密码			立即注册