Simulating gravitational interaction in Python
Since the dawn of humanity, humans have been fascinated by the starry sky, constantly gazing upward, trying to understand what is up there and how celestial bodies move. Fortunately, we now have a better understanding of the cosmos and can accurately simulate celestial movements.
What’s the plan?
We’ll explore the key steps to create a gravitational interaction simulator between bodies. To accomplish this, we need to address an initial value problem in ordinary differential equations, although we won’t delve into the solution in this discussion. In this scenario, the bodies are treated as point-like rather than extended. From a theoretical perspective, Newton’s second law of mechanics and gravitational force are essential. However, these equations require a numerical method for solving since our problem can’t be solved analytically. In this instance, we’ve opted for the fourth-order Runge-Kutta method. The selection is based on its ease of implementation in our context and its widespread use in numerical methods.
To achieve this simulation, two coupled fourth-order Runge-Kutta methods are necessary, one for velocities and the other for positions. It’s crucial to highlight that Newton’s second law of mechanics is a second-order equation, necessitating the use of two coupled Runge-Kutta methods.
With the simulator, it is possible to simulate the Sun-Earth-Moon system (not to scale)
or a random three-body system
or a more than three bodies system
A brief summary of the necessary theory
These explanations will not be rigorous; therefore, I encourage you to review the references to delve deeper. In physics class, we understand that the change in position and velocity can be expressed as follows:
Note that the notation we’ll be using considers bold letters as vectors, representing the co-ordinates (x, y, z) at this moment. The dot above vectors indicates the derivative with respect to time, so if r is the position then r dot represents velocity. The first equation is the famous Newton’s second law of mechanics, relating the change in velocity of a body with mass (m) to the force exerted on it. The second equation, in (1), relates the change in position of a body with mass (m) to its velocity. These two equations are the ones our fourth-order Runge-Kutta methods will approximate. However, before delving into the numerical method, we need to introduce the interaction. As mentioned earlier, this interacion is the gravitational force
This equation may be more complicated than the one you remember; however, it simply tells us that the force that a body i feels from other bodies j depends on the distance (r) between them. G represents the gravitational constant. Then, by substituting this equation into (1) and renaming the equations, we obtain:
Notice that the two equations are coupled, meaning you need to solve them simultaneously. You can’t solve the first equation from (3) to determine the velocity and then apply it to the second one to find the position.
On the other hand, from the literature (again, for a deeper understanding, check the references), we can apply the Runge-Kutta method to a system of higher-order ordinary equations, which is precisely what we have in equation (3), 2-order to be precise. In general, if we have an independent variable t in a real interval [a, b] as follows:
Where N > 0 and h = (a-b)/N. Also, the alphas represent the initial values of the variable m at point 0. Then we can calculate the variable m at point j+1:
Where the k’s are the intermediate values needed to progress to the next step. Applying this result to our problem:
and
Now, we are expanding the notation. Previously, we stated that vectors hold values for either position or velocity in the co-ordinates (x, y, z). However, from this point onward, the vectors are arrays of size 1x(3N), meaning they contain information about the coordinate values for all the bodies. Finally, the subscript n indicates the current step, and n+1 denotes the step we aim to calculate. Also, as you can see, to calculate k2 for both position and velocity, we need k1 from both, and to calculate k3, again k2 is needed, and so on. This behavior happens because the equations (3) are coupled. These equations, (6) , (7) and (3) are the ones we’ll implement in our code.
To conclude this section, I want to emphasize that choosing initial values, the alphas, is not trivial for this problem. A small change in the initial positions and velocities will yield a completely different result, and there is no formula to determine these initial conditions.
From equations to code
The functional codes can be found in this GitHub repository. How to find the optimal number of steps for a given interval and the initial values of positions and velocities is beyond the scope of this explanation, but you can experiment with those parameters to observe the outcomes.
When all is said and done, it’s time to code. We’ll begin with the first equation of equation (3). We’ll create a function accelerations that is designed to compute the accelerations of a system of point-like bodies due to gravitational interactions. It iterates over each body (m), calculating the gravitational forces exerted on it by other bodies and updating the accelerations accordingly. The algorithm considers pairwise interactions between bodies, ensuring that each body's acceleration is influenced by the gravitational forces from all other bodies in the system.
def accelerations(t, r, m):
# Initialize a vector to store accelerations; r contains the coordinates of the n bodies.
a = np.zeros(len(r), np.longdouble)
# Calculate gravitational interaction among the n bodies.
i, j = 0, 0
while i < len(m): # Loop to obtain the acceleration of body i.
while j < len(m): # Loop to sum interactions between body i and bodies j.
if i == (len(m) - 1) and i == j: # Condition to end loops if we have reached the last body and the last interaction.
break
elif i == j: # Condition to avoid calculating the interaction of body i with itself.
j += 1
# Equation for gravitational interaction per unit mass i.
a[3*i:3*i+3] += -G * m[j] * (r[3*i:3*i+3] - r[3*j:3*j+3]) / (np.linalg.norm(r[3*i:3*i+3] - r[3*j:3*j+3])**3 + epsilon)
j += 1
i += 1
j = 0
return aDon’t forget the vectors contain information about the co-ordinates (x, y, z) of all the N bodies. Then, we write down equations (6) and (7), where accelerations is used
# Calculating the dynamics of the N bodies.
# 4 order RK for velocities and positions:
# Kiv : Velocities' ks.
# Kir : Positions' ks
k1v = dt*accelerations(t, r, masses)
k1r = dt*v
k2v = dt*accelerations(t + dt/2, r + k1r/2, masses)
k2r = dt*(v + k1v/2)
k3v = dt*accelerations(t + dt/2, r + k2r/2, masses)
k3r = dt*(v + k2v/2)
k4v = dt*accelerations(t + dt, r + k3r, masses)
k4r = dt*(v + k3v)
r += (k1r + 2*k2r + 2*k3r + k4r)/6
v += (k1v + 2*k2v + 2*k3v + k4v)/6
t += dtAt the end, as always, the time needs to be updated. And that’s it. These are the code blocks you have to program to simulate the gravitational interaction between bodies. Now, just store them in a file, and then you can plot them, animate them, or do whatever you like.
I’ve obtained the positions. Now, how do I animate them?
The animation was created using the well-known Matplotlib package, specifically the animation module (imported as animation from matplotlib). The following code segment is responsible for generating and saving the 3D animation. It utilizes Matplotlib's FuncAnimation class to create the animation and the PillowWriter for saving it as a GIF file.
# Arguments of fargs: (num,fps,time,bodies,data,nameBodies,colorBodies)
varAnimation = animation.FuncAnimation(fig, nBodies_Animation3D, \
fargs=(fps, times, numberBodies, positions, nameBodies, colorBodies), \
interval=60, frames=numPoints)
# Save animation.
writer_gif = animation.PillowWriter(fps=fps) # Function that will write the animation to be saved.
varAnimation.save(f"{folderPlot}\{fileName}.gif", writer=writer_gif) The animation.FuncAnimation line initializes the animation process, specifying the figure object to animate (fig), the updating function (nBodies_Animation3D which you can customize as desired; see the function below), and relevant parameters such as the frame interval (interval) and the total number of frames (frames). Additionally, it passes arguments that nBodies_Animation3Duses, such as frames per second (fps), time points (times), bodies (m), positions (positions), and visual properties for bodies such as names and colors (nameBodies, colorBodies).
As mentioned earlier, nBodies_Animation3D is responsible for updating a 3D plot during each frame of the animation. It iterates over the different bodies, plots their trajectories, and updates their end points. This is also the section where you can add customization to make your plot visually appealing.
def nBodies_Animation3D(num, fps, time, Bodies, data, nameBodies, colorBodies):
ax.clear() # Clears the figure to update the line, points, title, and axes.
j = 0
while j < Bodies: # While loop to plot the trajectories of different masses.
# loc creates a dataframe from the existing one. to_numpy converts a pandas dataframe to a numpy array.
# loc is inclusive with the indices.
rg = data.loc[:, 1 + 3*j : 3*j + 3].astype(float).to_numpy()
# Update trajectory line (num+1 due to how indices work in python).
# fps * num would mean drawing every fps plots, which speeds up the animation if there are too many points.
ax.plot3D(rg[:(fps * num) + 1, 0], rg[:(fps * num) + 1, 1], rg[:(fps * num) + 1, 2], \
linewidth=2, color=colorBodies[j])
# Update end point.
ax.scatter(rg[(fps * num), 0], rg[(fps * num), 1], rg[(fps * num), 2], \
marker='o', s=20, color=colorBodies[j])
j += 1
# Axis limits (optional).
# Set it to +- 10 in case the orbits are in a plane. For example, if they are in the x,y plane,
# then the limits of z would be (0,0), which cannot be plotted.
# To animate the sun-earth-moon system is better to use this.
# if np.amin(rg[:, 0 :: 3]) == np.amax(rg[:, 0 :: 3]): # Check if it's in the y,z plane.
# ax.set_xlim3d([-10, 10])
# ax.set_ylim3d([np.amin(rg[:, 1 :: 3]), np.amax(rg[:, 1 :: 3])])
# ax.set_zlim3d([np.amin(rg[:, 2 :: 3]), np.amax(rg[:, 2 :: 3])])
# elif np.amin(rg[:, 1 :: 3]) == np.amax(rg[:, 1 :: 3]): # Check if it's in the x,z plane.
# ax.set_xlim3d([np.amin(rg[:, 0 :: 3]), np.amax(rg[:, 0 :: 3])])
# ax.set_ylim3d([-10, 10])
# ax.set_zlim3d([np.amin(rg[:, 2 :: 3]), np.amax(rg[:, 2 :: 3])])
# elif np.amin(rg[:, 2 :: 3]) == np.amax(rg[:, 2 :: 3]): # Check if it's in the x,y plane.
# ax.set_xlim3d([np.amin(rg[:, 0 :: 3]), np.amax(rg[:, 0 :: 3])])
# ax.set_ylim3d([np.amin(rg[:, 1 :: 3]), np.amax(rg[:, 1 :: 3])])
# ax.set_zlim3d([-10, 10])
# else: # Not in one of the main planes; none of the coordinates are 0 throughout the trajectory.
# ax.set_xlim3d([np.amin(rg[:, 0 :: 3]), np.amax(rg[:, 0 :: 3])])
# ax.set_ylim3d([np.amin(rg[:, 1 :: 3]), np.amax(rg[:, 1 :: 3])])
# ax.set_zlim3d([np.amin(rg[:, 2 :: 3]), np.amax(rg[:, 2 :: 3])])
# Graph labels.
ax.set_title(str(format(time[fps * num], ".4g")) + ' days', fontsize=17, pad=0)
ax.set_xlabel(r'$x$ $(M_l)$', fontsize=10)
ax.set_ylabel(r'$y$ $(M_l)$', fontsize=10)
ax.set_zlabel(r'$z$ $(M_l)$', fontsize=10)
# ax.legend(loc='upper center', bbox_to_anchor=(0.5, 0.05), ncol=3)
# Remove grid, axes, or ticks.
plt.grid(False)
plt.axis('off')
ax.set_xticklabels([])
ax.set_yticklabels([])
ax.set_zticklabels([])These are the key components to create the GIFs you saw at the beginning. This serves as an initial approach to simulating interactions between bodies. You can progressively enhance it to make it more realistic. However, keep in mind that the running time of the simulation increases rapidly with more bodies or when incresing the simulation time. Consequently, more steps are needed for the simulation to converge to an accurate result. To optimize this, consider switching to faster languages like C/C++ or FORTRAN, or explore specialized numerical methods for this problem.
Note: The GIF of the Sun-Earth-Moon system exaggerated the distances because the distance between the Sun and Earth is 400 times greater than the distance between Earth and the Moon, so you’ll see the earth and moon orbits as one if you don’t exaggerated them. In the GitHub repository, you’ll find a script that creates a file with exaggerated distances to visualize the Moon’s orbit as the GIF of the Sun-Earth-Moon system.
References
- Barger, V. and Olsson, M. (1954). CLASSICAL MECHANICS : A Modern Perspective Second Edition.
- Hand, L. N. and Finch, J. D. (1998). Analytical Mechanics. Cambridge University Press.
- Chow, T. L. (2013). Classical Mechanics, Second Edition. Taylor & Francis.
- Burden, R. L. and Faires, J. D. (2010). Numerical Analysis. Cengage Learning.
- Epperson, J. F. (2013). An Introduction to Numerical Methods and Analysis. Wiley.
- Faranak, R., Fudziah, I., and Mohamed, S. (2013). Improved Runge-Kutta methods for solving ordinary differential equations. Sains Malaysiana, 42(11):1679–1687.
