{
 "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": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "2-element Vector{ElementPsp}:\n ElementPsp(H, psp=\"hgh/lda/h-q1\")\n ElementPsp(H, psp=\"hgh/lda/h-q1\")"
     },
     "metadata": {},
     "execution_count": 1
    }
   ],
   "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 * I(3)\n",
    "H = ElementPsp(:H, psp=load_psp(\"hgh/lda/h-q1\"));\n",
    "atoms = [H, H]"
   ],
   "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",
    "    model = model_LDA(lattice, atoms, [x[1:3], x[4:6]])\n",
    "    basis = PlaneWaveBasis(model; Ecut, kgrid)\n",
    "    global ψ, ρ\n",
    "    if isnothing(ρ)\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]; grad[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.061651e+00     6.219289e-01\n",
      " * time: 5.698204040527344e-5\n",
      "     1    -1.064048e+00     2.935204e-01\n",
      " * time: 1.7946128845214844\n",
      "     2    -1.066007e+00     4.880014e-02\n",
      " * time: 2.262877941131592\n",
      "     3    -1.066049e+00     4.221693e-04\n",
      " * time: 2.5166330337524414\n",
      "     4    -1.066049e+00     4.707030e-06\n",
      " * time: 2.6771340370178223\n",
      "     5    -1.066049e+00     1.747533e-07\n",
      " * time: 2.9325900077819824\n",
      "\n",
      "Optimal bond length for Ecut=5.00: 1.556 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.8.1"
  },
  "kernelspec": {
   "name": "julia-1.8",
   "display_name": "Julia 1.8.1",
   "language": "julia"
  }
 },
 "nbformat": 4
}