7_Properties_automatic_differentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "sufficient-southwest",
   "metadata": {},
   "source": [
    "# Computing properties and automatic differentiation\n",
    "\n",
    "Even though we have now spent about quite some notebooks discussing the SCF, its convergence and errors, the resulting quantity --- the DFT total energy --- is not very interesting."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "revolutionary-membership",
   "metadata": {},
   "source": [
    "- From a physical point of view the total energy is arbitrary and can always be shifted around.\n",
    "- What is of great interest, however, are **differences** or **changes** in the energy\n",
    "  in response to a perturbation.\n",
    "- In fact such responses often represent quantities, which are directly measurable\n",
    "  in experiment (or are at least they are closely linked to measurements).\n",
    "- A few examples:\n",
    "  * *dipole moment* (response of the energy to a change in electric field)\n",
    "  * *forces* (response to a change in atomic positions)\n",
    "  * ...\n",
    "\n",
    "Notice that \"response\" is just a different term for \"taking an energy derivative\". So that's a good opportunity to try some automatic differentation tools in combination with DFTK."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "immune-mills",
   "metadata": {},
   "source": [
    "## Computing forces\n",
    "\n",
    "Given a crystal with atomic positions $x$ and DFT ground state energy $E_\\text{SCF}$ (implicitly depending on $x$ via the external potential $V_\\text{ext}$), the force is defined as\n",
    "\n",
    "$$ \\text{Force} = - \\frac{d E_\\text{SCF}}{d x} $$\n",
    "\n",
    "For a silicon crystal one could define the DFT ground state in terms of a shift vector ($x$) between\n",
    "both silicon atoms as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "seventh-exchange",
   "metadata": {},
   "outputs": [],
   "source": [
    "using DFTK\n",
    "\n",
    "function compute_silicon(x)\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 + x]]\n",
    "    \n",
    "    T = eltype(x)\n",
    "    model  = model_DFT(Array{T}(lattice), atoms,[:lda_x, :lda_c_vwn]; symmetries=false)\n",
    "    basis  = PlaneWaveBasis(model; Ecut=13, kgrid=[2, 2, 2]);\n",
    "    \n",
    "    self_consistent_field(basis; tol=1e-14, callback=identity)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cellular-airplane",
   "metadata": {},
   "source": [
    "With this function we compute the silicon force in DFTK as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "occasional-residence",
   "metadata": {},
   "outputs": [],
   "source": [
    "ε = 1e-2randn(3)  # Slight distortion to get something non-boring\n",
    "scfres = compute_silicon(ε)\n",
    "compute_forces(scfres)[1][2]  # Select force on displaced silicon atom"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "frozen-married",
   "metadata": {},
   "source": [
    "The `compute_forces` uses analytical derivative expressions implemented manually inside DFTK. Implementing these by hand for a few standard properties and DFT models is fine. However, the number of DFT models is rather large and so is the number of interesting properties. Some properties even require higher order derivatives, e.g. 2nd and 3rd energy derivatives are also not unusual.\n",
    "\n",
    "Deriving and implementing all these derivatives is tedious, error-prone and time consuming. Let's see if we can at least replicate this force result using some of Julia's **automatic differentiation** tools.\n",
    "\n",
    "First we baseline with **finite differences**. Instead of running the SCF algorithm twice (once for $x$ and once for $x + \\epsilon$), we explicitly exploit the *Hellmann-Feynman theorem*,\n",
    "which states that near an SCF ground state\n",
    "\n",
    "$$ \\frac{d E_\\text{SCF}}{d x} = \\frac{\\partial E_\\text{SCF}}{\\partial x} + \\frac{\\partial E_\\text{SCF}}{\\partial \\Psi} \\frac{\\partial \\Psi}{\\partial x} = \\frac{\\partial E_\\text{SCF}}{\\partial x}$$\n",
    "\n",
    "i.e. that the dependency of $E_\\text{SCF}$ on the orbitals $\\Psi = \\{\\psi_i\\}$ does not need to be considered.\n",
    "\n",
    "With this simplication:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "determined-wrapping",
   "metadata": {},
   "outputs": [],
   "source": [
    "using FiniteDiff\n",
    "scfres = compute_silicon(ε)\n",
    "\n",
    "function recompute_silicon_energy(x)\n",
    "    T = eltype(x)\n",
    "    model = scfres.basis.model\n",
    "    \n",
    "    Si = ElementPsp(:Si, psp=load_psp(\"hgh/lda/si-q4\"))\n",
    "    atoms = [Si => [ones(3)/8, -ones(3)/8 + x]]\n",
    "    new_model  = model_DFT(Array{T}(model.lattice),\n",
    "                           atoms,\n",
    "                           [:lda_x, :lda_c_vwn];\n",
    "                           symmetries=false)\n",
    "    new_basis  = PlaneWaveBasis(new_model; Ecut=13, kgrid=[2, 2, 2]);\n",
    "    \n",
    "    ρ = DFTK.compute_density(new_basis, scfres.ψ, scfres.occupation)\n",
    "    energy_hamiltonian(new_basis, scfres.ψ, scfres.occupation; ρ=scfres.ρ).E.total\n",
    "end\n",
    "\n",
    "- FiniteDiff.finite_difference_gradient(recompute_silicon_energy, ε)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "transparent-graphics",
   "metadata": {},
   "source": [
    "Sort of agrees ... next we try `ForwardDiff` (tip of the hat to our GSoC Niklas Schmitz, who recently made this possible):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "palestinian-tender",
   "metadata": {},
   "outputs": [],
   "source": [
    "using ForwardDiff\n",
    "- ForwardDiff.gradient(recompute_silicon_energy, ε)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "mobile-australian",
   "metadata": {},
   "source": [
    "## Computing stresses\n",
    "\n",
    "To close off this section on property computation, let's use **forward-mode AD** to implement a property computation that is not yet available in DFTK.\n",
    "\n",
    "We will consider the computation of stresses. Denoting the `lattice` matrix by $\\textbf{L}$ for simplicty, the stress is computed as"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "typical-magnet",
   "metadata": {},
   "source": [
    "$$ \\text{Stress} = \\frac{1}{\\text{det}(\\mathbf{L})} \\left. \\frac{d E_\\text{SCF}}{d [(I + \\mathbf{M}) \\, \\mathbf{L}]} \\right|_{\\mathbf{M}=0} = \\frac{1}{\\text{det}(\\mathbf{L})} \\left. \\frac{\\partial E_\\text{SCF}}{\\partial [(I + \\mathbf{M}) \\, \\mathbf{L}]} \\right|_{\\mathbf{M}=0},\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "organic-sender",
   "metadata": {},
   "source": [
    "where again we made use of the Hellman-Feynman theorem."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "middle-ownership",
   "metadata": {},
   "source": [
    "**Exercise 1:** Use the following code fragment to implement stresses using `ForwardDiff`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pretty-antigua",
   "metadata": {},
   "outputs": [],
   "source": [
    "using ForwardDiff\n",
    "using DFTK\n",
    "scfres = compute_silicon(zeros(3))\n",
    "\n",
    "function recompute_silicon_energy_stresses(lattice)\n",
    "    atoms = scfres.basis.model.atoms\n",
    "    new_model  = model_DFT(lattice, atoms, [:lda_x, :lda_c_vwn]; symmetries=false)\n",
    "    new_basis  = PlaneWaveBasis(new_model; Ecut=13, kgrid=[2, 2, 2]);\n",
    "    ρ = DFTK.compute_density(new_basis, scfres.ψ, scfres.occupation)\n",
    "    energy_hamiltonian(new_basis, scfres.ψ, scfres.occupation; ρ=scfres.ρ).E.total\n",
    "end\n",
    "\n",
    "L = scfres.basis.model.lattice\n",
    "\n",
    "# Your code here"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "radio-relations",
   "metadata": {},
   "source": [
    "#### Takeaways:\n",
    "- Practical DFT simulations require not only energies, but also energy derivatives.\n",
    "- While analytical derivatives are fast, implementing higher order derivatives\n",
    "  by hand is very, very time-consuming.\n",
    "- Only a subset of possibilities is usually implemented in standard codes (like DFTK)\n",
    "- Automatic differentation provides an easy way to go beyond the intrinsic properties available in a code\n",
    "\n",
    "\n",
    "- In DFTK we are currently working on more complicated setups for automatic differentation (reverse mode, going beyond Hellmann-Feynman)."
   ]
  }
 ],
 "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
}