Loading spinner
N-Body Orbit Simulation
Oct. 2nd, 2022ยท48 min read

Introduction

This post will introduce the principles behind predicting the motion of gravitationally interacting objects. Using Python with some popular modules, we'll implement the principles from scratch.

Prior Knowledge

There are a few concepts that you should know before coming into this post, in general I try to explain everything clearly but it helps if you have prior exposure to these concepts.

  • NumPy Matrices: We will use NumPy as the primary data storage and computation backbone. If you're not familiar with NumPy matrices take a look at this resource.
  • Numerical Integration for ODEs: We use numerical integrators to solve this problem but we're building off prior knowledge of simpler cases from my blog post from before.
  • Gravitational Laws: To be precise, we will be using Newton's law of universal gravitation that you may have seen in a physics course or you can review it here.

N-Body Problem

The n-body problem is to predict the individual motions of a group of objects interacting gravitationally. More precisely, given the initial conditions of a set of objects the goal is to predict the state of these objects over time when influenced by the gravity of all neighboring objects. Fields such as astronomy, space flight, and others that rely on the understanding of the positions and trajectories of celestial bodies utilize solutions to this problem. What makes this difficult is that for systems of more than two bodies, the math cannot be solved explicitly and therefore must be simulated.

The next few sections will focus on concretely defining the core operations, principles, and theories that we will use moving into the implementation.

General Goals

As with many problems, there are different methodologies behind how you may solve it. When developing this solution I focused on a few core principles.

  • Readable: Many operations can be written to run marginally faster at the expense of readability. As this post is meant to educate I focus on making code that is as self-explanatory as possible.
  • Adaptive: Depending on what you are simulating it may not make sense to generalize the integrator to only work in 3D. I wanted to produce a solution that adapts depending on the inputs and only computes the number of dimensions given.

Notation

I'll summarize the notation used for vectors, variable assignments, and some other structures that might seem foreign but will be explained through reading the post.

NotationDescriptionDefinition
Vector
Vector from 1 to 2
Vector Magnitude
Unit Vector
Position Vector in 3D
Velocity Vector in 3D
Acceleration Vector in 3D
Vector Derivative
Object State
System State
State Derivative

Gravity Laws

In this section, I will give a brief overview of Newton'zs law of universal gravitation and extend the concept of gravitation between two bodies to all of our bodies.

Law of Universal Gravitation

Newton's law of gravitation states that any two bodies in space will be attracted to each other, where and are the masses of the two bodies, is the distance separating them, and is the gravitational constant. I've given the equation below in the scalar and vector form.

This equation only gives us the magnitude of the force, but it would be helpful to have the force as vector. If we multiply the prior with the unit vector pointing from body 1 to body 2 we can get the force acting on body 1 due to object 2 as a vector.

The force of attraction between these two bodies is always on the line between the two bodies and we know that the magnitude of the force acting on each body is the same. Given this we can get the force vector for the force acting on body 2 due to body 1 as the negated version of the force on body 1.

Force on One Body

Now that we know how to determine the gravitational force on a body due to a particular neighboring body, we can extend this to more than 2 bodies. Consider a set of objects with masses where ranges from 1 to n. The total gravitational force on a body is the sum of the gravitational forces due to all other bodies.

In plain english, the net force on a particular body is the sum of the force equation evaluated between the body of interest and all other bodies. In the equation, the is what maintains that we do not calculate the attraction of a body onto itself as a body cannot attract itself and we would be dividing by 0 when calculating the force.

Acceleration on One Body

From the net force it's pretty simple to determine the net acceleration. Using Newton's second law we can add an additional equality to our prior expression. With this we can move around some terms and determine the net acceleration.

We can actually use this to summarize the n-body problem as a system of equations and provide the differential equation that we intend to integrate in a later section.

Recall from the prior numerical integrator post that we integrated a similar differential equation for a mass-spring system but only in 1 dimension! Something that I find beautiful is that our code extends super well into the higher dimensionality cases.

Acceleration for All Bodies

Clearly, we need to evaluate the net acceleration for all the bodies in our system to know how each of them are going to move. The simple way to do this would be evaluate the prior expression for each of the bodies but in practice this is quite inefficient as there's repeated work being done. Let's see why this is the case and how we can improve it.

Brute Force Method

Let's consider computing the attraction force between two bodies as a basic operation. We want to minimize the number of times we have to do this to improve the speed of our code. In this simple case, let's call it the brute force method, we compute the net force for each body as in independent operation. To do so, we need to compute forces for each of the bodies. Overall the number of force calculations () scales with the number of bodies quadratically.

Python pseudocode for this would look something like the following where we can see that compute force would get called times. Don't focus too much on the details of the actual functions or variables but think more about how many times we get to the line where we compute force.

