{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Geometry optimization\n",
    "\n",
    "We use the DFTK and Optim packages in this example to find the minimal-energy\n",
    "bond length of the ``H_2`` molecule. We setup ``H_2`` in an\n",
    "LDA model just like in the Tutorial for silicon."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using DFTK\n",
    "using Optim\n",
    "using LinearAlgebra\n",
    "using Printf\n",
    "\n",
    "kgrid = [1, 1, 1]       # k-Point grid\n",
    "Ecut = 5                # kinetic energy cutoff in Hartree\n",
    "tol = 1e-8              # tolerance for the optimization routine\n",
    "a = 10                  # lattice constant in Bohr\n",
    "lattice = a * Diagonal(ones(3))\n",
    "H = ElementPsp(:H, psp=load_psp(\"hgh/lda/h-q1\"));"
   ],
   "metadata": {},
   "execution_count": 1
  },
  {
   "cell_type": "markdown",
   "source": [
    "We define a blochwave and a density to be used as global variables so that we\n",
    "can transfer the solution from one iteration to another and therefore reduce\n",
    "the optimization time."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "ψ = nothing\n",
    "ρ = nothing"
   ],
   "metadata": {},
   "execution_count": 2
  },
  {
   "cell_type": "markdown",
   "source": [
    "First, we create a function that computes the solution associated to the\n",
    "position ``x \\in \\mathbb{R}^6`` of the atoms in reduced coordinates\n",
    "(cf. Reduced and cartesian coordinates for more\n",
    "details on the coordinates system).\n",
    "They are stored as a vector: `x[1:3]` represents the position of the\n",
    "first atom and `x[4:6]` the position of the second.\n",
    "We also update `ψ` and `ρ` for the next iteration."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "function compute_scfres(x)\n",
    "    atoms = [H => [x[1:3], x[4:6]]]\n",
    "    model = model_LDA(lattice, atoms)\n",
    "    basis = PlaneWaveBasis(model, Ecut; kgrid=kgrid)\n",
    "    global ψ, ρ\n",
    "    if ρ === nothing\n",
    "        ρ = guess_density(basis)\n",
    "    end\n",
    "    scfres = self_consistent_field(basis; ψ=ψ, ρ=ρ,\n",
    "                                   tol=tol / 10, callback=info->nothing)\n",
    "    ψ = scfres.ψ\n",
    "    ρ = scfres.ρ\n",
    "    scfres\n",
    "end;"
   ],
   "metadata": {},
   "execution_count": 3
  },
  {
   "cell_type": "markdown",
   "source": [
    "Then, we create the function we want to optimize: `fg!` is used to update the\n",
    "value of the objective function `F`, namely the energy, and its gradient `G`,\n",
    "here computed with the forces (which are, by definition, the negative gradient\n",
    "of the energy)."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "function fg!(F, G, x)\n",
    "    scfres = compute_scfres(x)\n",
    "    if G != nothing\n",
    "        grad = compute_forces(scfres)\n",
    "        G .= -[grad[1][1]; grad[1][2]]\n",
    "    end\n",
    "    scfres.energies.total\n",
    "end;"
   ],
   "metadata": {},
   "execution_count": 4
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now, we can optimize on the 6 parameters `x = [x1, y1, z1, x2, y2, z2]` in\n",
    "reduced coordinates, using `LBFGS()`, the default minimization algorithm\n",
    "in Optim. We start from `x0`, which is a first guess for the coordinates. By\n",
    "default, `optimize` traces the output of the optimization algorithm during the\n",
    "iterations. Once we have the minimizer `xmin`, we compute the bond length in\n",
    "cartesian coordinates."
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iter     Function value   Gradient norm \n",
      "     0    -1.061170e+00     6.234894e-01\n",
      " * time: 4.482269287109375e-5\n",
      "     1    -1.065563e+00     4.035462e-02\n",
      " * time: 2.2183380126953125\n",
      "     2    -1.065592e+00     7.038154e-04\n",
      " * time: 2.5513339042663574\n",
      "     3    -1.065592e+00     4.580753e-06\n",
      " * time: 2.8047499656677246\n",
      "     4    -1.065592e+00     1.451640e-06\n",
      " * time: 2.940534830093384\n",
      "\n",
      "Optimal bond length for Ecut=5.00: 1.557 Bohr\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "x0 = vcat(lattice \\ [0., 0., 0.], lattice \\ [1.4, 0., 0.])\n",
    "xres = optimize(Optim.only_fg!(fg!), x0, LBFGS(),\n",
    "                Optim.Options(show_trace=true, f_tol=tol))\n",
    "xmin = Optim.minimizer(xres)\n",
    "dmin = norm(lattice*xmin[1:3] - lattice*xmin[4:6])\n",
    "@printf \"\\nOptimal bond length for Ecut=%.2f: %.3f Bohr\\n\" Ecut dmin"
   ],
   "metadata": {},
   "execution_count": 5
  },
  {
   "cell_type": "markdown",
   "source": [
    "We used here very rough parameters to generate the example and\n",
    "setting `Ecut` to 10 Ha yields a bond length of 1.523 Bohr,\n",
    "which [agrees with ABINIT](https://docs.abinit.org/tutorial/base1/).\n",
    "\n",
    "!!! note \"Degrees of freedom\"\n",
    "    We used here a very general setting where we optimized on the 6 variables\n",
    "    representing the position of the 2 atoms and it can be easily extended\n",
    "    to molecules with more atoms (such as ``H_2O``). In the particular case\n",
    "    of ``H_2``, we could use only the internal degree of freedom which, in\n",
    "    this case, is just the bond length."
   ],
   "metadata": {}
  }
 ],
 "nbformat_minor": 3,
 "metadata": {
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.6.1"
  },
  "kernelspec": {
   "name": "julia-1.6",
   "display_name": "Julia 1.6.1",
   "language": "julia"
  }
 },
 "nbformat": 4
}