{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": "#### Hydrogen Atomic Orbitals\n\nThis document is a Jupyter notebook. If this is the first time you've worked with one, please take a moment to briefly read [these instructions](https://jupyter-notebook.readthedocs.io/en/stable/notebook.html#structure-of-a-notebook-document) explaining how to use it. In particular, make sure you understand the difference between a Markdown cell used for text/images/equations (such as the current one) and a code cell, which executes python code. When you do problem sets in this course, you will be expected to submit Jupyter Notebook files that have a mix of explanatory text and working code.\n\nIn this notebook, we will be visualizing the atomic orbitals of the hydrogen atom and learning about code vectorization along the way. Before starting, you should have reviewed the [Week 2 Background](https://leeping.github.io/che155/weeks/2.html). The equation for a hydrogen atomic orbital is\n\n$$ \\psi_{n\\ell m}(r,\\theta,\\phi) = \\sqrt{ \\left( \\frac{2}{n} \\right)^3 \\frac{ (n-\\ell-1)! }{2n[(n+\\ell)!]^3} } e^{-r/n} \\left(\\frac{r}{n}\\right)^\\ell L^{2\\ell+1}_{n-\\ell-1}(r) Y_\\ell^m(\\theta,\\phi) $$\n\nusing spherical coordinates defined as\n\n
\n\"Diagram\n
\n\n|           Cartesian           |                     Spherical                               |           Domain                     |\n| --- | --- | --- |\n| $x = r \\sin \\theta \\cos \\phi$ | $r = \\sqrt{x^2 + y^2 + z^2}$ | $r \\in [0,\\infty)$ |\n| $y = r \\sin \\theta \\sin \\phi$ | $\\theta = \\arccos \\frac{z}{r} = \\arccos \\frac{z}{\\sqrt{x^2+y^2+z^2}}$ | $\\theta \\in [0,\\pi] $ |\n| $z = r \\cos \\theta $ | $\\phi = \\arctan \\frac{y}{x}$ | $\\phi \\in [0,2\\pi)$ |\n\n## Getting started\n\nLet's start by writing a Python function that will calculate the value of the wavefunction at a point in space. To do so, the function requires $r$, $\\theta$, $\\phi$, $n$, $\\ell$, and $m$. We also see that we need to evaluate factorials, a generalized Laguerre polynomial $L^{2\\ell+1}_{n-\\ell-1}(r)$, and a spherical harmonic function $Y_\\ell^m(\\theta,\\phi)$. In your earlier experience with python, you have probably encountered the [math.factorial](https://docs.python.org/3.7/library/math.html#math.factorial) function, but the standard python libraries do not have pre-existing implementations of $L^{2\\ell+1}_{n-\\ell-1}(r)$ and $Y_\\ell^m(\\theta,\\phi)$. These functions are, however, available in SciPy, so you can save a lot of time using a premade implementation rather than writing your own. In addition, the functions available in SciPy are often based on well-known and studied implementations, and are much more likely to be efficient and bug-free than one you would come up with yourself. Where possible, always try to make use of SciPy and NumPy, as we will do in this course.\n\nYou will want to become comfortable reading the [NumPy/SciPy documentation](https://docs.scipy.org/doc/). This is where you will find lists of the functions that are available, explanations of how the functions work, and example code. On the documentation page, you will find links to documentation for different versions of SciPy and NumPy, so you need to select the version you are using. If you're not sure, it is easy to check. Execute the cell below to see what version you are using." }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "import scipy as sp\n", "\n", "print(f'NumPy version {np.__version__}, SciPy version {sp.__version__}')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The functions we want are in the [`scipy.special` module](https://docs.scipy.org/doc/scipy/reference/tutorial/special.html), which contains [many other useful functions](https://docs.scipy.org/doc/scipy/reference/special.html#module-scipy.special) as well. In particular, we want:\n", "\n", "- [`scipy.special.factorial`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.factorial.html#scipy.special.factorial)\n", "- [`scipy.special.sph_harm`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.sph_harm.html#scipy.special.sph_harm)\n", "- [`scipy.special.eval_genlaguerre`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.eval_genlaguerre.html#scipy.special.eval_genlaguerre)\n", "\n", "With these in hand, we can write the definition of the hydrogen atomic orbital function. To save a bit of writing, we'll import the `scipy.special` module as `sps` to save some characters.\n", "\n", "It is **very** important to read the documentation of these functions carefully, as they will explain what the arguments mean. The generalized Laguerre function has the signature `eval_genlaguerre(l,alpha,X)`, and this corresponds to $L_n^\\alpha(X)$. Our hydrogen wavefunction involves has the term $L^{2\\ell+1}_{n-\\ell-1}(r)$, so for our purpose, `n` is $n-\\ell-1$, `alpha` is $2\\ell+1$, and `X` is $r$. For the spherical harmonics, we have $Y_l^m(\\theta,\\phi)$, and the function signature is `sph_harm(n,m,theta,phi)` for $Y_m^n(\\theta,\\phi)$. Checking the documentation, the definitions of all of the parameters are exactly backwards with our definitions, so`m` = $\\ell$, `n` = $m$, `theta`=$\\phi$ and `phi`=$\\theta$. Note that in the SciPy 1.15, `sph_harm` will be deprecated and replaced by `sph_harm_y`, which uses the standard scientific conventions. This is why it is always important to pay attention to the details of the function documentation." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import scipy.special as sps\n", "\n", "def h_orbital(r,theta,phi,n,l,m):\n", " pf = ( (2/n)**3. * sps.factorial(n-l-1) / (2 * n * sps.factorial(n+l)**3) )**0.5\n", " return pf * np.exp(-r/n) *(r/n)**l * sps.sph_harm(m,l,phi,theta) * sps.eval_genlaguerre(n-l-1,2*l+1,r)" ] }, { "cell_type": "markdown", "metadata": { "tags": [] }, "source": [ "#### Important: Writing efficient code\n", "\n", "At this point you may be wondering why we're using `sps.factorial` instead of `math.factorial`. What we'll soon see is that using the NumPy/SciPy versions of these functions, together with NumPy's [`ndarray` structure](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html) allows for us to write *faster* and *more efficient* code.\n", "\n", "To demonstrate this, we'll write some code using standard python data structures, and compare to the same code using a `numpy.ndarray`. We'll evaluate a $2p$ orbital with $m=0$ from 0 to 10 $a_0$ in steps of 0.1 $a_0$, holding $\\theta=0$ and $\\phi=0$." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "out = [] # create list to store output\n", "\n", "for i in range(0,101): #loop from 0-100; divide by 10 to get desired value\n", " out.append(h_orbital(i/10.,0,0,2,1,0))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "To do the same thing with an `ndarray`, we will first use the [`numpy.linspace`](https://numpy.org/doc/stable/reference/generated/numpy.linspace.html) function to create an array of $r$ values, and then we will call our `h_orbital` function with that array as the `r` argument. It may seem like this shouldn't work. After all, our function was written to expect a single value of `r`, and now we're giving it an array. However, all of NumPy and SciPy's functions are [universal functions](https://numpy.org/doc/stable/reference/ufuncs.html), and when they see an array as an argument, they *automatically* operate element-by-element on the array. Furthermore, because arrays are of fixed size and shape, the code is [vectorized](https://numpy.org/doc/stable/glossary.html#term-vectorization), and NumPy under the hood calls C-based implementations of loops that are much, much faster than python's `for` loops. This code accomplishes the same thing as before." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "r = np.linspace(0,10,101) #create an array of 101 points evenly spaced from 0 to 10 (inclusive)\n", "\n", "out2 = h_orbital(r,0,0,2,1,0)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "To test the efficiency of the code, we can use the ipython [`%timeit`](https://ipython.readthedocs.io/en/stable/interactive/magics.html#magic-timeit) magic command, which executes a cell over and over for several seconds and computes the average amount of time it takes for the code to complete. To compare the efficiency, we'll just see how long it takes to compute the function, not worrying about capturing the output." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%%timeit\n", "\n", "for i in range(0,101): #loop from 0-100; divide by 10 to get desired value\n", " h_orbital(i/10.,0,0,2,1,0)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%%timeit\n", "\n", "h_orbital(r,0,0,2,1,0)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "You can see that using the `ndarray` data structure to perform the vectorized computation is about 100 times faster in this test compared to a python `for` loop. Feel free to vary the sizes of the loops and compare! If we had used `math.factorial`, this vectorization would not work, because `math.factorial` is not a universal function. It can only take a single number as an argument, and will raise an error if you give it an array. Compare:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "test_x = np.linspace(1,10,10)\n", "\n", "print(sps.factorial(test_x))\n", "\n", "import math\n", "try:\n", " print(math.factorial(test_x))\n", "except TypeError as e:\n", " print(f'TypeError: {e}')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Vectorization may not seem like such a big deal right now. Both implementations run effectively instantaneously. However, remember that the hydrogen orbitals are 3D functions, so we need to vary $\\theta$ and $\\phi$ in addition to $r$. Let's do that comparison by calculating 101 values of $\\theta$ at each $r$, and 101 values of $\\phi$ at each $r$ and $\\theta$. The total number of function evaluations is 101x101x101=1030301. First, standard python." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%%timeit\n", "\n", "#Standard python -- this will take about 3-4 minutes!\n", "\n", "for r in range(0,101):\n", " for t in range(0,101):\n", " for p in range(0,101):\n", " h_orbital(r/10,t*np.pi/100,p*2*np.pi/100,2,1,0)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "To do the same calculation with NumPy's `ndarray` structures, we can take advantage of [broadcasting](https://numpy.org/doc/stable/user/basics.broadcasting.html). Broadcasting is how NumPy operations automatically loop over arrays. We were using this before to evaluate $\\psi$ at all values of $r$ in our 1D array. As explained in the documentation if we have 2 variables that we want to evaluate in a function `func(x,y)`, such as a set of 100 x coordinates and a set of 100 y coordinates, the behavior depends on how we set up the x and y arrays. If both arrays are 1D, that is, their shapes are both `(100,)` then the function is evaluated 100 times. The first time is at `x[0],y[0]`, the second at `x[1],y[1]`, and so on. The output in that case is a 1D array that is the same shape as both input arrays: `(100,)`.\n", "\n", "However, if we want to evaluate at every possible combination of `x` and `y`, we we need to make `x` and `y` each 2D arrays with different shapes. We've seen that when the input arrays have the same shape, the output has the same shape as the inputs. When NumPy encounters an operation with differently-shaped arrays, it follows broadcasting rules to make the shapes compatible if possible. Importantly, the broadcasting rules state that when comparing the axis size for arrays, they either must match exactly, or one of them must be size 1. If we set up `x` as a 2D array with 100 entries in the first dimension, and 1 entry in the second dimension, we have a 2D array with shape `(100,1)`. The only difference between this 2D array and our original 1D array is that the 2D array has a second axis of length 1. The data inside the 2D array is the exact same as the 1D array. According to the broadcasting rules, if we set up `y` as a 1x100 array (shape `(1,100)`), when we call `func(x,y)`, the output will be a 100x100 array (shape `(100,100)`). This is because NumPy looks at each axis in the arrays starting from the last axis, and each axis of size 1 is stretched to match the longer axis.\n", "\n", "| variable | axis 1 size | axis 2 size |\n", "| --- | --- | --- |\n", "| `x` | 100 | 1 |\n", "| `y` | 1 | 100 |\n", "| `result` | 100 | 100 |\n", "\n", "To illustrate this, let's look at a simple example. We'll create 2 arrays: `a` and `b`, each with the numbers 1-5, then show how NumPy uses broadcasting." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def test_function(x,y):\n", " print(x*y)\n", " \n", "a = np.asarray([1,2,3,4,5])\n", "b = np.asarray([1,2,3,4,5])\n", "\n", "print(f'a shape: {a.shape}, b shape {b.shape}') #show that a and b are 1D arrays of shape (5,)\n", "\n", "test_function(a,b)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "In this case, `a` and `b` are each 1D arrays, and so NumPy loops over both arrays *together*, producing another 1D array. NumPy saw that `a` and `b` only contained 1 axis, and the axis sizes matched, so the output has an axis of that same size.\n", "\n", "| variable | axis size |\n", "| --- | --- |\n", "| `a` | 5 |\n", "| `b` | 5 |\n", "| `result` | 5 |\n", "\n", "If, instead, we want to evaluate the function at every combination of `a` and `b`, we need to input 2D arrays. This is where broadcasting comes in. If one is shaped `(5,1)` and the other `(1,5)`, Numpy will broadcast the two arrays together into a resultant array that has shape `(5,5)`, and contains every combination of the two values. When NumPy operates on the two arrays, it first compares the size of axis 2. It finds that axis is size 1 for `a` and size 5 for `b`, so `a` is stretched over all values of b to make the result have axis 2 size = 5. A similar thing occurs for axis 1.\n", "\n", "| variable | axis 1 size | axis 2 size |\n", "| --- | --- | --- |\n", "| `a` | 5 | 1 |\n", "| `b` | 1 | 5 |\n", "| `result` | 5 | 5 |" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "a2 = a.reshape((5,1))\n", "b2 = b.reshape((1,5))\n", "\n", "print(f'a2 shape: {a2.shape}, b2 shape {b2.shape}')\n", "\n", "test_function(a2,b2)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "In our case, we want to have a 3D grid of $(r,\\theta,\\phi)$ coordinates, each with 101 points. If we just create 3 1D arrays, each with 101 points, and call `h_orbital`, the result will just be a 1D array of length 101. The function will not be evaluated at all possible combinations of $r$, $\\theta$, and $\\phi$.\n", "\n", "| variable | axis 1 size | axis 2 size | axis 3 size |\n", "| --- | --- | --- | --- |\n", "| $r$ | 101 | - | - |\n", "| $\\theta$ | 101 | - | - |\n", "| $\\phi$ | 101 | - | - |\n", "| `result` | 101 | - | - |\n", "\n", "Instead, we need to create a set of 3D arrays that cover our grid. Following the broadcasting rules, we can create `rg` with shape `(101,1,1)`, `tg` with shape `(1,101,1)`, and `pg` with shape `(1,1,101)` for $r$, $\\theta$, and $\\phi$. Then a call to `h_orbital(rg,tg,pg,2,1,0)` will return a 3D array with shape `(101,101,101)`, where the function has been evaluated at each point on the grid.\n", "\n", "| variable | axis 1 size | axis 2 size | axis 3 size |\n", "| --- | --- | --- | --- |\n", "| `rg` | 101 | 1 | 1 |\n", "| `tg` | 1 | 101 | 1 |\n", "| `pg` | 1 | 1 | 101 |\n", "| `result` | 101 | 101 | 101 |\n", "\n", "The function [`numpy.meshgrid`](https://numpy.org/doc/stable/reference/generated/numpy.meshgrid.html) can be used to conveniently construct a multidimensional grid like this from 1D arrays, such as the ones made by `numpy.linspace`. This is convenient, as it automatically sets up the dimensions that will work correctly with broadcasting. Read the documentation to understand why `indexing` and `sparse` are used below. You can also change them and see what happens." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "r = np.linspace(0,10,101)\n", "t = np.linspace(0,np.pi,101)\n", "p = np.linspace(0,2.*np.pi,101)\n", "\n", "print(f'r shape = {r.shape}, t shape = {t.shape}, p shape = {p.shape}')\n", "\n", "rg,tg,pg = np.meshgrid(r,t,p,indexing='ij',sparse=True)\n", "\n", "print(f'rg shape = {rg.shape}, tg shape = {tg.shape}, pg shape = {pg.shape}')\n", "\n", "out3 = h_orbital(rg,tg,pg,2,1,0)\n", "\n", "print (f'out3 shape = {out3.shape}, total number of elements: {out3.size}')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "If, instead, we did not use `numpy.meshgrid` and just used `r`, `t`, and `p` directly, the function would only be evaluated at 101 points, not 101x101x101 as we desire." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "out4 = h_orbital(r,t,p,2,1,0)\n", "\n", "print(f'out4 shape = {out4.shape}, total number of elements: {out4.size}')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now with the grids set up, we can do a timing comparison." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%%timeit\n", "\n", "h_orbital(rg,tg,pg,2,1,0)" ] }, { "cell_type": "markdown", "metadata": { "tags": [] }, "source": [ "In my trial this was about 3000 times faster than the python `for` loops! The take-home point of this section is that whenever possible, you should avoid using python's built-in `for` or `while` loops to loop through NumPy arrays. In fact, many other loop applications can also be avoided by clever use of `ndarray` and broadcasting, and in general this approach will be much faster. You can always test the performance of a code snippet with the `%timeit` magic function that we have been using; it's often a good way to identify performance bottlenecks and test ways to improve them.\n", "\n", "## Introducing `matplotlib`\n", "\n", "For data visualization, [`matplotlib`](https://matplotlib.org/) is one of the most commonly-used python packages, and it is the one we will use in this course. If this is the first time you are using `matplotlib`, take a moment to read through this basic [usage guide](https://matplotlib.org/tutorials/introductory/usage.html), which walks through the very basics of creating simple plots. Below we will begin generating plots, and the focus will mostly be on the hydrogen orbitals and the NumPy details, not on every matplotlib call. The usage guide will provide more explanation of many of the basics, and the API documentation for the [Axes](https://matplotlib.org/api/axes_api.html) and [Figure](https://matplotlib.org/api/_as_gen/matplotlib.figure.Figure.html) objects have many more details and options.\n", "\n", "To begin, let's work with $s$ orbitals, which are spherically symmetric. This means that an $s$ orbital has the same value for every $\\theta$ and $\\phi$, and therefore the wavefunction only depends on $r$. While we could simply use 1D arrays for $r$ and $\\psi$, we will still do the calculation in 3D so that our code is flexible. For this calculation, we'll use 1000 $r$ points, and 100 $\\theta$ and $\\pi$ points." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "#Specifying endpoint tells linspace whether to include the end value in the array\n", "r = np.linspace(0,10.,1000,endpoint=False)\n", "t = np.linspace(0,np.pi,100,endpoint=True)\n", "p = np.linspace(0,2*np.pi,100,endpoint=False)\n", "\n", "rg, tg, pg = np.meshgrid(r,t,p,indexing='ij',sparse=True)\n", "\n", "psi_1s = h_orbital(rg,tg,pg,1,0,0)\n", "\n", "print(f'psi_1s shape: {psi_1s.shape}')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "As you read the `matplotlib` documentation, you will see that it references the \"pyplot API\" and the \"object-oriented API.\" While both are perfectly valid, we will use the object-oriented API here. This is mostly to prevent errors that may happen when you use a notebook environment and execute cells in different orders. First we need to import the library and create a new set of `Figure` and `Axes` objects. The `Figure` is like the container for one or more plots, and the unfortunately-named `Axes` object (not to be confused with an `Axis` object!) is one of the plots. We can create the `Figure` and `Axes` with another unfortunately-named command: [`pyplot.subplots`](https://matplotlib.org/api/_as_gen/matplotlib.pyplot.subplots.html?highlight=subplots#matplotlib.pyplot.subplots). Although that function can create a figure with multiple subplots, when called without arguments, it actualy produces a figure with only one plot, which is what we want in this case." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import matplotlib.pyplot as plt\n", "\n", "fig, ax = plt.subplots()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We'll start by making a 1D visualization of the $1s$ orbital as a function of $r$ at a single value of $\\theta$ and $\\phi$. To do this, we will take [slices](https://numpy.org/doc/stable/reference/arrays.indexing.html#basic-slicing-and-indexing) of the `psi_1s` array. We're also using the [`np.real`](https://numpy.org/doc/stable/reference/generated/numpy.real.html) function to take only the real part of the wavefunction. The $1s$ orbital is entirely real anyways, so this does not discard any information; later we will deal with other wavefunctions that are complex.\n", "\n", "To make a graph of a 1D function, we can use [`Axes.plot`](https://matplotlib.org/api/_as_gen/matplotlib.axes.Axes.plot.html), which by default produces a line graph for xy data. `Axes.plot` can take either 1 or 2 arrays as arguments. If only 1 array is present, then it is interpreted as the y data for the plot, and the x values will be the point number in the array (0, 1, 2, 3, ...). If 2 arrays are present, then they must be the same size, and the first array contains the x values and the second array contains the corresponding y values. Here we take our 1D slice of `psi_1s` at $\\theta=0$, $\\phi=0$, and plot it against a 1D slice of `rg`. In the call to `Axes.plot`, we start with the x data, then the y data, and then we can specify additional arguments that customize the appearance of the curve. More details on the options that are available are in the documentation link above; here we'll just add a label that will later show up in a plot legend." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ax.cla() #clear the axes object in case things have already been plotted on it\n", "ax.plot(rg[:,0,0],np.real(psi_1s[:,0,0]),label=r'$\\psi_{1s}$')\n", "ax.set_xlabel('$r$ ($a_0$)')\n", "ax.set_ylabel(r'$\\psi$')\n", "fig #Since the figure was already plotted above, we need to tell iPython to show it again with the updated info" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "A plot can have more than one curve on it. Here we'll calculate the $2s$ orbital, plot it, and add a legend." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "psi_2s = h_orbital(rg,tg,pg,2,0,0)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ax.plot(rg[:,0,0],np.real(psi_2s[:,0,0]),label=r'$\\psi_{2s}$')\n", "ax.legend()\n", "fig" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "To better see the node structure, we can add a dashed horizontal black line." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ax.axhline(0,0,1,linestyle='dotted',color='black')\n", "fig" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Or, perhaps better in this case, we can change the way the [axis spines](https://matplotlib.org/api/spines_api.html) are set up. The `matplotlib` documentation has many examples that demonstrate its capabilities in a more user-friendly way than the API documentation directly. [This example](https://matplotlib.org/examples/pylab_examples/spine_placement_demo.html) shows an overview of some ways that axis spine positioning works. Let's create another figure for comparison. Because later we're going to add a second subplot to the figure, this time we'll create the figure with [`pyplot.figure`](https://matplotlib.org/api/_as_gen/matplotlib.pyplot.figure.html), telling it to make a figure that is 8 inches wide by 4 inches tall, then call [`Figure.add_subplot`](https://matplotlib.org/api/_as_gen/matplotlib.figure.Figure.html#matplotlib.figure.Figure.add_subplot) with the argument `1,2,1` to say that the `Axes` object will be part of a figure with 1 row of plots, 2 columns of plots, and this plot is at position 1." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "fig2 = plt.figure(figsize=(8,4))\n", "ax2 = fig2.add_subplot(1,2,1)\n", "ax2.plot(rg[:,0,0],np.real(psi_1s[:,0,0]),label=r'$\\psi_{1s}$')\n", "ax2.plot(rg[:,0,0],np.real(psi_2s[:,0,0]),label=r'$\\psi_{2s}$')\n", "ax2.set_xlabel('$r$ ($a_0$)')\n", "ax2.set_ylabel(r'$\\psi$')\n", "\n", "ax2.spines['left'].set_position('zero')\n", "ax2.spines['right'].set_color('none')\n", "ax2.spines['bottom'].set_position('zero')\n", "ax2.spines['top'].set_color('none')\n", "ax2.xaxis.set_ticks_position('bottom')\n", "ax2.yaxis.set_ticks_position('left')\n", "ax2.legend()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "There are many, many ways to customize the appearance of plots, annotate them with additional information, or make them more sophisticated. It's impossible in this class to do any more than scratch the surface. Be sure to check out the [matplotlib gallery](https://matplotlib.org/gallery.html) to see on overview of the kinds of things that are possible. Each example has source code available that you can draw from as you make your own plots. We will take a look at other plotting styles and plot types as we go forward, but you should always go back to the gallery as you learn more!\n", "\n", "Just like in ordinary python, you can put expressions into the call to `Axes.plot`. Here we'll plot the squares of the wavefunctions. Recall that for a complex number $z = a + bi$, its square is $z^*z$, where $z^*$ is the complex conjugate of $z$: $z^* = a - bi$. The square is therefore real\n", "\n", "$$ z^*z = (a - bi)(a + bi) = a^2 + b^2 $$\n", "\n", "however, when Numpy operates on an array of complex numbers, the output is also an array of complex numbers, even if the imaginary part is 0. We'll calculate the square, then take the real part since the imaginary part is zero anyways mathematically.\n", "\n", "We'll add a subplot to `fig2` with the square using [`Figure.add_subplot`](https://matplotlib.org/api/_as_gen/matplotlib.figure.Figure.html#matplotlib.figure.Figure.add_subplot), and add a grid for the fun of it." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ax3 = fig2.add_subplot(1,2,2)\n", "\n", "ax3.plot(rg[:,0,0],np.real(np.conj(psi_1s[:,0,0])*psi_1s[:,0,0]),label=r'$\\psi_{1s}^2$')\n", "ax3.plot(rg[:,0,0],np.real(np.conj(psi_2s[:,0,0])*psi_2s[:,0,0]),label=r'$\\psi_{2s}^2$')\n", "ax3.set_xlabel('$r$ ($a_0$)')\n", "ax3.set_ylabel(r'$\\psi^2$')\n", "\n", "ax3.spines['left'].set_position('zero')\n", "ax3.spines['right'].set_color('none')\n", "ax3.spines['bottom'].set_position('zero')\n", "ax3.spines['top'].set_color('none')\n", "ax3.xaxis.set_ticks_position('bottom')\n", "ax3.yaxis.set_ticks_position('left')\n", "ax3.grid()\n", "ax3.legend()\n", "\n", "fig2.tight_layout() #recalculates subplot sizes to try and prevent labels from overlapping\n", "fig2" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Sometimes, it's useful to manually adjust the plot ranges. On the right subplot, it's hard to make out the node in the $2s$ orbital. We can zoom in with [`Axes.set_ylim`](https://matplotlib.org/api/_as_gen/matplotlib.axes.Axes.set_ylim.html) (there is an equivalent [`Axes.set_xlim`](https://matplotlib.org/api/_as_gen/matplotlib.axes.Axes.set_xlim.html) for the x axis as well)." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ax3.set_ylim(0,0.002)\n", "fig2" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.9.19" } }, "nbformat": 4, "nbformat_minor": 4 }