Combination Method

In this method, we will try to leverage the property that each force we compute can be negated to provide a second force that we can use. In that case, we can resolve two forces in our sets of forces that we need to calculate and reduce our force calculations by half.

In practice, all this entails is computing the forces for all combinations of bodies! Looking at the equation for number of combinations we can see that this reduces to the prior assertion for operation count. To be precise, we'll be looking at all combinations of size 2 () from the set of all our bodies. I will refer to the set of combinations the pairs of bodies () throughout the rest of the post.

Python pseudocode for this would look something like the following where we reduce the number of calls to compute_force. The biggest takeaway here is that by looking at the combinations of bodies we can avoid repeat force calculations.

Operation Count Comparison

When comparing the different operation counts as the number of bodies increases we can see that there is a formidable gain in using the combination method over the brute force method.

Operation Count Comparison
Operation Count Comparison

Those are familiar with big- notation may have noticed the both the implementations are still quadratic in nature and are therefore the same time complexity. However, if you were to measure the actual runtime we would see a gain. This is an example where trying to compare time complexity to runtime can get messy.

Runtime Comparison

Like I mentioned before, the time complexity can be a good indicator of how we would expect the runtime to scale but there are some factors in play that can make them different.

  • In big- we drop constants and lower order terms
    • Ex: =
    • Constants can be play a large role in runtime (consider vs )
  • We may adjust an algorithm where each "iteration" takes longer to run
    • We reduced iteration count in our case by 50% but each iteration has more work

Plotting the runtimes of our two different variants we get the following graph.

Comparison of Runtimes
Comparison of Runtimes

We can see that when we normalize the time and theoretical iteration count on the brute force method the times line up for the brute force (blue) meaning that the trend is correct. When we look at the combination method what we see is that the runtimes do not follow the expected trend, why is this? In our case, our algo has more work to do per iteration than the brute force method so even though we expect the runtime to be 50% of the brute force we see that it's actually around ~63%.

All of our analysis is still correct but it's important to understand the limitations of big- in terms of comparing such similar algorithms. If we managed to develop an algorithm that ran in time complexity we would not be nearly as concerned about what the true value of the constants are.

System Energies

Clearly the system starts with a set amount of energy based on the positions, masses, and velocities that manifest as kinetic energy and potential energy. We don't need to compute these to execute a simulation but as we expect the system to maintain energy this can be a good check to see how we're doing.

Kinetic Energy

Kinetic is really easy! We just need to sum up the result of a classic physics equation.

Potential Energy

The gravitational potential energy between two objects is something that is really familiar! Just as we're only looking at the combinations of bodies to determine our forces we will do the same here where and will be the indices for our bodies out of our set of combinations ().

System State

I have been throwing around the word state a bit loosely, let's try to refine our understanding of what this means and why it's useful for us. The 30 second answer is that a state is a matrix that maintains all of the important values that describe the aspects of our system such as the positions and velocities of all of our objects.

States are a joint representation of everything going on in our system. They help us simplify our integration operations and keep track of our variables.

Structure

As I mentioned, a state is merely a matrix but the structure of how its constructed and what it contains is the important part. For our physical modeling cases, the state will contain the position and velocities of all the objects where each row is the state of one object and each column is a variable of importance (think a certain position/velocity).

The prior would be the state for a 2D case with two objects. Don't confuse with , the capital is the state and has nothing to do with the position variable .

State Derivative

Along with the state, there is a accompanying state derivative , it might sound scary but it isn't. Throughout the integration process we make incremental changes to the state, but doing so for each value independently would be annoying and inefficient. The state derivative is a way to encode the corresponding derivative for each of the values in the state so that when we need to do update steps we can proceed with them as matrix sums. For our purposes, the structure would look like the below for the same case as before.

Advantages

For simple systems, it might be easier to just build out a numerical integrator for each of the variables of interest, but for large n-body systems this would get messy quickly. Consider an example in 3D with 50 bodies, this would mean that we have 300 variables of interest and therefore 300 integrators! When we represent this system in the state matrix form, there are still the 300 values being updated but it's reduced to a series of sums between the entire matrices which is much cleaner (and can be vectorized).

Simple Example

Let's do a simple two-body example in 2D to see how this is going to play out in our code! Let's consider two bodies with the properties in the below table.

BodyPosition Velocity
1
2

State Representation

If we follow the general structure shown before it's quite evident how this would be constructed. Recall that each row contains the variables needed to describe a single object's condition in our system.

State Derivative

The goal is to fill out the prior structure which requires us to compute the acceleration on each body. Let's assume that the only force involved in this case is the gravitational attraction between the bodies, for the sake of cleaner math let's assume and that the masses are both .

