DFT_Calculator_and_Analaysis_4.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyOhQzXO8gsI5PphFF8BpIEX",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Ash100/DaS/blob/main/DFT_Calculator_and_Analaysis_4.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "I am **Dr. Ashfaq Ahmad**, working in the field of Bioinformatics and Structure Biology. This notebook is designed from the published work of Jacob S. Hirschi and can be found and cited from here *https://doi.org/10.1021/acs.jchemed.3c00535*\n",
        "\n",
        "If you are interested in learning with me, you can subscribe and join my youtube channel here **https://www.youtube.com/@Bioinformaticsinsights**"
      ],
      "metadata": {
        "id": "qx880u38kQ7C"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "GIgZ20uzkPXu"
      },
      "outputs": [],
      "source": [
        "#@title Install/Import Packages\n",
        "\n",
        "# import standard anaconda packages\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import pandas as pd\n",
        "%config InlineBackend.figure_format = 'svg'\n",
        "from scipy import sparse\n",
        "from IPython.display import Markdown, display, clear_output\n",
        "from time import sleep\n",
        "import ipywidgets as widgets\n",
        "import plotly\n",
        "import plotly.graph_objects as go\n",
        "from plotly.subplots import make_subplots\n",
        "\n",
        "# Keep this commented out until ipyvolume is fixed\n",
        "#try:\n",
        "#  import ipyvolume as ipv\n",
        "#except:\n",
        "#  !pip install -q ipyvolume==v0.6.3\n",
        "#  import ipyvolume as ipv\n",
        "\n",
        "from google.colab import output\n",
        "output.enable_custom_widget_manager()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### PAH Specifications\n",
        "\n",
        "<center>\n",
        "\n",
        "| PAH        | Rings | Dimensions (Bohr) | $\\pi$ electrons | Fuse-type    |\n",
        "|:----------:|:-----:|:-----------------:|:---------------:|:------------:|\n",
        "| benzene    |   1   |     8 x 8 x 3     |       6         |   linear     |\n",
        "| naphthalene|   2   |    12 x 8 x 3     |      10         |   linear     |\n",
        "| anthracene |   3   |    16 x 8 x 3     |      14         |   linear     |\n",
        "| tetracene  |   4   |    20 x 8 x 3     |      18         |   linear     |\n",
        "| pentacene  |   5   |    24 x 8 x 3     |      22         |   linear     |\n",
        "| hexacene   |   6   |    28 x 8 x 3     |      26         |   linear     |\n",
        "| heptacene  |   7   |    32 x 8 x 3     |      30         |   linear     |\n",
        "| ---        |  ---  |     ---           |      ---        |     ---      |\n",
        "| perylene   |   5   |   20 x 12 x 3     |      20         | non-linear   |\n",
        "| coronene   |   7   |   20 x 16 x 3     |      24         | non-linear   |\n",
        "\n",
        "<center/>"
      ],
      "metadata": {
        "id": "xEUpaERtkbKN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **DFT Calculator**\n",
        "\n",
        "### Define GUI widgets\n",
        "# Box length widgets\n",
        "length_labels = widgets.Label(value='Box Lengths (Bohr): ')\n",
        "lx_slider_dft = widgets.IntSlider(value=12,min=3,max=32,step=1,description='lx',disabled=False,readout_format='d',continuous_update=False)\n",
        "ly_slider_dft = widgets.IntSlider(value=8,min=3,max=32,step=1,description='ly',disabled=False,readout_format='d',continuous_update=False)\n",
        "lz_slider_dft = widgets.IntSlider(value=3,min=3,max=32,step=1,description='lz',disabled=False,readout_format='d',continuous_update=False)\n",
        "\n",
        "# Hamiltonian form GUI widgets\n",
        "functional_dropdown = widgets.Dropdown(options=[('LDA', 'LDA'), ('PBE (GGA)', 'PBE')],\n",
        "                                       value='LDA',\n",
        "                                       description='Fuctional:',\n",
        "                                       disabled=False,\n",
        "                                       layout = widgets.Layout(width='200px'))\n",
        "hartree_dropdown    = widgets.Dropdown(options=[('On', True), ('Off', False)],\n",
        "                                       value=True,\n",
        "                                       description='Hartree:',\n",
        "                                       disabled=False,\n",
        "                                       layout = widgets.Layout(width='175px'))\n",
        "exchange_dropdown   = widgets.Dropdown(options=[('On', True), ('Off', False)],\n",
        "                                       value=True,\n",
        "                                       description='Exchange:',\n",
        "                                       disabled=False,\n",
        "                                       layout = widgets.Layout(width='175px'))\n",
        "correlation_dropdown = widgets.Dropdown(options=[('On', True), ('Off', False)],\n",
        "                                       value=True,\n",
        "                                       description='Correlation:',\n",
        "                                       disabled=False,\n",
        "                                       layout = widgets.Layout(width='175px'))\n",
        "\n",
        "# Number of electrons widget\n",
        "num_elec_label    = widgets.Label(value='Number Electrons: ')\n",
        "num_elec_dropdown = widgets.Dropdown(options=np.arange(2, 31, 2),\n",
        "                                    value=10, description=' ',\n",
        "                                    disabled=False,\n",
        "                                     layout = widgets.Layout(width='175px'))\n",
        "\n",
        "converge_parameter_label = widgets.Label(value='Convergence Settings: ')\n",
        "max_iter_dropdown        = widgets.Dropdown(options=np.arange(5, 105, 5),\n",
        "                                    value=30, description='max iter',\n",
        "                                    disabled=False,\n",
        "                                     layout = widgets.Layout(width='150px'))\n",
        "\n",
        "density_mix_dropdown     = widgets.Dropdown(options=['OFF', '50/50', '60/40', '70/30', '80/20', '90/10'],\n",
        "                                       value='70/30',\n",
        "                                       description='density mix (%): ',\n",
        "                                       disabled=False,\n",
        "                                       layout = widgets.Layout(width='200px'),\n",
        "                                       style = {'description_width': 'initial'})\n",
        "\n",
        "e_tol_dropdown = widgets.Dropdown(options=['-3', '-4', '-5', '-6'],\n",
        "                                       value='-4',\n",
        "                                       description='energy tol (10^ Ha): ',\n",
        "                                       disabled=False,\n",
        "                                       layout = widgets.Layout(width='200px'),\n",
        "                                        style = {'description_width': 'initial'})\n",
        "\n",
        "# save figure buttons for subsequent analysis cells\n",
        "filename_text = widgets.Text(description='Filename (.png): ',value='energy_log',style={'description_width': 'initial'})\n",
        "save_button   = widgets.Button(description='Save Image')\n",
        "\n",
        "den_filename_text = widgets.Text(description='Filename (.png): ',value='density_slice',style={'description_width': 'initial'})\n",
        "den_save_button   = widgets.Button(description='Save Image')\n",
        "\n",
        "orbital_filename_text = widgets.Text(description='Filename (.png): ',value='orbital_energies',style={'description_width': 'initial'})\n",
        "orbital_save_button   = widgets.Button(description='Save Image')\n",
        "\n",
        "run_scf_button = widgets.Button(description='Run SCF')\n",
        "run_scf_button.style.font_weight = 'bold'\n",
        "reset_button   = widgets.Button(description='Reset')\n",
        "\n",
        "lib_dropdown = widgets.Dropdown(options=[('ipyvolume', 'ipyvol'), ('plotly', 'plotly')], value='plotly',\n",
        "                                description='3D Plotting Library: ', disabled=False, layout = widgets.Layout(width='225px'),\n",
        "                                style = {'description_width': 'initial'})\n",
        "lib_dropdown.disabled = True\n",
        "\n",
        "density_analysis_button = widgets.Button(description='Analyze Density')\n",
        "run_scf_button.style.font_weight = 'bold'\n",
        "\n",
        "side_select = widgets.SelectionSlider(\n",
        "options=['lx', 'ly', 'lz'],\n",
        "    value='lx',\n",
        "    description='box side: ',\n",
        "    disabled=False,\n",
        "    continuous_update=False,\n",
        "    orientation='horizontal',\n",
        "    readout=True\n",
        ")\n",
        "\n",
        "scf_ui            = widgets.HBox([run_scf_button, reset_button])\n",
        "box_length_iso_ui = widgets.HBox([length_labels, lx_slider_dft, ly_slider_dft, lz_slider_dft])\n",
        "parameter_ui      = widgets.HBox([num_elec_label, num_elec_dropdown])\n",
        "ham_ui            = widgets.HBox([functional_dropdown, hartree_dropdown,\n",
        "                                  exchange_dropdown, correlation_dropdown])\n",
        "convergence_ui    = widgets.HBox([converge_parameter_label, max_iter_dropdown,\n",
        "                                  density_mix_dropdown, e_tol_dropdown])\n",
        "filename_text     = widgets.Text(description='Filename (.png): ', value='energy_log',\n",
        "                                 style={'description_width': 'initial'})\n",
        "save_button       = widgets.Button(description='Save Image')\n",
        "\n",
        "def hamiltonian_display(functional, har, ex, cor):\n",
        "  '''Generates LaTeX equation of effective single-particle Hamiltonian.\n",
        "  Args:\n",
        "    functional (str): XC potential (LDA or PBE).\n",
        "    har (bool): Hartree potential on or off.\n",
        "    ex (bool): Exchange term on or off.\n",
        "    cor (bool): Correlation term on or off.\n",
        "  '''\n",
        "  ham = '## $$ \\hat{h}_i = \\hat{T}_{\\\\text{kin}, i} +'\n",
        "  if har == True:\n",
        "    ham += 'v_{\\\\text{Ha}}(n(r))'\n",
        "  else:\n",
        "    ham += '0'\n",
        "  if ex == True:\n",
        "    if functional == 'LDA':\n",
        "      ham += '+ v_{\\\\text{X}}^{\\\\text{LDA}}(n(r))'\n",
        "    elif functional == 'PBE':\n",
        "      ham += '+ v_{\\\\text{X}}^{\\\\text{PBE}}(n(r), \\\\nabla n(r))'\n",
        "  else:\n",
        "    ham += '+ 0'\n",
        "  if cor == True:\n",
        "    if functional == 'LDA':\n",
        "      ham += '+ v_{\\\\text{C}}^{\\\\text{LDA}}(n(r))'\n",
        "    elif functional == 'PBE':\n",
        "      ham += '+ v_{\\\\text{C}}^{\\\\text{PBE}}(n(r), \\\\nabla n(r))'\n",
        "  else:\n",
        "    ham += '+ 0'\n",
        "  ham += '$$'\n",
        "  display(Markdown(ham))\n",
        "\n",
        "output            = widgets.interactive_output(hamiltonian_display, {'functional':functional_dropdown,\n",
        "                                                 'har':hartree_dropdown,\n",
        "                                                 'ex':exchange_dropdown,\n",
        "                                                 'cor':correlation_dropdown})\n",
        "\n",
        "# various logs to keep tract of SCF loop progression\n",
        "energy_log     = [0] # zero needed for first loop; removed after end of for loop\n",
        "ener_diff_log  = []\n",
        "converged      = None\n",
        "exch_log       = []\n",
        "cor_log        = []\n",
        "har_log        = []\n",
        "kin_log        = []\n",
        "eigenvalue_log = []\n",
        "eigenstate_log = []\n",
        "density_log    = []\n",
        "grid_points    = []\n",
        "grid_lines     = []\n",
        "\n",
        "def energy_plot(den_log, ener_log, converge_state, show_fig=True, save_fig=False, filename=None):\n",
        "  '''Generates iteration number vs total energy plot.\n",
        "  Args:\n",
        "    den_log (list): List of numpy arrays with electron density.\n",
        "    ener_log (list): List of total energy values in Ha.\n",
        "    converge_state (bool): SCF loop ended in a converged (True) or unconverged (False) state.\n",
        "    show_fig (bool): Display figure to output (True) or not (False).\n",
        "    save_fig (bool): Save a .png file of diagram.\n",
        "    filename (str): Filename to save to (w/o extension).\n",
        "  '''\n",
        "  fig = plt.figure(figsize=(6, 4))\n",
        "  ax = fig.add_subplot(1, 1, 1)\n",
        "  plt.title('Convergence Plot')\n",
        "  plt.scatter(0, energy_log[0], color='#F97306', label='noninter Energy')\n",
        "  plt.plot(np.arange(1,len(energy_log[1:]) + 1), energy_log[1:], 'o-', label='DFT Energy')\n",
        "  plt.legend(loc='upper right')\n",
        "  plt.text(0.785, 0.800, f'iterations: {len(density_log) - 1}',\n",
        "          horizontalalignment='center', verticalalignment='center',\n",
        "          transform=ax.transAxes)\n",
        "  plt.text(0.785, 0.730,f'            energy (Ha): {energy_log[-1]:.5f}',\n",
        "          size=10.5, horizontalalignment='center', verticalalignment='center',\n",
        "          transform= ax.transAxes)\n",
        "  if converge_state == True:\n",
        "    plt.text(0.79, 0.660,f'converged',\n",
        "            size=10.5, horizontalalignment='center', verticalalignment='center',\n",
        "            transform= ax.transAxes, weight='bold')\n",
        "  elif converge_state == False:\n",
        "    plt.text(0.79, 0.660,f'unconverged',\n",
        "            size=10.5, horizontalalignment='center', verticalalignment='center',\n",
        "            transform= ax.transAxes, weight='bold')\n",
        "  plt.ylabel('Energy (Ha)')\n",
        "  plt.xlabel('iteration #')\n",
        "  plt.tight_layout()\n",
        "  if save_fig:\n",
        "    plt.savefig(f'{filename}.png', dpi = 800)\n",
        "  if show_fig:\n",
        "    plt.show()\n",
        "  else:\n",
        "    plt.close()\n",
        "\n",
        "def edge_cleaner(func_3d, nx, ny, nz, num_edges=1):\n",
        "  '''Sets outermost layer (or two) of grid values of a 3D function to zero.\n",
        "  This is needed for the density because of the discontinuities in the numeric 2nd derivatives at the box edges.\n",
        "  Args:\n",
        "    func_3d (np.array): Flattened 3D grid of points.\n",
        "    nx, ny, nz (int): Number of grid points in each dimension.\n",
        "    num_edges (1 or 2): Number of border layers to set to zero.\n",
        "  Returns:\n",
        "    A flattened array of grid points.\n",
        "  '''\n",
        "  func_3d = func_3d.reshape(nx, ny, nz)\n",
        "  if num_edges == 1:\n",
        "    func_3d[0,:,:]  = 0\n",
        "    func_3d[-1,:,:] = 0\n",
        "    func_3d[:,0,:]  = 0\n",
        "    func_3d[:,-1,:] = 0\n",
        "    func_3d[:,:,0]  = 0\n",
        "    func_3d[:,:,-1] = 0\n",
        "  elif num_edges == 2:\n",
        "    func_3d[0,:,:]  = 0\n",
        "    func_3d[-1,:,:] = 0\n",
        "    func_3d[:,0,:]  = 0\n",
        "    func_3d[:,-1,:] = 0\n",
        "    func_3d[:,:,0]  = 0\n",
        "    func_3d[:,:,-1] = 0\n",
        "    func_3d[1,:,:]  = 0\n",
        "    func_3d[-2,:,:] = 0\n",
        "    func_3d[:,1,:]  = 0\n",
        "    func_3d[:,-2,:] = 0\n",
        "    func_3d[:,:,1]  = 0\n",
        "    func_3d[:,:,-2] = 0\n",
        "  return func_3d.flatten()\n",
        "\n",
        "def hard_walls(potential, nx, ny, nz):\n",
        "  '''Sets potential of outermost grid points to 1,000. Used for testing.\n",
        "  Args:\n",
        "    potential (np.array): KS (or other) potential as a flattened array.\n",
        "    nx, ny, nz (int): Grid points in each direction.\n",
        "  Returns:\n",
        "    Flattened array of grid points.\n",
        "  '''\n",
        "  potential         = potential.reshape(nx, ny, nz)\n",
        "  potential[0,:,:]  = 1000\n",
        "  potential[-1,:,:] = 1000\n",
        "  potential[:,0,:]  = 1000\n",
        "  potential[:,-1,:] = 1000\n",
        "  potential[:,:,0]  = 1000\n",
        "  potential[:,:,-1] = 1000\n",
        "  return potential.flatten()\n",
        "\n",
        "def integ_3d(func_3d, dx, dy, dz):\n",
        "  '''Integrates a 3D function over all defined space.\n",
        "  Args:\n",
        "    func_3d (np.array): 3D grid of points.\n",
        "    dx, dy, dz (float): Differential volume element in each dimension.\n",
        "  Returns:\n",
        "    Integrated 3D function as scalar value.\n",
        "  '''\n",
        "  return np.sum(func_3d * dx * dy * dz)\n",
        "\n",
        "def norm_psi_and_den(e_vecs, occ_states, dx, dy, dz):\n",
        "  '''Normalizes raw eigenvectors from the solver and finds electron density.\n",
        "  Args:\n",
        "    e_vecs (np.array): Array of eigenvectors from the eigenvalue solver.\n",
        "    occ_states (int): Number of occupied KS states (i.e. # electrons / 2).\n",
        "    dx, dy, dz (float): Differential volume element in each dimension.\n",
        "  Returns:\n",
        "    norm_psi (np.array): Array of normalized eigenvectors.\n",
        "    el_den (np.array): Electron density as array.\n",
        "  '''\n",
        "  norm_psi = np.zeros_like(e_vecs)\n",
        "  el_den   = np.zeros_like(e_vecs[:,0])\n",
        "  for i in range(e_vecs.shape[1]):\n",
        "      norm_psi[:,i] = e_vecs[:,i]/np.sqrt(integ_3d(e_vecs[:,i]**2, dx, dy, dz))\n",
        "  for i in range(occ_states):\n",
        "      el_den += 2* norm_psi[:,i]**2\n",
        "  return norm_psi, el_den\n",
        "\n",
        "def noninter_kin_e(norm_eigenvecs, occ_states, kin_mat, dx, dy, dz, nx, ny, nz):\n",
        "  '''Finds noninteracting KS kinetic energy.\n",
        "  Args:\n",
        "    norm_eigenvecs (np.array): Array of normalized eigenvectors as columns/rows.\n",
        "    occ_states (int): Number of occupied KS states (i.e. # electrons / 2).\n",
        "    kin_mat (sparse array): Kinetic energy operator as matrix.\n",
        "    dx, dy, dz (float): Differential volume element in each dimension.\n",
        "    nx, ny, nz (int): Number of grid points in each dimension.\n",
        "  Returns:\n",
        "    Scalar energy value in Ha.\n",
        "  '''\n",
        "  kin_energy_values = []\n",
        "  for eig in norm_eigenvecs.T[:occ_states]:\n",
        "    inner_prod   = eig*kin_mat.dot(eig)\n",
        "    inner_prod   = edge_cleaner(inner_prod, nx, ny, nz, num_edges=1)\n",
        "    orbital_k_en = integ_3d(inner_prod, dx, dy, dz)\n",
        "    kin_energy_values.append(orbital_k_en)\n",
        "  return sum(kin_energy_values)\n",
        "\n",
        "def hartree(den, kin_oper, dx, dy, dz, nx, ny, nz):\n",
        "  '''Uses Poisson's equation to find Hartree potential from density.\n",
        "  Args:\n",
        "    den (np.array): Electron density.\n",
        "    kin_oper (np.array): Kinetic energy operator as sparse matrix.\n",
        "    dx, dy, dz (float): Differential volume element in each dimension.\n",
        "    nx, ny, nz (int): Number of grid points in each dimension.\n",
        "  Returns:\n",
        "    v_ha_flat (np.array): Hartree potential as flattened array.\n",
        "    v_ha_ener (float): Hartree energy (in Ha); needed to compute total energy.\n",
        "  '''\n",
        "  clean_den = np.ma.array(den, mask= abs(den) < 0.000001) # Mask the low-density points\n",
        "  clean_den = np.ma.filled(clean_den, fill_value=0.0) # Fill with zeros\n",
        "  den       = clean_den\n",
        "  den       = edge_cleaner(den, nx, ny, nz, num_edges=1) # Set edges to zero\n",
        "  v_ha_flat = sparse.linalg.cg(-2*kin_oper,-4.*np.pi*den)[0]\n",
        "  v_ha_flat = edge_cleaner(v_ha_flat, nx, ny, nz, num_edges=1)\n",
        "  v_ha_ener = (1/2)*integ_3d(v_ha_flat*den, dx, dy, dz)\n",
        "  return v_ha_flat, v_ha_ener\n",
        "\n",
        "def lda_exchange(den, dx, dy, dz):\n",
        "  '''Finds LDA exchange potential and energy from density.\n",
        "  Args:\n",
        "    den (np.array): Electron density.\n",
        "    dx, dy, dz (float): Differential volume element in each dimension.\n",
        "  Returns:\n",
        "    exch_pot_lda (np.array): LDA exchange potential.\n",
        "    exch_ener_lda (float): LDA exchnage energy (in Ha).\n",
        "  '''\n",
        "  exch_pot_lda = -(3/4)*(3/np.pi)**(1/3)*(den)**(1/3)\n",
        "  clean_den = np.ma.array(den, mask= abs(den) < 0.000001)\n",
        "  clean_den = np.ma.filled(clean_den, fill_value=0.0)\n",
        "  den       = clean_den\n",
        "  exch_ener_lda = -(3/4)*(3/np.pi)**(1/3)*integ_3d(den**(4/3), dx, dy, dz)\n",
        "  return exch_pot_lda, exch_ener_lda\n",
        "\n",
        "def lda_correlation(den, dx, dy, dz):\n",
        "  '''Finds LDA correlation potential and energy from density.\n",
        "  Args:\n",
        "    den (np.array): Electron density.\n",
        "    dx, dy, dz (float): Differential volume element in each dimension.\n",
        "  Returns:\n",
        "    corr_pot (np.array): LDA correlation potential.\n",
        "    corr_en (float): LDA correlation energy (in Ha).\n",
        "  '''\n",
        "  # a, b, c fitting constants used in paper: 10.1063/1.4958669\n",
        "  a, b, c   = (np.log(2)-1)/(2*np.pi**2), 20.4562557, (4*np.pi/3)**(1/3)\n",
        "  corr_pot  = a*np.log(1 + b*c*den**(1/3) + b*(c**2)*den**(2/3))\n",
        "  clean_den = np.ma.array(den, mask= abs(den) < 0.000001)\n",
        "  clean_den = np.ma.filled(clean_den, fill_value=0.0)\n",
        "  den       = clean_den\n",
        "  corr_en   = integ_3d(den*corr_pot, dx, dy, dz)\n",
        "  return corr_pot, corr_en\n",
        "\n",
        "def RDG(den, der_1st, nx, ny, nz):\n",
        "  '''Finds dimensionless reduced density gradient (needed for PBE exchange).\n",
        "  Args:\n",
        "    den (np.array): Electron density.\n",
        "    der_1st (np.array): 1st derivative operator as sparse matrix.\n",
        "    nx, ny, nz (int): Number of grid points in each dimension.\n",
        "  Returns:\n",
        "    RDG (np.array): Reduced density gradient as matrix.\n",
        "  '''\n",
        "  clean_den = np.ma.array(den, mask = abs(den) < 0.000001)\n",
        "  den       = clean_den\n",
        "  RDG       = (2*3**(1/3)*np.pi**(2/3))**(-1) * abs(der_1st.dot(den)) * den**(-4/3)\n",
        "  RDG       = edge_cleaner(RDG, nx, ny, nz, num_edges=1) #set the edges to zero\n",
        "  RDG       = np.ma.filled(RDG, fill_value=0.0) # return zeros where density is very small\n",
        "  return RDG\n",
        "\n",
        "def pbe_exchange(den, D1st, dx, dy, dz, nx, ny, nz):\n",
        "  '''Finds PBE exchange potential and energy from density + gradient.\n",
        "  Args:\n",
        "    den (np.array): Electron density.\n",
        "    D1st (np.array): 1st derivative operator as sparse matrix.\n",
        "    dx, dy, dz (float): Differential volume element in each dimension.\n",
        "    nx, ny, nz (int): Number of grid points in each dimension.\n",
        "  Returns:\n",
        "    exch_pot_pbe (np.array): PBE exchange potential.\n",
        "    exch_ener_pbe (float): PBE exchange energy (in Ha).\n",
        "  '''\n",
        "  # kappa and mu constants used in the paper, 10.1103/PhysRevLett.77.3865\n",
        "  kappa, mu = 0.804, 0.2195149727645171\n",
        "  s         = RDG(den, D1st, nx, ny, nz)\n",
        "  F_xs      = 1 + kappa - kappa * (1 + mu * s**2 / kappa)**(-1) # exch enhancement factor\n",
        "  exch_pot_pbe  = F_xs * -(3/4)*(3/np.pi)**(1/3)*((den)**(1/3))\n",
        "  clean_den     = np.ma.array(den, mask= abs(den) < 0.000001)\n",
        "  clean_den     = np.ma.filled(clean_den, fill_value=0.0)\n",
        "  den           = clean_den\n",
        "  exch_ener_pbe = integ_3d(den*exch_pot_pbe, dx, dy, dz)\n",
        "  return exch_pot_pbe, exch_ener_pbe\n",
        "\n",
        "def cor_den_grad(den, der_1st, nx, ny, nz):\n",
        "  '''Finds dimensionless correlation density gradient (used in PBE correlation).\n",
        "  Args:\n",
        "    den (np.array): Electron density.\n",
        "    der_1st (np.array): 1st derivative operator as sparse matrix.\n",
        "    nx, ny, nz (int): Number of grid points in each dimension.\n",
        "  Returns:\n",
        "    t (np.array): Correlation density gradient.\n",
        "  '''\n",
        "  d_g       = abs(der_1st.dot(den))\n",
        "  clean_den = np.ma.array(den, mask = abs(den) < 0.000001)\n",
        "  den       = clean_den\n",
        "  t         = (d_g*np.pi**(1/6))/(4*3**(1/6)*den**(7/6))\n",
        "  t         = edge_cleaner(t, nx, ny, nz, num_edges=1)\n",
        "  t         = np.ma.filled(t, fill_value=0.0)\n",
        "  return t\n",
        "\n",
        "def pbe_correlation(den, der_1st, dx, dy, dz, nx, ny, nz):\n",
        "  '''Finds PBE correlation potential and energy from density + gradient.\n",
        "  Args:\n",
        "    den (np.array): Electron density.\n",
        "    der_1st (np.array): 1st derivative operator as sparse matrix.\n",
        "    dx, dy, dz (float): Differential volume element in each dimension.\n",
        "    nx, ny, nz (int): Number of grid points in each dimension.\n",
        "  Returns:\n",
        "    cor_pot_pbe (np.array): PBE correlation potential.\n",
        "    cor_ener_pbe (float): PBE correlation energy (in Ha).\n",
        "  '''\n",
        "  lda_c_pot = lda_correlation(den, dx, dy, dz)[0]\n",
        "  # beta and gamma constants used in the paper, 10.1103/PhysRevLett.77.3865\n",
        "  beta, gamma = 0.06672455060314922, 0.031090690869654894\n",
        "  lda_c_pot = np.ma.array(lda_c_pot, mask = abs(lda_c_pot) < 0.000001)\n",
        "  A = (beta/gamma)*((np.exp(-lda_c_pot/gamma)-1)**(-1))\n",
        "  t = cor_den_grad(den, der_1st, nx, ny, nz)\n",
        "  H = gamma*np.log(1+(beta/gamma)*t**2*((1+A*(t**2))/(1+A*(t**2)+(A**2)*(t**4))))\n",
        "  cor_pot_pbe = lda_c_pot + H\n",
        "  cor_ener_pbe = integ_3d(den*cor_pot_pbe, dx, dy, dz)\n",
        "  return cor_pot_pbe, cor_ener_pbe\n",
        "\n",
        "def reset(b):\n",
        "  ''' Resets the calculator with button click. '''\n",
        "  clear_output()\n",
        "  # re-enable the GUI\n",
        "  lx_slider_dft.disabled        = False\n",
        "  ly_slider_dft.disabled        = False\n",
        "  lz_slider_dft.disabled        = False\n",
        "  num_elec_dropdown.disabled    = False\n",
        "  functional_dropdown.disabled  = False\n",
        "  hartree_dropdown.disabled     = False\n",
        "  exchange_dropdown.disabled    = False\n",
        "  correlation_dropdown.disabled = False\n",
        "  run_scf_button.disabled       = False\n",
        "  max_iter_dropdown.disabled    = False\n",
        "  density_mix_dropdown.disabled = False\n",
        "  e_tol_dropdown.disabled       = False\n",
        "\n",
        "  # reset the logs\n",
        "  global energy_log, ener_diff_log, converged,exch_log, cor_log, har_log, kin_log, eigenvalue_log, eigenstate_log, density_log, grid_points, grid_lines\n",
        "  energy_log     = [0] # zero needed for first difference; removed after end of for loop\n",
        "  ener_diff_log  = []\n",
        "  converged      = False\n",
        "  exch_log       = []\n",
        "  cor_log        = []\n",
        "  har_log        = []\n",
        "  kin_log        = []\n",
        "  eigenvalue_log = []\n",
        "  eigenstate_log = []\n",
        "  density_log    = []\n",
        "  grid_points    = []\n",
        "  grid_lines     = []\n",
        "\n",
        "  # regenerate the display\n",
        "  display(Markdown('## <center> Single Particle Hamiltonian <center/> '),\n",
        "                output,\n",
        "                Markdown('<br>'),\n",
        "                ham_ui,\n",
        "                Markdown('<br>'),\n",
        "                box_length_iso_ui,\n",
        "                Markdown('<br>'),\n",
        "                parameter_ui,\n",
        "                Markdown('<br>'),\n",
        "                convergence_ui,\n",
        "                Markdown('<br>'),\n",
        "                scf_ui,\n",
        "                Markdown('---'))\n",
        "\n",
        "reset_button.on_click(reset) # link button to function\n",
        "\n",
        "def run_scf(b):\n",
        "  ''' Run the calculation with widget parameters. '''\n",
        "  lx_slider_dft.disabled        = True\n",
        "  ly_slider_dft.disabled        = True\n",
        "  lz_slider_dft.disabled        = True\n",
        "  num_elec_dropdown.disabled    = True\n",
        "  functional_dropdown.disabled  = True\n",
        "  hartree_dropdown.disabled     = True\n",
        "  exchange_dropdown.disabled    = True\n",
        "  correlation_dropdown.disabled = True\n",
        "  run_scf_button.disabled       = True\n",
        "  max_iter_dropdown.disabled    = True\n",
        "  density_mix_dropdown.disabled = True\n",
        "  e_tol_dropdown.disabled       = True\n",
        "\n",
        "  occ_states = int(num_elec_dropdown.value/2)\n",
        "  nextra = 2 # two unoccupied states calculated in addition to occupied (arbitrary number)\n",
        "  max_iter = int(max_iter_dropdown.value)\n",
        "\n",
        "  # define the energy tolerance\n",
        "  e_tol = 10**(int(e_tol_dropdown.value))\n",
        "\n",
        "  nx, ny, nz = (5 * lx_slider_dft.value), (5 * ly_slider_dft.value), (5 * lz_slider_dft.value)\n",
        "  grid_points.append([nx, ny, nz])\n",
        "  xp, yp, zp = np.linspace(0, lx_slider_dft.value, nx), np.linspace(0, ly_slider_dft.value, ny), np.linspace(0, lz_slider_dft.value, nz)\n",
        "  grid_lines.append([xp, yp, zp])\n",
        "\n",
        "  # Construction of 1st and 2nd derivative operators through finite-difference method\n",
        "\n",
        "  diag1x = np.ones(nx)/(xp[1])\n",
        "  diag1y = np.ones(nx)/(yp[1])\n",
        "  diag1z = np.ones(nx)/(zp[1])\n",
        "\n",
        "  # calculation of d/dx, d/dy, and d/dz as sparse matrices:\n",
        "  D1x = sparse.spdiags(np.array([-diag1x, diag1x]), np.array([0,1]), nx, nx)\n",
        "  D1y = sparse.spdiags(np.array([-diag1y, diag1y]), np.array([0,1]), ny, ny)\n",
        "  D1z = sparse.spdiags(np.array([-diag1z, diag1z]), np.array([0,1]), nz, nz)\n",
        "\n",
        "  # overall 1st derivative for functions on this grid:\n",
        "  D1st = sparse.kronsum(D1z,sparse.kronsum(D1y,D1x))\n",
        "\n",
        "  diagx = np.ones(nx)/(xp[1]**2)\n",
        "  diagy = np.ones(ny)/(yp[1]**2)\n",
        "  diagz = np.ones(nz)/(zp[1]**2)\n",
        "\n",
        "  # calculation of d^2/dx^2, d^2/dy^2, and d^2/dz^2 as sparse matrices:\n",
        "  Dx = sparse.spdiags(np.array([diagx, -2*diagx, diagx]), np.array([-1,0,1]), nx, nx)\n",
        "  Dy = sparse.spdiags(np.array([diagy, -2*diagy, diagy]), np.array([-1,0,1]), ny, ny)\n",
        "  Dz = sparse.spdiags(np.array([diagz, -2*diagz, diagz]), np.array([-1,0,1]), nz, nz)\n",
        "\n",
        "  # construct the sparse matrix\n",
        "  T = -1/2 * sparse.kronsum(Dz,sparse.kronsum(Dy,Dx))\n",
        "\n",
        "  columns = ['E_total', 'delta_E', 'E_KE', 'E_x', 'E_c', 'E_ha']\n",
        "  df = pd.DataFrame(columns =['#i', 'E_tot'.center(7, \"-\"),\n",
        "                              'delta_E'.center(7, \"-\"),\n",
        "                              'E_KE'.center(7, \"-\"),\n",
        "                              'E_x'.center(7, \"-\"),\n",
        "                              'E_c'.center(7, \"-\"),\n",
        "                              'E_ha'.center(7, \"-\")])\n",
        "  display(df) # dataframe with column headings\n",
        "\n",
        "  # start the SCF loop\n",
        "  for it_num in range(max_iter + 1):\n",
        "    if it_num == 0:\n",
        "      placeholder_density = np.zeros(T.shape[0]) ## placeholder density of zero for the first iteration; removed after loop\n",
        "      density_log.append(placeholder_density)\n",
        "\n",
        "    density = density_log[-1]\n",
        "    density = np.ma.array(density, mask= abs(density) < 0.000001)\n",
        "    density = np.ma.filled(density, fill_value=0.0)\n",
        "    density = edge_cleaner(density, nx, ny, nz, num_edges=1)\n",
        "\n",
        "    ### Hartree on or off\n",
        "    if hartree_dropdown.value == True:\n",
        "      if it_num == 0:\n",
        "        har_pot  = np.zeros(T.shape[0])\n",
        "        har_ener = 0\n",
        "      elif it_num != 0:\n",
        "        har_pot, har_ener = hartree(density, T, xp[1], yp[1], zp[1], nx, ny, nz)\n",
        "    elif hartree_dropdown.value == False:\n",
        "      har_pot  = np.zeros(T.shape[0])\n",
        "      har_ener = 0\n",
        "\n",
        "    ### exchange lda, pbe, or off\n",
        "    if exchange_dropdown.value == True:\n",
        "      if functional_dropdown.value == 'LDA':\n",
        "        if it_num == 0:\n",
        "          exch_pot  = np.zeros(T.shape[0])\n",
        "          exch_ener = 0\n",
        "        elif it_num != 0:\n",
        "          exch_pot, exch_ener = lda_exchange(density, xp[1], yp[1], zp[1])\n",
        "      elif functional_dropdown.value == 'PBE':\n",
        "        if it_num == 0:\n",
        "          exch_pot  = np.zeros(T.shape[0])\n",
        "          exch_ener = 0\n",
        "        elif it_num != 0:\n",
        "          exch_pot, exch_ener = pbe_exchange(density, D1st, xp[1], yp[1], zp[1], nx, ny, nz)\n",
        "    elif exchange_dropdown.value == False:\n",
        "      exch_pot  = np.zeros(T.shape[0])\n",
        "      exch_ener = 0\n",
        "\n",
        "    ### correlation lda, pbe, or off\n",
        "    if correlation_dropdown.value == True:\n",
        "      if functional_dropdown.value == 'LDA':\n",
        "        if it_num == 0:\n",
        "          cor_pot  = np.zeros(T.shape[0])\n",
        "          cor_ener = 0\n",
        "        elif it_num != 0:\n",
        "          cor_pot, cor_ener = lda_correlation(density, xp[1], yp[1], zp[1])\n",
        "      elif functional_dropdown.value == 'PBE':\n",
        "        if it_num == 0:\n",
        "          cor_pot  = np.zeros(T.shape[0])\n",
        "          cor_ener = 0\n",
        "        elif it_num != 0:\n",
        "          cor_pot, cor_ener = pbe_correlation(density, D1st, xp[1], yp[1], zp[1], nx, ny, nz)\n",
        "    elif correlation_dropdown.value == False:\n",
        "      cor_pot  = np.zeros(T.shape[0])\n",
        "      cor_ener = 0\n",
        "\n",
        "    v_diagonal = har_pot + exch_pot + cor_pot\n",
        "    v_clean    = np.ma.array(v_diagonal, mask = abs(density) < 0.000001)\n",
        "    v_clean    = np.ma.filled(v_clean, fill_value=0.0)\n",
        "\n",
        "    # write the hamiltonian\n",
        "    V = sparse.diags(v_clean)\n",
        "\n",
        "    H = T + V\n",
        "\n",
        "    # finds the eigenvalues and eigenvectors (k is number of levels)\n",
        "    eigenenergies, raw_evecs = sparse.linalg.eigsh(H, k=occ_states+nextra, which='SM',mode='cayley')\n",
        "\n",
        "    normalized_evecs, new_density = norm_psi_and_den(raw_evecs, occ_states, xp[1], yp[1], zp[1])\n",
        "\n",
        "    new_density = edge_cleaner(new_density, nx, ny, nz, num_edges=1)\n",
        "\n",
        "    # density mixing section 'OFF', '50/50', '60/40', etc.\n",
        "    if it_num == 0:\n",
        "      density_log.append(new_density)\n",
        "    elif it_num != 0:\n",
        "      if density_mix_dropdown.value == 'OFF':\n",
        "        new_new_density = 0.0*density + 1.0*new_density\n",
        "      elif density_mix_dropdown.value == '50/50':\n",
        "        new_new_density = 0.5*density + 0.5*new_density\n",
        "      elif density_mix_dropdown.value == '60/40':\n",
        "        new_new_density = 0.6*density + 0.4*new_density\n",
        "      elif density_mix_dropdown.value == '70/30':\n",
        "        new_new_density = 0.7*density + 0.3*new_density\n",
        "      elif density_mix_dropdown.value == '80/20':\n",
        "        new_new_density = 0.8*density + 0.2*new_density\n",
        "      elif density_mix_dropdown.value == '90/10':\n",
        "        new_new_density = 0.9*density + 0.1*new_density\n",
        "      density_log.append(new_new_density)\n",
        "\n",
        "    # noninteracting kinetic energy\n",
        "    kin_ener = noninter_kin_e(normalized_evecs, occ_states, T, xp[1], yp[1], yp[1], nx, ny, nz)\n",
        "\n",
        "    # total KS-DFT energy and delta\n",
        "    total_ener = kin_ener + har_ener + exch_ener + cor_ener\n",
        "    delta_ener = total_ener - energy_log[-1]\n",
        "\n",
        "    # update the logs\n",
        "    energy_log.append(total_ener)\n",
        "    ener_diff_log.append(delta_ener)\n",
        "    kin_log.append(kin_ener)\n",
        "    exch_log.append(exch_ener)\n",
        "    cor_log.append(cor_ener)\n",
        "    har_log.append(har_ener)\n",
        "    eigenvalue_log.append(eigenenergies)\n",
        "\n",
        "    dm = pd.DataFrame(columns =[str(it_num).zfill(2), \"{:.5f}\".format(total_ener).zfill(7),\n",
        "                                \"{:.5f}\".format(delta_ener).zfill(7),\n",
        "                                \"{:.5f}\".format(kin_ener).zfill(7), \"{:.5f}\".format(exch_ener).zfill(7),\n",
        "                                \"{:.5f}\".format(cor_ener).zfill(7), \"{:.5f}\".format(har_ener).zfill(7)])\n",
        "    display(dm) # display total energy and contributions as dataframe\n",
        "\n",
        "    # Converged exit loop\n",
        "    if abs(delta_ener) < e_tol:\n",
        "      converged = True\n",
        "      latest_energies = eigenvalue_log[-1]\n",
        "      HL_Gap          = latest_energies[-2] - latest_energies[-3]\n",
        "      HL_Gap_eV       = round(HL_Gap*27.2114, 3)\n",
        "      print('   ')\n",
        "      print('Converged!')\n",
        "      print(f'delta E = {abs(delta_ener):.6f} < {e_tol:.6f}')\n",
        "      print(f'HOMO-LUMO Gap is {HL_Gap:.5f} Ha ({HL_Gap_eV:.3f} eV)')\n",
        "      print(f\"Final energy of {energy_log[-1]:.5f} Ha reached in {it_num} iterations\")\n",
        "      break\n",
        "\n",
        "    # Unconverged exit loop\n",
        "    if it_num == max_iter:\n",
        "      converged = False\n",
        "      print('   ')\n",
        "      print(f'delta E = {abs(delta_ener):.6f} > {e_tol:.6f}')\n",
        "      print(\"Not converged :( \")\n",
        "      break\n",
        "\n",
        "  del density_log[0] # remove placeholder values\n",
        "  del energy_log[0]\n",
        "\n",
        "  energy_plot(density_log, energy_log, converged)\n",
        "\n",
        "  def save_energy_plot(b):\n",
        "    ''' Save a .png file of energy plot. '''\n",
        "    energy_plot(density_log, energy_log, converged, show_fig=False,\n",
        "                save_fig=True, filename=filename_text.value)\n",
        "  save_button.on_click(save_energy_plot)\n",
        "  display(widgets.HBox([filename_text, save_button]))\n",
        "\n",
        "run_scf_button.on_click(run_scf) # link scf button + scf function\n",
        "\n",
        "# generate the display\n",
        "display(Markdown('## <center> Single Particle Hamiltonian <center/> '),\n",
        "                output,\n",
        "                Markdown('<br>'),\n",
        "                ham_ui,\n",
        "                Markdown('<br>'),\n",
        "                box_length_iso_ui,\n",
        "                Markdown('<br>'),\n",
        "                parameter_ui,\n",
        "                Markdown('<br>'),\n",
        "                convergence_ui,\n",
        "                Markdown('<br>'),\n",
        "                scf_ui,\n",
        "                Markdown('---'))"
      ],
      "metadata": {
        "id": "ALmv1OCVkeCB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Electron Density Slices**\n",
        "\n",
        "# making sure the uer has run the SCF loop first\n",
        "try:\n",
        "  if len(density_log) != 0:\n",
        "    color1, color2 = \"#D4CC47\", \"#7C4D8B\"\n",
        "\n",
        "    def hex_to_RGB(hex_str):\n",
        "      \"\"\"Convert hex to RGB color.\"\"\"\n",
        "      return [int(hex_str[i:i+2], 16) for i in range(1,6,2)]\n",
        "\n",
        "    def get_color_gradient(c1, c2, n):\n",
        "      '''Produces a color gradient given two input colors.\n",
        "      Args:\n",
        "        c1, c2 ('str'): Colors in hex format.\n",
        "        n (int): Number of output colors.\n",
        "      Returns:\n",
        "        Gradient of n colors.\n",
        "      '''\n",
        "      assert n > 1\n",
        "      c1_rgb = np.array(hex_to_RGB(c1))/255\n",
        "      c2_rgb = np.array(hex_to_RGB(c2))/255\n",
        "      mix_pcts = [x/(n-1) for x in range(n)]\n",
        "      rgb_colors = [((1-mix)*c1_rgb + (mix*c2_rgb)) for mix in mix_pcts]\n",
        "      return [\"#\" + \"\".join([format(int(round(val*255)), \"02x\") for val in item]) for item in rgb_colors]\n",
        "\n",
        "    colors = get_color_gradient(color1, color2, len(density_log))\n",
        "\n",
        "    nx = grid_points[0][0]\n",
        "    ny = grid_points[0][1]\n",
        "    nz = grid_points[0][2]\n",
        "\n",
        "    def density_analysis(index, show_fig=True, save_fig=False, filename=None):\n",
        "      '''Generates plot of planar average of the electron density.\n",
        "      Args:\n",
        "        index (str): Box side to consider, i.e. 'lx', 'ly', or 'lz'.\n",
        "        show_fig (bool): Whether to display the figure.\n",
        "        save_fig (bool): Whether to save a .png file.\n",
        "        filename (str): Filename to save to (w/o) extension.\n",
        "      '''\n",
        "      fig = plt.figure()\n",
        "      ax = fig.add_subplot(1, 1, 1)\n",
        "      xp = grid_lines[0][0]\n",
        "      yp = grid_lines[0][1]\n",
        "      zp = grid_lines[0][2]\n",
        "      for i in range(len(density_log)):\n",
        "        density_it = density_log[i].reshape(nx,ny,nz)\n",
        "        if index == 'lx':\n",
        "          lin_avg_lx = []\n",
        "          density_slice = density_it[0,:,:]\n",
        "          den_slice_sum = np.zeros_like(density_slice)\n",
        "          for j in range(nx):\n",
        "            lin_avg_lx.append(density_it[j,:,:].mean())\n",
        "          plt.plot(xp,lin_avg_lx,label=f'iteration: {i}',color=colors[i])\n",
        "          plt.xlabel('l$_x$ (Bohr)')\n",
        "        if index == 'ly':\n",
        "          lin_avg_ly = []\n",
        "          density_slice = density_it[:,0,:]\n",
        "          den_slice_sum = np.zeros_like(density_slice)\n",
        "          for j in range(ny):\n",
        "            lin_avg_ly.append(density_it[:,j,:].mean())\n",
        "          plt.plot(yp,lin_avg_ly,label=f'iteration: {i}',color=colors[i])\n",
        "          plt.xlabel('l$_y$ (Bohr)')\n",
        "        if index == 'lz':\n",
        "          lin_avg_lz = []\n",
        "          density_slice = density_it[:,:,0]\n",
        "          den_slice_sum = np.zeros_like(density_slice)\n",
        "          for j in range(nz):\n",
        "            lin_avg_lz.append(density_it[:,:,j].mean())\n",
        "          plt.plot(zp,lin_avg_lz,label=f'iteration: {i}',color=colors[i])\n",
        "          plt.xlabel('l$_z$ (Bohr)')\n",
        "      plt.text(1.27, 1.05,f'Final Energy (Ha): {energy_log[-1]:.5f}',\n",
        "          size=10.5, horizontalalignment='center', verticalalignment='center',\n",
        "          transform= ax.transAxes)\n",
        "      plt.title('Planar Average')\n",
        "      plt.ylabel('Density (e/Bohr$^{3}$)')\n",
        "      plt.legend(bbox_to_anchor=(1.05, 1.0), ncol=3)\n",
        "\n",
        "      # control events\n",
        "      if save_fig:\n",
        "        plt.savefig(f'{filename}.png', dpi=800)\n",
        "      if show_fig:\n",
        "        plt.show()\n",
        "      else:\n",
        "        plt.close()\n",
        "\n",
        "    def save_den_plot(b):\n",
        "      ''' Button click to save file with widget parameters. '''\n",
        "      density_analysis(side_select.value, show_fig=False, save_fig=True,\n",
        "                      filename=den_filename_text.value)\n",
        "\n",
        "    den_save_button.on_click(save_den_plot) # link button + function\n",
        "\n",
        "    # link widget to function\n",
        "    planar_den_output = widgets.interactive_output(density_analysis, {'index':side_select,\n",
        "                                                                      'show_fig':widgets.fixed(True),\n",
        "                                                                      'save_fig':widgets.fixed(False),\n",
        "                                                                      'filename':widgets.fixed(None)})\n",
        "    # generate display\n",
        "    display(side_select,\n",
        "            planar_den_output,\n",
        "            widgets.HBox([den_filename_text, den_save_button]))\n",
        "\n",
        "  else:\n",
        "    display(Markdown('### Please run the above SCF loop first.'))\n",
        "except NameError:\n",
        "    display(Markdown('### Please run the above SCF loop first.'))"
      ],
      "metadata": {
        "id": "p6dx0mfrkm0w"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Isodensity Surface Comparison**\n",
        "\n",
        "\n",
        "# using density_log as a proxy to see if SCF has been run\n",
        "assert len(density_log) != 0, 'Please run the above SCF loop first.'\n",
        "\n",
        "\n",
        "noninter_den     = density_log[0].reshape(nx,ny,nz)\n",
        "noninter_den_iso = noninter_den.mean()*2\n",
        "converged_den    = density_log[-1].reshape(nx,ny,nz)\n",
        "\n",
        "# Create a subplot with two 3D isosurfaces, vertically oriented\n",
        "fig = make_subplots(rows=2, cols=1, specs=[[{'type': 'isosurface'}], [{'type': 'isosurface'}]])\n",
        "\n",
        "fig.update_layout(width=800, height=800)\n",
        "\n",
        "# Data for the noninteracting density isosurface\n",
        "xp, yp, zp = grid_lines[0][0], grid_lines[0][1], grid_lines[0][2]\n",
        "X, Y, Z    = np.meshgrid(xp, yp, zp, indexing='ij')\n",
        "isosurface_1 = go.Isosurface(\n",
        "        x=X.flatten(),\n",
        "        y=Y.flatten(),\n",
        "        z=Z.flatten(),\n",
        "        value=noninter_den.flatten(),\n",
        "        colorscale='BlueRed',\n",
        "        isomin=-noninter_den_iso,\n",
        "        isomax=noninter_den_iso,\n",
        "        surface_count=2,\n",
        "        showscale=False,\n",
        "        caps=dict(x_show=False, y_show=False, z_show=False))\n",
        "\n",
        "# Data for the KS-optimized density isosurface\n",
        "isosurface_2 = go.Isosurface(\n",
        "        x=X.flatten(),\n",
        "        y=Y.flatten(),\n",
        "        z=Z.flatten(),\n",
        "        value=converged_den.flatten(),\n",
        "        colorscale='BlueRed',\n",
        "        isomin=-noninter_den_iso,\n",
        "        isomax=noninter_den_iso,\n",
        "        surface_count=2,\n",
        "        showscale=False,\n",
        "        caps=dict(x_show=False, y_show=False, z_show=False))\n",
        "\n",
        "# Add isosurfaces to the subplot\n",
        "fig.add_trace(isosurface_1, row=1, col=1)\n",
        "fig.add_trace(isosurface_2, row=2, col=1)\n",
        "\n",
        "# Update the layout and add subtitles\n",
        "fig.update_layout(\n",
        "    scene=dict(\n",
        "        xaxis_title='lx',\n",
        "        yaxis_title='ly',\n",
        "        zaxis_title='lz',\n",
        "    ),\n",
        "    scene2=dict(\n",
        "        xaxis_title='lx',\n",
        "        yaxis_title='ly',\n",
        "        zaxis_title='lz',\n",
        "    ),\n",
        "    annotations=[\n",
        "        dict(\n",
        "            text=f'Electron density @ <b>iteration 0</b> <br>isovalue: {noninter_den_iso:.5f} e/Bohr<sup>3<sup>',\n",
        "            x=0.5,\n",
        "            y=1.1,\n",
        "            xref='paper',\n",
        "            yref='paper',\n",
        "            showarrow=False,\n",
        "            font=dict(size=18),\n",
        "        ),\n",
        "        dict(\n",
        "            text=f'Electron density @ <b>iteration {len(density_log)-1}</b> <br>isovalue: {noninter_den_iso:.5f} e/Bohr<sup>3<sup>',\n",
        "            width=800,\n",
        "            height=400,\n",
        "            x=0.5,\n",
        "            y=0.47,\n",
        "            xref='paper',\n",
        "            yref='paper',\n",
        "            showarrow=False,\n",
        "            font=dict(size=18),\n",
        "        ),\n",
        "    ]\n",
        ")\n",
        "\n",
        "# Show the subplot\n",
        "fig.show()\n",
        "\n",
        "\n",
        "# ipyvolume version of the keep this commented until the library is fixed\n",
        "#print('\\033[1m==============================\\033[0m')\n",
        "#print(\"\\033[1m\\033[94mElectron density @ iteration 0\\033[0m\")\n",
        "#print('\\033[1m==============================\\033[0m')\n",
        "#print(f'isovalue:\\033[0m {isoval:.5f} e/Bohr**3')\n",
        "#ipv.clear()\n",
        "#fig = ipv.figure(width=500, height=500)\n",
        "#ipv.plot_isosurface(noninter_den,level=isoval,controls=True,\n",
        "#                    description='non-interacting density')\n",
        "#ipv.squarelim()\n",
        "#ipv.xyzlabel(r'lx','ly','lz')\n",
        "#ipv.style.box_off()\n",
        "#ipv.show()\n",
        "#display(Markdown('---'))\n",
        "#\n",
        "## KS-optimized density isosurface\n",
        "#print('\\033[1m==============================\\033[0m')\n",
        "#print(f\"\\033[1m\\033[94mElectron density @ iteration {len(density_log)-1}\\033[0m\")\n",
        "#print('\\033[1m==============================\\033[0m')\n",
        "#print(f'isovalue:\\033[0m {isoval:.5f} e/Bohr**3')\n",
        "#ipv.clear()\n",
        "#fig2 = ipv.figure(width=500, height=500)\n",
        "#ipv.plot_isosurface(converged_den,level=isoval,controls=True,\n",
        "#                    description='converged density')\n",
        "#ipv.squarelim()\n",
        "#ipv.xyzlabel(r'lx','ly','lz')\n",
        "#ipv.style.box_off()\n",
        "#ipv.show()\n"
      ],
      "metadata": {
        "id": "C8wKCyJIkrsw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title **Kohn-Sham Orbitals**\n",
        "\n",
        "# using density_log as a proxy to see if SCF has been run\n",
        "assert len(density_log) != 0, 'Please run the above SCF loop first.'\n",
        "\n",
        "\n",
        "# iteration number slider widget\n",
        "it_num_slider = widgets.IntSlider(\n",
        "    value=0,\n",
        "    min=0,\n",
        "    max=len(eigenvalue_log)-1,\n",
        "    step=1,\n",
        "    description='Iteration #: ',\n",
        "    disabled=False,\n",
        "    continuous_update=False,\n",
        "    orientation='horizontal',\n",
        "    readout=True,\n",
        "    readout_format='d'\n",
        ")\n",
        "\n",
        "# find the max orbital energy for the plot limits\n",
        "ener_array = np.array(eigenvalue_log)\n",
        "max_ener   = ener_array.max()\n",
        "min_ener   = ener_array.min()\n",
        "\n",
        "def eigenstate_analysis(index, show_fig=True, save_fig=False, filename=None):\n",
        "  '''Generates an interactive orbital energy diagram from above KS-DFT calculation.\n",
        "  Args:\n",
        "    index (int): Iteration number of SCF loop to consider.\n",
        "    show_fig (bool): Whether to display the figure.\n",
        "    save_fig (bool): Whether to save a .png file.\n",
        "    filename (str): Filename to save to (w/o) extension.\n",
        "  '''\n",
        "  clear_output(wait=True)\n",
        "\n",
        "  occ2 = len(eigenvalue_log[0])-2\n",
        "  y_DFT = eigenvalue_log[index]\n",
        "  HL_Gap = y_DFT[-2] - y_DFT[-3]\n",
        "  HL_Gap_eV = round(HL_Gap*27.2114, 3)\n",
        "  x_DFT = 1.5*np.ones(y_DFT.shape[0])\n",
        "  fig = plt.figure(figsize=(4, 6))\n",
        "  ax = fig.add_subplot(1, 1, 1)\n",
        "  plt.title('Orbital Energies')\n",
        "  plt.ylabel(\"Energy (Ha)\",labelpad=7)\n",
        "  plt.scatter(x_DFT[:-2], y_DFT[:-2], marker=0, s=1200, linewidths=4, color='green', label='occupied')\n",
        "  plt.scatter(x_DFT[-2:], y_DFT[-2:], marker=0, s=1200, linewidths=4, color='#F97306', label='virtual')\n",
        "  plt.legend(bbox_to_anchor=(1.3, 1))\n",
        "  plt.rcParams[\"legend.markerscale\"] = 0.35\n",
        "  plt.xticks([])\n",
        "  plt.xlim([-0.1, 3.3])\n",
        "  plt.ylim([min_ener - 0.03, max_ener + 0.03])\n",
        "  plt.text(1.07, 0.855, f'DFT iteration #{index}',\n",
        "        horizontalalignment='center', verticalalignment='center',\n",
        "        transform=ax.transAxes)\n",
        "  plt.text(1.07, 0.810, f'     Energy: {energy_log[index]:.5f} Ha',\n",
        "        horizontalalignment='center', verticalalignment='center',\n",
        "        transform=ax.transAxes)\n",
        "  plt.text(1.07, 0.765, f'      H-L Gap: {HL_Gap:.4f} Ha',\n",
        "      horizontalalignment='center', verticalalignment='center',\n",
        "      transform=ax.transAxes)\n",
        "  plt.text(1.07, 0.720, f'                       ({HL_Gap_eV:.3f} eV)',\n",
        "    horizontalalignment='center', verticalalignment='center',\n",
        "    transform=ax.transAxes)\n",
        "  annotationsDFT = ['HOMO-' + str(i) for i in range(occ2)]\n",
        "  annotationsDFT[0] = 'HOMO'\n",
        "  annotationsDFT.insert(0,'LUMO')\n",
        "  annotationsDFT.insert(0,'LUMO+1')\n",
        "  for i, label in enumerate(reversed(annotationsDFT)):\n",
        "      plt.annotate(label, (x_DFT[i] + 0.005, y_DFT[i]),size=8)\n",
        "      plt.text(x_DFT[i]-0.9, y_DFT[i], \"{:.3f}\".format(y_DFT[i]), size=8)\n",
        "  ax.spines['left'].set_position(('axes', .16))\n",
        "  ax.spines['top'].set_visible(False)\n",
        "  ax.spines['right'].set_visible(False)\n",
        "  ax.spines['bottom'].set_visible(False)\n",
        "\n",
        "  # control events\n",
        "  if save_fig:\n",
        "    plt.tight_layout()\n",
        "    plt.savefig(f'{filename}.png', dpi=800)\n",
        "  if show_fig:\n",
        "    plt.show()\n",
        "  else:\n",
        "    plt.close()\n",
        "\n",
        "def save_orbital_plot(b):\n",
        "  ''' Saves a .png of orbital energy diagram with widget parameters. '''\n",
        "  eigenstate_analysis(it_num_slider.value, show_fig=False, save_fig=True,\n",
        "                      filename=orbital_filename_text.value)\n",
        "\n",
        "orbital_save_button.on_click(save_orbital_plot) # link button + function\n",
        "\n",
        "# link widgets to function\n",
        "eigenstate_output = widgets.interactive_output(eigenstate_analysis, {'index':it_num_slider,\n",
        "                                                                  'show_fig':widgets.fixed(True),\n",
        "                                                                  'save_fig':widgets.fixed(False),\n",
        "                                                                  'filename':widgets.fixed(None)})\n",
        "\n",
        "# generate display\n",
        "display(it_num_slider,\n",
        "        eigenstate_output,\n",
        "        widgets.HBox([orbital_filename_text, orbital_save_button]))\n"
      ],
      "metadata": {
        "id": "yVsa6P8ak14e"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "## **References**\n",
        "\n",
        "* Hohenberg, P.; Kohn, W. Inhomogeneous Electron Gas. Phys. Rev. **1964**, *136*, B864-B871. DOI: [10.1103/PhysRev.136.B864](https://doi.org/10.1103/PhysRev.136.B864).\n",
        "\n",
        "*  Kohn, W.; Sham, L. J. Self-Consistent Equations Including Exchange and Correlation Effects. *Physical Review* **1965**, *140*, A1133-A1138. DOI: [10.1103/PhysRev.140.A1133](https://doi.org/10.1103/PhysRev.140.A1133).\n",
        "\n",
        "*  Baseden, K. A.; Tye, J. W. Introduction to Density Functional Theory: Calculations by Hand on the Helium Atom. *Journal of Chemical Education* **2014**, *91*, 2116–2123. DOI: [10.1021/ed5004788](https://doi.org/10.1021/ed5004788).\n",
        "\n",
        "*  Chachiyo, T. Communication: Simple and Accurate Uniform Electron Gas Correlation Energy for the Full Range of Densities. *The Journal of Chemical Physics* **2016**, *145*, 021101. DOI: [10.1063/1.4958669](https://doi.org/10.1063/1.4958669).\n",
        "\n",
        "* Halpern, A. M.; Ge, Y.; Glendening, E. D. Visualizing Solutions of the One-Dimensional Schrödinger Equation Using a Finite Difference Method. *Journal of Chemical Education* **2022**, *99*, 3053–3060. DOI: [10.1021/acs.jchemed.2c00557](https://doi.org/10.1021/acs.jchemed.2c00557).\n"
      ],
      "metadata": {
        "id": "CO8Dmgj2k6I6"
      }
    }
  ]
}