{
 "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.846872559925                   -0.70    5.0         \n",
      "  2   -7.852329074596       -2.26       -1.53    1.0   22.6ms\n",
      "  3   -7.852613512883       -3.55       -2.56    1.2   23.9ms\n",
      "  4   -7.852645956388       -4.49       -2.87    2.8   31.2ms\n",
      "  5   -7.852646464554       -6.29       -3.12    1.0   23.2ms\n",
      "  6   -7.852646676639       -6.67       -3.94    1.0   23.0ms\n",
      "  7   -7.852646686168       -8.02       -5.06    1.5   25.3ms\n",
      "  8   -7.852646686721       -9.26       -5.37    2.5   31.6ms\n",
      "  9   -7.852646686729      -11.09       -6.00    1.0   65.5ms\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.846862056955                   -0.70           4.8         \n",
      "  2   -7.852524744014       -2.25       -1.64   0.80    2.2    258ms\n",
      "  3   -7.852635224137       -3.96       -2.73   0.80    1.0   21.3ms\n",
      "  4   -7.852646499059       -4.95       -3.23   0.80    2.2   29.5ms\n",
      "  5   -7.852646673992       -6.76       -4.05   0.80    1.2   22.4ms\n",
      "  6   -7.852646686371       -7.91       -4.81   0.80    1.8   26.1ms\n",
      "  7   -7.852646686724       -9.45       -5.57   0.80    1.8   26.2ms\n",
      "  8   -7.852646686730      -11.24       -7.03   0.80    2.0   28.3ms\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.376673e+01     3.446661e+00\n",
      " * time: 0.05784797668457031\n",
      "     1     1.035668e+00     1.764910e+00\n",
      " * time: 0.27242517471313477\n",
      "     2    -1.780741e+00     2.153443e+00\n",
      " * time: 0.2955482006072998\n",
      "     3    -3.884170e+00     1.924643e+00\n",
      " * time: 0.32851099967956543\n",
      "     4    -5.096888e+00     1.758705e+00\n",
      " * time: 0.3610820770263672\n",
      "     5    -6.838055e+00     9.559709e-01\n",
      " * time: 0.3939549922943115\n",
      "     6    -7.465364e+00     5.330211e-01\n",
      " * time: 0.4918639659881592\n",
      "     7    -7.681976e+00     3.708038e-01\n",
      " * time: 0.5148661136627197\n",
      "     8    -7.757692e+00     1.084613e-01\n",
      " * time: 0.5376181602478027\n",
      "     9    -7.806085e+00     1.598147e-01\n",
      " * time: 0.5601329803466797\n",
      "    10    -7.829592e+00     6.875602e-02\n",
      " * time: 0.5826570987701416\n",
      "    11    -7.841884e+00     6.256555e-02\n",
      " * time: 0.6049520969390869\n",
      "    12    -7.846960e+00     5.524165e-02\n",
      " * time: 0.6271171569824219\n",
      "    13    -7.848877e+00     4.607307e-02\n",
      " * time: 0.6494541168212891\n",
      "    14    -7.850422e+00     2.534032e-02\n",
      " * time: 0.6719121932983398\n",
      "    15    -7.851783e+00     1.470556e-02\n",
      " * time: 0.6945559978485107\n",
      "    16    -7.852414e+00     7.602897e-03\n",
      " * time: 0.7171990871429443\n",
      "    17    -7.852586e+00     5.228225e-03\n",
      " * time: 0.7398731708526611\n",
      "    18    -7.852628e+00     2.276126e-03\n",
      " * time: 0.7624790668487549\n",
      "    19    -7.852641e+00     9.418454e-04\n",
      " * time: 0.7851331233978271\n",
      "    20    -7.852645e+00     6.488569e-04\n",
      " * time: 0.8080360889434814\n",
      "    21    -7.852646e+00     3.455321e-04\n",
      " * time: 0.830517053604126\n",
      "    22    -7.852646e+00     2.399337e-04\n",
      " * time: 0.853262186050415\n",
      "    23    -7.852647e+00     1.632441e-04\n",
      " * time: 0.8759369850158691\n",
      "    24    -7.852647e+00     1.049781e-04\n",
      " * time: 0.898604154586792\n",
      "    25    -7.852647e+00     4.875373e-05\n",
      " * time: 0.9212710857391357\n",
      "    26    -7.852647e+00     2.425934e-05\n",
      " * time: 0.9436380863189697\n",
      "    27    -7.852647e+00     1.430669e-05\n",
      " * time: 0.9660019874572754\n",
      "    28    -7.852647e+00     8.628782e-06\n",
      " * time: 0.9886350631713867\n",
      "    29    -7.852647e+00     5.037319e-06\n",
      " * time: 1.0112099647521973\n",
      "    30    -7.852647e+00     3.506866e-06\n",
      " * time: 1.0340790748596191\n",
      "    31    -7.852647e+00     1.753959e-06\n",
      " * time: 1.0567450523376465\n",
      "    32    -7.852647e+00     1.271976e-06\n",
      " * time: 1.0790660381317139\n",
      "    33    -7.852647e+00     5.539523e-07\n",
      " * time: 1.1015121936798096\n",
      "    34    -7.852647e+00     3.544319e-07\n",
      " * time: 1.1238820552825928\n",
      "    35    -7.852647e+00     2.011782e-07\n",
      " * time: 1.14650297164917\n",
      "    36    -7.852647e+00     1.115273e-07\n",
      " * time: 1.1691629886627197\n",
      "    37    -7.852647e+00     8.662922e-08\n",
      " * time: 1.191802978515625\n",
      "    38    -7.852647e+00     3.329043e-08\n",
      " * time: 1.2141931056976318\n",
      "    39    -7.852647e+00     2.182318e-08\n",
      " * time: 1.2366251945495605\n",
      "    40    -7.852647e+00     1.549553e-08\n",
      " * time: 1.2692999839782715\n",
      "    41    -7.852647e+00     6.963544e-09\n",
      " * time: 1.2917990684509277\n",
      "    42    -7.852647e+00     6.493402e-09\n",
      " * time: 1.3149960041046143\n",
      "    43    -7.852647e+00     2.580704e-09\n",
      " * time: 1.337494134902954\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.846853629199                   -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.852645914146                   -1.64         \n",
      "  2   -7.852646686730       -6.11       -3.71    1.90s\n",
      "  3   -7.852646686730      -13.29       -7.25    140ms\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|  = 8.327201480445901e-7\n",
      "|ρ_newton - ρ_scfv| = 3.5148077303630417e-7\n",
      "|ρ_newton - ρ_dm|   = 7.169235209148582e-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.9.0"
  },
  "kernelspec": {
   "name": "julia-1.9",
   "display_name": "Julia 1.9.0",
   "language": "julia"
  }
 },
 "nbformat": 4
}