{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 13 - Hydrogen functions\n",
    "Start with some imports fro Symbolic Python library:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sympy.physics.hydrogen import R_nl\n",
    "from sympy.functions.special.spherical_harmonics import Ynm\n",
    "from sympy import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define some variables, radial, polar, azimuthal, time, and two frequencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(r, theta, phi, t, w1, w2)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "var(\"r theta phi t w1 w2\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Look at a few of the radial equations and the spherical harmonics\n",
    "Notice that instead of Ylm the name is Ynm... the arguments to the function are still the quantum numbers `l` and `m`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 2 e^{- r}$"
      ],
      "text/plain": [
       "2*exp(-r)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "R_nl(1, 0, r, 1)  # the n = 1, l = 0 radial function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{1}{2 \\sqrt{\\pi}}$"
      ],
      "text/plain": [
       "1/(2*sqrt(pi))"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Ynm(0,0,theta,phi).expand(func=True)  # the l = 0, m = 0 spherical harmonic"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Write the equation for the $|nlm\\rangle = |100\\rangle$ state. Use the sympy method `.expand(func=True)` to convert to the actual expression. To create this state, we combine the Radial function and the Ylm function. Make sure to set n, l, and m to the correct values. The fourth argument to `R_nl` is `Z` which we set to 1 since we are talking about a 1-proton nucleus.\n",
    "\n",
    "The combination of R_nl and Ynm should look like the following (replace N, L, and M with the appropriate values):\n",
    "\n",
    "`R_nl(N, L, r, 1)*Ynm(L, M, theta, phi).expand(func=True)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is the |100> state:\n",
    "psi100 = R_nl(1, 0, r, 1)*Ynm(0,0,theta,phi).expand(func=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{e^{- r}}{\\sqrt{\\pi}}$"
      ],
      "text/plain": [
       "exp(-r)/sqrt(pi)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "psi100  # check to see how it looks as an expression"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Integrating over all space\n",
    "Remember spherical coordinate integrals of function $f(r,\\theta,\\phi)$ over all space look like: $$\\int_0^\\infty\\int_0^\\pi\\int_0^{2\\pi}r^2\\sin(\\theta)drd\\theta d\\phi \\,\\,f(r,\\theta,\\phi)$$ so you alwasy need to add a factor of `r**2*sin(theta)` and then integrate `r` from 0 to infinity, `theta` from $0-\\pi$ and `phi` from $0-2\\pi$. As a check, you should integrate the square of the `psi100` wavefunction over all space to see that it equals 1 (i.e. it is normalized)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 1$"
      ],
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "integrate(r**2*sin(theta) * (psi100)**2 ,(r,0,oo),(theta,0,pi),(phi,0,2*pi))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now do the $|210\\rangle$ state:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi210 = R_nl(2, 1, r, 1)*Ynm(1,0,theta,phi).expand(func=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\sqrt{2} r e^{- \\frac{r}{2}} \\cos{\\left(\\theta \\right)}}{8 \\sqrt{\\pi}}$"
      ],
      "text/plain": [
       "sqrt(2)*r*exp(-r/2)*cos(theta)/(8*sqrt(pi))"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "psi210  # check how it looks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note, if you compare these to listed solutions (for example at http://hyperphysics.phy-astr.gsu.edu/hbase/quantum/hydwf.html#c3) you see that there are not any factors of $a_0$. This is because the `R_nl` function is defined in units of $a_0$. $a_0$ is the Bohr Radius: http://en.wikipedia.org/wiki/Bohr_radius"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle - \\frac{r e^{- \\frac{r}{2}} e^{i \\phi} \\sin{\\left(\\theta \\right)}}{8 \\sqrt{\\pi}}$"
      ],
      "text/plain": [
       "-r*exp(-r/2)*exp(I*phi)*sin(theta)/(8*sqrt(pi))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "psi211 = R_nl(2, 1, r, 1)*Ynm(1,1,theta,phi).expand(func=True)\n",
    "psi211"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now calculate $\\langle z \\rangle$:\n",
    "To calculate $\\langle z \\rangle$ we need to convert to spherical coordinates: $z = r\\cos\\theta$. The terms in the following integral are the $r^2\\sin\\theta$ then $z$ (in spherical coords) then the wave function squared."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "expect = integrate(r**2*sin(theta)* (r*cos(theta)) * (psi100*psi100),(r,0,oo),(theta,0,pi),(phi,0,2*pi))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 0$"
      ],
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expect"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "No surprise, the average z position of the electron in the hydrogen atom is 0."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Now for problem 13.21\n",
    "find $\\langle z \\rangle(t)$. Use the same integral, but add a time-dependent piece to each term in the wavefunction, add them together and multiply by the complex conjugate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "psi = 1/sqrt(2)*(psi100*exp(1j*w1*t) + psi210*exp(1j*w2*t))\n",
    "psi_conj = 1/sqrt(2)*(psi100*exp(-1j*w1*t) + psi210*exp(-1j*w2*t))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "outer = (psi*psi_conj).simplify()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{\\left(\\sqrt{2} r e^{r + 1.0 i t w_{1}} \\cos{\\left(\\theta \\right)} + 8 e^{\\frac{r}{2} + i t w_{2}}\\right) \\left(\\sqrt{2} r e^{r + 1.0 i t w_{2}} \\cos{\\left(\\theta \\right)} + 8 e^{\\frac{r}{2} + i t w_{1}}\\right) e^{- 3 r - 1.0 i t w_{1} - 1.0 i t w_{2}}}{128 \\pi}$"
      ],
      "text/plain": [
       "(sqrt(2)*r*exp(r + 1.0*I*t*w1)*cos(theta) + 8*exp(r/2 + I*t*w2))*(sqrt(2)*r*exp(r + 1.0*I*t*w2)*cos(theta) + 8*exp(r/2 + I*t*w1))*exp(-3*r - 1.0*I*t*w1 - 1.0*I*t*w2)/(128*pi)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "expect2 = integrate(r**2 * sin(theta) * (r*cos(theta)) * outer,(r,0,oo),(theta,0,pi),(phi,0,2*pi))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 2 \\pi \\left(- \\frac{\\left(- 27390260577042432 \\sqrt{2} \\pi^{11} e^{3.0 i t w_{1}} e^{5.0 i t w_{2}} - 141371459682435072 \\pi^{11} e^{4.0 i t w_{1}} e^{4.0 i t w_{2}} - 27390260577042432 \\sqrt{2} \\pi^{11} e^{5.0 i t w_{1}} e^{3.0 i t w_{2}}\\right) e^{- 4.0 i t w_{1}} e^{- 4.0 i t w_{2}}}{415989582513831936 \\pi^{12}} + \\frac{\\left(27390260577042432 \\sqrt{2} \\pi^{11} e^{3.0 i t w_{1}} e^{5.0 i t w_{2}} - 141371459682435072 \\pi^{11} e^{4.0 i t w_{1}} e^{4.0 i t w_{2}} + 27390260577042432 \\sqrt{2} \\pi^{11} e^{5.0 i t w_{1}} e^{3.0 i t w_{2}}\\right) e^{- 4.0 i t w_{1}} e^{- 4.0 i t w_{2}}}{415989582513831936 \\pi^{12}}\\right)$"
      ],
      "text/plain": [
       "2*pi*(-(-27390260577042432*sqrt(2)*pi**11*exp(3.0*I*t*w1)*exp(5.0*I*t*w2) - 141371459682435072*pi**11*exp(4.0*I*t*w1)*exp(4.0*I*t*w2) - 27390260577042432*sqrt(2)*pi**11*exp(5.0*I*t*w1)*exp(3.0*I*t*w2))*exp(-4.0*I*t*w1)*exp(-4.0*I*t*w2)/(415989582513831936*pi**12) + (27390260577042432*sqrt(2)*pi**11*exp(3.0*I*t*w1)*exp(5.0*I*t*w2) - 141371459682435072*pi**11*exp(4.0*I*t*w1)*exp(4.0*I*t*w2) + 27390260577042432*sqrt(2)*pi**11*exp(5.0*I*t*w1)*exp(3.0*I*t*w2))*exp(-4.0*I*t*w1)*exp(-4.0*I*t*w2)/(415989582513831936*pi**12))"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expect2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\frac{64 \\sqrt{2} \\left(e^{2.0 i t w_{1}} + e^{2.0 i t w_{2}}\\right) e^{- 1.0 i t \\left(w_{1} + w_{2}\\right)}}{243}$"
      ],
      "text/plain": [
       "64*sqrt(2)*(exp(2.0*I*t*w1) + exp(2.0*I*t*w2))*exp(-1.0*I*t*(w1 + w2))/243"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expect2.simplify()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to interpret this result. First you should show that this expression is simply a constant amplitude factor times $\\cos((w2-w1)t)$, in other words $\\langle z \\rangle$ oscillates at frequency `w2-w1`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Your assignment:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Explore other combinations of states and draw conclusions about the z behavior from the results. You may not be able to get these expressions to simplify, but the important thing is to look for the time dependence and simplify that part."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Does $\\langle z \\rangle$ oscillate for any combination of two Hydrogen states $|nlm\\rangle$?\n",
    "- Are there restrictions on what n values give oscillating $\\langle z \\rangle$ expressions? (hint, to keep it simple, always let one state be the n=1 state)\n",
    "- How does $\\langle z \\rangle$ change with different l and m values are used in the state?\n",
    "\n",
    "Hints for interpreting your results:\n",
    "- What are the relavant frequencies in your expression for $\\langle z \\rangle$ and why?\n",
    "- Simplify one of your $\\langle z \\rangle$ expressions and write the time dependence in terms of the frequencies w2 and w1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}