{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Comparison of DFT solvers"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "We compare four different approaches for solving the DFT minimisation problem,\n",
    "namely a density-based SCF, a potential-based SCF, direct minimisation and Newton."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "First we setup our problem"
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "1.0e-6"
     },
     "metadata": {},
     "execution_count": 1
    }
   ],
   "cell_type": "code",
   "source": [
    "using DFTK\n",
    "using LinearAlgebra\n",
    "\n",
    "a = 10.26  # Silicon lattice constant in Bohr\n",
    "lattice = a / 2 * [[0 1 1.];\n",
    "                   [1 0 1.];\n",
    "                   [1 1 0.]]\n",
    "Si = ElementPsp(:Si, psp=load_psp(\"hgh/lda/Si-q4\"))\n",
    "atoms     = [Si, Si]\n",
    "positions = [ones(3)/8, -ones(3)/8]\n",
    "\n",
    "model = model_LDA(lattice, atoms, positions)\n",
    "basis = PlaneWaveBasis(model; Ecut=5, kgrid=[3, 3, 3])\n",
    "\n",
    "# Convergence we desire in the density\n",
    "tol = 1e-6"
   ],
   "metadata": {},
   "execution_count": 1
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Density-based self-consistent field"
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n     Energy            log10(ΔE)   log10(Δρ)   Diag\n",
      "---   ---------------   ---------   ---------   ----\n",
      "  1   -7.846773131013                   -0.70    4.5\n",
      "  2   -7.852308137955       -2.26       -1.53    1.0\n",
      "  3   -7.852614082992       -3.51       -2.55    1.5\n",
      "  4   -7.852645885494       -4.50       -2.87    2.8\n",
      "  5   -7.852646493085       -6.22       -3.18    1.0\n",
      "  6   -7.852646678669       -6.73       -4.01    1.0\n",
      "  7   -7.852646686573       -8.10       -5.02    1.8\n",
      "  8   -7.852646686726       -9.82       -5.57    2.0\n",
      "  9   -7.852646686728      -11.60       -5.70    1.5\n",
      " 10   -7.852646686730      -11.82       -6.85    1.0\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "scfres_scf = self_consistent_field(basis; tol);"
   ],
   "metadata": {},
   "execution_count": 2
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Potential-based SCF"
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n     Energy            log10(ΔE)   log10(Δρ)   α      Diag\n",
      "---   ---------------   ---------   ---------   ----   ----\n",
      "  1   -7.846864266995                   -0.70           4.5\n",
      "  2   -7.852526326642       -2.25       -1.64   0.80    2.0\n",
      "  3   -7.852636361040       -3.96       -2.73   0.80    1.0\n",
      "  4   -7.852646544511       -4.99       -3.26   0.80    2.2\n",
      "  5   -7.852646674713       -6.89       -4.15   0.80    1.0\n",
      "  6   -7.852646686465       -7.93       -4.77   0.80    1.5\n",
      "  7   -7.852646686715       -9.60       -5.61   0.80    1.5\n",
      "  8   -7.852646686729      -10.85       -7.04   0.80    2.2\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "scfres_scfv = DFTK.scf_potential_mixing(basis; tol);"
   ],
   "metadata": {},
   "execution_count": 3
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Direct minimization\n",
    "Note: Unlike the other algorithms, tolerance for this one is in the energy,\n",
    "thus we square the density tolerance value to be roughly equivalent."
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iter     Function value   Gradient norm \n",
      "     0     1.378900e+01     3.227591e+00\n",
      " * time: 0.42667388916015625\n",
      "     1     1.278543e+00     1.644298e+00\n",
      " * time: 0.6292400360107422\n",
      "     2    -1.877766e+00     1.892585e+00\n",
      " * time: 0.653986930847168\n",
      "     3    -3.780167e+00     1.658428e+00\n",
      " * time: 0.6890289783477783\n",
      "     4    -5.233362e+00     1.544238e+00\n",
      " * time: 0.7242648601531982\n",
      "     5    -6.921043e+00     1.014799e+00\n",
      " * time: 0.759943962097168\n",
      "     6    -7.069459e+00     1.266764e+00\n",
      " * time: 0.7841019630432129\n",
      "     7    -7.583674e+00     8.041395e-01\n",
      " * time: 0.8088908195495605\n",
      "     8    -7.684095e+00     6.327646e-01\n",
      " * time: 0.8331918716430664\n",
      "     9    -7.767708e+00     1.731340e-01\n",
      " * time: 0.8577549457550049\n",
      "    10    -7.800157e+00     1.143894e-01\n",
      " * time: 0.8820569515228271\n",
      "    11    -7.821996e+00     9.443334e-02\n",
      " * time: 0.9064679145812988\n",
      "    12    -7.837750e+00     6.033945e-02\n",
      " * time: 0.93076491355896\n",
      "    13    -7.848451e+00     3.127138e-02\n",
      " * time: 0.9552218914031982\n",
      "    14    -7.850966e+00     1.596878e-02\n",
      " * time: 0.979496955871582\n",
      "    15    -7.851912e+00     1.613052e-02\n",
      " * time: 1.0040638446807861\n",
      "    16    -7.852331e+00     7.899294e-03\n",
      " * time: 1.0284030437469482\n",
      "    17    -7.852554e+00     5.478427e-03\n",
      " * time: 1.0526719093322754\n",
      "    18    -7.852623e+00     2.638148e-03\n",
      " * time: 1.0768299102783203\n",
      "    19    -7.852639e+00     1.768830e-03\n",
      " * time: 1.1010689735412598\n",
      "    20    -7.852644e+00     1.082814e-03\n",
      " * time: 1.1251659393310547\n",
      "    21    -7.852645e+00     7.240267e-04\n",
      " * time: 1.1493990421295166\n",
      "    22    -7.852646e+00     4.168570e-04\n",
      " * time: 1.17364501953125\n",
      "    23    -7.852646e+00     2.096628e-04\n",
      " * time: 1.1980388164520264\n",
      "    24    -7.852647e+00     1.377217e-04\n",
      " * time: 1.2223830223083496\n",
      "    25    -7.852647e+00     5.699591e-05\n",
      " * time: 1.2468678951263428\n",
      "    26    -7.852647e+00     2.823652e-05\n",
      " * time: 1.2714238166809082\n",
      "    27    -7.852647e+00     2.061149e-05\n",
      " * time: 1.295900821685791\n",
      "    28    -7.852647e+00     1.266611e-05\n",
      " * time: 1.3199939727783203\n",
      "    29    -7.852647e+00     8.176777e-06\n",
      " * time: 1.3451869487762451\n",
      "    30    -7.852647e+00     5.212700e-06\n",
      " * time: 1.3694109916687012\n",
      "    31    -7.852647e+00     2.245190e-06\n",
      " * time: 1.3938159942626953\n",
      "    32    -7.852647e+00     1.355626e-06\n",
      " * time: 1.418238878250122\n",
      "    33    -7.852647e+00     8.394806e-07\n",
      " * time: 1.4426288604736328\n",
      "    34    -7.852647e+00     4.435495e-07\n",
      " * time: 1.4671118259429932\n",
      "    35    -7.852647e+00     2.446567e-07\n",
      " * time: 1.5639889240264893\n",
      "    36    -7.852647e+00     1.597361e-07\n",
      " * time: 1.5888140201568604\n",
      "    37    -7.852647e+00     1.063875e-07\n",
      " * time: 1.613497018814087\n",
      "    38    -7.852647e+00     7.377490e-08\n",
      " * time: 1.6378049850463867\n",
      "    39    -7.852647e+00     4.399803e-08\n",
      " * time: 1.6623530387878418\n",
      "    40    -7.852647e+00     2.336608e-08\n",
      " * time: 1.686720848083496\n",
      "    41    -7.852647e+00     1.536328e-08\n",
      " * time: 1.7115130424499512\n",
      "    42    -7.852647e+00     8.031708e-09\n",
      " * time: 1.7361149787902832\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "scfres_dm = direct_minimization(basis; tol=tol^2);"
   ],
   "metadata": {},
   "execution_count": 4
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Newton algorithm"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Start not too far from the solution to ensure convergence:\n",
    "We run first a very crude SCF to get close and then switch to Newton."
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n     Energy            log10(ΔE)   log10(Δρ)   Diag\n",
      "---   ---------------   ---------   ---------   ----\n",
      "  1   -7.846636964707                   -0.70    4.5\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "scfres_start = self_consistent_field(basis; tol=0.5);"
   ],
   "metadata": {},
   "execution_count": 5
  },
  {
   "cell_type": "markdown",
   "source": [
    "Remove the virtual orbitals (which Newton cannot treat yet)"
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "n     Energy            log10(ΔE)   log10(Δρ)\n",
      "---   ---------------   ---------   ---------\n",
      "  1   -7.852645931821                   -1.64\n",
      "  2   -7.852646686730       -6.12       -3.72\n",
      "  3   -7.852646686730      -13.36       -7.27\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "ψ = DFTK.select_occupied_orbitals(basis, scfres_start.ψ, scfres_start.occupation).ψ\n",
    "scfres_newton = newton(basis, ψ; tol);"
   ],
   "metadata": {},
   "execution_count": 6
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Comparison of results"
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|ρ_newton - ρ_scf|  = 2.2919382795221063e-7\n",
      "|ρ_newton - ρ_scfv| = 2.8835937239523406e-7\n",
      "|ρ_newton - ρ_dm|   = 1.2974206445550981e-9\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "println(\"|ρ_newton - ρ_scf|  = \", norm(scfres_newton.ρ - scfres_scf.ρ))\n",
    "println(\"|ρ_newton - ρ_scfv| = \", norm(scfres_newton.ρ - scfres_scfv.ρ))\n",
    "println(\"|ρ_newton - ρ_dm|   = \", norm(scfres_newton.ρ - scfres_dm.ρ))"
   ],
   "metadata": {},
   "execution_count": 7
  }
 ],
 "nbformat_minor": 3,
 "metadata": {
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.8.4"
  },
  "kernelspec": {
   "name": "julia-1.8",
   "display_name": "Julia 1.8.4",
   "language": "julia"
  }
 },
 "nbformat": 4
}