{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Practical error bounds for the forces\n",
    "\n",
    "This is a simple example showing how to compute error estimates for the forces\n",
    "on a ``{\\rm TiO}_2`` molecule, from which we can either compute asymptotically\n",
    "valid error bounds or increase the precision on the computation of the forces.\n",
    "\n",
    "The strategy we follow is described with more details in [^CDKL2021] and we\n",
    "will use in comments the density matrices framework. We will also needs\n",
    "operators and functions from\n",
    "[`src/scf/newton.jl`](https://dftk.org/blob/master/src/scf/newton.jl).\n",
    "\n",
    "[^CDKL2021]:\n",
    "    E. Cancès, G. Dusson, G. Kemlin, and A. Levitt\n",
    "    *Practical error bounds for properties in plane-wave electronic structure\n",
    "    calculations* Preprint, 2021. [arXiv](https://arxiv.org/abs/2111.01470)"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using DFTK\n",
    "using LinearAlgebra\n",
    "using ForwardDiff\n",
    "using LinearMaps\n",
    "using IterativeSolvers"
   ],
   "metadata": {},
   "execution_count": 1
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Setup\n",
    "We setup manually the ``{\\rm TiO}_2`` configuration from\n",
    "[Materials Project](https://materialsproject.org/materials/mp-2657/)."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "Ti = ElementPsp(:Ti, psp=load_psp(\"hgh/lda/ti-q4.hgh\"))\n",
    "O = ElementPsp(:O, psp=load_psp(\"hgh/lda/o-q6.hgh\"))\n",
    "atoms = [Ti => [[0.5, 0.5, 0.5],\n",
    "                [0.0, 0.0, 0.0]],\n",
    "         O  => [[0.19542, 0.80458, 0.5],\n",
    "                [0.80458, 0.19542, 0.5],\n",
    "                [0.30458, 0.30458, 0.0],\n",
    "                [0.69542, 0.69542, 0.0]]]\n",
    "lattice = [[8.79341  0.0      0.0];\n",
    "           [0.0      8.79341  0.0];\n",
    "           [0.0      0.0      5.61098]];"
   ],
   "metadata": {},
   "execution_count": 2
  },
  {
   "cell_type": "markdown",
   "source": [
    "We apply a small displacement to one of the ``\\rm Ti`` atoms to get nonzero\n",
    "forces."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "el, pos = atoms[1]\n",
    "atoms[1] = Pair(el, pos .+ [[-0.22, 0.28, 0.35] / 10, [0, 0, 0]]);"
   ],
   "metadata": {},
   "execution_count": 3
  },
  {
   "cell_type": "markdown",
   "source": [
    "We build a model with one k-point only, not too high `Ecut_ref` and small\n",
    "tolerance to limit computational time. These parameters can be increased for\n",
    "more precise results."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "model = model_LDA(lattice, atoms)\n",
    "kgrid = [1, 1, 1]\n",
    "Ecut_ref = 35\n",
    "basis_ref = PlaneWaveBasis(model; Ecut=Ecut_ref, kgrid)\n",
    "tol = 1e-8;"
   ],
   "metadata": {},
   "execution_count": 4
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Computations"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "We compute the reference solution ``P_*`` from which we will compute the\n",
    "references forces."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "scfres_ref = self_consistent_field(basis_ref, tol=tol, callback=info->nothing)\n",
    "ψ_ref, _ = DFTK.select_occupied_orbitals(basis_ref, scfres_ref.ψ,\n",
    "                                         scfres_ref.occupation);"
   ],
   "metadata": {},
   "execution_count": 5
  },
  {
   "cell_type": "markdown",
   "source": [
    "We compute a variational approximation of the reference solution with\n",
    "smaller `Ecut`. `ψr`, `ρr` and `Er` are the quantities computed with `Ecut`\n",
    "and then extended to the reference grid.\n",
    "\n",
    "!!! note \"Choice of convergence parameters\"\n",
    "    Be careful to choose `Ecut` not too close to `Ecut_ref`.\n",
    "    Note also that the current choice `Ecut_ref = 35` is such that the\n",
    "    reference solution is not converged and `Ecut = 15` is such that the\n",
    "    asymptotic regime (crucial to validate the approach) is barely established."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "Ecut = 15\n",
    "basis = PlaneWaveBasis(model; Ecut=Ecut, kgrid)\n",
    "scfres = self_consistent_field(basis, tol=tol, callback=info->nothing)\n",
    "ψr = DFTK.transfer_blochwave(scfres.ψ, basis, basis_ref)\n",
    "ρr = compute_density(basis_ref, ψr, scfres.occupation)\n",
    "Er, hamr = energy_hamiltonian(basis_ref, ψr, scfres.occupation; ρ=ρr);"
   ],
   "metadata": {},
   "execution_count": 6
  },
  {
   "cell_type": "markdown",
   "source": [
    "We then compute several quantities that we need to evaluate the error bounds."
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "- Compute the residual ``R(P)``, and remove the virtual orbitals, as required\n",
    "  in [`src/scf/newton.jl`](https://github.com/JuliaMolSim/DFTK.jl/blob/fedc720dab2d194b30d468501acd0f04bd4dd3d6/src/scf/newton.jl#L121)."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "res = DFTK.compute_projected_gradient(basis_ref, ψr, scfres.occupation)\n",
    "res, occ = DFTK.select_occupied_orbitals(basis_ref, res, scfres.occupation)\n",
    "ψr, _ = DFTK.select_occupied_orbitals(basis_ref, ψr, scfres.occupation);"
   ],
   "metadata": {},
   "execution_count": 7
  },
  {
   "cell_type": "markdown",
   "source": [
    "- Compute the error ``P-P_*`` on the associated orbitals ``ϕ-ψ`` after aligning\n",
    "  them: this is done by solving ``\\min |ϕ - ψU|`` for ``U`` unitary matrix of\n",
    "  size ``N\\times N`` (``N`` being the number of electrons) whose solution is\n",
    "  ``U = S(S^*S)^{-1/2}`` where ``S`` is the overlap matrix ``ψ^*ϕ``."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "function compute_error(basis, ϕ, ψ)\n",
    "    map(zip(ϕ, ψ)) do (ϕk, ψk)\n",
    "        S = ψk'ϕk\n",
    "        U = S*(S'S)^(-1/2)\n",
    "        ϕk - ψk*U\n",
    "    end\n",
    "end\n",
    "err = compute_error(basis_ref, ψr, ψ_ref);"
   ],
   "metadata": {},
   "execution_count": 8
  },
  {
   "cell_type": "markdown",
   "source": [
    "- Compute ``{\\bm M}^{-1}R(P)`` with ``{\\bm M}^{-1}`` defined in [^CDKL2021]:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "P = [PreconditionerTPA(basis_ref, kpt) for kpt in basis_ref.kpoints]\n",
    "map(zip(P, ψr)) do (Pk, ψk)\n",
    "    DFTK.precondprep!(Pk, ψk)\n",
    "end\n",
    "function apply_M(φk, Pk, δφnk, n)\n",
    "    DFTK.proj_tangent_kpt!(δφnk, φk)\n",
    "    δφnk = sqrt.(Pk.mean_kin[n] .+ Pk.kin) .* δφnk\n",
    "    DFTK.proj_tangent_kpt!(δφnk, φk)\n",
    "    δφnk = sqrt.(Pk.mean_kin[n] .+ Pk.kin) .* δφnk\n",
    "    DFTK.proj_tangent_kpt!(δφnk, φk)\n",
    "end\n",
    "function apply_inv_M(φk, Pk, δφnk, n)\n",
    "    DFTK.proj_tangent_kpt!(δφnk, φk)\n",
    "    op(x) = apply_M(φk, Pk, x, n)\n",
    "    function f_ldiv!(x, y)\n",
    "        x .= DFTK.proj_tangent_kpt(y, φk)\n",
    "        x ./= (Pk.mean_kin[n] .+ Pk.kin)\n",
    "        DFTK.proj_tangent_kpt!(x, φk)\n",
    "    end\n",
    "    J = LinearMap{eltype(φk)}(op, size(δφnk, 1))\n",
    "    δφnk = cg(J, δφnk, Pl=DFTK.FunctionPreconditioner(f_ldiv!),\n",
    "              verbose=false, reltol=0, abstol=1e-15)\n",
    "    DFTK.proj_tangent_kpt!(δφnk, φk)\n",
    "end\n",
    "function apply_metric(φ, P, δφ, A::Function)\n",
    "    map(enumerate(δφ)) do (ik, δφk)\n",
    "        Aδφk = similar(δφk)\n",
    "        φk = φ[ik]\n",
    "        for n = 1:size(δφk,2)\n",
    "            Aδφk[:,n] = A(φk, P[ik], δφk[:,n], n)\n",
    "        end\n",
    "        Aδφk\n",
    "    end\n",
    "end\n",
    "Mres = apply_metric(ψr, P, res, apply_inv_M);"
   ],
   "metadata": {},
   "execution_count": 9
  },
  {
   "cell_type": "markdown",
   "source": [
    "We can now compute the modified residual ``R_{\\rm Schur}(P)`` using a Schur\n",
    "complement to approximate the error on low-frequencies[^CDKL2021]:\n",
    "\n",
    "$$\n",
    "\\begin{bmatrix}\n",
    "(\\bm \\Omega + \\bm K)_{11} & (\\bm \\Omega + \\bm K)_{12} \\\\\n",
    "0 & {\\bm M}_{22}\n",
    "\\end{bmatrix}\n",
    "\\begin{bmatrix}\n",
    "P_{1} - P_{*1} \\\\ P_{2}-P_{*2}\n",
    "\\end{bmatrix} =\n",
    "\\begin{bmatrix}\n",
    "R_{1} \\\\ R_{2}\n",
    "\\end{bmatrix}.\n",
    "$$"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "- Compute the projection of the residual onto the high and low frequencies:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "resLF = DFTK.transfer_blochwave(res, basis_ref, basis)\n",
    "resHF = res - DFTK.transfer_blochwave(resLF, basis, basis_ref);"
   ],
   "metadata": {},
   "execution_count": 10
  },
  {
   "cell_type": "markdown",
   "source": [
    "- Compute ``{\\bm M}^{-1}_{22}R_2(P)``:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "e2 = apply_metric(ψr, P, resHF, apply_inv_M);"
   ],
   "metadata": {},
   "execution_count": 11
  },
  {
   "cell_type": "markdown",
   "source": [
    "- Compute the right hand side of the Schur system:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "# Rayleigh coefficients needed for `apply_Ω`\n",
    "Λ = map(enumerate(ψr)) do (ik, ψk)\n",
    "    Hk = hamr.blocks[ik]\n",
    "    Hψk = Hk * ψk\n",
    "    ψk'Hψk\n",
    "end\n",
    "ΩpKe2 = DFTK.apply_Ω(e2, ψr, hamr, Λ) .+ DFTK.apply_K(basis_ref, e2, ψr, ρr, occ)\n",
    "ΩpKe2 = DFTK.transfer_blochwave(ΩpKe2, basis_ref, basis)\n",
    "rhs = resLF - ΩpKe2;"
   ],
   "metadata": {},
   "execution_count": 12
  },
  {
   "cell_type": "markdown",
   "source": [
    "- Solve the Schur system to compute ``R_{\\rm Schur}(P)``: this is the most\n",
    "  costly step, but inverting ``\\bm{\\Omega} + \\bm{K}`` on the small space has\n",
    "  the same cost than the full SCF cycle on the small grid."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "ψ, _ = DFTK.select_occupied_orbitals(basis, scfres.ψ, scfres.occupation)\n",
    "e1 = DFTK.solve_ΩplusK(basis, ψ, rhs, occ; tol_cg=tol)\n",
    "e1 = DFTK.transfer_blochwave(e1, basis, basis_ref)\n",
    "res_schur = e1 + Mres;"
   ],
   "metadata": {},
   "execution_count": 13
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Error estimates"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "We start with different estimations of the forces on the first ``\\rm Ti`` atom,\n",
    "in direction 1.\n",
    "- Forces computed with the reference solution:"
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F(P_*) = -0.299855737211441\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "f_ref = compute_forces(scfres_ref)\n",
    "# [1][1][1] means that we look at the forces on the first atom type (Ti), the\n",
    "# first atom of this type, in the first direction.\n",
    "println(\"F(P_*) = $(f_ref[1][1][1])\")"
   ],
   "metadata": {},
   "execution_count": 14
  },
  {
   "cell_type": "markdown",
   "source": [
    "- Forces computed with the variational approximation:"
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F(P) = 0.034342967212890096\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "f = compute_forces(scfres)\n",
    "println(\"F(P) = $(f[1][1][1])\")"
   ],
   "metadata": {},
   "execution_count": 15
  },
  {
   "cell_type": "markdown",
   "source": [
    "We then try to improve ``F(P)`` using the first order linearization:\n",
    "\n",
    "$$\n",
    "F(P) = F(P_*) + {\\rm d}F(P)\\cdot(P-P_*).\n",
    "$$"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "To this end, we use the `ForwardDiff.jl` package to compute ``{\\rm d}F(P)``\n",
    "using automatic differentiation."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "function df(basis, occupation, ψ, δψ, ρ)\n",
    "    δρ = DFTK.compute_δρ(basis, ψ, δψ, occupation)\n",
    "    ForwardDiff.derivative(ε -> compute_forces(basis, ψ.+ε.*δψ, occupation; ρ=ρ+ε.*δρ), 0)\n",
    "end;"
   ],
   "metadata": {},
   "execution_count": 16
  },
  {
   "cell_type": "markdown",
   "source": [
    "- Computation of the forces by a linearization argument if we have access to\n",
    "  the actual error ``P-P_*``:"
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F(P) - df(P)⋅(P-P_*) = -0.3200261541003476\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "df_err = df(basis_ref, occ, ψr, DFTK.proj_tangent(err, ψr), ρr)\n",
    "println(\"F(P) - df(P)⋅(P-P_*) = $(f[1][1][1]-df_err[1][1][1])\")"
   ],
   "metadata": {},
   "execution_count": 17
  },
  {
   "cell_type": "markdown",
   "source": [
    "- Computation of the forces by a linearization argument when replacing the\n",
    "  error ``P-P_*`` by the modified residual:"
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "F(P) - df(P)⋅Rschur(P) = -0.2127458246942039\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "df_schur = df(basis_ref, occ, ψr, res_schur, ρr)\n",
    "println(\"F(P) - df(P)⋅Rschur(P) = $(f[1][1][1]-df_schur[1][1][1])\")"
   ],
   "metadata": {},
   "execution_count": 18
  },
  {
   "cell_type": "markdown",
   "source": [
    "Then, we estimate the error on the forces made by the different computations\n",
    "above:"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "- Relative error on the forces with no post-processing:"
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|F(P) - F(P_*)| / |F(P_*)| = 0.4156911952038415\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "println(\"|F(P) - F(P_*)| / |F(P_*)| = $(norm(f-f_ref)/norm(f_ref))\")"
   ],
   "metadata": {},
   "execution_count": 19
  },
  {
   "cell_type": "markdown",
   "source": [
    "- Relative error made by the linearization ``F(P) - {\\rm d}F(P)⋅(P-P_*)`` if\n",
    "  we had access to the actual error ``P-P_*``, which is not the case in\n",
    "  practice (we are aiming at reaching this precision):"
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|F(P) - dF(P)⋅(P-P_*) - F(P_*)| / |F(P_*)| = 0.08590972360440276\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "println(\"|F(P) - dF(P)⋅(P-P_*) - F(P_*)| / |F(P_*)| = $(norm(f-df_err-f_ref)/norm(f_ref))\")"
   ],
   "metadata": {},
   "execution_count": 20
  },
  {
   "cell_type": "markdown",
   "source": [
    "- Relative error made by replacing ``P-P_*`` by the modified residual\n",
    "  ``R_{\\rm Schur}(P)`` (computable in practice) in the linearization (note\n",
    "  how closer we are to the previous one):"
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "|F(P) - dF(P)⋅Rschur(P) - F(P_*)| / |F(P_*)| = 0.14723660777408942\n"
     ]
    }
   ],
   "cell_type": "code",
   "source": [
    "println(\"|F(P) - dF(P)⋅Rschur(P) - F(P_*)| / |F(P_*)| = $(norm(f-df_schur-f_ref)/norm(f_ref))\")"
   ],
   "metadata": {},
   "execution_count": 21
  }
 ],
 "nbformat_minor": 3,
 "metadata": {
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.7.1"
  },
  "kernelspec": {
   "name": "julia-1.7",
   "display_name": "Julia 1.7.1",
   "language": "julia"
  }
 },
 "nbformat": 4
}