notebook/quantum-mechanics/msoft.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7f7a3068",
   "metadata": {},
   "source": [
    "# **The Multiple Split Operator Fourier Transform (MSOFT) method**\n",
    "\n",
    "**Authors:** Taylor Baird and Sara Bonella\n",
    "\n",
    "<i class=\"fa fa-home fa-2x\"></i><a href=\"../index.ipynb\" style=\"font-size: 20px\"> Go back to index</a>\n",
    "\n",
    "\n",
    "**Source code:**  https://github.com/osscar-org/quantum-mechanics/blob/master/notebook/quantum-mechanics/msoft.ipynb\n",
    "\n",
    "This notebook carries out the numerical solution of the 1D time-dependent Schrödinger equation for nuclear evolution on multiple electronic potential energy surfaces via the Multiple Split Operator Fourier Transform (MSOFT) method. \n",
    "\n",
    "<hr style=\"height:1px;border:none;color:#cccccc;background-color:#cccccc;\" />"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3736f929-a80f-486f-99d5-16dcb42f2895",
   "metadata": {},
   "source": [
    "# **Goals**\n",
    "* Understand the different steps in the MSOFT algorithm and how they are translated to code. \n",
    "* Familiarize yourself with the key assumptions underlying the algorithm and how their validity depends on the values of the parameters used in the simulation.\n",
    "* Use the MSOFT algorithm to investigate the nonadiabatic dynamics of a nuclear wavepacket propagating in two different 2-level systems (a double harmonic potential and the Tully potential)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f2cc0f1",
   "metadata": {},
   "source": [
    "# **Background Theory**\n",
    "\n",
    "[More in the background theory](./theory/theory_msoft.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4ea8caf",
   "metadata": {},
   "source": [
    "## **Tasks and exercises**\n",
    "\n",
    "1. Investigate the dependence on timestep of the stability of the dynamics (try moving the slider for $dt$ and monitor the behaviour of the various control properties of the simulation).\n",
    "    <details>\n",
    "    <summary style=\"color: red\">Solution</summary>\n",
    "        One may observe that as the timestep employed in the simulation is increased, the conservation of total energy of the system degrades until eventually the timestep is so large that the dynamics becomes totally unstable. The reason why this integration scheme does not conserve total energy may be attributed to the non-commutativity of the split-operator propagator with the Hamiltonian. It is worth noting, however, that norm conservation is maintained even as one increases the timestep. This latter fact is due to the unitarity of the propagator in the split-operator scheme.\n",
    "    </details>\n",
    "\n",
    "2. What dictates the maximum size of the timestep that we can use for the MSOFT algorithm (consider the main assumptions/approximations that are made when formulating the propagation scheme for MSOFT).\n",
    "    <details>\n",
    "    <summary style=\"color: red\">Solution</summary>\n",
    "    Recall that the central approximation used in the derivation of the MSOFT propagation scheme is the Trotter product formula: $e^{A+B}  \\underbrace{=}_{P \\to \\infty} (e^{\\frac{A}{P}} e^{\\frac{B}{P}})^{P}$. This approximation becomes more and more accurate the larger we make the value of $P$. In our specific case we have $e^{\\frac{it}{\\hbar} (\\hat{T} + \\hat{V} )}$ and we approximate this via the product \n",
    "        $(e^{\\frac{idt}{\\hbar}\\hat{T} } e^{\\frac{idt}{\\hbar}\\hat{V} })^{N_{\\text{steps}}}$ where $N_{\\text{steps}}\\cdot dt = t$.\n",
    "    </details>\n",
    "\n",
    "3. Why is the use of the MSOFT algorithm unfeasible for larger systems (think about how much data is needed to represent the state of the system on a computer and how many operations are required to carry out the propagation)?\n",
    "    <details>\n",
    "    <summary style=\"color: red\">Solution</summary>\n",
    "    In order to implement the MSOFT algorithm on a computer it is necessary to discretize the nuclear wavefunction. To do so, we must introduce a grid of points that make up the space throughout which the wavefunction extends. Say that for each dimension of the grid we use $N$ grid points. For a system in $d$ dimensions, this means that we shall require $N^d$ points to represent the wavefunction of a single nucleus. Now, if we want to instead consider a system of say $n$ nuclei then we find that a total number of $N^{nd}$ grid points are required. In other words, the amount of data required to represent our system scales exponentially with increasing system size. Moreover, since we must loop over each of these grid points to carry out the time evolution of our system - the number of operations required also scales exponentially. This is what renders the MSOFT algorithm unsuitable for large systems.\n",
    "    </details>\n",
    "\n",
    "4. Investigate the effect of increasing the coupling strength between different electronic states on the evolution of the two state populations (vary the $C$ slider for each potential).\n",
    "    <details>\n",
    "    <summary style=\"color: red\">Solution</summary>\n",
    "        One can note that by increasing the value of $C$ the probability of inducing a transition of the nuclear wavepacket from one electronic potential energy surface to the other is modulated. You can trace this back to the equations of motion given in the theory section.\n",
    "    </details>\n",
    "\n",
    "<hr style=\"height:1px;border:none;color:#cccccc;background-color:#cccccc;\" />"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af938e9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Notebook demonstrating the use of the MSOFT algorithm for quantum dynamics.\n",
    "\"\"\"\n",
    "%matplotlib widget\n",
    "import numpy as np\n",
    "from math import pi\n",
    "from numpy.fft import fft as fft\n",
    "from numpy.fft import ifft as ifft\n",
    "\n",
    "from ipywidgets import Button, FloatSlider, HBox, VBox, Accordion, Label, IntProgress, Dropdown, Layout\n",
    "from IPython.display import display, Markdown\n",
    "import ipywidgets as widgets\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import animation\n",
    "import matplotlib.ticker as ticker\n",
    "from IPython.display import clear_output\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e04d5db5",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MSOFT(object):\n",
    "    \"\"\"\n",
    "    Class to encapsulate the various ingredients of the MSOFT method.\n",
    "    x: real space grid\n",
    "    psi_0: inital total wavefunction, including real and imaginary parts of both electronic states\n",
    "    h: the potential in diabatic representation\n",
    "    dt: the time interval\n",
    "    hbar: the Plank constant (default value 1.0 as atomic unit)\n",
    "    m: the mass (default value 2000 au)\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, x, psi_0, h, dt, hbar=1.0, m = 1.0):\n",
    "        self.N = len(x);\n",
    "        self.hbar = hbar;\n",
    "        self.m = m;\n",
    "        \n",
    "        dx = x[1] -x [0]; #real space grid spacing\n",
    "        dk = (2*pi)/(self.N*dx); # reciprocal space grid spacing \n",
    "        \n",
    "        self.x = x;\n",
    "        self.dx = dx;\n",
    "        self.dk = dk;\n",
    "        self.t = 0.0;\n",
    "        self.dt = dt;\n",
    "        self.tot_steps=100000\n",
    "\n",
    "#         self.tot_steps=int(1e3/self.dt)\n",
    "        self.tarr=self.dt*np.arange(self.tot_steps)\n",
    "        self.m = m;\n",
    "        \n",
    "        self.k_x = np.zeros(self.N)\n",
    "        for sx in range(int(self.N/2)):\n",
    "            self.k_x[sx]=2*pi*sx/(dx*self.N)\n",
    "        for sx in range(int(self.N/2),self.N):\n",
    "            self.k_x[sx]=2*pi*(sx-self.N)/(dx*self.N)\n",
    "\n",
    "        self.psi_x = np.zeros((self.N,2), dtype=np.complex128) # two components (one for each electronic\n",
    "        self.psi_k = np.zeros((self.N,2), dtype=np.complex128) # state)\n",
    "\n",
    "        self.P1 = np.array([np.nan]*self.tot_steps) # to track populations of each \n",
    "        self.P2 = np.array([np.nan]*self.tot_steps) # electronic state\n",
    "        \n",
    "        self.ekint = np.array([np.nan]*self.tot_steps)\n",
    "        self.epot  = np.array([np.nan]*self.tot_steps)\n",
    "\n",
    "        self.psi_x = psi_0;\n",
    "        self._periodic_bc()\n",
    "\n",
    "        if callable(h):\n",
    "            self.h = h(x) \n",
    "        else:\n",
    "            self.h = h  \n",
    "                \n",
    "        # Diagonalize the diabatic coupling matrix here and store eigenvalues.\n",
    "        self._init_pot_prop()\n",
    "              \n",
    "    def _init_pot_prop(self):\n",
    "        \n",
    "        LX=self.x[-1]\n",
    "        self.u = np.zeros(shape=(self.N,2),dtype=np.complex128)\n",
    "        self.E_eigvals = np.zeros((self.N,2))\n",
    "        self.De = np.zeros(shape=(self.N,2,2))\n",
    "        self.DeT = np.zeros(shape=(self.N,2,2))\n",
    "\n",
    "        for i in range(self.N):\n",
    "            self.E_eigvals[i][0] = (self.h[i][0][0]+self.h[i][1][1]+np.sqrt((self.h[i][0][0]-self.h[i][1][1])*(self.h[i][0][0]-self.h[i][1][1])+4.*self.h[i][0][1]*self.h[i][1][0]))/2.;\n",
    "            self.E_eigvals[i][1] = (self.h[i][0][0]+self.h[i][1][1]-np.sqrt((self.h[i][0][0]-self.h[i][1][1])*(self.h[i][0][0]-self.h[i][1][1])+4.*self.h[i][0][1]*self.h[i][1][0]))/2.\n",
    "            if(i<=self.N/2):\n",
    "        #    /* 1st column = 1st eigenvector */\n",
    "                self.De[i][0][0] = self.h[i][0][1]/(self.E_eigvals[i][0]-self.h[i][0][0]);\n",
    "                self.De[i][1][0] = 1;\n",
    "                norm_fac = np.sqrt(self.De[i][0][0]*self.De[i][0][0]+1.);\n",
    "                self.De[i][0][0] /= norm_fac;\n",
    "                self.De[i][1][0] /= norm_fac;\n",
    "                #   /* 2nd column = 2nd eigenvector */\n",
    "                self.De[i][0][1] = -self.De[i][1][0];\n",
    "                self.De[i][1][1] = self.De[i][0][0]\n",
    "\n",
    "            else:            \n",
    "                self.De[i][0][1] = self.h[i][0][1]/(self.E_eigvals[i][1]-self.h[i][0][0]);\n",
    "                self.De[i][1][1] = 1;\n",
    "                norm_fac = np.sqrt(self.De[i][0][1]*self.De[i][0][1]+1.);\n",
    "                self.De[i][0][1] /= norm_fac;\n",
    "                self.De[i][1][1] /= norm_fac;\n",
    "                \n",
    "                self.De[i][1][0] = -self.De[i][0][1];\n",
    "                self.De[i][0][0] = self.De[i][1][1];\n",
    "\n",
    "            self.DeT[i][0][0] = self.De[i][0][0];\n",
    "            self.DeT[i][1][1] = self.De[i][1][1];\n",
    "            self.DeT[i][1][0] = self.De[i][0][1];\n",
    "            self.DeT[i][0][1] = self.De[i][1][0];\n",
    "        \n",
    "            self.u[i][0] = np.exp(-0.5j*self.dt*self.E_eigvals[i][0])\n",
    "            self.u[i][1] = np.exp(-0.5j*self.dt*self.E_eigvals[i][1])\n",
    "\n",
    "        \n",
    "    def _periodic_bc(self):\n",
    "        self.psi_x[-1][0] = self.psi_x[0][0]\n",
    "        self.psi_x[-1][1] = self.psi_x[0][1]\n",
    "\n",
    "               \n",
    "    def _half_pot_prop(self, ft=True):\n",
    "        if ft == True:\n",
    "     \n",
    "            self.psi_x[:,0] = ifft(self.psi_k[:,0])\n",
    "            self.psi_x[:,1] = ifft(self.psi_k[:,1])\n",
    "        \n",
    "        \n",
    "        self.psi_x=np.einsum('ijk,ik->ij',self.De,self.psi_x)\n",
    "        self.psi_x=self.psi_x*self.u\n",
    "        self.psi_x=np.einsum('ijk,ik->ij',self.DeT,self.psi_x)\n",
    "\n",
    "\n",
    "        self._periodic_bc()\n",
    "\n",
    "    def _full_kinetic_prop(self):\n",
    "        self.psi_k[:] = fft(self.psi_x[:],axis=0)\n",
    "        self.psi_k[:] = self.psi_k[:]*np.exp(-1.0j*self.k_x**2*self.dt/(2.0*self.m))[:,np.newaxis]            \n",
    "        \n",
    "    def _get_kinetic_energy(self, Nsteps):\n",
    "        #ke=0.5*(1./self.m)*sum(np.ravel(np.einsum('i,ik->ik',self.k_x**2,np.absolute(self.psi_k)**2)))*self.dk\n",
    "        psi_k_0=fft(self.psi_x[:,0])/self.N\n",
    "        psi_k_1=fft(self.psi_x[:,1])/self.N\n",
    "        #ke=0.5*(1./self.m)*sum((np.conj(psi_k_0)*psi_k_0*self.k_x**2 + np.conj(psi_k_1)*psi_k_1*self.k_x**2 ))*(1/.self.N*self.dx)\n",
    "        ke=0.\n",
    "        for sx in range(self.N):\n",
    "            k=self.k_x[sx]\n",
    "            ke+= 0.5*(1./self.m)*(k**2)*(np.linalg.norm(psi_k_0[sx])**2 + np.linalg.norm(psi_k_1[sx])**2)\n",
    "        ke*=(self.dx*self.N)\n",
    "        # ke=0.5*(1./self.m)*sum((np.conj(psi_k[:,0])*psi_k[:,0]*self.k_x**2 + np.conj(psi_k[:,1])*psi_k[:,1]*self.k_x**2 ))*self.dk\n",
    "        self.ekint[int(self.t/self.dt)] = np.real(ke)\n",
    "        \n",
    "    def _get_potential_energy(self, Nsteps):\n",
    "        epot = sum(np.conj(self.psi_x[:,0])*self.h[:,0,0]*self.psi_x[:,0] + \\\n",
    "                        np.conj(self.psi_x[:,0])*self.h[:,0,1]*self.psi_x[:,1] + \\\n",
    "                        np.conj(self.psi_x[:,1])*self.h[:,0,1]*self.psi_x[:,0] + \\\n",
    "                       np.conj(self.psi_x[:,1])*self.h[:,1,1]*self.psi_x[:,1] )*self.dx\n",
    "        self.epot[int(self.t/self.dt)] = np.real(epot)\n",
    "\n",
    "\n",
    "    def _pop_states(self,Nsteps):\n",
    "        # get the relative populations of each electronic state\n",
    "        # no multiplication by dx b/c of definition of psi_x.\n",
    "        pops = np.linalg.norm(self.psi_x,axis=0)**2\n",
    "        self.P1[int(self.t/self.dt)] = pops[0] # interpolate populations\n",
    "        self.P2[int(self.t/self.dt)] = pops[1] # between updates\n",
    "        \n",
    "    def _get_norm(self):\n",
    "        self.norm = np.linalg.norm(self.psi_x)\n",
    "             \n",
    "    def evolution(self, Nsteps=1):\n",
    "\n",
    "        for i in range(Nsteps):\n",
    "            self._half_pot_prop(ft = False)\n",
    "            self._full_kinetic_prop()         \n",
    "            self._half_pot_prop(ft = True)\n",
    "            \n",
    "        self._get_kinetic_energy(Nsteps)\n",
    "        self._get_potential_energy(Nsteps)\n",
    "        self._pop_states(Nsteps)\n",
    "        self._get_norm()\n",
    "        self.t += self.dt*Nsteps        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8056c5b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def gauss_x(x, a, x0, k0):\n",
    "    \"\"\"\n",
    "    a gaussian wave packet of width a, centered at x0, with momentum k0\n",
    "    \"\"\"\n",
    "    gauss = np.exp(-1.*(x-x0)**2/(4.*a**2)) +0.j\n",
    "    gauss*= np.exp(1.j*k0*(x-x0))\n",
    "    return gauss/np.linalg.norm(gauss)\n",
    "\n",
    "def double_harmonic(x,A,B,b,C,D):\n",
    "\n",
    "    NX = len(x)\n",
    "    h = np.zeros(shape=(NX,2,2))\n",
    "   \n",
    "    LX=x[-1]\n",
    "    for sx,x_i in enumerate(x):\n",
    "        \n",
    "        h[sx][0][0] = A*(x_i-B)*(x_i-B)\n",
    "        h[sx][0][1] = C*np.exp(-D*(x_i-0.5*LX)*(x_i-0.5*LX))\n",
    "        h[sx][1][0] =  h[sx][0][1]\n",
    "        h[sx][1][1] = A*(x_i-b)*(x_i-b)\n",
    "        \n",
    "    return h\n",
    "\n",
    "def tully_pot(x,A,B,C,D):\n",
    "   \n",
    "    NX = len(x)\n",
    "    h = np.zeros(shape=(NX,2,2))    \n",
    "          \n",
    "    for sx in range(NX):\n",
    "        x_i = -0.5*LX+dx*sx\n",
    "        \n",
    "        h[sx][0][0] = A*(1.-np.exp(-B*abs(x_i)))*np.sign(x_i)\n",
    "        h[sx][0][1] = C*np.exp(-D*x_i**2)\n",
    "        h[sx][1][0] =  h[sx][0][1]\n",
    "        h[sx][1][1] = -h[sx][0][0]\n",
    "         \n",
    "        if (sx==int(NX/2)):\n",
    "            h[int(NX/2)][0][0] = 0.\n",
    "            h[int(NX/2)][0][1] = C\n",
    "            h[int(NX/2)][1][0] = C\n",
    "            h[int(NX/2)][1][1] = 0. \n",
    "    return h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fd44a50",
   "metadata": {},
   "outputs": [],
   "source": [
    "######################################################################\n",
    "# Create the animation\n",
    "\n",
    "# specify time steps and duration\n",
    "dt = 0.25\n",
    "N_steps = 50\n",
    "t_max = 120\n",
    "frames = int(t_max / float(N_steps * abs(dt)))\n",
    "\n",
    "# Specify frequency at which observables are computed \n",
    "NECAL = 10\n",
    "NNCAL = 1000\n",
    "\n",
    "# specify constants\n",
    "hbar = 1.0   # planck's constant\n",
    "m = 2000.      # particle mass\n",
    "\n",
    "# specify range in x coordinate\n",
    "LX=50.\n",
    "N=512\n",
    "dx = LX/N\n",
    "x = dx * np.arange(N)\n",
    "\n",
    "# specify potential\n",
    "A=0.001\n",
    "B=20.\n",
    "b=30.\n",
    "C=0.005\n",
    "D=1.0\n",
    "\n",
    "# specify initial momentum and quantities derived from it\n",
    "x0 = 19.0\n",
    "p0 = 20.\n",
    "s0 = 1.0\n",
    "dp2 = p0 * p0 \n",
    "d = hbar / np.sqrt(2 * dp2)\n",
    "\n",
    "state_idx = 0\n",
    "k0 = p0 / hbar\n",
    "v0 = p0 / m\n",
    "psi_x0 = np.zeros((N,2), dtype=np.complex128) # two components (one for each electronic state)\n",
    "psi_x0[:,state_idx] = gauss_x(x, s0, x0, p0) # IC is a Gaussian wavepacket on 1st state and 0 component on second state\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b699f3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "##################################################################################\n",
    "# All code below this point relates to visualization and interactivity           #\n",
    "# and can be ignored when considering only the specifics of the MSOFT algorithm. #\n",
    "##################################################################################\n",
    "\n",
    "style = {'description_width': 'initial'}\n",
    "\n",
    "pot_select = Dropdown(\n",
    "    options=['1. Double Harmonic', '2. Tully'],\n",
    "    index = 0,\n",
    "    description='Potential type:',\n",
    "    disabled=False,\n",
    "    style = style\n",
    ")\n",
    "\n",
    "state_select = Dropdown(\n",
    "    options=['1', '2'],\n",
    "    index = 0,\n",
    "    description='Choose initial electronic state.',\n",
    "    disabled=False,\n",
    "    style = style\n",
    ")\n",
    "\n",
    "layout_hidden  = widgets.Layout(visibility = 'hidden')\n",
    "layout_visible = widgets.Layout(visibility = 'visible')\n",
    "\n",
    "# Set params of double harmonic PES (pot1)\n",
    "\n",
    "s_dharm_A = FloatSlider(value = A, min = 0.0, max = 10.*A, step=0.001,readout_format='.5f', description = 'A: ')\n",
    "s_dharm_B = FloatSlider(value = B, min = 0.0, max = 10.*B, description = 'B: ')\n",
    "s_dharm_b = FloatSlider(value = b, min = 0.0, max = 10.*b, description = 'b: ')\n",
    "s_dharm_C = FloatSlider(value = C, min = 0.0, max = 0.5, step=0.001, readout_format='.4f', description = 'C: ')\n",
    "s_dharm_D = FloatSlider(value = D, min = 0.0, max = 10.*D, description = 'D: ')\n",
    "\n",
    "s_dharm_A.layout = layout_visible\n",
    "s_dharm_B.layout = layout_visible\n",
    "s_dharm_b.layout = layout_visible\n",
    "s_dharm_C.layout = layout_visible\n",
    "s_dharm_D.layout = layout_visible\n",
    "\n",
    "# Set params for Tully PES (pot2)\n",
    "\n",
    "s_tully_A = FloatSlider(value = 0.01, min = 0.01, max = 0.25,step=0.01, description = 'A: ');\n",
    "s_tully_B = FloatSlider(value = 1.6, min = 0.0, max = 3.2, description = 'B: ');\n",
    "s_tully_C = FloatSlider(value = 0.005, min = 0.0, max = 0.05,step=0.001, description = 'C: ');\n",
    "s_tully_D = FloatSlider(value = 1.0, min = 0.0, max = 10.0, description = 'D: ');\n",
    "\n",
    "\n",
    "s_tully_A.layout = layout_hidden \n",
    "s_tully_B.layout = layout_hidden \n",
    "s_tully_C.layout = layout_hidden \n",
    "s_tully_D.layout = layout_hidden \n",
    "\n",
    "#Show the potential image\n",
    "file1 = open(\"images/dharm_pot.png\", \"rb\")\n",
    "image1 = file1.read()\n",
    "file2 = open(\"images/tully_potential.png\", \"rb\")\n",
    "image2 = file2.read()\n",
    "\n",
    "pot_img = widgets.Image(\n",
    "    value=image1,\n",
    "    format='png',\n",
    "    width=700,\n",
    "    height=700,\n",
    ")\n",
    "\n",
    "l = Layout(flex='0 1 auto', height='80px', min_height='80px', width='auto')\n",
    "pot_eqn=Label(value=r'\\(h(x_i)=\\begin{pmatrix} A(1-e^{-B|\\frac{LX}{2}-x_i|})\\cdot \\text{sgn}(\\frac{LX}{2}-x_i) & Ce^{-D(\\frac{LX}{2}-x_i)^2} \\\\Ce^{-D(\\frac{LX}{2}-x_i)^2} & -A(1-e^{-B|\\frac{LX}{2}-x_i|})\\cdot \\text{sgn}(\\frac{LX}{2}-x_i)\\end{pmatrix}\\)',layout=l)\n",
    "#pot_eqn=Label(value=r'\\(h(x_i)=\\begin{pmatrix} h(x_i)=A(x_i-B)^2 & Ce^{(-D(x_i-\\frac{LX}{2})^2)} \\\\Ce^{(-D(x_i-\\frac{LX}{2})^2)} & A(x_i-b)^2\\end{pmatrix}\\)',layout=l)\n",
    "\n",
    "\n",
    "pot_img.layout = layout_visible\n",
    "\n",
    "def pot_change(change):\n",
    "    global h_x\n",
    "\n",
    "    if pot_select.index == 0:\n",
    "        s_dharm_A.layout = layout_visible\n",
    "        s_dharm_B.layout = layout_visible\n",
    "        s_dharm_b.layout = layout_visible\n",
    "        s_dharm_C.layout = layout_visible\n",
    "        s_dharm_D.layout = layout_visible\n",
    "        s_tully_A.layout = layout_hidden \n",
    "        s_tully_B.layout = layout_hidden \n",
    "        s_tully_C.layout = layout_hidden \n",
    "        s_tully_D.layout = layout_hidden \n",
    "        pot_img.layout = layout_visible\n",
    "        pot_img.value = image1\n",
    "        pot_eqn.value=r'\\(h(x_i)=\\begin{pmatrix} A(1-e^{-B|\\frac{LX}{2}-x_i|})\\cdot \\text{sgn}(\\frac{LX}{2}-x_i) & Ce^{-D(\\frac{LX}{2}-x_i)^2} \\\\Ce^{-D(\\frac{LX}{2}-x_i)^2} & -A(1-e^{-B|\\frac{LX}{2}-x_i|})\\cdot \\text{sgn}(\\frac{LX}{2}-x_i)\\end{pmatrix}\\)'\n",
    "\n",
    "        \n",
    "    elif pot_select.index == 1:\n",
    "      \n",
    "        s_dharm_A.layout = layout_hidden\n",
    "        s_dharm_B.layout = layout_hidden\n",
    "        s_dharm_b.layout = layout_hidden\n",
    "        s_dharm_C.layout = layout_hidden\n",
    "        s_dharm_D.layout = layout_hidden\n",
    "        s_tully_A.layout = layout_visible\n",
    "        s_tully_B.layout = layout_visible\n",
    "        s_tully_C.layout = layout_visible\n",
    "        s_tully_D.layout = layout_visible\n",
    "        pot_img.layout = layout_visible\n",
    "        pot_img.value = image2         \n",
    "        pot_eqn.value = r'\\(h(x_i)=\\begin{pmatrix} A(x_i-B)^2 & Ce^{(-D(x_i-\\frac{LX}{2})^2)} \\\\Ce^{(-D(x_i-\\frac{LX}{2})^2)} & A(x_i-b)^2\\end{pmatrix}\\)'\n",
    "        \n",
    "        \n",
    "def state_change(change):\n",
    "    global state_idx\n",
    "    if state_select.index==0:\n",
    "        state_idx = 0\n",
    "    else:\n",
    "        state_idx = 1\n",
    "    \n",
    "state_select.observe(state_change, names='value', type='change');\n",
    "pot_select.observe(pot_change, names='value', type='change');\n",
    "\n",
    "def on_pot_update(event):\n",
    "    global h_x\n",
    "    if pot_select.index == 0:\n",
    "        h_x = double_harmonic(x, s_dharm_A.value,s_dharm_B.value,s_dharm_b.value,s_dharm_C.value,s_dharm_D.value)\n",
    "    else:\n",
    "        h_x = tully_pot(x,s_tully_A.value, s_tully_B.value, s_tully_C.value, s_tully_D.value)\n",
    "        ax1.set_ylim(-0.05,0.25)\n",
    "        ax1.set_xlim(20,50)\n",
    "        tot_time=1800.*(30/p0) # scale x-axis depending on momentum\n",
    "        ax2.set_xlim(0,tot_time)\n",
    "        ax3.set_xlim(0,tot_time)\n",
    "        ax4.set_xlim(0,tot_time)\n",
    "\n",
    "\n",
    "pot_update = Button(description=\"Update potential\")\n",
    "pot_update.on_click(on_pot_update)\n",
    "\n",
    "# The accordion for the potential\n",
    "label1 = Label(\"(Press update potential and then update parameters to modify the potential.)\")\n",
    "pot_accordion = Accordion(children=[VBox([pot_select, pot_img, pot_eqn, state_select,\n",
    "                   HBox([s_dharm_A,s_dharm_B,s_dharm_b,s_dharm_C,s_dharm_D]),\n",
    "                   HBox([s_tully_A,s_tully_B,s_tully_C,s_tully_D]),\n",
    "                   HBox([pot_update, label1])])], selected_index = None)\n",
    "\n",
    "pot_accordion.set_title(0, \"Select potential and set parameters\")\n",
    "display(pot_accordion)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b19b0ff6",
   "metadata": {},
   "outputs": [],
   "source": [
    "w_mass = FloatSlider(value = 2000, min = 0.0, max=5000, description=\"mass: \") \n",
    "w_dt = FloatSlider(value = 0.25, min = 0.001, max = 100.0, step=0.01, description=\"dt: \")\n",
    "w_p0 = FloatSlider(value = 20.0, min = 7.5, max = 30.0, step=0.5, description=\"p0: \")\n",
    "\n",
    "def on_mass_change(change):\n",
    "    global m\n",
    "    m = w_mass.value;\n",
    "    \n",
    "def on_dt_change(change):\n",
    "    global dt\n",
    "    dt = w_dt.value;\n",
    "    \n",
    "def on_p0_change(change):\n",
    "    global p0\n",
    "    p0 = w_p0.value;\n",
    "    \n",
    "def on_init_change(b):\n",
    "    global S,h_x,p0,state_idx    \n",
    "    # two components (one for each electronic state)\n",
    "    psi_x0 = np.zeros((N,2), dtype=np.complex128) \n",
    "    # default IC is a Gaussian wavepacket on 1st state and 0 component on second state but \n",
    "    # user can choose the reverse of this.\n",
    "    psi_x0[:,state_idx] = gauss_x(x, s0, x0, p0) \n",
    "    S = MSOFT(x = x, dt = dt, psi_0=psi_x0, h=h_x, hbar=hbar, m=m)\n",
    "    S.evolution(1)\n",
    "    \n",
    "    w_mass.observe(on_mass_change, names='value')\n",
    "w_dt.observe(on_dt_change, names='value')\n",
    "w_p0.observe(on_p0_change, names='value')\n",
    "w_init = Button(description=\"Update parameters\");\n",
    "w_init.on_click(on_init_change)\n",
    "\n",
    "\n",
    "label2 = Label(\"(Press update parameters to modify the simulation parameters.)\")\n",
    "\n",
    "para_accordion = Accordion(children=[VBox([HBox([w_mass, w_dt, w_p0]), \n",
    "                                           HBox([w_init, label2])])], \n",
    "                           selected_index = None)\n",
    "para_accordion.set_title(0, \"Set simulation parameters\")\n",
    "display(para_accordion)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30286f7c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "The animation shall consist of the stationary plots of the potential in the adiabatic and diabatic picture \n",
    "along with the squared moduli of the components of the wavepackets on each of the potential energy surfaces.\n",
    "\"\"\"\n",
    "\n",
    "# set double harmonic potential as default\n",
    "h_x = double_harmonic(x,A,B,b,C,D) \n",
    "\n",
    "S = MSOFT(x = x, dt = dt, psi_0=psi_x0, h=h_x, hbar=hbar, m=m)\n",
    "S.evolution(1)\n",
    "\n",
    "######################################################################\n",
    "# Set up plot\n",
    "\n",
    "fig = plt.figure(figsize=(10, 10))\n",
    "fig.canvas.header_visible = False\n",
    "\n",
    "# plotting limits\n",
    "xlim = (0, 50)\n",
    "tlim = (0,S.tot_steps)\n",
    "\n",
    "ymin = 0.\n",
    "if pot_select.index == 0:\n",
    "    ymax = 1.0 \n",
    "    ylims=(ymin - 0.2 * (ymax - ymin),\n",
    "                            ymax + 0.2 * (ymax - ymin))\n",
    "else:\n",
    "    xlim=(20,50)\n",
    "    ylims=(-0.2,0.2)\n",
    "    \n",
    "ax1 = fig.add_subplot(221, xlim=xlim,\n",
    "                      ylim=ylims)\n",
    "psi1_x_line, = ax1.plot([], [], c='r', label=r'$|c_1(x)|$', linewidth=1.2)\n",
    "psi2_x_line, = ax1.plot([], [], c='magenta', label=r'$|c_2(x)|$', linewidth=1.2)\n",
    "\n",
    "V1_ad_x_line, = ax1.plot([], [], c='purple', label=r'$V_{ad}(x)$')\n",
    "V2_ad_x_line, = ax1.plot([], [], c='cyan')#, label=r'$V_{ad}(x)$')\n",
    "V1_di_x_line, = ax1.plot([], [], c='green', label=r'$V_{di}(x)$')\n",
    "V2_di_x_line, = ax1.plot([], [], c='green')#, label=r'$V_{di}(x)$')\n",
    "\n",
    "psi1_x_line.set_visible(True)\n",
    "psi2_x_line.set_visible(True)\n",
    "\n",
    "title = ax1.set_title(\"\")\n",
    "ax1.legend(prop=dict(size=8), ncol=4, loc=1)\n",
    "ax1.set_xlabel('$x$')\n",
    "\n",
    "ax2 = fig.add_subplot(223, xlim=(0,S.tot_steps*S.dt), ylim=(-2,2)) #axis for state populations\n",
    "P1_line, = ax2.plot([], [], c='r', label=r'Population (state 1)', linewidth=1.8)\n",
    "P2_line, = ax2.plot([], [], c='b', label=r'Population (state 2)', linewidth=1.8)\n",
    "\n",
    "ax3 = fig.add_subplot(222, xlim=(0,S.tot_steps*S.dt),ylim=(-0.05,0.1)) # axis for energy plots\n",
    "ax4 = fig.add_subplot(224, xlim=(0,S.tot_steps*S.dt), ylim=(-2,2)) # axis for norm plot\n",
    "\n",
    "ax3.legend(prop=dict(size=12))\n",
    "ax3.set_xlabel('$t$')\n",
    "ax3.set_ylabel(r'Energy')\n",
    "ke_line, = ax3.plot([], [], c='r', label=r'Kinetic energy', linewidth=2.0)\n",
    "pe_line, = ax3.plot([], [], c='b', label=r'Potential energy', linewidth=2.0)\n",
    "tot_e_line, = ax3.plot([], [], c='black', label=r'Total energy', linewidth=2.0)\n",
    "\n",
    "ax3.legend(prop=dict(size=12))\n",
    "ax4.set_xlabel('$t$')\n",
    "ax4.set_ylabel(r'Norm')\n",
    "norm_line, = ax4.plot([], [], c='r', label=r'Norm', linewidth=2.0)\n",
    "ax4.legend(prop=dict(size=12))\n",
    "ax2.legend(prop=dict(size=12))\n",
    "ax2.set_xlabel('$t$')\n",
    "ax2.set_ylabel(r'$Population$')\n",
    "\n",
    "V1_ad_x_line.set_data(S.x, S.E_eigvals[:,0])\n",
    "V2_ad_x_line.set_data(S.x, S.E_eigvals[:,1])\n",
    "V1_di_x_line.set_data(S.x, S.h[:,0,0])\n",
    "V2_di_x_line.set_data(S.x, S.h[:,1,1])\n",
    "\n",
    "psi1_x_line.set_data(S.x, 1. * abs(S.psi_x[:,0])+S.E_eigvals[:,0])\n",
    "psi2_x_line.set_data(S.x, 1. * abs(S.psi_x[:,1])+S.E_eigvals[:,1])\n",
    "\n",
    "current_step = int(S.t/S.dt)\n",
    "P1_mask = np.isfinite(S.P1) # masks for plotting intermittant data\n",
    "P2_mask = np.isfinite(S.P2)\n",
    "\n",
    "P1_line.set_data(S.tarr[P1_mask], S.P1[P1_mask])\n",
    "P2_line.set_data(S.tarr[P2_mask], -1.*S.P2[P2_mask])\n",
    "\n",
    "norm = np.sqrt(S.P1+S.P2)\n",
    "norm_mask = np.isfinite(norm)\n",
    "norm_line.set_data(S.tarr[norm_mask], norm[norm_mask])\n",
    "ke_mask = np.isfinite(S.ekint)\n",
    "ke_line.set_data(S.tarr[ke_mask], S.ekint[ke_mask])\n",
    "\n",
    "pe_mask = np.isfinite(S.epot)\n",
    "pe_line.set_data(S.tarr[pe_mask], S.epot[pe_mask])\n",
    "tot_e_line.set_data(S.tarr[pe_mask], S.epot[pe_mask]+S.ekint[ke_mask])\n",
    "\n",
    "# relabelling population graph axis to reflect positive populations\n",
    "@ticker.FuncFormatter\n",
    "def major_formatter(x, pos):\n",
    "    label = str(-x) if x < 0 else str(x)\n",
    "    return label\n",
    "ax2.yaxis.set_major_formatter(major_formatter)\n",
    "\n",
    "title.set_text(\"t = %-5i\" % S.t)\n",
    "plt.show()\n",
    "\n",
    "######################################################################\n",
    "# Functions to Animate the plot\n",
    "\n",
    "pause = True\n",
    "\n",
    "def init():\n",
    "    psi1_x_line.set_data([], [])\n",
    "    psi2_x_line.set_data([], [])\n",
    "    V1_ad_x_line.set_data([],[])\n",
    "    V2_ad_x_line.set_data([],[])\n",
    "    V1_di_x_line.set_data([],[])\n",
    "    V2_di_x_line.set_data([],[])\n",
    "    \n",
    "    P1_line.set_data([], [])\n",
    "    P2_line.set_data([], [])\n",
    "    \n",
    "    ke_line.set_data([], [])\n",
    "    pe_line.set_data([], [])\n",
    "\n",
    "    norm_line.set_data([], [])   \n",
    "    title.set_text(\"\")\n",
    "    \n",
    "    return (psi1_x_line, psi2_x_line, V1_ad_x_line, P1_line, title)\n",
    "\n",
    "def animate(i):\n",
    "    \n",
    "    global S,pause\n",
    "    \n",
    "    current_time = datetime.now()\n",
    "    delta_time=(current_time-init_time).total_seconds()\n",
    "    animation_dur = 50\n",
    "    if not pause and (delta_time<60.0):\n",
    "\n",
    "    # in the case of the double harmonic potential, if the animation has proceeded\n",
    "        # for longer than the duration of interest then restart it.\n",
    "        if(((S.t/S.dt)+animation_dur) >= S.tot_steps):\n",
    "            on_init_change(None)\n",
    "\n",
    "        if(pot_select.index==1):\n",
    "        # if considering the Tully potential, stop sim after wavepacket reaches end of region\n",
    "        # of interest\n",
    "            if(S.psi_x[-1][0] > 0.0005 or S.psi_x[-1][1] > 0.0005):\n",
    "                on_init_change(None)\n",
    "        \n",
    "        S.evolution(animation_dur)\n",
    "        \n",
    "        V1_ad_x_line.set_data(S.x, S.E_eigvals[:,0])\n",
    "        V2_ad_x_line.set_data(S.x, S.E_eigvals[:,1])\n",
    "        V1_di_x_line.set_data(S.x, S.h[:,0,0])\n",
    "        V2_di_x_line.set_data(S.x, S.h[:,1,1])\n",
    "        \n",
    "        psi1_x_line.set_data(S.x, abs(S.psi_x[:,0])+S.E_eigvals[:,0])\n",
    "        psi2_x_line.set_data(S.x, abs(S.psi_x[:,1])+S.E_eigvals[:,1])\n",
    "        \n",
    "        current_step = int(S.t/S.dt)\n",
    "        P1_mask = np.isfinite(S.P1)\n",
    "        P2_mask = np.isfinite(S.P2)\n",
    "        \n",
    "        P1_line.set_data(S.tarr[P1_mask], S.P1[P1_mask])\n",
    "        P2_line.set_data(S.tarr[P2_mask], -1.*S.P2[P2_mask])\n",
    "     \n",
    "        norm = np.sqrt(S.P1+S.P2)\n",
    "        norm_mask = np.isfinite(norm)\n",
    "        norm_line.set_data(S.tarr[norm_mask], norm[norm_mask])\n",
    "        \n",
    "        ke_mask = np.isfinite(S.ekint)\n",
    "        ke_line.set_data(S.tarr[ke_mask], S.ekint[ke_mask])\n",
    "        \n",
    "        pe_mask = np.isfinite(S.epot)\n",
    "        pe_line.set_data(S.tarr[pe_mask], S.epot[pe_mask])\n",
    "        tot_e_line.set_data(S.tarr[pe_mask], S.epot[pe_mask]+S.ekint[ke_mask])\n",
    "\n",
    "        norm_text.value=\"Norm: \"+str(norm[norm_mask][-1])\n",
    "        title.set_text(\"t = %-5i\" % S.t)\n",
    "        return (psi1_x_line, V1_ad_x_line, P1_line, title)\n",
    "    else:\n",
    "        pause=True\n",
    "        button_pause.description = \"Play\"\n",
    "        anim.event_source.stop()\n",
    "       \n",
    "def onClick(event):\n",
    "    global pause\n",
    "    global init_time\n",
    "    init_time=datetime.now()\n",
    "    \n",
    "    pause ^= True\n",
    "    if button_pause.description == \"Pause\":\n",
    "        button_pause.description = \"Play\"\n",
    "        anim.event_source.stop()\n",
    "    else:\n",
    "        button_pause.description = \"Pause\"\n",
    "        anim.event_source.start()\n",
    "\n",
    "        \n",
    "init_time=datetime.now()\n",
    "\n",
    "anim = animation.FuncAnimation(fig, animate, init_func=init, frames=frames, interval=1, blit=True)\n",
    "button_pause = Button(description=\"Play\");\n",
    "button_pause.on_click(onClick)\n",
    "display(button_pause);\n",
    "norm_text = widgets.Text()\n",
    "display(norm_text)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af0064fb-0f8f-4545-81ab-c19c03039870",
   "metadata": {},
   "source": [
    "<hr style=\"height:1px;border:none;color:#cccccc;background-color:#cccccc;\" />\n",
    "\n",
    "# Legend (How to use the interactive visualization)\n",
    "\n",
    "## Visualization\n",
    "\n",
    "The visualization consists of 4 subplots (clockwise from the top):\n",
    "1. The main animation pane. This shows the evolution of the two components of the nuclear wavefunction (red and magenta curves) on their respective electronic potential energy surfaces (PESs). The potential energy surfaces are plotted in both the adiabatic (purple) and diabatic bases (green).\n",
    "\n",
    "2. Plot of the kinetic, potential and total energies vs simulation timestep.\n",
    "\n",
    "3. Evolution of the populations on the two different electronic PESs.\n",
    "\n",
    "4. Evolution of the total norm of the nuclear wavepacket.\n",
    "\n",
    "Below the plots there is also a display field indicating the current value of the norm of the total wavefunction.\n",
    "\n",
    "## Interactivity controls\n",
    "\n",
    "* The type of potential may be selected by first expanding the \"Select potential and set parameters\" drawer. The two options are: a double harmonic potential and a Tully model potential.\n",
    "* One can choose on which electronic state the Gaussian nuclear wavepacket is initialized by using the \"Choose initial electronic state\" dropdown.\n",
    "* The parameters of each potential may be varied using the various sliders.\n",
    "* To update the potential, the \"Update potential\" button should be pressed, followed by the \"Update parameters\" button in the \"Set simulation parameters\" drawer.\n",
    "\n",
    "* The various simulation parameters (nuclear mass, timestep (dt), and initial momentum of the nuclear wavepacket (p0)) may be set using the sliders in the \"Set simulation parameters\" drawer.\n",
    "\n"
   ]
  }
 ],
 "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}