source/xyg3_cpks.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# XYG3 的 CP-KS 方程计算相关"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "通过上一节，我们基本了解了 CP-HF 方程的实际代码过程．这一节，我们会引入 DFT 的内容，首先对 B3LYP 的极化率作说明，以了解 CP-KS 方程的具体解法；随后对 XYG3 的极化率作叙述，以了解 XYG3 电场下的梯度计算流程．"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 准备工作"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 分子计算"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 环境搭建\n",
    "import psi4\n",
    "import numpy as np\n",
    "\n",
    "# 引入 DIIS 模块\n",
    "import sys\n",
    "sys.path.append(\"include\")\n",
    "from DIIS_helper import DIIS_helper\n",
    "\n",
    "# 简化矩阵输出\n",
    "np.set_printoptions(8, linewidth=100, suppress=True)\n",
    "\n",
    "# 输出文件\n",
    "psi4.set_output_file(\"output.dat\", False)\n",
    "\n",
    "# 设置内存 0.5 GB\n",
    "psi4.set_memory(int(5e8))\n",
    "\n",
    "# 设置分子坐标\n",
    "mol = psi4.geometry(\"\"\"\n",
    "    O  0.000000000000    -0.000000000000    -0.079135765807\n",
    "    H  0.000000000000     0.707106781187     0.627971015380\n",
    "    H  0.000000000000    -0.707106781187     0.627971015380\n",
    "    symmetry c1\n",
    "\"\"\")\n",
    "\n",
    "# 设置计算选项\n",
    "psi4.set_options({\n",
    "    'basis':               '6-31g',\n",
    "    'scf_type':            'pk',\n",
    "    'mp2_type':            'conv',\n",
    "    'e_convergence':        1e-8,\n",
    "    'd_convergence':        1e-8,\n",
    "    'dft_spherical_points': 590,\n",
    "    'dft_radial_points':    99,\n",
    "})\n",
    "\n",
    "# 波函数信息\n",
    "scf_e, scf_wfn = psi4.energy(\"B3LYP\", molecule=mol, return_wfn=True)\n",
    "\n",
    "# 验证结果\n",
    "psi4.compare_values(scf_e, -76.3771828949, 6, 'B3LYP Energy')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 中间变量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 电子积分引擎\n",
    "mints = psi4.core.MintsHelper(scf_wfn.basisset())\n",
    "# 积分\n",
    "T = np.asarray(mints.ao_kinetic())  # AO 基组动能积分\n",
    "V = np.asarray(mints.ao_potential())  # AO 基组电子-核势能积分\n",
    "eri = np.asarray(mints.ao_eri())  # AO 基组双电子排斥积分\n",
    "\n",
    "# DFT 积分引擎\n",
    "V_pot = scf_wfn.V_potential()  # DFT 积分引擎\n",
    "\n",
    "# 数值量\n",
    "nocc = scf_wfn.nalpha()\n",
    "nbf = mints.nbf()\n",
    "nmo = scf_wfn.nmo()\n",
    "nvir = nmo - nocc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# B3LYP 导出变量\n",
    "F = np.asarray(scf_wfn.Fa())  # AO 基组 Fock 矩阵\n",
    "D = np.asarray(scf_wfn.Da())  # AO 基组单电子密度\n",
    "C = np.asarray(scf_wfn.Ca())  # AO <-> MO 轨道系数 C_{up}\n",
    "Co = np.asarray(scf_wfn.Ca_subset('AO', 'OCC'))  # 占据轨道系数\n",
    "Cv = np.asarray(scf_wfn.Ca_subset('AO', 'VIR'))  # 未占轨道系数\n",
    "e = np.asarray(scf_wfn.epsilon_a())  # 轨道能级\n",
    "eo = np.asarray(scf_wfn.epsilon_a_subset('AO', 'OCC'))  # 占据轨道能级\n",
    "ev = np.asarray(scf_wfn.epsilon_a_subset('AO', 'VIR'))  # 未占轨道能级\n",
    "dip_psi4 = mints.ao_dipole()\n",
    "dip = np.array([np.asarray(mat) for mat in dip_psi4])  # AO 基组偶极矩积分\n",
    "dip_vo = np.einsum(\"guv, ua, vi -> gai\", dip, Cv, Co, optimize=True)  # 分子轨道 ai 的偶极矩积分"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PT2 导出变量\n",
    "# 轨道能之差的张量 D_{ij}^{ab} -> d[i, a, j, b]\n",
    "d_ovov = (eo.reshape(-1, 1, 1, 1) - ev.reshape(-1, 1, 1) + eo.reshape(-1, 1) - ev)\n",
    "# 轨道能之差 D_i^a -> d_vo[a, i]\n",
    "d_vo = - ev.reshape(-1, 1) + eo\n",
    "# 全 MO 基组 ERI <pq|rs> = (pr|qs) = g_{pq}^{rs} = g_{rs}^{pq} -> g[p, r, q, s]\n",
    "g_pqrs = np.einsum(\"up, vr, uvkl, kq, ls -> prqs\", C, C, eri, C, C, optimize=True)\n",
    "g_ovov = g_pqrs[:nocc, nocc:, :nocc, nocc:]  # <ij|ab> = g_{ij}^{ab}\n",
    "t_ovov = g_ovov / d_ovov  # 轨道对激发振幅"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### XYG3 非自洽泛函有关准备"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义 XYG3 非自洽泛函\n",
    "def build_xyg3_nc_superfunctional(name, npoints, deriv, restricted):\n",
    "    \n",
    "    sup = psi4.core.SuperFunctional.blank()\n",
    "    sup.set_name('XYG3NC')\n",
    "    sup.set_description('    XYG3 Non-Consistent Functional without MP2 Part\\n')\n",
    "    \n",
    "    lda_x = psi4.core.LibXCFunctional(\"XC_LDA_X\", restricted)\n",
    "    lda_x.set_alpha(-0.0140)\n",
    "    sup.add_x_functional(lda_x)\n",
    "    \n",
    "    gga_x = psi4.core.LibXCFunctional(\"XC_GGA_X_B88\", restricted)\n",
    "    gga_x.set_alpha(0.2107)\n",
    "    sup.add_x_functional(gga_x)\n",
    "    \n",
    "    lyp_c = psi4.core.LibXCFunctional(\"XC_GGA_C_LYP\", restricted)\n",
    "    lyp_c.set_alpha(0.6789)\n",
    "    sup.add_c_functional(lyp_c)\n",
    "    \n",
    "    sup.set_x_alpha(0.8033)\n",
    "    \n",
    "    return sup\n",
    "\n",
    "# 计算 XYG3 能量\n",
    "nscf_e, nscf_wfn = psi4.energy(\"SCF\", dft_functional=build_xyg3_nc_superfunctional, return_wfn=True)\n",
    "# 非自洽部分 DFT 积分引擎\n",
    "Vn_pot = nscf_wfn.V_potential()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# XYG3 非自洽 Fock 矩阵\n",
    "Vn = psi4.core.Matrix(nbf, nbf)\n",
    "Vn_pot.set_D([scf_wfn.Da()])\n",
    "Vn_pot.compute_V([Vn])\n",
    "Fn = T + V\n",
    "Fn += 2 * np.einsum(\"uvkl, kl -> uv\", eri, D)\n",
    "Fn -= 0.8033 * np.einsum(\"ukvl, kl -> uv\", eri, D)\n",
    "Fn += Vn\n",
    "Fn_pq = C.T @ Fn @ C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CP-KS 方程变量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diis = None\n",
    "USE_DIIS = True\n",
    "if USE_DIIS:\n",
    "    diis = DIIS_helper()\n",
    "MAX_ITER = 100\n",
    "CONV = 1.e-8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## B3LYP 极化率"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "B3LYP 所涉及到的 CP-KS 方程与 CP-HF 唯一的不同是在四脚标张量 $\\textbf{A}$ 中存在 DFT 所特有的响应张量 $\\textbf{G}$；$\\textbf{G}$ 相当于 DFT 势函数的导数，或者说 DFT Kernel 的二阶导数．我们通过构建张量 $\\textbf{A}$ 的过程来了解张量 $\\textbf{G}$ 的实际应用方式．\n",
    "\n",
    "我们知道，在 HF 方法中，可以使用直接 CP-HF 方法，也可以使用迭代法求解．在 CP-KS 方程中，一般不直接导出 $\\mathbf{G}$ 张量，因此这里始终会使用迭代法求解．迭代求解 B3LYP 极化率所需要的旋转矩阵 $U_{ai}^g$ 的公式是\n",
    "\n",
    "\\begin{equation}\n",
    "U_{ai}^g = \\frac{1}{\\varepsilon_i - \\varepsilon_a} \\left( A_{ij}^{ab} U_{bj}^g - \\mu_{ai}^g \\right)\n",
    "\\end{equation}\n",
    "\n",
    "其中，$A_{ij}^{ab}$ 的形式比 Hartree-Fock 下多了响应张量 $G_{ij}^{ab}$，并且要在 Exchange 积分贡献上乘上 B3LYP 泛函所指定的 Hartree-Fock 型 Exchange 贡献的系数 $c_\\mathrm{x}$ (对于 B3LYP 来说 $c_\\mathrm{x} = 0.2$)：\n",
    "\n",
    "\\begin{equation}\n",
    "A_{pr}^{qs} = 4 (pq|rs) - c_\\mathrm{x} (pr|qs) - c_\\mathrm{x} (ps|qr) + 4 G_{pr}^{qs}\n",
    "\\end{equation}\n",
    "\n",
    "在迭代求解的过程中，我们不需要真正地使用四脚标的张量，因此得到下述的表达式：\n",
    "\n",
    "\\begin{equation}\n",
    "A_{ij}^{ab} U_{bj} = C_{\\mu i} C_{\\nu a} (4 J_{\\mu \\nu} [D_{\\mu \\nu}] - c_\\mathrm{x} K_{\\nu \\mu} [D_{\\mu \\nu}] - c_\\mathrm{x} K_{\\mu \\nu} [D_{\\mu \\nu}] + 4 G_{\\mu \\nu}[D_{\\mu \\nu}])\n",
    "\\end{equation}\n",
    "\n",
    "其中，上述的广义密度定义为\n",
    "\n",
    "\\begin{equation}\n",
    "D_{\\mu \\nu} = C_{\\mu b} U_{bj}^g C_{\\nu j}\n",
    "\\end{equation}\n",
    "\n",
    "上述的广义密度在实际应用中，可以由 `psi4.core.JK.D` 方法给出；而 $G_{\\mu \\nu}[D_{\\mu \\nu}]$ 则可以由 DFT 积分引擎的 `psi4.core.VBase.compute_Vx` 在代入广义密度的情况下给出．下面我们就进行 CP-KS 方程的求解．似乎在 Psi4 中，DFT 积分的耗时比较多，所以下述代码的执行需要花一些时间．"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# JK 引擎初始化\n",
    "jk = psi4.core.JK.build(scf_wfn.basisset())\n",
    "jk.initialize()\n",
    "\n",
    "# 定义左矢为 Cv，右矢为空矩阵\n",
    "C_right_list = []\n",
    "for g in range(3):\n",
    "    jk.C_left_add(psi4.core.Matrix.from_array(Cv))\n",
    "    mat = psi4.core.Matrix(nbf, nvir)\n",
    "    C_right_list.append(np.asarray(mat))\n",
    "    jk.C_right_add(mat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 初猜\n",
    "U = - dip_vo / d_vo\n",
    "# 旧矩阵\n",
    "U_old = np.copy(U)\n",
    "# DIIS 初始化\n",
    "diis = []\n",
    "if USE_DIIS:\n",
    "    for _ in range(3):\n",
    "        diis.append(DIIS_helper())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "for it in range(1, MAX_ITER + 1):\n",
    "    # 更新 JK 引擎的右矢为 U_bj C_lj\n",
    "    for g in range(3):\n",
    "        C_right_list[g][:] = Co @ U[g].T\n",
    "\n",
    "    # 计算 J[D]、K[D] 积分\n",
    "    jk.compute()\n",
    "    for g in range(3):\n",
    "        J = np.asarray(jk.J()[g])\n",
    "        K = np.asarray(jk.K()[g])\n",
    "        \n",
    "        # 构建 DFT 的 G 响应\n",
    "        Vx = psi4.core.Matrix(nbf, nbf)  # 空矩阵储存 G 响应\n",
    "        D_jk = jk.D()[g]  # 相当于 Cv @ U[g] @ Co.T\n",
    "        V_pot.compute_Vx([D_jk], [Vx])  # 计算 G 响应到 Vx 矩阵中\n",
    "        Vx = np.asarray(Vx)  # 转换为 NumPy 矩阵\n",
    "        \n",
    "        # 可以更新 U 矩阵了\n",
    "        # Unew_ai^g = [ C_ui C_va (4 * JD_uv - cx KD_vu - cx KD_uv + 4 GD_uv) - mu_ai^g ] / (e_i - e_a)\n",
    "        U[g] = Cv.T @ (4 * J - 0.2 * K.T - 0.2 * K + 4 * Vx) @ Co - dip_vo[g]\n",
    "        U[g] /= d_vo\n",
    "\n",
    "    # 若使用了 DIIS 加速，则执行下述插值\n",
    "    if USE_DIIS:\n",
    "        for g in range(3):\n",
    "            diis[g].add(U[g], U[g] - U_old[g])\n",
    "            U[g] = diis[g].extrapolate()\n",
    "\n",
    "    # 检查收敛情况\n",
    "    rms = np.linalg.norm(U - U_old)\n",
    "    # print('CPHF Iteration {:3d}: RMS = {:14.10f}'.format(it, rms))\n",
    "\n",
    "    # 判断是否收敛\n",
    "    if (rms < CONV):\n",
    "        print(\"CPHF Converged in {:3d} iterations!\".format(it))\n",
    "        break\n",
    "    else:\n",
    "        U_old = np.copy(U)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "由此，我们可以给出 B3LYP 的极化率张量："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha_b3lyp = np.einsum(\"gai, fai -> fg\", U, dip_vo) * 4\n",
    "alpha_b3lyp.round(decimals=6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "不过上述的计算结果可能无法很好地与 Guassian 的结果对上，并且差距相当可观．Gaussian 输入卡：[B3LYP](include/b3lyp_dipole.gjf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha_b3lyp_gaussian = np.array([1.4146668, 7.2595695, 6.4526498])\n",
    "(np.diag(alpha_b3lyp) - alpha_b3lyp_gaussian).round(decimals=7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这可能与格点积分的选取有很大的关系．下面是在 Gaussian 中更改格点的精度所产生的 B3LYP 的极化率，不同格点精度的极化率的差基本上与上面的误差的数量级相同．因此，我们认为我们还是基本重复出了 B3LYP 的极化率．"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Grid UltraFine -      Fine: \",\n",
    "      np.array([1.4146668, 7.2595695, 6.4526498]) - np.array([1.4146945, 7.2598606, 6.4524672]))\n",
    "print(\"Grid SuperFine - UltraFine: \",\n",
    "      np.array([1.414654, 7.2595522, 6.4525927]) - np.array([1.4146668, 7.2595695, 6.4526498]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## XYG3 偶极矩与自然轨道"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "**提示**\n",
    "\n",
    "之后会约定，在符号右上标 $\\mathrm{s}$ 的为自洽场泛函 (XYG3 中则表示 B3LYP)，而上标 $\\mathrm{n}$ 的则为非自洽场泛函 (XYG3 中则表示其自身)．\n",
    "\n",
    "在后文中，会不加说明地代入 XYG3 的 Hartree-Fock 型 Exchange 积分缩放系数 $c_\\mathrm{x}^\\mathrm{n} = 0.8033$ 与 B3LYP 型的 $c_\\mathrm{x}^\\mathrm{s} = 0.2$，以及 XYG3 的 MP2 型 Correlation 积分缩放系数 $c_\\mathrm{c}^\\mathrm{n} = 0.3211$．\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "XYG3 的偶极矩求取方式与 MP2 非常类似，区别在于 CP-HF 方程更变为 CP-KS 方程，以及 MP2 拉格朗日矩阵更变为 XYG3 拉格朗日矩阵．由于 CP-KS 方程的导出前提是 $\\partial_\\xi F_{pq}^\\mathrm{s} = 0$，因此 XYG3 的 CP-KS 方程形式没有变化；但拉格朗日矩阵的导出则是从能量的梯度所产生的，因此其形式与 MP2 会有所不同．"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### XYG3 拉格朗日矩阵的构建"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "首先，我们写出 XYG3 拉格朗日量：\n",
    "\n",
    "\\begin{equation}\n",
    "L_{ai}^\\mathrm{n} = \\frac{1}{2} P_{jk}^{(2)} A^\\mathrm{s,}{}_{aj}^{ik} + \\frac{1}{2} P_{bc}^{(2)} A^\\mathrm{s,}{}_{ab}^{ic} - c_\\mathrm{c}^\\mathrm{n} T_{jk}^{ab} g_{ib}^{jk} + c_\\mathrm{c}^\\mathrm{n}  T_{ij}^{bc} g_{aj}^{bc} + F_{ai}^\\mathrm{n}\n",
    "\\end{equation}\n",
    "\n",
    "为了程序书写上的简便，并且区分 DFT 和与 MP2 拉格朗日矩阵类似的项区分开，我们定义张量\n",
    "\n",
    "\\begin{equation}\n",
    "a_{pr}^{qs} = 4 (pq|rs) - c_\\mathrm{x} (pr|qs) - c_\\mathrm{x} (ps|qr) = A^\\mathrm{s,}{}_{pr}^{qs} - 4 G^\\mathrm{s,}{}_{pr}^{qs}\n",
    "\\end{equation}\n",
    "\n",
    "同时定义弛豫密度的部分贡献\n",
    "\n",
    "\\begin{equation}\n",
    "p_{pq}^{(2)} = \\left\\{\n",
    "\\begin{matrix}\n",
    "P_{ij}^{(2)} = - c_\\mathrm{c}^\\mathrm{n} T_{ik}^{ab} t_{jk}^{ab}, & p = i \\text{ and } q = j \\\\\n",
    "P_{ab}^{(2)} = c_\\mathrm{c}^\\mathrm{n} T_{ij}^{ac} t_{ij}^{bc}, & p = a \\text{ and } q = b \\\\\n",
    "0, & \\text{otherwise}\n",
    "\\end{matrix}\n",
    "\\right.\n",
    "\\end{equation}\n",
    "\n",
    "那么我们可以重新写拉格朗日量为\n",
    "\n",
    "\\begin{equation}\n",
    "L_{ai}^\\mathrm{n} = \\frac{1}{2} p_{pq}^{(2)} a_{ap}^{iq} - c_\\mathrm{c}^\\mathrm{n} T_{jk}^{ab} g_{ib}^{jk} + c_\\mathrm{c}^\\mathrm{n} T_{ij}^{bc} g_{aj}^{bc} + 2 G_{ai}^\\mathrm{s} [p_{pq}^{(2)}] + F_{ai}^\\mathrm{n}\n",
    "\\end{equation}\n",
    "\n",
    "我们首先构造张量 $a_{pr}^{qs}$ 与弛豫密度部分贡献 $p_{pq}^{(2)}$："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T_ovov = 2 * t_ovov - t_ovov.swapaxes(1,3)  # 闭壳层的 T 张量\n",
    "P2_oo = - 0.3211 * np.einsum(\"iakb, jakb -> ij\", T_ovov, t_ovov)  # 占据-占据弛豫密度\n",
    "P2_vv = 0.3211 * np.einsum(\"iajc, ibjc -> ab\", T_ovov, t_ovov)  # 非占-非占弛豫密度\n",
    "p2_pq = np.block([\n",
    "    [P2_oo, np.zeros((nocc, nvir))],\n",
    "    [np.zeros((nvir, nocc)), P2_vv],\n",
    "])  # 弛豫密度部分贡献\n",
    "# A 张量的非 DFT 贡献\n",
    "a = 4 * g_pqrs - 0.2 * np.einsum(\"prqs -> pqrs\", g_pqrs) - 0.2 * np.einsum(\"psqr -> pqrs\", g_pqrs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "随后就可以构建 $L_{ai}^\\mathrm{n}$ 了："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 第一段：非 DFT 贡献\n",
    "L_vo = 0.5 * np.einsum(\"pq, aipq -> ai\", p2_pq, a[nocc:, :nocc, :, :])\n",
    "L_vo += - 0.3211 * np.einsum(\"jakb, ijbk -> ai\", T_ovov, g_pqrs[:nocc, :nocc, nocc:, :nocc])\n",
    "L_vo += 0.3211 * np.einsum(\"ibjc, abjc -> ai\", T_ovov, g_pqrs[nocc:, nocc:, :nocc, nocc:])\n",
    "# 第二段：DFT 贡献\n",
    "Gs_vo = psi4.core.Matrix(nbf, nbf)\n",
    "V_pot.compute_Vx([psi4.core.Matrix.from_array(C @ p2_pq @ C.T)], [Gs_vo])\n",
    "L_vo += 2 * Cv.T @ Gs_vo @ Co\n",
    "# 第三段：非自洽 Fock 矩阵\n",
    "L_vo += Fn_pq[nocc:, :nocc]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### XYG3 CP-KS 方程"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "随后解下述 CP-KS 方程即可：\n",
    "\n",
    "\\begin{equation}\n",
    "P_{ai}^{(2)}  = \\frac{1}{\\varepsilon_i - \\varepsilon_a} \\big( A^\\mathrm{s,}{}_{ij}^{ab} P_{bj}^{(2)} + L_{ai} \\big)\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# JK 引擎\n",
    "jk = psi4.core.JK.build(scf_wfn.basisset())\n",
    "jk.initialize()\n",
    "jk.C_left_add(psi4.core.Matrix.from_array(Cv))\n",
    "mat = psi4.core.Matrix(nbf, nvir)\n",
    "C_right = np.asarray(mat)\n",
    "jk.C_right_add(mat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 初猜与收敛设定\n",
    "P2_vo = L_vo / d_vo\n",
    "P2_vo_old = np.copy(P2_vo)\n",
    "diis = None\n",
    "if USE_DIIS:\n",
    "    diis = DIIS_helper()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "# CP-HF 方程\n",
    "for it in range(1, MAX_ITER + 1):\n",
    "    C_right[:] = Co @ P2_vo.T\n",
    "    jk.compute()\n",
    "    J = np.asarray(jk.J()[0])\n",
    "    K = np.asarray(jk.K()[0])\n",
    "    Vx = psi4.core.Matrix(nbf, nbf)\n",
    "    D_jk = jk.D()[0]\n",
    "    V_pot.compute_Vx([D_jk], [Vx])\n",
    "    Vx = np.asarray(Vx)\n",
    "    P2_vo = Cv.T @ (4 * J - 0.2 * K.T - 0.2 * K + 4 * Vx) @ Co + L_vo\n",
    "    P2_vo /= d_vo\n",
    "\n",
    "    # 若使用了 DIIS 加速，则执行下述插值\n",
    "    if USE_DIIS:\n",
    "        diis.add(P2_vo, P2_vo - P2_vo_old)\n",
    "        P2_vo = diis.extrapolate()\n",
    "\n",
    "    # 检查收敛情况\n",
    "    rms = np.linalg.norm(P2_vo - P2_vo_old)\n",
    "    # print('CPHF Iteration {:3d}: RMS = {:14.10f}'.format(it, rms))\n",
    "\n",
    "    # 判断是否收敛\n",
    "    if (rms < CONV):\n",
    "        print(\"CPHF Converged in {:3d} iterations!\".format(it))\n",
    "        break\n",
    "    else:\n",
    "        P2_vo_old = np.copy(P2_vo)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最后重构弛豫密度矩阵\n",
    "\n",
    "\\begin{equation}\n",
    "P_{pq}^{(2)} =\n",
    "\\begin{pmatrix}\n",
    "P_{ij}^{(2)} & P_{ia}^{(2)} \\\\\n",
    "P_{ai}^{(2)} & P_{ab}^{(2)}\n",
    "\\end{pmatrix}\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P2_pq = np.block([\n",
    "    [P2_oo, P2_vo.T],\n",
    "    [P2_vo, P2_vv]\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "P2 = C @ P2_pq @ C.T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### XYG3 自然轨道占据数与偶极矩"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "有了弛豫密度后，许多自然轨道占据数与偶极矩就非常容易给出．自然轨道表示如下："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义自洽场密度\n",
    "D_pq = np.zeros_like(D)\n",
    "for i in range(5):\n",
    "    D_pq[i, i] = 1\n",
    "# 对自洽场密度与弛豫密度的和进行对角化\n",
    "no_occ, no_coef = np.linalg.eigh(P2_pq + D_pq)\n",
    "# 从而得到自然轨道的占据数\n",
    "no_occ[::-1] * 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "B3LYP 贡献的电子云偶极矩与原子核偶极矩表示如下．Gaussian 输入卡：[B3LYP](include/b3lyp_dipole.gjf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 原子核偶极矩\n",
    "mol_geo = np.asarray(mol.geometry())\n",
    "neu_charge = []\n",
    "for neu in range(mol_geo.shape[0]):\n",
    "    neu_charge.append(mol.charge(neu))\n",
    "dip_mol = np.zeros(3)\n",
    "for neu in range(mol_geo.shape[0]):\n",
    "    for ind in range(3):\n",
    "        dip_mol[ind] += neu_charge[neu] * mol_geo[neu][ind]\n",
    "\n",
    "# B3LYP 偶极矩\n",
    "dip_b3lyp = 2 * np.einsum(\"fuv, uv -> f\", dip, D)\n",
    "dip_b3lyp = dip_b3lyp.round(decimals=6)\n",
    "# 验证结果\n",
    "psi4.compare_values((dip_mol + dip_b3lyp)[2], 1.031112, 6, 'B3LYP Energy')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "XYG3 泛函贡献的偶极矩部分则可以写为弛豫部分与 B3LYP 贡献部分的加和："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dip_xyg3 = 2 * np.einsum(\"fuv, uv -> f\", dip, P2)\n",
    "dip_xyg3 = dip_xyg3.round(decimals=6)\n",
    "dip_tot_xyg3 = dip_mol + dip_b3lyp + dip_xyg3\n",
    "dip_tot_xyg3"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}