{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 04 The Closed-Shell MP2 energy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unlike the Hartree-Fock energy, correlation energies like the MP2 (second-order Moller-Plesset perturbation) energy are usually expressed in terms of MO-basis quantities (integrals, MO energies). The most expensive part of the calculation is the transformation of the two-electron integrals from the AO to the MO basis representation. The purpose of this project is to understand this transformation and fundamental techniques for its efficient implementation. The theoretical background and a concise set of instructions for this project may be found [here](https://github.com/CrawfordGroup/ProgrammingProjects/blob/master/Project%2304/project4-instructions.pdf)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": [
     "thebe-init",
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# Following os.chdir code is only for thebe (live code), since only in thebe default directory is /home/jovyan\n",
    "import os\n",
    "if os.getcwd().split(\"/\")[-1] != \"Project_04\":\n",
    "    os.chdir(\"source/Project_04\")\n",
    "from solution_04 import Molecule as SolMol\n",
    "\n",
    "from pyscf import gto\n",
    "import numpy as np\n",
    "import scipy.linalg\n",
    "from scipy.linalg import fractional_matrix_power\n",
    "from typing import Tuple\n",
    "np.set_printoptions(precision=7, linewidth=120, suppress=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": [
     "thebe-init",
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# Solution mol only uses PySCF approach\n",
    "sol_mole = SolMol()\n",
    "sol_mole.construct_from_dat_file(\"input/h2o/STO-3G/geom.dat\")\n",
    "sol_mole.obtain_mol_instance(basis=\"STO-3G\")\n",
    "sol_mole.obtain_nao()\n",
    "sol_mole.obtain_nocc()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Molecule Object Initialization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this project, we may use the updated Molecule initialization. Since these are technical details, we toggle the code and explanation below. Most of the method functions are illustrated in [Project 01](../Project_01/Project_01.ipynb) and [Project 03](../Project_03/Project_03.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": [
     "thebe-init",
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "class Molecule:\n",
    "\n",
    "    def __init__(self):\n",
    "        # Project 03 existed\n",
    "        self.atom_charges = NotImplemented  # type: np.ndarray\n",
    "        self.atom_coords = NotImplemented  # type: np.ndarray\n",
    "        self.natm = NotImplemented  # type: int\n",
    "        self.mol = NotImplemented  # type: gto.Mole\n",
    "        self.nao = NotImplemented  # type: int\n",
    "        self.charge = 0  # type: int\n",
    "        self.nocc = NotImplemented  # type: int\n",
    "        # Project 04 added\n",
    "        self.mo_coeff = NotImplemented  # type: np.ndarray\n",
    "        self.mo_energy = NotImplemented  # type: np.ndarray\n",
    "        self.eri_ao = NotImplemented  # type: np.ndarray\n",
    "        self.eri_mo = NotImplemented  # type: np.ndarray\n",
    "        self.energy_rhf = NotImplemented  # type: np.ndarray\n",
    "        self.energy_corr = NotImplemented  # type: np.ndarray\n",
    "\n",
    "    def construct_from_dat_file(self, file_path: str):\n",
    "        with open(file_path, \"r\") as f:\n",
    "            dat = np.array([line.split() for line in f.readlines()][1:])\n",
    "            self.atom_charges = np.array(dat[:, 0], dtype=float).astype(int)\n",
    "            self.atom_coords = np.array(dat[:, 1:4], dtype=float)\n",
    "            self.natm = self.atom_charges.shape[0]\n",
    "\n",
    "    def obtain_mol_instance(self, basis: str, verbose=0):\n",
    "        mol = gto.Mole()\n",
    "        mol.unit = \"Bohr\"\n",
    "        mol.atom = \"\\n\".join([(\"{:3d} \" + \" \".join([\"{:25.18f}\"] * 3)).format(chg, *coord) for chg, coord in zip(self.atom_charges, self.atom_coords)])\n",
    "        mol.basis = basis\n",
    "        mol.charge = self.charge\n",
    "        mol.spin = 0\n",
    "        mol.verbose = verbose\n",
    "        self.mol = mol.build()\n",
    "\n",
    "    def eng_nuclear_repulsion(self) -> float:\n",
    "        return self.mol.energy_nuc()\n",
    "\n",
    "    def obtain_nao(self):\n",
    "        self.nao = self.mol.nao_nr()\n",
    "\n",
    "    def obtain_nocc(self):\n",
    "        assert (self.atom_charges.sum() - self.charge) % 2 == 0\n",
    "        self.nocc = (self.atom_charges.sum() - self.charge) // 2\n",
    "\n",
    "    def integral_ovlp(self) -> np.ndarray:\n",
    "        return self.mol.intor(\"int1e_ovlp\")\n",
    "\n",
    "    def integral_kin(self) -> np.ndarray:\n",
    "        return self.mol.intor(\"int1e_kin\")\n",
    "\n",
    "    def integral_nuc(self) -> np.ndarray:\n",
    "        return self.mol.intor(\"int1e_nuc\")\n",
    "\n",
    "    def get_hcore(self) -> np.ndarray:\n",
    "        return self.integral_kin() + self.integral_nuc()\n",
    "\n",
    "    def integral_eri(self) -> np.ndarray:\n",
    "        return self.mol.intor(\"int2e\")\n",
    "\n",
    "    def integral_ovlp_m1d2(self) -> np.ndarray:\n",
    "        return fractional_matrix_power(self.integral_ovlp(), -1/2)\n",
    "\n",
    "    def get_fock(self, dm: np.ndarray) -> np.ndarray:\n",
    "        return self.get_hcore() + (self.eri_ao * dm).sum(axis=(-1, -2)) - 0.5 * (self.eri_ao * dm[:, None, :]).sum(axis=(-1, -3))\n",
    "\n",
    "    def get_coeff_from_fock_diag(self, fock: np.ndarray) -> np.ndarray:\n",
    "        return scipy.linalg.eigh(fock, self.integral_ovlp())[1]\n",
    "\n",
    "    def make_rdm1(self, coeff: np.ndarray) -> np.ndarray:\n",
    "        return 2 * coeff[:, :self.nocc] @ coeff[:, :self.nocc].T\n",
    "\n",
    "    def get_updated_dm(self, dm: np.ndarray) -> np.ndarray:\n",
    "        return self.make_rdm1(self.get_coeff_from_fock_diag(self.get_fock(dm)))\n",
    "\n",
    "    def eng_total(self, dm: np.ndarray) -> float:\n",
    "        return (0.5 * (self.get_hcore() + self.get_fock(dm)) * dm).sum() + self.eng_nuclear_repulsion()\n",
    "\n",
    "    def scf_process(self, dm_guess: np.ndarray=None) -> Tuple[float, np.ndarray]:\n",
    "        eng, dm = 0., np.zeros((self.nao, self.nao)) if dm_guess is None else np.copy(dm_guess)\n",
    "        max_iter, thresh_eng, thresh_dm = 64, 1e-10, 1e-8\n",
    "        for epoch in range(max_iter):\n",
    "            eng_next, dm_next = self.eng_total(dm), self.get_updated_dm(dm)\n",
    "            if np.abs(eng_next - eng) < thresh_eng and np.linalg.norm(dm_next - dm) < thresh_dm:\n",
    "                eng, dm = eng_next, dm_next\n",
    "                break\n",
    "            eng, dm = eng_next, dm_next\n",
    "        return eng, dm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{toggle}\n",
    "\n",
    "The meaning of additional attributes of `Molecule` are:\n",
    "- `mo_coeff` $C_{\\mu p}$ Molecular orbital coefficient\n",
    "- `mo_energy` $\\varepsilon_p$ Molecular orbital energy (from fock diagonalization)\n",
    "- `eri_ao` $(\\mu \\nu | \\kappa \\lambda)$ Atomic orbital electron repulsion integral\n",
    "- `eri_mo` $(p q | r s)$ Molecular orbital electron repulsion integral\n",
    "- `energy_rhf` $E_\\mathrm{tot}^\\mathsf{RHF}$ Total energy of RHF\n",
    "- `energy_corr` $E_\\mathrm{corr}$ Correlation energy, in this project it is MP2's form $E_\\mathrm{corr}^\\mathsf{MP2}$\n",
    "\n",
    "`mo_coeff` $C_{\\mu p}$, `mo_energy` $e_p$, `eri_ao` $(\\mu \\nu | \\kappa \\lambda)$, `energy_rhf` $E_\\mathrm{tot}^\\mathsf{RHF}$ is already calculated in SCF process, as illustrated in [Project 03](../Project_03/Project_03.ipynb). To store these quantities as attribute of `Molecule` is out of consideration of convenience and efficiency (one does not need to evaluate atomic eri every time when it is called).\n",
    "\n",
    "The other two quantities, `eri_mo` $(pq|rs)$ and `energy_corr` $E_\\mathrm{corr}$ is our main concern in this project.\n",
    "\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Two-Electron Integrals"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Mulliken-ordered integrals are defined as:\n",
    "\n",
    "$$\n",
    "(\\mu \\nu | \\kappa \\lambda) = \\int \\phi_\\mu (\\boldsymbol{r}_1) \\phi_\\nu (\\boldsymbol{r}_1) \\frac{1}{|\\boldsymbol{r}_1 - \\boldsymbol{r}_2|} \\phi_\\kappa (\\boldsymbol{r}_2) \\phi_\\lambda (\\boldsymbol{r}_2) \\, \\mathrm{d} \\boldsymbol{r}_1 \\, \\mathrm{d} \\boldsymbol{r}_2\n",
    "$$\n",
    "\n",
    "We have already read and used this integral in [Project 03](../Project_03/Project_03.ipynb#step-3-two-electron-integrals). However we still does not have a function storing this integral to attribute of `Molecule`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "def obtain_eri_ao(mole: Molecule):\n",
    "    # Attribute Modification: `eri_ao` atomic orbital electron repulsion integral\n",
    "    raise NotImplementedError(\"Exactly 1 line of code\")\n",
    "\n",
    "Molecule.obtain_eri_ao = obtain_eri_ao"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(7, 7, 7, 7)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sol_mole.obtain_eri_ao()\n",
    "sol_mole.eri_ao.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Obtain the SCF Intermediate Results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have already implemented an SCF process, but results are not stored in attributes of `Molecule` currently. We need to have a function to store RHF total energy $E_\\mathrm{tot}^\\mathsf{RHF}$, molecular orbital coefficients $C_{\\mu p}$ and molecular orbital energy $\\varepsilon_p$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{admonition} Hint 1: Process of obtaining SCF intermediates\n",
    ":class: dropdown\n",
    "\n",
    "- Obtain SCF energy and density matrix from SCF calculation (`scf_process`)\n",
    "- Obtain molecular orbital coefficients and energy from diagonalization of fock matrix (see how `get_coeff_from_fock_diag` works)\n",
    "\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "def obtain_scf_intermediates(mole: Molecule, dm_guess: np.ndarray=None):\n",
    "    # Attribute Modification: `energy_rhf` Total energy of RHF\n",
    "    # Attribute Modification: `mo_energy` Molecular orbital energies\n",
    "    # Attribute Modification: `mo_coeff` Molecular orbital coefficients\n",
    "    raise NotImplementedError(\"About 3~5 lines of code\")\n",
    "\n",
    "Molecule.obtain_scf_intermediates = obtain_scf_intermediates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(-74.94207992819237,\n",
       " array([-20.2628916,  -1.2096974,  -0.5479647,  -0.4365272,  -0.3875867,   0.4776187,   0.5881393]),\n",
       " array([[-0.9944346, -0.2391589,  0.       , -0.0936832, -0.       ,  0.1116399,  0.       ],\n",
       "        [-0.024097 ,  0.8857356, -0.       ,  0.4795859,  0.       , -0.6695791, -0.       ],\n",
       "        [ 0.       ,  0.       , -0.6072848, -0.       ,  0.       ,  0.       , -0.9192343],\n",
       "        [-0.0031615,  0.0858962,  0.       , -0.7474314, -0.       , -0.7384886, -0.       ],\n",
       "        [ 0.       , -0.       , -0.       ,  0.       , -1.       ,  0.       , -0.       ],\n",
       "        [ 0.0045937,  0.1440396, -0.4529977, -0.3294712, -0.       ,  0.7098495,  0.7324607],\n",
       "        [ 0.0045937,  0.1440396,  0.4529977, -0.3294712,  0.       ,  0.7098495, -0.7324607]]))"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sol_mole.obtain_scf_intermediates()\n",
    "sol_mole.energy_rhf, sol_mole.mo_energy, sol_mole.mo_coeff"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: MO ERI Transformation (Naïve Algorithm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Molecular orbital (MO) electron repulsion integral (ERI) could be defined as\n",
    "\n",
    "$$\n",
    "(pq|rs) = \\int \\phi_p (\\boldsymbol{r}_1) \\phi_q (\\boldsymbol{r}_1) \\frac{1}{|\\boldsymbol{r}_1 - \\boldsymbol{r}_2|} \\phi_r (\\boldsymbol{r}_2) \\phi_s (\\boldsymbol{r}_2) \\, \\mathrm{d} \\boldsymbol{r}_1 \\, \\mathrm{d} \\boldsymbol{r}_2\n",
    "$$\n",
    "\n",
    "just similar to definition of AO ERI $(\\mu \\nu | \\kappa \\lambda)$. Note that $r$ in subscript means any (occupied or virtual) molecular orbital, not related to electron coordinate vector $\\boldsymbol{r}_1$ or $\\boldsymbol{r}_2$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The most straightforward expression of the AO/MO integral transformation is"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "(pq|rs) = \\sum_\\mu \\sum_\\nu \\sum_\\kappa \\sum_\\lambda C_{\\mu p} C_{\\nu q} (\\mu \\nu | \\kappa \\lambda) C_{\\kappa r} C_{\\lambda s}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This approach is easy to implement (hence the word \"[naïve](https://stackoverflow.com/q/5700575/9647779)\" above), but is expensive (time-costly) due to its $O(N^8)$ computational order. Nevertheless, you should start with this algorithm to get your code working, and run timings for the test cases below to get an idea of its computational cost."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{admonition} Naïve transformation code\n",
    ":class: dropdown\n",
    "\n",
    "Here's a code block that demonstrates how to carry out the AO to MO two-electron integral transformation using a single $O(N^8)$ step. Note that the original AO- and transformed MO-basis integrals are stored in a 4-dimensional 8-fold-permutational-symmetric tensors.\n",
    "\n",
    "Indices `p`, `q`, `r`, `s` denotes molecular orbital $p$, $q$, $r$, $s$. Indices `u`, `v`, `k`, `l` denotes atomic orbital $\\mu$, $\\nu$, $\\kappa$, $\\lambda$. In this program, we take the convenience that `mole.nao` $n_\\mathrm{AO} = n_\\mathrm{MO}$, i.e. number of atomic orbitals is equal to molecular orbitals.\n",
    "\n",
    "```python\n",
    "for p in range(mole.nao):\n",
    "    for q in range(mole.nao):\n",
    "        for r in range(mole.nao):\n",
    "            for s in range(mole.nao):\n",
    "                for u in range(mole.nao):\n",
    "                    for v in range(mole.nao):\n",
    "                        for k in range(mole.nao):\n",
    "                            for l in range(mole.nao):\n",
    "                                eri_mo[p, q, r, s] += eri_ao[u, v, k, l] * coeff[u, p] * coeff[v, q] * coeff[k, r] * coeff[l, s]\n",
    "```\n",
    "\n",
    "Spectacular, isn't it :-)\n",
    "\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{admonition} Multiple nested for-loop in one line of code\n",
    ":class: dropdown\n",
    "\n",
    "There could be some trick to remove some of nests. One could create an numpy array which contains multiple loop indices (ordinary nested list is also okay). For example, `p`, `q`, `r`, `s` could be looped by following code:\n",
    "\n",
    "```python\n",
    "def repeat_loop(dim: int, nested: int) -> np.ndarray:\n",
    "    return np.array([np.tile(np.repeat(np.arange(dim), dim**(nested - 1 - i)), dim**i) for i in range(nested)]).T\n",
    "```\n",
    "```python\n",
    ">>> repeat_loop(mole.nao, 4)\n",
    "array([[0, 0, 0, 0],\n",
    "       [0, 0, 0, 1],\n",
    "       [0, 0, 0, 2],\n",
    "       ...,\n",
    "       [6, 6, 6, 4],\n",
    "       [6, 6, 6, 5],\n",
    "       [6, 6, 6, 6]])\n",
    "```\n",
    "\n",
    "Then run the following code:\n",
    "\n",
    "```python\n",
    "for p, q, r, s in repeat_loop(mole.nao, 4):\n",
    "    for u, v, k, l in repeat_loop(mole.nao, 4):\n",
    "        eri_mo[p, q, r, s] += eri_ao[u, v, k, l] * coeff[u, p] * coeff[v, q] * coeff[k, r] * coeff[l, s]\n",
    "```\n",
    "\n",
    "The result `eri_mo` should be the same. However, simplify code in this way is at the cost of memory ($O(N^4)$ memory consumption). So in most cases, it is unbearable to use the following memory-costly code, although even more concise:\n",
    "\n",
    "```python\n",
    "for p, q, r, s, u, v, k, l in repeat_loop(mole.nao, 8):\n",
    "    eri_mo[p, q, r, s] += eri_ao[u, v, k, l] * coeff[u, p] * coeff[v, q] * coeff[k, r] * coeff[l, s]\n",
    "```\n",
    "\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{admonition} Timing in Jupyter\n",
    ":class: dropdown\n",
    "\n",
    "In Jupyter (or IPython), one can use [magic command](https://ipython.readthedocs.io/en/stable/interactive/magics.html). There are two magic commands in timing:\n",
    "- [%%time](https://ipython.readthedocs.io/en/stable/interactive/magics.html#magic-time) Time of one execution.\n",
    "- [%%timeit](https://ipython.readthedocs.io/en/stable/interactive/magics.html#magic-timeit) Average time of multiple execution.\n",
    "\n",
    "In most cases, `%%time` is more useful and percise. However, `%%time` is only suitable when execution time is relatively quick (0~1000 ms). For this step, we may use `%time`.\n",
    "\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "def get_eri_mo_naive(mole: Molecule) -> np.ndarray:\n",
    "    # Naive algorithm\n",
    "    # Output: MO electron repulsion integral\n",
    "    raise NotImplementedError(\"About 5~15 lines of code\")\n",
    "\n",
    "Molecule.get_eri_mo_naive = get_eri_mo_naive"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 19.9 s, sys: 78.1 ms, total: 19.9 s\n",
      "Wall time: 20.9 s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[ 0.166246 ,  0.0218153, -0.       ,  0.0244196, -0.       ,  0.0027303,  0.       ],\n",
       "       [ 0.0218153,  0.0126143, -0.       , -0.0184892,  0.       , -0.0184186, -0.       ],\n",
       "       [-0.       , -0.       ,  0.0055846,  0.       , -0.       , -0.       , -0.0002852],\n",
       "       [ 0.0244196, -0.0184892,  0.       , -0.0089254, -0.       ,  0.0012042, -0.       ],\n",
       "       [-0.       ,  0.       , -0.       , -0.       ,  0.0046733, -0.       , -0.       ],\n",
       "       [ 0.0027303, -0.0184186, -0.       ,  0.0012042, -0.       ,  0.0180084,  0.       ],\n",
       "       [ 0.       , -0.       , -0.0002852, -0.       , -0.       ,  0.       ,  0.0037402]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "sol_mole.get_eri_mo_naive()[0, 3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This algorithm is not suitable if basis gets even larger."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: MO ERI Transformation (Smarter Algorithm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that none of the orbital coefficients $\\mathbf{C}$ in the above expression have any indices in common. Thus, the summation could be rearranged such that:\n",
    "\n",
    "$$\n",
    "(pq|rs) = \\sum_\\lambda C_{\\lambda s} \\left[ \\sum_\\kappa C_{\\kappa r} \\left[ \\sum_\\nu C_{\\nu q} \\left[ \\sum_\\mu C_{\\mu p} (\\mu \\nu | \\kappa \\lambda) \\right] \\right] \\right]\n",
    "$$\n",
    "\n",
    "This means that each summation within brackets could be carried out separately, starting from the innermost summation over $\\mu$, if we store the results at each step. This reduces the $O(N^8)$ algorithm above to four $O(N^5)$ steps.\n",
    "\n",
    "After you have the noddy algorithm working and timed, modify it to use this smarter algorithm and time the test cases again. Enjoy!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{admonition} Hint 1: Smarter algorithm step-by-step\n",
    ":class: dropdown\n",
    "\n",
    "Although this algorithm can be illustrated in one line of code, it is better to be splited into 4 lines of formulas.\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "(p \\nu | \\kappa \\lambda) &= \\sum_\\mu C_{\\mu p} (\\mu \\nu | \\kappa \\lambda) \\\\\n",
    "(pq | \\kappa \\lambda) &= \\sum_\\nu C_{\\nu q} (p \\nu | \\kappa \\lambda) \\\\\n",
    "(pq | r \\lambda) &= \\sum_\\kappa C_{\\kappa r} (pq | \\kappa \\lambda) \\\\\n",
    "(pq | rs) &= \\sum_\\lambda C_{\\lambda s} (pq | r \\lambda) \\\\\n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "To calculate $(pq|\\kappa \\lambda)$, for example, could be\n",
    "\n",
    "```python\n",
    "loop_indices = repeat_loop(mole.nao, 4)\n",
    "tmp_1 = np.zeros((nao, nao, nao, nao))\n",
    "for u, v, k, l in loop_indices:\n",
    "    for p in range(nao):\n",
    "        tmp_1[p, v, k, l] += eri_ao[u, v, k, l] * coeff[u, p]\n",
    "tmp_2 = np.zeros((nao, nao, nao, nao))\n",
    "for p, v, k, l in loop_indices:\n",
    "    for q in range(nao):\n",
    "        tmp_2[p, q, k, l] += tmp_1[p, v, k, l] * coeff[v, q]\n",
    "```\n",
    "\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "def get_eri_mo_smarter(mole: Molecule):\n",
    "    # Smarter algorithm\n",
    "    # Output: MO electron repulsion integral\n",
    "    raise NotImplementedError(\"About 5 + 3~5 * 4 lines of code\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 93.8 ms, sys: 0 ns, total: 93.8 ms\n",
      "Wall time: 103 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[ 0.166246 ,  0.0218153, -0.       ,  0.0244196, -0.       ,  0.0027303,  0.       ],\n",
       "       [ 0.0218153,  0.0126143, -0.       , -0.0184892,  0.       , -0.0184186, -0.       ],\n",
       "       [-0.       , -0.       ,  0.0055846,  0.       , -0.       ,  0.       , -0.0002852],\n",
       "       [ 0.0244196, -0.0184892,  0.       , -0.0089254, -0.       ,  0.0012042, -0.       ],\n",
       "       [-0.       ,  0.       , -0.       , -0.       ,  0.0046733, -0.       , -0.       ],\n",
       "       [ 0.0027303, -0.0184186,  0.       ,  0.0012042, -0.       ,  0.0180084,  0.       ],\n",
       "       [ 0.       , -0.       , -0.0002852, -0.       , -0.       ,  0.       ,  0.0037402]])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "sol_mole.get_eri_mo_smarter()[0, 3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "97 ms ± 801 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "sol_mole.get_eri_mo_smarter()[0, 3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: MO ERI Transformation (`numpy.einsum`)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`numpy.einsum` ([NumPy API](https://numpy.org/doc/stable/reference/generated/numpy.einsum.html)) provides an extremely concise way to do tensor contraction, with sufficient support of efficiency.\n",
    "\n",
    "We still use the following formula:\n",
    "\n",
    "$$\n",
    "(pq|rs) = \\sum_{\\mu \\nu \\kappa \\lambda} C_{\\mu p} C_{\\nu q} (\\mu \\nu | \\kappa \\lambda) C_{\\kappa r} C_{\\lambda s}\n",
    "$$\n",
    "\n",
    "and molecular orbital basis ERI could be calculated as\n",
    "\n",
    "```python\n",
    "eri_mo = np.einsum(\"up, vq, uvkl, kr, ls -> pqrs\", coeff, coeff, eri_ao, coeff, coeff, optimize=True)\n",
    "```\n",
    "\n",
    "where `optimize=True` tells `numpy.einsum` to use the optimized (smarter) algorithm. This optional parameter is critical \n",
    "\n",
    "Finally, store molecular orbital ERI into attribute `eri_mo` in `Molecule`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{admonition} Hint 1: Contraction path of numpy.einsum\n",
    ":class: dropdown\n",
    "\n",
    "To see how `numpy.einsum` do the contraction route, we can use function `numpy.einsum_path` ([NumPy API](https://numpy.org/doc/stable/reference/generated/numpy.einsum_path.html)). Learning this function is optional to this project.\n",
    "\n",
    "```\n",
    ">>> print(np.einsum_path(\"up, vq, uvkl, kr, ls -> pqrs\", coeff, coeff, eri_ao, coeff, coeff)[1])\n",
    "  Complete contraction:  up,vq,uvkl,kr,ls->pqrs\n",
    "         Naive scaling:  8\n",
    "     Optimized scaling:  5\n",
    "      Naive FLOP count:  2.882e+07\n",
    "  Optimized FLOP count:  1.345e+05\n",
    "   Theoretical speedup:  214.373\n",
    "  Largest intermediate:  2.401e+03 elements\n",
    "--------------------------------------------------------------------------\n",
    "scaling                  current                                remaining\n",
    "--------------------------------------------------------------------------\n",
    "   5               uvkl,up->klpv                      vq,kr,ls,klpv->pqrs\n",
    "   5               klpv,vq->klpq                         kr,ls,klpq->pqrs\n",
    "   5               klpq,kr->lpqr                            ls,lpqr->pqrs\n",
    "   5               lpqr,ls->pqrs                               pqrs->pqrs\n",
    "```\n",
    "\n",
    "From the output, we know that the naive algorithm complexity is $O(N^8)$, the optimized (smarter) algorithm is $O(N^5)$. The first contraction path, for example, could be expressed as\n",
    "\n",
    "$$\n",
    "\\mathtt{Tensor1}_{\\kappa \\lambda p \\nu} = \\sum_{\\mu} (\\mu \\nu | \\kappa \\lambda) C_{\\mu p}\n",
    "$$\n",
    "\n",
    "The last contraction path could be\n",
    "\n",
    "$$\n",
    "(pq|rs) = \\sum_{\\lambda} \\mathtt{Tensor3}_{\\lambda pqr} C_{\\lambda s}\n",
    "$$\n",
    "\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "def get_eri_mo_einsum(mole: Molecule):\n",
    "    # Use numpy.einsum to automatically find optimal contraction path\n",
    "    # Output: MO electron repulsion integral\n",
    "    raise NotImplementedError(\"About 1~3 lines of code\")\n",
    "\n",
    "Molecule.get_eri_mo_einsum = get_eri_mo_einsum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "def obtain_eri_mo(mole: Molecule):\n",
    "    # Attribute Modification: `eri_mo` Molecular orbital electron repulsion integral\n",
    "    raise NotImplementedError(\"Exactly 1 line of code\")\n",
    "\n",
    "Molecule.obtain_eri_mo = obtain_eri_mo"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 0 ns, sys: 0 ns, total: 0 ns\n",
      "Wall time: 1.49 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[ 0.166246 ,  0.0218153, -0.       ,  0.0244196, -0.       ,  0.0027303,  0.       ],\n",
       "       [ 0.0218153,  0.0126143, -0.       , -0.0184892,  0.       , -0.0184186, -0.       ],\n",
       "       [-0.       , -0.       ,  0.0055846,  0.       , -0.       ,  0.       , -0.0002852],\n",
       "       [ 0.0244196, -0.0184892,  0.       , -0.0089254, -0.       ,  0.0012042, -0.       ],\n",
       "       [-0.       ,  0.       , -0.       , -0.       ,  0.0046733, -0.       , -0.       ],\n",
       "       [ 0.0027303, -0.0184186,  0.       ,  0.0012042, -0.       ,  0.0180084,  0.       ],\n",
       "       [ 0.       , -0.       , -0.0002852, -0.       , -0.       ,  0.       ,  0.0037402]])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "sol_mole.get_eri_mo_einsum()[0, 3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "416 µs ± 2.87 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "sol_mole.get_eri_mo_einsum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6: Compute the MP2 Energy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Restricted MP2 correlation energy is (Helgaker, et al., eq 14.4.56)\n",
    "\n",
    "$$\n",
    "E_\\mathrm{corr}^\\mathsf{MP2} = \\sum_{iajb} \\frac{(ia|jb) \\big[ 2 (ia|jb) - (ib|ja) \\big]}{\\varepsilon_i + \\varepsilon_j - \\varepsilon_a - \\varepsilon_b}\n",
    "$$\n",
    "\n",
    "where $i$ and $j$ denote doubly-occupied orbitals and $a$ and $b$ unoccupied orbitals, and the denominator involves the MO energies.\n",
    "\n",
    "Finally, store MP2 correlation energy into attribute `energy_corr` in `Molecule`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{admonition} Hint 1: Transpose of tensor\n",
    ":class: dropdown\n",
    "\n",
    "We know that for matrix, we can use `numpy.transpose`, or simply `numpy.ndarray.T` to transpose the matrix. However, tensor have more than 2 dimensions, thus transpose of tensor needs more parameters. For example, if we want to transpose `tsr` $T_{iajb}$ to $T_{ibja}$, we can either use `numpy.transpose` ([NumPy API](https://numpy.org/doc/stable/reference/generated/numpy.transpose.html)) or `numpy.swapaxes` ([NumPy](https://numpy.org/doc/stable/reference/generated/numpy.swapaxes.html)).\n",
    "\n",
    "```python\n",
    ">>> tsr.transpose(0, 3, 2, 1)  # `numpy.transpose` approach\n",
    ">>> tsr.swapaxes(1, 3)         # `numpy.swapaxes`  approach\n",
    "```\n",
    "\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{admonition} Hint 2: Denominator\n",
    ":class: dropdown\n",
    "\n",
    "You can make another 4-dimension tensor $D_{ij}^{ab}$, which is stored in $(i, a, j, b)$:\n",
    "\n",
    "$$\n",
    "D_{ij}^{ab} = \\varepsilon_i - \\varepsilon_a + \\varepsilon_j - \\varepsilon_b\n",
    "$$\n",
    "\n",
    "Using NumPy boardcasting ([NumPy doc](https://numpy.org/doc/stable/user/theory.broadcasting.html#array-broadcasting-in-numpy)), this can be done by\n",
    "\n",
    "```python\n",
    "D_iajb = e[:nocc, None, None, None] - e[None, nocc:, None, None] + e[None, None, :nocc, None] - e[None, None, None, nocc:]\n",
    "```\n",
    "\n",
    "Utilizing this tensor, one can calculate restricted MP2 correlation energy with no for-loops.\n",
    "\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "def eng_mp2_corr(mole: Molecule):\n",
    "    # Output: (Restricted) MP2 correlation energy\n",
    "    raise NotImplementedError(\"About 4~10 lines of code\")\n",
    "\n",
    "Molecule.eng_mp2_corr = eng_mp2_corr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [],
   "source": [
    "def obtain_mp2_corr(mole: Molecule):\n",
    "    # Attribute Modification: `energy_corr` Post-HF (in this project is RMP2) correlation energy\n",
    "    raise NotImplementedError(\"Exactly 1 line of code\")\n",
    "\n",
    "Molecule.obtain_mp2_corr = obtain_mp2_corr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SCF total       energy:     -74.94207993\n",
      "MP2 correlation energy:      -0.04914964\n",
      "MP2 total       energy:     -74.99122956\n"
     ]
    }
   ],
   "source": [
    "sol_mole.obtain_eri_mo()\n",
    "sol_mole.obtain_mp2_corr()\n",
    "print(\"SCF total       energy: {:16.8f}\".format(sol_mole.energy_rhf))\n",
    "print(\"MP2 correlation energy: {:16.8f}\".format(sol_mole.energy_corr))\n",
    "print(\"MP2 total       energy: {:16.8f}\".format(sol_mole.energy_rhf + sol_mole.energy_corr))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Cases"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The input structures, integrals, etc. for these examples are found in the [input directory](https://github.com/ajz34/PyCrawfordProgProj/tree/master/source/Project_04/). You can also use PySCF approach and simply ignore the integral files."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Water/STO-3G** ([Directory](https://github.com/ajz34/PyCrawfordProgProj/tree/master/source/Project_04/input/h2o/STO-3G))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SCF total       energy:     -74.94207993\n",
      "MP2 correlation energy:      -0.04914964\n",
      "MP2 total       energy:     -74.99122956\n"
     ]
    }
   ],
   "source": [
    "sol_mole = SolMol()\n",
    "sol_mole.construct_from_dat_file(\"input/h2o/STO-3G/geom.dat\")\n",
    "sol_mole.obtain_mol_instance(basis=\"STO-3G\")\n",
    "sol_mole.print_solution_04()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Water/DZ** ([Directory](https://github.com/ajz34/PyCrawfordProgProj/tree/master/source/Project_04/input/h2o/DZ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SCF total       energy:     -75.97787898\n",
      "MP2 correlation energy:      -0.15270988\n",
      "MP2 total       energy:     -76.13058885\n"
     ]
    }
   ],
   "source": [
    "sol_mole = SolMol()\n",
    "sol_mole.construct_from_dat_file(\"input/h2o/DZ/geom.dat\")\n",
    "sol_mole.obtain_mol_instance(basis=\"DZ\")\n",
    "sol_mole.print_solution_04()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Water/DZP** ([Directory](https://github.com/ajz34/PyCrawfordProgProj/tree/master/source/Project_04/input/h2o/DZP)) For this molecule, some additional code is included. See [Project 03](../Project_03/Project_03.ipynb#Test-Cases) for details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SCF total       energy:     -76.00882179\n",
      "MP2 correlation energy:      -0.22251923\n",
      "MP2 total       energy:     -76.23134103\n"
     ]
    }
   ],
   "source": [
    "sol_mole = SolMol()\n",
    "sol_mole.construct_from_dat_file(\"input/h2o/DZP/geom.dat\")\n",
    "sol_mole.obtain_mol_instance(basis=\"DZP_Dunning\")\n",
    "sol_mole.mol._basis[\"H\"][2][1][0] = 0.75\n",
    "sol_mole.mol.cart = True\n",
    "sol_mole.mol.build()\n",
    "sol_mole.print_solution_04()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Methane/STO-3G** ([Directory](https://github.com/ajz34/PyCrawfordProgProj/tree/master/source/Project_04/input/ch4/STO-3G))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "tags": [
     "hide-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SCF total       energy:     -39.72685032\n",
      "MP2 correlation energy:      -0.05604667\n",
      "MP2 total       energy:     -39.78289699\n"
     ]
    }
   ],
   "source": [
    "sol_mole = SolMol()\n",
    "sol_mole.construct_from_dat_file(\"input/ch4/STO-3G/geom.dat\")\n",
    "sol_mole.obtain_mol_instance(basis=\"STO-3G\")\n",
    "sol_mole.print_solution_04()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Szabo, A.; Ostlund, N. S. *Modern Quantum Chemistry: Introduction to Advanced Electronic Structure Theory* Dover Publication Inc., 1996.\n",
    "\n",
    "    ISBN-13: 978-0486691862\n",
    "\n",
    "- Helgaker, T.; Jørgensen, P.; Olsen, J. *Molecular Electronic-Structure Theory* John Wiley & Sons, ltd., 2000\n",
    "\n",
    "    ISBN-13: 978-1118531471"
   ]
  }
 ],
 "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": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}