{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 02 - Performing unitary transformations\n",
    "\n",
    ":::{important}\n",
    "This is a JupyterLab notebook. We assume you execute all the cells in this notebook in sequential order. This is relevant as some cells depend on the output as generated in a previous cell. This notebook assumes that the following packages have been installed:\n",
    "\n",
    "* pyqint (`conda install -c ifilot pyqint`)\n",
    "* pytessel (`conda install -c ifilot pytessel`)\n",
    "* jupyterlab-myst (`pip install jupyterlab-myst`)\n",
    ":::\n",
    "\n",
    "## Introduction\n",
    "\n",
    "In the previous notebook we observed how a Hartree-Fock calculation can be performed and how to read out the eigenvalues and -vectors corresponding to the molecular orbital energies and coefficients. In this notebook, we will expand on this knowledge and show that unitary transformation on the set of occupied molecular orbitals will not change the overall electronic energy. Furthermore, by systematic variation of the occupied molecular orbitals, i.e. by mixing them among each other, we can find so-called localized molecular orbitals which have unique properties very much in line to our chemical intuition.\n",
    "\n",
    "## Learning goals\n",
    "\n",
    "In this notebook, you will learn\n",
    "\n",
    "* How the total electronic energy and the electronic energy of a molecular orbital is defined.\n",
    "* How to apply unitary transformations to the occupied molecular orbitals and why these will not change the overall electronic energy.\n",
    "* How to optimize the molecular orbitals to maximize a metric.\n",
    "* How to perform the Foster-Boys localization strategy.\n",
    "\n",
    "## Calculating the total electronic energy and the energies per orbitals\n",
    "\n",
    "In the previous notebook, the molecular orbital energies and the total electronic energy were simply handed to us (via the results library). Suppose we do **not** have access to this result library but we have at least access to the coefficient matrix and the Hamiltonian (Fock) matrix, how can be retrieve the same information?\n",
    "\n",
    "To find the molecular orbital energies, we can utilize the following equation:\n",
    "\n",
    "$$\n",
    "\\mathbf{E} = \\mathbf{C}^{\\intercal} \\mathbf{H} \\mathbf{C}\n",
    "$$\n",
    "\n",
    "which will yield the molecular orbital energies as the values on the diagonal of the matrix $\\mathbf{E}$.\n",
    "\n",
    "In the script below, we will explicitly demonstrate this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pyqint import Molecule,HF\n",
    "import numpy as np\n",
    "\n",
    "mol = Molecule('CO')\n",
    "mol.add_atom('C', 0, 0, -1.08232106)\n",
    "mol.add_atom('O', 0, 0, 1.08232106)\n",
    "res = HF().rhf(mol, 'sto3g')\n",
    "\n",
    "# explicitly calculate the orbital energies using the expression and show that\n",
    "# these are consistent with the \"orbe\" array\n",
    "orbe = np.diag(res['orbc'].transpose() @ res['fock'] @ res['orbc'])\n",
    "print('id       before        after')\n",
    "print('----------------------------')\n",
    "for i in range(0,len(res['orbe'])):\n",
    "    print('%2i %12.4f %12.4f' % (i+1, orbe[i], res['orbe'][i]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In quite a similar fashion, we can reproduce the total electronic energy. First, we calculate the density matrix $\\mathbf{P}$ via\n",
    "\n",
    "$$\n",
    "\\mathbf{P} = 2 \\sum_{k=1} \\sum_{k=1}^{N/2} C_{ik}C_{jk}\n",
    "$$\n",
    "\n",
    "Next, we calculate the inner product between every row of the density matrix and every column of sum of the Fock, kinetic and nuclear attraction matrices and add the repulsion of the nuclei to obtain the total electronic energy, as given by\n",
    "\n",
    "$$\n",
    "E = \\frac{1}{2} \\sum_{ij} P_{ji} \\left(T_{ij} + V_{ij} + H_{ij} \\right) + E_{\\textrm{nuc-nuc-rep}}\n",
    "$$\n",
    "\n",
    "The procedure is demonstrated in the script below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# grab the coefficient matrix\n",
    "C = res['orbc']\n",
    "\n",
    "# calculate occupancy vector\n",
    "occ = (np.arange(len(res['orbe'])) < res['nelec']//2) * 2 # list of number of electrons per atomic orbital\n",
    "\n",
    "# construct density matrix P\n",
    "# this might look like black voodoo magic, but einsum might actually the best invention since sliced bread,\n",
    "# have a look here: https://numpy.org/doc/stable/reference/generated/numpy.einsum.html\n",
    "P = np.einsum('ik,jk,k->ij', C, C, occ)\n",
    "\n",
    "# calculate the total electronic energy\n",
    "E = 0.5 * np.einsum('ij,ji', P, res['fock'] + res['kinetic'] + res['nuclear']) + res['enucrep']\n",
    "\n",
    "print(res['energy'])\n",
    "print(E)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 01\n",
    "\n",
    "Calculate the total electronic energy for the CH4 molecule and show that the procedure is consistent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pyqint import MoleculeBuilder,HF\n",
    "mol = MoleculeBuilder().from_name('ch4') # load molecule directly from the molecule class\n",
    "mol.name = 'CH4'                         # give the molecule an appropriate name\n",
    "res = HF().rhf(mol, 'sto3g')\n",
    "\n",
    "#!BEGIN ANSWER\n",
    "# grab the coefficient matrix\n",
    "C = res['orbc']\n",
    "\n",
    "# calculate occupancy vector\n",
    "occ = (np.arange(len(res['orbe'])) < res['nelec']//2) * 2 # list of number of electrons per atomic orbital\n",
    "\n",
    "# construct density matrix P\n",
    "P = np.einsum('ik,jk,k->ij', C, C, occ)\n",
    "\n",
    "# calculate the total electronic energy\n",
    "E = 0.5 * np.einsum('ij,ji', P, res['fock'] + res['kinetic'] + res['nuclear']) + res['enucrep']\n",
    "\n",
    "print(res['energy'])\n",
    "print(E)\n",
    "#END ANSWER"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unitary transformations\n",
    "\n",
    "Hopefully you are fully convinced that we can reproduce orbital energies and total electronic energies purely on the basis of the kinetic, nuclear, fock and coefficient matrices, the total number of electrons and the positions of the nuclei (or simply the corresponding energy term resulting from the repulsion between the nuclei).\n",
    "\n",
    ":::{note}\n",
    "You might never have fully appreciated this, but the kinetic, nuclear and Fock matrices capture the result of an operation in Hilbert space. That Hilbert space is spanned by the basis used in the calculation. If the basis is not going to change, that representation remains **valid**.\n",
    ":::\n",
    "\n",
    "We are now going to mix the **occupied** molecular orbital among each other. Since the basis set representing the molecular orbitals remains the same, the only matrix that is going to change is the coefficient matrix. We will explore what these transformations will do to the orbital energies and to the total electronic energies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import scipy\n",
    "from scipy.stats import ortho_group\n",
    "\n",
    "# recalculate CO molecule, because this will be our go-to toy system (I am deeply passionate about CO as a molecular system, please ask me over a beer)\n",
    "mol = Molecule('CO')\n",
    "mol.add_atom('C', 0, 0, -1.08232106)\n",
    "mol.add_atom('O', 0, 0, 1.08232106)\n",
    "res = HF().rhf(mol, 'sto3g')\n",
    "\n",
    "# determine number of occupied molecular orbitals\n",
    "N = res['nelec'] // 2\n",
    "\n",
    "# list of number of electrons per atomic orbital\n",
    "occ = (np.arange(len(res['orbe'])) < res['nelec']//2) * 2\n",
    "\n",
    "# produce 1000 different coefficient matrices, all connected via unitary transformations\n",
    "nr_iterations = 1000\n",
    "E_arr = np.zeros(nr_iterations)\n",
    "orbe_arr = np.zeros((nr_iterations, len(res['orbe'])))\n",
    "for i in range(0,nr_iterations):\n",
    "    \n",
    "    # randomly generate orthogonal transformation matrix\n",
    "    orthomat = ortho_group.rvs(N)\n",
    "    identitymat = np.identity(len(res['orbc']) - N)\n",
    "    T = scipy.linalg.block_diag(orthomat, identitymat)\n",
    "    \n",
    "    # perform unitary transformation\n",
    "    C = res['orbc'] @ T\n",
    "    \n",
    "    # build density matrix\n",
    "    P = np.einsum('ik,jk,k->ij', C, C, occ)\n",
    "    \n",
    "    # calculate total electronic energy\n",
    "    E = 0.5 * np.einsum('ij,ji', P, res['fock'] + res['kinetic'] + res['nuclear']) + res['enucrep']\n",
    "    \n",
    "    # store into array\n",
    "    E_arr[i] = E\n",
    "    \n",
    "    # calculate orbital energies (only occupied orbitals)\n",
    "    orbe = np.diag(C.transpose() @ res['fock'] @ C)\n",
    "        \n",
    "    # sort the orbital energies from low to high\n",
    "    orbe = np.sort(orbe)\n",
    "    \n",
    "    # and store them in the array\n",
    "    orbe_arr[i,:] = orbe\n",
    "\n",
    "# plot total electronic energies as function of number of random iterations\n",
    "figure, ax = plt.subplots(1,2, figsize=(12,4))\n",
    "ax[0].plot(E_arr, 'o', alpha=0.5)\n",
    "ax[0].set_ylim(res['energy']-10, res['energy']+10)\n",
    "\n",
    "# plot orbital energies as function of random iterations\n",
    "for i in range(0,N):\n",
    "    ax[1].plot(orbe_arr[:,i], 'o', alpha=0.5)\n",
    "\n",
    "# plot the original orbital energies into the figure\n",
    "ax[1].hlines(res['orbe'][:N], 0, nr_iterations, linestyle='--', color='black', alpha=0.5)\n",
    "    \n",
    "# give the figures some axis labels\n",
    "ax[0].set_xlabel('Iteration number [-]')\n",
    "ax[0].set_ylabel('Total electronic energy [eV]')\n",
    "ax[1].set_xlabel('Iteration number [-]')\n",
    "ax[1].set_ylabel('Molecular orbital energy')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us digest the result as shown above. The figure on the left hand side tells us that unitary transformations among the occupied molecular orbitals will not change the total electronic energy of the system. In other words, the total electronic energy is an **invariant** under a unitary transformation. In contrast, the orbital energies can change under an arbitrary unitary transformation and are surely not invariant.\n",
    "\n",
    "Rather than performing a *random* unitary transformation, let us seek among the infinite set of unitary transformations for a very specific solution. The solution we are aiming to find is that solution that maximizes the following metric\n",
    "\n",
    "$$\n",
    "\\mathcal{M} = \\sum_{i \\in \\textrm{occ}} \\left< \\phi_{i} | \\vec{r} | \\phi_{i} \\right>^{2}.\n",
    "$$\n",
    "\n",
    "This metric corresponds to the sum of squares of the dipole moment magnitudes. For the canonical solutions, i.e. the solution that diagonalizes the Fock matrix, you will find a very small value for this metric as all the orbitals are maximally delocalized or spread out. This spreading out of the orbitals ensures that the orbital centroids lie roughly in the center of the molecule, giving relatively small values for this metric. We are going to search for the situation wherein the orbitals are very **localized**. This localization implies that an orbital is essentially confined to some small corner of the molecule. With the exception of a few orbitals, the majority of the orbitals will have their centroid at a relatively large distance from the center of the molecule and consequently, the metric will be large.\n",
    "\n",
    "Let me show you the procedure for the CO molecule."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pyqint import PyQInt\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy\n",
    "\n",
    "# recalculate CO molecule\n",
    "mol = Molecule('CO')\n",
    "mol.add_atom('C', 0, 0, -1.08232106)\n",
    "mol.add_atom('O', 0, 0, 1.08232106)\n",
    "res = HF().rhf(mol, 'sto3g')\n",
    "\n",
    "# list of number of electrons per atomic orbital\n",
    "occ = (np.arange(len(res['orbe'])) < res['nelec']//2) * 2\n",
    "\n",
    "# first calculate the dipole tensor, note that since the basis functions are not going to change, this\n",
    "# dipole tensor is **fixed** throughout the calculation, we only have to calculate it once\n",
    "Norb = len(res['orbe'])\n",
    "cgfs = res['cgfs']\n",
    "dipol = np.zeros((Norb,Norb,3))\n",
    "integrator = PyQInt()\n",
    "for i,cgf1 in enumerate(cgfs):\n",
    "    for j,cgf2 in enumerate(cgfs):\n",
    "        for k in range(0,3):\n",
    "            dipol[i,j,k] = integrator.dipole(cgf1, cgf2, k, 0.0)\n",
    "\n",
    "#\n",
    "# Here we have to define three functions to aid us in our epic quest:\n",
    "#\n",
    "# 1. calculate r2: Calculate the sum of the squared dipole norm given the dipole moment tensor and the coefficient matrix\n",
    "# 2. evaluate_orbital_centroids: Returns the negative of the sum of squared dipole norms given a rotation alpha among two molecular orbitals\n",
    "# 3. perform_foster_boys_mixing: The routine that does the r2 optimization for us, internally it uses the optimize library of scipy\n",
    "#\n",
    "            \n",
    "def calculate_r2(C, dipol, occ):\n",
    "    \"\"\"\n",
    "    Calculate r2 from coefficient matrix\n",
    "    \"\"\"\n",
    "    dipolest = np.einsum('ji,ki,jkl->il', C, C, dipol)\n",
    "    r2 = np.einsum('ij,i->', dipolest**2, occ)\n",
    "\n",
    "    return r2\n",
    "\n",
    "def evaluate_orbital_centroids(alpha, C, i, j, dipol, occ):\n",
    "    \"\"\"\n",
    "    Evaluate the result of r2 after a 2x2 rotation\n",
    "    between solutions i and j\n",
    "    \"\"\"\n",
    "    Cnew = C.copy()\n",
    "    Cnew[:,i] = np.cos(alpha) * C[:,i] + np.sin(alpha) * C[:,j]\n",
    "    Cnew[:,j] = -np.sin(alpha) * C[:,i] + np.cos(alpha) * C[:,j]\n",
    "    r2 = calculate_r2(Cnew, dipol, occ)\n",
    "\n",
    "    return -r2\n",
    "\n",
    "def perform_foster_boys_mixing(C, dipol, occ):\n",
    "    \"\"\"\n",
    "    Successively perform n(n-1)/2 optimization among the occupied\n",
    "    molecular orbitals using the scipy minimize function\n",
    "    \"\"\"\n",
    "    r2start = calculate_r2(C, dipol, occ)\n",
    "    r2final = r2start\n",
    "\n",
    "    N = np.sum(occ)//2\n",
    "    for i in range(N):\n",
    "        for j in range(i+1, N):\n",
    "            \n",
    "            # optimize local pair of MOs on the interval -(pi,pi)\n",
    "            res = scipy.optimize.minimize(evaluate_orbital_centroids,\n",
    "                                          0.0,\n",
    "                                          args=(C, i, j, dipol, occ),\n",
    "                                          bounds=[(-np.pi, np.pi)],\n",
    "                                          tol=1e-12)\n",
    "\n",
    "            alpha = res.x[0]\n",
    "            Cnew = C.copy()\n",
    "            Cnew[:,i] = np.cos(alpha) * C[:,i] + np.sin(alpha) * C[:,j]\n",
    "            Cnew[:,j] = -np.sin(alpha) * C[:,i] + np.cos(alpha) * C[:,j]\n",
    "\n",
    "            r2 = calculate_r2(Cnew, dipol, occ)\n",
    "            if(r2 > r2start):\n",
    "                C = Cnew.copy()\n",
    "                r2final = r2start\n",
    "\n",
    "    return C, r2final\n",
    "\n",
    "def construct_random_orthogonal_transformation(C, nelec, nops=200):\n",
    "    \"\"\"\n",
    "    Construct unitary transformation matrix via a series of two-dimensional\n",
    "    unitary transformations among the occupied molecular orbitals\n",
    "    \"\"\"\n",
    "    for i in range(0,nops):\n",
    "        Cnew = C.copy()\n",
    "\n",
    "        # randomly pick two numbers among the occupied orbitals\n",
    "        n = np.random.choice(range(nelec//2), size=2, replace=False)\n",
    "\n",
    "        # and mix them by a random angle\n",
    "        gamma = np.random.uniform() * 2.0 * np.pi\n",
    "\n",
    "        for j in range(0,len(C)):\n",
    "            Cnew[j,n[0]] = np.cos(gamma) * C[j,n[0]] + np.sin(gamma) * C[j,n[1]]\n",
    "            Cnew[j,n[1]] = -np.sin(gamma) * C[j,n[0]] + np.cos(gamma) * C[j,n[1]]\n",
    "\n",
    "        C = Cnew.copy()\n",
    "\n",
    "    return Cnew\n",
    "\n",
    "old_r2 = 0.0     # start assuming orbital centroids lie at origin\n",
    "r2_arr = []      # store r2 values in an array\n",
    "diff = 1\n",
    "nriter = 0\n",
    "max_iter = 10000 # always set a maximum number of iterations in a while loop ;-)\n",
    "\n",
    "# start by generating a random orthogonal transformation (the algorithm depends on a proper initial 'shaking')\n",
    "C = construct_random_orthogonal_transformation(C, N)\n",
    "\n",
    "# keep on iterating until no changes in r2 are encountered\n",
    "while diff > 1e-10 and nriter < max_iter:\n",
    "    C, r2 = perform_foster_boys_mixing(C, dipol, occ)\n",
    "    r2_arr.append(r2)\n",
    "    diff = np.abs(r2 - old_r2)\n",
    "    old_r2 = r2\n",
    "    nriter += 1\n",
    "\n",
    "# plot the convergence\n",
    "plt.figure(figsize=(3,3))\n",
    "plt.plot(r2_arr, 'o-')\n",
    "plt.xlabel('Number of iterations')\n",
    "plt.ylabel(r'$\\sum_{i} \\left< \\psi_{i} |\\vec{r}| \\psi_{i} \\right> ^{2}$')\n",
    "plt.grid(linestyle='--', color='black', alpha=0.5)\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# print the MO energies after mixing\n",
    "orbe = np.diag(C.transpose() @ res['fock'] @ C)\n",
    "orbe = np.sort(orbe)\n",
    "print('id    localized    canonical')\n",
    "print('----------------------------')\n",
    "for i in range(Norb):\n",
    "    print('%2i %12.5f %12.5f' % ((i+1),orbe[i],res['orbe'][i]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the algorithm converges relatively rapidly. This is not always the case, but for small molecules, not a lot of iterations are necessary. \n",
    "\n",
    "Perhaps more interesting are the molecular orbital energies. In the listing as shown at the end of the script, the second column corresponds to the molecular orbital energies after the orbital localization and the third column are the molecular orbital energies *prior* to mixing, i.e., the energies for the canonical molecular orbitals. Remember that only the first 7 molecular orbitals are occupied, the other 3 are unoccupied and virtual. Because we do not mix among the virtual orbitals, the energies match prior and after mixing.\n",
    "\n",
    "Looking at the canonical molecular orbitals, you will find that there is one double-degerate pair among the occupied orbitals. In contrast, for the localized orbitals, you will find a triple-degenerate set. To interpret these results, it can be of great advantage to produce isosurfaces of the localized molecular orbitals. This will be your next exercise.\n",
    "\n",
    "### Exercise 02\n",
    "\n",
    "Produce the isosurface of all the occupied molecular orbitals after the Foster-Boys localization.\n",
    "\n",
    ":::{hint}\n",
    "* Have a look at the previous tutorial. We can readily recycle the `build_isosurface` script.\n",
    "* Note that the orbital localization has already been performed, you can readily use the `C` variable to produce the isosurfaces.\n",
    "* Recall that if you use the `build_isosurface` script that the `ply` files are written in the same folder as where the notebook resides.\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from pytessel import PyTessel\n",
    "\n",
    "def build_isosurface(filename, cgfs, coeff, isovalue):\n",
    "    \"\"\"\n",
    "    Build the isosurface of a molecular orbital. Note that the dimensions of \n",
    "    the unit cell are hardcoded to be 10x10x10 centered at the origin.\n",
    "    \n",
    "    Good trial isovalues are +/- 0.1\n",
    "    \n",
    "    filename: output filename, should end in .ply\n",
    "    cgfs:     list of CGF objects\n",
    "    coeff:    list of linear expansion coefficients\n",
    "    isovalue: isovalue of the isosurface\n",
    "    \"\"\"\n",
    "    sz = 100\n",
    "    integrator = PyQInt()\n",
    "    grid = integrator.build_rectgrid3d(-5, 5, sz)\n",
    "    scalarfield = np.reshape(integrator.plot_wavefunction(grid, coeff, cgfs), (sz, sz, sz))\n",
    "    unitcell = np.diag(np.ones(3) * 10.0)\n",
    "\n",
    "    pytessel = PyTessel()\n",
    "    vertices, normals, indices = pytessel.marching_cubes(scalarfield.flatten(), scalarfield.shape, unitcell.flatten(), isovalue)\n",
    "    pytessel.write_ply(filename, vertices, normals, indices)\n",
    "    \n",
    "for i in range(0,N):\n",
    "    print('Building MO %i' % (i+1))\n",
    "    build_isosurface('mo_local_co_%02i_pos.ply' % (i+1), cgfs, C[:,i], 0.1)\n",
    "    build_isosurface('mo_local_co_%02i_neg.ply' % (i+1), cgfs, C[:,i], -0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If everything went according to plan, you should see isosurfaces very similar to the picture as shown below.\n",
    "\n",
    "![isosurfaces of the localized molecular orbitals of CO](https://ifilot.pages.tue.nl/pyqint-han-sur-lesse-jupyterlab/_images/MO_CO_FB.png)\n",
    "\n",
    ":::{note}\n",
    "The image above is grabbed from the internet and might require an active internet connection.\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 03\n",
    "\n",
    "* Perform a Foster-Boys orbital localization procedure on the CH4 molecule.\n",
    "* Analyze the result. How many degenerate orbitals do you find? What is their interpretation?\n",
    "\n",
    ":::{note}\n",
    "Although I did not previously emphasize this, you really want to perform this procedure on the optimized geometry for best results. As such, I have added a snippet of example code below how you can perform a geometry optimization using PyQInt.\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pyqint import MoleculeBuilder, GeometryOptimization\n",
    "\n",
    "# calculate CH4 molecule\n",
    "mol = MoleculeBuilder().from_name('ch4')\n",
    "res = GeometryOptimization().run(mol, 'sto3g')['data'] # perform a geometry optimization and take the data of the last ionic step\n",
    "\n",
    "# extract object files\n",
    "Norb = len(res['orbe'])\n",
    "N = res['nelec']//2\n",
    "C = res['orbc']\n",
    "occ = (np.arange(len(res['orbe'])) < res['nelec']//2) * 2\n",
    "\n",
    "# Objective: construct the dipole tensor for CH4\n",
    "Norb = len(res['orbe'])\n",
    "cgfs = res['cgfs']\n",
    "dipol = np.zeros((Norb,Norb,3))\n",
    "integrator = PyQInt()\n",
    "for i,cgf1 in enumerate(cgfs):\n",
    "    for j,cgf2 in enumerate(cgfs):\n",
    "        for k in range(0,3):\n",
    "            dipol[i,j,k] = integrator.dipole(cgf1, cgf2, k, 0.0)\n",
    "\n",
    "old_r2 = 0.0     # start assuming orbital centroids lie at origin\n",
    "r2_arr = []      # store r2 values in an array\n",
    "diff = 1\n",
    "nriter = 0\n",
    "max_iter = 10000 # always set a maximum number of iterations in a while loop ;-)\n",
    "\n",
    "# Objective: add here the code to perform the Foster-Boys localization\n",
    "C = construct_random_orthogonal_transformation(C, N)\n",
    "\n",
    "while diff > 1e-10 and nriter < max_iter:\n",
    "    C, r2 = perform_foster_boys_mixing(C, dipol, occ)\n",
    "    r2_arr.append(r2)\n",
    "    diff = np.abs(r2 - old_r2)\n",
    "    old_r2 = r2\n",
    "    nriter += 1\n",
    "\n",
    "# Objective: print the molecular orbital energies\n",
    "orbe = np.diag(C.transpose() @ res['fock'] @ C)\n",
    "orbe = np.sort(orbe)\n",
    "print('id    localized    canonical')\n",
    "print('----------------------------')\n",
    "for i in range(Norb):\n",
    "    print('%2i %12.5f %12.5f' % ((i+1),orbe[i],res['orbe'][i]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise 04\n",
    "\n",
    "* Visualize the localized molecular orbitals of CH4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for i in range(0,N):\n",
    "    print('Building MO %i' % (i+1))\n",
    "    build_isosurface('mo_local_ch4_%02i_pos.ply' % (i+1), cgfs, C[:,i], 0.1)\n",
    "    build_isosurface('mo_local_ch4_%02i_neg.ply' % (i+1), cgfs, C[:,i], -0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "If everything went according to plan, you should see isosurfaces very similar to the picture as shown below.\n",
    "\n",
    "![isosurfaces of the localized molecular orbitals of CH4](https://ifilot.pages.tue.nl/pyqint-han-sur-lesse-jupyterlab/_images/MO_CH4_FB.png)\n",
    "\n",
    ":::{note}\n",
    "The image above is grabbed from the internet and might require an active internet connection.\n",
    ":::"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Exercise 05\n",
    "\n",
    "* Produce the canonical and localized orbitals of ethylene and compare their orbital energies. **Rather than performing the Foster-Boys routine by yourself, look into the [manual of PyQInt](https://pyqint.imc-tue.nl/user_interface.html#orbital-localization-foster-boys) and use the pre-written routine.**\n",
    "* Produce the isosurfaces of these orbitals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pyqint import FosterBoys\n",
    "\n",
    "molname = 'ethylene'\n",
    "mol = MoleculeBuilder().from_name('ethylene')\n",
    "res = HF().rhf(mol, 'sto3g')\n",
    "\n",
    "# Objective: Perform the Foster-Boys localization using the pre-written routine\n",
    "resfb = FosterBoys(res).run()\n",
    "\n",
    "# Objective: print the molecular orbital energies\n",
    "print('id    localized    canonical')\n",
    "print('----------------------------')\n",
    "for i in range(Norb):\n",
    "    print('%2i %12.5f %12.5f' % ((i+1),res['orbe'][i],resfb['orbe'][i]))\n",
    "\n",
    "for i in range(0,len(res['orbe'])):\n",
    "    print('Building MO %i' % (i+1))\n",
    "    build_isosurface('mo_canonical_ethylene_%02i_pos.ply' % (i+1), res['cgfs'], res['orbc'][:,i], 0.1)\n",
    "    build_isosurface('mo_canonical_ethylene_%02i_neg.ply' % (i+1), res['cgfs'], res['orbc'][:,i], -0.1)\n",
    "    build_isosurface('mo_localized_ethylene_%02i_pos.ply' % (i+1), res['cgfs'], resfb['orbc'][:,i], 0.1)\n",
    "    build_isosurface('mo_localized_ethylene_%02i_neg.ply' % (i+1), res['cgfs'], resfb['orbc'][:,i], -0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "If everything went according to plan, you should see isosurfaces very similar to the picture as shown below.\n",
    "\n",
    "#### Canonical orbitals of ethylene\n",
    "![isosurfaces of the localized molecular orbitals of ethylene](https://ifilot.pages.tue.nl/pyqint-han-sur-lesse-jupyterlab/_images/MO_ethylene_CAN.png)\n",
    "\n",
    "#### Localized orbitals of ethylene\n",
    "![isosurfaces of the localized molecular orbitals of ethylene](https://ifilot.pages.tue.nl/pyqint-han-sur-lesse-jupyterlab/_images/MO_ethylene_FB.png)\n",
    "\n",
    ":::{note}\n",
    "The image above is grabbed from the internet and might require an active internet connection.\n",
    ":::"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}