imaginary_time_examples.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from scipy.special import erf\n",
    "from scipy.sparse import spdiags\n",
    "from scipy.sparse.linalg import eigsh\n",
    "import time\n",
    "import tqdm\n",
    "from Laplacians import CylinderLaplacian\n",
    "from poisson_solvers import PoissonCylindrical\n",
    "from FockMatrices import CylinderFockMatrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Verfgau(d, Z=1.0, mu=1.0):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "    d: d=|r-A| where A is the nucleus position\n",
    "    Z: nuclear charge\n",
    "    mu: screening parameter\n",
    "    \"\"\"\n",
    "    c = 0.923 + 1.568 * mu\n",
    "    alpha = 0.2411 + 1.405 * mu\n",
    "    long_range = erf(mu * Z * d) / (Z * d)\n",
    "    return -(Z**2) * (\n",
    "        c * np.exp(-(alpha**2) * Z**2 * d**2) + long_range\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "rho_max = 15\n",
    "drho = 0.1\n",
    "n_rho = int(rho_max / drho) - 1\n",
    "rho = np.linspace(0, rho_max, n_rho + 2)\n",
    "rho_inner = rho[1:-1]\n",
    "\n",
    "z_max = 15\n",
    "dz = 0.1\n",
    "n_z = int(2 * z_max / dz) - 1\n",
    "z = np.linspace(-z_max, z_max, n_z + 2)\n",
    "z_inner = z[1:-1]\n",
    "\n",
    "Laplacian = CylinderLaplacian(rho, z, m=0)\n",
    "Poisson = PoissonCylindrical(Laplacian, rho, z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 9999/9999 [00:19<00:00, 520.55it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Energy imaginary time : -0.49825781\n",
      "Energy diagonalization: -0.49826612\n",
      "E_imag-E_diag         : 0.00000832\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAZoAAAEGCAYAAABcolNbAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/OQEPoAAAACXBIWXMAAAsTAAALEwEAmpwYAAAl6UlEQVR4nO3de5hcZZ3u/e9dfUx3p5NO0kk6Z0LCIUA4pBlBBdQEXw8Mwa2ydRg3OjpsZYvOzGacsPe13/Ewe4vDdkZH1JHBGeK8iigyF4igQFBhwEE6HBIgQBIIpJPOOeR86MPv/aNWJ5Wmk3QnVbUqXffnuuqqdXhq1W8p9M3zrFXPUkRgZmZWKJm0CzAzs6HNQWNmZgXloDEzs4Jy0JiZWUE5aMzMrKAq0y6g1IwZMyamTZuWdhlmZieUxYsXb4qI5v72OWj6mDZtGm1tbWmXYWZ2QpH02uH2eejMzMwKykFjZmYF5aAxM7OCctCYmVlBOWjMzKygHDRmZlZQDhozMysoB02ePLlqCzf96kW6e/zYBTOzXA6aPHnm9Tf49q9Xsmt/V9qlmJmVFAdNnjTUZidZ2LXPQWNmlstBkyf1NQ4aM7P+OGjypKGmAoAdex00Zma5HDR50lBTBcCufd0pV2JmVlocNHlSn/RodnrozMzsEA6aPBme9GgcNGZmh3LQ5EnjsOzNANv2dKZciZlZaXHQ5EljbRWVGbF55760SzEzKykOmjzJZMSo+mo279yfdilmZiXFQZNHYxpq2OQejZnZIVIPGkmjJD0oaXny3nSEto2S2iXdnKwPl/RMzmuTpG8k+2ok3SFphaQnJE0r9LmMGe6gMTPrK/WgARYAiyJiJrAoWT+crwCP9K5ExI6IOKf3BbwG3JXs/iSwNSJmAH8PfK0QxecaU1/NJg+dmZkdohSCZj6wMFleCFzRXyNJc4BxwAOH2X8KMBZ4tJ/j3gnMlaT8lNy/3h5NhGdwNjPrVQpBMy4iOpLldWTD5BCSMsDXgeuPcJyPAHfEwb/yE4HVABHRBWwDRvf3QUnXSGqT1LZx48ZjOwtgfGMt+7p62LrbtzibmfWqLMaXSHoIGN/Prv+ZuxIRIam/7sC1wH0R0X6ETslHgI8dS30RcQtwC0Bra+sxd0cmjBwGwNo39jCqvvpYD2NmNqQUJWgiYt7h9klaL6klIjoktQAb+ml2IXCRpGuBBqBa0s6IWJAc42ygMiIW53xmDTAZaJdUCYwANufplPo1qSkbNO1b93DmxBGF/CozsxNGKQyd3QNcnSxfDdzdt0FEXBURUyJiGtnhsx/0hkzio8DtRzjuh4CHo8AXT3J7NGZmllUKQXMjcKmk5cC8ZB1JrZJuHeAxruTNQfN9YLSkFcBfcOS72fKiqa6KYVUVrHHQmJkdUJShsyOJiM3A3H62twGf6mf7bcBtfbZN76fdXuDD+apzICQxYWQta7Y6aMzMepVCj2ZImdhUx9ptDhozs14OmjybOLKWdvdozMwOcNDk2ZRR9WzZtZ8de/1bGjMzcNDk3bTRdQC8tnl3ypWYmZUGB02eTR1dD8CqzbtSrsTMrDQ4aPJs2hj3aMzMcjlo8qyuupKxw2t4dZN7NGZm4KApiGmj63nNQ2dmZoCDpiCmjaljlYfOzMwAB01BTB1dz8Yd+9i5ryvtUszMUuegKYCTxmTvPPPwmZmZg6Ygpia/pVm1ycNnZmYOmgLo7dG8snFnypWYmaXPQVMAddWVTBw5jOUbHDRmZg6aApkxtoEVDhozMwdNocwc28DKjTvp7inoQz3NzEqeg6ZAZoxtYF9Xjx+CZmZlz0FTIDPHNQCwfMOOlCsxM0uXg6ZAZjQPB/B1GjMrew6aAhlRV0Xz8BrfeWZmZc9BU0Azmn3nmZmZg6aAZo7LBk2E7zwzs/LloCmgGWMb2Lmvi/Xb96VdiplZalIPGkmjJD0oaXny3nSEto2S2iXdnKwPl/RMzmuTpG8k+z4uaWPOvk8V6ZQOmDE2e+fZS+t955mZla/UgwZYACyKiJnAomT9cL4CPNK7EhE7IuKc3hfwGnBXTvs7cvbfWoDaj+j08Y0AvLRue7G/2sysZJRC0MwHFibLC4Er+mskaQ4wDnjgMPtPAcYCj+a/xGPTVF/N+MZalnW4R2Nm5asUgmZcRHQky+vIhskhJGWArwPXH+E4HyHbg8m98v5BSUsk3Slp8uE+KOkaSW2S2jZu3HgMp3B4p7UMZ1mHezRmVr6KEjSSHpL0XD+v+bntkpDo7xata4H7IqL9CF/zEeD2nPWfA9MiYjbwIAd7TW8SEbdERGtEtDY3Nw/4vAbi9JZGVm7cyf6unrwe18zsRFFZjC+JiHmH2ydpvaSWiOiQ1AJs6KfZhcBFkq4FGoBqSTsjYkFyjLOByohYnPOdm3M+fyvwt/k4l8E6bfxwOruDFRt2MmtCYxolmJmlqhSGzu4Brk6Wrwbu7tsgIq6KiCkRMY3s8NkPekMm8VEO7c2QhFavy4Fl+Sx6oGa1ZMPlRd8QYGZlqig9mqO4EfiJpE+SvWvsSgBJrcCnI2IgtyVfCbyvz7bPSboc6AK2AB/PW8WDcNKYeqorM75OY2ZlK/WgSYa45vazvQ14U8hExG3AbX22Te+n3Q3ADfmq81hVVmQ4ZVwDL67znWdmVp5KYehsyDt9fKN7NGZWthw0RXB6SyObdu5nw469aZdiZlZ0DpoiOD25IeCFte7VmFn5cdAUwRkTs0Hz3JptKVdiZlZ8DpoiaKytYvqYepa0O2jMrPw4aIrkrEkjWOoejZmVIQdNkcyeNJKObXt9Q4CZlR0HTZHMnjQCgKUePjOzMuOgKZJZLY1khK/TmFnZcdAUSX1NJTPGNvg6jZmVHQdNEZ01cSRL2rdx6CNzzMyGNgdNEc2eNIJNO/exbrtvCDCz8uGgKaLeGwKeXe3hMzMrHw6aIjq9pZGqCvHM6jfSLsXMrGgcNEVUW1XBmRNH8NRrW9MuxcysaBw0RTZnShPPtr/B/q6etEsxMysKB02RzZnaxL6uHl7w82nMrEw4aIrsvKlNACz28JmZlQkHTZGNa6xlUtMwX6cxs7LhoEnBeVOaaHtti3+4aWZlwUGTgjlTm1i/fR9rt/mHm2Y29DloUjDH12nMrIw4aFJw2vjh1FVX0LZqS9qlmJkVnIMmBZUVGeZMbeL3rzpozGzoK4mgkTRK0oOSlifvTUdo2yipXdLNOds+KmmppCWSfilpzGCPW2wXTB/Ni+t2sHnnvrRLMTMrqJIIGmABsCgiZgKLkvXD+QrwSO+KpErgm8A7I2I2sAT47DEct6gumD4KwL0aMxvySiVo5gMLk+WFwBX9NZI0BxgHPJC7OXnVSxLQCKwdzHHTMHvSSIZVVfAfr2xOuxQzs4IqlaAZFxEdyfI6smFyCEkZ4OvA9bnbI6IT+AywlGzAzAK+P9DjJse+RlKbpLaNGzce77kMSFVFhtZpTfzOQWNmQ1zRgkbSQ5Ke6+c1P7ddZH/F2N8vGa8F7ouI9j7HrSIbNOcCE8gOnd3Q98NHOC4RcUtEtEZEa3Nz8zGd37G4YPpoXl6/k02+TmNmQ1hlsb4oIuYdbp+k9ZJaIqJDUguwoZ9mFwIXSboWaACqJe0EfpYcf2VyrJ9w8FrMQI6bmgtPHg1kr9O876yWlKsxMyuMUhk6uwe4Olm+Gri7b4OIuCoipkTENLLDZz+IiAXAGmCWpN6uyKXAsoEeN01nTRxBXXUFv1vp4TMzG7pKJWhuBC6VtByYl6wjqVXSrUf6YESsBb4EPCJpCXAO8H+OdNxSUVWR4fxpo3ydxsyGNHlix0O1trZGW1tb0b7ve79dyVfvf5H/uGEu40fUFu17zczySdLiiGjtb1+p9GjK1sWnZEf8Hnm5OHe7mZkVm4MmZaeNH87Y4TX8drmDxsyGJgdNyiRx0cxm/n35Jrp7PIxpZkOPg6YEXHJqM9v2dLKk/Y20SzEzyzsHTQm4aMYYJHjk5U1pl2JmlncOmhLQVF/N7Ikj+O3LJfV7UjOzvBh00Eiql1RRiGLK2cWnNPPM6jfYtrsz7VLMzPLqqEEjKSPpjyT9QtIG4EWgQ9ILkm6SNKPwZQ59l5zSTE/Aoyt895mZDS0D6dH8GjiZ7ESV4yNickSMBd4O/AfwNUl/XMAay8I5k0fSVFfFQy+sT7sUM7O8GsikmvOSqfgPERFbyE5o+bNkBmU7DpUVGd512jgefGEdnd09VFX48pmZDQ1H/WvWGzKSHj9aGzs+l84ay/a9XbSt2pp2KWZmeTOY/2x+00Rcki7KYy1l76KZzVRXZnjQw2dmNoQMJmhOlfRvkv5G0kckvRO4rUB1laX6mkredvJoHly2Dk92amZDxWCC5lWy0++vBOYAnyI7Pb/l0aWzxrN6yx5eXr8z7VLMzPJiME/Y3B8RTwJPFqoYg7mnj4V/g4eWrefU8cPTLsfM7LgNpkdzScGqsAPGNdZy9uSR/PK5dWmXYmaWFwP5waYAImLH0dpYfrz/rPEsXbON1zfvTrsUM7PjNqAfbEq6TtKU3I2SqiW9S9JC4OrClFee3ndWCwD3Ll2bciVmZsdvIEHzHqAbuF3S2mTqmVeA5cBHgW9ExG0FrLHsTGqq49wpI/nFko60SzEzO24DuRng5Ij4DvCdZAaAMcCeiHijoJWVufef1cLf/GIZr27axUlj6tMux8zsmA2kR/OvOctXR0RHb8hIqitIVcb7Z2eHz36xxMNnZnZiG0jQ5F7ov7bPvkfzWIvlaBkxjNapTdzr4TMzO8ENJGhyf6Le9+4yz/xYQJfNbuHFdTtYvv6wN/yZmZW8gQTFeEkfl3Qubw6a454nRdIoSQ9KWp68Nx2hbaOkdkk352z7qKSlkpZI+qWkMcn2L0paI+mZ5PW+46212N43u4WKjLjr6TVpl2JmdswGEjRfJDvlzDeAScldZz+T9L/J3hhwvBYAiyJiJrAoWT+crwCP9K5IqgS+CbwzImYDS4DP5rT/+4g4J3ndl4dai2rs8FounjmGf3tqDd09nvvMzE5MA3lMwC0RcV1EXBIRY4B3A7cC28n5o38c5gMLk+WFwBX9NZI0BxgHPJC7OXnVJz8abQSG1NXzD86ZxLrte3l85aa0SzEzOyaDvsYSEe0RcX9EfC0i8vFkzXER0XvFex3ZMDmEpAzwdeD6PrV0Ap8BlpINmFnA93OafDYZUvvnIw3JlbJ5p4+jsbaSny1uT7sUM7NjUpSL+ZIekvRcP6/5ue0iOzd+f2NE1wL3RcQhf22T3/V8BjgXmEB26OyGZPd3yT6C+hygg2xQHa6+ayS1SWrbuHHjsZ1kgdRWVXDZ2RP45fPr2LHXz5czsxPPYGZvPmYRMe9w+yStl9QSER2SWoAN/TS7ELhI0rVAA1AtaSfZR0kTESuTY/2E5BpPRBx4epikfwLuPUJ9twC3ALS2tpbcxZAPnjeJHz3xOvcvXceV509Ouxwzs0EphduT7+HgXGlXA3f3bRARV0XElIiYRnb47AcRsQBYA8yS1Jw0vRRYBpCEVq8PAM8VpvzCO2/KSKaPqeeni1enXYqZ2aCVQtDcCFwqaTkwL1lHUqukW4/0wYhYS/bha49IWkJ2mOz/JLv/tve2Z+CdwJ8XqP6Ck8SV50/myVVb/ZsaMzvhyI8MPlRra2u0tbWlXcabbN65jwu/+jB/9JYpfPHyM9Iux8zsEJIWR0Rrf/tKoUdjAzC6oYb3njWenz3Vzu79XWmXY2Y2YA6aE8hVb5nKjr1d/PzZIfVTITMb4hw0J5DzpzVxyrgGfvjE62mXYmY2YA6aE4gkrnrLVJa0b2NJ+xtpl2NmNiAOmhPMB86bSH11Bf/y2Kq0SzEzGxAHzQmmsbaKK8+fzM+fXcu6bXvTLsfM7KgcNCegP3nbSfREcNvjq9IuxczsqBw0J6DJo+p4z5nj+dETr7Frn291NrPS5qA5QX3y7dPZvreLn7Z5WhozK20OmhPUnKlNnDdlJN9/7FU/FM3MSpqD5gR2zcXTWb1lD/cu8Q84zax0OWhOYO+eNZ5TxjXwrYdXuFdjZiXLQXMCy2TEde+ayYoNO7n/uY6jf8DMLAUOmhPc+85q4eTmem5+eAU97tWYWQly0JzgKpJezYvrdvDAC+uP/gEzsyJz0AwBl81u4aQx9Xxz0XL3asys5DhohoDKigx/Nm8myzq2c48fIWBmJcZBM0T84ewJzGpp5P8+8BL7urrTLsfM7AAHzRCRyYgF7z2N9q17+OF/+Hk1ZlY6HDRDyEUzx/C2GaP51sPL2b63M+1yzMwAB82QIom/es9pbN3dyT/+ZmXa5ZiZAQ6aIWf2pJFccc4Ebn30VV7bvCvtcszMHDRD0Q3vO52qCvGVe19IuxQzMwfNUDSusZbr5s7koWUb+PWLG9Iux8zKXOpBI2mUpAclLU/em47QtlFSu6Sbc7b9Z0lLJD0v6Ws522sk3SFphaQnJE0r8KmUlD9520lMH1PPl37+vG93NrNUpR40wAJgUUTMBBYl64fzFeCR3hVJo4GbgLkRcQYwXtLcZPcnga0RMQP4e+BrfQ82lFVXZvjry89g1ebdfO+3r6RdjpmVsVIImvnAwmR5IXBFf40kzQHGAQ/kbJ4OLI+Ijcn6Q8AH+znuncBcScpf2aXvklOauWx2Czc/vIIVG3akXY6ZlalSCJpxEdE7x/06smFyCEkZ4OvA9X12rQBOlTRNUiXZkJqc7JsIrAaIiC5gGzC6vwIkXSOpTVLbxo0b+2tywvri5WdQV1PBF+5c4mfWmFkqihI0kh6S9Fw/r/m57SIigP7+Gl4L3BcR7X3abwU+A9wBPAqsAgZ9QSIibomI1ohobW5uHuzHS9qYhhr++g9n8dTrb/Cvv1uVdjlmVoYqi/ElETHvcPskrZfUEhEdklqA/m6TuhC4SNK1QANQLWlnRCyIiJ8DP0+OdQ0Hg2YN2d5Ne9LbGQFszt9ZnTiuOGcidz+zlr/91UvMPX0ck0fVpV2SmZWRUhg6uwe4Olm+Gri7b4OIuCoipkTENLLDZz+IiAUAksYm701kez639nPcDwEPJz2msiOJ//2Bs6iQ+PM7nqGruyftksysjJRC0NwIXCppOTAvWUdSq6Rbj/jJrG9KegF4DLgxIl5Otn8fGC1pBfAXHPlutiFv4shhfOWKM2l7bSvf9fQ0ZlZEKtP/yD+s1tbWaGtrS7uMgvn8j5/m3iUd3PnpCzl3ymF/smRmNiiSFkdEa3/7SqFHY0X05flnMr6xlj+74xl27utKuxwzKwMOmjIzYlgV3/jIOazespsv3Pks7tGaWaE5aMrQ+dNG8VfvOY37lq7j1kdfTbscMxviHDRl6pqLp/PeM8dz4y9f5Hcry/KubzMrEgdNmZLETR8+m2mj67ju9qfo2LYn7ZLMbIhy0JSxhppKvvexOezt7OGTt7X55gAzKwgHTZmbMXY4377qPF5av4PrfvSUf8xpZnnnoDEuOaWZL88/g1+/tJEv3/uC70Qzs7wqylxnVvquestUXtu8m1seeYUJI4fx6UtOTrskMxsiHDR2wIL3nEbHtr3ceP+LNNRU8scXTE27JDMbAhw0dkAmI/7uyrPZva+L/3X3c9TXVPCBcyelXZaZneB8jcYOUVWR4dtXnccFJ43m+p8u4ZfPdRz9Q2ZmR+CgsTeprargn65uZfakEfy3Hz3NPc+uTbskMzuBOWisXw01lfzrJ9/CnKlNfP7HT/OTttVpl2RmJygHjR1WQ00lCz/xB7x9xhi+cOcSFj6+Ku2SzOwE5KCxIxpWXcGtV7dy6axx/PU9z3Pj/S/S0+Pf2ZjZwDlo7KhqKiv47lXncdVbpvCPv13J5378NHs7u9Muy8xOEL692QaksiLD31xxJlNG1fHV+19k3ba9fO9jcxjdUJN2aWZW4tyjsQGTxH+95GS+/UfnsWTNNi771r/zzOo30i7LzEqcg8YG7f2zW7jrM28lI3HlP/6O23//etolmVkJc9DYMTlz4gjuve7tvGX6KG64aynX//RZdvkxA2bWDweNHbOm+mpu+8Qf8Ll3zeBnT7Xz/n941ENpZvYmDho7LhUZ8RfvPpUf/+kFdHYHH/zu4/zDouV+ro2ZHeCgsbx4y/TR3Pf5i7hsdgt/9+DLfOA7j/P82m1pl2VmJSD1oJE0StKDkpYn701HaNsoqV3SzTnb/rOkJZKel/S1nO0fl7RR0jPJ61OFPpdyN2JYFd/8yLnc/Efn0rFtL5ff/BhfvX8Ze/b7Nzdm5Sz1oAEWAIsiYiawKFk/nK8Aj/SuSBoN3ATMjYgzgPGS5ua0vyMizkletxagduvHZbMnsOgvLuHDcybxvd++wv/zjUd44Pl1fnKnWZkqhaCZDyxMlhcCV/TXSNIcYBzwQM7m6cDyiNiYrD8EfLAwZdpgjKir4sYPzub2P72A6soM1/zrYq669QmWdWxPuzQzK7JSCJpxEdH70JN1ZMPkEJIywNeB6/vsWgGcKmmapEqyITU5Z/8Hk2G1OyVN5jAkXSOpTVLbxo0bD9fMjsGFJ4/m/s9fxJfnn8ELHdt5/z88yg13LaFj2560SzOzIilK0Eh6SNJz/bzm57aL7NhKf+Mr1wL3RUR7n/Zbgc8AdwCPAquA3gsCPwemRcRs4EEO9preJCJuiYjWiGhtbm4+xrO0w6mqyPBfLpzGb69/J59420ncubidS276DV+853k2bN+bdnlmVmBKe9xc0kvAOyKiQ1IL8JuIOLVPmx8CFwE9QANQDXwnIhb0aXcNMCMivtBnewWwJSJGHK2e1tbWaGtrO65zsiNbvWU33/71Cn66uJ3KjPjjC6bypxdNZ/yI2rRLM7NjJGlxRLT2u68EguYmYHNE3ChpATCqb1D0af9xoDUiPpusj42IDcndar8GroyIlyW19A7JSfoA8FcRccHR6nHQFM/rm3fzrYeXc9fTaxBw+dkT+NRF05k1oTHt0sxskI4UNKUwe/ONwE8kfRJ4DbgSQFIr8OmIONptyd+UdHay/OWIeDlZ/pyky4EuYAvw8bxXbsdlyug6bvrw2Xxu7kz++bFXuePJ1dz19BreNmM0n3jrSbzj1GYqK0rhMqKZHY/UezSlxj2a9Gzb08ntv3+df3nsVdZv38f4xlqubJ3EledPZlJTXdrlmdkRlPTQWalx0KSvs7uHRcs2cPvvX+eR5dm7AC+e2cx/Om8i804fR31NKXTEzSyXg2YQHDSlpX3rbn7y5Gp+uridjm17GVZVwdzTx3L52RO45NRmaior0i7RzHDQDIqDpjT19ARPrtrCPc+u5b6lHWzd3cnw2krecepY5p0+lnecMpYRdVVpl2lWthw0g+CgKX2d3T08tmITv1jSwcMvbmDzrv1UZMT505qYd/o4Lj6lmZljG5CUdqlmZcNBMwgOmhNLd0/wzOo3WLRsPYuWbeCl9TsAaB5ew1tPHp28xjB5lG8mMCskB80gOGhObKu37ObxlZt4bMVmHl+5mU079wEwqWkYc6Y2cd6U7Ou0luFU+dZps7xx0AyCg2boiAiWb9jJYys28cQrW3jq9a1s2JENntqqDLMnjeTcKSM5Y8IIZrU0ctKYeioyHm4zOxYOmkFw0AxdEcGaN/bw1Otv8NRrW3n69a280LGdzu7svwO1VRlOG9/IrAmNzGpp5LTxwzm5uYGm+uqUKzcrfQ6aQXDQlJd9Xd2s2LCTF9Zu54WO7Qfed+ztOtBmVH0108fUc3JzAyePzb6fNKaeiU3DfHu1WaLUp6AxS01NZQVnTBjBGRMOzrcaEbRv3cOKDTtZubH3tYtFL67njrb9h3x+XGMNk5vqmNQ0jMmjsu+TmuqYOHIY4xprGVbtIDJz0Jj1IYnJo+qYPKqOd5429pB9b+zez8qNu1i1aRftW/eweutu2rfu5slVW7nn2bX09BkgaKytZFxjLeMaaxnbWJNdHp59bx5eQ1N9NaPqqhkxrIqMrw/ZEOWgMRuEkXXVzJlazZypTW/a19ndw7pte1m9dTdrtu5hw459bNi+l/Xb97F+x16eeGUXG3bsPXBNKFdG0FRXfSB4RtUny/VVNNVVM7y2ksbaKobXVjG8tjJ5ZZdrq9xrstLmoDHLk6qKzIGe0OH09ARbd+9n/fZ9bNixlzd2d7Jl13627t7P5l372bprP1t27eeVTTvZ8lonW3fvp7tvN6mP6soMjTnBU19dSV11BcOqK6irrqCuujK7XJXd1rt9WFVlsj+7rbaqguqKDDVVGWoqKqiuzFBdmfGdeHbcHDRmRZTJiNENNYxuqGEWR3/uTk9PsGNfFzv2drJjbxc79naxfU8nO/blrO/tZPueg2127euiY1snezq72b2/i937u9nb2d1vT2ogKjOiJgmdmsqDAVRzyHs2pKoqRGVFhsqMsq8KUZnJhlVVhajIZJJ3UZW0613ur03vZzOCjEQmWa6QkHRwX0ZkpGR7tu3h9lUk65lMcszefRkOtBHZd+CQdZH9jIRnnhgEB41ZCctkxIhhVYwYdvzzuHV297B7fzd79h8MoGwYdbNnfxf7unrY19XD/kPeuw9ZP7Ctu4d9nT0H3rft6WR/Vw9d3T109wSdPT10dwedPZFd7+6hqzsO7BtKN7v2BpCk5D0bTBzYfmhQ9bYjd72fY3DIZ958jAPf3RuIHNx2oLYDNeqQdXTo/t42n587kz88e0K+/qc5wEFjViaqKjKMGJbJS2gdr57eMOoJOpMA6uruoasn6OoOunqyy53d2VDq7gl6IugJsu89QXdEn31BTw8Hl/vZl/1M0J20yy7nHDf5XAQEcSAQI/muoPdz2WUiCA5tn7vOgfU37zvk2Ic5Br3rRzh+ttXBY/U6uK3/NodkfbJSqH82HDRmVnSZjKjJ+CaGcuHJnszMrKAcNGZmVlAOGjMzKygHjZmZFZSDxszMCspBY2ZmBeWgMTOzgnLQmJlZQfnBZ31I2gi8dowfHwNsymM5JwKfc3nwOZeH4znnqRHR3N8OB00eSWo73BPmhiqfc3nwOZeHQp2zh87MzKygHDRmZlZQDpr8uiXtAlLgcy4PPufyUJBz9jUaMzMrKPdozMysoBw0ZmZWUA6aPJH0HkkvSVohaUHa9RSapMmSfi3pBUnPS/p82jUVg6QKSU9LujftWopB0khJd0p6UdIySRemXVOhSfrz5J/p5yTdLqk27ZryTdI/S9og6bmcbaMkPShpefLelK/vc9DkgaQK4NvAe4FZwEclzUq3qoLrAv57RMwCLgD+WxmcM8DngWVpF1FE3wR+GRGnAWczxM9d0kTgc0BrRJwJVAAfSbeqgrgNeE+fbQuARRExE1iUrOeFgyY//gBYERGvRMR+4MfA/JRrKqiI6IiIp5LlHWT/AE1Mt6rCkjQJeD9wa9q1FIOkEcDFwPcBImJ/RLyRalHFUQkMk1QJ1AFrU64n7yLiEWBLn83zgYXJ8kLginx9n4MmPyYCq3PW2xnif3RzSZoGnAs8kXIphfYN4AtAT8p1FMtJwEbgX5Lhwlsl1addVCFFxBrg/wKvAx3Atoh4IN2qimZcRHQky+uAcfk6sIPGjoukBuBnwJ9FxPa06ykUSZcBGyJicdq1FFElcB7w3Yg4F9hFHodTSlFyXWI+2ZCdANRL+uN0qyq+yP7uJW+/fXHQ5McaYHLO+qRk25AmqYpsyPwwIu5Ku54CextwuaRVZIdG3yXp/0u3pIJrB9ojoreneifZ4BnK5gGvRsTGiOgE7gLemnJNxbJeUgtA8r4hXwd20OTHk8BMSSdJqiZ78fCelGsqKEkiO3a/LCL+Lu16Ci0iboiISRExjez/vw9HxJD+L92IWAeslnRqsmku8EKKJRXD68AFkuqSf8bnMsRvgMhxD3B1snw1cHe+DlyZrwOVs4jokvRZ4Fdk71L554h4PuWyCu1twMeApZKeSbb9j4i4L72SrACuA36Y/AfUK8AnUq6noCLiCUl3Ak+RvbPyaYbgVDSSbgfeAYyR1A78NXAj8BNJnyT7qJQr8/Z9noLGzMwKyUNnZmZWUA4aMzMrKAeNmZkVlIPGzMwKykFjZmYF5aAx64eknUX4jscL/R19vu+K3IlPJX1Z0rxi1mDlybc3m/VD0s6IaEi7joGQVBkRXQNodxtwb0TcWfiqzA5yj8bsCCS9Q9JvJd0t6RVJN0q6StLvJS2VdHLS7g8lPZFMPvmQpHHJ9ubk2R7PJ5NSviZpTLJvZ853/CbnuS8/TH6VjqT/V9KTybNRbsnZ/htJ35DUBvxPSa8mUwIhqTF3Pdn2VuBy4CZJz0g6WdJtkj6U7F8l6avJvjZJ50n6laSVkj6dc5y/TOpZIulLRfi/wIYAB43Z0Z0NfBo4nexsCKdExB+QfVzAdUmbfwcuSCaf/DHZWZ4h+4vrhyPiDLJzhU05zHecC/wZ2ecZTSc78wLAzRFxfvJslGHAZTmfqY6I1oj4EvAbso8wgOwUOXclc3UBEBGPk51i5C8j4pyIWNlPDa9HxDnAo2SfV/Ihss8a+hKApHcDM8k+FuMcYI6kiw9zPmYHOGjMju7J5Pk7+4CVQO+08UuBacnyJOBXkpYCfwmckWx/O9ngISJ+CWw9zHf8PiLaI6IHeCbnuO9MekpLgXflHBfgjpzlWzk4PcwngH8Z5DnCwfn5lgJPRMSOiNgI7JM0Enh38nqa7BQtp5ENHrMj8lxnZke3L2e5J2e9h4P/Dn0L+LuIuEfSO4AvHsd3dAOVySOEv0P2aY+rJX0RyH2s8K7ehYh4TNK05LsrIuI5Bi/3vPqecyUg4KsR8b1jOLaVMfdozPJjBAcfDXF1zvbHSCYnTIaeBvMc9t5Q2ZQ89+dDR2n/A+BHHL43swMYPojv7+tXwJ8ktSBpoqSxx3E8KxMOGrP8+CLwU0mLgU05278EvFvSc8CHyT65cMdADpg8NvmfgOfI/pF/8igf+SHZILv9MPt/DPxlcsPCyQOpoU89D5ANst8lQ3l3cnzBZWXCtzebFZCkGqA7eZTEhWSfVnlOgb7rQ8D8iPhYIY5vdqx8jcassKaQfcZHBtgP/GkhvkTSt4D3Au8rxPHNjod7NGZmVlC+RmNmZgXloDEzs4Jy0JiZWUE5aMzMrKAcNGZmVlD/P0/37VNnej41AAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\"\"\"\n",
    "Solve for the Hydrogen atom using imaginary time propagation and compare with exact diagonalization.\n",
    "\"\"\"\n",
    "V = np.zeros((n_rho, n_z))\n",
    "psi0 = np.zeros((n_rho, n_z))\n",
    "# Setup potential and initial guess for the orbitals.\n",
    "for i in range(n_rho):\n",
    "    for j in range(n_z):\n",
    "        r = np.sqrt(rho_inner[i] ** 2 + z_inner[j] ** 2)\n",
    "        V[i, j] = Verfgau(r, Z=1, mu=1)\n",
    "        psi0[i, j] = np.sqrt(rho_inner[i]) * np.exp(-0.75 * r)\n",
    "\n",
    "V = V.flatten()\n",
    "H = -0.5 * Laplacian.L + spdiags([V], diags=[0])\n",
    "# Normalize the initial guess.\n",
    "psi0 = psi0.ravel()\n",
    "norm_psi0 = np.sum(np.abs(psi0) ** 2) * drho * dz\n",
    "psi0 /= np.sqrt(norm_psi0)\n",
    "\n",
    "dt = 0.001\n",
    "n_t = int(1e4)\n",
    "imag_time = np.linspace(0, n_t * dt, n_t)\n",
    "\n",
    "energy = np.zeros(n_t)\n",
    "\n",
    "for n in tqdm.tqdm(range(n_t - 1)):\n",
    "    energy[n] = np.sum(psi0 * H.dot(psi0)) * drho * dz\n",
    "    k1 = -H.dot(psi0)\n",
    "    k2 = -H.dot(psi0 + 0.5 * dt * k1)\n",
    "    k3 = -H.dot(psi0 + 0.5 * dt * k2)\n",
    "    k4 = -H.dot(psi0 + dt * k3)\n",
    "    psi0 += dt / 6 * (k1 + 2 * k2 + 2 * k3 + k4)\n",
    "    psi0 /= np.sqrt(np.sum(np.abs(psi0) ** 2) * drho * dz)\n",
    "\n",
    "energy[n_t - 1] = np.sum(psi0 * H.dot(psi0)) * drho * dz\n",
    "\n",
    "# Solve by diagonalizing the Hamiltonian for comparison with imaginary time propagation.\n",
    "eps, C = eigsh(H, k=5, which=\"SA\")\n",
    "\n",
    "\n",
    "print(f\"Energy imaginary time : {energy[-1]:.8f}\")\n",
    "print(f\"Energy diagonalization: {eps[0]:.8f}\")\n",
    "print(f\"E_imag-E_diag         : {np.abs(energy[-1] - eps[0]):.8f}\")\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(imag_time, energy)\n",
    "plt.xlabel(\"Imaginary time\")\n",
    "plt.ylabel(r\"$E(\\tau)$\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 9999/9999 [01:43<00:00, 96.36it/s] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Energy imaginary time : -0.59438627\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAZoAAAEGCAYAAABcolNbAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/OQEPoAAAACXBIWXMAAAsTAAALEwEAmpwYAAAoLklEQVR4nO3deXxV9Z3/8dcnOwl72AIIAQQBERUjolalirZWq3a0jtNpxVZr25lOf13s1La/X2s70xlau02n1SlqW+xqVarUOgqCVusedkhAdgiEJCwJhJDt3s/vj3vAiDeQwL059ybv5+ORxz3ne77n3M8RzJuzfY+5OyIiIsmSEXYBIiLSvSloREQkqRQ0IiKSVAoaERFJKgWNiIgkVVbYBaSaQYMGeXFxcdhliIiklaVLl+5x98HxlilojlFcXExpaWnYZYiIpBUz29beMp06ExGRpFLQiIhIUiloREQkqRQ0IiKSVAoaERFJKgWNiIgklYJGRESSSs/RiIh0IXenJeK0RqO0Rp1IxIm4E4m+86c16kTdaY0En9F394m4E40es8ydSDRKJMoxn7F+7hD1WD93iATfE3U4v3gAl4yP+8zlKVHQiEi35O40tUZpaonS1BqJTbdGaGyJHp1+x7I2040tR/pHaW6N0hKJ/bRGnObgsyUSpSXqtBxZHky3RqO0HFkeZ53WaOq+A+zTl41T0IhI9+PuNLZEOdjYwoHGVg42tlDf1EpDc4TDzREONbdyuDlCwzHTR5YdmW4IphuaIxxuidDcGj2luswgNyuDnMwMcrIyyMrIIDvLyM7IIDszg6xMIzszg+xMIycrg/zMDHKCtqygPTtYJyvjyDbeXicrMzafmWFkZRgZRz7NyMoMPjMyyMyAzLafZmRk8M5lFttO2593bTOYz8wwMgwyzIKfYDrDEvQn+m6hB42ZDQQeAYqBrcDN7r4/Tr8IsDqY3e7u1wXtlwPfB3KApcDt7t5qZjOBJ4EtwTrz3f3bSdsRkR4sEnXqDrewv6GZ2oZm9h86Mt3CgcYWDja2Hv08ePTz7emO/is/JyuD/JxM8rMz6ZWTSUFuFr2yMxnUO4f8nPxYW04meTmZ5GVlkpudQW5WJrlZGbGf7Ezygs+jbUG/vGPasjMNs+T98u1JQg8a4G5gsbvPMbO7g/mvxOl32N3PadtgZhnAPOAKd3/LzL4NzAYeCrq85O7XJq90ke4pGnVqD7dQc7Ap9lPfeHR636EWahua2RcEyf6GZuoOt9DeW+EzDHrnZtEnL5s+eVn0zctmWN88xg95u+3tz9jygtws8oMgyc+JhUp+diZZmbp/KR2lQtBcD8wMpucBLxA/aOIpBJrd/a1gfhHwVd4OGhE5Rn1TK5W1h9lV18iu2sNU1h6m6kATNfVNR8NkT31T3KOMvOwMCgty6Z+fzYD8HEb078XAghz65+cwIGjrn5/NwIKco9O9c7N0ZNDDpULQDHX3ymB6NzC0nX55ZlYKtAJz3P0JYA+QZWYl7l4K3ASc1madC81sJbALuMvd18bbsJndCdwJMGrUqFPdH5FQ1TW0sG3fIbbubWDHvoZYmAShsqv2MAcaW9/RP8NgUO9cBveJ/Uwc1ochfXMZ3DuXwX3yjrYP7pNLQU6mQkM6rUuCxsyeA4bFWfT1tjPu7mbW3sna0e6+08zGAkvMbLW7bzKzW4AfmVkusBCIBP2XBevUm9kHgCeA8fE27O5zgbkAJSUlqXtLiEigtqGZjdX1bN3bwPa9sVDZtq+BbXsPUdvQ8o6+/fOzGd6vFyMH9GL6mIEU9evF8P55DO/fi+H9ezGkTy7ZOiUlSdQlQePus9pbZmZVZlbk7pVmVgRUt7ONncHnZjN7ATgX2OTurwKXBNu6CpgQ9DvQZt2nzew+Mxvk7nsStV8iybb/UDNvVR1kQ3U9G458VtdTc7DpaJ8MgxEDejF6YAHXnFXE6MJ8RhcWMLown9MG5FOQmwonLqQnS4W/gQuIXcCfE3w+eWwHMxsANLh7k5kNAi4GvhcsG+Lu1cERzVeA7wTtw4Cq4ChpOrFREPZ2xQ6JdFY06mzf18DaXQdYs6uOtbsOULarjj31zUf7FORkcvrQPlw2YTAThvbm9CG9KS4sYOSAfHKydEQiqSsVgmYO8Eczux3YBtwMYGYlwKfd/Q5gEvBzM4sSC4w57l4WrP9lM7s2aL/f3ZcE7TcBnzGzVuAwcIt7e/fFiHQdd6di/2GWbd/Pyh11rNlVR/muAxxsil07ycowJgztw3vPGMIZw/pw+pDeTBjah6J+ebo+ImnJ9Lv3nUpKSlyvcpZEamqNsGbnAZZt28+y7ftZum0/1cGpr7zsDCYV9WXK8H6cObwvU0b0Y/zQ3uRmZYZctUjnmNlSdy+JtywVjmhEupWWSJSVO2p5eeNeXtm0h+Xba2mOxJ5SHzUwn4vGFXLe6AFMGz2AM4b20bMh0u0paEROkbtTXnmQv22s4ZVNe3ljyz4amiOYwZnD+zL7otGcN3og00b3Z0ifvLDLFelyChqRk9DYEuGVTXtYXF7NknXVVNY1AjBucAE3nTeSi8YVMmNsIf3zc0KuVCR8ChqRDtp/qJln1+5mYVkVr2zaQ2NLlIKcTN4zfhBfmDWBSycMZlg/HbGIHEtBI3Ic+w81s7BsN0+tquSVTXuJRJ3TBvbilvNHccWkIUwfM1AX7kVOQEEjcozGlgjPrt3N/GU7eXnjHlqjzqiB+dx56ViuOauIM4f31W3GIp2goBEhdkF/ZUUdj5buYMHKXRxsbGVE/17cfskYrj1rOFNGKFxETpaCRnq0g40tPLa0gt+9vp0N1fXkZmVw9ZRhfLjkNC4cW5jUl0GJ9BQKGumRNlYf5OFXt/H40goONUc4+7T+/MeHzuLas4vom5cddnki3YqCRnoMd+evb9Xw0N+28NKGPeRkZnDt2UXcdlExU0f2D7s8kW5LQSPdXiTqPLNmN/e9sJG1uw4wrG8ed101gVumj2JQ79ywyxPp9hQ00m21RKL8aflO/ueFTWzec4ixgwr43o1TueHcERrtWKQLKWik24lGnadWV/LDhevZureByUV9+dlHpvH+KcPI1MV9kS6noJFuw9154a0a7n1mPWWVB5g4rA8P3FrCrElDdGuySIgUNNItlFce4Ft/Xstrm/cxamA+P/77c/jg2cN1BCOSAhQ0ktZqG5r54aK3+M1r2+jXK5tvX38mt5w/StdgRFKIgkbSUjTq/O6N7Xx/4XoOHG7hYzNG84UrJ2i0ZJEUpKCRtLOxup67H19F6bb9zBg7kHuuO5OJw/qGXZaItENBI2mjJRLl53/dxE8WbyQ/N5MffPhs/m7aCF3oF0lxChpJC+WVB/jiH1dSXnmAa6YWcc8Hz2RwHz1sKZIOFDSS0qJR51evbGXOM+vom5fNzz92Hu87c1jYZYlIJyhoJGVVH2zkrkdX8eJbNcyaNITv3jiVQg0ZI5J2FDSSkl7aUMPn/7CC+qZW/u2GKXz0glG6FiOSphQ0klKiUef+v27i+wvXM2FIHx751AxOH9In7LJE5BQoaCRlHGhs4Ut/XMmisiquO3s4c248i/wc/RUVSXf6v1hSwuaaeu6YV8r2fQ1849rJfPziYp0qE+kmFDQSutc37+XOXy8lK8P43SdnMH3MwLBLEpEEUtBIqOYvq+Arj69i1MB8fnnbdEYV5oddkogkmIJGQuHu/GTxRn703FtcOLaQ//noefTLzw67LBFJgpQY4tbMBprZIjPbEHwOaKffKDNbaGblZlZmZsVB+xgze93MNprZI2aWE7TnBvMbg+XFXbdX0p5o1PnWn8v40XNvceO0kcz7xHSFjEg3lhJBA9wNLHb38cDiYD6eh4F73X0SMB2oDtq/C/zI3U8H9gO3B+23A/uD9h8F/SRErZEoX35sFb96ZSufvGQM3//wVA3pL9LNpcr/4dcD84LpecANx3Yws8lAlrsvAnD3endvsNitSZcDj8VZv+12HwOuMN3KFJqm1gif/d1yHl9WwRevnMDXPjBJd5aJ9ACpEjRD3b0ymN4NDI3TZwJQa2bzzWy5md1rZplAIVDr7q1BvwpgRDA9AtgBECyvC/q/g5ndaWalZlZaU1OTuL2So5pbo3zmN8t4Zu1uvnHtZD53xXiFjEgP0WU3A5jZc0C80RC/3nbG3d3MPE6/LOAS4FxgO/AIcBvw5KnW5u5zgbkAJSUl8b5bTkFLJMpnf7eMJeuq+c6HpvCPF4wOuyQR6UJdFjTuPqu9ZWZWZWZF7l5pZkW8fe2lrQpghbtvDtZ5ApgB/ALob2ZZwVHLSGBnsM5O4DSgwsyygH7A3kTtk5xYayTK5x9ZwcKyKu754GSFjEgPlCqnzhYAs4Pp2cQ/SnmTWKAMDuYvB8rc3YHngZvirN92uzcBS4L+0gWiUedfH1vFX1ZV8rUPTOS2i8eEXZKIhCBVgmYOcKWZbQBmBfOYWYmZPQjg7hHgLmCxma0GDHggWP8rwBfNbCOxazAPBe0PAYVB+xdp/242SYLvPF3O/OU7+dKVE7jz0nFhlyMiITH9A/+dSkpKvLS0NOwy0t4DL27mO0+X8/GLi/nGtZN14V+kmzOzpe5eEm9ZqhzRSDfy5IqdfOfpcq45q4j/d41CRqSnU9BIQr2yaQ93PbqSC8YM5Ac3n01GhkJGpKdT0EjCbNt7iM/8ZhnFhQXMvbWEvOzMsEsSkRSgoJGEONjYwh3zSjGDh2afT79eGrtMRGI0erOcsmjU+cIjK9i85xC//oSG+heRd9IRjZyyHyxaz3Pl1Xzj2slcdPqgsMsRkRSjoJFT8lxZFT97fhO3nH8at16op/5F5N0UNHLSKvY38KVHVzJlRF/uue5M3cYsInEpaOSkNLdG+ezvlhONOj/7yDTdYSYi7dLNAHJSvvfMOlbsqOVnH5nG6MKCsMsRkRSmIxrptCXrqnjwb1u49cLRXDO1KOxyRCTFKWikU/bWN/Gvj61m4rA+fO0Dk8IuR0TSgE6dSYe5O1//0xoOHG7h17dP13UZEekQHdFIh81ftpNn1u7mi1dNYFJR37DLEZE0oaCRDtlZe5h7FqxlevFAPnnJ2LDLEZE0oqCRE3J37n58FVF3fnDz2WRqRGYR6QQFjZzQn5bv5KUNe7j76omcNlDjmIlI5yho5Lj2HWrm354qY9qo/vzjBRpiRkQ6T0Ejx/XvT5VR39TKnBun6iVmInJSFDTSrhffqmH+8p185rJxTBjaJ+xyRCRNKWgkrsaWCP/3iTWMHVzAP7339LDLEZE0pgc2Ja6H/raF7fsa+O0dF+jBTBE5JTqikXeprDvMT5ds5Oopw7hYLzITkVOkoJF3+c+n1xF111hmIpIQChp5hze27GPByl186rJxemZGRBJCQSNHRaLONxesZXi/PD5z2biwyxGRbkJBI0fNX1ZBeeUBvvqBSfTK0Q0AIpIYChoBYrcz/3DRW5w9sh/X6mVmIpJAoQeNmQ00s0VmtiH4HNBOv1FmttDMys2szMyKg/YxZva6mW00s0fMLCdov83MasxsRfBzRxfuVtp5+NWtVNY18pWrJ2KmEQBEJHFCDxrgbmCxu48HFgfz8TwM3Ovuk4DpQHXQ/l3gR+5+OrAfuL3NOo+4+znBz4PJKT/91TW08LPnN3HZhMFcNE63M4tIYqVC0FwPzAum5wE3HNvBzCYDWe6+CMDd6929wWL/9L4ceOx468vx3f/XTRxobOEr758Ydiki0g2lQtAMdffKYHo3MDROnwlArZnNN7PlZnavmWUChUCtu7cG/SqAEW3Wu9HMVpnZY2Z2WtL2II3trmvkly9v4UPnjGDycL01U0QSr0uCxsyeM7M1cX6ub9vP3R3wOJvIAi4B7gLOB8YCt53ga/8MFLv7VGARbx81xavvTjMrNbPSmpqaju9YN3DfCxuJRJ0vXDkh7FJEpJvqkrHO3H1We8vMrMrMity90syKePvaS1sVwAp33xys8wQwA/gF0N/MsoKjmpHAzuA797ZZ/0Hge8epby4wF6CkpCRe0HVLlXWH+cMbO/hwyUg9nCkiSZMKp84WALOD6dnAk3H6vEksUAYH85cDZcER0PPATceuH4TWEdcB5QmuO+3d/8Imou7800yNziwiyZMKQTMHuNLMNgCzgnnMrMTMHgRw9wix02aLzWw1YMADwfpfAb5oZhuJXbN5KGj/nJmtNbOVwOc48am2HmV3XSN/eGMHN52noxkRSa7QXxMQnOK6Ik57KXBHm/lFwNQ4/TYTu9352PavAl9NaLHdyP0vbCTqzj/rXTMikmSpcEQjXWx3XSO/f3MHN07T0YyIJJ+Cpgd64KXNRKI6mhGRrqGg6WHqGlr4/Rvb+eDUIkYV6mhGRJJPQdPD/Ob1bTQ0R7jzUr0GQES6hoKmB2lsifDLl7dy6YTBGgVARLpMp4PGzAqC4V8kzTyxfCd76pv41KVjwy5FRHqQEwaNmWWY2UfM7C9mVg2sAyqDofrvNTNdUU4D0agz98XNTBnRl4vGFYZdjoj0IB05onkeGEfsmZRh7n6auw8B3gO8BnzXzD6axBolARaVV7F5zyE+dek4vW9GRLpURx7YnOXuLcc2uvs+4HHgcTPLTnhlklAP/W0LIwf04uopw8IuRUR6mBMe0RwJGTN75UR9JDWVVx7gjS37uPXC0WRl6v4PEelanfmtk3dsg5ldksBaJEnmvbKVvOwMbi7RK3lEpOt1ZqyzM8zsT8BaYA1QRWz4fT2QkcJqG5p5YsVOPnTuCPrn54Rdjoj0QJ0Jmi3AfwBTgPOA4cC3klGUJM4jb+6gsSXKrRcWh12KiPRQnQmaZnd/k9i7YSQNRKLOr1/bxvQxA5lUpAc0RSQcnblGc1nSqpCkWFxeRcX+w9x2UXHYpYhID9aRBzYNwN0PnqiPpJZfv7aNon55XDV5aNiliEgP1qEHNs3sX8xsVNtGM8sxs8vNbB5vv4pZUsT2vQ28tGEPt5w/Src0i0ioOnKN5v3AJ4Dfm9kYoJbYrc6ZwELgx+6+PGkVykl5pHQ7GQY3nz8y7FJEpIfrSNCMc/f7gPuCEQAGAYfdvTaplclJa41EebS0gplnDKGoX6+wyxGRHq4j51R+3WZ6trtXHgkZM9Obs1LQ8+trqD7YxC3n6wFNEQlfR4Km7YX+fzpm2UsJrEUS5A9vbGdwn1zeO3FI2KWIiHQoaLzN9LF3l+kqc4qprDvM8+ur+fB5I8nWTQAikgI6co1mmJndBqzk3UHj7+4uYXq0tIKow9/rtJmIpIiOBM09xIac+Tgw0szKgHJiL0AblLzSpLOiUeePpTu4+PRCRhcWhF2OiAjQgaBx97lt581sJHAWMBV4MUl1yUl4fcs+KvYf5svvOyPsUkREjurMWGcAuHsFUAH8b+LLkVPxp+UVFORkctVkvdxMRFKHrhZ3E40tEZ5evZurzyqiV05m2OWIiByloOkmFpZVUd/Uyt+dOyLsUkRE3kFB0038aVkFw/vlMWNsYdiliIi8Q+hBY2YDzWyRmW0IPge002+UmS00s3IzKzOz4qD9s2a20czczAa16W9m9pNg2Sozm9ZFu9Tlag428eKGPVx/7ggyMjSQtoikltCDBrgbWOzu44HFwXw8DwP3uvskYDpQHbS/DMwCth3T/2pgfPBzJ3B/gutOGQtW7iISdZ02E5GUlApBcz0wL5ieB9xwbAczmwxkufsiAHevd/eGYHq5u29tZ7sPe8xrQH8zK0pC/aH70/IKzhrRj/FD+4RdiojIu6RC0Ax198pgejcQ7y1dE4BaM5tvZsvN7F4zO9GtVSOAHW3mK4K2bmVD1UHW7DzAh3Q0IyIpqtPP0ZwMM3sOiPdwx9fbzri7m1m8YW2ygEuAc4HtwCPAbcBDCarvTmKn1xg1atQJeqeWP6/cRYbBB88eHnYpIiJxdUnQuPus9paZWZWZFbl7ZXBqqzpOtwpghbtvDtZ5ApjB8YNmJ9B2wK+RQVu8+uYCcwFKSkrSZvw2d+epVZXMGFvI4D65YZcjIhJXKpw6W8Dbr4KeDTwZp8+bxK6xDA7mLwfKOrDdW4O7z2YAdW1O0XULZZUH2LznENdO1dGMiKSuVAiaOcCVZraB2N1jcwDMrMTMHgRw9whwF7DYzFYTG0X6gaDf58ysgtgRy6oj6wBPA5uBjUHfY9+lk/aeWlVJZobx/ikackZEUpe5p82Zoi5RUlLipaWlYZdxQu7Opfc+T3FhAb++/YKwyxGRHs7Mlrp7SbxlqXBEIydh9c46duw7zAd12kxEUpyCJk09taqS7EzjfWfqtJmIpDYFTRpyd/6yqpJLxg+mX3522OWIiByXgiYNLdtey87aw1w7tVsOdCAi3YyCJg39ZVUlOZkZzJocbxAFEZHUoqBJM+7Os2t3c+mEQfTN02kzEUl9Cpo0s2bnAXbWHuYq3QQgImlCQZNmnl27mwyDWZN02kxE0oOCJs08u3Y308cMZGBBTtiliIh0iIImjWyuqWdDdb2enRGRtKKgSSPPrq0C0PUZEUkrCpo08uza3Zw1oh8j+vcKuxQRkQ5T0KSJ3XWNrNhRy/vO1E0AIpJeFDRpYlHZbgBdnxGRtKOgSRPPrq1i7KACTh/SO+xSREQ6RUGTBuoaWnht816uOnMYZhZ2OSIinaKgSQNL1lfRGnVdnxGRtKSgSQOLy6sZ1DuXs0f2D7sUEZFOU9CkuJZIlL++VcPlEweTkaHTZiKSfhQ0Ka50634ONrZy+USdNhOR9KSgSXFL1lWRk5nBe8YPCrsUEZGToqBJcYvXVXPB2IH0zs0KuxQRkZOioElhW/YcYnPNIa6YOCTsUkRETpqCJoUtWVcNoOszIpLWFDQpbMm6KsYP6c2owvywSxEROWkKmhR1sLGF1zfv4/JJOm0mIulNQZOiXtqwh9aoc4VOm4lImlPQpKjF5dX065XNtFH9wy5FROSUKGhSUCTqvLC+mplnDCYrU39EIpLeQv8tZmYDzWyRmW0IPge002+UmS00s3IzKzOz4qD9s2a20czczAa16T/TzOrMbEXw840u2qVTtrKilr2HmrlctzWLSDcQetAAdwOL3X08sDiYj+dh4F53nwRMB6qD9peBWcC2OOu85O7nBD/fTnDdSbOkvJrMDOOyCYPDLkVE5JSlQtBcD8wLpucBNxzbwcwmA1nuvgjA3evdvSGYXu7uW7um1K6xeF01540eQP/8nLBLERE5ZakQNEPdvTKY3g3Eu81qAlBrZvPNbLmZ3WtmmR3Y9oVmttLM/tfMzmyvk5ndaWalZlZaU1NzEruQOJV1hymvPKDTZiLSbXTJAFpm9hwQ72X3X2874+5uZh6nXxZwCXAusB14BLgNeOg4X7sMGO3u9Wb2AeAJYHy8ju4+F5gLUFJSEu/7u8wL62NBp6ARke6iS4LG3We1t8zMqsysyN0rzayIt6+9tFUBrHD3zcE6TwAzOE7QuPuBNtNPm9l9ZjbI3fec7H50hefXVTOify/GD+kddikiIgmRCqfOFgCzg+nZwJNx+rwJ9DezI1fHLwfKjrdRMxtmZhZMTye2r3sTUnGSNLVGeHnjHmaeMZigdBGRtJcKQTMHuNLMNhC7e2wOgJmVmNmDAO4eAe4CFpvZasCAB4J+nzOzCmAksOrIOsBNwBozWwn8BLjF3UM9LXYipVv3c6g5wnvP0GkzEek+Qn/JibvvBa6I014K3NFmfhEwNU6/nxALkmPbfwr8NKHFJtnz66rJyczgotMLwy5FRCRhUuGIRgJL1sdecpafE3r+i4gkjIImRWzbG3vJmU6biUh3o6BJEUdua36vbmsWkW5GQZMinl9fTXFhPmMGFYRdiohIQiloUsDh5givbtrLTJ02E5FuSEGTAl7bvJem1qhOm4lIt6SgSQHPr68mLzuDC8YMDLsUEZGEU9CEzN1Zsq6ai8cNIi+7I+OEioikFwVNgpRu3cedD5dSWXe4U+ttqjlExf7DzNRpMxHpphQ0CVJzsImFZVXUNrR0ar0X1sfGEJ2pl5yJSDeloEmQnKzYf8rm1min1nt+fTXjh/TmtIH5yShLRCR0CpoEyc6M/adsiXQ8aOqbWnljyz7dbSYi3ZqCJkFO5ojm5Y17aIk4M8/QaTMR6b4UNAlyJGiaOnFEs6S8mt65WZSM1m3NItJ9KWgSJOfIqbMOHtFEo87idVVcdsbgoyElItId6Tdcghw9ddbBI5oVFbXsqW/mqslDk1mWiEjoFDQJcuSIpqPXaJ4rqyIzw5g5QTcCiEj3pqBJkOyszt11tqisiunFA+mXn53MskREQqegSZDOHNFs3XOIDdX1zNJpMxHpARQ0CXL0rrMOBM1z5VUAXDlJQSMi3Z+CJkFyO3EzwHPlVZwxtA+jCjUagIh0fwqaBMnNyiAzwzjU1HrcfvsONfPm1v3MmqybAESkZ1DQJIiZ0Ts3i4ONxw+aZ9fuJhJ1rp5S1EWViYiES0GTQH3ysqg/QdA8tWoXYwYVcObwvl1UlYhIuBQ0CdQ7N4sDxwmaPfVNvLppL9ecVYSZdWFlIiLhUdAkUN+8bA42tv8+mmfW7CbqcM1UnTYTkZ5DQZNAffKyqD/OzQB/XrmLsYMLmDisTxdWJSISLgVNAg0oyGFvfXPcZdv3NvD6ln3ccM4InTYTkR5FQZNARf3yqD7YSGucZ2keXboDM7jpvJEhVCYiEp7Qg8bMBprZIjPbEHwOaKffKDNbaGblZlZmZsVB+2/NbL2ZrTGzX5hZdtBuZvYTM9toZqvMbFqy92VYvzyiDjX1Te9oj0Sdx5ZWcOn4wQzv3yvZZYiIpJTQgwa4G1js7uOBxcF8PA8D97r7JGA6UB20/xaYCJwF9ALuCNqvBsYHP3cC9yel+jaK+uUBsKu28R3ti8urqKxr5O/PPy3ZJYiIpJxUCJrrgXnB9DzghmM7mNlkIMvdFwG4e727NwTTT3sAeAM4cm7qeuDhYNFrQH8zS+rtXsWFBQBsqq5/R/vcFzczckAvvXtGRHqkVAiaoe5eGUzvBuL9Np4A1JrZfDNbbmb3mllm2w7BKbOPAc8ETSOAHW26VARt72Jmd5pZqZmV1tTUnPSOFBcWUJCTydpddUfb3tiyj9Jt+7n9PWPIykyF/9wiIl0rqyu+xMyeA4bFWfT1tjPu7mbmcfplAZcA5wLbgUeA24CH2vS5D3jR3V/qbH3uPheYC1BSUhLv+zskI8OYOrI/r27eC8Re1/zvfyljWN88bjl/1MluVkQkrXVJ0Lj7rPaWmVmVmRW5e2Vwaqs6TrcKYIW7bw7WeQKYQRA0ZvZNYDDwqTbr7ATaXhQZGbQl1funDOObC9aybPt+XlhXzaqKOv7rlnPolZN54pVFRLqhVDiXswCYHUzPBp6M0+dNYtdYBgfzlwNlAGZ2B/A+4B/cve19xQuAW4O7z2YAdW1O0SXNh6aNYEifXG7+n1f5yZKN3DhtJNedPTzZXysikrIsdg09xALMCoE/AqOAbcDN7r7PzEqAT7v7HUG/K4EfAAYsBe5092Yzaw3WOxhscr67f9tiT0X+FHg/0AB83N1LT1RPSUmJl5aesNtxbag6yC9e3sKEoX249cJiMjP0gKaIdG9mttTdS+IuCztoUk0igkZEpKc5XtCkwqkzERHpxhQ0IiKSVAoaERFJKgWNiIgklYJGRESSSkEjIiJJpaAREZGkUtCIiEhS6YHNY5hZDbGRBk7GIGBPAstJB9rnnkH73DOcyj6PdvfB8RYoaBLIzErbezK2u9I+9wza554hWfusU2ciIpJUChoREUkqBU1izQ27gBBon3sG7XPPkJR91jUaERFJKh3RiIhIUiloREQkqRQ0CWJm7zez9Wa20czuDrueZDOz08zseTMrM7O1ZvZ/wq6pK5hZppktN7Onwq6lq5hZfzN7zMzWmVm5mV0Ydk3JZGZfCP5OrzGz35tZXtg1JYOZ/cLMqs1sTZu2gWa2yMw2BJ8DEvFdCpoEMLNM4GfA1cBk4B/MbHK4VSVdK/Ald58MzAD+uQfsM8D/AcrDLqKL/RfwjLtPBM6mG++/mY0APgeUuPsUIBO4JdyqkuZXxF5139bdwGJ3Hw8sDuZPmYImMaYDG919s7s3A38Arg+5pqRy90p3XxZMHyT2y2dEuFUll5mNBK4BHgy7lq5iZv2AS4GHANy92d1rQy0q+bKAXmaWBeQDu0KuJync/UVg3zHN1wPzgul5wA2J+C4FTWKMAHa0ma+gm//SbcvMioFzgddDLiXZfgz8KxANuY6uNAaoAX4ZnDJ80MwKwi4qWdx9J/B9YDtQCdS5+8Jwq+pSQ929MpjeDQxNxEYVNHJKzKw38DjweXc/EHY9yWJm1wLV7r407Fq6WBYwDbjf3c8FDpGg0ympKLgmcT2xgB0OFJjZR8OtKhwee/YlIc+/KGgSYydwWpv5kUFbt2Zm2cRC5rfuPj/sepLsYuA6M9tK7NTo5Wb2m3BL6hIVQIW7HzlafYxY8HRXs4At7l7j7i3AfOCikGvqSlVmVgQQfFYnYqMKmsR4ExhvZmPMLIfYxcMFIdeUVGZmxM7bl7v7D8OuJ9nc/avuPtLdi4n9+S5x927/L1133w3sMLMzgqYrgLIQS0q27cAMM8sP/o5fQTe++SGOBcDsYHo28GQiNpqViI30dO7eamafBZ4ldpfKL9x9bchlJdvFwMeA1Wa2Imj7mrs/HV5JkiT/Avw2+EfUZuDjIdeTNO7+upk9BiwjdmflcrrpUDRm9ntgJjDIzCqAbwJzgD+a2e3EXpdyc0K+S0PQiIhIMunUmYiIJJWCRkREkkpBIyIiSaWgERGRpFLQiIhIUiloRNphZvVd8B2vJPs7jvm+G9oOfmpm3zazWV1Zg/Q8ur1ZpB1mVu/uvcOuoyPMLMvdWzvQ71fAU+7+WPKrEonREY3ICZjZTDP7q5k9aWabzWyOmf2jmb1hZqvNbFzQ74Nm9now+ORzZjY0aB8cvNtjbTAo5TYzGxQsq2/zHS+0ee/Lb4Mn0zGzb5jZm8H7Uea2aX/BzH5sZqXA181sSzAsEGbWt+180HYRcB1wr5mtMLNxZvYrM7spWL7VzP4zWFZqZtPM7Fkz22Rmn26znS8H9awys291wR+BpDkFjUjHnA18GphEbESECe4+ndgrA/4l6PM3YEYw+OQfiI30DLEnrpe4+5nExgob1c53nAt8ntg7jcYSG30B4Kfufn7wfpRewLVt1slx9xJ3/xbwArHXGEBsmJz5wXhdALj7K8SGGPmyu5/j7pvi1LDd3c8BXiL2vpKbiL1v6FsAZnYVMJ7YqzHOAc4zs0vb2R8RQEEj0lFvBu/gaQI2AUeGjl8NFAfTI4FnzWw18GXgzKD9PcSCB3d/Btjfzne84e4V7h4FVrTZ7nuDI6XVwOVttgvwSJvpB3l7eJiPA7/s5D7C22P0rQZed/eD7l4DNJlZf+Cq4Gc5sWFaJhILHpF2aawzkY5pajMdbTMf5e3/j/4b+KG7LzCzmcA9p/AdESAreI3wfcTe+LjDzO4B2r5a+NCRCXd/2cyKg+/OdPc1dF7b/Tp2n7MAA/7T3X9+EtuWHkpHNCKJ04+3Xw8xu037ywSDEwannjrzHvYjobInePfPTSfo/zDwO9o/mjkI9OnE9x/rWeATQS2Y2QgzG3IK25MeQEEjkjj3AI+a2VJgT5v2bwFXmdka4MPE3lx4sCMbDF6b/ACwhtgv+TdPsMpviQXZ79tZ/gfgy8ENC+M6UsMx9SwkFmSvBqfyHuPUgkt6AN3eLJJkZpYLRILXSVxI7G2V5yTpu24Crnf3jyVj+yInQ9doRJJvFLF3fGQAzcAnk/ElZvbfwNXAB5KxfZGTpSMaERFJKl2jERGRpFLQiIhIUiloREQkqRQ0IiKSVAoaERFJqv8PIAuXcRH7FOcAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\"\"\"\n",
    "Solve for the H2 molecule with an internuclear distance of 0.7 a.u. using imaginary time propagation.\n",
    "\"\"\"\n",
    "\n",
    "V_H2 = np.zeros((n_rho, n_z))\n",
    "V_H2 = np.zeros((n_rho, n_z))\n",
    "psi_H2 = np.zeros((n_rho, n_z))\n",
    "Z = 1\n",
    "dist = 0.7\n",
    "# Setup potential and initial guess for the orbitals.\n",
    "for i in range(n_rho):\n",
    "    for j in range(n_z):\n",
    "        r_p = np.sqrt(rho_inner[i] ** 2 + (z_inner[j] - dist) ** 2)\n",
    "        r_m = np.sqrt(rho_inner[i] ** 2 + (z_inner[j] + dist) ** 2)\n",
    "        V_H2[i, j] = Verfgau(r_p, Z) + Verfgau(r_m, Z)\n",
    "        psi_H2[i, j] = np.sqrt(rho_inner[i]) * np.exp(-Z * r_p) + np.sqrt(\n",
    "            rho_inner[i]\n",
    "        ) * np.exp(-Z * r_m)\n",
    "\n",
    "V_H2 = V_H2.flatten()\n",
    "H = -0.5 * Laplacian.L + spdiags([V_H2], diags=[0])\n",
    "\n",
    "# Normalize the initial guess.\n",
    "psi_H2 = psi_H2.ravel()\n",
    "norm_psi_H2 = np.sum(np.abs(psi_H2) ** 2) * drho * dz\n",
    "psi_H2 /= np.sqrt(norm_psi_H2)\n",
    "\n",
    "dt = 0.001\n",
    "n_t = int(1e4)\n",
    "imag_time = np.linspace(0, n_t * dt, n_t)\n",
    "energy_H2 = np.zeros(n_t)\n",
    "\n",
    "# For two-particle systems there is no exchange term.\n",
    "def _F_psi(psi, W_direct):\n",
    "    return H.dot(psi) + W_direct * psi\n",
    "\n",
    "\n",
    "for n in tqdm.tqdm(range(n_t - 1)):\n",
    "    # Compute direct potential\n",
    "    P = np.abs(psi_H2) ** 2 / np.sqrt(2 * np.pi)\n",
    "    P = P.reshape((n_rho, n_z))\n",
    "    P = np.einsum(\"i, ij->ij\", 1 / rho_inner, P)\n",
    "    f = -4 * np.pi * np.einsum(\"i, ij->ij\", np.sqrt(rho_inner), P)\n",
    "    w_ii = Poisson.solve(f, inhomgenous_dirichlet=True)\n",
    "    W_ii_direct = np.einsum(\"i, ij->ij\", 1 / np.sqrt(rho_inner), w_ii)\n",
    "    W_direct = W_ii_direct.ravel() / np.sqrt(2 * np.pi)\n",
    "\n",
    "    F_psi = _F_psi(psi_H2, W_direct)\n",
    "    energy_H2[n] = np.sum(psi_H2 * F_psi) * drho * dz\n",
    "\n",
    "    k1 = -F_psi\n",
    "    k2 = -_F_psi(psi_H2 + 0.5 * dt * k1, W_direct)\n",
    "    k3 = -_F_psi(psi_H2 + 0.5 * dt * k2, W_direct)\n",
    "    k4 = -_F_psi(psi_H2 + dt * k3, W_direct)\n",
    "    psi_H2 += dt / 6 * (k1 + 2 * k2 + 2 * k3 + k4)\n",
    "    psi_H2 /= np.sqrt(np.sum(np.abs(psi_H2) ** 2) * drho * dz)\n",
    "\n",
    "F_psi = _F_psi(psi_H2, W_direct)\n",
    "energy_H2[n_t - 1] = np.sum(psi_H2 * F_psi) * drho * dz\n",
    "\n",
    "print(f\"Energy imaginary time : {energy_H2[-1]:.8f}\")\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(imag_time, energy_H2)\n",
    "plt.xlabel(\"Imaginary time\")\n",
    "plt.ylabel(r\"$E(\\tau)$\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 9999/9999 [04:16<00:00, 38.95it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Imag time final energy: -0.29522, -2.36165\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAY0AAAEGCAYAAACZ0MnKAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/OQEPoAAAACXBIWXMAAAsTAAALEwEAmpwYAAAYcElEQVR4nO3dfZAkdZ3n8fe3qrun5xEYBoaHAQcHXEUQXHs58Wl9YA3wFNwNn28NVveW2LjT27240MAj4tT749a73TM0XL2Tc1dxV/E8lYBQVh58WJ91BoHlUXkQZBBnBgGZYZiH7vreH5k9XV3d1ZM901VZ0/N+RVRU5i9/+ctvwUx9JjMrMyMzkSSpikbdBUiSDh2GhiSpMkNDklSZoSFJqszQkCRVNlR3Ab22Zs2aXL9+fd1lSNIh46abbno0M4+ZbdmiD43169ezadOmusuQpENGRDzYbZmHpyRJlRkakqTKDA1JUmWGhiSpMkNDklSZoSFJqszQkCRVZmh0ccOdW7hv2466y5CkgbLoL+47UH/22eKCwI+99fm85szjaTai5ookqX7uaezHu6+8mZf9j2/xkRt/zj1bttddjiTVyj2NLo5YOszrzjqeF29Yw2d/+CAf/cY9fOTGe9hwzHJe9qxjeNGGNfyrZ65m1ehw3aVKUt8YGnMYajS44MzjueDM49n65C6uu+PXXH/nFj7/41/y6e8/QLMRPPeEVZy17kiet+4IzjrpSDYcs8JDWZIWLUOjomNXjfL2c9fz9nPXs3t8gp8++AQ/uO9RNj7wGFfd/DD/8KPi/l7LR5qctnYlz1q7gtOOXcmpa1fwrLUrOeGIUSIME0mHNkPjACwZanLuhqM5d8PRALRayf2P7uCWh37LbZuf4OdbdvDNu7fyxU2b960zOtxg3VHLOOmopZy0ehknHbWseF+9lONWjXLUshEa7qFIGnCGxgJoNIJTj13Jqceu5A0vWLev/bGn9nDv1h38fMt2Hnj0KR56fCcPPfY0mx58nO27xqeNMdQIjl25hGNWjbJ25RKOXbWEtStHOXbVEo5cNsLq5SMctWyYI5eNcOTSYYaa/oZBUv8ZGj20evkI55yymnNOWT1j2W937i1DZCdbntzF1u272fLkbrZu38WDv9nJxgce4/Gde7uOvXJ0iNXLRzhyWREmRywdZsWSIVaMDrFqtJwu51cuGWLl6DArRou2laNDLBlqeLhM0rwZGjU5YtkwRyw7gjNOPKJrn93jE2zbvpsndu7l8Z17eHznXp7YuYfHntozre03O/bwi0efYseucbbvHmfPeGu/228ELB1usnSkyehwk6XDTZZNTo809y1bOtw2PdJkdKjJyFCDkaEGS8rXyFCDkeZU+0hzann7/MhQg6FGGFbSIczQ6CIz6y6BJUNN1h21jHVHzW+93eMT7Ng1zo7d42zfVbx27B5nx+697Ng1zpO7xtm1d4KdeyZ4eu8Eu8r3yfnHntrD0+X0vve9EyzEf5IIpoXKUKPBUDMYagRDzSJUhpsNmo1guBmzLh9qNhhuBEPNoNlozNEvGG40aDSCZkCzDKxmI2hGFO0NaHS2lfNRrjPVNzr6Ti2PLu2NtvdGFNtqRDF2lPMBbW0GqgabobEILRlqsmRFk6NXLFmwMTOT3eMtnt4zwZ6JFnvGW+weL94n54vpie7L2uZ3l9PjEy3GJ5LxVjLearF3IploJXvL9olWsmt8gvGJom2iVfSdnN47Uaw3MZHsbU2NdShrDxdiar4zXPa1d85TzjcgmGpn37hF+2RotfeLaGvv6DcZdJNtlNsptlfOz2iLjmVT609OTy7sXHcyPifbaF9337Lp4825PaaH8sxl89xe2dBtWft47abqi475uZdPrR9d+rfVHzA63ORNYyex0AwNVRIRjA4Xh68GXWZOC5dWCyYyaWXSaiUT5fLJ9olWsWyiNTXdSqa1T1svk4lWx/Kcap/Zt3hlQlKM3cpyvtxW7msr59v6Ub539kumt2fmrPNFv86xshxrZr/ct2yq3+RnyfK/b/EOSTkBbW2Tn61cztSe+7TlbcuyHKB9/Bn9s2OsWcab2htub8tpyypvb1+f9tpnjjeo1qxYYmhIVUQUh6aGmhwSIadD32TQzhZCMD14pq83+/IZYdvRnznWm1ynV4c6DQ1JOkiTh4TKuTpL6Tl/7C9JqszQkCRVZmhIkiozNCRJlRkakqTKDI0uBvwn2JJUC0NjDt7RQZKmMzQkSZUZGpKkygwNSVJlhoYkqTJDQ5JUmaEhSaqs9tCIiNURcUNE3FO+z/qcuoiYiIhbytc1/a5TkjQAoQFcCnwjM08DvlHOz+bpzDy7fF3Yv/IkSZMGITQuAq4op68AXl9fKW28JFySZhiE0FibmY+U078G1nbpNxoRmyLiRxHx+rkGjIhLyr6btm3bdsCFxSJ/mIokzVdfntwXETcCx82y6LL2mczMiOj2b/xnZObDEfFM4JsRcVtm3jdbx8y8HLgcYGxszH0GSVogfQmNzDyv27KI2BIRx2fmIxFxPLC1yxgPl+/3R8S3gecDs4aGJKk3BuHw1DXAxeX0xcDVnR0i4qiIWFJOrwFeDNzZtwolScBghMaHgD+IiHuA88p5ImIsIj5V9nkOsCkibgW+BXwoMw0NSeqzvhyemktm/gZ41Sztm4B/W07/ADizz6VJkjoMwp6GJOkQYWhIkiozNCRJlRkaXXhxhyTNZGjMwWeES9J0hoYkqTJDQ5JUmaEhSarM0JAkVWZoSJIqMzQkSZUZGpKkygyNLjK9vE+SOhkac/DaPkmaztCQJFVmaEiSKjM0JEmVGRqSpMoMDUlSZYaGJKkyQ0OSVJmhIUmqzNDowuvBJWkmQ2MOPu5VkqYzNCRJlRkakqTKDA1JUmWGhiSpMkNDklSZoSFJqszQkCRVZmhIkiozNLrwEeGSNJOhMYfwknBJmsbQkCRVZmhIkiqrPTQi4o0RcUdEtCJibI5+50fEzyLi3oi4tJ81SpIKtYcGcDvwR8B3unWIiCbwceAC4HTgrRFxen/KkyRNGqq7gMy8C/Z70vkc4N7MvL/s+wXgIuDOnhcoSdpnEPY0qjgReKhtfnPZNquIuCQiNkXEpm3btvW8OEk6XPRlTyMibgSOm2XRZZl59UJvLzMvBy4HGBsb84oLSVogfQmNzDzvIId4GDipbX5d2SZJ6qPaz2lUtBE4LSJOoQiLtwBv6+UGb33/q2l4bZ8kTVP7OY2I+MOI2AycC3wtIq4r20+IiGsBMnMceBdwHXAX8MXMvKOXdY0MNRhq1v6fR5IGSu17Gpl5FXDVLO2/Al7TNn8tcG0fS5MkdfCf0pKkygwNSVJlhoYkqTJDQ5JUmaEhSarM0JAkVWZoSJIqMzQkSZUZGpKkygwNSVJl8w6NiFhePklPknSY2W9oREQjIt4WEV+LiK3A3cAjEXFnRPx1RJza+zIlSYOgyp7Gt4ANwPuA4zLzpMw8FngJ8CPgv0fEH/ewRknSgKhyl9vzMnNvZ2NmPgZ8GfhyRAwveGWSpIGz3z2NycCIiB/sr48kaXGbz4nw0c6GiHjpAtYiSRpw83kI0+9ExFXAHcDtwBbgUxTnOyRJh4H5hMYvgP8GnAG8ADgB+GAvipIkDab5hMaezNwIbOxVMZKkwTafcxq/37MqJEmHhCoX9wVAZm7fX59F5aYr4KGf1F2FJA2UShf3RcS7I+Lk9saIGImIV0bEFcDFvSmvRl+/FO66pu4qJGmgVDmncT7wTuDKiDgFeILi57dN4HrgI5l5c88qlCQNjCqhsSEzPwF8orzyew3wdGY+0dPKBkFm3RVI0kCpcnjqH9qmL87MRyYDIyKW9aSqgbD4TtNI0sGqEhrt357/rmPZdxewFknSgKsSGu3HaDr/+b14H+K0CH8QJkkHq8o5jeMi4k+AW5kZGh70l6TDSJXQ+ADFbUPeAayLiDuBuygexrSmd6UNAE+ES9I0+w2NzLy8fT4i1gFnAs8DvtOjugaAh6ckqdN87j0FQGZuBjYD/7Tw5Qwa9zQkqd3iPZF9sDwRLkkzGBqSpMoMjbl4IlySpjE0uvLwlCR1qj00IuKNEXFHRLQiYmyOfg9ExG0RcUtEbOpPde5pSFK7ef96qgduB/4I+GSFvq/IzEd7XI8kqYvaQyMz7wIYuOc4DVg5kjQIaj88NQ8JXB8RN0XEJf3ZooenJKldX/Y0IuJG4LhZFl2WmVdXHOYlmflwRBwL3BARd2fmrFekl6FyCcDJJ588W5cK3NWQpE59CY3MPG8Bxni4fN8aEVcB59DlNiblrU8uBxgbG3N3QZIWyCFxeCoilkfEyslp4NUUJ9B7zLyRpHa1h0ZE/GFEbAbOBb4WEdeV7SdExLVlt7XA9yLiVuAnwNcy8+s9Lqynw0vSoWgQfj11FXDVLO2/Al5TTt8PnNXn0jwRLkkdat/TGFzuaUhSJ0NDklSZoTEnD09JUjtDoxtPhEvSDIbGXDwRLknTGBpduachSZ0MDUlSZYbGnDw8JUntDI1uPBEuSTMYGnPxRLgkTWNodOWehiR1MjQkSZUZGnPy8JQktTM0uvFEuCTNYGhIkiozNObir6ckaRpDoysPT0lSJ0NjTu5pSFI7Q6MbT4RL0gyGhiSpMkNjLp4Il6RpDI2uPDwlSZ0MjTm5pyFJ7QyNbjwRLkkzGBqSpMoMjbl4dEqSpjE0uvLwlCR1MjTm5K6GJLUzNLrxRLgkzWBoSJIqMzTm4hXhkjSNodGVh6ckqZOhIUmqzNCYk4enJKmdodGNR6ckaQZDYy6eCJekaWoPjYj464i4OyL+JSKuiogju/Q7PyJ+FhH3RsSlfais95uQpENM7aEB3ACckZnPA34OvK+zQ0Q0gY8DFwCnA2+NiNP7WqUkqf7QyMzrM3O8nP0RsG6WbucA92bm/Zm5B/gCcFEfquv9JiTpEFJ7aHR4J/BPs7SfCDzUNr+5bOsdbyMiSTMM9WMjEXEjcNwsiy7LzKvLPpcB48DnFmB7lwCXAJx88skHPpAnwiVpmr6ERmaeN9fyiPgT4LXAqzJn/aZ+GDipbX5d2dZte5cDlwOMjY0d4De/exqS1Kn2w1MRcT7wXuDCzNzZpdtG4LSIOCUiRoC3ANf0q0ZJUqH20AD+FlgJ3BARt0TE/waIiBMi4lqA8kT5u4DrgLuAL2bmHb0vzcNTktSuL4en5pKZp3Zp/xXwmrb5a4Fr+1WXJ8IlaaZB2NMYXJ4Il6RpDI2u3NOQpE6GhiSpMkNjTh6ekqR2hkY3ngiXpBkMDUlSZYbGXPz1lCRNY2h05eEpSepkaMzJPQ1JamdodOOJcEmawdCQJFVmaHQzvBR2b6+7CkkaKIZGN0efCo/eU3cVkjRQDI1ujnk2PPFLePrxuiuRpIFhaHSz/iVAwgPfq7sSSRoYhkY3634PRlbCz6+ruxJJGhiGRjfNYXjOa+HOq2Hv03VXI0kDwdCYy9lvg91Pwu1fqbsSSRoIhsZc1r8U1p4J3/swtCbqrkaSamdozCUCfv+98Jt74aZP112NJNXO0Nif57wOnvkKuOH98PgDdVcjSbUyNPYnAl73UWg04fNvhqefqLsiSaqNoVHFUc+AN/8j/OY+uOJ1sH1L3RVJUi0MjapOeRm89cri/MYnXwo/+3rdFUlS3xka83HaH8Cf3gDL1sCVb4YrLoT7/xlarbork6S+GKq7gEPOcWfAJd+CjX8H3/2f8NkL4chnwOkXFifMT34hjCyvu0pJ6onIRf4c7LGxsdy0aVNvBt+zE+7+Ktx6ZXGPqok9QMDRG2DtGbD6FFh1YvFasRZGV8GSVcX70Gh/H/SUCdkqXq2JcnqibT5naZuczy7rtMr1ZlunVeyBzbpOaz+17K9PhXGzNfW5i4n9zFfpk9Pe5rfOAWynW9uCW+A/hwv+5/owqm8ha1uyCl774QMsI27KzLHZlrmncTBGlsHz3lS89jwFD/4QHt4Ev74NHrkV7v4atPbOvm5jGJoj0Bwq3hvDU9NQflnk1Jc5WX5vZNuyub5QO75UF+Wja6P4VVs0IMr3RrP8ixdtfwHL9/3NV+lTaR0OYJ1u893aFsiC/6Nxgcc7rOpb4NqWHb2w45UMjYUyshxOO694TWq14Klt8ORm2LGtuCXJrt8W77u3w8Te8rWnCJeJ8XJvhbYvvkbHl2Dbe6MxyxdmY/prxpdqt3XKL9tZ12nOPWaVcaeN0fEFP+s67WN2q8VH8kr9Zmj0UqMBK9cWL0laBPz1lCSpMkNDklSZoSFJqszQkCRVZmhIkiozNCRJlRkakqTKDA1JUmWL/t5TEbENePAAV18DPLqA5RwK/MyL3+H2ecHPPF/PyMxjZluw6EPjYETEpm437Vqs/MyL3+H2ecHPvJA8PCVJqszQkCRVZmjM7fK6C6iBn3nxO9w+L/iZF4znNCRJlbmnIUmqzNCQJFVmaMwiIs6PiJ9FxL0RcWnd9fRaRJwUEd+KiDsj4o6I+Iu6a+qXiGhGxM0R8dW6a+mHiDgyIr4UEXdHxF0RcW7dNfVaRPzH8s/17RFxZUSM1l3TQouIv4+IrRFxe1vb6oi4ISLuKd+PWohtGRodIqIJfBy4ADgdeGtEnF5vVT03DvynzDwdeCHw7w+DzzzpL4C76i6ijz4KfD0znw2cxSL/7BFxIvAfgLHMPANoAm+pt6qe+AxwfkfbpcA3MvM04Bvl/EEzNGY6B7g3M+/PzD3AF4CLaq6ppzLzkcz8aTm9neKL5MR6q+q9iFgH/GvgU3XX0g8RcQTwMuDvADJzT2Y+UWtR/TEELI2IIWAZ8Kua61lwmfkd4LGO5ouAK8rpK4DXL8S2DI2ZTgQeapvfzGHwBTopItYDzwd+XHMp/fAR4L1Aq+Y6+uUUYBvw6fKQ3KciYnndRfVSZj4M/A3wS+AR4LeZeX29VfXN2sx8pJz+NbB2IQY1NLRPRKwAvgz8ZWY+WXc9vRQRrwW2ZuZNddfSR0PA7wL/KzOfDzzFAh2yGFTlcfyLKALzBGB5RPxxvVX1XxbXVizI9RWGxkwPAye1za8r2xa1iBimCIzPZeZX6q6nD14MXBgRD1AcgnxlRPxjvSX13GZgc2ZO7kV+iSJEFrPzgF9k5rbM3At8BXhRzTX1y5aIOB6gfN+6EIMaGjNtBE6LiFMiYoTipNk1NdfUUxERFMe578rMD9ddTz9k5vsyc11mrqf4f/zNzFzU/wLNzF8DD0XE75RNrwLurLGkfvgl8MKIWFb+OX8Vi/zkf5trgIvL6YuBqxdi0KGFGGQxyczxiHgXcB3FLy3+PjPvqLmsXnsx8Hbgtoi4pWz7z5l5bX0lqUfeDXyu/AfR/cA7aq6npzLzxxHxJeCnFL8SvJlFeEuRiLgSeDmwJiI2A+8HPgR8MSL+lOLxEG9akG15GxFJUlUenpIkVWZoSJIqMzQkSZUZGpKkygwNSVJlhoYOCxGxow/b+EGvt9Gxvde331gyIv5rRJzXzxp0+PEntzosRMSOzFxRdx1VRMRQZo5X6PcZ4KuZ+aXeVyUV3NPQYSUiXh4R/xwRV0fE/RHxoYj4NxHxk4i4LSI2lP1eFxE/Lm/sd2NErC3bjymfTXBHecO/ByNiTblsR9s2vt323IrPlVcjExH/JSI2ls92uLyt/dsR8ZGI2ARcFhG/KG/tQkSsap8v214EXAj8dUTcEhEbIuIzEfGGcvkDEfFX5bJNEfG7EXFdRNwXEX/eNs57ynr+JSI+2If/BTrEGRo6HJ0F/DnwHIor4Z+VmedQ3CL93WWf7wEvLG/s9wWKu+FCcaXtNzPzuRT3bjq5yzaeD/wlxTNZnklx1T3A32bm75XPdlgKvLZtnZHMHMvMDwLfprhtOxS3OflKee8kADLzBxS3iXhPZp6dmffNUsMvM/Ns4LsUz1t4A8XzUj4IEBGvBk6jeBzA2cALIuJlXT6PBBgaOjxtLJ8hshu4D5i8VfZtwPpyeh1wXUTcBrwHeG7Z/hKKECEzvw483mUbP8nMzZnZAm5pG/cV5R7MbcAr28YF+L9t059i6hYf7wA+Pc/PCFP3TLsN+HFmbs/MbcDuiDgSeHX5upniNhvPpggRqSvvPaXD0e626VbbfIupvxMfAz6cmddExMuBDxzENiaAofIxo5+geIrcQxHxAaD90aNPTU5k5vcjYn257WZm3s78tX+uzs88BATwV5n5yQMYW4cp9zSk2R3B1C3xL25r/z7ljd/Kwzvzee7yZEA8Wj675A376f9Z4PN038vYDqycx/Y7XQe8s6yFiDgxIo49iPF0GDA0pNl9APh/EXET8Ghb+weBV0fE7cAbKZ6Itr3KgOWjVf8PcDvFF/bG/azyOYpQurLL8i8A7ylP1m+oUkNHPddThNIPy8NlX+LgQkiHAX9yK81DRCwBJspb6J9L8RS8s3u0rTcAF2Xm23sxvnQgPKchzc/JFM8oaAB7gD/rxUYi4mPABcBrejG+dKDc05AkVeY5DUlSZYaGJKkyQ0OSVJmhIUmqzNCQJFX2/wF61Dm+FCO7xwAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "\"\"\"\n",
    "Solve for LiH molecule with an internuclear distance of 1.525 a.u. using imaginary time propagation.\n",
    "\"\"\"\n",
    "\n",
    "V = np.zeros((n_rho, n_z))\n",
    "psi0 = np.zeros((n_rho, n_z))\n",
    "psi1 = np.zeros((n_rho, n_z))\n",
    "Z1 = 3.0\n",
    "Z2 = 1.0\n",
    "dist = 1.525\n",
    "\n",
    "# Setup potential and initial guess for the orbitals.\n",
    "for i in range(n_rho):\n",
    "    for j in range(n_z):\n",
    "        r_p = np.sqrt(rho_inner[i] ** 2 + (z_inner[j] - dist) ** 2)\n",
    "        r_m = np.sqrt(rho_inner[i] ** 2 + (z_inner[j] + dist) ** 2)\n",
    "        V[i, j] = Verfgau(r_p, Z=Z1) + Verfgau(r_m, Z=Z2)\n",
    "        psi0[i, j] = np.sqrt(rho_inner[i]) * np.exp(-Z1 * r_p)\n",
    "        psi1[i, j] = np.sqrt(rho_inner[i]) * np.exp(-Z2 * r_m)\n",
    "\n",
    "V = V.flatten()\n",
    "Fock = CylinderFockMatrix(Laplacian, V)\n",
    "\n",
    "##############################################################################\n",
    "# Orthonormalize the orbitals\n",
    "Psi = np.zeros((2, n_rho * n_z))\n",
    "norm_psi0 = np.sum(np.abs(psi0.ravel()) ** 2) * drho * dz\n",
    "psi0 /= np.sqrt(norm_psi0)\n",
    "\n",
    "norm_psi1 = np.sum(np.abs(psi1.ravel()) ** 2) * drho * dz\n",
    "psi1 /= np.sqrt(norm_psi1)\n",
    "\n",
    "Psi[0] = psi0.ravel()\n",
    "Psi[1] = psi1.ravel()\n",
    "\n",
    "S = np.zeros((2, 2))\n",
    "for i in range(2):\n",
    "    for j in range(2):\n",
    "        S[i, j] = np.sum(Psi[i].conj() * Psi[j]) * drho * dz\n",
    "eps, C = np.linalg.eigh(S)\n",
    "Psi = np.dot(C, Psi)\n",
    "for i in range(2):\n",
    "    Psi[i] /= np.sqrt(np.sum(np.abs(Psi[i]) ** 2) * drho * dz)\n",
    "##############################################################################\n",
    "dt = 0.001\n",
    "n_t = int(1e4)\n",
    "imag_time = np.linspace(0, n_t * dt, n_t)\n",
    "energy = np.zeros((2, n_t))\n",
    "##############################################################################\n",
    "\n",
    "W_direct = np.zeros((2, n_rho * n_z))\n",
    "for n in tqdm.tqdm(range(n_t - 1)):\n",
    "\n",
    "    # Compute direct potential\n",
    "    for i in range(2):\n",
    "        P = np.abs(Psi[i]) ** 2 / np.sqrt(2 * np.pi)\n",
    "        P = P.reshape((n_rho, n_z))\n",
    "        P = np.einsum(\"i, ij->ij\", 1 / rho_inner, P)\n",
    "        f = -4 * np.pi * np.einsum(\"i, ij->ij\", np.sqrt(rho_inner), P)\n",
    "        w_ii = Poisson.solve(f, inhomgenous_dirichlet=True)\n",
    "        W_ii_direct = np.einsum(\"i, ij->ij\", 1 / np.sqrt(rho_inner), w_ii)\n",
    "        W_direct[i] = W_ii_direct.ravel()\n",
    "\n",
    "    # Compute exchange potential\n",
    "    P_10 = Psi[1] * Psi[0] / np.sqrt(2 * np.pi)\n",
    "    P_10 = P_10.reshape((n_rho, n_z))\n",
    "    P_10 = np.einsum(\"i, ij->ij\", 1 / rho_inner, P_10)\n",
    "    f = -4 * np.pi * np.einsum(\"i, ij->ij\", np.sqrt(rho_inner), P_10)\n",
    "    w_10 = Poisson.solve(f, inhomgenous_dirichlet=False)\n",
    "    W_10 = np.einsum(\"i, ij->ij\", 1 / np.sqrt(rho_inner), w_10)\n",
    "    W_10 = W_10.ravel()\n",
    "\n",
    "    F_psi = Fock.F_Psi(Psi, W_direct, W_10)\n",
    "    energy[0, n] = np.sum(Psi[0] * F_psi[0]) * drho * dz\n",
    "    energy[1, n] = np.sum(Psi[1] * F_psi[1]) * drho * dz\n",
    "\n",
    "    k1 = -F_psi\n",
    "    k2 = -Fock.F_Psi(Psi + 0.5 * dt * k1, W_direct, W_10)\n",
    "    k3 = -Fock.F_Psi(Psi + 0.5 * dt * k2, W_direct, W_10)\n",
    "    k4 = -Fock.F_Psi(Psi + dt * k3, W_direct, W_10)\n",
    "    Psi += dt / 6 * (k1 + 2 * k2 + 2 * k3 + k4)\n",
    "\n",
    "    S = np.zeros((2, 2))\n",
    "    for i in range(2):\n",
    "        for j in range(2):\n",
    "            S[i, j] = np.sum(Psi[i].conj() * Psi[j]) * drho * dz\n",
    "    eps, C = np.linalg.eigh(S)\n",
    "    Psi = np.dot(C, Psi)\n",
    "    for i in range(2):\n",
    "        Psi[i] /= np.sqrt(np.sum(np.abs(Psi[i]) ** 2) * drho * dz)\n",
    "\n",
    "\n",
    "# Compute direct potential\n",
    "for i in range(2):\n",
    "    P = np.abs(Psi[i]) ** 2 / np.sqrt(2 * np.pi)\n",
    "    P = P.reshape((n_rho, n_z))\n",
    "    P = np.einsum(\"i, ij->ij\", 1 / rho_inner, P)\n",
    "    f = -4 * np.pi * np.einsum(\"i, ij->ij\", np.sqrt(rho_inner), P)\n",
    "    w_ii = Poisson.solve(f, inhomgenous_dirichlet=True)\n",
    "    W_ii_direct = np.einsum(\"i, ij->ij\", 1 / np.sqrt(rho_inner), w_ii)\n",
    "    W_direct[i] = W_ii_direct.ravel()\n",
    "\n",
    "# Compute exchange potential\n",
    "P_10 = Psi[1] * Psi[0] / np.sqrt(2 * np.pi)\n",
    "P_10 = P_10.reshape((n_rho, n_z))\n",
    "P_10 = np.einsum(\"i, ij->ij\", 1 / rho_inner, P_10)\n",
    "f = -4 * np.pi * np.einsum(\"i, ij->ij\", np.sqrt(rho_inner), P_10)\n",
    "w_10 = Poisson.solve(f, inhomgenous_dirichlet=False)\n",
    "W_10 = np.einsum(\"i, ij->ij\", 1 / np.sqrt(rho_inner), w_10)\n",
    "W_10 = W_10.ravel()\n",
    "\n",
    "F_psi = Fock.F_Psi(Psi, W_direct, W_10)\n",
    "energy[0, n_t - 1] = np.sum(Psi[0] * F_psi[0]) * drho * dz\n",
    "energy[1, n_t - 1] = np.sum(Psi[1] * F_psi[1]) * drho * dz\n",
    "\n",
    "print(f\"Imag time final energy: {energy[0,-1]:.5f}, {energy[1,-1]:.5f}\")\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(imag_time, energy[0, :])\n",
    "plt.plot(imag_time, energy[1, :])\n",
    "plt.xlabel(\"Imaginary time\")\n",
    "plt.ylabel(r\"$E(\\tau)$\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}