Guide to Matplotlib Plotting¶
Matplotlib is the most established Python plotting library and has a rich history and set of documentation. It is geared more toward static plots, so without some of the fancy interaction and bells and whistles of other plotting libraries, but if you just need a simple plot of something, you can't really go wrong with Matplotlib.
We can begin by importing the matplotlib libraries and numpy for handling vectors. The %matplotlib is a special ipython magic that sets up the notebook to display the plots directly in the notebook itself.
%matplotlib inline
import matplotlib.pyplot as plt
plt.style.use('seaborn-darkgrid') # This is my own preference but I like this style better than the default
import numpy as np
Parametric Plotting (and Plotting in general)¶
In many instances, particularly with multidimensional problems, you may find your solutions in terms of x, y, and z with respect to time, but want to plot the spatial dependence (say, x vs y). In such cases, it can be use to plot things parametrically.
The basic form of the expression is:
plt.plot( x(t), y(t), other options )
In order to easily plot $x(t)$ and $y(t)$, I find it useful to define them as a function. Here we'll look at the simple example of projectile motion, where we're launching a projectile at 100 m/s at an angle of $45^\circ$. We'll define our constants and the projectile motion solutions for $x$ and $y$:
# Defining initial conditions
ri = (0,0) # Initial position in m
vi = 100 # Initial velocity magnitude in m/s
ai = 45*np.pi/180 # Initial angle in radians
g = 9.8 # Acceleration in m/s^2
def x(t,ai):
return ri[0] + vi*np.cos(ai)*t
def y(t,ai):
return ri[1] + vi*np.sin(ai)*t - g/2*t**2
Note that here I also made x and y function of the initial angle. If I always wanted to use the same firing angle, I didn't need to do this, but since I wanted to vary that later on, I accounted for it up front here. Otherwise I would have needed to redefine the function later to update the new angle.
I'll mention here that there is an alternative way to define functions which is a bit more concise, but maybe not as clear. Instead of the def functions, I could have used:
x = lambda t, ai: ri[0] + vi*np.cos(ai)*t
y = lambda t, ai: ri[1] + vi*np.sin(ai)*t - g/2*t**2
These are called lambda functions and work identically to the above functions. You can use whichever you like.
Now that we have our solutions, we just need to define the time span we want to plot over and then we can go ahead and plot up the trajectory:
# Specify your range of times. Here we go from 0 to 10 with 100 data points
t = np.linspace(0,10,100)
plt.plot(x(t,ai),y(t,ai))
plt.show()
Plotting Options¶
You can of course specify some properties about the line you want to plot:
plt.plot(x(t,ai),y(t,ai), linestyle='--', color='red', linewidth=3)
plt.show()
It is always a good idea to label things!
plt.plot(x(t,ai),y(t,ai), linestyle='--', color='red', linewidth=3)
plt.xlabel('Horizontal Distance')
plt.ylabel('Vertical Height')
plt.grid() #toggles the grid on or off
plt.show()
Putting multiple plots on the same graph is as simple as calling multiple plt.plot calls. In Jupyter notebooks, if you want plots to be on the same graph, they MUST be called in the same cell (as far as I understand). No cross-cell plot combinations!
ai = 45*np.pi/180
plt.plot(x(t,ai),y(t,ai), linestyle='--', color='red', linewidth=3, label='45$^\circ$')
ai = 20*np.pi/180
plt.plot(x(t,ai),y(t,ai), linewidth=3, label='20$^\circ$')
plt.xlabel('Horizontal Distance')
plt.ylabel('Vertical Height')
plt.legend(loc='upper left')
plt.show()
Note that the use of legends requires that you have labeled your plots when you created them!
I can also make the plot larger and specify the plot limits:
plt.figure(figsize=(12,8))
ai = 45*np.pi/180
plt.plot(x(t,ai),y(t,ai), linestyle='--', color='red', linewidth=3, label='45$^\circ$')
ai = 20*np.pi/180
plt.plot(x(t,ai),y(t,ai), '.', markersize=10, label='20$^\circ$')
plt.xlabel('Horizontal Distance')
plt.ylabel('Vertical Height')
plt.legend(loc='upper left')
plt.ylim([0,300])
plt.xlim([0,800])
plt.show()
Normal Plotting¶
If you just wanted a normal plot (non-parametric) then it is as simple as plotting against your independent variable:
plt.figure(figsize=(12,8))
plt.plot(t, x(t,ai), linewidth=3, color='green', label='X-Position') # x vs t
plt.plot(t, y(t,ai), linewidth=3, color='gold', label='Y-Position') # y vs t
plt.xlabel('Time')
plt.ylabel('Distance')
plt.legend(loc='upper left')
plt.show()
Plotting 3 Variables ¶
In many cases, we don't have just one dimensional data. We might know something about the potential energy at a certain point in space. Which means we could have 2 (or 3) independent variables along with the dependent variable. Here we'll look at methods to visualize that sort of data.
The setup¶
Before we can plot anything, we need to worry about generating the data to plot! Let's look at the function: $$ U(x,y) = (1-(x^2 + y^3))e^{-(x^2+y^2)/2} $$ where $x$ and $y$ are in dimensions of meters and some constant of 1 out front gets things in units of joules. We want to visualize the potential energy over some space in $x$ and $y$. We can define our function as we normally do, except now it is a function of two variables:
def U(x,y):
return (1-(x**2+y**3))*np.exp(-(x**2+y**2)/2)
To visualize this, we need to calculate the potential energy at each point in space. Say we want to visualize from -3m to 3m in 1cm increments in both the $x$ and $y$ direction. One way we could calculate all the needed potentials is by using a nested for loop, looping over all the possible points in both the $x$ and $y$ directions. And this would totally work. But we can simplify this by using Numpy's meshgrid function:
xs = np.arange(-3,3, 0.01) # Measurement points in X
ys = np.arange(-3,3, 0.01) # Measurement points in Y
Xgrid, Ygrid = np.meshgrid( xs, ys) # Generating the grid
print(Xgrid)
Meshgrid can be difficult to understand, so I displayed Xgrid above to help. All meshgrid really does is create two matrices: one which holds the x values of each point and one that holds the y values at each point. In the above we can see that Xgrid has the same value in each column, because each point in a column is at the same $x$ coordinate. Numpy can handle functions of matrices (arrays), so this means that we can simply pass both results of meshgrid to our function to return the potential energy at all our desired points:
AllPotentials = U(Xgrid, Ygrid)
Heat/Height Maps¶
It can be tough to check our answer without visualizing it in some way (hence part of the power of visualization!). To start we'll just visualize things with a simple heat/height map:
plt.figure(figsize=(10,8))
plt.imshow(AllPotentials,
cmap = 'magma', #setting the color map
origin = 'lower', # IMPORTANT, or else the upper left will be your origin!
extent = [-5,5,-5,5], # imshow doesn't know the boundaries otherwise
)
plt.colorbar()
plt.title('Potential Energy in Joules')
plt.xlabel('X (m)')
plt.ylabel('Y (m)')
plt.show()
Imshow has a few oddities that I described above in the code comments. In particular, be wary of the origin flip and the fact that you must set an extent to get the proper axis ticks. Try commenting out any of those rows to see what they change (But realize that the origin flip won't be visible here because of the symmetry.) Also note that making an imshow image without a colorbar is akin to failing to include units in your calculation! Don't do it!
Contour Plots¶
In some cases, it can be cleaner or more clear to use a contour plot instead of a height map. The data preparation is identical to that of the height map, so we are free to just construct the plot.
plt.figure(figsize=(8,8))
cs = plt.contour( # Note we are specifically naming this to reference in the next command
Xgrid, # array of x positions
Ygrid, # array of y positions
U(Xgrid, Ygrid), # the function or 'z' values
np.arange(-1,2,.2), # A number of desired contours or a list of exactly what contours
cmap='viridis', # a colormap, if want single color use colors='black'
)
cs.clabel( # adds text labels to our contours
inline = True, # inserts the label in a gap in the contour
fontsize = 8, # sets the label size
fmt = '%0.0f J', # sets the format of the string. Could usually just take default
)
plt.xlabel('X (m)')
plt.ylabel('Y (m)')
plt.title('Contour Plot of Potential Energy')
plt.show()
Some benefits of contour plots over height maps include not having to specify the origin or extent. Plus in many situations they more clearly show the exact levels without having to rely on an imprecise color bar.
Often the nicest results can be had by using both, plotting the height map and then overlaying a contour plot on top of that. For example:
plt.figure(figsize=(10,8))
plt.imshow(U(X,Y), cmap='coolwarm', extent=[-3,3,-3,3], origin='lower')
plt.colorbar(label='Potential Energy (J)')
plt.title('Some Mysterious Potential Energy Plot')
plt.xlabel('X (m)')
plt.ylabel('Y (m)')
plt.grid(False)
cs = plt.contour(X, Y, U(X,Y), 10, colors='black')
cs.clabel(inline = True, fontsize=8, fmt='%0.1f J')
plt.show()
3D Plots¶
One of the last ways we can represent data in this format is with an actual 3D plot. These come with some caveats however:
- In many cases, a 3D plot might look pretty but might not actually convey the important information as well
- There are some slightly different steps and syntax to making a 3D plot
Those stated, there are plenty on instances where the overall data just can't be understood in 2D slices, so 3D plots have their place. Documentation for 3D plots can be found here. To get started, we'll need to import the 3D axis toolkit:
from mpl_toolkits.mplot3d import Axes3D
Now, the 3D plotting isn't too different, but it does require a little different format from what we've used above. So, to make a 3D plot of our potential $U$:
with plt.style.context(('default')): # using the default style here, as some others look weird in 3D
plt.figure(figsize=(10,6)) # Set the figure size as per normal
ax = plt.subplot( # We need to create a named 'subplot' of 1 here to specify the projection
111, # This is shorthand for a 1x1 grid of plots, where this is the 1st one
projection='3d', # This important bit, the projection must be 3D!
)
ax.plot_surface( # Command to plot a surface. Note it has the above name as the start!
Xgrid, # Our x values array
Ygrid, # Our y values array
U(Xgrid,Ygrid), # Our z values array
cmap='viridis', # Choose your colormap
alpha=.9, # alpha is how opaque the surface is. This makes it 10% transparent
)
# Need to label axes differently in 3D plots
ax.set_xlabel('X (m)')
ax.set_ylabel('Y (m)')
ax.set_zlabel('U (J)')
ax.set_title('Surface Plot of $U$', pad=30)
# If you want to rotate the view around, elev is tilt and azim is swivel
ax.view_init(elev=30, azim=45)
# Fun extra, you can plot contours (or height maps) on various planes
# Uncomment the below if you want to see the effect
ax.contour(Xgrid, Ygrid, U(Xgrid,Ygrid), zdir='z', offset=-1, cmap='viridis')
plt.show()
Vector Plots ¶
All the above plotting methods work great for when you only have 1 dependent variable to show. But what about in the case of vector functions? Here you are returning a vector, which would have 2 to 3 or more components. We could plot the magnitude of the vector for instance as a contour plot, but we couldn't show both the $x$ and $y$ components of the vector in a contour plot. One method to nicely display this type of data is with what are commonly called vector plots, and in matplotlib are known as quiver plots.
Let us take the vector function we had in class $\vec F (\vec{r})$: $$ \vec F(\vec r) = 3y \,\boldsymbol{\hat{x}} - x \boldsymbol{\hat{y}} $$ At each point in space we have two vector components to show. Our initial setup will look very similar to what we did above, except our function will now return two values: both the $x$ and $y$ component.
def F(r):
rx = r[0] # Just defining these to be perfectly clear
ry = r[1]
return 3*ry, -rx
Since we want to visualize the function over a 2D space, we'll still need to use meshgrid to generate our coordinates. Here we'll go over the interval from -2 to 2 in steps of 0.2.
xs = np.arange(-2,2,0.2)
ys = np.arange(-2,2,0.2)
Xgrid, Ygrid = np.meshgrid(xs, ys)
Now we can run everything through our function as before. An important point: we made our function F be a function of just a single vector $\vec r$, so when passing it components we need to do so as a list or tuple:
Fxs, Fys = F((Xgrid,Ygrid))
Now that we have our $x$ and $y$ components of $F$ all through the desired space, we can go ahead and use a quiver plot to visualize the vector field.
plt.figure(figsize=(8,8))
plt.quiver(Xgrid, # The array of x coordinates
Ygrid, # The array of y coordinates
Fxs, # The array of x components of F
Fys, # The array of y components of F
)
plt.xlabel('X Values')
plt.ylabel('Y Values')
plt.title(r'Quiver Plot of $\vec F(\vec a)$')
plt.show()
You can add extra arguments if you want to color the arrows by their magnitude, set where the arrow starts from, or to tweak the automatic scaling. I'll include an example below, or check out the documentation.
def G(b):
bx, by = b
return by*np.cos(bx), bx*np.sin(by)
xs = np.linspace(0,2*np.pi,20)
ys = np.linspace(0,2*np.pi, 20)
X,Y = np.meshgrid(xs, ys)
Gx, Gy = G((X,Y))
plt.figure(figsize=(10,10))
Q = plt.quiver(X,Y,Gx,Gy,np.hypot(Gx,Gy), cmap='magma') # The 5th entry is what determines color. Here I base it off the magnitude
plt.quiverkey(Q, .87, .89, 10, r'10 units', coordinates='figure') # Adds the little key at the top right
plt.xlabel('X Axis')
plt.ylabel('Y Axis')
plt.title(r'$\vec G(\vec b)$ Vector Field')
plt.show()
