{
 "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": [],
   "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\n",
    "tol = 1e-12\n",
    "is_converged = DFTK.ScfConvergenceDensity(tol);"
   ],
   "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.846832249821                   -0.70    4.8\n",
      "  2   -7.852321060476       -2.26       -1.53    1.0\n",
      "  3   -7.852646061824       -3.49       -2.52    3.2\n",
      "  4   -7.852646678176       -6.21       -3.37    2.5\n",
      "  5   -7.852646686299       -8.09       -4.73    1.5\n",
      "  6   -7.852646686727       -9.37       -5.44    3.5\n",
      "  7   -7.852646686730      -11.54       -6.27    1.2\n",
      "  8   -7.852646686730      -13.01       -7.71    2.2\n",
      "  9   -7.852646686730      -14.57       -7.74    2.2\n",
      " 10   -7.852646686730      -15.05       -8.71    2.5\n",
      " 11   -7.852646686730   +    -Inf       -8.94    1.0\n",
      " 12   -7.852646686730   +  -15.05      -10.67    2.2\n",
      " 13   -7.852646686730   +    -Inf       -9.47    1.0\n",
      " 14   -7.852646686730      -15.05       -9.24    1.0\n",
      " 15   -7.852646686730      -14.75       -9.10    1.0\n",
      " 16   -7.852646686730   +  -14.35       -8.50    2.0\n",
      " 17   -7.852646686730   +  -15.05       -9.80    1.5\n",
      " 18   -7.852646686730      -14.75      -10.57    1.8\n",
      " 19   -7.852646686730      -14.75      -10.95    1.8\n",
      " 20   -7.852646686730   +  -15.05      -10.82    1.0\n",
      " 21   -7.852646686730   +  -15.05      -11.71    1.0\n",
      " 22   -7.852646686730      -14.75      -12.45    2.0\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "scfres_scf = self_consistent_field(basis; is_converged);"
   ],
   "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.846826529307                   -0.70           4.8\n",
      "  2   -7.852525112695       -2.24       -1.63   0.80    2.0\n",
      "  3   -7.852635705254       -3.96       -2.72   0.80    1.0\n",
      "  4   -7.852646446816       -4.97       -3.27   0.80    2.2\n",
      "  5   -7.852646672886       -6.65       -4.10   0.80    1.2\n",
      "  6   -7.852646686341       -7.87       -4.72   0.80    1.5\n",
      "  7   -7.852646686721       -9.42       -5.57   0.80    1.8\n",
      "  8   -7.852646686729      -11.08       -6.98   0.80    1.8\n",
      "  9   -7.852646686730      -12.41       -7.57   0.80    3.2\n",
      " 10   -7.852646686730      -14.75       -7.99   0.80    1.8\n",
      " 11   -7.852646686730      -15.05       -9.12   0.80    1.0\n",
      " 12   -7.852646686730      -14.75       -9.60   0.80    3.2\n",
      " 13   -7.852646686730   +  -15.05      -10.10   0.80    1.0\n",
      " 14   -7.852646686730   +  -15.05      -10.82   0.80    1.2\n",
      " 15   -7.852646686730   +    -Inf      -11.97   0.80    2.2\n",
      " 16   -7.852646686730   +    -Inf      -12.64   0.80    2.5\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "scfres_scfv = DFTK.scf_potential_mixing(basis; is_converged);"
   ],
   "metadata": {},
   "execution_count": 3
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Direct minimization"
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iter     Function value   Gradient norm \n",
      "     0     1.403903e+01     3.503994e+00\n",
      " * time: 0.44902586936950684\n",
      "     1     1.342235e+00     1.778317e+00\n",
      " * time: 0.651695966720581\n",
      "     2    -1.828507e+00     1.896481e+00\n",
      " * time: 0.6792130470275879\n",
      "     3    -3.754821e+00     1.820452e+00\n",
      " * time: 0.7940008640289307\n",
      "     4    -5.188438e+00     1.584808e+00\n",
      " * time: 0.8335838317871094\n",
      "     5    -6.918980e+00     8.684053e-01\n",
      " * time: 0.8725728988647461\n",
      "     6    -6.992047e+00     9.572685e-01\n",
      " * time: 0.899763822555542\n",
      "     7    -7.689852e+00     7.397375e-01\n",
      " * time: 0.9269559383392334\n",
      "     8    -7.759766e+00     8.800055e-01\n",
      " * time: 0.9543449878692627\n",
      "     9    -7.778897e+00     7.374077e-01\n",
      " * time: 0.9812648296356201\n",
      "    10    -7.785095e+00     1.073822e+00\n",
      " * time: 1.0081868171691895\n",
      "    11    -7.798981e+00     8.160733e-01\n",
      " * time: 1.0350289344787598\n",
      "    12    -7.819909e+00     4.539929e-01\n",
      " * time: 1.0741620063781738\n",
      "    13    -7.824708e+00     6.909547e-01\n",
      " * time: 1.1010818481445312\n",
      "    14    -7.843710e+00     2.596031e-01\n",
      " * time: 1.1279468536376953\n",
      "    15    -7.846046e+00     2.959001e-01\n",
      " * time: 1.15488600730896\n",
      "    16    -7.851371e+00     2.329684e-02\n",
      " * time: 1.1817529201507568\n",
      "    17    -7.852230e+00     8.761840e-03\n",
      " * time: 1.2086799144744873\n",
      "    18    -7.852569e+00     6.798428e-03\n",
      " * time: 1.23555588722229\n",
      "    19    -7.852604e+00     4.930048e-03\n",
      " * time: 1.262833833694458\n",
      "    20    -7.852625e+00     2.813465e-03\n",
      " * time: 1.289719820022583\n",
      "    21    -7.852641e+00     1.634270e-03\n",
      " * time: 1.3174378871917725\n",
      "    22    -7.852645e+00     9.289906e-04\n",
      " * time: 1.3442578315734863\n",
      "    23    -7.852646e+00     4.364427e-04\n",
      " * time: 1.3711888790130615\n",
      "    24    -7.852646e+00     3.232755e-04\n",
      " * time: 1.3981859683990479\n",
      "    25    -7.852647e+00     1.153466e-04\n",
      " * time: 1.4249730110168457\n",
      "    26    -7.852647e+00     6.613080e-05\n",
      " * time: 1.451746940612793\n",
      "    27    -7.852647e+00     4.115242e-05\n",
      " * time: 1.4787218570709229\n",
      "    28    -7.852647e+00     2.545155e-05\n",
      " * time: 1.505707025527954\n",
      "    29    -7.852647e+00     1.371649e-05\n",
      " * time: 1.5330719947814941\n",
      "    30    -7.852647e+00     8.255275e-06\n",
      " * time: 1.5598359107971191\n",
      "    31    -7.852647e+00     3.840623e-06\n",
      " * time: 1.586789846420288\n",
      "    32    -7.852647e+00     2.238605e-06\n",
      " * time: 1.6139118671417236\n",
      "    33    -7.852647e+00     1.239133e-06\n",
      " * time: 1.6414918899536133\n",
      "    34    -7.852647e+00     6.781613e-07\n",
      " * time: 1.6685888767242432\n",
      "    35    -7.852647e+00     3.939354e-07\n",
      " * time: 1.6955578327178955\n",
      "    36    -7.852647e+00     2.162569e-07\n",
      " * time: 1.7229039669036865\n",
      "    37    -7.852647e+00     1.289076e-07\n",
      " * time: 1.7498118877410889\n",
      "    38    -7.852647e+00     6.530468e-08\n",
      " * time: 1.7768030166625977\n",
      "    39    -7.852647e+00     3.634828e-08\n",
      " * time: 1.8040690422058105\n",
      "    40    -7.852647e+00     2.074358e-08\n",
      " * time: 1.831233024597168\n",
      "    41    -7.852647e+00     1.169768e-08\n",
      " * time: 1.858158826828003\n",
      "    42    -7.852647e+00     7.731757e-09\n",
      " * time: 1.8854949474334717\n",
      "    43    -7.852647e+00     4.208548e-09\n",
      " * time: 1.913322925567627\n",
      "    44    -7.852647e+00     4.207983e-09\n",
      " * time: 1.988473892211914\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "scfres_dm = direct_minimization(basis; tol);"
   ],
   "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.846843511052                   -0.70    4.5\n",
      "  2   -7.852322679588       -2.26       -1.53    1.0\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "scfres_start = self_consistent_field(basis; tol=1e-1);"
   ],
   "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.852646686713                   -2.54\n",
      "  2   -7.852646686730      -10.77       -6.03\n",
      "  3   -7.852646686730   +  -15.05      -12.64\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|  = 9.899854352494933e-14\n",
      "|ρ_newton - ρ_scfv| = 2.1712353697319453e-13\n",
      "|ρ_newton - ρ_dm|   = 2.420188489527844e-10\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.3"
  },
  "kernelspec": {
   "name": "julia-1.8",
   "display_name": "Julia 1.8.3",
   "language": "julia"
  }
 },
 "nbformat": 4
}