8_Direct_minimisation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "local-transmission",
   "metadata": {},
   "source": [
    "# Direct minimisation\n",
    "\n",
    "- Recall that we introduced DFT as the minimisation problem\n",
    "  $$ \\displaystyle \\text{min}_{\\{\\psi_i\\}} \\mathcal{E}_\\text{DFT}(\\{\\psi_i\\}) $$\n",
    "where the energy expression $\\mathcal{E}_\\text{DFT}(\\{\\psi_i\\})$ is a known analytic function and the $\\{\\psi_i\\}$ are orthonormal orbitals.\n",
    "\n",
    "\n",
    "- Taking a closer look at this minimisation problem one might rightfully wonder why one should bother going through the SCF procedure at all. Could one not just minimise the energy wrt. the orbitals?\n",
    "- The answer is yes and leads to a procedure called **direct minimisation** (DM). In practice DM provides an alternative route to access a DFT ground state."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "alien-microphone",
   "metadata": {},
   "source": [
    "While the basic principle of a DM algorithms is relatively straightforward\n",
    "a practical implementation faces a few subtle points that require some thought:\n",
    "\n",
    "- Due to the orthogonality constraints on the orbitals $\\{\\psi_i\\}$,\n",
    "  the set of admissible orbitals does not form a vector space,\n",
    "  but much rather the unknowns of $\\{\\psi_i\\}$ belong to a Stiefel manifold.\n",
    "  This needs to be taken into account\n",
    "  (via appropriate projections) in order to get the correct minimum.\n",
    "  Fortunately `Optim.jl` can do that out of the box ...\n",
    "\n",
    "- DM in combination with metals (i.e. without a band gap) is tricky,\n",
    "  because of possible degeneracies of the orbitals at the Fermi level ($i = N$).\n",
    "  This problem manifests when computing the gradient of $\\mathcal{E}_\\text{DFT}$\n",
    "  in such a setting, where thus special care is needed to not run into\n",
    "  \"division by zero\" issues.\n",
    "  \n",
    "- In this workbook we will restrict ourseves to insulators,\n",
    "  which have a band gap and thus don't feature this numerical problem.\n",
    "  Furthermore we will restrict ourselves to a single $k$-Point\n",
    "  as going beyond that requires a little more bookkeeping\n",
    "  (see the [DFTK implementation](https://github.com/JuliaMolSim/DFTK.jl/blob/master/src/scf/direct_minimization.jl)).\n",
    "  \n",
    "All right, now let's consider computing the gradient.\n",
    "Since the total energy can also be written as\n",
    "$$ \\mathcal{E}_\\text{DFT}(\\{\\psi_i\\}) = 2 \\sum_i \\int \\psi_i^\\ast H \\psi_i + E_\\text{nuclear} $$\n",
    "(where $E_\\text{nuclear}$ is a constant), the **gradient of $\\mathcal{E}_\\text{DFT}$**\n",
    "can be represented as $4 H \\psi_i $.\n",
    "\n",
    "We can proceed to a simple implementation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "instrumental-causing",
   "metadata": {},
   "outputs": [],
   "source": [
    "using DFTK\n",
    "using Optim\n",
    "using LineSearches\n",
    "\n",
    "# Standard silicon setup\n",
    "a = 10.26\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 => [ones(3)/8, -ones(3)/8]]\n",
    "model  = model_DFT(lattice, atoms, [:lda_x, :lda_c_vwn])\n",
    "basis  = PlaneWaveBasis(model; Ecut=10, kgrid=[1, 1, 1]);\n",
    "\n",
    "# One unit cell has 2 Silicon atoms.\n",
    "# In the model we use (where only valence electrons are explictly treated)\n",
    "# this makes 8 electrons, which requires 4 bands with 2 electrons each:\n",
    "occupation = [2.0, 2.0, 2.0, 2.0]\n",
    "\n",
    "# We specify a random initial guess for the 4 orbitals:\n",
    "n_G = length(G_vectors(only(basis.kpoints)))\n",
    "ψ0 = Matrix(qr(randn(ComplexF64, n_G, 4)).Q);\n",
    "\n",
    "# Function to compute energies and gradients\n",
    "function fg!(E, G, ψ)\n",
    "    ρ = compute_density(basis, [ψ], [occupation])\n",
    "    energies, H = energy_hamiltonian(basis, [ψ], [occupation]; ρ=ρ)\n",
    "\n",
    "    if G !== nothing\n",
    "        # Optim expects the gradient in G\n",
    "        G .= 4 * (H.blocks[1] * ψ)\n",
    "    end\n",
    "    energies.total\n",
    "end\n",
    "\n",
    "# Select a quasi-Newton algorithm with backtracking linesearches\n",
    "# to avoid to many costly gradient evaluations\n",
    "algorithm = Optim.LBFGS(manifold=Optim.Stiefel(),\n",
    "                        linesearch=LineSearches.BackTracking())\n",
    "\n",
    "# Set some convergence options in Optim:\n",
    "options = Optim.Options(; allow_f_increases=true, show_trace=true, x_tol=1e-6)\n",
    "\n",
    "# Run the direct minimisation\n",
    "res = Optim.optimize(Optim.only_fg!(fg!), ψ0, algorithm, options)\n",
    "\n",
    "@show res.minimum"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "thrown-highland",
   "metadata": {},
   "source": [
    "#### More details\n",
    "- [Geometry of algorithms with orthogonality constraints](https://doi.org/10.1137/S0895479895290954)\n",
    "- [Convergence analysis of direct minimisation and self-consistent field iterations](https://doi.org/10.1137/20M1332864)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.6.1",
   "language": "julia",
   "name": "julia-1.6"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}