We know from before that the force for the other body in the force pair is just the negated version of this force so we can also compute the force for body 1.

As the masses are just 1, the value of the acceleration and the force are the same so we can go ahead and fill out the state derivative.

Example Euler Integration Step

Now that we have the state and the state derivative, we can actually proceed with doing a simple Euler integration step. Let's consider a step size of , we can compute the new state with the following.

In this post, we will be using a better integrator then Euler, but this simple example presents the core workflow of how we will represent our system and the role that the state derivative plays in propagating the state forward.

Integrator Refresher

Although it helps to have read my post on integrators already, I'll do a short summary of the usage of the RK4 integrator that we are using in this post. If you are already comfortable doing the example cases from that post you can skip ahead, but I would recommend you read regardless as I also wrote this example as a precursor to the later structure.

General Form

The basic form of the RK4 integrator for a given state is as follows. The integrator will estimate what the new state will be one later (in our case our unit is time so this would be a timestep). In simpler terms, given the values of our system at a time, we can use the integrator to see how these values would change if we progress one timestep into the future.

This is best defined as a function in code as follows.

Mass-Spring System

Let's go ahead and consider a simple differential equation for a mass and spring system with no damping. Note that the smaller in this equation is not a state but just a single variable.

Our initial conditions for the system will be the position and velocity at . The state and the corresponding state derivative would be the following.

Let's go ahead and write our differential equation into a function that will return the state derivative for a given state. In general, when your system of equations changes this will manifest as a change to the function that returns the state derivative. For example, if this system also had a damper we would have to include the changes to the ODE in the calculation for xdd.

We can now execute a simple simulation and plot the results against the known solution to this which is for these constants.

This should yield the following plot where we can see that the simulated points match up with the true result quite well!

Mass Spring Position
Mass Spring Position

I point out this example because the form of the ODE is more similar to the n-body problem then you may have thought. The prior formulation of the state and state derivative are very similar to what we do for the n-body case.

Executing a Simulation

The last step before we can move onto the implementation is just seeing how we stack up all of our prior operations to actually run a simulation.

Number of Iterations

The first step is to determine how many simulation steps we're going to be taking. This is trivially done with the given end time and the time step . We use the ceiling function here to run a simulation slightly longer than needed as opposed to too short.

Initializations

We have a few additional initializations that make sense to run here and not in init.

History

We'll store the history of the states in one large history matrix of size [Iterations +1, N, D]. Essentially it's a stack of (Iterations + 1) states indexed by the simulation iteration.

Body Pairs

Across multiple different operations we utilize the sets of body pairs based on the combination. It's helpful to pre compute them so that it's not a repeated operation.

Energies

Similarly to the state history, we'll initialize a history of energies of size [Iterations + 1, 3] as we store kinetic, potential, and total energy. This step is done after the body pairs as the energy calculations depend on the pairs.

Running Iterations

We now have everything we need to execute a simulation! All we do now is call our integrator function to step the state ahead by on timestep and save the new state into our history. Along the way we also compute the energy for each of these new states and store that too.

Implementation

Time for the fun stuff, in this section we'll go over how to take these concepts and implement them in Python. Let's go ahead and start with importing the Python modules that we intend to use.

Simulation Class

Many of the operations and variables we use are reused across the simulation, so it helps to encapsulate it all into a class to clean up many of the references. Below is the framework that we'll fill out in the later steps, so if you want to follow along go ahead and copy this before we start.

  • init: Initializes some of the reused variables across our class
  • get_energy: Returns the kinetic and potential energy of the system at a given state
  • get_state_deriv: Determines the state derivative of a given state
  • rk4: Integrates the input state to a new state one timestep ahead
  • run_simulation: Calls all requisite functions to produce a sequence of states up to a time T

The main benefit of doing this as a class is that we can provide instances for different types of simulations or different initial conditions. Also there are many variables that are static and used by different functions so this reduces the amount of repeat work (think about how many times we use self.pairs).

Init

Some of this may not be immediately clear but you'll see why we're doing this as we develop more of the core functions.

System Energy

The prior formulations maps really well into the code! While there are faster ways to compute these by using vectorization and maintaining matrices as opposed to iterating but this is way easier to understand which is why I opted to present it this way.

State Derivative

Take some time to really understand this section as it's the meat of the conversation today. I personally think it's much easier to understand what is happening when looking at the code.

RK4 Integrator

This is a direct copy from my prior work. The purpose of the integrator is to take in a current state and advance the state to where it should be one timestep (dt) later. This of course relies on being able to compute the state derivative which is what the function we pass in as evaluate is needed for.

