docs/source/derivonce/grad_gga.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GGA 泛函核坐标梯度"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们已经将 RHF 与 MP2 的核坐标梯度表达式求得了；但我们仍然没有达到求解 XYG3 型泛函的目标。从 RHF 方法到 MP2 方法是一个突跃，我们需要掌握 U 矩阵、A 张量的计算方式 (CP-HF 方程)，以及对 MP2 方法公式相当繁杂的推导。XYG3 型泛函的绝大多数公式推导，都能从 MP2 公式中获得。另一个突跃会是这篇文档所述的从 RHF 到 GGA 方法；在这个过程中，我们需要对交换相关能的梯度作推演与计算。\n",
    "\n",
    "这一节，我们首先以 B3LYP 为例 (GGA 在通篇文档中代表的是使用了 GGA 泛函的计算方法，也包括杂化泛函)，计算 GGA 自洽场的核坐标梯度。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 准备工作"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib notebook\n",
    "\n",
    "from pyscf import gto, scf, dft, dft, lib\n",
    "import numpy as np\n",
    "from functools import partial\n",
    "import warnings\n",
    "from matplotlib import pyplot as plt\n",
    "from pyxdh.Utilities import NucCoordDerivGenerator, DipoleDerivGenerator, NumericDiff, GridHelper, KernelHelper\n",
    "from pyxdh.DerivOnce import GradSCF\n",
    "\n",
    "np.einsum = partial(np.einsum, optimize=[\"greedy\", 1024 ** 3 * 2 / 8])\n",
    "np.allclose = partial(np.allclose, atol=1e-6, rtol=1e-4)\n",
    "np.set_printoptions(5, linewidth=150, suppress=True)\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<pyscf.gto.mole.Mole at 0x7f6dc5311b50>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mol = gto.Mole()\n",
    "mol.atom = \"\"\"\n",
    "O  0.0  0.0  0.0\n",
    "O  0.0  0.0  1.5\n",
    "H  1.0  0.0  0.0\n",
    "H  0.0  0.7  1.0\n",
    "\"\"\"\n",
    "mol.basis = \"6-31G\"\n",
    "mol.verbose = 0\n",
    "mol.build()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "需要注意到，我们在这里需要用 DFT 模块，它的定义还包含格点积分。由于格点与分子直接挂钩，因此我们使用下述 `mol_to_grids` 定义从分子到 (75, 302) 格点的函数："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mol_to_grids(mol, atom_grid=(75, 302)):\n",
    "    grids = dft.Grids(mol)\n",
    "    grids.atom_grid = atom_grid\n",
    "    grids.becke_scheme = dft.gen_grid.stratmann\n",
    "    grids.prune = None\n",
    "    grids.build()\n",
    "    return grids\n",
    "grids = mol_to_grids(mol)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们也需要从分子到计算实例的程序 `mol_to_scf`，它主要用来生成 pyxdh 的 GGA 梯度计算实例，以及数值梯度的实例："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mol_to_scf(mol):\n",
    "    scf_eng = dft.RKS(mol)\n",
    "    scf_eng.grids = mol_to_grids(mol)\n",
    "    scf_eng.xc = \"B3LYPg\"\n",
    "    scf_eng.conv_tol = 1e-10\n",
    "    return scf_eng.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "GGA 与 RHF 一样，都使用 `GradSCF` 来实例化。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "gradh = GradSCF({\"scf_eng\": mol_to_scf(mol)})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "nmo, nao, natm, nocc, nvir = gradh.nao, gradh.nao, gradh.natm, gradh.nocc, gradh.nvir\n",
    "mol_slice = gradh.mol_slice\n",
    "so, sv, sa = gradh.so, gradh.sv, gradh.sa\n",
    "C, Co, Cv, e, eo, ev, D = gradh.C, gradh.Co, gradh.Cv, gradh.e, gradh.eo, gradh.ev, gradh.D\n",
    "H_0_ao, S_0_ao, eri0_ao, F_0_ao = gradh.H_0_ao, gradh.S_0_ao, gradh.eri0_ao, gradh.F_0_ao\n",
    "H_0_mo, S_0_mo, eri0_mo, F_0_mo = gradh.H_0_mo, gradh.S_0_mo, gradh.eri0_mo, gradh.F_0_mo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_natm_3(mat: np.ndarray):\n",
    "    shape = list(mat.shape)\n",
    "    shape = [int(shape[0] / 3), 3] + shape[1:]\n",
    "    return mat.reshape(shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "H_1_ao, S_1_ao, eri1_ao = to_natm_3(gradh.H_1_ao), to_natm_3(gradh.S_1_ao), to_natm_3(gradh.eri1_ao)\n",
    "H_1_mo, S_1_mo, eri1_mo = to_natm_3(gradh.H_1_mo), to_natm_3(gradh.S_1_mo), to_natm_3(gradh.eri1_mo)\n",
    "U_1 = to_natm_3(gradh.U_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "但与 RHF 不同的是，我们需要进行格点积分。我们定义 `grdh` 为格点的辅助助手，`kerh` 为泛函核的格点。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "grdh = GridHelper(mol, grids, D)\n",
    "kerh = KernelHelper(grdh, \"B3LYPg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "同时，我们也需要杂化系数 `cx` $c_\\mathrm{x}$："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.2"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cx = gradh.cx\n",
    "cx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def grad_generator(mol):\n",
    "    scf_eng = mol_to_scf(mol)\n",
    "    config = {\"scf_eng\": scf_eng}\n",
    "    return GradSCF(config)\n",
    "gradn = NucCoordDerivGenerator(mol, grad_generator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GGA 能量梯度"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "首先我们回顾 GGA (B3LYP) 的能量计算公式：\n",
    "\n",
    "$$\n",
    "E_\\mathrm{elec} = h_{\\mu \\nu} D_{\\mu \\nu} + \\frac{1}{2} D_{\\mu \\nu} (\\mu \\nu | \\kappa \\lambda) D_{\\kappa \\lambda} - \\frac{c_\\mathrm{x}}{4} D_{\\mu \\nu} (\\mu \\kappa | \\nu \\lambda) D_{\\kappa \\lambda} + f \\rho\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "其中，前三项分别是 Hamiltonian Core、Coulomb、Exchange 积分对总能量的贡献，使用了 Einstein Summation 进行了符号的简化；而第四项是交换相关能，所用的是本文档特化的简化，即\n",
    "\n",
    "$$\n",
    "\\sum_{w} w_g f_g \\rho_g \\Leftarrow f \\rho \\Rightarrow \\int f[\\rho] \\rho(\\boldsymbol{r}) \\, \\mathrm{d} \\boldsymbol{r}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "两种一般来说更容易接受的写法中，右边一种是积分方式，即泛函核 $f[\\rho]$ 与电子态密度 $\\rho(\\boldsymbol{r})$ 乘积的积分；而左边则是将积分元 $\\mathrm{d} \\boldsymbol{r}$ 离散化为带权重格点，随后对这些格点求和。我们简化的主要目的是让公式能与程序能作对应。我们不妨用下述两行代码，验证一下 B3LYP 下的分子的电子态能量 $E_\\mathrm{elec}$："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-189.26221747920502"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gradh.scf_eng.energy_elec()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-189.2622174792067"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(\n",
    "    + np.einsum(\"uv, uv -> \", H_0_ao, D)\n",
    "    + 0.5 * np.einsum(\"uv, uvkl, kl -> \", D, eri0_ao, D)\n",
    "    - 0.25 * cx * np.einsum(\"uv, ukvl, kl -> \", D, eri0_ao, D)\n",
    "    + np.einsum(\"g, g -> \", kerh.exc, grdh.rho_0)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 符号定义与交换相关能全导数"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们 [曾经](../qcbasic/basic_gga.ipynb) 对轨道、密度和泛函格点的符号作过补充定义，这里列举如下："
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "**记号说明：轨道函数或格点**\n",
    "\n",
    "* $\\phi$ 统一代表原子轨道函数，以电子坐标为自变量\n",
    "\n",
    "* $\\phi_\\mu$ 代表原子轨道 $\\mu$ 所对应的原子轨道函数\n",
    "\n",
    "* $\\phi_{r \\mu} = \\partial_r \\phi_\\mu$ 代表原子轨道在电子坐标分量 $r$ 下的偏导数\n",
    "\n",
    "* $\\phi_{r w \\mu} = \\partial_r \\partial_w \\phi_\\mu$ 代表原子轨道在电子坐标分量 $r$ 与 $w$ 下的二阶偏导数\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "**记号说明：密度函数或格点**\n",
    "\n",
    "* $\\rho$ 代表电子态密度密度\n",
    "\n",
    "* $\\rho_r = \\partial_r \\rho$\n",
    "\n",
    "* $\\rho_{rw} = \\partial_r \\partial_w \\rho$\n",
    "\n",
    "* $\\gamma = \\rho_r \\rho_r$ 表示密度梯度量\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "**记号说明：泛函格点**\n",
    "\n",
    "* $f$ 代表泛函核；泛函核满足关系：在函数图景下 $E_\\mathrm{xc} = \\int f[\\rho] \\rho(\\boldsymbol{r}) \\, \\mathrm{d} \\boldsymbol{r}$，或格点积分下，$E_\\mathrm{xc} = f \\rho$\n",
    "\n",
    "* $f_\\rho = \\partial_\\rho (f \\rho)$；**注意**不是 $\\partial_\\rho f$，**这种记号可能引起歧义但足够简洁**\n",
    "\n",
    "* $f_\\gamma = \\partial_\\gamma (f \\rho)$\n",
    "\n",
    "* $f_{\\rho \\gamma} = \\partial_\\rho \\partial_\\gamma (f \\rho)$，其它高阶导数同理\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们以后仍然使用这些定义。但我们还会补充定义下述符号："
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "**记号补充：导数相关与原子核相关格点**\n",
    "\n",
    "* $\\phi_\\mu^\\mathbb{A} = \\partial_\\mathbb{A} \\phi_\\mu$ 表示轨道在 $\\mathbb{A}$ 下的导数\n",
    "\n",
    "* $\\phi_{t \\mu_A}$ 表示在 $\\phi_\\mu$ 关于电子坐标分量 $t$ 下的导数，但若 $\\mu$ 作为 Gaussian 函数的中心并非 $A$ 原子核，则值为零\n",
    "\n",
    "* $\\rho^\\mathbb{A}_r$ 是 $\\partial_r \\rho$ 的 Skeleton 导数，而 **并非** $\\partial_\\mathbb{A} \\partial_r \\rho$\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "现在我们对 $f \\rho$ 进行导数计算。依据链式法则，并且留意到 $\\gamma = \\rho_r \\rho_r$ (对 $r$ 坐标分量求和)，我们可以导出，\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\frac{\\partial (f \\rho)}{\\partial \\mathbb{A}}\n",
    "&= \\frac{\\partial (f \\rho)}{\\partial \\rho} \\frac{\\partial \\rho}{\\partial \\mathbb{A}} + \\frac{\\partial (f \\rho)}{\\partial \\gamma} \\frac{\\partial \\gamma}{\\partial \\mathbb{A}} \\\\\n",
    "&= \\frac{\\partial (f \\rho)}{\\partial \\rho} \\frac{\\partial \\rho}{\\partial \\mathbb{A}} + 2 \\rho_r \\frac{\\partial (f \\rho)}{\\partial \\gamma} \\frac{\\partial \\rho_r}{\\partial \\mathbb{A}} \\\\\n",
    "&= f_\\rho \\partial_\\mathbb{A} \\rho + 2 f_\\gamma \\rho_r \\partial_\\mathbb{A} \\rho_r\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "但我们不能再推演下去了，因为我们尚不知道如何计算 $\\partial_\\mathbb{A} \\rho$ 与 $\\partial_\\mathbb{A} \\rho_r$。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 密度格点的导数"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们曾经提及过，对于密度矩阵 $D_{\\mu \\nu}$，其关于 $\\partial_\\mathbb{A}$ 的导数为\n",
    "\n",
    "$$\n",
    "\\frac{\\partial D_{\\mu \\nu}}{\\partial \\mathbb{A}} = 2 U_{mi}^\\mathbb{A} (C_{\\mu m} C_{\\nu i} + C_{\\mu i} C_{\\nu m})\n",
    "$$\n",
    "\n",
    "上面的所有的导数结果都是非 Skeleton 的。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "而在密度泛函中，密度并非是由密度矩阵 $D_{\\mu \\nu}$ 所表示，而是密度格点 $\\rho$ 表示。密度格点是通过下式给出的，这我们以前也有所提及：\n",
    "\n",
    "$$\n",
    "\\rho = \\phi_\\mu \\phi_\\nu D_{\\mu \\nu}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "由于显式地引入了轨道，因此密度的格点 **存在** Skeleton 导数。我们下面就具体地讨论 $\\mathbb{A} = A_t$ 即被求导量为核坐标分量的情况。为此，我们先列举下述结论：\n",
    "\n",
    "$$\n",
    "\\phi_\\mu^{A_t} = \\partial_{A_t} \\phi_\\mu = - \\partial_t \\phi_{\\mu_A} = - \\phi_{t \\mu_A}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "上述结论已经在 [RHF Skeleton 导数](grad_rhf_skeleton.ipynb#解析导数-(2)-最终表达式) 文档中有较为详细的论述了，这里就不展开了。因此，\n",
    "\n",
    "$$\n",
    "\\frac{\\partial \\rho}{\\partial A_t} = - \\phi_{t \\mu_A} \\phi_\\nu D_{\\mu \\nu} - \\phi_\\mu \\phi_{t \\nu_A} D_{\\mu \\nu} + 4 \\phi_\\mu \\phi_\\nu U_{mi}^{A_t} C_{\\mu m} C_{\\nu i}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "其中，最后与 U 矩阵有关的部分我们单独考虑，并定义密度格点的 Skeleton 导数 `A_rho_1` $\\rho^{A_t}$ 为上式的前两项 (维度为 $(A, t, r, g)$，其中最后一维度为格点维度)\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\rho^{A_t}\n",
    "&= - \\phi_{t \\mu_A} \\phi_\\nu D_{\\mu \\nu} - \\phi_\\mu \\phi_{t \\nu_A} D_{\\mu \\nu} \\\\\n",
    "&= - 2 \\phi_{t \\mu_A} \\phi_\\nu D_{\\mu \\nu}\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "对程序要作补充的是，尽管 $D_{\\mu \\nu}$ 处并没有写成 $D_{\\mu_A \\nu}$，但由于前面 $\\phi_{t \\mu_A}$ 中要求 $\\mu$ 必须要在 $A$ 原子核上，因此在实际写程序的时候确实要用 $D_{\\mu_A \\nu}$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "A_rho_1 = np.zeros((natm, 3, grdh.ngrid))\n",
    "for A in range(natm):\n",
    "    sA = mol_slice(A)\n",
    "    A_rho_1[A] = - 2 * np.einsum(\"tgu, gv, uv -> tg\", grdh.ao_1[:, :, sA], grdh.ao_0, D[sA])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "**任务 (1)**\n",
    "\n",
    "请证明上述等式的第二个等号。\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "pyxdh 中，`A_rho_1` 是用来计算 $\\rho^{A_t}$："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(A_rho_1, grdh.A_rho_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 密度梯度格点的导数"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "处理 $\\partial_{A_t} \\rho_r$ 的方式也是一样的。我们先需要回顾一下 $\\rho_r$ 的定义：\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\rho_r = \\frac{\\partial \\rho}{\\partial r}\n",
    "&= \\phi_{r \\mu} \\phi_\\nu D_{\\mu \\nu} + \\phi_\\mu \\phi_{r \\nu} D_{\\mu \\nu} \\\\\n",
    "&= 2 \\phi_{r \\mu} \\phi_\\nu D_{\\mu \\nu}\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "那么，\n",
    "\n",
    "$$\n",
    "\\frac{\\partial \\rho_r}{\\partial A_t} = - 2 \\phi_{tr \\mu_A} \\phi_\\nu D_{\\mu \\nu} - 2 \\phi_{r \\mu} \\phi_{t \\nu_A} D_{\\mu \\nu} + 4 \\phi_{r \\mu} \\phi_\\nu U_{mi}^{A_t} (C_{\\mu m} C_{\\nu i} + C_{\\mu i} C_{\\nu m})\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们定义 `A_rho_2` $\\partial_r^{A_t}$ (维度为 $(A, t, r, g)$) 为\n",
    "\n",
    "$$\n",
    "\\rho_r^{A_t} = - 2 \\phi_{tr \\mu_A} \\phi_\\nu D_{\\mu \\nu} - 2 \\phi_{r \\mu} \\phi_{t \\nu_A} D_{\\mu \\nu}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "A_rho_2 = np.zeros((natm, 3, 3, grdh.ngrid))\n",
    "for A in range(natm):\n",
    "    sA = mol_slice(A)\n",
    "    A_rho_2[A]  = - 2 * np.einsum(\"trgu, gv, uv -> trg\", grdh.ao_2[:, :, :, sA], grdh.ao_0, D[sA])\n",
    "    A_rho_2[A] += - 2 * np.einsum(\"rgu, tgv, uv -> trg\", grdh.ao_1, grdh.ao_1[:, :, sA], D[:, sA])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在 pyxdh 中，有 `A_rho_2` 与之对应："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(A_rho_2, grdh.A_rho_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-info\">\n",
    "\n",
    "**任务 (2)**\n",
    "\n",
    "我们曾经不加证明地在 $\\partial_{A_t} \\rho$ 表达式中，利用到\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\partial_{A_t} \\rho\n",
    "&\\leftarrow 2 \\phi_\\mu \\phi_\\nu U_{mi}^{A_t} (C_{\\mu m} C_{\\nu i} + C_{\\mu i} C_{\\nu m}) \\\\\n",
    "&= 4 \\phi_\\mu \\phi_\\nu U_{mi}^{A_t} C_{\\mu m} C_{\\nu i}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "但对于 $\\partial_{A_t} \\rho_r$，我们并没有作简化：\n",
    "\n",
    "$$\n",
    "\\partial_{A_t} \\rho_r \\leftarrow 4 \\phi_{r \\mu} \\phi_\\nu U_{mi}^{A_t} (C_{\\mu m} C_{\\nu i} + C_{\\mu i} C_{\\nu m})\n",
    "$$\n",
    "\n",
    "请简述原因并用程序验证。\n",
    "\n",
    "后文可能会经常作一些与对称性有关的变换，譬如 $\\partial_{A_t} \\rho_r$ 的 U 导数还可以表示为\n",
    "\n",
    "$$\n",
    "\\partial_{A_t} \\rho_r \\leftarrow 4 (\\phi_{r \\mu} \\phi_\\nu + \\phi_\\mu \\phi_{r \\nu}) U_{mi}^{A_t} C_{\\mu m} C_{\\nu i}\n",
    "$$\n",
    "\n",
    "读者可能需要熟悉和适应这种变化。\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### U 导数与 Fock 矩阵的关系"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们再回到\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\partial_{A_t} E_\\mathrm{elec} \\xleftarrow{\\text{GGA part}} \\frac{\\partial (f \\rho)}{\\partial A_t}\n",
    "&= f_\\rho \\partial_{A_t} \\rho + 2 f_\\gamma \\rho_r \\partial_{A_t} \\rho_r \\\\\n",
    "&= f_\\rho \\rho^{A_t} + 2 f_\\gamma \\rho_r \\rho_r^{A_t} + (f_\\rho \\phi_\\mu \\phi_\\nu + 2 f_\\gamma \\rho_r \\phi_{r \\mu} \\phi_\\nu + 2 f_\\gamma \\rho_r \\phi_\\mu \\phi_{r \\nu}) \\cdot 2 (C_{\\mu m} C_{\\nu i} + C_{\\mu i} C_{\\nu p}) U_{mi}^{A_t}\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们会发现，上式中出现了\n",
    "\n",
    "$$\n",
    "v_{\\mu \\nu}^\\mathrm{xc} [\\rho] = f_\\rho \\phi_\\mu \\phi_\\nu + 2 f_\\gamma \\rho_r (\\phi_{r \\mu} \\phi_{\\nu} + \\phi_{\\mu} \\phi_{r \\nu})\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "因此，U 矩阵导数部分有\n",
    "\n",
    "$$\n",
    "\\partial_{A_t} E_\\mathrm{elec} \\xleftarrow{\\text{GGA part}} \\partial_{A_t} (f \\rho) \\xleftarrow{\\text{U derivative}} v_{\\mu \\nu}^\\mathrm{xc} \\cdot 2 (C_{\\mu m} C_{\\nu i} + C_{\\mu i} C_{\\nu p}) U_{mi}^{A_t} = v_{\\mu \\nu}^\\mathrm{xc} \\partial_{A_t} D_{\\mu \\nu}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "当我们联系到类似于 HF 部分的贡献为\n",
    "\n",
    "$$\n",
    "E_\\mathrm{elec} \\xleftarrow{\\text{HF part}} h_{\\mu \\nu} D_{\\mu \\nu} + \\frac{1}{2} D_{\\mu \\nu} (\\mu \\nu | \\kappa \\lambda) D_{\\kappa \\lambda} - \\frac{c_\\mathrm{x}}{4} D_{\\mu \\nu} (\\mu \\kappa | \\nu \\lambda) D_{\\kappa \\lambda}\n",
    "$$\n",
    "\n",
    "及其 U 导数\n",
    "\n",
    "$$\n",
    "\\partial_{A_t} E_\\mathrm{elec} \\xleftarrow{\\text{HF part}} \\big( h_{\\mu \\nu} + (\\mu \\nu | \\kappa \\lambda) D_{\\kappa \\lambda} - \\frac{c_\\mathrm{x}}{2} (\\mu \\kappa | \\nu \\lambda) D_{\\kappa \\lambda} \\big) \\partial_{A_t} D_{\\mu \\nu}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "因此，电子态总能量全部的 U 导数可以写为\n",
    "\n",
    "$$\n",
    "\\partial_{A_t} E_\\mathrm{elec} \\xleftarrow{\\text{U derivative}} \\big( h_{\\mu \\nu} + (\\mu \\nu | \\kappa \\lambda) D_{\\kappa \\lambda} - \\frac{c_\\mathrm{x}}{2} (\\mu \\kappa | \\nu \\lambda) D_{\\kappa \\lambda} + v_{\\mu \\nu}^\\mathrm{xc} \\big) \\partial_{A_t} D_{\\mu \\nu} = F_{\\mu \\nu} \\partial_{A_t} D_{\\mu \\nu} = - 2 F_{ij} S_{ij}^{A_t}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们再将 $\\partial_{A_t} E_\\mathrm{elec}$ 的 Skeleton 导数部分列举如下：\n",
    "\n",
    "$$\n",
    "\\partial_{A_t} E_\\mathrm{elec} \\xleftarrow{\\text{Skeleton derivative}} h_{\\mu \\nu}^{A_t} D_{\\mu \\nu} + \\frac{1}{2} D_{\\mu \\nu} (\\mu \\nu | \\kappa \\lambda)^{A_t} D_{\\kappa \\lambda} - \\frac{c_\\mathrm{x}}{4} D_{\\mu \\nu} (\\mu \\kappa | \\nu \\lambda)^{A_t} D_{\\kappa \\lambda} + f_\\rho \\rho^{A_t} + 2 f_\\gamma \\rho_r \\rho_r^{A_t}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 电子态能量总导数"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "有了上面的推导之后，我们就可以一口气地将所有电子态贡献项列出："
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\partial_{A_t} E_\\mathrm{elec} = h_{\\mu \\nu}^{A_t} D_{\\mu \\nu} + \\frac{1}{2} D_{\\mu \\nu} (\\mu \\nu | \\kappa \\lambda)^{A_t} D_{\\kappa \\lambda} - \\frac{c_\\mathrm{x}}{4} D_{\\mu \\nu} (\\mu \\kappa | \\nu \\lambda)^{A_t} D_{\\kappa \\lambda} + f_\\rho \\rho^{A_t} + 2 f_\\gamma \\rho_r \\rho_r^{A_t} - 2 F_{ij} S_{ij}^{A_t}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-2.27471, -0.79557, -9.07091],\n",
       "       [-0.37246, -2.30276, 10.1379 ],\n",
       "       [ 2.70067, -0.03745, -0.61219],\n",
       "       [-0.05351,  3.13578, -0.4548 ]])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "E_1 = (\n",
    "    + np.einsum(\"Atuv, uv -> At\", H_1_ao, D)\n",
    "    + 0.5 * np.einsum(\"uv, Atuvkl, kl -> At\", D, eri1_ao, D)\n",
    "    - 0.25 * cx * np.einsum(\"uv, Atukvl, kl -> At\", D, eri1_ao, D)\n",
    "    + np.einsum(\"g, Atg -> At\", kerh.fr, A_rho_1)\n",
    "    + 2 * np.einsum(\"g, rg, Atrg -> At\", kerh.fg, grdh.rho_1, A_rho_2)\n",
    "    - 2 * np.einsum(\"ij, Atij -> At\", F_0_mo[so, so], S_1_mo[:, :, so, so])\n",
    ")\n",
    "E_1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们可以用数值导数来验证上述结果："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[-2.2747 , -0.79556, -9.07092],\n",
       "       [-0.37246, -2.30276, 10.1379 ],\n",
       "       [ 2.70066, -0.03746, -0.61219],\n",
       "       [-0.0535 ,  3.13578, -0.4548 ]])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nd_E_0 = NumericDiff(gradn, lambda gradh: gradh.scf_eng.energy_elec()[0]).derivative\n",
    "nd_E_0.reshape(natm, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(E_1, nd_E_0.reshape(natm, 3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 参考任务解答"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 任务 (1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "由于待证等式对 $\\mu, \\nu$ 求和，那么我们将 $\\phi_\\mu \\phi_{t \\nu_A} D_{\\mu \\nu}$ 中的 $\\mu, \\nu$ 角标对换一下，并且利用 $D_{\\mu \\nu}$ 的对称性，就能立即得到 $\\phi_{t \\mu_A} \\phi_\\nu D_{\\mu \\nu}$。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 任务 (2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "首先，我们需要说明\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\partial_{A_t} \\rho\n",
    "&\\leftarrow 2 \\phi_\\mu \\phi_\\nu U_{mi}^{A_t} (C_{\\mu m} C_{\\nu i} + C_{\\mu i} C_{\\nu m}) \\\\\n",
    "&= 4 \\phi_\\mu \\phi_\\nu U_{mi}^{A_t} C_{\\mu m} C_{\\nu i}\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "程序如下："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(\n",
    "    + 2 * np.einsum(\"gu, gv, Atmi, um, vi -> Atg\", grdh.ao_0, grdh.ao_0, U_1[:, :, :, so], C, Co)\n",
    "    + 2 * np.einsum(\"gu, gv, Atmi, ui, vm -> Atg\", grdh.ao_0, grdh.ao_0, U_1[:, :, :, so], Co, C),\n",
    "    + 4 * np.einsum(\"gu, gv, Atmi, um, vi -> Atg\", grdh.ao_0, grdh.ao_0, U_1[:, :, :, so], C, Co)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "证明应当是很简单的：我们只要根据求和角标可交换，交换一下 $\\mu, \\nu$ 即可。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "但我们注意到\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "\\partial_{A_t} \\rho_r\n",
    "&\\leftarrow 4 \\phi_{r \\mu} \\phi_\\nu U_{mi}^{A_t} (C_{\\mu m} C_{\\nu i} + C_{\\mu i} C_{\\nu m}) \\\\\n",
    "&\\not\\equiv 8 \\phi_{r \\mu} \\phi_\\nu U_{mi}^{A_t} C_{\\mu m} C_{\\nu i}\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(\n",
    "    + 4 * np.einsum(\"rgu, gv, Atmi, um, vi -> Atgr\", grdh.ao_1, grdh.ao_0, U_1[:, :, :, so], C, Co)\n",
    "    + 4 * np.einsum(\"rgu, gv, Atmi, ui, vm -> Atgr\", grdh.ao_1, grdh.ao_0, U_1[:, :, :, so], Co, C),\n",
    "    + 8 * np.einsum(\"rgu, gv, Atmi, um, vi -> Atgr\", grdh.ao_1, grdh.ao_0, U_1[:, :, :, so], C, Co)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "之所以上述不恒等号成立，我们仍然是先看看，交换 $\\phi_{r \\mu} \\phi_\\nu U_{mi}^{A_t} C_{\\mu i} C_{\\nu m}$ 一项中的 $\\mu, \\nu$ 角标后，得到 $\\phi_\\mu \\phi_{r \\nu} U_{mi}^{A_t} C_{\\mu m} C_{\\nu i}$；它并不等价于 $\\phi_{r \\mu} \\phi_\\nu U_{mi}^{A_t} C_{\\mu m} C_{\\nu i}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们一般来说，总是希望将表达式化简来推导公式或编写程序；但这类看起来相当微妙的相等或不等关系，在处理的时候需要当心。"
   ]
  }
 ],
 "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.8.1"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "243.6px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}