Tour of matplotlib

This notebook covers roughly the same ground as this video.

In this notebook we plot various graphs. For this we need the matplotlib library as well as numpy. The standard way to import these is as below. There is comprehensive documentation for matplotlib at matplotlib.org, including tutorials, examples and detailed references. Also, you can just tell Google Gemini in English what you want to do, and it will often be able to supply the right matplotlib commands.

import math
import numpy as np
import matplotlib.pyplot as plt

We first make a graph of $y=\sin(x)$ from $x=0$ to $x=2\pi$. We will calculate 100 points on this graph, and then matplotlib will draw a smooth curve joining them. We first use np.linspace to generate an array xs of 100 equally spaced $x$ values between $0$ and $2\pi$. (You should think of xs as the plural of x; this is a fairly standard naming convention.) We then want to calculate the corresponding $y$ values. Because np.sin() is a universal function, we can just write ys = np.sin(xs) to calculate all the $y$ values at the same time. We then call plt.plot() to draw the graph.

xs = np.linspace(0, 2*np.pi, 100)  # x coordinates of points on the graph
ys = np.sin(xs)                    # y coordinates of points on the graph
plt.plot(xs, ys)

For the simplest plots we can just call plt.plot() as in the previous example. However, we will usually prefer a slightly different style which gives more control over the result. In this alternative style, we usually start with the line fig, ax = plt.subplots(). After that, we call various methods of the ax object, such as ax.plot(), ax.set_title() and so on.

In the example below:

• We first plot $y=\sin(\theta)$ for $\theta=0\dotsb2\pi$. By supplying some additional keyword arguments in the plot command, we make the curve orange, twice as thick as usual, and dashed. We also add the label $\sin(\theta)$. We enter this using LaTeX notation, which as usual contains backslashes. Computers often treat backslashes in various special ways, which can mess things up. To prevent the special treatment of backslashes we use the prefix r (short for "raw") when entering this string. Note also that we need to call ax.legend() later for this label to take effect.
• We then plot y=\sin(\theta)\sin(20\theta) in red. We could use color = 'red' as in the previous example, but matplotlib provides much shorter abbreviations for some of the most common options. In particular, we can just write 'r' for color = 'red'. Note also that if we call ax.plot() repeatedly then the resulting curves are all combined into a single picture. If we want to start a new picture we need to call plt.show() and then fig, ax = plt.subplots() again.
• We call ax.set_xlabel() to label the $x$-axis as $\theta$. Here again the string has a backslash so we use the prefix r.
• We call ax.set_title() to set the title which will appear at the top of the plot.
• By default, matplotlib would put labelled tick marks on the horizontal axis at $\theta=0$, $\theta=1$, $\theta=2$ and so on. That is not appropriate here: it is much better to have tick marks at multiples of $\pi/2$. We arrange this by calling ax.set_xticks().
• By default, matplotlib will put the $x$-axis at the bottom of the picture, at $y=-1.1$ or so. In this context it is better to put the $x$-axis at $y=0$ instead. We can refer to the $x$-axis as ax.spines['bottom'], and put it in the right place by calling ax.spines['bottom'].set_position('zero').
• Finally, we call ax.legend(). This tries to arrange all labels in a sensible way so that everything is visible.
• At the end we call plt.show() to display the plot. This is not really necessary because the Jupyter notebook framework causes plt.show() to be called automatically when it gets to the end of the cell. However, it is conceptually important. Before calling plt.show(), we have just queued up a list of instructions about what we want in the plot. No action is taken on these instructions until we get to plt.show(). The instructions will be processed in a sensible order even if they were given in a different order. Occasionally we will want to call plt.show() explicitly before the end of the cell. That way we can generate several separate plots in the same cell, for example.

ts = np.linspace(0, 2*np.pi, 629)
ys = np.sin(ts)
fig, ax = plt.subplots()
ax.plot(ts, ys, 
        color = 'orange',       # note the American spelling: color not colour
        linewidth = 2,          # twice as thick as usual
        linestyle = '--',       # dashed line
        label=r'$\sin(\theta)$' # Use r'...' for strings with backslashes
        )
ax.plot(ts, ys * np.sin(20*ts), 'r')
ax.set_xlabel(r'$\theta$') # Use r'...' for strings with backslashes
ax.set_title('A simple plot')
ax.set_xticks(np.linspace(0, 2*np.pi, 5),['0', r'$\pi/2$', r'$\pi$', r'$3\pi/2$', r'$2\pi$'])
ax.spines['bottom'].set_position('zero') # Draw x-axis at y=0 instead of at the bottom of the graph
ax.spines['top'].set_visible(False) # Don't draw top axis
ax.legend() # Without this, the label sin(theta) will not be shown
plt.show()

We next have an example where we make four separate graphs arranged in a grid. We start with fig, ax = plt.subplots(2, 2, figsize=(10, 10)), which specifies a $2\times 2$ grid of overall size $10\times 10$ (as compared with the usual default of $6.4\times 4.8$). Now ax is no longer a single object, but instead a $2\times 2$ array of objects, one for each subplot. We draw the top left subplot by calling methods of ax[0,0], then we draw the top right subplot by calling methods of ax[0,1], and so on.

We have added the word None at the bottom of the cell below to prevent the printing of an irrelevant message.

