{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hydrogen-like Atomic Orbitals"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Theory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Wavefunctions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The goal of this notebook is to visualise the structure of the atomic orbitals for hydrogenoid atoms, in the simplified non-relativistic case (not accounting for any correction in the wavefunctions and in the energies). The expression of these orbitals can be found by solving the Schrödinger stationary equation with a radial potential $V(r)=-\\frac{1}{4\\pi\\varepsilon_0}\\frac{Ze^2}{r}$ emanating from Coulomb's law. \n",
    "$$\n",
    "\\left[-\\frac{\\hbar^2}{2 \\mu}\\nabla^2 + V(r)\\right] \\psi(r, \\theta, \\varphi)= E \\psi(r, \\theta, \\varphi),\n",
    "$$\n",
    "where $\\mu=\\frac{m_e m_N}{m_e + m_N}$ the reduced mass, $m_e$ and $m_N$ the mass of the electron and of the nucleus respectively (generally $\\mu = m_e$). From the expression of the Laplacian $\\nabla^2$ in the spherical coordinates \n",
    "$$\n",
    "\\nabla^2\\equiv\\frac{1}{r^2}\\frac{\\partial}{\\partial r}\\left(r^2\\frac{\\partial}{\\partial r}\\right)+\\frac{1}{r^2}\\left[\\frac{1}{\\sin\\theta}\\frac{\\partial}{\\partial\\theta}\\left(\\sin\\theta\\frac{\\partial}{\\partial\\theta}\\right)+\\frac{1}{\\sin^2\\theta}\\frac{\\partial^2}{\\partial\\phi^2}\\right] = \\frac{1}{r^2}\\frac{\\partial}{\\partial r}r^2\\frac{\\partial}{\\partial r}-\\frac{1}{\\hbar^2r^2}\\hat{L}^2\n",
    "$$\n",
    "where $\\hat{L}$ is the angular momentum operator. Due to the spherical symmetry of the potential, we can split the wave function into a radial and an angular part $\\psi_{n\\ell m}(r,\\theta,\\varphi)=R_{n\\ell}(r)Y^m_\\ell(\\theta, \\varphi)$, where $R_{n\\ell}$ is the radial solution and $Y^m_\\ell$ are the spherical harmonics (solution to the angular equation). The quantum numbers $n\\ell m$ emerge from the imposition of boundary condition on $r$ and periodicity for $\\theta$ and $\\varphi$ respectively.\n",
    "\n",
    "Following the separation of variables leads to the radial eigenvalue equation, which has a simpler form for $\\chi(r)= r\\cdot R(r)$\n",
    "$$\n",
    "-\\frac{\\hbar^2}{2\\mu}\\frac{d^2\\chi\\left(r\\right)}{dr^2}+\\left(V\\left(r\\right)+\\frac{\\hbar^2}{2\\mu}\\frac{\\ell\\left(\\ell+1\\right)}{r^2}\\right)\\chi\\left(r\\right)=E\\chi\\left(r\\right)\n",
    "$$\n",
    "and the angular equation for which the spherical harmonics are solutions\n",
    "$$\n",
    "\\hat{L}^2Y_\\ell^m\\left(\\theta,\\varphi\\right)=\\hbar^2\\ell\\left(\\ell+1\\right)Y_\\ell^m\\left(\\theta,\\varphi\\right)\n",
    "$$\n",
    "where we use the normalised spherical harmonics $\\int_0^\\pi\\mathrm{d}\\theta\\int_0^{2\\pi}\\mathrm{d}\\varphi \\left(Y_{\\ell'}^{m'}\\right)^*Y_\\ell^{m} = \\delta_{\\ell \\ell'} \\delta_{mm'}$\n",
    "$$\n",
    "Y_\\ell^m( \\theta , \\varphi ) = (-1)^m \\sqrt{\\frac{(2\\ell+1)}{4\\pi}\\frac{(\\ell-m)!}{(\\ell+m)!}} \\, P_{\\ell}^m ( \\cos{\\theta} ) \\, e^{i m \\varphi }\n",
    "$$\n",
    "This gives us two equations to solve to obtain the form of the orbitals, and we can show that the solution is\n",
    "$$\n",
    "\\psi_{n\\ell m}(r,\\theta,\\varphi) = \\sqrt{ \\left( \\frac{2Z}{n a^*} \\right)^3 \\frac{ (n-\\ell-1)! }{2n(n+\\ell)!} } e^{-Zr/na^*} \\left(\\frac{2Zr}{na^*}\\right)^\\ell L^{(2\\ell+1)}_{n-\\ell-1}\\left(\\frac{2Zr}{na^*}\\right) Y_\\ell^m(\\theta,\\varphi),\n",
    "$$\n",
    "where\n",
    "* $a^* = \\frac{4\\pi\\varepsilon_0\\hbar^2}{\\mu e^2} = \\frac{m_{\\mathrm{e}}}{\\mu} a_0$ is the effective Bohr radius, with $a_0$ the Bohr radius.\n",
    "* $L_{n-\\ell-1}^{(2 \\ell+1)}(x)$ are generalized Laguerre polynomials (nontrivial solutions of Laugerre's differential equation)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Energies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The energies which are the eigenvalues for the stationary Schrödinger equation depend only on the radial quantum number $n$ and are given by\n",
    "$$\n",
    "E_{n} = -\\left(\\frac{Z^2 \\mu e^4}{32 \\pi^2\\epsilon_0^2\\hbar^2}\\right)\\frac{1}{n^2} = -\\left(\\frac{Z^2\\hbar^2}{2\\mu (a^*)^2}\\right)\\frac{1}{n^2}.\n",
    "$$\n",
    "\n",
    "also the quantum numbers can only take a particular set of values :\n",
    "* $n=1,2,3,4, \\dots$ associated to the energy level of the orbital.\n",
    "* $\\ell=0,1,2,\\dots,n-1$ associated to the total angular momentum $\\hat{L}^2 \\psi_{n\\ell m} = \\hbar^2\\ell(\\ell+1)\\psi_{n\\ell m}$.\n",
    "* $m=-\\ell,-\\ell+1,\\ldots,0,\\ldots,\\ell-1,\\ell$ associated to the quantization axis of the angular momentum (magnetic quantum number) $\\hat{L}_z \\psi_{n\\ell m} = \\hbar m\\psi_{n\\ell m}$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Magnetic dipole moment\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the magnetic dipole moment of the electron orbiting the nucleus ($g_L$-factor is 1 when $m_N\\gg m_e$) is related to the angular momentum through $\\vec{m}_L= -g_L\\frac{e}{2\\mu}\\vec{L}= -g_L\\frac{\\mu_B}{\\hbar}\\vec{L}$, with $g_L=1-\\frac{m_e}{m_N}$. As the angular momentum, this operator is quantized and we have a clear connection along the quantization axis, where\n",
    "$$m_z=-g_L \\mu_B m,$$ \n",
    "for $\\mu_B=\\frac{2m_e}{e\\hbar}$ the Bohr magneton and $m$ the magnetic quantum number. We can also define $\\gamma=g_L \\mu_B$, the gyromagnetic ratio."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Electric dipole moment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "The electric dipole moment of the electron is more complex, we can derive its expression based on the multipole expansion.\n",
    "$$\n",
    "\\vec d = -e\\int_V \\mathrm{d}V\\,\\vec r\\rho(\\vec r) = -e\\int_V \\mathrm{d}V\\,\\vec r|\\psi_{n\\ell m}(\\vec r)|^2 = \\langle\\psi_{n\\ell m}(\\vec r)|-e\\hat{\\vec r}|\\psi_{n\\ell m}(\\vec r)\\rangle\n",
    "$$\n",
    "\n",
    "Note :\n",
    "We can define other multipole moments (quadripole, octopole) through \n",
    "$$\n",
    "Q_{i_1\\cdots i_n} = \\int \\mathrm{d}V\\, T_{i_1\\cdots i_n}(\\vec r) \\rho(\\vec r),\n",
    "$$\n",
    "where the tensor $T_{i_1\\cdots i_n}(\\vec r)$ is defined by\n",
    "$$\n",
    "T_{i_1\\cdots i_n}(\\vec r) \\equiv |\\vec r|^{2n+1} \\frac{\\partial}{\\partial r_{i_1}'}\\cdots \\frac{\\partial}{\\partial r_{i_n}'}\\left.\\frac{1}{|\\vec r - \\vec r'|}\\right|_{\\vec r'=0}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plots & Animations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import math\n",
    "# plotly for the 3D plots\n",
    "import plotly.graph_objs as go\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import Normalize\n",
    "from matplotlib import colormaps # requires matplotlib at least 3.5.0\n",
    "from scipy.special import sph_harm, assoc_laguerre\n",
    "from scipy.integrate import trapz\n",
    "from scipy import constants\n",
    "# scikit-image package for marching_cubes (find isosurfaces in 3d volumetric data)\n",
    "from skimage import measure\n",
    "\n",
    "# VScode\n",
    "%matplotlib widget\n",
    "# web based notebook\n",
    "# %matplotlib inline\n",
    "\n",
    "# plt.rc('text', usetex=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Defining some utility functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cart2sph(X, Y, Z, invert_xy=False):\n",
    "    xy2 = X**2 + Y**2\n",
    "    R = np.sqrt(xy2 + Z**2)    \n",
    "    Theta = np.arctan2(np.sqrt(xy2), Z)\n",
    "    if invert_xy:\n",
    "        Phi = np.arctan2(X, Y)\n",
    "    else:\n",
    "        Phi = np.arctan2(Y, X)\n",
    "    # if np.size(Phi) == 1:\n",
    "    #     if Phi < 0:\n",
    "    #         Phi += np.pi * 2\n",
    "    # else:\n",
    "    #     Phi[Phi < 0] += np.pi * 2\n",
    "    return R, Theta, Phi\n",
    "\n",
    "def sph2cart(r, theta, phi):\n",
    "    return r*np.sin(theta)*np.cos(phi), r*np.sin(theta)*np.sin(phi), r*np.cos(theta)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Definition of orbitals"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we define functions to compute the different parts of the orbitals (note that we use the Condon-Shortley convention for the spherical harmonics). About the spherical harmonics, the usual one $Y_m^\\ell(\\theta, \\varphi)$ are complex and maybe difficult to give a clear interpretation, therefore we provide also the possibility to compute the real orbitals $Y_{\\ell m}(\\theta, \\varphi)$.\n",
    "$$\n",
    "Y_{\\ell m} = \\sqrt{2} \\, (-1)^m \\, \\Im [{Y_\\ell^{|m|}}] \\qquad\\qquad \\text{if } m<0\\\\\n",
    "Y_{\\ell m} = Y_\\ell^0 \\qquad\\qquad\\qquad\\qquad\\qquad \\text{if } m=0\\\\\n",
    "Y_{\\ell m} = \\sqrt{2} \\, (-1)^m \\, \\Re [{Y_\\ell^m}] \\qquad\\qquad \\text{if } m>0.\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_sph_harm(l, m, theta, phi, real_orbital=True):\n",
    "    if abs(m) > l:\n",
    "        raise ValueError(\"The quantum numbers are not compatible, we need |m| <= l\")\n",
    "    # in SciPy's sph_harm function the azimuthal coordinate, theta, comes before the polar coordinate, phi.\n",
    "    # the Condon-Shortley phase is included in the associated Legendre functions.\n",
    "    Y = sph_harm(abs(m), l, phi, theta)\n",
    "\n",
    "    if real_orbital:\n",
    "        # Linear combination of Y_l,m and Y_l,-m to create the real form.\n",
    "        if m < 0:\n",
    "            Y = np.sqrt(2) * (-1)**m * Y.imag\n",
    "        elif m > 0:\n",
    "            Y = np.sqrt(2) * (-1)**m * Y.real\n",
    "    else:\n",
    "        if m < 0:\n",
    "            Y = (-1)**m*Y.conj()\n",
    "    return Y\n",
    "\n",
    "def compute_radial_part(n, l, r, a_0=1, m_e_over_m_N=0, Z=1):\n",
    "    if n < 1:\n",
    "        raise ValueError(\"The quantum numbers are incorrect, we need n >= 1\")\n",
    "    if n <= l:\n",
    "        raise ValueError(\"The quantum numbers are not compatible, we need n > l\")\n",
    "    a_mu = a_0 * (1 + m_e_over_m_N)\n",
    "    factor = 2.*Z/(n*a_mu)\n",
    "    rho = factor * r\n",
    "    prefactor = np.sqrt(factor**3*math.factorial(n-l-1)/(2.*n*math.factorial(n+l)))\n",
    "    R = prefactor * assoc_laguerre(rho, n-l-1, 2*l+1) * np.exp(-rho/2.) * rho**l\n",
    "    return R\n",
    "\n",
    "def compute_sign_radial_part(n, l, r, a_0=1, m_e_over_m_N=0, Z=1):\n",
    "    a_mu = a_0 * (1 + m_e_over_m_N)\n",
    "    factor = 2.*Z/(n*a_mu)\n",
    "    rho = factor * r\n",
    "    R = np.sign(assoc_laguerre(rho, n-l-1, 2*l+1))\n",
    "    return R\n",
    "\n",
    "def compute_orbitals(n, l, m, r, theta, phi, a_0=1, m_e_over_m_N=0, Z=1, real_orbital=True):\n",
    "    R = compute_radial_part(n, l, r, a_0, m_e_over_m_N, Z)\n",
    "    Y = compute_sph_harm(l, m, theta, phi, real_orbital)\n",
    "    return R * Y\n",
    "\n",
    "def check_normalization_radial_part(n, l, r, a_0=1, m_e_over_m_N=0, Z=1):\n",
    "    R = compute_radial_part(n, l, r, a_0, m_e_over_m_N, Z)\n",
    "    print(f\"Norm of the Radial part : {trapz((r*R)**2, r)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Definition of basic plotting fucntions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plotting spherical harmonics for given l, m\n",
    "def plot_sph_harm(ax, l, m, theta, phi, real_orbital=True):\n",
    "    \"\"\"Plot the spherical harmonic of degree l and order m on Axes ax.\"\"\"\n",
    "    Y = compute_sph_harm(l, m, theta, phi, real_orbital=real_orbital)\n",
    "    xyz = np.array([np.sin(theta) * np.sin(phi),\n",
    "                    np.sin(theta) * np.cos(phi), \n",
    "                    np.cos(theta)])\n",
    "    Yx, Yy, Yz = np.abs(Y) * xyz\n",
    "\n",
    "    # Colour the plotted surface according to the sign of Y.\n",
    "    cmap = plt.cm.ScalarMappable(cmap=plt.get_cmap('PRGn'))\n",
    "    cmap.set_clim(-0.5, 0.5)\n",
    "\n",
    "    ax.plot_surface(Yx, Yy, Yz, facecolors=cmap.to_rgba(Y.real), rstride=2, cstride=2)\n",
    "\n",
    "    if real_orbital:\n",
    "        ax.set_title(r'real_orbital (real) Spherical Harmonic with l={}, m={} : $Y_{{{},{}}}(\\theta, \\varphi)$'.format(l, m, l, m))\n",
    "    else:\n",
    "        ax.set_title(r'Laplace Spherical Harmonic with l={}, m={} : $Y_{{{}}}^{{{}}}(\\theta, \\varphi)$'.format(l, m, l, m))\n",
    "    ax_lim = 0.5\n",
    "    ax.set_xlim(-ax_lim, ax_lim)\n",
    "    ax.set_ylim(-ax_lim, ax_lim)\n",
    "    ax.set_zlim(-ax_lim, ax_lim)\n",
    "    ax.axis('off')\n",
    "\n",
    "# plotting radial solution for given n, l\n",
    "def plot_radial_part(ax, n, l, r, a_0=1, m_e_over_m_N=0, Z=1):\n",
    "    \"\"\"Plot the spherical harmonic of degree l and order m on Axes ax.\"\"\"\n",
    "    R = compute_radial_part(n, l, r, a_0, m_e_over_m_N, Z)\n",
    "    ax.plot(r, np.abs(r*R))\n",
    "    # ax.plot(r, np.abs(R))\n",
    "    ax.grid(True)\n",
    "    ax.set_title(r'Radial distribution with n={}, l={} : $r^2\\cdot |R_{{{},{}}}(r)|^2$'.format(n, l, n, l))\n",
    "    # ax.axis('off')\n",
    "\n",
    "# plotting both spherical harmonics and radial solution for given n, l, m\n",
    "def plot_components(fig, n, l, m, r, theta, phi, a_0=1, m_e_over_m_N=0, Z=1, real_orbital=True):\n",
    "    theta, phi = np.meshgrid(theta, phi)\n",
    "    ax0 = fig.add_subplot(1, 2, 1, projection='3d')\n",
    "    plot_sph_harm(ax0, l, m, theta, phi, real_orbital)\n",
    "    ax1 = fig.add_subplot(1, 2, 2)\n",
    "    plot_radial_part(ax1, n, l, r, a_0, m_e_over_m_N, Z)\n",
    "    # plt.savefig('Y{}_{}.png'.format(l, m))\n",
    "    plt.show()\n",
    "\n",
    "# plotting the 3d orbital as an isosurface\n",
    "def plot_orbitals(n, l, m, x_, y_, z_, a_0=1, m_e_over_m_N=0, Z=1, real_orbital=True, colorscale=\"hsv\", isoprob=None):\n",
    "    dxdydz = np.array([x_[1]-x_[0], y_[1]-y_[0], z_[1]-z_[0]])\n",
    "    x_, y_, z_ = np.meshgrid(x_, y_, z_)\n",
    "    r, theta, phi = cart2sph(x_, y_, z_)\n",
    "    \n",
    "    psi = compute_orbitals(n, l, m, r, theta, phi, a_0, m_e_over_m_N, Z, real_orbital)\n",
    "    magnitude_sq = abs(psi)**2\n",
    "    # isosurface value\n",
    "    if isoprob is None:\n",
    "        isoprob = np.max(magnitude_sq)/10\n",
    "    print(isoprob)\n",
    "    verts, faces, _, values = measure.marching_cubes(magnitude_sq, isoprob, spacing=(1, 1, 1))\n",
    "    verts *= dxdydz\n",
    "    verts -= np.mean(verts)\n",
    "    x, y, z = verts.T\n",
    "    I, J, K = faces.T\n",
    "    means = np.zeros((3, len(I)))\n",
    "    for idx in range(len(I)):\n",
    "        for i in range(3):\n",
    "            means[i, idx] += np.sum([verts[faces[idx, j], i] for j in range(3)])\n",
    "    means /= 3\n",
    "    norm = Normalize(vmin=-np.pi, vmax=np.pi)\n",
    "    \n",
    "    # 2 more runs needed for the colors (phase)\n",
    "    R_m, Theta_m, Phi_m = cart2sph(means[0], means[1], means[2], invert_xy=True)\n",
    "    R, Theta, Phi = cart2sph(x, y, z, invert_xy=True)\n",
    "    values_faces = np.angle(compute_sign_radial_part(n, l, R_m, a_0=a_0, m_e_over_m_N=m_e_over_m_N, Z=Z)*compute_sph_harm(l, m, Theta_m, Phi_m, real_orbital))\n",
    "    values_vertices = np.angle(compute_sign_radial_part(n, l, R, a_0=a_0, m_e_over_m_N=m_e_over_m_N, Z=Z)*compute_sph_harm(l, m, Theta, Phi, real_orbital))\n",
    "\n",
    "    face_colors = colormaps[colorscale](norm(values_faces))  \n",
    "    vertices_colors = colormaps[colorscale](norm(values_vertices))\n",
    "\n",
    "    fig = go.Figure(go.Mesh3d(x=x, y=y, z=z, colorscale=colorscale, vertexcolor=vertices_colors, facecolor=face_colors, i=I, j=J, k=K, name=''))\n",
    "    # adding the colorbar\n",
    "    fig.add_scatter3d(x=x, y=y, z=z, mode='markers', opacity=0, marker=dict(size=0.01, color=values_vertices, colorscale=colorscale, showscale=True, colorbar_thickness=24))\n",
    "    fig.update_layout(width=600, height=600, font_size=11)\n",
    "    fig.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameter definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if a_0 =1, the radius is given in terms of Bohr radius units\n",
    "a_0 = 1\n",
    "# ratio of electron mass over nucleus mass (usually m_N >> m_e)\n",
    "m_e_over_m_N = 0\n",
    "# atomic number (number of protons in the nucleus)\n",
    "Z = 1\n",
    "# colorscale for the phases (use a periodic colorscale such as \"hsv\" or \"twilight\")\n",
    "# colorscale_phase = \"hsv\"\n",
    "colorscale_phase = \"twilight\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting spherical harmonics & radial solution"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we start plotting each component separately. This plot is useful to ensure that the range of values of $r$ is sufficient to showcase the main aspects of the behaviour of the wavefunction. For this, we can check the normalisation of the radial wavefunction.\n",
    "\n",
    "You should obtain pictures similar to [these ones](https://en.wikipedia.org/wiki/Spherical_harmonics#/media/File:Spherical_Harmonics.png) for the real spherical harmonics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = np.linspace(0, 50*a_0, 300)\n",
    "theta = np.linspace(0, np.pi, 100)\n",
    "phi = np.linspace(0, 2*np.pi, 100)\n",
    "fig = plt.figure(figsize=(12, 5))\n",
    "\n",
    "n, l, m = 4, 3, 3\n",
    "# n, l, m = 5, 4, 0\n",
    "# n, l, m = 2, 0, 0\n",
    "real_orbital = False\n",
    "\n",
    "# just checking the normalisation of the radial part\n",
    "check_normalization_radial_part(n, l, r, a_0=a_0, m_e_over_m_N=m_e_over_m_N, Z=Z)\n",
    "\n",
    "plot_components(fig, n, l, m, r, theta, phi, a_0=a_0, m_e_over_m_N=m_e_over_m_N, Z=Z, real_orbital=real_orbital)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = np.linspace(0, 50*a_0, 300)\n",
    "theta = np.linspace(0, np.pi, 100)\n",
    "phi = np.linspace(0, 2*np.pi, 100)\n",
    "fig = plt.figure(figsize=(12, 5))\n",
    "\n",
    "n, l, m = 4, 3, 3\n",
    "# n, l, m = 5, 4, 0\n",
    "# n, l, m = 2, 0, 0\n",
    "real_orbital = True\n",
    "\n",
    "# just checking the normalisation of the radial part\n",
    "check_normalization_radial_part(n, l, r, a_0=a_0, m_e_over_m_N=m_e_over_m_N, Z=Z)\n",
    "\n",
    "plot_components(fig, n, l, m, r, theta, phi, a_0=a_0, m_e_over_m_N=m_e_over_m_N, Z=Z, real_orbital=real_orbital)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting the 3d orbital"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we plot a 3d representation of the wavefunction, so combining both the spherical harmonics and the radial part. In this case, we compare to [these ones](https://en.wikipedia.org/wiki/File:Atomic_orbitals_spdf_m-eigenstates_and_superpositions.png), where the complex (original) wavefunctions should have isosurfaces similar to these right orbitals, while the real orbitals should be similar to the left ones (use \"hsv\" colorscale to get the same color code).\n",
    "\n",
    "You can also play with the isoprob (value for the isosurface), the span of the box in which the wavefucntion is represented, the spacing between the points of the mesh and the values of $n, \\ell, m$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n, l, m = 4, 2, -2\n",
    "n, l, m = 3, 2, 1\n",
    "# n, l, m = 1, 0, 0\n",
    "isoprob = 2e-5\n",
    "span = 50\n",
    "dx = 0.5\n",
    "real_orbital = False\n",
    "x = y = z = np.arange(-span/2,span/2, dx)\n",
    "\n",
    "plot_orbitals(n, l, m, x, y, z, a_0=a_0, m_e_over_m_N=m_e_over_m_N, Z=Z, real_orbital=real_orbital, colorscale=colorscale_phase, isoprob=isoprob)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n, l, m = 4, 2, -2\n",
    "# n, l, m = 1, 0, 0\n",
    "isoprob = 2e-5\n",
    "span = 50\n",
    "dx = 0.5\n",
    "real_orbital = True\n",
    "x = y = z = np.arange(-span/2,span/2, dx)\n",
    "\n",
    "plot_orbitals(n, l, m, x, y, z, a_0=a_0, m_e_over_m_N=m_e_over_m_N, Z=Z, real_orbital=real_orbital, colorscale=colorscale_phase, isoprob=isoprob)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Animation of orbital transition"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this section, we plot an animation of the time evolution of the wavefunction going from one eigenstate to another. We consider the Rabi oscillation with a detuning $\\Delta = \\omega_D - \\omega_{21}$ where $\\omega_{21} = \\omega_2 - \\omega_1$ and $\\omega_D$ is the drive frequency. We also define the Rabi frequency $\\Omega_R$ and the detuned / modified Rabi frequency $\\Omega = \\sqrt{\\Omega_R^2 + \\Delta^2}$. The time evolution of the wavefunction in the interaction picture is now given by (see Semiclassical Atom-Field interaction exercise or https://doi.org/10.1038/s41586-022-04948-y)\n",
    "$$\n",
    "\\psi_I(r, \\theta, \\varphi, t) = \\left[\\cos\\left(\\frac{\\Omega t}{2}\\right)-i\\frac{\\Delta}{\\Omega}\\sin\\left(\\frac{\\Omega t}{2}\\right)\\right]e^{i\\Delta t/2} \\cdot \\psi_{n_1\\ell_1 m_1}(r, \\theta, \\varphi) +i\\frac{\\Omega_R}{\\Omega}\\sin\\left(\\frac{\\Omega t}{2}\\right)e^{-i\\Delta t/2} \\cdot \\psi_{n_2\\ell_2 m_2}(r, \\theta, \\varphi)\n",
    "$$\n",
    "In the original picture, we have to restore an additional phase factor related to the energy of each state\n",
    "$$\n",
    "\\psi_S(r, \\theta, \\varphi, t) = \\left[\\cos\\left(\\frac{\\Omega t}{2}\\right)-i\\frac{\\Delta}{\\Omega}\\sin\\left(\\frac{\\Omega t}{2}\\right)\\right]e^{i(\\Delta/2 - \\omega_1) t}\\cdot \\psi_{n_1\\ell_1 m_1}(r, \\theta, \\varphi) +i\\frac{\\Omega_R}{\\Omega}\\sin\\left(\\frac{\\Omega t}{2}\\right)e^{-i(\\Delta/2+ \\omega_2) t} \\cdot \\psi_{n_2\\ell_2 m_2}(r, \\theta, \\varphi)\n",
    "$$\n",
    "We plot this evolution in a certain number of frames $N_{\\mathrm{frames}}$ and over a total time $T_{\\mathrm{max}}$. We assume our temporal unit to be the (detuned) Rabi period $T_R = 2\\pi / \\Omega$ (period of the square modulus of the wavefunction $|\\psi(r, \\theta, \\varphi, t)|^2$), meaning that we will count time $\\tau=t/T_R$ in units of Rabi period $T_R$ for $N_{\\mathrm{frames}}$ values of time and a range of $t=0, ... T_{\\mathrm{max}}$. This allows us to define every frequencies as ratio with $\\Omega$ and gives the following form (omitting the global phase factor $e^{2\\pi i\\tau (\\Delta/2 - \\omega_1)/\\Omega}$)\n",
    "$$\n",
    "\\psi_S(r, \\theta, \\varphi, \\tau) = \\left[\\cos\\left(\\pi\\tau\\right)-i\\frac{\\Delta}{\\Omega}\\sin\\left(\\pi\\tau\\right)\\right]\\cdot \\psi_{n_1\\ell_1 m_1}(r, \\theta, \\varphi) +i\\frac{\\Omega_R}{\\Omega}\\sin\\left(\\pi\\tau\\right)e^{- 2\\pi i\\tau (\\Delta+ \\omega_{21})/\\Omega} \\cdot \\psi_{n_2\\ell_2 m_2}(r, \\theta, \\varphi)\n",
    "$$\n",
    "Note that when we are on resonance, the detuning vanishes $\\Delta=0$ and the evolution simplifies to\n",
    "$$\n",
    "\\psi_S(r, \\theta, \\varphi, \\tau) = \\cos\\left(\\pi\\tau\\right)\\cdot \\psi_{n_1\\ell_1 m_1}(r, \\theta, \\varphi) +i\\sin\\left(\\pi\\tau\\right)e^{- 2\\pi i\\tau \\omega_{21}/\\Omega} \\cdot \\psi_{n_2\\ell_2 m_2}(r, \\theta, \\varphi)\n",
    "$$\n",
    "In order to perform a certain number of Rabi oscillation, we pick $T_{\\mathrm{max}} = k\\cdot T_R$ or equivalently $\\tau=0, ... k$. To return to the original state, we need an integer number of full period ($k=1$ to simplify), while in order to end on second state at the end of the simulation, we need to pick $k=1/2$. \n",
    "Note that the final state after half a period $\\tau=1/2$ or $t=T_R/2$ is\n",
    "$$\n",
    "\\psi_S(r, \\theta, \\varphi, \\tau=1/2) =-i\\frac{\\Delta}{\\Omega}\\cdot \\psi_{n_1\\ell_1 m_1}(r, \\theta, \\varphi) +i\\frac{\\Omega_R}{\\Omega}e^{- i\\pi(\\Delta+ \\omega_{21})/\\Omega} \\cdot \\psi_{n_2\\ell_2 m_2}(r, \\theta, \\varphi)\\underset{\\Delta\\to 0}{\\longrightarrow}i e^{- i\\pi\\omega_{21}/\\Omega}\\cdot \\psi_{n_2\\ell_2 m_2}(r, \\theta, \\varphi)\n",
    "$$\n",
    "\n",
    "For the plots, we introduce the ratio `Delta_over_Omega_R` to account for the detuning. This will determine how well you will transition to the second state and probabilty to still be in the original state even after half a period $T_R/2$. Indeed\n",
    "$$\n",
    "|\\langle\\psi_{n_1\\ell_1 m_1}|\\psi(r, \\theta, \\varphi, t=T_R/2)\\rangle|^2 = \\frac{\\Delta^2}{\\Omega^2} = \\frac{1}{1+\\frac{1}{\\Delta^2/\\Omega_R^2}} \\underset{\\Delta\\to 0}{\\longrightarrow} 0 \n",
    "$$\n",
    "while the probabilty to be in the new state is \n",
    "$$\n",
    "|\\langle\\psi_{n_2\\ell_2 m_2}|\\psi(r, \\theta, \\varphi, t=T_R/2)\\rangle|^2 = \\frac{\\Omega_R^2}{\\Omega^2} = \\frac{1}{1+\\Delta^2/\\Omega_R^2}\\underset{\\Delta\\to 0}{\\longrightarrow} 1\n",
    "$$\n",
    "\n",
    "Note that in the original Schrödinger picture, the relative phase has another component which is the transition energies between the atomic orbitals. This energy is given by the well-known Rydberg formula with $R_\\infty$ the Rydberg constant, $Z$ the atomic number of the atom\n",
    "$$\n",
    "\\Delta E_{21} = \\hbar\\omega_{21} = \\frac{h c}{\\lambda_{21}} = -h c R_{\\infty}\\left(\\frac{1}{n_2^2}-\\frac{1}{n_1^2}\\right)\\frac{Z^2}{1+\\frac{m_e}{m_N}} = -h c R_{N}\\left(\\frac{1}{n_2^2}-\\frac{1}{n_1^2}\\right)Z^2\n",
    "$$\n",
    "Note that here, $hcR_\\infty$ is the unit of energy. It is important to notice that in this representation, the quantity that matters is the ratio of the unit energy to the energy associated to the Rabi frequency :\n",
    "$$\n",
    "\\frac{hcR_{\\infty}}{\\hbar\\Omega_R} = \\frac{hcR_\\infty}{\\langle n_2,\\ell_2,m_2|e\\hat{\\boldsymbol{r}}\\cdot \\boldsymbol{E_0}|n_1,\\ell_1,m_1\\rangle}\n",
    "$$\n",
    "In https://doi.org/10.1038/s41586-022-04948-y, there is a ratio of $\\approx 300$ if we use the transition of one electron from 1s ($n_1=1$) to 4p ($n_2=4$), and this is for Helium with $Z=2$. Accounting for these 2 effects in the transition energy and using the \"energy unit\" instead of the  we have $\\frac{hcR_\\infty}{\\hbar\\Omega_R}\\approx 80$. This ratio depends clearly on the amplitude of the electric field, and will be one more free parameter `hcR_inf_over_Omega_R` to control the evolution in the original frame (this ratio has no impact on the interaction picture, where it can be set to 0). For interpretation, this parameters sort of controls the frame in which we place ourselves. Using the example above, setting `hcR_inf_over_Omega_R` to 0 gives the interaction picture, while setting it to 80 gives the original picture."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# computing each step needed for the animated version of the transition between 2 orbitals\n",
    "def compute_orbital_transitions(n1, l1, m1, n2, l2, m2, x_, y_, z_, t, Omega_R, Delta, hcR_inf_over_Omega_R=0, a_0=1, m_e_over_m_N=0, Z=1, real_orbital=True, colorscale=\"hsv\", isoprob=None):\n",
    "    dxdydz = np.array([x_[1]-x_[0], y_[1]-y_[0], z_[1]-z_[0]])\n",
    "    x_, y_, z_ = np.meshgrid(x_, y_, z_)\n",
    "    r, theta, phi = cart2sph(x_, y_, z_)\n",
    "\n",
    "    w1 = - hcR_inf_over_Omega_R * Omega_R * (Z / n1)**2 / (1 + m_e_over_m_N)\n",
    "    w2 = - hcR_inf_over_Omega_R * Omega_R * (Z / n2)**2 / (1 + m_e_over_m_N)\n",
    "    Omega = np.sqrt(Omega_R**2 + Delta**2)\n",
    "    a = (np.cos(Omega*t/2) - 1j * Delta / Omega * np.sin(Omega*t/2)) * np.exp(1j*(Delta/2-w1)*t)\n",
    "    b = 1j * Omega_R / Omega * np.sin(Omega*t/2) * np.exp(-1j*(Delta/2+w2)*t)\n",
    "\n",
    "    def psi_(psi1, psi2, a, b):\n",
    "        return  a * psi1 + b * psi2\n",
    "    \n",
    "    psi1 = compute_orbitals(n1, l1, m1, r, theta, phi, a_0, m_e_over_m_N, Z, real_orbital)\n",
    "    psi2 = compute_orbitals(n2, l2, m2, r, theta, phi, a_0, m_e_over_m_N, Z, real_orbital)\n",
    "    psi = psi_(psi1, psi2, a, b)\n",
    "    magnitude_sq = abs(psi)**2\n",
    "    # isosurface value\n",
    "    if isoprob is None:\n",
    "        isoprob = np.max(magnitude_sq)/10\n",
    "    # print(isoprob)\n",
    "    verts, faces, _, _ = measure.marching_cubes(magnitude_sq, isoprob, spacing=(1, 1, 1))\n",
    "    verts *= dxdydz\n",
    "    verts -= np.mean(verts)\n",
    "    x, y, z = verts.T\n",
    "    I, J, K = faces.T\n",
    "    means = np.zeros((3, len(I)))\n",
    "    for idx in range(len(I)):\n",
    "        for i in range(3):\n",
    "            means[i, idx] += np.sum([verts[faces[idx, j], i] for j in range(3)])\n",
    "    means /= 3\n",
    "    norm = Normalize(vmin=-np.pi, vmax=np.pi)\n",
    "    \n",
    "    R_m, Theta_m, Phi_m = cart2sph(means[0], means[1], means[2], invert_xy=True)\n",
    "    R, Theta, Phi = cart2sph(x, y, z, invert_xy=True)\n",
    "    psi1 = compute_orbitals(n1, l1, m1, R_m, Theta_m, Phi_m, a_0, m_e_over_m_N, Z, real_orbital)\n",
    "    psi2 = compute_orbitals(n2, l2, m2, R_m, Theta_m, Phi_m, a_0, m_e_over_m_N, Z, real_orbital)\n",
    "    values_faces = np.angle(psi_(psi1, psi2, a, b))\n",
    "    psi1 = compute_orbitals(n1, l1, m1, R, Theta, Phi, a_0, m_e_over_m_N, Z, real_orbital)\n",
    "    psi2 = compute_orbitals(n2, l2, m2, R, Theta, Phi, a_0, m_e_over_m_N, Z, real_orbital)\n",
    "    values_vertices = np.angle(psi_(psi1, psi2, a, b))\n",
    "\n",
    "    face_colors = colormaps[colorscale](norm(values_faces))  \n",
    "    vertices_colors = colormaps[colorscale](norm(values_vertices))\n",
    "\n",
    "    return x, y, z, I, J, K, vertices_colors, face_colors\n",
    "\n",
    "# compiling the animated version of the transition between 2 orbitals\n",
    "def plot_orbital_transitions(n1, l1, m1, n2, l2, m2, x_, y_, z_, nbr_Rabi_period=1/2., N_frames=20, Delta_over_Omega_R=0, hcR_inf_over_Omega_R=0, a_0=1, m_e_over_m_N=0, Z=1, real_orbital=True, colorscale=\"hsv\", isoprob=None):\n",
    "    frames = []\n",
    "    data0 = None\n",
    "    Omega = 2*np.pi/N_frames\n",
    "    Omega_R = Omega / np.sqrt(1+Delta_over_Omega_R**2)\n",
    "    Delta = Omega_R * Delta_over_Omega_R\n",
    "    for n in range(N_frames+1):\n",
    "        tau = n * nbr_Rabi_period\n",
    "        x, y, z, I, J, K, vertices_colors, face_colors = compute_orbital_transitions(n1, l1, m1, n2, l2, m2, x_, y_, z_, tau, Omega_R, Delta, hcR_inf_over_Omega_R=hcR_inf_over_Omega_R, a_0=a_0, m_e_over_m_N=m_e_over_m_N, Z=Z, real_orbital=real_orbital, colorscale=colorscale, isoprob=isoprob)\n",
    "        data=[\n",
    "            go.Mesh3d(x=x, y=y, z=z, colorscale=colorscale, vertexcolor=vertices_colors, facecolor=face_colors, i=I, j=J, k=K)\n",
    "        ]\n",
    "        if n == 0:\n",
    "            data0 = data\n",
    "        frames.append(go.Frame(data=data))\n",
    "\n",
    "    fig = go.Figure(\n",
    "        data=data0,\n",
    "        layout=go.Layout(\n",
    "            autosize=False,\n",
    "            width=700,\n",
    "            height=600,\n",
    "            margin=dict(l=40, r=40, b=40, t=40),\n",
    "            updatemenus=[{\n",
    "                \"type\": \"buttons\",\n",
    "                \"direction\": \"right\",\n",
    "                \"buttons\": [\n",
    "                    {\n",
    "                        \"args\": [None, {\"frame\": {\"duration\": 500, \"redraw\": True}, \"fromcurrent\": True, \"transition\": {\"duration\": 300, \"easing\": \"quadratic-in-out\"}}],\n",
    "                        \"label\": \"Play\",\n",
    "                        \"method\": \"animate\"\n",
    "                    },\n",
    "                    {\n",
    "                        \"args\": [[None], {\"frame\": {\"duration\": 0, \"redraw\": False}, \"mode\": \"immediate\", \"transition\": {\"duration\": 0}}],\n",
    "                        \"label\": \"Pause\",\n",
    "                        \"method\": \"animate\"\n",
    "                    }\n",
    "                ],\n",
    "                \"direction\": \"left\",\n",
    "                \"pad\": {\"r\": 10, \"t\": 87},\n",
    "                \"showactive\": False,\n",
    "                \"type\": \"buttons\",\n",
    "                \"x\": 0.1,\n",
    "                \"xanchor\": \"right\",\n",
    "                \"y\": 0,\n",
    "                \"yanchor\": \"top\"\n",
    "            }],\n",
    "        ),\n",
    "        frames=frames,\n",
    "    )\n",
    "\n",
    "    fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Transition $211 \\to 321$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n1, l1, m1 = 2, 1, 1\n",
    "n2, l2, m2 = 3, 2, 1\n",
    "isoprob = 2e-5\n",
    "span = 50\n",
    "dx = 0.5\n",
    "N_frames = 10\n",
    "Delta_over_Omega_R=0\n",
    "nbr_Rabi_period = 1/2.\n",
    "real_orbital = True\n",
    "hcR_inf_over_Omega_R = 0\n",
    "x = y = z = np.arange(-span/2,span/2, dx)\n",
    "\n",
    "plot_orbital_transitions(n1, l1, m1, n2, l2, m2, x, y, z, nbr_Rabi_period=nbr_Rabi_period, N_frames=N_frames, Delta_over_Omega_R=Delta_over_Omega_R, hcR_inf_over_Omega_R=hcR_inf_over_Omega_R, a_0=a_0, m_e_over_m_N=m_e_over_m_N, Z=Z, real_orbital=real_orbital, colorscale=colorscale_phase, isoprob=isoprob)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Transition $100 \\to 211$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n1, l1, m1 = 1, 0, 0\n",
    "n2, l2, m2 = 2, 1, 0\n",
    "isoprob = 2e-5\n",
    "span = 50\n",
    "dx = 0.5\n",
    "N_frames = 10\n",
    "Delta_over_Omega_R=0\n",
    "nbr_Rabi_period = 1/2.\n",
    "real_orbital = True\n",
    "hcR_inf_over_Omega_R = 0 # original picture\n",
    "x = y = z = np.arange(-span/2,span/2, dx)\n",
    "\n",
    "plot_orbital_transitions(n1, l1, m1, n2, l2, m2, x, y, z, nbr_Rabi_period=nbr_Rabi_period, N_frames=N_frames, Delta_over_Omega_R=Delta_over_Omega_R, hcR_inf_over_Omega_R=hcR_inf_over_Omega_R, a_0=a_0, m_e_over_m_N=m_e_over_m_N, Z=Z, real_orbital=real_orbital, colorscale=colorscale_phase, isoprob=isoprob)"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
