{
 "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   Δtime\n",
      "---   ---------------   ---------   ---------   ----   ------\n",
      "  1   -7.846818949646                   -0.70    4.5         \n",
      "  2   -7.852317784126       -2.26       -1.53    1.0   36.6ms\n",
      "  3   -7.852617645060       -3.52       -2.56    1.5   40.2ms\n",
      "  4   -7.852645858657       -4.55       -2.89    2.5   47.8ms\n",
      "  5   -7.852646504983       -6.19       -3.21    1.2    170ms\n",
      "  6   -7.852646677529       -6.76       -4.23    1.0   35.2ms\n",
      "  7   -7.852646686666       -8.04       -4.95    2.2   49.7ms\n",
      "  8   -7.852646686723      -10.24       -5.44    1.2   38.2ms\n",
      "  9   -7.852646686729      -11.28       -5.71    1.2   37.8ms\n",
      " 10   -7.852646686730      -11.96       -6.60    1.0   37.8ms\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   Δtime\n",
      "---   ---------------   ---------   ---------   ----   ----   ------\n",
      "  1   -7.846925223248                   -0.70           4.8         \n",
      "  2   -7.852529038765       -2.25       -1.63   0.80    2.0    346ms\n",
      "  3   -7.852636784512       -3.97       -2.72   0.80    1.0   32.2ms\n",
      "  4   -7.852646502333       -5.01       -3.29   0.80    2.2   43.6ms\n",
      "  5   -7.852646682143       -6.75       -4.13   0.80    1.2   34.4ms\n",
      "  6   -7.852646686366       -8.37       -4.83   0.80    1.2   33.7ms\n",
      "  7   -7.852646686722       -9.45       -5.89   0.80    2.0   42.1ms\n",
      "  8   -7.852646686730      -11.13       -6.67   0.80    2.0   41.5ms\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.381312e+01     3.058315e+00\n",
      " * time: 0.5095341205596924\n",
      "     1     9.712065e-01     1.717580e+00\n",
      " * time: 0.7737569808959961\n",
      "     2    -1.935801e+00     1.916885e+00\n",
      " * time: 0.8076651096343994\n",
      "     3    -3.728783e+00     1.842075e+00\n",
      " * time: 0.8577439785003662\n",
      "     4    -5.139869e+00     1.920588e+00\n",
      " * time: 0.9129831790924072\n",
      "     5    -6.710954e+00     1.367151e+00\n",
      " * time: 0.9677610397338867\n",
      "     6    -7.452315e+00     6.212700e-01\n",
      " * time: 1.0199639797210693\n",
      "     7    -7.676187e+00     4.592504e-01\n",
      " * time: 1.0547881126403809\n",
      "     8    -7.753233e+00     2.034475e-01\n",
      " * time: 1.0879709720611572\n",
      "     9    -7.788749e+00     7.622030e-02\n",
      " * time: 1.1221320629119873\n",
      "    10    -7.815776e+00     6.348562e-02\n",
      " * time: 1.1550359725952148\n",
      "    11    -7.833394e+00     5.858491e-02\n",
      " * time: 1.1883571147918701\n",
      "    12    -7.842064e+00     4.015647e-02\n",
      " * time: 1.2197380065917969\n",
      "    13    -7.850047e+00     1.962371e-02\n",
      " * time: 1.253767967224121\n",
      "    14    -7.851678e+00     1.101054e-02\n",
      " * time: 1.2899491786956787\n",
      "    15    -7.852405e+00     6.464870e-03\n",
      " * time: 1.3239161968231201\n",
      "    16    -7.852580e+00     5.886294e-03\n",
      " * time: 1.3566551208496094\n",
      "    17    -7.852629e+00     1.622977e-03\n",
      " * time: 1.390160083770752\n",
      "    18    -7.852642e+00     1.000632e-03\n",
      " * time: 1.4240131378173828\n",
      "    19    -7.852645e+00     4.842165e-04\n",
      " * time: 1.4560701847076416\n",
      "    20    -7.852646e+00     3.415662e-04\n",
      " * time: 1.4911260604858398\n",
      "    21    -7.852646e+00     2.065074e-04\n",
      " * time: 1.5242431163787842\n",
      "    22    -7.852647e+00     1.141188e-04\n",
      " * time: 1.5592751502990723\n",
      "    23    -7.852647e+00     6.159923e-05\n",
      " * time: 1.593641996383667\n",
      "    24    -7.852647e+00     5.006607e-05\n",
      " * time: 1.6258301734924316\n",
      "    25    -7.852647e+00     2.300981e-05\n",
      " * time: 1.657628059387207\n",
      "    26    -7.852647e+00     1.126302e-05\n",
      " * time: 1.6911160945892334\n",
      "    27    -7.852647e+00     5.934221e-06\n",
      " * time: 1.724276065826416\n",
      "    28    -7.852647e+00     3.654801e-06\n",
      " * time: 1.7565810680389404\n",
      "    29    -7.852647e+00     2.198017e-06\n",
      " * time: 1.7894861698150635\n",
      "    30    -7.852647e+00     1.361253e-06\n",
      " * time: 1.8241140842437744\n",
      "    31    -7.852647e+00     6.561026e-07\n",
      " * time: 1.8576619625091553\n",
      "    32    -7.852647e+00     3.630724e-07\n",
      " * time: 1.891535997390747\n",
      "    33    -7.852647e+00     1.658644e-07\n",
      " * time: 1.927828073501587\n",
      "    34    -7.852647e+00     1.058262e-07\n",
      " * time: 1.9621570110321045\n",
      "    35    -7.852647e+00     5.879169e-08\n",
      " * time: 1.9952220916748047\n",
      "    36    -7.852647e+00     4.118221e-08\n",
      " * time: 2.033506155014038\n",
      "    37    -7.852647e+00     2.182735e-08\n",
      " * time: 2.0687620639801025\n",
      "    38    -7.852647e+00     1.155588e-08\n",
      " * time: 2.1025850772857666\n",
      "    39    -7.852647e+00     9.669505e-09\n",
      " * time: 2.1523780822753906\n",
      "    40    -7.852647e+00     6.365953e-09\n",
      " * time: 2.185908079147339\n",
      "    41    -7.852647e+00     6.365953e-09\n",
      " * time: 2.368324041366577\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   Δtime\n",
      "---   ---------------   ---------   ---------   ----   ------\n",
      "  1   -7.846757130857                   -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(Δρ)   Δtime\n",
      "---   ---------------   ---------   ---------   ------\n",
      "  1   -7.852645641253                   -1.64         \n",
      "  2   -7.852646686730       -5.98       -3.68    2.26s\n",
      "  3   -7.852646686730      -12.84       -7.06    328ms\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.6589651775334955e-7\n",
      "|ρ_newton - ρ_scfv| = 8.547977389425358e-8\n",
      "|ρ_newton - ρ_dm|   = 1.0808370690073526e-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.5"
  },
  "kernelspec": {
   "name": "julia-1.8",
   "display_name": "Julia 1.8.5",
   "language": "julia"
  }
 },
 "nbformat": 4
}