{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hydrogen Orbitals Part 2\n",
    "\n",
    "In this notebook, we'll look at a few more ways to plot data, and will work on refactoring our `h_orbital` code into a python class to make our code more modular and reusable. This will also make it easy for us to transform into Cartesian coordinates instead of spherical coordinates, and allow us to plot orbitals from different atoms on the same graph. Along the way, we'll discuss good coding practices."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating the HAtom class\n",
    "\n",
    "Classes are data structures that contain both variables and functions. When using a class in your code, you create an **object** which is an instance of the class. The class definition tells what the object is made from and how it is built, and the object itself is what interacts with the rest of your code. Using classes and objects in your code has several benefits:\n",
    "\n",
    "- **Data organization**: The variables associated with an object are bundled together. As we'll see, if we want to plot mutliple atoms, we might want to keep track of each atom's center position and quantum numbrers. Instead of keeping track of everything in a master list, we can instead create each atom as an object that keeps track of its own data. It also keeps related code together so that when there is a problem, you know where to look to fix it.\n",
    "- **Reusability**: Once you have created a class that you like, you can import that class into other code that you write. Many of the python data structures you've used (strings, figures, axes) are classes.\n",
    "- **Flexibility**: Classes can be \"extended\" by creating a subclass that inherits all of the variables and functions from the parent class. This makes it easy to add new functionality without having to rewrite everything from scratch.\n",
    "\n",
    "We'll start by building a basic class for the hydrogen atom. If you have never built a python class, have a look at this [basic tutorial](http://openbookproject.net/thinkcs/python/english3e/classes_and_objects_I.html) before reading the code below. Particularly pay attention to the meanings of `__init__` and `self`!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import scipy.special as sps\n",
    "\n",
    "class HAtom:\n",
    "    \"\"\" Class for representing a hydrogen atom\n",
    "    \n",
    "    Here we can provide examples of how to use the class, and list the\n",
    "    important functions a user of the class may want to access. Since this\n",
    "    is a tutorial, we will omit the details here and show it in code.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, n : int, l : int, m : int) -> None:\n",
    "        \"\"\" Constructor: sets the n, l, and m quantum numbers \"\"\"\n",
    "        \n",
    "        self.n = n\n",
    "        self.l = l\n",
    "        self.m = m\n",
    "        \n",
    "    def orbital(self, r : np.ndarray, theta : np.ndarray, phi : np.ndarray) -> np.ndarray:\n",
    "        \"\"\"Calculates the hydrogen orbital on the $r$, $\\theta$, $\\phi$ grid\n",
    "        \n",
    "        Note: Here we could add a detailed description of the how the orbital\n",
    "        is represented, noting that the angular convention we're using is opposite\n",
    "        that of sp.special.sph_harm.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        r : np.ndarray\n",
    "            $r$ values, in atomic units\n",
    "        theta : np.ndarray\n",
    "            colatitudinal angle in radians, in the range [0,pi]\n",
    "        phi : np.ndarray\n",
    "            azimuthal angle in radians, in the range [0,2pi)\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        np.ndarray\n",
    "            3D array containing hydrogen atomic orbital evaluated at r, theta, phi\n",
    "        \"\"\"\n",
    "        \n",
    "        pf = ( (2./self.n)**3. * sps.factorial(self.n-self.l-1) / (2. * self.n * sps.factorial(self.n+self.l)) )**0.5\n",
    "        radial = np.exp(-r/2.) * r**self.l * sps.eval_genlaguerre(self.n-self.l-1, 2*self.l+1, r)\n",
    "        angular = sps.sph_harm(self.m, self.l, phi, theta)\n",
    "        return pf * radial * angular\n",
    "    \n",
    "    def __repr__(self) -> str:\n",
    "        return f'{type(self).__name__}(n={self.n}, l={self.l}, m={self.m})'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The HAtom class consists of three functions:\n",
    "\n",
    "- `__init__` is the function that initializes an instance of the class. When we create an object, we need to provide the arguments to this function, such as `h1 = HAtom(1,0,0)`.\n",
    "- `orbital` is the same function as `h_orbital` from the previous notebook. However, we do not need to provide the `n`, `l`, and `m` values when we call this function because they are already stored when we created the object.\n",
    "- `__repr__` is a special python function that tells the interpreter how to represent the object when it is entered as a command or in the print function. See the next code block for an example.\n",
    "\n",
    "All of the information in quotes are [docstrings](https://www.python.org/dev/peps/pep-0257/). When creating a class, it is good practice to describe what the class is used for, and to document all of the functions and attributes that a class user needs to know about. Writing detailed comments in your code is **always** a good idea, especially when the code is complicated. Good comments will save you time in the long run, and will make it much easier for people to work on your code (especially your future self!).\n",
    "\n",
    "Finally, we've included [type hints](https://www.python.org/dev/peps/pep-0484/) into the function signature, which lets the user know what kinds of input the functions expect and what kind of output they produce. The arguments to `__init__` are all integers, and it does not return anything, while the `orbital` function expects arrays and returns an array. Python as a language uses [Duck Typing](https://en.wikipedia.org/wiki/Duck_typing), so it doesn't enforce that the parameters you pass to a function be of the types specifiec by the type hints. They are mostly used to inform the end user of how you, as the programmer, intended for the function to work, even if python allows it to be used in other ways that you did not intend. There are also some external tools that can be used to analyze your code and warn you if a parameter you pass to a function doesn't match the function's type hints, and that can be useful to find bugs in large projects.\n",
    "\n",
    "To create an `HAtom` object we need to provide the arguments to the `__init__` function like this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h1 = HAtom(2,1,0)\n",
    "print(h1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we have created an `HAtom` named `h1` and specified its quantum numbers. When calling `print`, python looks in the class for a `__str__` function, and if it doesn't find one, it uses `__repr__`. We implemented only `__repr__` so that the interpreter prints the string when it encounters an `HAtom` object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Cartesian coordinates in the HAtom class\n",
    "\n",
    "The plots we made in the previous notebook are fine, but it is sometimes difficult to visualize how a slice in $(r,\\phi)$ space maps onto normal $(x,y,z)$ space. It would be more convenient to make our graphs in Cartesian coordinates, while still using spherical coordinates to calculate the orbital. We'll write a function called `to_spherical` that converts Cartesian coordinates to spherical coordinates, and then write a new function in the `HAtom` class that makes use of it.\n",
    "\n",
    "Normally, to modify the class you would just return to the earlier cell, make changes, and then re-run it. But for instructional purposes, in this notebook we'll make new versions of the class as we go. An important point to note about class definitions in Jupyter notebooks is that if you modify a class, any objects of that class that you have already created still refer to the **old version** of the class. You must recreate the object to use the updated class definition, as shown below.\n",
    "\n",
    "Finally, just for ease of future use, we will include all of the import statements that are needed for the class in the same cell as the definition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import scipy.special as sps\n",
    "from typing import Tuple\n",
    "\n",
    "def to_spherical(x : np.ndarray, y : np.ndarray, z : np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:\n",
    "    \"\"\"Converts the Cartesian coordinates x,y,z to r,theta,phi\n",
    "\n",
    "    This function expects x, y, and z to be broadcastable 3D arrays,\n",
    "    such as those generated by np.meshgrid. The output arrays are\n",
    "    always dense 3D grids, and correspond to r, theta, and phi\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x : np.ndarray\n",
    "        Cartesian x values, in a.u.\n",
    "    x : np.ndarray\n",
    "        Cartesian y values, in a.u.\n",
    "    x : np.ndarray\n",
    "        Cartesian z values, in a.u.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    np.ndarray\n",
    "        3D grid of r values\n",
    "    np.ndarray\n",
    "        3D grid of theta values\n",
    "    np.ndarray\n",
    "        3D grid of phi values\n",
    "    \"\"\"\n",
    "\n",
    "    r = ( x**2. + y**2. + z**2. )**0.5\n",
    "    theta = np.arccos(z/r)\n",
    "    phi = np.arctan(y/x)\n",
    "\n",
    "    return r, theta, phi\n",
    "\n",
    "class HAtom:\n",
    "    \"\"\" Class for representing a hydrogen atom\n",
    "    \n",
    "    Here we can provide examples of how to use the class, and list the\n",
    "    important functions a user of the class may want to access. Since this\n",
    "    is a tutorial, we will omit the details here and show it in code.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, n : int, l : int, m : int):\n",
    "        \"\"\" Constructor: sets the n, l, and m quantum numbers \"\"\"\n",
    "        \n",
    "        self.n = n\n",
    "        self.l = l\n",
    "        self.m = m\n",
    "    \n",
    "    def orbital_xyz(self, x : np.ndarray, y : np.ndarray, z : np.ndarray) -> np.ndarray:\n",
    "        \"\"\"Calculates the hydrogen orbital on the x,y,z grid\n",
    "        \n",
    "        The x, y, and z values are converted to spherical coordinates,\n",
    "        then evaluated\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        x : np.ndarray\n",
    "            Cartesian x values, in a.u.\n",
    "        x : np.ndarray\n",
    "            Cartesian y values, in a.u.\n",
    "        x : np.ndarray\n",
    "            Cartesian z values, in a.u.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        np.ndarray\n",
    "            3D array containing hydrogen atomic orbital evaluated at x, y, z\n",
    "        \"\"\"\n",
    "        \n",
    "        r, theta, phi = to_spherical(x, y, z)\n",
    "        return self.orbital(r, theta, phi)\n",
    "        \n",
    "    def orbital(self, r : np.ndarray, theta : np.ndarray, phi : np.ndarray) -> np.ndarray:\n",
    "        \"\"\"Calculates the hydrogen orbital on the $r$, $\\theta$, $\\phi$ grid\n",
    "        \n",
    "        Note: Here we could add a detailed description of the how the orbital\n",
    "        is represented, noting that the angular convention we're using is opposite\n",
    "        that of sp.special.sph_harm.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        r : np.ndarray\n",
    "            $r$ values, in atomic units\n",
    "        theta : np.ndarray\n",
    "            colatitudinal angle in radians, in the range [0,pi)\n",
    "        phi : np.ndarray\n",
    "            azimuthal angle in radians, in the range [0,2pi]\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        np.ndarray\n",
    "            3D array containing hydrogen atomic orbital evaluated at r, theta, phi\n",
    "        \"\"\"\n",
    "        \n",
    "        pf = ( (2./self.n)**3. * sps.factorial(self.n-self.l-1) / (2. * self.n * sps.factorial(self.n+self.l)) )**0.5\n",
    "        radial = np.exp(-r/2.) * r**self.l * sps.eval_genlaguerre(self.n-self.l-1, 2*self.l+1, r)\n",
    "        angular = sps.sph_harm(self.m, self.l, phi, theta)\n",
    "        return pf * radial * angular\n",
    "    \n",
    "    def __repr__(self):\n",
    "        return f'{type(self).__name__}(n={self.n}, l={self.l}, m={self.m})'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We've added the `to_spherical` function into the global namespace because it doesn't depend on anything in the `HAtom` class, then added `HAtom.orbital_xyz` that makes use of that function to do the conversion. Let's try it out. Note that since we have redefined the class, currently our `h1` object is still pointing to the old class definition and does not have an `orbital_xyz`function. We need to recreate the object now that the class definition has changed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xyz = np.linspace(-10,10,51)\n",
    "x, y, z = np.meshgrid(xyz,xyz,xyz,sparse=True)\n",
    "\n",
    "# This throws an exception because orbital_xyz wasn't in the class\n",
    "# definition when h1 was created\n",
    "print(h1)\n",
    "try:\n",
    "    psi = h1.orbital_xyz(x,y,z)\n",
    "except Exception as e:\n",
    "    print(f'{type(e).__name__}: {e}')\n",
    "    \n",
    "h1 = HAtom(2,1,1)\n",
    "print(h1)\n",
    "\n",
    "try:\n",
    "    psi = h1.orbital_xyz(x,y,z)\n",
    "except Exception as e:\n",
    "    print(f'{type(e).__name__}: {e}')\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code completed without raising an exception, but there are several warnings about invalid values and divide by zero! If you already know why, good. Here we'll take a moment to figure out how to track this down. The RuntimeWarning is generated by `numpy`, which has an error handler that is invoked when problems with floating-point operations happen. A warning like this allows the program to continue executing, but it prints the warning message to alert you to the fact that the results may not match what you expect.\n",
    "\n",
    "You can view and change numpy's error handling with [`numpy.geterr`](https://numpy.org/doc/stable/reference/generated/numpy.geterr.html) and [`numpy.seterr`](https://numpy.org/doc/stable/reference/generated/numpy.seterr.html). We can force numpy to halt the code when it encounters divide by zero by calling `np.seterr(all='raise')`. This tells us exactly where the problem is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [],
   "source": [
    "print(np.geterr())\n",
    "\n",
    "old_settings = np.seterr(all='raise')\n",
    "try:\n",
    "    psi = h1.orbital_xyz(x,y,z)\n",
    "finally:\n",
    "    np.seterr(**old_settings) # restore original settings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you may have supected, we have a problem at the origin. When `x=0`, `y=0`, and `z=0`, there are issues with the coordinate system because $\\theta$ and $\\phi$ are ill-defined, and the formulas we've used involve division by `r` and `x`. Let's take a look at some plots of the orbital we've tried to make using [`Axes.imshow`](https://matplotlib.org/api/_as_gen/matplotlib.axes.Axes.imshow.html), which is designed to make 2D plots with evenly spaced data along the two axes. Note in the code below the use of [`pyplot.subplots`](https://matplotlib.org/api/_as_gen/matplotlib.pyplot.subplots.html) to make the figure and a 2D array of `Axes` objects, which is then flattened and zipped with the functions and labels to create the plots in a loop. You can change `z_index` and rerun the cell to take a different cross section through the orbital, and explore how the [`imshow` interpolation methods](https://matplotlib.org/gallery/images_contours_and_fields/interpolation_methods.html) change the images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.cm\n",
    "\n",
    "fig, axes = plt.subplots(2,2, figsize=(12,12))\n",
    "\n",
    "funcs = [np.real,np.imag,np.absolute,np.angle]\n",
    "labels = ['Real','Imag','Mag','Phase']\n",
    "z_index = 25\n",
    "\n",
    "for ax,f,label in zip(axes.flatten(),funcs,labels):\n",
    "    d = f(psi[:,:,z_index])\n",
    "    \n",
    "    #this line ensures that the color scales are symmetric around 0\n",
    "    #When used with the RdBu color map, this ensures that blue is positive,\n",
    "    #red is negative, and white is 0.\n",
    "    #Becasue there are NaNs where there was division by 0, use nanmax\n",
    "    norm = matplotlib.cm.colors.Normalize(vmax=np.nanmax(np.absolute(d)), vmin=-np.nanmax(np.absolute(d)))\n",
    "    \n",
    "    \n",
    "    im = ax.imshow(d,cmap='RdBu',origin='lower',norm=norm,\n",
    "                   interpolation='none',extent=[xyz[0],xyz[-1],xyz[0],xyz[-1]])\n",
    "    cb = fig.colorbar(im,ax=ax)\n",
    "    cb.set_label(label)\n",
    "    ax.set_title(f'z = {xyz[z_index]}')\n",
    "    ax.set_xlabel('x ($a_0$)')\n",
    "    ax.set_ylabel('y ($a_0$)')\n",
    "    \n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This visualization lets us see that there are is a problem not only at $r=0$, but also at $x=0$: the sign of the imaginary part is abruptly changing, which we can also see in the phase plot. Note that the apparent discontinuity at $y=0$ in the phase plot corresponds to a change from $-\\pi \\to \\pi$, which is the normal expected behavior for a function that is $2\\pi$-periodic.\n",
    "\n",
    "To fix our spherical coordinate conversion, we should just set $\\theta = 0$ when $r = 0$, and we need to be careful at $x=0$ when calculating $\\phi$. For the first problem, you might think that we can calculate `r`, then use `if np.isclose(r,0.0): theta = 0.0`. However, that won't work because `theta` needs to be an array, not a single number. Instead, we can use the [`np.where`] function to accomplish the same thing, as shown in the code below.\n",
    "\n",
    "The $\\phi$ problem is so common that a special version of the arctangent function has been created to deal with it: the two-argument tangent function [`np.arctan2`](https://numpy.org/doc/stable/reference/generated/numpy.arctan2.html). With these corrections in place, the new `to_spherical` function becomes the following."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "raises-exception"
    ]
   },
   "outputs": [],
   "source": [
    "def to_spherical(x : np.ndarray, y : np.ndarray, z : np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:\n",
    "    \"\"\"Converts the Cartesian coordinates x,y,z to r,theta,phi\n",
    "\n",
    "    This function expects x, y, and z to be broadcastable 3D arrays,\n",
    "    such as those generated by np.meshgrid. The output arrays are\n",
    "    always dense 3D grids, and correspond to r, theta, and phi\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x : np.ndarray\n",
    "        Cartesian x values, in a.u.\n",
    "    x : np.ndarray\n",
    "        Cartesian y values, in a.u.\n",
    "    x : np.ndarray\n",
    "        Cartesian z values, in a.u.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    np.ndarray\n",
    "        3D grid of r values\n",
    "    np.ndarray\n",
    "        3D grid of theta values\n",
    "    np.ndarray\n",
    "        3D grid of phi values\n",
    "    \"\"\"\n",
    "\n",
    "    r = ( x**2. + y**2. + z**2. )**0.5\n",
    "    theta = np.where(np.isclose(r,0.0),np.zeros_like(r),np.arccos(z/r))\n",
    "    phi = np.arctan2(y,x)\n",
    "\n",
    "    return r, theta, phi\n",
    "\n",
    "\n",
    "psi = h1.orbital_xyz(x,y,z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You might be surprised to see that we're still getting a warning message about an invalid value in the divide function. This is due to the way that the `np.where` function works: first both `np.zeros_like` and `np.arccos` are evaluated at **all** values of `r` to build 2 arrays (we'll call them `a` and `b`). Then, `np.where` evaluates the condition `np.isclose(r,0.0)` for each value of `r`, and chooses the value from array `a` if the condition is true or array `b` if it is false. Because `np.arccos` is evaluated at `r=0`, there is a divide by zero, so a warning is brought up. Since we know that this is unavoidable, we can suppress the warning ising the python [`warnings`](https://docs.python.org/3/library/warnings.html) module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "\n",
    "def to_spherical(x : np.ndarray, y : np.ndarray, z : np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:\n",
    "    \"\"\"Converts the Cartesian coordinates x,y,z to r,theta,phi\n",
    "\n",
    "    This function expects x, y, and z to be broadcastable 3D arrays,\n",
    "    such as those generated by np.meshgrid. The output arrays are\n",
    "    always dense 3D grids, and correspond to r, theta, and phi\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x : np.ndarray\n",
    "        Cartesian x values, in a.u.\n",
    "    x : np.ndarray\n",
    "        Cartesian y values, in a.u.\n",
    "    x : np.ndarray\n",
    "        Cartesian z values, in a.u.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    np.ndarray\n",
    "        3D grid of r values\n",
    "    np.ndarray\n",
    "        3D grid of theta values\n",
    "    np.ndarray\n",
    "        3D grid of phi values\n",
    "    \"\"\"\n",
    "\n",
    "    r = ( x**2. + y**2. + z**2. )**0.5\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.simplefilter('ignore')\n",
    "        theta = np.where(np.isclose(r,0.0),np.zeros_like(r),np.arccos(z/r))\n",
    "    phi = np.arctan2(y,x)\n",
    "\n",
    "    return r, theta, phi\n",
    "\n",
    "\n",
    "psi = h1.orbital_xyz(x,y,z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the warning is suppressed. To make sure that this works properly, lets replot the orbital using the code above. (here you could just re-execute the earlier cell, but to make the notebook sequential that cell is copied below and edited to remove the code that handled NaNs). We can see that our earlier discontinuities are gone. The magnitude plot correctly shows that the wavefunction is 0 at the origin, and we see that the phase varies smoothly with $\\phi$, which is the angle measured in the $xy$ plane."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.cm\n",
    "\n",
    "fig, axes = plt.subplots(2,2, figsize=(12,12))\n",
    "\n",
    "funcs = [np.real,np.imag,np.absolute,np.angle]\n",
    "labels = ['Real','Imag','Mag','Phase']\n",
    "z_index = 25\n",
    "\n",
    "for ax,f,label in zip(axes.flatten(),funcs,labels):\n",
    "    d = f(psi[:,:,z_index])\n",
    "    norm = matplotlib.cm.colors.Normalize(vmax=abs(d).max(), vmin=-abs(d).max())\n",
    "    im = ax.imshow(d,cmap='RdBu',origin='lower',norm=norm,\n",
    "                   interpolation='none',extent=[xyz[0],xyz[-1],xyz[0],xyz[-1]])\n",
    "    cb = fig.colorbar(im,ax=ax)\n",
    "    cb.set_label(label)\n",
    "    ax.set_title(f'z = {xyz[z_index]}')\n",
    "    ax.set_xlabel('x ($a_0$)')\n",
    "    ax.set_ylabel('y ($a_0$)')\n",
    "    \n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Changing the atom's position and orbital\n",
    "\n",
    "So far we've assumed that the atom is located at the origin of our coordinate system. Let's change that so that we can visualize orbitals from multiple hydrogen atoms at different places, and be able to look at different orbitals from the same atom more easily. This will highlight a key advantage of using classes.\n",
    "\n",
    "We can add `x0`, `y0`, and `z0` variables to the `HAtom` class so that each atom keeps track of its own location. Then, when we convert `x`, `y`, and `z` to spherical coordinates, we need to subtract the atom's location coordinate from the Cartesian coordinate, like `x_ = (x-x0)`. This means that the `to_spherical` function needs to know `x0`, `y0`, and `z0`. We can either modify the function so that it takes 6 arguments or, better, make the function a part of the class itself, which we'll do here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import scipy.special as sps\n",
    "from typing import Tuple\n",
    "import warnings\n",
    "\n",
    "class HAtom:\n",
    "    \"\"\" Class for representing a hydrogen atom\n",
    "    \n",
    "    Here we can provide examples of how to use the class, and list the\n",
    "    important functions a user of the class may want to access. Since this\n",
    "    is a tutorial, we will omit the details here and show it in code.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, n : int, l : int, m : int, x0 : float = 0.0, y0 : float = 0.0, z0 : float = 0.0):\n",
    "        \"\"\" Constructor: sets the quantum numbers and position\n",
    "        \n",
    "        Quantum numbers are required. If position is not specified,\n",
    "        defaults to (0.0,0.0,0.0)\n",
    "        \"\"\"\n",
    "        \n",
    "        self.set_orbital(n,l,m)\n",
    "        self.set_position(x0,y0,z0)\n",
    "        \n",
    "    def set_orbital(self, n : int, l : int, m : int) -> None:\n",
    "        \"\"\" Sets quantum numbers \"\"\"\n",
    "        \n",
    "        self.n = n\n",
    "        self.l = l\n",
    "        self.m = m\n",
    "        \n",
    "    def set_position(self, x0 : float, y0 : float, z0 : float) -> None:\n",
    "        \"\"\" Sets position in a.u\"\"\"\n",
    "        \n",
    "        self.x0 = x0\n",
    "        self.y0 = y0\n",
    "        self.z0 = z0\n",
    "        \n",
    "    def to_spherical(self, x : np.ndarray, y : np.ndarray, z : np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:\n",
    "        \"\"\"Converts the Cartesian coordinates x,y,z to r,theta,phi\n",
    "\n",
    "        This function expects x, y, and z to be broadcastable 3D arrays,\n",
    "        such as those generated by np.meshgrid. The output arrays are\n",
    "        always dense 3D grids, and correspond to r, theta, and phi. The\n",
    "        atom's center position is taken into account.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        x : np.ndarray\n",
    "            Cartesian x values, in a.u.\n",
    "        x : np.ndarray\n",
    "            Cartesian y values, in a.u.\n",
    "        x : np.ndarray\n",
    "            Cartesian z values, in a.u.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        np.ndarray\n",
    "            3D grid of r values\n",
    "        np.ndarray\n",
    "            3D grid of theta values\n",
    "        np.ndarray\n",
    "            3D grid of phi values\n",
    "        \"\"\"\n",
    "        \n",
    "        x_ = x-self.x0\n",
    "        y_ = y-self.y0\n",
    "        z_ = z-self.z0\n",
    "        \n",
    "        r = ( x_**2. + y_**2. + z_**2. )**0.5\n",
    "        with warnings.catch_warnings():\n",
    "            warnings.simplefilter('ignore')\n",
    "            theta = np.where(np.isclose(r,0.0),np.zeros_like(r),np.arccos(z_/r))\n",
    "        phi = np.arctan2(y_,x_)\n",
    "\n",
    "        return r, theta, phi\n",
    "    \n",
    "    def orbital_xyz(self, x : np.ndarray, y : np.ndarray, z : np.ndarray) -> np.ndarray:\n",
    "        \"\"\"Calculates the hydrogen orbital on the x,y,z grid\n",
    "        \n",
    "        The x, y, and z values are converted to spherical coordinates,\n",
    "        then evaluated\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        x : np.ndarray\n",
    "            Cartesian x values, in a.u.\n",
    "        x : np.ndarray\n",
    "            Cartesian y values, in a.u.\n",
    "        x : np.ndarray\n",
    "            Cartesian z values, in a.u.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        np.ndarray\n",
    "            3D array containing hydrogen atomic orbital evaluated at x, y, z\n",
    "        \"\"\"\n",
    "        \n",
    "        r, theta, phi = self.to_spherical(x, y, z)\n",
    "        return self.orbital(r, theta, phi)\n",
    "        \n",
    "    def orbital(self, r : np.ndarray, theta : np.ndarray, phi : np.ndarray) -> np.ndarray:\n",
    "        \"\"\"Calculates the hydrogen orbital on the $r$, $\\theta$, $\\phi$ grid\n",
    "        \n",
    "        Note: Here we could add a detailed description of the how the orbital\n",
    "        is represented, noting that the angular convention we're using is opposite\n",
    "        that of sp.special.sph_harm.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        r : np.ndarray\n",
    "            $r$ values, in atomic units\n",
    "        theta : np.ndarray\n",
    "            colatitudinal angle in radians, in the range [0,pi)\n",
    "        phi : np.ndarray\n",
    "            azimuthal angle in radians, in the range [0,2pi]\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        np.ndarray\n",
    "            3D array containing hydrogen atomic orbital evaluated at r, theta, phi\n",
    "        \"\"\"\n",
    "        \n",
    "        pf = ( (2./self.n)**3. * sps.factorial(self.n-self.l-1) / (2. * self.n * sps.factorial(self.n+self.l)) )**0.5\n",
    "        radial = np.exp(-r/2.) * r**self.l * sps.eval_genlaguerre(self.n-self.l-1, 2*self.l+1, r)\n",
    "        \n",
    "        return pf * radial * self.angular(theta,phi)\n",
    "    \n",
    "    def angular(self, theta : np.ndarray, phi : np.ndarray) -> np.ndarray:\n",
    "        \"\"\"Calculates the angular part of the hydrogen wavefuntion\"\"\"\n",
    "        \n",
    "        return sps.sph_harm(self.m, self.l, phi, theta)\n",
    "        \n",
    "    \n",
    "    def __repr__(self):\n",
    "        return f'{type(self).__name__}(n={self.n}, l={self.l}, m={self.m}) at (x,y,z) = ({self.x0},{self.y0},{self.z0})'\n",
    "    \n",
    "h1 = HAtom(1,0,0)\n",
    "\n",
    "print(h1)\n",
    "\n",
    "h1.set_orbital(2,1,1)\n",
    "h1.set_position(3.0,-1.5,0.0)\n",
    "psi = h1.orbital_xyz(x,y,z)\n",
    "\n",
    "print(h1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We've made a few modifications to the class:\n",
    "\n",
    "- `to_spherical` is now part of the class. Its function signature now has the first argument `self`, which refers to the instance of the class, and allows the function to access the object's `x0`, `y0`, and `z0` values.\n",
    "- `__init__` now allows optional argupemts that specify `x0`, `y0`, and `z0`\n",
    "- A new `set_orbital` function makes it easy to change all the quantum numbers at once, instead of having to do `h1.n = 2`, `h1.l = 1`, `h1.m = m`\n",
    "- A new `set_position` function does the same for `x0`, `y0`, and `z0`\n",
    "- The `orbital_xyz` function calls `self.to_spherical` instead of the `to_spherical` function in the global namespace\n",
    "- Moved the calculation of the angular part of the orbital to its own function named `angular`, which will be useful in our next step.\n",
    "- The `__repr__` function now prints the atom's position in addition to its quantum numbers\n",
    "\n",
    "We can verify that the class is working by plotting the wavefunction on the same grid as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.cm\n",
    "\n",
    "fig, axes = plt.subplots(2,2, figsize=(12,12))\n",
    "\n",
    "funcs = [np.real,np.imag,np.absolute,np.angle]\n",
    "labels = ['Real','Imag','Mag','Phase']\n",
    "z_index = 25\n",
    "\n",
    "for ax,f,label in zip(axes.flatten(),funcs,labels):\n",
    "    d = f(psi[:,:,z_index])\n",
    "    norm = matplotlib.cm.colors.Normalize(vmax=abs(d).max(), vmin=-abs(d).max())\n",
    "    im = ax.imshow(d,cmap='RdBu',origin='lower',norm=norm,\n",
    "                   interpolation='none',extent=[xyz[0],xyz[-1],xyz[0],xyz[-1]])\n",
    "    cb = fig.colorbar(im,ax=ax)\n",
    "    cb.set_label(label)\n",
    "    ax.set_title(f'z = {xyz[z_index]}')\n",
    "    ax.set_xlabel('x ($a_0$)')\n",
    "    ax.set_ylabel('y ($a_0$)')\n",
    "    \n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Real forms of the orbitals\n",
    "\n",
    "We now have the orbitals in Cartesian space, but they're still complex. The normal forms of the orbitals that we use in chemistry are converted into real numbers by manipulating the angular portion as shown in the table below. The resultant real orbitals are completely equivalent to the complex ones.\n",
    "\n",
    "| m           |                           Real Form |\n",
    "| --- | --- |\n",
    "| $m < 0$ | $(-1)^m \\sqrt{2}\\, \\text{Im}\\left(Y_{\\ell}^{|m|}\\right) $ |\n",
    "| $m = 0$ | $Y_\\ell^0 $ |\n",
    "| $m > 0$ | $(-1)^m\\sqrt{2}\\, \\text{Re}\\left(Y_\\ell^{\\phantom{|}m}\\right) $ |\n",
    "\n",
    "We'll add this by creating a subclass of `HAtom` that instead calculates the real form of the wavefunction. Just so that the code is all together, the `HAtom` class is copied again in the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import scipy.special as sps\n",
    "from typing import Tuple\n",
    "import warnings\n",
    "\n",
    "class HAtom:\n",
    "    \"\"\" Class for representing a hydrogen atom\n",
    "    \n",
    "    Here we can provide examples of how to use the class, and list the\n",
    "    important functions a user of the class may want to access. Since this\n",
    "    is a tutorial, we will omit the details here and show it in code.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, n : int, l : int, m : int, x0 : float = 0.0, y0 : float = 0.0, z0 : float = 0.0):\n",
    "        \"\"\" Constructor: sets the quantum numbers and position\n",
    "        \n",
    "        Quantum numbers are required. If position is not specified,\n",
    "        defaults to (0.0,0.0,0.0)\n",
    "        \"\"\"\n",
    "        \n",
    "        self.set_orbital(n,l,m)\n",
    "        self.set_position(x0,y0,z0)\n",
    "        \n",
    "    def set_orbital(self, n : int, l : int, m : int) -> None:\n",
    "        \"\"\" Sets quantum numbers \"\"\"\n",
    "        \n",
    "        self.n = n\n",
    "        self.l = l\n",
    "        self.m = m\n",
    "        \n",
    "    def set_position(self, x0 : float, y0 : float, z0 : float) -> None:\n",
    "        \"\"\" Sets position in a.u\"\"\"\n",
    "        \n",
    "        self.x0 = x0\n",
    "        self.y0 = y0\n",
    "        self.z0 = z0\n",
    "        \n",
    "    def to_spherical(self, x : np.ndarray, y : np.ndarray, z : np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:\n",
    "        \"\"\"Converts the Cartesian coordinates x,y,z to r,theta,phi\n",
    "\n",
    "        This function expects x, y, and z to be broadcastable 3D arrays,\n",
    "        such as those generated by np.meshgrid. The output arrays are\n",
    "        always dense 3D grids, and correspond to r, theta, and phi. The\n",
    "        atom's center position is taken into account.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        x : np.ndarray\n",
    "            Cartesian x values, in a.u.\n",
    "        x : np.ndarray\n",
    "            Cartesian y values, in a.u.\n",
    "        x : np.ndarray\n",
    "            Cartesian z values, in a.u.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        np.ndarray\n",
    "            3D grid of r values\n",
    "        np.ndarray\n",
    "            3D grid of theta values\n",
    "        np.ndarray\n",
    "            3D grid of phi values\n",
    "        \"\"\"\n",
    "        \n",
    "        x_ = x-self.x0\n",
    "        y_ = y-self.y0\n",
    "        z_ = z-self.z0\n",
    "        \n",
    "        r = ( x_**2. + y_**2. + z_**2. )**0.5\n",
    "        with warnings.catch_warnings():\n",
    "            warnings.simplefilter('ignore')\n",
    "            theta = np.where(np.isclose(r,0.0),np.zeros_like(r),np.arccos(z_/r))\n",
    "        phi = np.arctan2(y_,x_)\n",
    "\n",
    "        return r, theta, phi\n",
    "    \n",
    "    def orbital_xyz(self, x : np.ndarray, y : np.ndarray, z : np.ndarray) -> np.ndarray:\n",
    "        \"\"\"Calculates the hydrogen orbital on the x,y,z grid\n",
    "        \n",
    "        The x, y, and z values are converted to spherical coordinates,\n",
    "        then evaluated\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        x : np.ndarray\n",
    "            Cartesian x values, in a.u.\n",
    "        x : np.ndarray\n",
    "            Cartesian y values, in a.u.\n",
    "        x : np.ndarray\n",
    "            Cartesian z values, in a.u.\n",
    "            \n",
    "        Returns\n",
    "        -------\n",
    "        np.ndarray\n",
    "            3D array containing hydrogen atomic orbital evaluated at x, y, z\n",
    "        \"\"\"\n",
    "        \n",
    "        r, theta, phi = self.to_spherical(x, y, z)\n",
    "        return self.orbital(r, theta, phi)\n",
    "        \n",
    "    def orbital(self, r : np.ndarray, theta : np.ndarray, phi : np.ndarray) -> np.ndarray:\n",
    "        \"\"\"Calculates the hydrogen orbital on the $r$, $\\theta$, $\\phi$ grid\n",
    "        \n",
    "        Note: Here we could add a detailed description of the how the orbital\n",
    "        is represented, noting that the angular convention we're using is opposite\n",
    "        that of sp.special.sph_harm.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        r : np.ndarray\n",
    "            $r$ values, in atomic units\n",
    "        theta : np.ndarray\n",
    "            colatitudinal angle in radians, in the range [0,pi)\n",
    "        phi : np.ndarray\n",
    "            azimuthal angle in radians, in the range [0,2pi]\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        np.ndarray\n",
    "            3D array containing hydrogen atomic orbital evaluated at r, theta, phi\n",
    "        \"\"\"\n",
    "        \n",
    "        pf = ( (2./self.n)**3. * sps.factorial(self.n-self.l-1) / (2. * self.n * sps.factorial(self.n+self.l)) )**0.5\n",
    "        radial = np.exp(-r/2.) * r**self.l * sps.eval_genlaguerre(self.n-self.l-1, 2*self.l+1, r)\n",
    "        \n",
    "        return pf * radial * self.angular(theta,phi)\n",
    "    \n",
    "    def angular(self, theta : np.ndarray, phi : np.ndarray) -> np.ndarray:\n",
    "        \"\"\"Calculates the angular part of the hydrogen wavefuntion\"\"\"\n",
    "        \n",
    "        return sps.sph_harm(self.m, self.l, phi, theta)\n",
    "        \n",
    "    \n",
    "    def __repr__(self):\n",
    "        return f'{type(self).__name__}(n={self.n}, l={self.l}, m={self.m}) at (x,y,z) = ({self.x0},{self.y0},{self.z0})'\n",
    "    \n",
    "class HAtomReal(HAtom):\n",
    "    \"\"\" HAtom that calculates real orbitals \"\"\"\n",
    "    \n",
    "    def angular(self, theta : np.ndarray, phi : np.ndarray) -> np.ndarray:\n",
    "        \"\"\" Overrides HAtom.angular to calculate real spherical harmonics \"\"\"\n",
    "        \n",
    "        if self.m > 0:\n",
    "            return (-1)**self.m * np.sqrt(2.) * np.real(sps.sph_harm(self.m, self.l, phi, theta))\n",
    "        elif self.m < 0:\n",
    "            return (-1)**self.m * np.sqrt(2.) * np.imag(sps.sph_harm(-self.m, self.l, phi, theta))\n",
    "        else:\n",
    "            return np.real(super().angular(theta,phi))\n",
    "    \n",
    "    \n",
    "h1 = HAtomReal(2,1,1)\n",
    "print(h1)\n",
    "psi = h1.orbital_xyz(x,y,z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The real forms of the orbitals are those you are familiar with from general chemistry. The code below generates a 3x3 grid of plots taking different cross sections of the $2p$ orbitals. Each column is one of the orbitals, and each row corresponds to a slice in either the $xy$, $xz$, or $yz$ planes. From the plots, you can see that $m=-1$ is now the $p_y$ orbital, $m=0$ is the $p_z$ orbital, and $m=1$ is the $p_x$ orbital. Feel free to make adjustments to customize the plots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.cm\n",
    "\n",
    "fig, axes = plt.subplots(3,3, figsize=(12,12))\n",
    "\n",
    "index = 25\n",
    "n = 2\n",
    "interp = 'none'\n",
    "\n",
    "#Note: when imshow gets a 2D array, it interprets the first axis as row data (i.e., Y axis)\n",
    "#and the second as column data (X axis). Because we used the default meshgrid, our data\n",
    "#are organized such that axis 0 = y, axis 1 = x, and axis 2 = z.\n",
    "#When taking slices, the following images are made:\n",
    "#[:,:,a] : y vs x, z = a\n",
    "#[:,a,:] : y vs z, x = a\n",
    "#[a,:,:] : x vs z, y = a\n",
    "\n",
    "#if z vs z is desired, then take the transpose: [a,:,:].T, as done in the third loop below\n",
    "\n",
    "print(f'x shape: {x.shape}, y shape: {y.shape}, z shape: {z.shape}')\n",
    "\n",
    "for ax,m in zip(axes[0],[-1,0,1]):\n",
    "    h1.set_orbital(n,1,m)\n",
    "    d = h1.orbital_xyz(x,y,z)[:,:,index]\n",
    "    cmax = max(abs(d).max(),1e-7)\n",
    "    norm = matplotlib.cm.colors.Normalize(vmax=cmax, vmin=-cmax)\n",
    "    im = ax.imshow(d,cmap='RdBu',origin='lower',norm=norm,\n",
    "                   interpolation=interp,extent=[xyz[0],xyz[-1],xyz[0],xyz[-1]])\n",
    "    cb = fig.colorbar(im,ax=ax)\n",
    "    ax.set_title(f'(${n},1,{m}$), z = ${xyz[index]}$ $a_0$')\n",
    "    ax.set_xlabel('x ($a_0$)')\n",
    "    ax.set_ylabel('y ($a_0$)')\n",
    "    ax.set_xticks([-10,-5,0,5,10])\n",
    "    ax.set_yticks([-10,-5,0,5,10])\n",
    "    \n",
    "\n",
    "for ax,m in zip(axes[1],[-1,0,1]):\n",
    "    h1.set_orbital(n,1,m)\n",
    "    d = h1.orbital_xyz(x,y,z)[:,index,:]\n",
    "    cmax = max(abs(d).max(),1e-7)\n",
    "    norm = matplotlib.cm.colors.Normalize(vmax=cmax, vmin=-cmax)\n",
    "    im = ax.imshow(d,cmap='RdBu',origin='lower',norm=norm,\n",
    "                   interpolation=interp,extent=[xyz[0],xyz[-1],xyz[0],xyz[-1]])\n",
    "    cb = fig.colorbar(im,ax=ax)\n",
    "    ax.set_title(f'(${n},1,{m}$), x = ${xyz[index]}$ $a_0$')\n",
    "    ax.set_xlabel('z ($a_0$)')\n",
    "    ax.set_ylabel('y ($a_0$)')\n",
    "    ax.set_xticks([-10,-5,0,5,10])\n",
    "    ax.set_yticks([-10,-5,0,5,10])\n",
    "    \n",
    "for ax,m in zip(axes[2],[-1,0,1]):\n",
    "    h1.set_orbital(n,1,m)\n",
    "    d = h1.orbital_xyz(x,y,z)[index,:,:].T #transpose to get z vs x\n",
    "    cmax = max(abs(d).max(),1e-7)\n",
    "    norm = matplotlib.cm.colors.Normalize(vmax=cmax, vmin=-cmax)\n",
    "    im = ax.imshow(d,cmap='RdBu',origin='lower',norm=norm,\n",
    "                   interpolation=interp,extent=[xyz[0],xyz[-1],xyz[0],xyz[-1]])\n",
    "    cb = fig.colorbar(im,ax=ax)\n",
    "    ax.set_title(f'(${n},1,{m}$), y = ${xyz[index]}$ $a_0$')\n",
    "    ax.set_xlabel('x ($a_0$)')\n",
    "    ax.set_ylabel('z ($a_0$)')\n",
    "    ax.set_xticks([-10,-5,0,5,10])\n",
    "    ax.set_yticks([-10,-5,0,5,10])\n",
    "    \n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Intro to 3D plotting\n",
    "\n",
    "The `matplotlib` library has support for 3D plotting through the [`mplot3d` interface](https://matplotlib.org/mpl_toolkits/mplot3d/tutorial.html). It is primarily designed to generate 3D plots where $z = f(x,y)$; that is, a 3D visualization of 2D functions. It is not a complete 3D graphics/rendering solution. Our situation is that we have a true 3D function, since $\\psi = f(x,y,z)$. Typical visualizations of atomic orbitals are made by plotting [isosurfaces](https://en.wikipedia.org/wiki/Isosurface). Unfortunately, there is no one-line `matplotlib` plotting mechanism that can find isosurfaces, so in order to make such a plot we have to calculate the isosurface ourselves. The details are beyond the scope of the course, but some example code is shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.cm\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "from mpl_toolkits.mplot3d.art3d import Poly3DCollection\n",
    "import skimage.measure as skm\n",
    "\n",
    "xyz = np.linspace(-10,10,51)\n",
    "x,y,z = np.meshgrid(xyz,xyz,xyz,sparse=True)\n",
    "\n",
    "\n",
    "sp = abs(xyz[1]-xyz[0])\n",
    "center = (xyz[-1]-xyz[0])/2\n",
    "\n",
    "def add_isosurface_sym(ax,psi,level,spacing,alpha=0.5):\n",
    "    for ll in [-level,level]:\n",
    "        if float(ll) <= float(psi.min()) or float(ll) >= float(psi.max()):\n",
    "            continue\n",
    "        verts, faces, _, vals = skm.marching_cubes(psi,level=ll,spacing=(spacing,spacing,spacing))\n",
    "        verts -= np.asarray([[center,center,center]])\n",
    "        mesh = Poly3DCollection(verts[faces])\n",
    "        mesh.set_edgecolor('k')\n",
    "        mesh.set_linewidth(0.05)\n",
    "        if ll > 0:\n",
    "            mesh.set_facecolor('#0000ff')\n",
    "        else:\n",
    "            mesh.set_facecolor('#ff0000')\n",
    "        mesh.set_alpha(alpha)\n",
    "        ax.add_collection3d(mesh)\n",
    "\n",
    "fig = plt.figure(figsize=(13,4.2))\n",
    "for plot in [1,2,3]:\n",
    "    ax = fig.add_subplot(1,3,plot, projection='3d')\n",
    "\n",
    "    if plot == 1:\n",
    "        h1.set_orbital(3,1,0)\n",
    "        title = '$3p\\,z$'\n",
    "    elif plot == 2:\n",
    "        h1.set_orbital(3,2,0)\n",
    "        title = '$3d\\,z^2$'\n",
    "    else:\n",
    "        h1.set_orbital(4,3,0)\n",
    "        title = '$4f\\,z^3$'\n",
    "        \n",
    "    psi = h1.orbital_xyz(x,y,z)\n",
    "\n",
    "    l = psi.max()\n",
    "    for level in [0.25*l,0.5*l,0.75*l]:\n",
    "        add_isosurface_sym(ax,psi,level,sp,0.3)\n",
    "\n",
    "    ticks = [xyz[0], xyz[0]+(xyz[-1]-xyz[0])/4,xyz[0]+(xyz[-1]-xyz[0])/2,xyz[0]+3*(xyz[-1]-xyz[0])/4,xyz[-1]]\n",
    "    ax.set_xlim(-10,10)\n",
    "    ax.set_ylim(-10,10)\n",
    "    ax.set_zlim(-10,10)\n",
    "    ax.set_xlabel('x ($a_0$)')\n",
    "    ax.set_ylabel('y ($a_0$)')\n",
    "    ax.set_zlabel('z ($a_0$)')\n",
    "    ax.set_xticks(ticks)\n",
    "    ax.set_yticks(ticks)\n",
    "    ax.set_zticks(ticks)\n",
    "    ax.set_title(title)\n",
    "    \n",
    "fig.tight_layout()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Although `matplotlib` has functions for controlling the angle and orentation of the figure, often it is more convenient to use an interactive plot to orient the figure as desired. This notebook is using the `inline` backend of matplotlib which only generates static figures, but the `notebook` backend (on some platforms) will generate interactive plots that can be zoomed, panned, and reoriented.\n",
    "\n",
    "Another plotting utility that is even more convenient than `matplotlib` for isosurface plots is [`plotly`](https://plotly.com/python/), which generates interactive plots in the web browser. For graphs that you wish to be interactive, it is a good option. First you need to install `plotly`. Instructions are available at the link, but if you're using the Anaconda python distribution, you can just run `conda install plotly`. The code below sets up a basic isosurface plot using plotly; it's much simpler than the equivalent `matplotlib` code, and the interactive plot can be zoomed and panned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import plotly.graph_objects as go\n",
    "\n",
    "xyz = np.linspace(-10,10,51)\n",
    "x,y,z = np.meshgrid(xyz,xyz,xyz,sparse=False)\n",
    "\n",
    "h1 = HAtomReal(2,1,0)\n",
    "h1.set_position(0,0,0)\n",
    "psi1 = h1.orbital_xyz(x,y,z)\n",
    "psimax = abs(psi1).max()\n",
    "\n",
    "fig= go.Figure(data=go.Isosurface(\n",
    "    x=x.flatten(),\n",
    "    y=y.flatten(),\n",
    "    z=z.flatten(),\n",
    "    value=psi1.flatten(),\n",
    "    colorscale='RdBu',\n",
    "    isomin=-.75*psimax,\n",
    "    isomax=.75*psimax,\n",
    "    surface_count=6,\n",
    "    opacity=0.5,\n",
    "    caps=dict(x_show=False,y_show=False,z_show=False)\n",
    "))\n",
    "\n",
    "fig.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One approximation of a chemical bond is just the sum (or difference) of the orbitals on the atoms being bonded together. As a final example, we'll use `plotly` to visualize the overlap between a $2p_x$ orbital from one atom with the $3d_{xz}$ orbital of another atom. We'll do this by creating two `HAtomReal` objects. The first one is placed at $z=3$ and the second at $z=-3$. After calculating each orbital on the grid, we add the two together and plot. Here we'll also add markers showing the positions of the atoms, a line connecting them, and labels telling what the orbitals are."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import plotly.graph_objects as go\n",
    "import plotly.subplots\n",
    "\n",
    "xyz = np.linspace(-10,10,51)\n",
    "x,y,z = np.meshgrid(xyz,xyz,xyz,sparse=False)\n",
    "\n",
    "h1 = HAtomReal(2,1,1)\n",
    "h1.set_position(0,0,3)\n",
    "psi1 = h1.orbital_xyz(x,y,z)\n",
    "\n",
    "\n",
    "h2 = HAtomReal(3,2,1)\n",
    "h2.set_position(0,0,-3)\n",
    "psi2 = h2.orbital_xyz(x,y,z)\n",
    "\n",
    "psi = psi1 + psi2\n",
    "psimax = abs(psi).max()\n",
    "\n",
    "\n",
    "\n",
    "fig= go.Figure(data=go.Isosurface(\n",
    "    x=x.flatten(),\n",
    "    y=y.flatten(),\n",
    "    z=z.flatten(),\n",
    "    value=psi.flatten(),\n",
    "    colorscale='RdBu',\n",
    "    isomin=-.75*psimax,\n",
    "    isomax=.75*psimax,\n",
    "    surface_count=6,\n",
    "    opacity=0.5,\n",
    "    caps=dict(x_show=False,y_show=False,z_show=False)\n",
    "))\n",
    "\n",
    "fig.add_trace(go.Scatter3d(\n",
    "    x = [h1.x0,h2.x0],\n",
    "    y = [h1.y0,h2.y0],\n",
    "    z = [h1.z0,h2.z0],\n",
    "    marker=dict(color='black'),\n",
    "    line=dict(color='black',width=5)\n",
    "))\n",
    "\n",
    "fig.update_layout(\n",
    "    autosize=False,\n",
    "    width=800,\n",
    "    height=700,\n",
    "    margin=dict(l=20,r=20,b=20,t=20),\n",
    "    title=dict(\n",
    "        text=\"Orbital overlap\",\n",
    "        y=0.9,\n",
    "        x=0.5,\n",
    "        xanchor='center',\n",
    "        yanchor='top'\n",
    "    ),\n",
    "    scene=dict(\n",
    "        xaxis=dict(\n",
    "            title_text=r'x (a<sub>0</sub>)'\n",
    "        ),\n",
    "        yaxis=dict(\n",
    "            title_text=r'y (a<sub>0</sub>)'\n",
    "        ),\n",
    "        zaxis=dict(\n",
    "            title_text=r'z (a<sub>0</sub>)'\n",
    "        ),\n",
    "        camera=dict(\n",
    "            up=dict(x=0,y=0,z=1),\n",
    "            eye=dict(x=0.2,y=2.5,z=0.2)\n",
    "        ),\n",
    "        annotations = [dict(\n",
    "            showarrow=False,\n",
    "            x = h1.x0,\n",
    "            y = h1.y0,\n",
    "            z = h1.z0,\n",
    "            text = \"2px\",\n",
    "            xanchor='left',\n",
    "            xshift = 10,\n",
    "            opacity=1.0\n",
    "        ),dict(\n",
    "            showarrow=False,\n",
    "            x = h2.x0,\n",
    "            y = h2.y0,\n",
    "            z = h2.z0,\n",
    "            text = \"3dxz\",\n",
    "            xanchor='left',\n",
    "            xshift = 10,\n",
    "            opacity=1.0\n",
    "    )])\n",
    ")\n",
    "\n",
    "\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Key Takeaways\n",
    "\n",
    "- [`numpy.ndarray`](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html) objects are the preferred way to represent numerical data, and `numpy` versions of math functions should be used because they allow for fast, [vectorized](https://numpy.org/doc/stable/glossary.html#term-vectorization) code and support [broadcasting](https://numpy.org/doc/stable/user/basics.broadcasting.html).\n",
    "- When working with floating point numbers, be aware of numerical stability and artifacts, especially when comparing floating point numbers or when dividing them (division by 0!)\n",
    "- Classes and subclasses can be a useful way to make your code more modular, flexible, and reusable\n",
    "- [`matplotlib`](https://matplotlib.org/) is a standard library for visualizing data. It works by creating a [`Figure`](https://matplotlib.org/api/_as_gen/matplotlib.figure.Figure.html#matplotlib.figure.Figure) object and addint one or more [`Axes`](https://matplotlib.org/api/axes_api.html#matplotlib.axes.Axes) objects that represent individual graphs to display data. We looked at the following plotting methods\n",
    "    - [`Axes.plot`](https://matplotlib.org/api/_as_gen/matplotlib.axes.Axes.plot.html) makes line graphs, good for looking at a 1D function or a set of (x,y) data points. A related method is [`Axes.scatter`](https://matplotlib.org/api/_as_gen/matplotlib.axes.Axes.scatter.html).\n",
    "    - [`Axes.pcolormesh`](https://matplotlib.org/api/_as_gen/matplotlib.axes.Axes.pcolormesh.html) can show data of the form $f(x,y)$, and is useful when the data are on an irregualr grid. When the data are evenly spaced, [`Axes.imshow`](https://matplotlib.org/api/_as_gen/matplotlib.axes.Axes.imshow.html) is preferred.\n",
    "    - For 3D plotting, we looked briefly at the [`mplot3d` interface](https://matplotlib.org/mpl_toolkits/mplot3d/tutorial.html) and [`plotly`](https://plotly.com/python/), the latter of which is especially useful for making interactive plots on the web.\n",
    "\n",
    "Finally, though we did not show it in these notebooks, figures made by `matplotlib` can be saved to an image file using [`Figure.savefig`](https://matplotlib.org/3.2.2/api/_as_gen/matplotlib.figure.Figure.html#matplotlib.figure.Figure.savefig), which allows writing to a .png, .pdf, .eps., o .ps file (depending on the backend in use).\n",
    "\n"
   ]
  }
 ],
 "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": 4
}