How to draw a direction field in Python

Photo by USGS on Unsplash

This past semester I taught Linear Algebra and Differential Equations (one course that combines those two subjects) for engineering students. We placed a strong emphasis on using professional computing tools to handle a lot of the work. And when I use computing tools, I prefer to use open-source software whenever it makes sense, so I decided we’d use SymPy and Jupyter notebooks as our main platform.

My hope was that students would do all of their work simply by opening a Jupyter notebook, loading SymPy, and using only standard library or SymPy commands. Sadly, that wasn’t to be the case: While SymPy has some sophisticated methods for symbolic solutions to DEs, it doesn’t have any functionality for numerical solutions, or graphical solutions like direction fields. (Which is not a failing of SymPy; it’s called SymPy for a reason, after all.)

I was hoping to find a way to make direction fields within Python if not within SymPy. It turns out it’s possible, but not obvious. I found some tutorials about this (here and here for example) but it took a lot of remixing to make the instructions clear enough for engineering students who are Python newbies. To save you all the same work, this article will distill those instructions and explain what’s going on under the hood.

I’m going to go into a lot of detail on how to code this and what the code is doing. If you just want the code without the details, skip to the end, where I’ve put a big code block at the end for you to copy from.

What is a direction field?

To understand the Python behind making a direction field it’s helpful to think about what one of these things is, and how you’d make one by hand if you needed to.

Start with a first-order differential equation, dy/dx = f(x,y). A direction field for the DE is a collection of line segments (or arrows) where the slope of the line segment at a point equals the derivative at that point. For example if you have the DE dy/dx = x — y then the segment drawn at (1,2) would have slope -1; the segment drawn at the origin would have slope 0; and so on. This is what the direction field looks like for this particular DE if you drew segments at various points for x and y between -5 and 5:

Direction field for dy/dx = x — y on -5 < x < 5, -5 < y < 5

Most of the time direction fields normalize the lengths of the segments to length 1, as above. (I drew this with this website, not Python, because it’s easier; more about this later.) A direction field therefore shows you a birds-eye view of the “trajectories” of all possible particular solutions to the DE.

So how would you draw one of these if you had no computer? You would:

1. Decide where in the xy-plane you want the direction field; then
2. Decide how many segments you want to draw and locate the “tails” of the arrows that produce this many segments; then
3. At each tail, find dy/dx there and draw a length-one line segment that has that slope.

It’s fairly straightforward to draw a direction field by hand, but you wouldn’t want to make one that had a lot of arrows in it. That’s where the computer comes in.

Step 1: Making a grid

Making a direction field in Python is basically the same process as if you were doing it by hand. We just have to translate each step into the proper commands. To get this going in Python, we need to import two packages first:

import numpy as np
import matplotlib.pyplot as plt 

NumPy is Python’s numerical computing package, and Matplotlib is for drawing graphs of things.

First we do a combination of steps 1 and 2 of the “by hand” process: We’re going to plot a grid of dots in the xy-plane where the tails of our arrows are going to go. The way this works in NumPy is to first create an array of x-coordinates for those dots using the arange function, then an array of the y-coordinates, then combine them into a grid using the meshgrid function.

For example, if we wanted to draw a grid of dots with x and y between -5 and 5:

x = np.arange(-5,5,1) 
y = np.arange(-5,5,1)
X, Y = np.meshgrid(x,y) 

The first and second arguments in arange are the start and end points; the third argument tells you how far apart individual points are spaced. For example the first line above produces an array with entries -5, -4, -3, …, 2, 3, 4. It’s like the range command in the standard library, but it produces an array instead of an “iterable”.

The meshgrid command produces a data structure that holds pairs of pints with the coordinates specified by the first two lines. The actual output is a list of arrays that has length 2: The first array gives the x-coordinates of the dots and the second array gives the y-coordinates. In the code, we unpack the output into the variables X and Y, and note the capital letters. These are not the same things as x and y!

The output of meshgrid

The meshgrid is the data structure that represents the grid of dots internally. If we were to stop here and plot the results, here’s what we’d have:

Plotting the results of meshgrid

Step 2: Adding the differential equation

Now we need to draw arrows at these dots, with the proper slopes. Let’s look at the differential equation from earlier, dy/dx = x-y. (If your DE is not solved with dy/dx isolated on the left, do that first.) We start with the slope:

dy = X - Y
dx = np.ones(dy.shape) 

