docs/source/efficiency/mp2_eng.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 重新学习 MP2 能量笔记"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在 pyxdh 中，所有张量全部都储存在内存中，包括原子轨道基组的双电子积分。但在现实问题中，这对内存量与磁盘量上，是不允许的。\n",
    "\n",
    "所有的 MP2 算法，都多少会使用内存降低方法。我们需要逐个地考虑它们。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 初始化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from functools import partial\n",
    "np.einsum = partial(np.einsum, optimize=[\"greedy\", 1024 ** 3 * 2 / 8])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pyscf import gto, scf, mp, ao2mo, lib\n",
    "from pyscf.ao2mo import _ao2mo\n",
    "from functools import partial\n",
    "import h5py\n",
    "\n",
    "np.set_printoptions(4, suppress=True, linewidth=180)\n",
    "np.einsum = partial(np.einsum, optimize=[\"greedy\", 1024 ** 3 * 2 / 8])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<pyscf.gto.mole.Mole at 0x7f8865ac2be0>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mol = gto.Mole()\n",
    "mol.atom = \"\"\"\n",
    "N  0.  0.  0.\n",
    "H  1.5 0.  0.2\n",
    "H  0.1 1.2 0.\n",
    "H  0.  0.  1.\n",
    "\"\"\"\n",
    "mol.basis = \"6-31G\"\n",
    "mol.verbose = 0\n",
    "mol.build()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "nocc, nmo, nao, nbas = mol.nelec[0], mol.nao, mol.nao, mol.nbas\n",
    "nvir = nmo - nocc\n",
    "so, sv = slice(0, nocc), slice(nocc, nmo)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 最简单方法回顾"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最简单的方法是将原子轨道基的双电子积分完全储存到内存中，并作双电子积分转换。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "\\begin{align}\n",
    "(ia|jb) &= C_{\\mu i} C_{\\nu a} (\\mu \\nu | \\kappa \\lambda) C_{\\kappa j} C_{\\lambda b} \\\\\n",
    "D_{ij}^{ab} &= \\varepsilon_i - \\varepsilon_a + \\varepsilon_j - \\varepsilon_b \\\\\n",
    "t_{ij}^{ab} &= (ia|jb) / D_{ij}^{ab} \\\\\n",
    "T_{ij}^{ab} &= 2 t_{ij}^{ab} - t_{ij}^{ba} \\\\\n",
    "E_\\mathrm{corr} &= T_{ij}^{ab} t_{ij}^{ab} D_{ij}^{ab}\n",
    "\\end{align}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "mf_scf = scf.RHF(mol).run()\n",
    "C, e = mf_scf.mo_coeff, mf_scf.mo_energy\n",
    "Co, Cv = C[:, so], C[:, sv]\n",
    "eo, ev = e[so], e[sv]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-0.14554742350036615"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eri0_ao = mol.intor(\"int2e\")\n",
    "eri0_iajb = np.einsum(\"ui, va, uvkl, kj, lb -> iajb\", Co, Cv, eri0_ao, Co, Cv)\n",
    "D_iajb = eo[:, None, None, None] - ev[None, :, None, None] + eo[None, None, :, None] - ev[None, None, None, :]\n",
    "t_iajb = eri0_iajb / D_iajb\n",
    "T_iajb = 2 * t_iajb - t_iajb.swapaxes(-1, -3)\n",
    "np.einsum(\"iajb, iajb, iajb ->\", T_iajb, t_iajb, D_iajb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "整个过程最耗时部分，对于当前的问题反而是原子基组双电子积分 $(\\mu \\nu | \\kappa \\lambda)$ 的计算。这一步其实是 $O(N^4)$ 量级的。对于更大的分子，显然是原子轨道到分子轨道基转换的 $O(N^5)$ 量级更加耗时。因此不能说这里的的效率测评有意义。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10.4 ms ± 1.11 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 100\n",
    "eri0_ao = mol.intor(\"int2e\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.11 ms ± 232 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 100\n",
    "eri0_iajb = np.einsum(\"ui, va, uvkl, kj, lb -> iajb\", Co, Cv, eri0_ao, Co, Cv)\n",
    "D_iajb = eo[:, None, None, None] - ev[None, :, None, None] + eo[None, None, :, None] - ev[None, None, None, :]\n",
    "t_iajb = eri0_iajb / D_iajb\n",
    "T_iajb = 2 * t_iajb - t_iajb.swapaxes(-1, -3)\n",
    "np.einsum(\"iajb, iajb, iajb ->\", T_iajb, t_iajb, D_iajb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "通过 PySCF 默认路径的结果如下，可以验证上面结果的正确性。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-0.1455474235003661"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mf_mp2 = mp.MP2(mf_scf).run()\n",
    "mf_mp2.e_corr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "事实上，上面的双电子转换效率并不快。出于双电子积分本身的对称性、与 MP2 激发系数自身的对称性，PySCF 的默认程序 (incore 转换) 本身会更快一些。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3.02 ms ± 821 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 100\n",
    "mf_mp2 = mp.MP2(mf_scf).run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Incore 转换"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PySCF 默认使用 incore 转换。对于小分子体系而言，实际上原子基组的双电子积分 $(\\mu \\nu | \\kappa \\lambda)$ 是被储存下来的："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(7260,)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mf_scf._eri.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这个数值大小其实是对称性简化之后的双电子积分大小 (我们暂时规定了 $n_\\mathrm{AO} = n_\\mathrm{MO}$)\n",
    "\n",
    "$$\n",
    "\\frac{1}{2} \\left[ \\frac{n_\\mathrm{AO} (n_\\mathrm{AO} + 1)}{2} \\left( \\frac{n_\\mathrm{AO} (n_\\mathrm{AO} + 1)}{2} + 1 \\right) \\right]\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7260"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(nmo * (nmo + 1) // 2 * (nmo * (nmo + 1) // 2 + 1)) // 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "因此，这里的 incore 转换实际上要求内存占用至少包括原子基组双电子积分 $\\sim n_\\mathrm{AO}^4/8$。在此基础上，所有 $(ia|jb)$ 积分也全部储存，即额外内存 $n_\\mathrm{occ}^2 n_\\mathrm{vir}^2$。同时，转换过程还需要 $n_\\mathrm{occ} n_\\mathrm{vir} n_\\mathrm{AO}^2$ 大小。因此可以说是开销不计成本。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PySCF 中对 incore 的处理是非常简单粗暴的。直接提供简化或未简化的积分、以及用于缩并的四个系数矩阵就可以了。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(ao2mo.general(mf_scf._eri, (Co, Cv, Co, Cv)).reshape((nocc, nvir, nocc, nvir)), eri0_iajb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(ao2mo.general(eri0_ao, (Co, Cv, Co, Cv)).reshape((nocc, nvir, nocc, nvir)), eri0_iajb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "进行过转换的积分会用于给出 MP2 激发系数 $t_{ij}^{ab}$。但与 pyxdh 不同地是，pyxdh 使用维度 $(i, a, j, b)$ 储存激发系数，而 PySCF 使用维度 $(i, j, a, b)$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(mf_mp2.t2.swapaxes(-2, -3), t_iajb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Outcore 转换"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在 PySCF 中，outcore 转换必须要求 $(ia|jb)$ 大小的张量可以储存在内存中，因此需要比 $n_\\mathrm{occ}^2 n_\\mathrm{vir}^2$ 再大一些。它的实现可以非常简单，也可以非常复杂。\n",
    "\n",
    "但如果只求取 MP2 的能量，那么我们不一定真的需要这么大的张量。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 简单粗暴的 outcore"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "如果替换 `ao2mo.general` 的第一个参量为 `gto.Mole` 实例，那么会自动地使用 outcore 算法。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(ao2mo.general(mol, (Co, Cv, Co, Cv)).reshape((nocc, nvir, nocc, nvir)), eri0_iajb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.21 ms ± 489 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 100\n",
    "ao2mo.general(eri0_ao, (Co, Cv, Co, Cv))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "40.7 ms ± 4.35 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 10\n",
    "ao2mo.general(mol, (Co, Cv, Co, Cv))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 实际的 outcore (1)：壳层与原子对原子轨道的分割"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "之后几小节的讨论对象是 `mp.mp2._ao2mo_ovov`。首先，我们要学习或回顾原子轨道和轨道壳层的分割。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "不论任何量化程序，原子轨道的数量总是恒定的。它会用 $\\mu, \\nu, \\cdots$ 角标表示。但在实际的积分中，不是任何原子轨道的分割都可以带入计算。在 PySCF 中，积分的最小单元是壳层 (shell)。\n",
    "\n",
    "譬如，一个 $s$ 壳层对应 1 根原子轨道，一个 $p$ 壳层对应 3 根原子轨道。对于 $d$ 壳层，根据球谐 (`sph`) 或笛卡尔 (`cart`) 的不同，会分别取 5 或 6 根原子轨道。现在一般的基组都使用球谐轨道；笛卡尔基组已经很少使用了 (特殊的 6-31G 系列基组会使用)。\n",
    "\n",
    "这里顺便原子的分割。尽管 MP2 能量不需要用到，但核坐标梯度的程序会经常使用到。原子轨道之所以称为“原子轨道”，是因为它是以特定原子为中心的三维 Gaussian 函数。因此，原子轨道是可以对应到特定原子的。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "现在我们考虑程序的问题。壳层的分割可以用下面的函数表示："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[ 0  1  2  3  6  9 10 11 12 13 14 15]\n",
      "12\n"
     ]
    }
   ],
   "source": [
    "ao_loc = mol.ao_loc_nr()\n",
    "print(ao_loc)\n",
    "print(len(ao_loc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们知道，氨分子的 6-31G 基组有 15 根原子轨道。上面的结果意味着该分子有 12 个壳层。这 12 个壳层的起点体现在上面的列表了。\n",
    "\n",
    "第 3, 4 壳层 (起点分别为 3, 6) 不同于其他壳层；这一个壳层中包含 3 个原子轨道，因此可以推断为 p 轨道。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们这里回顾原子对原子轨道与轨道壳层的分割。PySCF 的原子分割是："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0,  5,  0,  9],\n",
       "       [ 5,  7,  9, 11],\n",
       "       [ 7,  9, 11, 13],\n",
       "       [ 9, 11, 13, 15]])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mol.aoslice_by_atom()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "上面的每一行代表每一个原子；第一个原子的壳层为 0-4，原子轨道为 0-8；第二个原子的壳层为 5-6，原子轨道 9-10；以此类推。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "现在遇到这样的现实的问题：假设氨分子非常大，我们的内存无法储存这种级别的分子。以什么标准对积分过程进行分割？\n",
    "\n",
    "在 outcore 过程中，第一步是先计算占据轨道的积分：\n",
    "\n",
    "$$\n",
    "(i \\nu | \\kappa j) = (\\mu \\nu | \\kappa \\lambda) C_{\\mu i} C_{\\lambda j}\n",
    "$$\n",
    "\n",
    "第二部才是计算最终的分子轨道积分：\n",
    "\n",
    "$$\n",
    "(ia|jb) = (i \\nu | \\kappa j) C_{\\nu a} C_{\\kappa b}\n",
    "$$\n",
    "\n",
    "由于第一步缩并过程需要完整的 $\\mu, \\lambda$，因此这两个原子轨道角标不能进行分割。待分割的角标是 $\\nu, \\kappa$。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "以多小的单位进行分割？在 PySCF 中，推荐的大小是 4 根原子轨道一组进行分割。由于电子积分的最小单位是壳层，因此如果基组包含 $d$ 轨道 (5 根原子轨道)，因此一般来说实际消耗会比 4 根轨道还要再大一些。\n",
    "\n",
    "之所以不使用更大的分割单位，我想是因为小一些的分割可以保证利用好原子轨道 ERI 积分的对称性；但不使用更小的分割，是为了减小积分函数的调用次数，避免不必要的调用时间。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PySCF 中的辅助函数 `ao2mo.outcore.balance_partition` 可以用于处理壳层的分割："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(0, 3, 3), (3, 4, 3), (4, 6, 4), (6, 10, 4), (10, 11, 1)]\n",
      "4\n"
     ]
    }
   ],
   "source": [
    "sh_ranges = ao2mo.outcore.balance_partition(ao_loc, 4)\n",
    "dmax = max(x[2] for x in sh_ranges)\n",
    "print(sh_ranges)\n",
    "print(dmax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "上述参量 `sh_ranges` 表示被分割的壳层。譬如以 `(4, 6, 4)` 为例，表示了第 2 个分割包含壳层 4 与壳层 5，该分割包含了 4 根轨道。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 实际的 outcore (2)：分割的原子轨道积分"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "随后的任务是分步地处理下述计算：\n",
    "\n",
    "$$\n",
    "(i \\nu_s | \\kappa_s j) = (\\mu \\nu_s | \\kappa_s \\lambda) C_{\\mu i} C_{\\lambda j}\n",
    "$$\n",
    "\n",
    "其中，下标 $s$ 表示分割。随后将 $(i \\nu_s | \\kappa_s j)$ 以维度 $(i, j, \\nu_s, \\kappa_s)$ 的形式储存到外部硬盘，并对该张量赋予名称 $s$。\n",
    "\n",
    "由于分割出来的原子轨道大小不超过壳层数，因此如果不考虑 $d, f$ 等高角动量轨道，那么内存消耗大约是 $16 n_\\mathrm{AO}^2$。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "一个额外的工作是，为了避免内存空间的重复分配耗时，我们先预置一块内存 `buf_eri`，专门用于储存临时积分 $(\\mu \\nu_s | \\kappa_s \\lambda)$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "buf_eri = np.empty((nao, dmax, dmax, nao))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "循环过程中，我们所需要使用的壳层分割和原子轨道分割列表分别用 `list_shell_slice` 与 `list_ao_slice` 储存："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "list_shell_slice, list_ao_slice = [], []\n",
    "for ip, (ish0, ish1, _) in enumerate(sh_ranges):\n",
    "    for jsh0, jsh1, nj in sh_ranges[:ip+1]:\n",
    "        list_shell_slice.append((ish0, ish1, jsh0, jsh1))\n",
    "        i0, i1 = ao_loc[ish0], ao_loc[ish1]\n",
    "        j0, j1 = ao_loc[jsh0], ao_loc[jsh1]\n",
    "        list_ao_slice.append((i0, i1, j0, j1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(0, 3, 0, 3) (0, 3, 0, 3)\n",
      "(3, 4, 0, 3) (3, 6, 0, 3)\n",
      "(3, 4, 3, 4) (3, 6, 3, 6)\n",
      "(4, 6, 0, 3) (6, 10, 0, 3)\n",
      "(4, 6, 3, 4) (6, 10, 3, 6)\n",
      "(4, 6, 4, 6) (6, 10, 6, 10)\n",
      "(6, 10, 0, 3) (10, 14, 0, 3)\n",
      "(6, 10, 3, 4) (10, 14, 3, 6)\n",
      "(6, 10, 4, 6) (10, 14, 6, 10)\n",
      "(6, 10, 6, 10) (10, 14, 10, 14)\n",
      "(10, 11, 0, 3) (14, 15, 0, 3)\n",
      "(10, 11, 3, 4) (14, 15, 3, 6)\n",
      "(10, 11, 4, 6) (14, 15, 6, 10)\n",
      "(10, 11, 6, 10) (14, 15, 10, 14)\n",
      "(10, 11, 10, 11) (14, 15, 14, 15)\n"
     ]
    }
   ],
   "source": [
    "for slice_shell, slice_ao in zip(list_shell_slice, list_ao_slice):\n",
    "    print(slice_shell, slice_ao)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "第一步完整的积分过程可以用下面的程序表示："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<KeysViewHDF5 ['0', '1', '10', '11', '12', '13', '14', '2', '3', '4', '5', '6', '7', '8', '9']>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ftmp = lib.H5TmpFile()\n",
    "count = 0\n",
    "for (ish0, ish1, jsh0, jsh1), (i0, i1, j0, j1) in zip(list_shell_slice, list_ao_slice):\n",
    "    # slice       mu          nu             kappa          lambda\n",
    "    shls_slice = (0, nbas) + (ish0, ish1) + (jsh0, jsh1) + (0, nbas)\n",
    "    # calculate  ( mu nu | kappa lambda ).  note that kappa is sliced, so aosym is none, i.e. ( mu nu | kappa lambda ) != ( mu nu | lambda kappa )\n",
    "    eri = mol.intor(\"int2e\", shls_slice=shls_slice, aosym=\"s1\", out=buf_eri)\n",
    "    # reshape to tensor (mu, nu, kappa, lambda) from one-dim array\n",
    "    eri.shape = (nao, (i1-i0), (j1-j0), nao)\n",
    "    # ( mu nu | kappa lambda ) -> ( i nu | kappa j )\n",
    "    tensor_ijvk = np.einsum(\"uvkl, ui, lj -> ijvk\", eri, Co, Co)\n",
    "    # dump result to h5py file\n",
    "    ftmp.create_dataset(str(count), data=tensor_ijvk)\n",
    "    count += 1\n",
    "ftmp.keys()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "现在我们验证其中一个积分。我们能看到第 4 个分割是 $\\nu_4$ 包含原子轨道 6-9，$\\kappa_4$ 包含原子轨道 3-5。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(6, 10, 3, 6)"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list_ao_slice[4]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "不妨直接地验证一下\n",
    "\n",
    "$$\n",
    "(i \\nu_4 | \\kappa_4 j) = (\\mu \\nu_4 | \\kappa_4 \\lambda) C_{\\mu i} C_{\\lambda j}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(np.einsum(\"uvkl, ui, kj -> ijvl\", mol.intor(\"int2e\"), Co, Co)[:, :, 6:10, 3:6], ftmp[\"4\"][()])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 实际的 outcore (3)：分子轨道的 ERI 积分导出"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最终的激发系数导出方式是\n",
    "\n",
    "$$\n",
    "(ia|jb) = (i \\nu | \\kappa j) C_{\\nu a} C_{\\kappa b}\n",
    "$$\n",
    "\n",
    "上面的步骤中，由于缩并时角标 $\\nu, \\kappa$ 的存在，因此这两个变量必须是连续的。可以通过内存操作被分割的角标是 $i, j$。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "上一步中，分割的大小被限制到 4 个原子轨道，是比较激进的限制方式；但这里，程序倾向于使用宽松的限制方式，因此 $i, j$ 的分割越小越好。这是因为上一步的分割限制涉及到因对称性所导致的浪费；如果分割太大，那么浪费也会越多。但这一步中，分割太小会导致大量硬盘 I/O，从而限制效率。\n",
    "\n",
    "在 PySCF 的实现中，假设内存有至少 $8 n_\\mathrm{occ} n_\\mathrm{AO}^2$。这要求了其中一个角标 $j$ 在内存中是连续储存的。不连续的部分就是 $i$，且最小的分割是  $n_\\mathrm{occblk} = 4$。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在演示的实例中，我们将分割的大小稍微缩小到 3，这是因为氨分子一共就 5 个占据轨道。因此，对于氨分子，两次分割分别是"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(0, 3, 3), (3, 5, 2)]"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "occblk = 3\n",
    "list_occ_slice = []\n",
    "for i in range(0, nocc, occblk):\n",
    "    num_blk = min(occblk, nocc-i)\n",
    "    list_occ_slice.append((i, i + num_blk, num_blk))\n",
    "list_occ_slice"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在实现积分转换之前，我们要能先把 $(i_s \\nu | \\kappa j)$ 储存到内存。为此，我们声明一块内存，用于储存这部分转换到一半的积分：(维度定为 $(i_s, j, \\nu, \\kappa)$)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "blk_eri = np.empty((occblk, nocc, nao, nao))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们需要写一个函数，以加载这部分积分。这一步是 I/O 关键步骤，因为尽管我们要求的内存量是 $n_\\mathrm{occblk} n_\\mathrm{occ} n_\\mathrm{AO}^2$，但要求单次分割的硬盘 I/O 量是 $n_\\mathrm{occ}^2 n_\\mathrm{AO}^2$，因此总共要求的硬盘 I/O 量是 $n_\\mathrm{occ}^3 n_\\mathrm{AO}^2 / n_\\mathrm{occblk}$。如果 $n_\\mathrm{occblk}$ 设置得太小，那么在 I/O 上就变相地变成五次方时间消耗，非常划不来。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load(i0, i1, eri):\n",
    "    for idx, (v0, v1, k0, k1) in enumerate(list_ao_slice):\n",
    "        eri[:i1-i0, :, v0:v1, k0:k1] = ftmp[str(idx)][i0:i1]\n",
    "        if v0 != k0:\n",
    "            dat = ftmp[str(idx)][:, i0:i1]\n",
    "            eri[:i1-i0, :, k0:k1, v0:v1] = dat.transpose((1, 0, 3, 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们验证一下这个程序的正确性。如果现在取第一个分割，即对 $i$ 取 0-2 号轨道，那么 `i0 = 0, i1 = 3`。我们需要事先准备好转换到一半的 `eri` $(i_s \\nu | \\kappa j)$ 的空间 (维度 $(i_s, j, \\nu, \\kappa)$)。这个函数会原地对 `eri` 内部的值直接改动，不返回结果。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eri = np.ndarray((3, nocc, nao, nao), buffer=blk_eri)\n",
    "load(0, 3, eri)\n",
    "np.allclose(np.einsum(\"uvkl, ui, kj -> ijvl\", mol.intor(\"int2e\"), Co, Co)[0:3], eri)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "随后，我们就可以考虑积分的过程了。尽管 $(i_s \\nu | \\kappa j)$ 的维度是 $(i_s, j, \\nu, \\kappa)$，但最终生成的 $(ia|jb)$ 的维度是 $(i, a, j, b)$。其过程如下。\n",
    "\n",
    "$$\n",
    "(i_s a|jb) = (i_s \\nu | \\kappa j) C_{\\nu a} C_{\\kappa b}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "feri = lib.H5TmpFile()\n",
    "h5_oovv = feri.create_dataset(\"ovov\", (nocc, nvir, nocc, nvir), \"f8\")\n",
    "for i0, i1, inum in list_occ_slice:\n",
    "    # Not using `blk_eri` directly, since i1-i0 = inum could be different in iterations\n",
    "    eri = np.ndarray((inum, nocc, nao, nao), buffer=blk_eri)\n",
    "    # load (i_s nu | kappa j) to `eri`\n",
    "    load(i0, i1, eri)\n",
    "    # calculate (i_s a | j b) tensor contraction\n",
    "    eri_mo = np.einsum(\"ijvl, va, lb -> iajb\", eri, Cv, Cv)\n",
    "    # dump to h5py file (hard-disk)\n",
    "    h5_oovv[i0:i1] = eri_mo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "最后，我们验证一下计算得到的 $(ia|jb)$ `h5_oovv` 是否确实是之前我们给出过的 `eri_iajb`："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(eri0_iajb, h5_oovv)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 浮点计算关键步骤：`_ao2mo.nr_e2` 函数"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "刚才一直强调硬盘 I/O 的消耗，但整个程序最耗时的步骤是浮点计算的张量缩并：\n",
    "\n",
    "```python\n",
    "# calculate (i_s a | j b) tensor contraction\n",
    "eri_mo = np.einsum(\"ijvl, va, lb -> iajb\", eri, Cv, Cv)\n",
    "```\n",
    "\n",
    "在 PySCF 中，它并非使用 Python 程序进行处理，而是用 C 程序进行计算。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "相同的工作可以使用 `np.einsum` 或 `lib.einsum` 完成，但对于这个特定的任务，`_ao2mo.nr_e2` 的处理效率非常高。文档的最后就讨论这个函数。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们现在假设遇到的问题是 $n_\\mathrm{occ} = 20$, $n_\\mathrm{vir}=100$, $n_\\mathrm{MO} = n_\\mathrm{AO} = 120$。不考虑对称性，我们随机一个 `eri` 与 `Cv` 矩阵，作为当前模型问题的矩阵。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "nao = nmo = 120\n",
    "nocc, nvir = 20, 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "eri = np.random.randn(nocc, nocc, nao, nao)\n",
    "Cv = np.random.randn(nao, nvir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们首先可以验证 `_ao2mo.nr_e2` 函数与 `np.einsum` 的作用效果相同。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.allclose(\n",
    "    _ao2mo.nr_e2(eri.reshape(nocc**2, nao**2), Cv, (0, nvir, 0, nvir), \"s1\", \"s1\").reshape(nocc, nocc, nvir, nvir),\n",
    "    np.einsum(\"ijvl, va, lb -> ijab\", eri, Cv, Cv)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "但从耗时上，`_ao2mo.nr_e2` 的消耗远远比 `np.einsum` 低："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22.4 ms ± 4 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 10\n",
    "_ao2mo.nr_e2(eri.reshape(nocc**2, nao**2), Cv, (0, nvir, 0, nvir), \"s1\", \"s1\").reshape(nocc, nocc, nvir, nvir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "47.1 ms ± 1.46 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 10\n",
    "np.einsum(\"ijvl, va, lb -> ijab\", eri, Cv, Cv)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在对程序没有更深了解的情况下，只能说 `_ao2mo.nr_e2` 是一个优秀的特化程序；可以的情况下就尽量使用。"
   ]
  }
 ],
 "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.3"
  },
  "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": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}