fig, ax = plt.subplots(2, 2, figsize=(10, 10))
xs = np.linspace(0, 1, 100)
ax[0,0].plot(xs, np.sin(4 * np.pi * xs))
ax[0,0].set_title('Sinusoid')
ax[0,1].plot(xs, 4 * xs * (1 - xs))
ax[0,1].set_title('Parabola')
ax[1,0].plot(xs, np.exp(-(4*xs-2)**2))
ax[1,0].set_title('Gaussian')
ax[1,1].plot(xs, np.exp(-5*xs) * np.sin(50*xs))
ax[1,1].set_title('Damped Sinusoid')
None

The next example is similar: we plot the graphs $\sin(n\pi x)$ for $n=2,\dotsc,7$ in a $2\times 3$ grid.

fig, ax = plt.subplots(2, 3, figsize=(15, 10)) # Now ax is an array of axes of shape 2 x 3
ax = ax.flatten() # Flatten the array to make it easier to loop through the six axes
                  # now we have ax[0] .. ax[5] instead of ax[0,0] .. ax[1,2]
for i in range(6):
    axi = ax[i]
    m = i + 2
    axi.plot(xs, np.sin(m * np.pi * xs))
    axi.set_title(r'$y=\sin(' + str(m) + r'\pi x$)')

We can also do parametric plots, where $x$ and $y$ are both given as functions of another variable $t$. For example, the cycloid curve is given by parametric equations $x=t-\sin(t)$ and $y=1-\cos(t)$. (This is the path traced by a point on the edge of a rolling wheel.)

ts = np.linspace(0, 8*np.pi, 1000)
xs = ts - np.sin(ts)
ys = 1 - np.cos(ts)
plt.axis('equal')
plt.plot(xs, ys)
plt.hlines(0, 0, 8*np.pi, color='gray', linestyle='--')

Similarly, we can plot the curve given by $x=\frac{1-t^2}{1+t^2}$ and $y=\frac{2t}{1+t^2}$. Can you explain why this has the shape that it does?

ts = np.linspace(-8, 8, 1000)
xs = (1 - ts**2) / (1 + ts**2)
ys = 2 * ts / (1 + ts**2)
plt.axis('equal')
plt.plot(xs, ys)

Sometimes we are given a mathematical relationship like $f(x, y)=0$, and we cannot solve it explicitly to give $y$ as a function of $x$, but we still want to plot the curve where the relationship is satisfied. For example, we can take $f(x,y)=y^2-(x^3-x)$ and plot the curve where $f(x,y)=0$ or equivalently $y^2=x^3-x$. (This is an example of an elliptic curve.) For this we need to use np.meshgrid(), which works as follows. We start by generating an array xr of $x$ values, like xr=[-2,-1.99,-1.98,...,1.98,1.99,2], and an array yr of $y$ values, of the same form. We then use np.meshgrid() to generate two-dimensional arrays x and y, with x[i,j]=xr[i] and y[i,j]=yr[j]. In other words, the (i,j)th point in our grid is [x[i,j],y[i,j]]=[xr[i],yr[j]]. We then put z=y**2 - (x**3 - x), so z[i,j] is the value of $f(x,y)$ at the (i,j)th point in the grid. We then call plt.contour(x, y, z, levels=[0]) to plot our curve. We add levels=[0] to say that we just want the curve where $f(x,y)=0$; if we left that out then matplotlib would draw several curves of the form $f(x,y)=c$ for various values of $c$ which it would choose. Note that it is unlikely that any of our $z$ values are exactly equal to $0$, but matplotlib is intelligent enough to look for places where he sign of $z$ changes and do some interpolation.

xr = np.linspace(-2, 2, 400)
yr = np.linspace(-2, 2, 400)
x, y = np.meshgrid(xr, yr)
z = y**2 - (x**3 - x)
plt.contour(x, y, z, levels=[0], colors='red')

Here is an example that we can do either implicitly or explicitly. Put \begin{align*} f(x,y) &= x\,\sin(x^2+y^2)-y\,\cos(x^2+y^2) \\ u(t) &= t\,\sin(t^2) \\ v(t) &= t\,\cos(t^2). \end{align*}
With a little mathematical analysis you can show that the curve where $f(x,y)=0$ can be parametrised as $(x,y)=(u(t),v(t))$. In the picture below, the left hand plot is done implicitly (using np.meshgrid() and plt.contour() as in the previous example) and the right hand plot is done parametrically.

# Setup
fig, ax = plt.subplots(1, 2)
fig.set_size_inches(10, 5)
# Implicit version, using ax[0] = left hand side
def f(x, y):
    return x * np.sin(x**2 + y**2) - y * np.cos(x**2 + y**2)
xr = np.linspace(-5, 5, 1000)
yr = np.linspace(-5, 5, 1000)
x, y = np.meshgrid(xr, yr)
ax[0].axis('equal')
ax[0].contour(x, y, f(x,y), levels=[0], colors=['red'])
# Parametric version, using ax[1] = right hand side
def u(t):
    return t * np.cos(t**2)
def v(t):
    return t * np.sin(t**2)
t = np.linspace(-5, 5, 1000)
ax[1].axis('equal')
ax[1].plot(u(t), v(t))