{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# TD S-eqn using SOFT\n", "\n", "**Authors:** Dou Du, Taylor James Baird and Sara Bonella\n", "\n", " Go back to index\n", "\n", "**Source code:** https://github.com/osscar-org/quantum-mechanics/blob/master/notebook/quantum-mechanics/soft.ipynb\n", "\n", "This notebook presents the split operator Fourier transform (SOFT) numerical method for solving the one-dimensional time-dependent Schrödinger equation.\n", "
" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## **Goals**\n", "* Understand the various steps in the SOFT 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 SOFT method to investigate various quantum mechanical features which appear in the dynamics of a quantum system.\n", "* Observe the manifestation of Heisenberg's uncertainty principle in the time evolution of a nuclear wavepacket." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## **Background theory**\n", "\n", "[**More in the background theory.**](./theory/theory_soft.ipynb)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## **Tasks and exercises**\n", "\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", "
    \n", " Solution\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 exactly 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", "
    \n", "
    2. \n", "
  2. What dictates the maximum size of the timestep that we can use for the SOFT algorithm (consider the main assumptions/approximations that are made when formulating the propagation scheme for SOFT).\n", "
    \n", " Solution\n", " Recall that the central approximation used in the derivation of the SOFT propagation scheme is in the truncation of the Trotter product formula: $e^{A+B} = \\lim\\limits_{P \\to \\infty} (e^{\\frac{A}{P}} e^{\\frac{B}{P}})^{P} \\approx (e^{\\frac{A}{N}} e^{\\frac{B}{N}})^{N}$ for $N$ sufficiently large. This approximation becomes more and more accurate the larger we make the value of $N$. 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", "
    \n", "
    3. \n", "
  3. Why is the use of the SOFT 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", "
    \n", " Solution\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 SOFT algorithm unsuitable for large systems.\n", "
    \n", "
    4. \n", "
  4. Observe the evolution of the wavefunction in both position and momentum space. Do you notice the manifestation of Heisenberg's uncertainty principle? Can you explain it?\n", "
    \n", " Solution\n", " The top figure shows the relation between position $x$ and the modulus of the wavefunction $|\\psi|$.\n", " On the other hand, the bottom figure shows the relation between momentum $p$ and the modulus of the wavefunction, $|\\psi|$. Using the Harmonic potential for example, one\n", " can clearly see that when the position of the particle is known with greater certainty (a sharp peak of the $|\\psi(x)|$), the $|\\psi(p)|$ will have\n", " a boarder shape. This remains true for all time.\n", "
    \n", "
    5. " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "metadata": {}, "source": [ "
    \n", "\n", "## Interactive visualization\n", "(be patient, it might take a few seconds to load)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%matplotlib widget\n", "\n", "import numpy as np\n", "from math import pi\n", "from scipy.fftpack import fft, ifft\n", "from ipywidgets import Button, FloatSlider, HBox, VBox, IntProgress\n", "from ipywidgets import Dropdown, Layout, Accordion, Label\n", "import ipywidgets as widgets\n", "import matplotlib.pyplot as plt\n", "from matplotlib import animation\n", "from datetime import datetime\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "class SOFT(object):\n", " \"\"\"\n", " This is the Split Operator Fourier Transform (SOFT)\n", " x: the grid space \n", " psi_0: inital wavefunction, included real and image parts\n", " V: the potential\n", " dt: the time interval\n", " hbar: the Plank constant (default value 1.0 as atomic unit)\n", " m: the mass (default value 1.0 as atomic unit)\n", " \"\"\"\n", " \n", " def __init__(self, x, psi_0, V, 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];\n", " dk = (2*pi)/(self.N*dx);\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", " self.tarr=self.dt*np.arange(self.tot_steps)\n", " self.m = m;\n", " \n", " self.k_x = -0.5*self.N*self.dk + dk*np.arange(self.N);\n", " self.psi_x = np.zeros((self.N), dtype=np.complex128)\n", " self.psi_k = np.zeros((self.N), dtype=np.complex128)\n", " self.psi_mod_x = np.zeros((self.N), dtype=np.complex128)\n", " self.psi_mod_k = np.zeros((self.N), dtype=np.complex128)\n", " \n", " self.ekint = np.array([np.nan]*self.tot_steps)\n", " self.epot = np.array([np.nan]*self.tot_steps)\n", " self.norm = np.array([np.nan]*self.tot_steps)\n", " \n", " self.psi_x = psi_0;\n", " self._periodic_bc();\n", " \n", " if callable(V):\n", " self.V = V(x)\n", " else:\n", " self.V = V\n", " \n", " self.psi_mod_x[:] = self.dx/np.sqrt(2*pi)*self.psi_x[:]*np.exp(-1.0j*self.k_x[0]*self.x);\n", " \n", " def _periodic_bc(self):\n", " self.psi_x[-1] = self.psi_x[0]\n", " \n", " def _half_pot_prop(self, ft=True):\n", " if ft == True:\n", " self.psi_mod_x[:] = ifft(self.psi_mod_k[:])\n", " self.psi_mod_x[:] = self.psi_mod_x[:]*np.exp(-1.0j*(self.dt/2.0)*self.V)\n", "\n", " def _full_kinetic_prop(self):\n", " self.psi_mod_k[:] = fft(self.psi_mod_x[:])\n", " self.psi_mod_k[:] = self.psi_mod_k[:]*np.exp(-1.0j*self.k_x**2*self.dt/(2.0*self.m))\n", " \n", " def _compute_psi(self):\n", " self.psi_x[:] = (np.sqrt(2*pi)/self.dx)*self.psi_mod_x*np.exp(1.0j*self.k_x[0]*self.x)\n", " self.psi_k[:] = self.psi_mod_k*np.exp(-1.0j*self.x[0]*self.dk*np.arange(self.N))\n", " \n", " def get_kinetic_energy(self):\n", " ke = 0.5*(1./self.m)*sum(self.k_x**2*np.absolute(self.psi_k)**2)*self.dk\n", " self.ekint[int(self.t/self.dt)] = ke\n", "\n", " \n", " def get_potential_energy(self):\n", " epot = sum(np.conj(self.psi_x)*self.V*self.psi_x)*self.dx\n", " self.epot[int(self.t/self.dt)] = epot\n", " \n", " def get_norm(self):\n", " self.norm[int(self.t/self.dt)] = sum(np.conj(self.psi_x)*self.psi_x)*self.dx\n", " \n", " \n", " def evolution(self, Nsteps=1):\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", " self._compute_psi()\n", " self.get_kinetic_energy()\n", " self.get_norm()\n", " self.get_potential_energy()\n", " self.psi_mod_x /= np.sqrt((abs(self.psi_mod_x)**2).sum()*2*pi/self.dx);\n", " self.t += self.dt*Nsteps\n", " " ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "######################################################################\n", "# Helper functions for gaussian wave-packets\n", "def gauss_x(x, a, x0, k0):\n", " \"\"\"\n", " a gaussian wave packet of width a, centered at x0, with momentum k0\n", " \"\"\"\n", " return ((a * np.sqrt(np.pi)) ** (-0.5)\n", " * np.exp(-0.5 * ((x - x0) * 1. / a) ** 2 + 1j * x * k0))\n", "\n", "\n", "def gauss_k(k, a, x0, k0):\n", " \"\"\"\n", " analytical fourier transform of gauss_x(x), above\n", " \"\"\"\n", " return ((a / np.sqrt(np.pi)) ** 0.5\n", " * np.exp(-0.5 * (a * (k - k0)) ** 2 - 1j * (k - k0) * x0))\n", "\n", "\n", "######################################################################\n", "# Utility functions for running the animation\n", "def theta(x):\n", " \"\"\"\n", " theta function :\n", " returns 0 if x<=0, and 1 if x>0\n", " \"\"\"\n", " x = np.asarray(x)\n", " y = np.zeros(x.shape)\n", " y[x > 0] = 1.0\n", " return y\n", "\n", "\n", "def square_barrier(x, width, height):\n", " return height * (theta(x + 0.5*width) - theta(x - 0.5*width))\n", "\n", "def parabola(x, a=1.0, x0 = 0.0):\n", " return a/1000.0*(x-x0)**2\n", "\n", "def morse(x, D = 1.0, b = 0.03, x0 = -60.0):\n", " return D*(np.exp(-2.0*b*(x-x0))-2*np.exp(-b*(x-x0))) + D\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "######################################################################\n", "# Create the animation\n", "\n", "# specify time steps and duration\n", "dt = 0.01\n", "N_steps = 50\n", "t_max = 120\n", "frames = int(t_max / float(N_steps * abs(dt)))\n", "\n", "# specify constants\n", "hbar = 1.0 # planck's constant\n", "m = 1.9 # particle mass\n", "\n", "# specify range in x coordinate\n", "N = 2 ** 11\n", "dx = 0.1\n", "x = dx * (np.arange(N) - 0.5 * N)\n", "\n", "# specify potential\n", "V0 = 1.5\n", "L = hbar / np.sqrt(2 * m * V0)\n", "a = 30 * L\n", "x0 = -60 * L\n", "\n", "# specify initial momentum and quantities derived from it\n", "p0 = np.sqrt(2 * m * 0.2 * V0)\n", "dp2 = p0 * p0 * 1. / 80\n", "d = hbar / np.sqrt(2 * dp2)\n", "\n", "k0 = p0 / hbar\n", "v0 = p0 / m\n", "psi_x0 = gauss_x(x, d, x0, k0)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "style = {'description_width': 'initial'}\n", "\n", "pot_select = Dropdown(\n", " options=['1. Box potential', '2. Morse potential', '3. Harmonic potential'],\n", " index = 0,\n", " description='Potential type:',\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 the depth and wide of the box potential (pot1)\n", "swide = FloatSlider(value = a, min = 0.0, max = 2*a, description = 'Width: ');\n", "sheight = FloatSlider(value = V0, min = 0.0, max = 2*V0, description = 'Height: ');\n", "swide.layout = layout_visible\n", "sheight.layout = layout_visible\n", "\n", "#Set the parameter for the Morse potential (pot2)\n", "pot2_D = FloatSlider(value = 1.0, min = 1.0, max = 5.0, description = 'D: ');\n", "pot2_b = FloatSlider(value = 0.03, min = 0.01, max = 0.10, step=0.01, description = 'b: ');\n", "pot2_x0 = FloatSlider(value = -60.0, min = -100.0, max = 100.0, description = 'x0: ');\n", "\n", "pot2_D.layout = layout_hidden\n", "pot2_b.layout = layout_hidden\n", "pot2_x0.layout = layout_hidden\n", "\n", "#Set the parameter for the Harmonic potential (pot3)\n", "pot3_a = FloatSlider(value = 1.0, min = 0.2, max = 2.0, description = 'a: ');\n", "pot3_x0 = FloatSlider(value = 0.0, min = -100.0, max = 100.0, description = 'x0: ');\n", "pot3_a.layout = layout_hidden\n", "pot3_x0.layout = layout_hidden\n", "\n", "#Show the potential image\n", "file1 = open(\"images/pot1.png\", \"rb\")\n", "image1 = file1.read()\n", "file2 = open(\"images/pot2.png\", \"rb\")\n", "image2 = file2.read()\n", "file3 = open(\"images/pot3.png\", \"rb\")\n", "image3 = file3.read()\n", "\n", "pot_img = widgets.Image(\n", " value=image1,\n", " format='png',\n", " width=300,\n", " height=400,\n", ")\n", "\n", "pot_img.layout = layout_visible\n", "\n", "def pot_change(change):\n", " global V_x\n", " if pot_select.index == 0:\n", " swide.layout = layout_visible\n", " sheight.layout = layout_visible\n", " pot2_D.layout = layout_hidden\n", " pot2_b.layout = layout_hidden\n", " pot2_x0.layout = layout_hidden\n", " pot3_a.layout = layout_hidden\n", " pot3_x0.layout = layout_hidden\n", " pot_img.layout = layout_visible\n", " pot_img.value = image1\n", " elif pot_select.index == 1:\n", " swide.layout = layout_hidden\n", " sheight.layout = layout_hidden\n", " pot2_D.layout = layout_visible\n", " pot2_b.layout = layout_visible\n", " pot2_x0.layout = layout_visible\n", " pot3_a.layout = layout_hidden\n", " pot3_x0.layout = layout_hidden\n", " pot_img.layout = layout_visible\n", " pot_img.value = image2 \n", " elif pot_select.index == 2:\n", " swide.layout = layout_hidden\n", " sheight.layout = layout_hidden\n", " pot2_D.layout = layout_hidden\n", " pot2_b.layout = layout_hidden\n", " pot2_x0.layout = layout_hidden\n", " pot3_a.layout = layout_visible\n", " pot3_x0.layout = layout_visible\n", " pot_img.layout = layout_visible\n", " pot_img.value = image3 \n", " \n", "pot_select.observe(pot_change, names='value', type='change');\n", "\n", "def on_pot_update(event):\n", " global V_x, S\n", " if pot_select.index == 0:\n", " V_x = square_barrier(x, swide.value, sheight.value)\n", " V_x[x < -98] = 100.\n", " V_x[x > 98] = 100.\n", " elif pot_select.index == 1:\n", " V_x = morse(x, pot2_D.value, pot2_b.value, pot2_x0.value)\n", " elif pot_select.index == 2:\n", " V_x = parabola(x, pot3_a.value, pot3_x0.value)\n", " V_x_line.set_data(x, V_x)\n", " on_init_change('para')\n", "\n", "pot_update = Button(description=\"Update potential\")\n", "pot_update.on_click(on_pot_update)\n", "\n", "label1 = Label(\"(Pause the simulation to change the potential)\")\n", "\n", "# The accordion for the potential\n", "pot_accordion = Accordion(children=[VBox([pot_select, pot_img,\n", " HBox([swide, sheight]),\n", " HBox([pot2_D, pot2_b]), pot2_x0,\n", " HBox([pot3_a, pot3_x0]),\n", " HBox([pot_update, label1])])], selected_index = None)\n", "\n", "pot_accordion.set_title(0, \"Select potential and set parameters\")\n", "display(pot_accordion)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "w_mass = FloatSlider(value = 1.0, min = 0.2, max=5.0, description=\"mass: \") \n", "w_dt = FloatSlider(value = 0.01, min = 0.01, max = 1.0, step=0.01, description=\"dt: \")\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", "w_mass.observe(on_mass_change, names='value')\n", "w_dt.observe(on_dt_change, names='value')" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "V_x = square_barrier(x, a, V0)\n", "V_x[x < -98] = 100\n", "V_x[x > 98] = 100\n", "\n", "S = SOFT(x = x, dt = dt, psi_0=psi_x0, V=V_x, hbar=hbar, m=m)\n", "S.evolution(1)\n", "\n", "def on_init_change(b):\n", " global S\n", " psi_x0 = gauss_x(x, d, x0, k0)\n", " S = SOFT(x = x, dt = dt, psi_0=psi_x0, V=V_x, hbar=hbar, m=m)\n", " S.evolution(1)\n", " setup_plot()\n", " \n", "def setup_plot():\n", " psi_x_line.set_data(S.x, 4 * abs(S.psi_x))\n", " psi_x_real.set_data(S.x, 4 * S.psi_x.real)\n", " psi_x_imag.set_data(S.x, 4 * S.psi_x.imag)\n", " V_x_line.set_data(S.x, S.V)\n", " psi_k_line.set_data(S.k_x, abs(S.psi_k))\n", " title.set_text(\"t = %-5i\" % S.t)\n", " \n", "w_init = Button(description=\"Update parameters\");\n", "w_init.on_click(on_init_change)\n", "\n", "label2 = Label(\"(Pause the simulation to change the parameters)\")\n", "\n", "para_accordion = Accordion(children=[VBox([HBox([w_mass, w_dt]), \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, "metadata": {}, "outputs": [], "source": [ "######################################################################\n", "# Set up plot\n", "\n", "fig = plt.figure(figsize=(10, 10))\n", "fig.canvas.header_visible = False\n", "\n", "# plotting limits\n", "xlim = (-100, 100)\n", "klim = (-5, 5)\n", "\n", "# top axes show the x-space data\n", "ymin = 0\n", "ymax = V0\n", "ax1 = fig.add_subplot(221, xlim=xlim,\n", " ylim=(ymin - 0.2 * (ymax - ymin),\n", " ymax + 0.2 * (ymax - ymin)))\n", "psi_x_line, = ax1.plot([], [], c='r', label=r'$|\\psi(x)|$', linewidth=1.2)\n", "psi_x_real, = ax1.plot([], [], c='b', label=r'$Re(\\psi)$', linewidth=0.8)\n", "psi_x_imag, = ax1.plot([], [], c='orange', label=r'$Im(\\psi)$', linewidth=0.8)\n", "V_x_line, = ax1.plot([], [], c='k', label=r'$V(x)$')\n", "\n", "psi_x_line.set_visible(True)\n", "\n", "\n", "title = ax1.set_title(\"\")\n", "ax1.legend(prop=dict(size=8), ncol=4, loc=1)\n", "psi_x_real.set_visible(False)\n", "psi_x_imag.set_visible(False)\n", "ax1.set_xlabel('$x$')\n", "ax1.set_ylabel(r'$|\\psi(x)|$')\n", "\n", "# bottom axes show the k-space data\n", "ymin = abs(S.psi_k).min()\n", "ymax = abs(S.psi_k).max()\n", "ax2 = fig.add_subplot(222, xlim=klim,\n", " ylim=(ymin - 0.2 * (ymax - ymin),\n", " ymax + 0.2 * (ymax - ymin)))\n", "psi_k_line, = ax2.plot([], [], c='r', label=r'$|\\psi(k)|$', linewidth=1.2)\n", "\n", "p0_line1 = ax2.axvline(-p0 / hbar, c='k', ls=':', label=r'$\\pm p_0$')\n", "p0_line2 = ax2.axvline(p0 / hbar, c='k', ls=':')\n", "\n", "\n", "ax2.legend(prop=dict(size=8), ncol=4, loc=1)\n", "ax2.set_xlabel('$p$')\n", "ax2.set_ylabel(r'$|\\psi(k)|$')\n", "\n", "# axis for energy plots\n", "ax3 = fig.add_subplot(223, xlim=(0,S.tot_steps*S.dt),ylim=(-2.,2.))\n", "# axis for norm plot\n", "ax4 = fig.add_subplot(224, xlim=(0,S.tot_steps*S.dt), ylim=(-2,2)) \n", "\n", "\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", "ax3.set_xlabel('$t$')\n", "ax3.set_ylabel(r'Energy')\n", "ax3.legend(prop=dict(size=6), ncol=4, loc=1)\n", "\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", "\n", "norm = S.norm\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", "V_x_line.set_data(S.x, S.V)\n", "psi_x_line.set_data(S.x, 4 * abs(S.psi_x))\n", "V_x_line.set_data(S.x, S.V)\n", "\n", "\n", "psi_k_line.set_data(S.k_x, abs(S.psi_k))\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", " psi_x_line.set_data([], [])\n", " psi_x_real.set_data([], [])\n", " psi_x_imag.set_data([], [])\n", " V_x_line.set_data([], [])\n", " psi_k_line.set_data([], [])\n", " ke_line.set_data([], [])\n", " pe_line.set_data([], [])\n", " tot_e_line.set_data([], [])\n", " norm_line.set_data([], [])\n", " title.set_text(\"\")\n", "\n", " return (psi_x_line, V_x_line, psi_k_line, title)\n", "\n", "def animate(i):\n", " global S,pause\n", " animation_dur = 50\n", " current_time = datetime.now()\n", " delta_time=(current_time-init_time).total_seconds()\n", " if not pause and (delta_time<30.0):\n", " if(((S.t/S.dt)+animation_dur) >= S.tot_steps):\n", " on_init_change(None)\n", " \n", " S.evolution(animation_dur)\n", " psi_x_line.set_data(S.x, 4 * abs(S.psi_x))\n", " psi_x_real.set_data(S.x, 4 * S.psi_x.real)\n", " psi_x_imag.set_data(S.x, 4 * S.psi_x.imag)\n", " V_x_line.set_data(S.x, S.V)\n", "\n", " psi_k_line.set_data(S.k_x, abs(S.psi_k))\n", " \n", " norm = S.norm\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", " \n", " title.set_text(\"t = %-5i\" % S.t)\n", "\n", " return (psi_x_line, V_x_line, psi_k_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", " # pot_update.disabled = False\n", " # w_init.disabled = False\n", " anim.event_source.stop()\n", " else:\n", " button_pause.description = \"Pause\"\n", " # pot_update.disabled = True\n", " # w_init.disabled = True\n", " anim.event_source.start()\n", "\n", "\n", "# call the animator.\n", "# blit=True means only re-draw the parts that have changed.\n", "init_time=datetime.now()\n", "\n", "anim = animation.FuncAnimation(fig, animate, init_func=init, frames=frames, interval=0, blit=True)\n", "\n", "psi1 = widgets.Checkbox(value=True, description=r'Abs', disabled=False, layout=Layout(width='25%'))\n", "psi2 = widgets.Checkbox(value=False, description=r'Re', disabled=False, layout=Layout(width='25%'))\n", "psi3 = widgets.Checkbox(value=False, description=r'Im', disabled=False, layout=Layout(width='25%'))\n", "# wfn_comp_box = VBox([HBox([psi1, psi2, psi3])\n", "# widgets.Text(\n", "# value='Wavefunction component',\n", "# placeholder='Type something',\n", "# description='String:',\n", "# disabled=True\n", "# )\n", "def on_psi1_change(b):\n", " psi_x_line.set_visible(psi1.value)\n", " \n", "def on_psi2_change(b):\n", " psi_x_real.set_visible(psi2.value)\n", " \n", "def on_psi3_change(b):\n", " psi_x_imag.set_visible(psi3.value)\n", "\n", "psi1.observe(on_psi1_change, names='value')\n", "psi2.observe(on_psi2_change, names='value')\n", "psi3.observe(on_psi3_change, names='value')\n", "\n", "button_pause = Button(description=\"Play\");\n", "button_pause.on_click(onClick)\n", "\n", "display(HBox([Label(\"Wavefunction component\"),psi1, psi2, psi3]))\n", "display(button_pause);" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "
    \n", "\n", "# Legend\n", "\n", "(How to use the interactive visualization)\n", "\n", "## Interactive figures\n", "\n", "The interactive figures consist of four subplots. Clockwise from the top left we have:\n", "1. The evolution of the wavefunction in the position representation, $x$. The chosen potential is also shown in this subplot.\n", "2. The evolution of the wavefunction in the momentum representation, $p$. The initial momentum $p_0$ (along with $-p_0$) are indicated with veritcal lines in this subplot. \n", "3. Plot of the kinetic, potential and total energies vs simulation timestep (zoom to better discern the evolution of each contribution).\n", "4. Evolution of the total norm of the nuclear wavepacket.\n", "\n", "\n", "## Controls\n", "\n", "Before running the simulations, users need to select a potential and set various simulation parameters.\n", "We offer three types of potentials, namely the box, Morse and Harmonic potentials.\n", "One can also tune the parameters of these potentials. Moreover, the nuclear mass, $m$, and the size of the time step, $\\Delta t$ may be modified via their corresponding sliders. Afterwards, \n", "one can click the \"Play\" button to run the simulation. One can also pause\n", "the simulation and reselect potential and reset parameters.\n", "\n", "By default, the figures only show the absolute value of the wavefunction,\n", "$|\\psi(x)|$. However, one can also select to show the real Re$(\\psi(x))$ and \n", "image Im$(\\psi(x))$ parts of the wavefunction by ticking the corresponding checkboxes." ] }, { "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.9.7" }, "voila": { "authors": "Dou Du, Sara Bonella and Giovanni Pizzi" } }, "nbformat": 4, "nbformat_minor": 4 }