A field guide to two-dimensional arrays
Usually on the blog we apply a function to each element of a one-dimensional array, like [1, 2, 3], and get back another array, like [10, 20, 30]. Then we use both lists to plot a series of (x, y) ordered pairs.
Today I’d like to expand on that process and apply functions to two-dimensional arrays. You can think of them like lists of lists. For example:
[[1, 2, 3], [4, 5, 6], [7, 8, 9]]
We can streamline things by building on a previous blog post. Back in August we modeled gravitational forces and plotted the motion of stars. Let’s revisit the universe we created then.
Pictured is a different neighborhood within the universe but the same physics apply. Think of the yellow body as our sun and the blue body as a comet making a rare visit. The arrow represents how much and in what direction the comet is accelerated by the gravitational force applied to it.
Remember, the magnitude of the force depends on the masses of the bodies and the square of the distance between them.
The motion makes total sense if you’ve played Asteroids, the arcade game. Even when you stop applying a force (thrust) your ship keeps moving. It isn’t like driving a car here on earth where friction will eventually bring you to a stop.
As our comet drifts away from the sun, the sun stops applying a force to it and it’s no longer accelerated, but it keeps drifting in the direction it was already going. This demonstrates Newton’s First Law: An object in motion will remain in motion until it’s acted on by an external force.
We’ve spent some time working with gravitational forces. Now let’s go a step further and plot a gravitational field. It might sound complicated but it’s not. It just means that when one mass approaches another, a force is applied to it. If a comet travels near the sun, it will enter the sun’s gravitational field and a force will be applied to the comet, causing its motion to change. This region of potential interaction is called a gravitational field.
The important concept is that the force depends on where exactly the comet is located. At each point within the field there is a certain potential for interaction. Our goal is to define and visualize each of the points.
We can visualize a field by drawing a bunch of arrows all over it. The arrows will point in whatever direction a body would experience a force at that point. And the bigger the arrow, the stronger the force.
Fortunately, Matplotlib has us covered when it comes to drawing a bunch of arrows in a grid pattern. This is called a quiver plot.
Prepare the data.
Let’s plot the “universe” pictured above with a yellow sun located at (20, 40). First define our constants, which are entirely arbitrary and not to scale.
STAR_POS_X = 20 STAR_POS_Y = 40 STAR_MASS = 4000 COMET_MASS = 100 G_CONSTANT = 1.0
We’re talking about quantities with magnitude and direction, which means we’re dealing with vectors. That means it’s time to import numpy.
This is where we begin working with two-dimensional arrays. Use meshgrid to create a sort of coordinate system for the universe.
import numpy as np
x_grid, y_grid = np.meshgrid(range(-100, 105, 5),
range(-100, 105, 5)
)
The returned arrays look like this:
x_grid = [[-100 -95 -90 ... 90 95 100]
[-100 -95 -90 ... 90 95 100]
[-100 -95 -90 ... 90 95 100]
...
[-100 -95 -90 ... 90 95 100]
[-100 -95 -90 ... 90 95 100]
[-100 -95 -90 ... 90 95 100]]
y_grid = [[-100 -100 -100 ... -100 -100 -100]
[ -95 -95 -95 ... -95 -95 -95]
[ -90 -90 -90 ... -90 -90 -90]
...
[ 90 90 90 ... 90 90 90]
[ 95 95 95 ... 95 95 95]
[ 100 100 100 ... 100 100 100]]
You should consider the arrays bound together. Each corresponding pair of values essentially forms an ordered pair, (x, y). Together the arrays define a grid that spans our whole universe.
These arrays are like the independent variable, x, on an old-fashioned line plot. If we were working with one-dimensional data, we would apply a function to the values of x and get back another series of values, y. In this case, we’ll apply a function to each point within the grid and get back a new two-dimensional array.
As we learned in the previous post, the equation for gravitational force between two bodies looks like this:
… where m1 and m2 are the masses of the bodies, d is the distance between their centers, and G is the gravitational constant.
But we’ll have to stretch our thinking to implement it for two-dimensional data.
First let’s calculate the denominator, distance. Remember, we’re finding the sun’s distance from each point in the grid.
distance_x = STAR_POS_X - x_grid distance_y = STAR_POS_Y - y_grid distance = np.sqrt(distance_x**2 + distance_y**2)
Notice how easy numpy makes it to “vectorize” functions across whole two-dimensional arrays. We don’t have to loop through each element and repeat the operation over and over again.
Since the sun exists at a specific point within the grid, (20, 40), distance at that point is zero, and we can’t divide by zero. We can use ma.masked_array to filter out values from the array and avoid an error. Let’s add a fudge factor up to 15 so we’ll have room to plot the sun.
distance = np.ma.masked_array(distance, distance <= 15)
Now that we know distance, use the equation above to calculate the force’s magnitude.
force = G_CONSTANT * STAR_MASS * COMET_MASS / distance**2
At this point, we don’t know the force’s direction so we couldn’t plot arrows. We could create a sort of “heat map” that uses color to indicate the gravitational force’s strength at each point. But I’ll save that for another post.
Today we need to know direction so use your trigonometry skills to find the x- and y-components of each arrow. We’re only considering one mass in this universe so every arrow will point directly at (20, 40).
angle = np.atan2(distance_y, distance_x) x_components = force * np.cos(angle) y_components = force * np.sin(angle)
The operations here could definitely be simplified. My linear algebra isn’t the strongest so I’m okay with writing a few extra lines of code.
Plot the data.
I’ll use a custom Matplotlib style that I’ll link at the bottom of this post. Create a figure and Axes instance for plotting.
import matplotlib.pyplot as plt
plt.style.use("grav_field.mplstyle")
fig, ax = plt.subplots()
quiver has four required parameters: x-position, y-position, x-component, and y-component. It needs to know where arrows are located along with their magnitude and direction. We pass in the two-dimensional arrays that we just created. pivot specifies which part of the arrow should touch the point—tip, tail, or middle. Tip works best for this application.
ax.quiver(x_grid, y_grid, x_components, y_components, pivot="tip")
Let’s also plot the sun as it appears in the GIF.
ax.scatter([STAR_POS_X], [STAR_POS_Y], color="#F5E342", s=STAR_MASS, edgecolor="black", linewidth=0.5)
Finally, set ticks, window limits, and save the figure.
ax.set(xticks=range(-100, 120, 20),
yticks=range(-100, 120, 20),
xlim=(-103, 103),
ylim=(-103, 103)
)
plt.savefig("grav_quiver.png")
The output.
Just as we expected, the gravitational field is strongest near the sun. Gravity’s strength depends on the square of the distance between the bodies, so it’s very strong when you pass nearby a large object but falls off rapidly as you drift further away. The comet in the GIF experiences a strong force (the large arrows) and is sharply accelerated.
I should point out that a gravitational field doesn’t have a fixed number of points like the grid of our plot. The field is continuous but we can’t plot an infinite number of arrows. We have to scale things down in order to visualize.
This is also an intentionally simple example that includes only one celestial body. All the arrows point toward the sun. But it wouldn’t be too complicated to add more bodies. You would simply calculate the field as we did above and repeat the process for every body, then add the arrays together.
As an example, here is a system with two equally massive stars:
Notice how the arrows tend to point toward the “average” of the two bodies. Directly between them, force vectors cancel each other out and there is minimal force applied. In a perfect system, some third body could hang out here without falling toward either star.
I hope this post demystifies meshgrid to some extent. Obviously it’s capable of a lot more but understanding just this much opens up a whole world of two-dimensional data. You can create quiver plots, heat maps, contour plots, and more.
Download the Matplotlib style.
Full code:
import numpy as np
import matplotlib.pyplot as plt
STAR_POS_X = 20
STAR_POS_Y = 40
STAR_MASS = 4000
COMET_MASS = 100
G_CONSTANT = 1.0
x_grid, y_grid = np.meshgrid(range(-100, 105, 5),
range(-100, 105, 5)
)
distance_x = STAR_POS_X - x_grid
distance_y = STAR_POS_Y - y_grid
distance = np.sqrt(distance_x**2 + distance_y**2)
distance = np.ma.masked_array(distance, distance <= 15)
force = G_CONSTANT * STAR_MASS * COMET_MASS / distance**2
angle = np.atan2(distance_y, distance_x)
x_components = force * np.cos(angle)
y_components = force * np.sin(angle)
plt.style.use("grav_field.mplstyle")
fig, ax = plt.subplots()
ax.quiver(x_grid, y_grid, x_components, y_components, pivot="tip")
ax.scatter([STAR_POS_X], [STAR_POS_Y], color="#F5E342", s=STAR_MASS, edgecolor="black", linewidth=0.5)
ax.set(xticks=range(-100, 120, 20),
yticks=range(-100, 120, 20),
xlim=(-103, 103),
ylim=(-103, 103)
)
plt.savefig("grav_quiver.png")