{
 "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.846813780799                   -0.70    4.5         \n",
      "  2   -7.852318691129       -2.26       -1.53    1.0   28.9ms\n",
      "  3   -7.852615622982       -3.53       -2.55    1.5   31.8ms\n",
      "  4   -7.852645998993       -4.52       -2.89    2.8   40.5ms\n",
      "  5   -7.852646527538       -6.28       -3.22    1.2   30.8ms\n",
      "  6   -7.852646680373       -6.82       -4.09    1.0   29.4ms\n",
      "  7   -7.852646686677       -8.20       -4.96    2.0   37.4ms\n",
      "  8   -7.852646686725      -10.32       -5.62    1.2   30.7ms\n",
      "  9   -7.852646686728      -11.53       -5.63    2.0   36.1ms\n",
      " 10   -7.852646686730      -11.76       -7.58    1.0   29.6ms\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.846834668358                   -0.70           4.8         \n",
      "  2   -7.852525624604       -2.24       -1.64   0.80    2.0    308ms\n",
      "  3   -7.852635591496       -3.96       -2.73   0.80    1.0   26.2ms\n",
      "  4   -7.852646468538       -4.96       -3.23   0.80    2.2   35.8ms\n",
      "  5   -7.852646670118       -6.70       -4.05   0.80    1.0   25.9ms\n",
      "  6   -7.852646686441       -7.79       -4.72   0.80    1.8   31.8ms\n",
      "  7   -7.852646686724       -9.55       -5.60   0.80    1.5   29.6ms\n",
      "  8   -7.852646686730      -11.28       -6.98   0.80    2.0   32.1ms\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.374359e+01     3.256674e+00\n",
      " * time: 0.46646595001220703\n",
      "     1     1.730458e+00     1.790498e+00\n",
      " * time: 0.6952509880065918\n",
      "     2    -1.335246e+00     1.951220e+00\n",
      " * time: 0.7235841751098633\n",
      "     3    -3.760175e+00     1.627614e+00\n",
      " * time: 0.7639119625091553\n",
      "     4    -5.251124e+00     1.598272e+00\n",
      " * time: 0.8038010597229004\n",
      "     5    -6.881863e+00     1.087604e+00\n",
      " * time: 0.8435521125793457\n",
      "     6    -7.186370e+00     1.148941e+00\n",
      " * time: 0.8716800212860107\n",
      "     7    -7.630205e+00     4.551217e-01\n",
      " * time: 0.8998050689697266\n",
      "     8    -7.744011e+00     1.423004e-01\n",
      " * time: 0.9278621673583984\n",
      "     9    -7.784470e+00     1.401288e-01\n",
      " * time: 0.9559991359710693\n",
      "    10    -7.810615e+00     1.166473e-01\n",
      " * time: 0.9841580390930176\n",
      "    11    -7.830250e+00     6.502726e-02\n",
      " * time: 1.0125031471252441\n",
      "    12    -7.841569e+00     6.423847e-02\n",
      " * time: 1.0405981540679932\n",
      "    13    -7.844808e+00     4.926604e-02\n",
      " * time: 1.0687329769134521\n",
      "    14    -7.845371e+00     6.485087e-02\n",
      " * time: 1.0972959995269775\n",
      "    15    -7.845371e+00     6.485087e-02\n",
      " * time: 1.5219321250915527\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.846815988933                   -0.70    4.8         \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.852645893540                   -1.64         \n",
      "  2   -7.852646686730       -6.10       -3.70    1.83s\n",
      "  3   -7.852646686730      -13.27       -7.24    193ms\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.2494199662991404e-7\n",
      "|ρ_newton - ρ_scfv| = 1.9158817774244348e-7\n",
      "|ρ_newton - ρ_dm|   = 0.019041239112697225\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
}