Generally speaking, the RK4 integrator is a good all-around integrator for various types of problems (to me it's reliable but likely not optimal depending on the case). As an example, when simulating Hamiltonian systems (often called conservative systems) one can use special integrators that preserve the constant energy nature of these systems. I won't get into it here, but I'm just trying to imply that there isn't one best integrator.

Running the Simulation

Now that we have all the requisite functions defined, we can go ahead and piece it together here in the main simulation call. As we are using a class to hold all of our information, I added some assertions to the start to ensure we have the required arguments to execute a simulation.

Simple Two Body & Verification

Before we step into larger and more interesting examples, it helps to validate that the work that we did is providing results that we think are correct. Let's consider a simple case of a small satellite orbiting in a circular low earth orbit (LEO), akin to a Starlink satellite.

Initial Conditions

For this simple case we can make a useful function to help us determine what the initial conditions would be for a circular orbit. When an object is orbiting a body, there is a balancing act between the centripetal force and gravitational force. A circular orbit is constantly in this balance, so we can set the force equal to each other and solve for the needed velocity for a given input radius.

Where is the mass of the body that our satellite is orbiting. Generalizing this for a satellite orbiting earth we can make a helpful function for setting up initial conditions.

With this function we can easily define our initial conditions in terms of an orbital altitude. The first body in our state is the earth at origin not moving (hence ) and the second body is our satellite.

Running Simulation

Now that we have the initial conditions and masses arranged, we can go ahead and initialize an instance of our NBody class and call a simulation. In this case given the expected period of the orbit, I opted for a 500 minute end time at a 1 second timestep to get a few periods without an excessive runtime.

Now that we have the simulation history, we can go ahead and plot the example with the following code.

Satellite Orbiting Earth
Satellite Orbiting Earth

Validation

Looks great! Let's dive a bit deeper and verify the orbital period and system energy. Let's consider some simple checks to see if the integrator is performing as we'd expect.

Orbital Period

For the period evaluation we can make a pretty simple plot to give us confidence that this works. First, let's determine what the period is supposed to be by computing how long it would take to traverse a circle with a radius of our orbit when traveling at our initial velocity. Then we can plot how far the satellite is from the initial position, this way when we return back to where we started we'll see a minimum in our graph that should reflect a period.

Evaluation of Period
Evaluation of Period

Everything seems in order for the periods so let's move onto looking at the energy of the system.

Energy Validation

Energy validation is even simpler, we already compute the energy alongside the history of states so we can just plot based on that. It's a given that we expect the total energy of the system to remain constant as there are no losses but what do we expect the kinetic energy and potential energy to look like?

As the satellite is essentially the only object moving (Earth is getting pulled slightly but it's negligible) and it's orbit is circular we would expect the velocity and therefore the kinetic energy to remain constant. This would also mean that the potential energy would be constant so we should expect 3 constant values for KE, PE, and TE.

Energy Trends For Two Body
Energy Trends For Two Body

Great! Seems like everything is working as expected so we can move on to some more interesting examples.

Special Three Body

In this section, we'll have fun with an N-body choreography with three bodies. The history of them is quite short and many fun initial conditions have already been found. We're focusing on one that should produce a stable figure 8!

It's pretty typical for these to use

Initial Conditions

I have provided the initial conditions and masses for this system in the form needed for our solution below. I chose this choreography as many of the other examples are quite unstable and may not provide good results.

Running Simulation

Just as before, we can construct an instance of our NBody class and proceed with the simulation!

We can go ahead and plot the results to see if we get the correct result.

Stable Figure 8 Choreography
Stable Figure 8 Choreography

That looks great! I took the liberty of producing an animation of their motion and the associated energy as I think it's a great way to appreciate this case.

Partial Solar System

Lastly, let's consider a partial solar system example. I limited this to a partial simulation as the periods of orbit for some of the outermost planets are quite long and I don't think it's reasonable for most people to run a full period of these orbits on their computers.

Initial Conditions

Much like the figure-8 case, I was able to find the initial conditions online and I have provided them in the format needed for my simulation structure below.

Running Simulation

The process here is identical to before, but given the scale of the problem notice how the timestep is much larger now. I try to scale my timestep based on the periods where over a period I want a certain number of steps. Playing with this more you can develop better intuition behind this.

Now that we have the history we can go ahead and produce a test plot.

Partial Solar System
Partial Solar System

In my animation, I did some extra work to try and scale the sizes of the planets with some respect to their actual sizes but I opted not to for the example here as it just adds unnecessary complexity for plotting the example for validation purposes.

Little challenge for you, plot the energy for this system and see if it makes sense to you.

Closing Remarks

This one was quite a bit to write but I hope you enjoyed reading it as much as I did making it. Let me know if you have any suggestions or comments and I'd be happy to consider them.

Where Are You?
Where Are You?