5_Solving_the_SCF_problem.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "psychological-praise",
   "metadata": {},
   "source": [
    "# Solving the SCF problem\n",
    "\n",
    "For the work in this notebook we will stick with the aluminum setup introduced before with one additional hinge: We will allow to make the problem harder or easier, by forming a supercell.\n",
    "\n",
    "- Since we are dealing with periodic problems there is no unique definition of the lattice. Clearly any duplication of the lattice along an axis is also a valid lattice.\n",
    "- This is exactly what a **supercell** is: An $n$-fold repetition in one of the lattice axes.\n",
    "\n",
    "The following code achieves this for aluminium:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "psychological-arrangement",
   "metadata": {},
   "outputs": [],
   "source": [
    "using DFTK\n",
    "using LinearAlgebra\n",
    "setup_threading(n_blas=1)\n",
    "\n",
    "function aluminium_setup(repeat=1; Ecut=12.0, kgrid=[2, 2, 2])\n",
    "    a = 7.65339\n",
    "    lattice = a * Matrix(I, 3, 3)\n",
    "    Al = ElementPsp(:Al, psp=load_psp(\"hgh/lda/al-q3\"))\n",
    "    atoms     = [Al, Al, Al, Al]\n",
    "    positions = [[0.0, 0.0, 0.0], [0.0, 0.5, 0.5], [0.5, 0.0, 0.5], [0.5, 0.5, 0.0]]\n",
    "\n",
    "    # Make supercell in ASE:\n",
    "    # We convert our lattice to the conventions used in ASE\n",
    "    # and then back ...\n",
    "    supercell = ase_atoms(lattice, atoms, positions) * (repeat, 1, 1)\n",
    "    lattice   = load_lattice(supercell)\n",
    "    positions = load_positions(supercell)\n",
    "    atoms = fill(Al, length(positions))\n",
    "\n",
    "    # Construct an LDA model and discretise\n",
    "    model = model_LDA(lattice, atoms, positions; temperature=1e-3, symmetries=false)\n",
    "    PlaneWaveBasis(model; Ecut, kgrid)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "parental-thickness",
   "metadata": {},
   "source": [
    "As we will see in this notebook the modelling of a system generally becomes harder if the system becomes larger. \n",
    "\n",
    "- This sounds like a trival statement as *per se* the cost per SCF step increases\n",
    "  as the system (and thus $N$) gets larger.\n",
    "- But there is more to it:\n",
    "  If one is not careful also the *number of SCF iterations* increases\n",
    "  as the system gets larger.\n",
    "  \n",
    "  \n",
    "- The aim of many tricks I will show in this workbook is to ensure the **number of SCF iterations remains constant** when the system size increases."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "latter-session",
   "metadata": {},
   "source": [
    "## Gaining insight\n",
    "\n",
    "In DFTK one can easily patch or extend the SCF procedure\n",
    "by replacing parts of the code with custom callback functions.\n",
    "\n",
    "We will use this in this notebook to construct our own SCF solver,\n",
    "but without needing to worry about all the nasty bits\n",
    "(proper normalisation, numerically stable formation of the density etc.).\n",
    "\n",
    "Before we do that, let's demonstrate first how to use callbacks to extract information from a running SCF:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "alone-stroke",
   "metadata": {},
   "outputs": [],
   "source": [
    "using DFTK\n",
    "using Plots\n",
    "\n",
    "p = plot(yaxis=:log)  # Setup an empty plot canvas\n",
    "density_differences = Float64[]\n",
    "function plot_callback(info)    \n",
    "    if info.stage == :finalize\n",
    "        # When done with the SCF: Plot it!\n",
    "        plot!(p, density_differences, label=\"|ρout - ρin|\", markershape=:x)\n",
    "    else\n",
    "        # Just add the density difference of this step\n",
    "        push!(density_differences, norm(info.ρout - info.ρin))\n",
    "    end\n",
    "        \n",
    "    info  # Pass info through to allow callback chaining\n",
    "end\n",
    "\n",
    "# Chain the custom callback with the default one\n",
    "# (printing the convergence table)\n",
    "callback = DFTK.ScfDefaultCallback() ∘ plot_callback\n",
    "    \n",
    "# Run the SCF and show the plot\n",
    "scfres = self_consistent_field(aluminium_setup(); tol=1e-12, callback, mixing=SimpleMixing())\n",
    "p"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "furnished-congress",
   "metadata": {},
   "source": [
    "**Live exercise 4.1:** Try making this problem harder by running on `aluminium_setup(2)` or `aluminium_setup(5)` (or higher if you can efford it). What do you observe in the plot?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "administrative-manufacturer",
   "metadata": {},
   "source": [
    "**Remark:** This callback allows to read or modify the full state of the SCF iteration,\n",
    "which is a valuable tool when debugging an SCF algorithm.\n",
    "When working in the REPL or with scripts one of my favourite callbacks is\n",
    "```julia\n",
    "using Infiltrator\n",
    "function infiltrate_callback(info)\n",
    "    @infiltrate info.n_iter == 5\n",
    "    info    \n",
    "end\n",
    "scfres = self_consistent_field(aluminium_setup(); tol=1e-6,\n",
    "                               callback=infiltrate_callback);\n",
    "```\n",
    "to take you right to the SCF state at a surprising iteration."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "herbal-concentration",
   "metadata": {},
   "source": [
    "Now we have all the tools in place ... let's do some numerics:\n",
    "\n",
    "## Building our own SCF 1: Fixed-point iterations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "incorporated-woman",
   "metadata": {},
   "source": [
    "As we saw before the self-consistent field procedure required to solve the DFT problem can be written as a fixed-point problem\n",
    "$$ F(\\rho) = \\rho $$\n",
    "where $F(\\rho) = D(V(\\rho))$ is the basic SCF step. That is the construction of the Kohn-Sham Hamiltonian $H(\\rho)$ given the density $\\rho$, followed its diagonalisation to obtain its eigenpairs $(\\varepsilon_{k i}, \\psi_{ki})$\n",
    "and from these a new density\n",
    "$$ \\rho(r) = \\sum_{k\\in\\text{BZ}} \\sum_i \\psi_{ki}(r) \\, \\psi_{ki}^\\ast(r).$$\n",
    "\n",
    "We will not be concerned with $F$ itself, which we will take for \"granted\" (i.e. delivered by DFTK).\n",
    "What we will consider, however, is multiple ways to solve the DFT fixed-point problem.\n",
    "\n",
    "The easiest are plain **fixed-point iterations**, i.e."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "binary-thinking",
   "metadata": {},
   "source": [
    "$$ \\rho_{n+1} = F(\\rho_n), $$\n",
    "starting from a hopefully good initial guess $\\rho_0$. DFTK automatically provides a reasonable\n",
    "guess density, such that we only need to take care of the iterations themselves.\n",
    "In the language of DFTK this algorithm is written as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "reported-hardware",
   "metadata": {},
   "outputs": [],
   "source": [
    "function fixed_point_iteration(F, ρ₀, maxiter; tol)\n",
    "    # F:        The SCF step function\n",
    "    # ρ₀:       The initial guess density\n",
    "    # maxiter:  The maximal number of iterations to be performed\n",
    "    # tol:      The selected convergence tolerance\n",
    "    \n",
    "    ρ  = ρ₀\n",
    "    Fρ = F(ρ)\n",
    "    for n = 1:maxiter \n",
    "        # If change less than tolerance, break iterations:\n",
    "        if norm(Fρ - ρ) < tol\n",
    "            break\n",
    "        end\n",
    "        ρ  = Fρ\n",
    "        Fρ = F(ρ)\n",
    "    end\n",
    "    \n",
    "    # Return some stuff DFTK needs ...\n",
    "    (fixpoint=ρ, converged=norm(Fρ - ρ) < tol)\n",
    "end;\n",
    "\n",
    "# use this algorithm inside DFTK's SCF for solving the aluminium problem\n",
    "# (the other parameters are needed to overwrite some DFTK defaults\n",
    "#  we don't want to use just yet).\n",
    "self_consistent_field(aluminium_setup(); solver=fixed_point_iteration, damping=1.0,\n",
    "                                         maxiter=40, mixing=SimpleMixing());"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "selective-zambia",
   "metadata": {},
   "source": [
    "As can be observed this algorithm is not very good and completely fails to converge. This is a known limitation of this algorithm, which is why it is not used in practice."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "composite-springfield",
   "metadata": {},
   "source": [
    "## Step 2: Damped iterations\n",
    "\n",
    "The next step is to introduce a so-called damping parameter $\\alpha$, which is given a value between $0$ and $1$. One now iterates as follows:\n",
    "$$ \\rho_{n+1} = \\rho_{n} + \\alpha (F(\\rho_n) - \\rho_n) $$\n",
    "In other words the update $F(\\rho_n) - \\rho_n$ proposed in the $n$-th SCF step is not fully taken, but scaled-down by the damping $\\alpha$.\n",
    "\n",
    "**Home work 4.2:** Modify `fixed_point_iteration` such that it supports this *damped* fixed-point iteration. Try different values for $\\alpha$ between $0$ and $1$ and estimate roughly the $\\alpha$ which gives fastest convergence. For which $\\alpha$ do you observe no convergence at all?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "arabic-haiti",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Your solution here ..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "radio-phase",
   "metadata": {},
   "source": [
    "**Remark:** The observations you make here are general. We will argue in the next notebook why every SCF converges (locally) if a small enough $\\alpha > 0$ is chosen."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "continuous-genius",
   "metadata": {},
   "source": [
    "## Step 3: Anderson acceleration\n",
    "\n",
    "The `fixed_point_iteration` function above (with the damping extension) already covers the main gist of standard DFT algorithms. To make things converge faster the next step to follow is Anderson acceleration, where not only $\\rho_n$ and $F(\\rho_n)$, but also older iterates are used to propose the next density.\n",
    "\n",
    "For Anderson one exploits that the update $R(\\rho) = F(\\rho) - \\rho$ is also the residual of the fixed-point problem $F(\\rho) = \\rho$, i.e. how far away we are from the fixed-point density. A good next density $\\rho_{n+1}$ therefore should be found by minimising an approximation for $R(\\rho_{n+1})$. Assuming the SCF was linear in the density (which it is not), a good idea is to find a linear combination of residuals\n",
    "$$\\min_{\\beta_i} \\left\\| \\sum_i \\beta_i R(\\rho_i) \\right\\|^2$$\n",
    "which has the smallest possible norm and to use these coefficients $\\beta_i$ to extrapolate the next\n",
    "density\n",
    "$$ \\rho_{n+1} =  \\sum_i \\beta_i (\\rho_i + \\alpha R(\\rho_i)) $$\n",
    "where you notice the \"standard\" damped fixed-point iteration in the summed terms.\n",
    "\n",
    "In terms of an algorithm this is "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cultural-individual",
   "metadata": {},
   "outputs": [],
   "source": [
    "function anderson_iteration(F, ρ₀, maxiter; tol)\n",
    "    # F:        The SCF step function\n",
    "    # ρ₀:       The initial guess density\n",
    "    # maxiter:  The maximal number of iterations to be performed\n",
    "    # tol:      The selected convergence tolerance\n",
    "    \n",
    "    converged = false\n",
    "    ρ  = ρ₀\n",
    "    ρs = []\n",
    "    Rs = []\n",
    "    for n = 1:maxiter\n",
    "        Fρ = F(ρ)\n",
    "        Rρ = Fρ - ρ\n",
    "        converged = norm(Rρ) < tol\n",
    "        converged && break\n",
    "        \n",
    "        ρnext = vec(ρ) .+ vec(Rρ)\n",
    "        if !isempty(Rs)\n",
    "            M = hcat(Rs...) .- vec(Rρ)\n",
    "            βs = -(M \\ vec(Rρ))\n",
    "            \n",
    "            for (iβ, β) in enumerate(βs)\n",
    "                ρnext .+= β .* (ρs[iβ] .- vec(ρ) .+ Rs[iβ] .- vec(Rρ))\n",
    "            end\n",
    "        end\n",
    "                    \n",
    "        push!(ρs, vec(ρ))\n",
    "        push!(Rs, vec(Rρ))\n",
    "        ρ = reshape(ρnext, size(ρ₀)...)\n",
    "    end\n",
    "    \n",
    "    # Return some stuff DFTK needs ...\n",
    "    (fixpoint=ρ, converged=converged)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cultural-neighborhood",
   "metadata": {},
   "source": [
    "To work with this algorithm we will use DFTK's intrinsic mechanism to choose a damping. The syntax for this is\n",
    "\n",
    "```julia\n",
    "repeat = 1\n",
    "self_consistent_field(aluminium_setup(repeat);\n",
    "                      solver=anderson_iteration,\n",
    "                      damping=0.8, maxiter=40,\n",
    "                      mixing=SimpleMixing());\n",
    "```\n",
    "to choose a damping of $\\alpha = 0.8$ and run for at most `maxiter` iterations.\n",
    "\n",
    "**Home work 4.3:** Pick $\\alpha = 0.8$ and make the problem harder by increasing `repeat` (e.g. `2`, `4`, `6`, `8`). Can you make Anderson fail to converge? What do you notice in terms of the number of iterations and runtimes?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "modular-delta",
   "metadata": {},
   "source": [
    "**Remark:** Anderson acceleration comes in many names and variants. It is sometimes also called Pulay mixing or DIIS (direct inversion of the iterative subspace). Closely related are also Broyden's methods. For more details on the relationship of these methods see this [review on SCF methods](http://doi.org/10.1088/1361-648x/ab31c0) as well as this [paper on the convergence analysis of DIIS methods](https://arxiv.org/abs/2002.12850)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "wireless-baseball",
   "metadata": {},
   "source": [
    "## Step 4: Using NLsolve\n",
    "\n",
    "Of course it is never a good idea to recode standard algorithms such as Anderson acceleration when they are already implemented in other Julia packages.\n",
    "In DFTK we actually rely on the Anderson solver from [NLsolve.jl](https://github.com/JuliaNLSolvers/NLsolve.jl), which is even a bit smarter than the simple version we coded up above.\n",
    "\n",
    "[Our code](https://github.com/JuliaMolSim/DFTK.jl/blob/master/src/scf/scf_solvers.jl#L16--L23) to invoke NLsolve is pretty much just:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sealed-basement",
   "metadata": {},
   "outputs": [],
   "source": [
    "using NLsolve\n",
    "function nlsolve_solver(F, ρ₀, maxiter; tol)\n",
    "    res = nlsolve(ρ -> F(ρ) - ρ, ρ₀; method=:anderson, m=10, xtol=tol,\n",
    "                  ftol=0.0, show_trace=false, iterations=maxiter)\n",
    "    (fixpoint=res.zero, converged=converged(res))\n",
    "end\n",
    "\n",
    "self_consistent_field(aluminium_setup(2); solver=nlsolve_solver, damping=0.8, maxiter=40, tol=1e-6);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "communist-yorkshire",
   "metadata": {},
   "source": [
    "## Step 5: Using preconditioned iterations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb51ed44",
   "metadata": {},
   "source": [
    "Anderson allows us to push the boundary for SCF methods, but for larger or more challenging systems it is not fully sufficient. The next ingredient for a stable SCF procedure is based on the insight we already had previously: The convergence properties of an SCF depends on the material which is modelled.\n",
    "As a result the ideal SCF procedure should be slightly different for each material.\n",
    "In our discussion so far we did not yet take that into account.\n",
    "\n",
    "The standard approach to include material specificity into the SCF is to employ *preconditioned* damped fixed-point iterations.\n",
    "To explain the idea, let us consider again a framework without Anderson acceleration.\n",
    "Preconditioned iterations are then\n",
    "$$ \\rho_{n+1} = \\rho_{n} + \\alpha P^{-1} (F(\\rho_n) - \\rho_n), $$\n",
    "where $P^{-1}$ is a preconditioner that hopefully improve convergence.\n",
    "To re-introduce Anderson around this iteration\n",
    "just replace the previous definition of $R$ by\n",
    "$R(\\rho) = P^{-1} (F(\\rho_n) - \\rho_n)$. That's it.\n",
    "\n",
    "As we will discuss in the next notebook the ideal preconditioning $P$\n",
    "depends on the dielectric properties of the material (e.g. whether it is a\n",
    "metal, insulator or semiconductor as well as other details). \n",
    "Finding a good $P$ is not always easy and for some cases satisfactory options are not yet known. For our aluminium case, however, we are lucky. The `KerkerMixing` implemented in DFTK provides a good $P$ for aluminium.\n",
    "\n",
    "You might wonder about the term *mixing*. In an interdisciplinary community like DFT, different scientists use different vocabulary and \"mixing\" is the \"physicists' term\" used for preconditioning.\n",
    "\n",
    "To use `KerkerMixing` with DFTK run the SCF as follows\n",
    "```julia\n",
    "self_consistent_field(basis; damping=0.8, mixing=KerkerMixing());\n",
    "```\n",
    "\n",
    "**Live exercise 4.4:** Try this setup for different values of `repeat` and check the number of iterations needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be905852",
   "metadata": {},
   "outputs": [],
   "source": [
    "repeat = 1\n",
    "kerker = self_consistent_field(aluminium_setup(repeat); damping=0.8, callback=identity, mixing=KerkerMixing());\n",
    "simple = self_consistent_field(aluminium_setup(repeat); damping=0.8, callback=identity, mixing=SimpleMixing());\n",
    "println(repeat, \" Kerker \", kerker.n_iter, \" Simple \", simple.n_iter)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "previous-impossible",
   "metadata": {},
   "source": [
    "**Remark:** Besides `KerkerMixing` DFTK also has some more preconditioner to offer. This includes `DielectricMixing` (best for semiconductors), `SimpleMixing` (which is $P = I$, i.e. no preconditioner at all, best for insulators) or `LdosMixing` (self-adapting, suitable for both metals *or* insulators *or* inhomogeneous mixtures). Note that `LdosMixing` is the default in DFTK (i.e. used if the `mixing` parameter is *not* supplied to `self_consistent_field`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "hollywood-walker",
   "metadata": {},
   "source": [
    "#### Takeaways:\n",
    "- Large systems require a matching preconditioner to converge in few SCF iterations\n",
    "- Anderson acceleration and/or small damping aids convergence\n",
    "- Provided Anderson is used one often gets away using a non-matching preconditioner\n",
    "\n",
    "\n",
    "- The callback infrastructure of DFTK's SCF allows to to modify a few more aspects of the iteration\n",
    "  very easily. See [this example](https://docs.dftk.org/stable/examples/custom_solvers/) in the documentation for details."
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Julia (4 threads) 1.7.3",
   "language": "julia",
   "name": "julia-(4-threads)-1.7"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}