一维谐振子定态 Schrödinger 方程的数值解法

本文链接：个人站 | 简书 | CSDN
版权声明：除特别声明外，本博客文章均采用 BY-NC-SA 许可协议。转载请注明出处。

前几天整理电脑的时候发现了本科上量子力学讨论班时做的一个 Slide，觉得挺有意思的。花了点时间整理成这篇博客。

一维谐振子

一个质量为 $m$ 的粒子，在一维势场 $V(x) = \dfrac12m\omega^2x^2$ 中运动。其哈密顿算符为
$\hat H = \frac{\hat p^2}{2m} + \frac12m\omega^2\hat x^2$
其中 $\hat x$ 为位置算符， $\hat p = -i\hbar\dfrac{\mathrm d}{\mathrm dx}$ 为动量算符。我们需要求解该体系的定态 Schrödinger 方程：
$\hat H\left|\psi\right> = E\left|\psi\right>$

一维谐振子是除了氢原子之外，为数不多的可以解析求解的体系。那么我们为什么要费劲求它的数值解呢？正因为绝大多数的量子体系都无法解析求解，数值方法才显得尤为重要。

有限差分法

回忆一下泰勒公式
$f(a+h) = f(a) + \frac{f'(a)}{1!}h+\frac{f''(a)}{2!}h^2+o(h^3)$
令 $h=-h$ ，有
$f(a-h) = f(a) - \frac{f'(a)}{1!}h+\frac{f''(a)}{2!}h^2+o(h^3)$
两式相加，可得
$\begin{aligned} f''(a) &= \frac{f(a-h)+f(a+h)-2f(a)}{h^2} + o(h^3)\\ &\approx \frac{f(a-h)+f(a+h) - 2f(a)}{h^2} \end{aligned}$

将 $\psi(x)$ 在 $x\in[-r, r]$ 区间离散化为
$\phi_i \equiv \psi(x_i) = \psi(i\Delta x - r),\quad i=0, 1, 2, \cdots, N$
其中 $N=2r/\Delta x$ ，则 Schrödinger 方程差分化为
$-\frac{\hbar^2}{2m}\frac{\phi_{i-1}+\phi_{i+1}-2\phi_i}{\Delta x^2}+\frac12m\omega^2x_i^2\phi_i = E\phi_i$
在这里，我们假设 $x<-r$ 或 $x>r$ 时 $\psi(x)\to 0$ 。这对能量较低的态是成立的。

将差分方程写成矩阵的形式为
$\left[ \begin{matrix} \frac{m\omega^2x_0^2}2+\frac{\hbar^2}{m\Delta x^2} & -\frac{\hbar^2}{2m\Delta x^2} & 0 & \cdots & 0\\ -\frac{\hbar^2}{2m\Delta x^2} & \frac{m\omega^2x_1^2}2+\frac{\hbar^2}{m\Delta x^2} & -\frac{\hbar^2}{2m\Delta x^2} & \cdots & 0\\ 0 & -\frac{\hbar^2}{2m\Delta x^2} & \frac{m\omega^2x_2^2}2+\frac{\hbar^2}{m\Delta x^2} & \cdots & 0\\ \vdots & \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & 0 & \cdots & \frac{m\omega^2x_N^2}2+\frac{\hbar^2}{m\Delta x^2} \end{matrix} \right] \left[ \begin{matrix} \phi_0\\ \phi_1\\ \phi_2\\ \vdots\\ \phi_N \end{matrix} \right]= E \left[ \begin{matrix} \phi_0\\ \phi_1\\ \phi_2\\ \vdots\\ \phi_N \end{matrix} \right]$

这样一来，问题就转化为求差分矩阵的特征值和特征向量。

QR 算法

QR 算法是一种常见的特征值算法。它利用了矩阵的 QR 分解，即将矩阵 $A$ 分解为一个正交矩阵 $Q$ 和一个上三角矩阵 $R$ 的乘积：
$A=QR$
为什么可以这么分解呢？我们回忆一下 Gram-Schmidt 正交化，将矩阵 $A$ 的列向量看作一组基，则可以通过一系列初等列变换获得一组标准正交基。反过来看，对这组标准正交基所组成的矩阵 $Q$ 作初等列变换也可以得到矩阵 $A$ 。我们知道，对矩阵作初等列变换相当于右乘初等矩阵，且由于 Gram-Schmidt 正交化不涉及列交换，这里用到的初等矩阵均为上三角矩阵。因此，矩阵 $A$ 可以由正交矩阵 $Q$ 右乘一个上三角矩阵 $R$ 得到。在实际应用中，除了 Gram-Schmidt 正交化，还可以用 Householder 变换、Givens 旋转等方法实现 QR 分解。

那么如何利用 QR 分解求解矩阵的特征值呢？

记 $A_0:=A$ ，对 $k=0, 1, 2, \cdots$ ，
$\begin{aligned} A_k &= Q_kR_k\\ A_{k+1} &:= R_kQ_k \end{aligned}$
即在每一步，对 $A_k$ 进行 QR 分解，再由分解后得到对 $Q_k$ 和 $R_k$ 计算 $A_{k+1}$ ，如此迭代。

注意到 $Q_k$ 是正交矩阵，有 $Q_k^{-1}=Q_k^\top$ ，故
$\begin{aligned} A_{k+1} &= R_kQ_k\\ &= Q_k^{-1} Q_kR_kQ_k\\ &= Q_k^{-1}A_kQ_k\\ &= Q_k^\top A_kQ_k \end{aligned}$
也就是说， $A_{k+1}$ 相似于 $A_k$ 。根据递推关系， $A_0, A_1, \cdots, A_k, \cdots$ 全都是相似的，这意味着所有的 $A_k$ 都有相同的特征值。在一定条件下， $A_k$ 会收敛为一个三角矩阵，特征值为其主对角元。

特别地，如果 $A$ 是一个实对称正定矩阵（我们所要求解的差分矩阵刚好是这种情况）， $A_k$ 将会收敛为一个对角矩阵 $\Lambda=diag\{\lambda_0,\lambda_1,\cdots,\lambda_N\}$ ，且 $[\lambda_0,\lambda_1,\cdots,\lambda_N]$ 依次递减。考虑到
$\begin{aligned} A=A_0 &= Q_0A_1Q_0^\top\\ &=(Q_0 Q_1\cdots Q_{k-1})A_k(Q_{k-1}^\top\cdots Q_1^\top Q_0^\top)\\ &=(Q_0 Q_1\cdots Q_{k-1})A_k(Q_0Q_1\cdots Q_{k-1})^\top\\ &=S_kA_kS_k^\top \end{aligned}$
当 $A_k$ 收敛时， $S_k$ 的列向量即为属于相应特征值的特征向量。

综上，对于实对称正定矩阵 $A$ ，我们令 $A_0=A, S_0=I$ ，对 $k=0, 1, 2, \cdots$ ，
$\begin{aligned} A_k &= Q_kR_k\\ A_{k+1} &:= R_kQ_k\\ S_{k+1} &:= S_kQ_k \end{aligned}$
当 $A_k$ 收敛时，就同时求得 $A$ 的特征值和特征向量。

Code

我们可以直接调用 numpy.linalg.qr 作 QR 分解：

import numpy as np

def qr_eig(A, iters=200, tol=1e-6):
    """
    使用 QR 算法求解实对称矩阵的特征值和特征向量
    
    Parameters
    ----------
    A : Array-like
        二维数组，表示待求解的实对称矩阵
    iters : int
        最大迭代次数
    tol : float
        提前退出循环的判断标准
    
    Returns
    -------
    (numpy.ndarray, numpy.ndarray)
        一维数组表示的特征值，二维数组表示的特征向量
    """
    A = np.asarray(A)
    S = np.eye(A.shape[0])
    for _ in range(iters):
        Q, R = np.linalg.qr(A)
        newA = np.dot(R, Q)
        newS = np.dot(S, Q)
        if np.abs(max(np.diag(newA)) - max(np.diag(A))) < tol:
            break
        A = newA
        S = newS
    return np.diag(A), S

接下来就可以求解一维谐振子了。为了方便起见，我们令 $\hbar=m=\omega=1$ 。

# 取 [-5, 5] 的区间，离散化为 N 个点
N = 101
r = 5
dx = 2 * r / (N - 1)
x = np.linspace(-r, r, N, endpoint=True)

# 初始化差分矩阵
A = np.diag(0.5*x**2 + 1/dx**2)
for i in range(N -1 ):
    A[i][i+1] = -0.5/dx**2
    A[i+1][i] = -0.5/dx**2

# 求解差分矩阵的特征值和特征向量
Lambda, S = qr_eig(A)

# 结果展示
# 打印最低的 5 个能级
print(Lambda[-1:-8:-1])
# 画出能量最低的三个态的波函数
plt.plot(x, S[:,-1], label=f'n=0, E={Lambda[-1]:.4f}')
plt.plot(x, S[:,-2], label=f'n=1, E={Lambda[-2]:.4f}')
plt.plot(x, S[:,-3], label=f'n=2, E={Lambda[-3]:.4f}')
plt.xlabel('x')
plt.ylabel('psi(x)')
plt.legend()

>>> [0.4996873  1.49843574 2.49593067 3.49217016 4.48715598]

numerical.png

与解析解的对比

一维谐振子的本征能量为：
$E_n = \left(n+\frac12\right)\hbar\omega,\qquad n=0, 1, 2,\cdots$
对应的本征态为：
$\psi_n(x) = \frac{1}{\sqrt{2^n n!}}\cdot\left(\frac{m\omega}{\pi\hbar}\right)^{1/4}\cdot\mathrm{e}^{-m\omega x^2/2\hbar}\cdot H_n\left(\sqrt{\frac{m\omega}{\hbar}}x\right)$
其中
$H_n(z) = (-)^n\mathrm{e}^{z^2}\frac{\mathrm d^n}{\mathrm dz^n}\left(\mathrm{e}^{-z^2}\right)$
为 Hermite 多项式。前三个 Hermite 多项式为：
$\begin{aligned} H_0(z) &= 1\\ H_1(z) &= 2z \\ H_2(z) &= 4z^2 - 2 \end{aligned}$
我们同样令 $\hbar=m=\omega=1$ ，则
$\begin{aligned} E_0&=\frac12, \qquad \psi_0(x) = \frac{1}{\pi^{1/4}}\cdot\mathrm{e}^{-x^2/2}\\ E_1&=\frac32, \qquad \psi_1(x) = \frac{1}{\sqrt 2}\cdot \frac{1}{\pi^{1/4}}\cdot\mathrm{e}^{-x^2/2}\cdot 2x\\ E_2&=\frac52, \qquad \psi_2(x) = \frac{1}{\sqrt 2}\cdot \frac{1}{\pi^{1/4}}\cdot\mathrm{e}^{-x^2/2}\cdot (2x^2-1) \end{aligned}$
画出来看看