content/notebooks/H-atom-4c.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "90e862bb-783a-4f97-a227-8f6e0b4012b8",
   "metadata": {},
   "source": [
    "# The Dirac equation for the Hydrogen atom\n",
    "\n",
    "In this notebook we will illustrate how one can solve the Dirac equation for a Hydrogen atom using the Multiwavelet framework provided by VAMPyR\n",
    "\n",
    "The Dirac equation can be coincisely written as follows:\n",
    "\n",
    "\\begin{equation}\n",
    "(\\beta m c^2+ c \\boldsymbol{\\alpha} \\cdot \\mathbf{p} + V) \\phi = \\epsilon \\phi \n",
    "\\end{equation}\n",
    "\n",
    "where $\\phi = (\\phi^{L\\alpha}, \\phi^{L\\beta}, \\phi^{S\\alpha}, \\phi^{S\\beta})^t$ represents a 4-component spinor, $\\boldsymbol{\\alpha} = \n",
    "\\begin{pmatrix}\n",
    "0_2 & \\boldsymbol{\\sigma} \\\\\n",
    "\\boldsymbol{\\sigma} & 0_2 & \n",
    "\\end{pmatrix}\n",
    "$ and\n",
    "$\\beta = \n",
    "\\begin{pmatrix}\n",
    "1_2 & 0_2 \\\\\n",
    "0_2 & -1_2\n",
    "\\end{pmatrix}\n",
    "$ are the 4x4 Dirac matrices, $\\boldsymbol{\\sigma}$ is a cartesian vector collecting the three 2x2 Pauli matrices, $\\mathbf{p}$ is the momentum operator, $c$ is the speed of light, $m$ is the electron mass and $V$ is the nuclear potential.\n",
    "\n",
    "As for the non-relativistic case,  the equation is first rewritten in its integral formulation:\n",
    "$$\\phi = \\frac{1}{2mc^2}(\\beta m c^2+ c \\boldsymbol{\\alpha} \\cdot \\mathbf{p} + \\epsilon) \\left[ G_\\mu \\star (V \\psi) \\right]$$\n",
    "\n",
    "where $G_\\mu(x) = \\frac{e^{-\\mu |x|}}{4 \\pi |x|}$ is the Helmholtz Green's kernel and $\\mu = \\sqrt{\\frac{m^2c^4-\\epsilon}{mc^2}}$. An initial guess can be obtained by taking a Slater orbital or a Gaussian function for the $\\psi^{L\\alpha}$ component and then applying the restricted kinetic balance:\n",
    "\n",
    "$$\n",
    "\\begin{pmatrix}\n",
    "\\phi^{S\\alpha} \\\\\n",
    "\\phi^{S\\beta}\n",
    "\\end{pmatrix}\n",
    "= \\frac{1}{2c}\\boldsymbol{\\sigma} \\cdot \\mathbf{p} \n",
    "\\begin{pmatrix}\n",
    "\\phi^{L\\alpha} \\\\\n",
    "0\n",
    "\\end{pmatrix}\n",
    "$$\n",
    "The guess obtained is then plugged into the integral form of the Dirac equation, which can then be iterated until convergence"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54c75006-0358-42c5-9075-37a5b8a909ba",
   "metadata": {},
   "source": [
    "We start by loading the relevant packages: the 3d version of `vampyr`, `numpy`, the `complex_function` class which deals with complex funtions and the `orbital` class which deals with 4-component spinors. Each complex function is handled as a pair of `function_tree`s and each spinor is handled as a 4c vector of complex functions. The `nuclear_potential` package is self-explanatory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "37e0e6b2-886e-4415-8612-d2a652f24c4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from vampyr import vampyr3d as vp\n",
    "from orbital4c import orbital as orb\n",
    "from orbital4c import nuclear_potential as nucpot\n",
    "from orbital4c import complex_fcn as cf\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fbe4360-9936-47d4-bb6f-d51bb3b34b69",
   "metadata": {},
   "source": [
    "As a reference, the exact Dirac energy for the ground state Hydrogen atom is computed in the function below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "0ef3d827-810b-4411-8ba3-bd63271033e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Exact Energy -0.5000066565989982\n"
     ]
    }
   ],
   "source": [
    "def analytic_1s(light_speed, n, k, Z):\n",
    "    alpha = 1/light_speed\n",
    "    gamma = orb.compute_gamma(k,Z,alpha)\n",
    "    tmp1 = n - np.abs(k) + gamma\n",
    "    tmp2 = Z * alpha / tmp1\n",
    "    tmp3 = 1 + tmp2**2\n",
    "    return light_speed**2 / np.sqrt(tmp3)\n",
    "\n",
    "light_speed = 137.03599913900001\n",
    "alpha = 1/light_speed\n",
    "k = -1\n",
    "l = 0\n",
    "n = 1\n",
    "m = 0.5\n",
    "Z = 1\n",
    "atom = \"H\"\n",
    "\n",
    "energy_1s = analytic_1s(light_speed, n, k, Z)\n",
    "print('Exact Energy',energy_1s - light_speed**2, flush = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6277d776-2236-496a-bdd5-41929d11ac3d",
   "metadata": {},
   "source": [
    "The `MultiResolutionAnalysis` object defining the simulation box is constructed and passed to the classes for complex functions and spinors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5ac41574-6f43-411d-841f-f8fdf0749b66",
   "metadata": {},
   "outputs": [],
   "source": [
    "mra = vp.MultiResolutionAnalysis(box=[-30,30], order=6)\n",
    "prec = 1.0e-4\n",
    "origin = [0.1, 0.2, 0.3]  # origin moved to avoid placing the nuclar charge on a node\n",
    "\n",
    "orb.orbital4c.light_speed = light_speed\n",
    "orb.orbital4c.mra = mra\n",
    "cf.complex_fcn.mra = mra"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5abdbe4-ba5b-49d3-ae7e-febab5ced4fe",
   "metadata": {},
   "source": [
    "We construct a starting guess by taking a simple Gaussian function and initialize the real part of the $\\phi^{L\\alpha}$ component of the spinor with it. Thereafter the restricted kinetic balance is employed. This is implemented in the `init_small_components` method of the `orbital` class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "95826623-b759-4318-afc5-e0c225cb5f1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "a_coeff = 3.0\n",
    "b_coeff = np.sqrt(a_coeff/np.pi)**3\n",
    "gauss = vp.GaussFunc(b_coeff, a_coeff, origin)\n",
    "gauss_tree = vp.FunctionTree(mra)\n",
    "vp.advanced.build_grid(out=gauss_tree, inp=gauss)\n",
    "vp.advanced.project(prec=prec, out=gauss_tree, inp=gauss)\n",
    "gauss_tree.normalize()\n",
    "\n",
    "spinor_H = orb.orbital4c()\n",
    "La_comp = cf.complex_fcn()\n",
    "La_comp.copy_fcns(real = gauss_tree)\n",
    "spinor_H.copy_components(La = La_comp)\n",
    "spinor_H.init_small_components(prec/10)\n",
    "spinor_H.normalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6a40ddd-7933-46ee-8128-96caf89f3682",
   "metadata": {},
   "source": [
    "The nuclear potential is defined and projected onto the `V_tree`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "4f0cda6f-217c-4941-8ce5-fa52576c2d98",
   "metadata": {},
   "outputs": [],
   "source": [
    "Peps = vp.ScalingProjector(mra, prec)\n",
    "f = lambda x: nucpot.coulomb_HFYGB(x, origin, Z, prec)\n",
    "V_tree = Peps(f)\n",
    "\n",
    "default_der = 'BS'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1678bc97-d85d-4a48-9aea-13901631534d",
   "metadata": {},
   "source": [
    "The orbital is optimized by iterating the integral version of the Dirac equation as follows:\n",
    "1. Application of the Dirac Hamiltonian $f^n = \\hat{h}_D \\phi^n = (\\beta m c^2+ c \\boldsymbol{\\alpha} \\cdot \\mathbf{p}) \\phi^n$\n",
    "2. Application of the potnetial operator $g^n = \\hat{V} \\phi^n$\n",
    "3. Sum $h^n = f^n + g^n$\n",
    "4. Expectation value of the energy $\\left\\langle \\phi^n | h^n \\right\\rangle$\n",
    "5. Calculation of the Helmholtz parameter $\\mu$\n",
    "6. Convolution with the Helmholtz kernel $t^n = G_\\mu \\star g^n$\n",
    "7. application of the shifted Dirac Hamiltonian $\\tilde{\\phi}^{n+1} = (\\hat{h}_D + \\epsilon) t^n$\n",
    "8. normalization of the new iterate\n",
    "9. calculation of the change in the orbital $\\delta \\phi^n = \\phi^{n+1}-\\phi^n$\n",
    "\n",
    "Once the orbital error is below the requested threshold the iteration is interrupted and the final energy is computed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "40a443b6-33f4-4771-8b06-ccd12c9da35d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Energy -0.14180720307558659\n",
      "Error 0.3648353655184581\n",
      "Energy -0.49611934473068686\n",
      "Error 0.0024839139737050653\n",
      "Energy -0.4997319277244969\n",
      "Error 0.00023304055836475225\n",
      "Energy -0.49998214000152075\n",
      "Error 2.3987883028320866e-05\n",
      "Final Energy -0.5000042447172746\n",
      "Exact Energy -0.5000066565989982\n",
      "Difference -2.411881723674014e-06\n"
     ]
    }
   ],
   "source": [
    "orbital_error = 1\n",
    "while orbital_error > prec:\n",
    "    # 1. Application of the Dirac Hamiltonian\n",
    "    hd_psi = orb.apply_dirac_hamiltonian(spinor_H, prec, der = default_der)\n",
    "    # 2. Application of the potnetial operator\n",
    "    v_psi = orb.apply_potential(-1.0, V_tree, spinor_H, prec)\n",
    "    # 3. Sum\n",
    "    add_psi = hd_psi + v_psi\n",
    "    # 4. Expectation value of the energy\n",
    "    energy = (spinor_H.dot(add_psi)).real\n",
    "    print('Energy',energy-light_speed**2)\n",
    "    # 5. Calculation of the Helmholtz parameter\n",
    "    mu = orb.calc_dirac_mu(energy, light_speed)\n",
    "    # 6. Convolution with the Helmholtz kernel\n",
    "    tmp = orb.apply_helmholtz(v_psi, mu, prec)\n",
    "    tmp.crop(prec/10)\n",
    "    # 7. application of the shifted Dirac Hamiltonian\n",
    "    new_orbital = orb.apply_dirac_hamiltonian(tmp, prec, energy, der = default_der) \n",
    "    new_orbital.crop(prec/10)\n",
    "    # 8. normalization of the new iterate\n",
    "    new_orbital.normalize()\n",
    "    delta_psi = new_orbital - spinor_H\n",
    "    # 9. calculation of the change in the orbital\n",
    "    orbital_error = (delta_psi.dot(delta_psi)).real\n",
    "    print('Error',orbital_error, flush = True)\n",
    "    spinor_H = new_orbital\n",
    "\n",
    "# Computing the final energy\n",
    "hd_psi = orb.apply_dirac_hamiltonian(spinor_H, prec, der = default_der)\n",
    "v_psi = orb.apply_potential(-1.0, V_tree, spinor_H, prec)\n",
    "add_psi = hd_psi + v_psi\n",
    "energy = (spinor_H.dot(add_psi)).real\n",
    "print('Final Energy',energy - light_speed**2)\n",
    "\n",
    "energy_1s = analytic_1s(light_speed, n, k, Z)\n",
    "\n",
    "print('Exact Energy',energy_1s - light_speed**2)\n",
    "print('Difference',energy_1s - energy)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}