What we’re doing here is specifying a formula for “rise” and “run” of each arrow that we will eventually apply to each dot in the grid. The dy is just the differential equation, since this tells you the rate at which y changes at each point. Note: We’re using X and Y, capitalized, because those are the coordinates in the grid; not x or y, lower case, even though it’s lower case on the left.

What’s happening with dx is that we’re specifying a constant “run” for each of the arrows, namely 1. The command np.ones creates an array filled with 1’s, and the dy.shape tells np.ones to make the length of that array the same as the length of dy.

To stop and take a peek at the output:

Output of dy

You can almost begin to see the direction field here just from the numbers. The output of dx is just like the above except all entries are 1.

If we plotted the arrows in the direction field now, they would all have different lengths because the change in the x-direction is constantly 1. The direction field would be hard to read because the arrows might overlap or be so short as to be invisible. So to make things more visually informative, let’s find unit vectors for each of these that have the same slope:

norm = np.sqrt(dx**2 + dy**2)
dyu = dy/norm
dxu = dx/norm

This is just the standard trick of dividing by the length of each arrow. Note that we never divide by 0 here because we specified the initial dx to be 1.

Step 3: Plotting the arrows

What do you call it when you have a bunch of arrows together? It’s called a quiver, and that’s the name of the Matplotlib command to plot a collection of vectors. These two lines complete the process:

plt.quiver(X,Y,dxu,dyu)
plt.show()

The quiver command takes four inputs: The first two tell Python where the arrows are supposed to go, and the second two give the rise and run of each one. There are optional additional arguments for things like color. Here’s the full documentation. And here’s the result of the code:

The completed direction field

Direction fields for systems

One nice thing about this process is that it’s very easy to extend it to making direction fields for autonomous systems of first-order DEs — where you have y and x both changing with respect to a third variable t but the t variable doesn’t appear in the DEs themselves. Such systems are fairly common in modeling applications. All we have to change is the dx, so that instead of being a constant array of 1’s, it matches the dx/dt equation in the system.

Here’s a direction field for the system dx/dt = x + 2y, dy/dt = x — y using the same parameters as before:

x = np.arange(-5,5,1) 
y = np.arange(-5,5,1)
X, Y = np.meshgrid(x,y)

dy = X - Y
dx = X + 2*Y
norm = np.sqrt(X**2 + Y**2) 
dyu = dy/norm
dxu = dx/norm

plt.quiver(X,Y,dxu,dyu,color="purple")
plt.show()

Note that this time we do encounter division by 0, at the origin; Python gives you a warning about it but produces the plot anyway, just leaving out those points.

Is Python the best tool for this job?

This is a question I asked myself a lot during the course. On the one hand, the code for generating the direction field in the end is pretty simple. But it does take some effort, and there’s a tendency in students to just copy and paste the code. And, while it’s not hard to create the direction field, it’s not easy to work with it once it’s produced. For example, it’s nontrivial to plot particular solutions to the DE on the direction field, which is something I want students to do. Even adding axes to the plot takes additional code; if this were a numerical methods course I’d expect students to handle this, but it isn’t.

So in practice, for teaching purposes, we used this website hosted at Bluffton College a lot more frequently. It’s interactive and intuitive, for example to see a particular solution in a direction field, just click on the initial condition. It also includes numerical solutions using several different algorithms alongside the direction fields. We reserved the Python stuff for their applications assignments, which were to be entirely done self-contained in a Jupyter notebook.

If there are simple ways to do this in Python with the direction fields I’ve described above, I’d love to know about this — leave a comment!

TL;DR

If you are just wanting the code for a direction field, here you go.

For a single differential equation (this uses dy/dx = x-y):

# Specify the grid of dots 
x = np.arange(-5,5,1) 
y = np.arange(-5,5,1)
X, Y = np.meshgrid(x,y)

# Create unit vectors at each dot with correct slope
dy = X - Y
dx = np.ones(dy.shape)
norm = np.sqrt(X**2 + Y**2) 
dyu = dy/norm
dxu = dx/norm

# Plot everything 
plt.quiver(X,Y,dxu,dyu)
plt.show()

For an autonomous system, see above.

Thanks for reading! Most of my writing is about higher education, grading practices, and productivity but I write about data science and computing here on Medium from time to time. I post something, somewhere every week and aggregate it at my website, rtalbert.org along with new material. I appreciate your attention and really value your feedback, so leave comments with insights and suggestions.