{ "cells": [ { "cell_type": "markdown", "id": "8653dd44", "metadata": {}, "source": [ "# Imaginary time projection method\n", "Target: Construct a code that numerically finds the ground state energy $E_0$ and the corresponding wave function $\\psi_0$ of the Schrodinger equation \n", "\\begin{equation}\n", "\\hat{H}\\psi_0(x) = E_0\\psi_0(x)\n", "\\end{equation}\n", "\n", "Consider the following equation \n", "\\begin{equation}\n", "-\\hbar\\frac{\\partial \\varphi(x,\\tau)}{\\partial \\tau} = \\hat{H}\\varphi(x,\\tau).\\label{eqn:imgt}\n", "\\end{equation}\n", "Its formal solution is\n", "\\begin{equation}\n", "\\varphi(x,\\tau) = e^{-\\tau \\hat{H}/\\hbar}\\varphi_0(x),\\label{eqn:sol}\n", "\\end{equation}\n", "where $\\varphi_0(x)$ is an arbitrary initial wave function, which should be set by initial conditions. Assume for a moment that we know all solutions of our target equation, i.e.\n", "\\begin{equation}\n", "\\hat{H}\\psi_n(x) = E_n\\psi_n(x).\n", "\\end{equation}\n", "The eigenstates $\\psi_n$ of any hermitian operator constitute a basis, so any function can be written as\n", "\\begin{equation}\n", "\\varphi_0(x) = \\sum_n c_n \\psi_n(x),\\label{eqn:expansion}\n", "\\end{equation}\n", "where $c_n$ are expansion coefficients (their values will not be needed in practice). Then, we get\n", "\\begin{equation}\n", "\\varphi(x,\\tau) = e^{-\\tau \\hat{H}/\\hbar}\\sum_n c_n \\psi_n(x)=\\sum_n c_ne^{-\\tau \\hat{H}/\\hbar}\\psi_n(x)=\\sum_n c_ne^{-\\tau E_n/\\hbar}\\psi_n(x).\n", "\\end{equation}\n", "The eigenenergies are organized as follows: $E_00$ were assumed to be negligible in comparison to $n=0$ contribution.\n", "The remaining unknown coefficient $C_0$ can be obtained from the normalization condition.\n", "\n", "Finally, the ground state energy can be found as\n", "\\begin{equation}\n", "E_0 = \\int \\psi_0^*(x)\\hat{H}\\psi_0(x) dx \n", "\\end{equation} " ] }, { "cell_type": "markdown", "id": "486587ea", "metadata": {}, "source": [ "# Implementation (one of the possibilities)\n", "To make use of the imaginary time evolution method we need a prescription how to execute the operation $e^{-\\tau \\hat{H}/\\hbar}\\varphi_0(x)$. We can do it as follow (I provide the simplest prescription, there are also more efficient prescriptions)\n", "\\begin{equation}\n", "\\varphi(x,\\tau)=e^{-\\tau \\hat{H}/\\hbar}\\varphi_0(x)=e^{-\\Delta\\tau \\hat{H}/\\hbar}e^{-\\Delta\\tau \\hat{H}/\\hbar}\\ldots e^{-\\Delta\\tau \\hat{H}/\\hbar}\\varphi_0(x),\n", "\\end{equation}\n", "which means that we can apply the elementary operation $e^{-\\Delta\\tau \\hat{H}/\\hbar}$ in a loop, until the result converge to the ground state (in practice with some precision $\\varepsilon$). If $\\Delta\\tau$ is sufficiently small we can use the Taylor expansion\n", "\\begin{equation}\n", "e^{-\\Delta\\tau \\hat{H}/\\hbar}\\varphi(x)\\approx (1-\\Delta\\tau \\hat{H}/\\hbar)\\varphi(x).\n", "\\label{eqn:singlestep}\n", "\\end{equation}\n", "One can also use higher order expansion." ] }, { "cell_type": "markdown", "id": "b24acbfe", "metadata": {}, "source": [ "# Wave function representation\n", "\n", "Next we need to represent the wave function. The function $\\varphi(x)$ can be represented as an array of numbers $\\varphi[k]$ of size $N$ where $\\varphi[k]\\equiv \\varphi(x_0+ka)$ with $k=0,1,\\ldots,N-1$. The $x$ domain spans interval $[x_0,x_0+a(N-1)]$, and $a$ is a space discretization step. \n", "\n", "The $x_0$, $a$ and $N$ will be the algorithm parameters, and one needs to select them in such a way to assure that the solution vanishes at the boundaries and at the same time the discretization is sufficient to resolve structures of the wave function. Then, the external potential is also represented as a vector $V(x)\\rightarrow V[k]$. The operation $V(x)\\varphi(x)$ is equivalent to element-by-element multiplication. Computation of derivatives can be executed either by using finite-difference formulas or by spectral methods. \n", "\n", "Integrals can be computed as follow (one can use more accurate formulas for integrals, like Simpson's quadrature)\n", "\\begin{equation}\n", "\\int \\varphi(x)dx\\approx \\sum_{k=0}^{N-1} \\varphi[k] a.\n", "\\end{equation}" ] }, { "cell_type": "markdown", "id": "b2894f00", "metadata": {}, "source": [ "# Computation of derivatives via spectral method\n", "Spectral representation\n", "\\begin{equation}\n", "f(x)=\\frac{1}{2\\pi}\\int\\tilde{f}(k)e^{ikx}dk\n", "\\end{equation}\n", "Then\n", "\\begin{equation}\n", "\\frac{d^n}{dx^n}f(x)=\\frac{1}{2\\pi}\\int(ik)^n\\tilde{f}(k)e^{ikx}dk=\\frac{1}{2\\pi}\\tilde{F}^{(n)}(k)e^{ikx}dk,\n", "\\end{equation}\n", "Thus, the algorithm is:\n", "1. Compute forward FT of f(x) to obtain f(k)\n", "2. Construct $\\tilde{F}^{(n)}(k)=(ik)^n\\tilde{f}(k)$\n", "3. Compute backward FT of $\\tilde{F}^{(n)}(k)$ to obtain $\\frac{d^n f(x)}{dx^n}$" ] }, { "cell_type": "markdown", "id": "2c1dc5e4", "metadata": {}, "source": [ "# Simple GPE code" ] }, { "cell_type": "code", "execution_count": 1, "id": "be272456", "metadata": {}, "outputs": [], "source": [ "# imports\n", "import numpy as np\n", "import matplotlib.pyplot as plt\n", "import matplotlib.cm as cm" ] }, { "cell_type": "code", "execution_count": 2, "id": "fafb4c39", "metadata": {}, "outputs": [], "source": [ "# Constants\n", "m=2.0 # mass of dimer (two particles) \n", "hbar=1.0 # h/2pi\n", "N=100.0 # number of dimers\n", "a_s=0.1 # scattering length\n", "U0 = 4.*np.pi*hbar**2*a_s/m # Coupling constant - unregularized\n", "\n", "# Domain - 2D lattice \n", "nx=32\n", "ny=32\n", "dx=1.0 # lattice spacing\n", "dy=1.0 # lattice spacing\n", "\n", "# make grid\n", "x = np.arange(-dx*nx/2, dx*nx/2, dx) # domain - x coordinate\n", "y = np.arange(-dy*ny/2, dy*ny/2, dy) # domain - y coordinate\n", "xv, yv = np.meshgrid(x, y, indexing='ij')" ] }, { "cell_type": "code", "execution_count": 3, "id": "9ee26076", "metadata": {}, "outputs": [], "source": [ "# Computation of derivatives\n", "kx = np.fft.fftfreq(nx, d=dx)*2.*np.pi # domain - momentum space (x)\n", "ky = np.fft.fftfreq(ny, d=dy)*2.*np.pi # domain - momentum space (y)\n", "kxv, kyv = np.meshgrid(kx, ky, indexing='ij')\n", "k2v = kxv**2 + kyv**2 # k^2 = kx^2 + ky^2\n", "\n", "def laplace2d(f):\n", " f_k=np.fft.fft2(f)\n", " f_k = f_k*(-1.0)*k2v\n", " f_k = np.fft.ifft2(f_k)\n", " return f_k\n", "\n", "def dfdx2d(f):\n", " f_k=np.fft.fft2(f)\n", " f_k = f_k*(1.0j)*kxv\n", " f_k = np.fft.ifft2(f_k)\n", " return f_k\n", "\n", "def dfdy2d(f):\n", " f_k=np.fft.fft2(f)\n", " f_k = f_k*(1.0j)*kyv\n", " f_k = np.fft.ifft2(f_k)\n", " return f_k" ] }, { "cell_type": "code", "execution_count": 4, "id": "434f408f", "metadata": {}, "outputs": [], "source": [ "# Function for making 2D plots\n", "def plot2d(f):\n", " ratio=nx/ny\n", " scale=4.0\n", " fig, ax = plt.subplots(figsize=(scale*ratio,1.3*scale*ratio))\n", "\n", " im = ax.imshow(np.transpose(f), interpolation='bilinear', extent=[-nx/2*dx, (nx/2-1)*dx, -ny/2*dy, (ny/2-1)*dy], aspect='auto')\n", " fig.colorbar(im, ax=ax, orientation=\"horizontal\")\n", "\n", " plt.tight_layout()\n", " plt.show()" ] }, { "cell_type": "code", "execution_count": 5, "id": "5741107f", "metadata": {}, "outputs": [], "source": [ "# smooth step function - will be needed later\n", "def smooth_from0to1(x):\n", " if x<=0.0: return 0.0\n", " if x>=1.0: return 1.0\n", " return 0.5*( 1.0+np.tanh( np.tan( 0.5*np.pi*( 2.0*x-1.0 ) ) ) )\n", "def smooth_step(x, x1, x2):\n", " r=np.ones(x.shape)\n", " for ix in np.arange(nx):\n", " for iy in np.arange(ny):\n", " if x[ix,iy]<=x1:\n", " r[ix,iy]=1.0\n", " elif x[ix,iy]<=x2:\n", " r[ix,iy]=1.0-smooth_from0to1((x[ix,iy]-x1)/(x2-x1))\n", " else:\n", " r[ix,iy]=0.0\n", " return r\n", " \n", "mask_fun=smooth_step(np.sqrt(xv**2+yv**2), 0.85*nx/2, 0.95*nx/2)\n", "#plot2d(mask_fun)" ] }, { "cell_type": "code", "execution_count": 6, "id": "a7bcf2af", "metadata": {}, "outputs": [], "source": [ "# External potential\n", "omega_x = 0.05\n", "omega_y = 0.05\n", "Vext = m*omega_x**2*xv**2/2 + m*omega_y**2*yv**2/2 # 2D harmonic potential\n", "#plot2d(Vext)" ] }, { "cell_type": "code", "execution_count": 7, "id": "20d86505", "metadata": {}, "outputs": [], "source": [ "# this function applies Hamitonian to the wave-function\n", "def apply_H(psi, Vext, nonlinear_term):\n", " return (-1.0*hbar**2/(2.*m))*laplace2d(psi) + Vext*psi + nonlinear_term*psi\n", "\n", "# # version with the rotating frame\n", "# def apply_H(psi, Vext, nonlinear_term):\n", "# Omega=0.05\n", "# omLz = Omega*(xv*(-1.j*hbar*dfdy2d(psi)) -yv*(-1.j*hbar*dfdx2d(psi)))\n", "# omLz = omLz*mask_fun\n", "# return (-1.0*hbar**2/(2.*m))*laplace2d(psi) + Vext*psi + nonlinear_term*psi -omLz" ] }, { "cell_type": "code", "execution_count": 8, "id": "90a2d154", "metadata": {}, "outputs": [], "source": [ "# this function computes energy\n", "def energy(psi):\n", " nonlin=U0*np.abs(psi)**2 # nonlinear term\n", " Hpsi=apply_H(psi,Vext,nonlin) # H*psi\n", " return np.real(np.sum(np.conj(psi)*Hpsi)*dx*dy)" ] }, { "cell_type": "code", "execution_count": 9, "id": "3518f596", "metadata": {}, "outputs": [], "source": [ "# imaginary time evolution step\n", "def img_step(dtau, psi):\n", " nonlin=U0*np.abs(psi)**2 # nonlinear term\n", " Hpsi=apply_H(psi,Vext,nonlin) # H*psi\n", " HHpsi=apply_H(Hpsi,Vext,nonlin) # H*H*psi\n", " HHHpsi=apply_H(HHpsi,Vext,nonlin) # H*H*H*psi\n", " return (psi -dtau*Hpsi/hbar +dtau**2*HHpsi/hbar**2/2. -dtau**3*HHHpsi/hbar**3/6.)" ] }, { "cell_type": "code", "execution_count": 10, "id": "2d6fec1e", "metadata": {}, "outputs": [], "source": [ "# normalize psi such that integral |psi|^2 dxdy = N\n", "def normalize(psi):\n", " norm=np.sum(np.abs(psi)**2)*dx*dy\n", " return psi*np.sqrt(N/norm)" ] }, { "cell_type": "markdown", "id": "988a1288", "metadata": {}, "source": [ "## Main code" ] }, { "cell_type": "code", "execution_count": 16, "id": "1476b63b", "metadata": {}, "outputs": [], "source": [ "# initial wave-function\n", "# psi=normalize(np.ones([nx,ny]))\n", "psi=normalize( np.random.random([nx,ny])+np.random.random([nx,ny])*1.j )" ] }, { "cell_type": "code", "execution_count": 17, "id": "67d3791e", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ " 66: E= 32.16168350 diff= 0.00000027\r" ] } ], "source": [ "# compute initial value of the energy\n", "Eold=energy(psi)\n", "Enew=Eold\n", "\n", "max_sweeps=1000\n", "nsteps=100\n", "dtau=0.01\n", "eps=1.0e-6\n", "for isweep in range(max_sweeps):\n", " # do imaginary time evolution for nsteps\n", " for istep in range(nsteps):\n", " psi=img_step(dtau, psi)\n", " psi=normalize(psi) # normalize after each step\n", " \n", " # check new energy\n", " Enew=energy(psi)\n", " print(\"%8d: E=%12.8f diff=%12.8f\" % (isweep, Enew,np.abs(Enew-Eold)), end='\\r')\n", " \n", " # break if energy change is smaller than the treshold\n", " if np.abs(Enew-Eold)" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "# plot density distribution\n", "plot2d(np.abs(psi)**2)" ] }, { "cell_type": "code", "execution_count": null, "id": "f6ee5ef4", "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.10.11" } }, "nbformat": 4, "nbformat_minor": 5 }