One of the uses of programming is to help us understand the real world through simulation. This technique is used in science, finance, and many other quantitative fields. As long as the “rules” which govern the real-world properties are known, you can write a computer program that explores the outcomes you get from following those rules. In this article, you’ll **simulate a 3D solar system in Python** using the popular visualisation library Matplotlib.

If you want to start with a simpler version of this project, you can read the first article in the Orbiting Planets Series. The first article deals with simulating orbiting planets in 2D and uses the relatively simple `turtle`

graphics module. This article is the second in the series and will define classes that are modelled on the ones used in the 2D version. However, you don’t need to have read and followed the first article. If you prefer, you can jump straight into the 3D version in this article.

By the end of this article, you’ll be able to create your own 3D solar system in Python with as many suns and planets as you wish. Here’s an example of a simple solar system with one sun and two planets:

You’ll also be able to turn on a 2D projection on the floor of the animation to show the 3D nature of the solar system better. Here’s the same solar system simulation, including the 2D projection:

### Outline of the article

Here’s an outline of this article so that you know what’s coming:

- A brief discussion about the
**gravitational attraction between two bodies**which you’ll need to use for simulating a 3D solar system in Python. - A brief introduction to
**vectors in 3D**. **Definition of classes for the solar system and the orbiting bodies**within it, such as suns and planets. You’ll write these classes in a step-by-step approach and test them with a simple solar system.- Addition of the
**option to show a 2D projection**of the orbiting bodies along with the 3D simulation. This 2D projection helps to visualise the motion in 3D. - Creation of a
**binary star system**.

You’ll use object-oriented programming and Matplotlib in this article. If you wish to read more about either topic, you can read:

Let’s make a start with simulating a 3D solar system in Python using Matplotlib.

## Let’s Talk About Gravity

Suns, planets, and other objects in a solar system are bodies that are in motion and that attract each other due to the gravitational force exerted between any two objects.

If the two objects have masses m_1 and m_2 and are a distance of r away, then you can calculate the gravitational force between them using the following equation:

F=G\frac{m_1m_2}{r^2}

The constant G is a gravitational constant. You’ll see how you’ll be able to ignore this constant in the version of the simulation you’ll write in this article in which you’ll use arbitrary units for mass and distance rather than kg and m.

Once you know the gravitational force between two objects, you can work out the acceleration a each object undergoes due to this gravitational force using the following formula:

F=ma

Using this acceleration, you can adjust the velocity of the moving object. When the velocity changes, both the speed and the direction of travel will change.

## Representing Points and Vectors in 3D

When simulating a 3D solar system in Python, you’ll need to represent the solar system as a region of space using three dimensions. Therefore, each point in this 3D space can be represented using three numbers, the *x*-, *y*-, and *z*-coordinates. For example, if you wish to place a sun in the centre of the solar system, you can represent the sun’s position as *(0, 0, 0)*.

You’ll also need to represent vectors in 3D space. A vector has both magnitude and direction. You’ll need vectors for quantities such as velocity, acceleration, and force since these quantities all have a direction as well as a magnitude.

I won’t be discussing vector algebra in detail in this article. Instead, I’ll state any results that you’ll need as and when you need them. You can read more about vectors and vector algebra if you wish.

To make working with vectors easier in the code, you can create a class to deal with them. Writing this class will serve as a quick refresher on classes and object-oriented programming. You can read about object-oriented programming in Python if you feel you need a more thorough explanation. Although you can also create a class to deal with points in 3D space, this isn’t necessary, and I won’t be creating one in this article.

## Creating The `Vector`

Class (aka Reviewing Classes)

*If you’re familiar with vectors and object-oriented programming, you can skip this section and just review the code at the end defining the Vector class.*

Create a new file called `vectors.py`

in which you’ll define the `Vector`

class. You’ll use this script to define the class and test it out. You can then delete the testing code at the end and leave just the class definition in this script:

# vectors.py class Vector: def __init__(self, x=0, y=0, z=0): self.x = x self.y = y self.z = z def __repr__(self): return f"Vector({self.x}, {self.y}, {self.z})" def __str__(self): return f"{self.x}i + {self.y}j + {self.z}k" # Testing Vector Class - TO BE DELETED test = Vector(3, 5, 9) print(test) print(repr(test)) test = Vector(2, 2) print(test) print(repr(test)) test = Vector(y=5, z=3) print(test) print(repr(test))

The `__init__()`

method for the `Vector`

class has three parameters representing the value along each axis. Each parameter has a default value of `0`

representing the origin for that axis. Although we prefer not to use single letter names in Python, `x`

, `y`

, and `z`

are appropriate as they represent the terms commonly used in mathematics for the Cartesian coordinate system.

You’ve also defined the two dunder methods to represent the object as a string:

`__repr__()`

returns an output intended for a programmer showing the class name. The output from`__repr__()`

can be used to recreate the object.`__str__()`

returns a non-programmer’s version of the string representation of the object. In this case, it returns a representation that’s commonly used in maths to represent vectors, using the unit vectors**i**,**j**, and**k**.

You can read more about the differences between the two types of string representations in the Snippets section at the end of Chapter 9 in The Python Coding Book.

The output from the testing code block is the following:

3i + 5j + 9k Vector(3, 5, 9) 2i + 2j + 0k Vector(2, 2, 0) 0i + 5j + 3k Vector(0, 5, 3)

### Making the `Vector`

class indexable

In this 3D solar system in Python project, it would be convenient if the `Vector`

class was indexable so that you can use the `[]`

notation with an index to extract one of the values. With the class in its current form, if you add `print(test[0])`

in your script, you’ll get a `TypeError`

saying that the `Vector`

object is not subscriptable. You can fix this by adding another dunder method to the class definition:

# vectors.py class Vector: def __init__(self, x=0, y=0, z=0): self.x = x self.y = y self.z = z def __repr__(self): return f"Vector({self.x}, {self.y}, {self.z})" def __str__(self): return f"{self.x}i + {self.y}j + {self.z}k" def __getitem__(self, item): if item == 0: return self.x elif item == 1: return self.y elif item == 2: return self.z else: raise IndexError("There are only three elements in the vector") # Testing Vector Class - TO BE DELETED test = Vector(3, 5, 9) print(test[0])

By defining `__getitem__()`

, you’ve made the `Vector`

class indexable. The first item in a vector is the value of *x*, the second is the value of *y*, and the third is the value of *z*. Any other index will raise an error. The output from the testing code block is the following:

3

`test[0]`

returns the first item in the vector, the value for *x*.

### Defining addition and subtraction in the `Vector`

class

You can define addition and subtraction for objects of a class by defining the `__add__()`

and `__sub__()`

dunder methods. These methods will enable you to use the `+`

and `-`

symbols to perform these operations. Without these dunder methods, using `+`

and `-`

raises a `TypeError`

.

To add or subtract two vectors, you can add or subtract each element of the vectors separately:

# vectors.py class Vector: def __init__(self, x=0, y=0, z=0): self.x = x self.y = y self.z = z def __repr__(self): return f"Vector({self.x}, {self.y}, {self.z})" def __str__(self): return f"{self.x}i + {self.y}j + {self.z}k" def __getitem__(self, item): if item == 0: return self.x elif item == 1: return self.y elif item == 2: return self.z else: raise IndexError("There are only three elements in the vector") def __add__(self, other): return Vector( self.x + other.x, self.y + other.y, self.z + other.z, ) def __sub__(self, other): return Vector( self.x - other.x, self.y - other.y, self.z - other.z, ) # Testing Vector Class - TO BE DELETED test = Vector(3, 5, 9) + Vector(1, -3, 2) print(test) test = Vector(3, 5, 9) - Vector(1, -3, 2) print(test)

Both `__add__()`

and `__sub__()`

return another `Vector`

object with each element equal to the addition or subtraction of the respective elements in the two original vectors. The output is the following:

4i + 2j + 11k 2i + 8j + 7k

You can do the same for multiplication and division, although these operations need more care when dealing with vectors.

### Defining scalar multiplication, dot product and scalar division in the `Vector`

class

You cannot just refer to ‘multiplication’ when dealing with vectors as there are different types of ‘multiplication’. In this project, you’ll only need scalar multiplication. Scalar multiplication is when a vector is multiplied by a scalar (which has a magnitude but no direction). However, in this subsection, you’ll also define the dot product of two vectors. You’d like to use the `*`

operator for both scalar multiplication and the dot product. Therefore, you can define the `__mul__()`

dunder method:

# vectors.py class Vector: def __init__(self, x=0, y=0, z=0): self.x = x self.y = y self.z = z def __repr__(self): return f"Vector({self.x}, {self.y}, {self.z})" def __str__(self): return f"{self.x}i + {self.y}j + {self.z}k" def __getitem__(self, item): if item == 0: return self.x elif item == 1: return self.y elif item == 2: return self.z else: raise IndexError("There are only three elements in the vector") def __add__(self, other): return Vector( self.x + other.x, self.y + other.y, self.z + other.z, ) def __sub__(self, other): return Vector( self.x - other.x, self.y - other.y, self.z - other.z, ) def __mul__(self, other): if isinstance(other, Vector): # Vector dot product return ( self.x * other.x + self.y * other.y + self.z * other.z ) elif isinstance(other, (int, float)): # Scalar multiplication return Vector( self.x * other, self.y * other, self.z * other, ) else: raise TypeError("operand must be Vector, int, or float") # Testing Vector Class - TO BE DELETED test = Vector(3, 5, 9) * Vector(1, -3, 2) print(test) test = Vector(3, 5, 9) * 3 print(test)

The result of using the `*`

operator will depend on whether the second operand, the one following the `*`

symbol, is a scalar or a vector. If the second operand, represented by the parameter `other`

, is of type `Vector`

, the dot product is calculated. However, if `other`

is of type `int`

or `float`

, the returned result is a new `Vector`

, scaled accordingly.

The output from the code above is the following:

6 9i + 15j + 27k

If you want scalar multiplication, the scalar needs to come *after* the `*`

symbol. If you attempt to run the statement `3*Vector(3, 5, 9)`

instead, a `TypeError`

will be raised since the `Vector`

class is not a valid operand for using `*`

with objects of type `int`

.

Two vectors cannot be divided. However, you can divide a vector by a scalar. You can use the `/`

operator with the `Vector`

class if you define the `__truediv__()`

dunder method:

# vectors.py class Vector: def __init__(self, x=0, y=0, z=0): self.x = x self.y = y self.z = z def __repr__(self): return f"Vector({self.x}, {self.y}, {self.z})" def __str__(self): return f"{self.x}i + {self.y}j + {self.z}k" def __getitem__(self, item): if item == 0: return self.x elif item == 1: return self.y elif item == 2: return self.z else: raise IndexError("There are only three elements in the vector") def __add__(self, other): return Vector( self.x + other.x, self.y + other.y, self.z + other.z, ) def __sub__(self, other): return Vector( self.x - other.x, self.y - other.y, self.z - other.z, ) def __mul__(self, other): if isinstance(other, Vector): # Vector dot product return ( self.x * other.x + self.y * other.y + self.z * other.z ) elif isinstance(other, (int, float)): # Scalar multiplication return Vector( self.x * other, self.y * other, self.z * other, ) else: raise TypeError("operand must be Vector, int, or float") def __truediv__(self, other): if isinstance(other, (int, float)): return Vector( self.x / other, self.y / other, self.z / other, ) else: raise TypeError("operand must be int or float") # Testing Vector Class - TO BE DELETED test = Vector(3, 6, 9) / 3 print(test)

And the output is:

1.0i + 2.0j + 3.0k

### Finding the magnitude of a vector and normalizing a vector

If you have a vector (x, y, z), you can find its magnitude using the expression \sqrt(x^2 +y^2 + z^2). You can also normalize a vector. Normalization gives a vector with the same direction but with a magnitude of `1`

. You can calculate the normalized vector by dividing each element of the vector by the vector’s magnitude.

You can define two new methods to complete the `Vector`

class:

# vectors.py import math class Vector: def __init__(self, x=0, y=0, z=0): self.x = x self.y = y self.z = z def __repr__(self): return f"Vector({self.x}, {self.y}, {self.z})" def __str__(self): return f"{self.x}i + {self.y}j + {self.z}k" def __getitem__(self, item): if item == 0: return self.x elif item == 1: return self.y elif item == 2: return self.z else: raise IndexError("There are only three elements in the vector") def __add__(self, other): return Vector( self.x + other.x, self.y + other.y, self.z + other.z, ) def __sub__(self, other): return Vector( self.x - other.x, self.y - other.y, self.z - other.z, ) def __mul__(self, other): if isinstance(other, Vector): # Vector dot product return ( self.x * other.x + self.y * other.y + self.z * other.z ) elif isinstance(other, (int, float)): # Scalar multiplication return Vector( self.x * other, self.y * other, self.z * other, ) else: raise TypeError("operand must be Vector, int, or float") def __truediv__(self, other): if isinstance(other, (int, float)): return Vector( self.x / other, self.y / other, self.z / other, ) else: raise TypeError("operand must be int or float") def get_magnitude(self): return math.sqrt(self.x ** 2 + self.y ** 2 + self.z ** 2) def normalize(self): magnitude = self.get_magnitude() return Vector( self.x / magnitude, self.y / magnitude, self.z / magnitude, ) # Testing Vector Class - TO BE DELETED test = Vector(3, 6, 9) print(test.get_magnitude()) print(test.normalize()) print(test.normalize().get_magnitude())

The testing code gives the following output:

11.224972160321824 0.2672612419124244i + 0.5345224838248488j + 0.8017837257372732k 1.0

The third output gives the magnitude of the normalized vector, showing that its magnitude is `1`

.

Depending on what IDE or other tools you’re using, you may get a warning when dividing `self.x`

, `self.y`

, and `self.z`

, such as in `__truediv__()`

and `normalize()`

. You don’t need to worry about this, but if you’d like to fix it, you can do so by changing the `__init__()`

signature to either of the following:

def __init__(self, x=0.0, y=0.0, z=0.0):

or

def __init__(self, x:float=0, y:float=0, z:float=0):

Both options let your IDE know that the arguments should be floats. In the second option, you’re using type-hinting to do so.

You can now delete the testing code at the end of this script so that all you have in `vectors.py`

is the class definition.

## Simulating a 3D Solar System in Python

Now, you can start working on the 3D solar system in Python. You’ll create two main classes:

`SolarSystem`

: this class takes care of the solar system, keeps track of how many bodies there are within it and the interactions between them.`SolarSystemBody`

: this class deals with each individual body in the solar system and the body’s movement.

You’ll use Matplotlib to create and visualise the solar system. You can install Matplotlib by using the following in the Terminal:

$ pip install matplotlib

or

$ python -m pip install matplotlib

The `Axes3D`

object in Matplotlib will ‘host’ the solar system. If you’ve used Matplotlib and mostly used 2D plots, you would have used (knowingly or unknowingly) the `Axes`

object. `Axes3D`

is the 3D equivalent of `Axes`

, as the name implies!

It’s time to get started writing and testing these classes. You can create two new files:

`solar_system_3d.py`

will contain the class definitions.`simple_solar_system.py`

will contain the code to create a solar system. You’ll use this file to test the classes as you write them, leading towards creating a simple solar system with one sun and two orbiting planets.

Next, you’ll start working on the `SolarSystem`

class.

### Setting up the `SolarSystem`

class

You’ll use arbitrary units throughout this project. This means that rather than using meters for distances and kilograms for masses, you’ll use quantities with no units. The parameter `size`

is used to define the size of the cube that will contain the solar system:

# solar_system_3d.py class SolarSystem: def __init__(self, size): self.size = size self.bodies = [] def add_body(self, body): self.bodies.append(body)

You define the `SolarSystem`

class with an `__init__()`

method which includes the parameter `size`

. You also define the `bodies`

attribute. This attribute is an empty list that will contain all the bodies within the solar system when you create them later on. The `add_body()`

method can be used to add orbiting bodies to the solar system.

The next step is to introduce Matplotlib. You can create a figure and a set of axes using the `subplots()`

function in `matplotlib.pyplot`

:

# solar_system_3d.py import matplotlib.pyplot as plt class SolarSystem: def __init__(self, size): self.size = size self.bodies = [] self.fig, self.ax = plt.subplots( 1, 1, subplot_kw={"projection": "3d"}, figsize=(self.size / 50, self.size / 50), ) self.fig.tight_layout() def add_body(self, body): self.bodies.append(body)

You call `plt.subplots()`

, which returns a figure and a set of axes. The values returned are assigned to the attributes `fig`

and `ax`

. You call `plt.subplots()`

with the following arguments:

- The first two arguments are
`1`

and`1`

to create a single set of axes in the figure. - The
`subplot_kw`

parameter has a dictionary as its argument, which sets the projection to 3D. This means the axes created are an`Axes3D`

object. `figsize`

sets the overall size of the figure containing the`Axes3D`

object.

You also call the method `tight_layout()`

. This is a method of the `Figure`

class in Matplotlib. This method reduces the margins at the edge of the figure.

You can try the code so far in the Console/REPL:

>>> import matplotlib.pyplot as plt >>> from solar_system_3d import SolarSystem >>> solar_system = SolarSystem(400) >>> plt.show() # if not using interactive mode

This gives a figure with an empty set of 3D axes:

You’ll use the `size`

parameter later to set the size of this cube. You’ll return to the `SolarSystem`

class later. For the time being, you can turn your attention to defining the `SolarSystemBody`

class.

### Setting up the `SolarSystemBody`

class

You can start creating the `SolarSystemBody`

class and its `__init__()`

method. I’m truncating the code in the `SolarSystem`

class definition in the code below for display purposes. In this and later code blocks, the lines containing `# ...`

indicate code you’ve already written earlier that’s not being displayed:

# solar_system_3d.py import matplotlib.pyplot as plt from vectors import Vector # class SolarSystem: # ... class SolarSystemBody: def __init__( self, solar_system, mass, position=(0, 0, 0), velocity=(0, 0, 0), ): self.solar_system = solar_system self.mass = mass self.position = position self.velocity = Vector(*velocity) self.solar_system.add_body(self)

The parameters in the `__init__()`

method are:

`solar_system`

enables you to link the body to a solar system. The argument should be of type`SolarSystem`

.`mass`

is an integer or float that defines the body’s mass. In this project, you’ll use arbitrary units, so you don’t need to use ‘real’ masses for stars and planets.`position`

is a point in 3D space defining the body’s position. It’s a tuple containing the*x*-,*y*-, and*z*-coordinates of the point. The default is the origin.`velocity`

defines the velocity of the body. Since the velocity of a moving body has magnitude and direction, it must be a vector. Although the argument needed when instantiating a`SolarSystemBody`

is a tuple, you can convert the tuple into a`Vector`

object when assigning it to the attribute`self.velocity`

.

You also call the `add_body()`

method you defined earlier in the `SolarSystem`

class to add this body to the solar system. Later on, you’ll add a bit more to the `__init__()`

method.

You can define another method in `SolarSystemBody`

to move the body using its current position and velocity:

# solar_system_3d.py import matplotlib.pyplot as plt from vectors import Vector # class SolarSystem: # ... class SolarSystemBody: def __init__( self, solar_system, mass, position=(0, 0, 0), velocity=(0, 0, 0), ): self.solar_system = solar_system self.mass = mass self.position = position self.velocity = Vector(*velocity) self.solar_system.add_body(self) def move(self): self.position = ( self.position[0] + self.velocity[0], self.position[1] + self.velocity[1], self.position[2] + self.velocity[2], )

The `move()`

method redefines the `position`

attribute based on the `velocity`

attribute. We’ve already discussed how you’re using arbitrary units for distance and mass. You’re also using arbitrary units for time. Each ‘time unit’ will be one iteration of the loop you’ll use to run the simulation. Therefore, `move()`

will shift the body by the amount required for one iteration, which is one time unit.

### Drawing the solar system bodies

You’ve already created the Matplotlib structures that will hold the solar system and all its bodies. Now, you can add a `draw()`

method to `SolarSystemBody`

to display the body on the Matplotlib plot. You can do this by drawing a marker.

Before you do so, you’ll need to define a few more attributes in `SolarSystemBody`

to control the colour and size of the markers that you’ll draw to represent the bodies:

# solar_system_3d.py import math import matplotlib.pyplot as plt from vectors import Vector # class SolarSystem: # ... class SolarSystemBody: min_display_size = 10 display_log_base = 1.3 def __init__( self, solar_system, mass, position=(0, 0, 0), velocity=(0, 0, 0), ): self.solar_system = solar_system self.mass = mass self.position = position self.velocity = Vector(*velocity) self.display_size = max( math.log(self.mass, self.display_log_base), self.min_display_size, ) self.colour = "black" self.solar_system.add_body(self) def move(self): self.position = ( self.position[0] + self.velocity[0], self.position[1] + self.velocity[1], self.position[2] + self.velocity[2], ) def draw(self): self.solar_system.ax.plot( *self.position, marker="o", markersize=self.display_size, color=self.colour )

The class attributes `min_display_size`

and `display_log_base`

set up the parameters for determining the size of the markers you’ll display on the 3D plot. You set a minimum size so that the marker you display is not too small, even for small bodies. You’ll use a logarithmic scale to convert from mass to marker size, and you set the base for this logarithm as another class attribute.

The `display_size`

instance attribute in the `__init__()`

method chooses between the calculated marker size and the minimum marker size you set. To determine the body’s display size in this project, you’re using its mass.

You also add the `colour`

attribute in `__init__()`

, which, for the time being, defaults to black.

To test these new additions, you can try the following in the Console/REPL:

>>> import matplotlib.pyplot as plt >>> from solar_system_3d import SolarSystem, SolarSystemBody >>> solar_system = SolarSystem(400) >>> plt.show() # if not using interactive mode >>> body = SolarSystemBody(solar_system, 100, velocity=(1, 1, 1)) >>> body.draw() >>> body.move() >>> body.draw()

The first call to `body.draw()`

draws the body at the origin since you’re using the default position for a solar system body. The call to `body.move()`

moves the body by the amount required for one ‘time unit’. Since the body’s velocity is `(1, 1, 1)`

, the body will move by one unit along each of the three axes. The second call to `body.draw()`

draws the solar system body in the second position. Note that the axes will automatically rescale when you do this. You’ll take care of this in the main code shortly.

## Moving Stars and Planets

You can return to the `SolarSystem`

class and link the solar system and its bodies further by adding two new methods to the class: `update_all()`

and `draw_all()`

:

# solar_system_3d.py import math import matplotlib.pyplot as plt from vectors import Vector class SolarSystem: def __init__(self, size): self.size = size self.bodies = [] self.fig, self.ax = plt.subplots( 1, 1, subplot_kw={"projection": "3d"}, figsize=(self.size / 50, self.size / 50), ) self.fig.tight_layout() def add_body(self, body): self.bodies.append(body) def update_all(self): for body in self.bodies: body.move() body.draw() def draw_all(self): self.ax.set_xlim((-self.size / 2, self.size / 2)) self.ax.set_ylim((-self.size / 2, self.size / 2)) self.ax.set_zlim((-self.size / 2, self.size / 2)) plt.pause(0.001) self.ax.clear() # class SolarSystemBody: # ...

The `update_all()`

method goes through each body in the solar system and moves and draws each body. The `draw_all()`

method sets the limits for the three axes using the solar system’s size and updates the plot through the `pause()`

function. This method also clears the axes, ready for the next plot.

You can start building a simple solar system and test the code you’ve written so far by creating a new script called `simple_solar_system.py`

:

# simple_solar_system.py from solar_system_3d import SolarSystem, SolarSystemBody solar_system = SolarSystem(400) body = SolarSystemBody(solar_system, 100, velocity=(1, 1, 1)) for _ in range(100): solar_system.update_all() solar_system.draw_all()

When you run this script, you’ll see a black body moving away from the centre of the plot:

You can change the perspective of the 3D plot so that you’re viewing the 3D axes directly along one of the axes. You can do so by setting both the azimuth and the elevation of the view to `0`

in `SolarSystem.__init__()`

:

# solar_system_3d.py import math import matplotlib.pyplot as plt from vectors import Vector class SolarSystem: def __init__(self, size): self.size = size self.bodies = [] self.fig, self.ax = plt.subplots( 1, 1, subplot_kw={"projection": "3d"}, figsize=(self.size / 50, self.size / 50), ) self.fig.tight_layout() self.ax.view_init(0, 0) def add_body(self, body): self.bodies.append(body) def update_all(self): for body in self.bodies: body.move() body.draw() def draw_all(self): self.ax.set_xlim((-self.size / 2, self.size / 2)) self.ax.set_ylim((-self.size / 2, self.size / 2)) self.ax.set_zlim((-self.size / 2, self.size / 2)) plt.pause(0.001) self.ax.clear() # class SolarSystemBody: # ...

Running `simple_solar_system.py`

now gives the following view:

The *x*-axis is now perpendicular to your screen. Since you’re displaying a 3D view on a 2D display, you’ll always have one direction which is perpendicular to the 2D plane you’re using to display the plot. This restriction can make it hard to distinguish when an object is moving along that axis. You can see this by changing the body’s velocity in `simple_solar_system.py`

to `(1, 0, 0)`

and running the script again. The body appears stationary since it’s only moving along the axis coming out of your screen!

### Helping with the 3D perspective

You can improve the 3D visualisation by changing the size of the marker depending on its *x*-coordinate. Objects closer to you appear larger, and objects further away appear smaller. You can make a change to the `draw()`

method in the `SolarSystemBody`

class:

# solar_system_3d.py # ... class SolarSystemBody: # ... def draw(self): self.solar_system.ax.plot( *self.position, marker="o", markersize=self.display_size + self.position[0] / 30, color=self.colour )

`self.position[0]`

represents the body’s position along the *x*-axis, which is the one perpendicular to the screen. The factor of `30`

you divide by is an arbitrary factor you can use to control how strong you want this effect to be.

Later in this tutorial, you’ll also add another feature that will help visualise the 3D motion of the stars and planets.

## Adding The Effects Of Gravity

You have a solar system with bodies that can move within it. The code so far works fine if you have a single body. But that’s not a very interesting solar system! If you have two or more bodies, they will interact through their mutual gravitational attraction.

Toward the beginning of this article, I briefly reviewed the physics you’ll need to deal with the gravitational force between two objects. Since you’re using arbitrary units in this project, you can ignore the gravitational constant *G* and simply work out the force due to gravity between two objects as:

F=\frac{m_1m_1}{r^2}

Once you know the force between two objects, since F=ma, you can work out the acceleration that each object is subject to using:

a=\frac{F}{m}

And once you know the acceleration, you can change the object’s velocity.

You can add two new methods, one in `SolarSystemBody`

and another in `SolarSystem`

, to work out the force and acceleration between any two bodies and to go through all the bodies in the solar system and work out the interactions between them.

### Working out the acceleration due to gravity

The first of these methods works out the gravitational force between two bodies, calculates the acceleration of each of the bodies and changes the velocities of the two bodies. You can split these tasks into three methods if you prefer, but in this example, I’ll put these tasks into a single method in `SolarSystemBody`

:

# solar_system_3d.py import math import matplotlib.pyplot as plt from vectors import Vector # class SolarSystem: # ... class SolarSystemBody: # ... def accelerate_due_to_gravity(self, other): distance = Vector(*other.position) - Vector(*self.position) distance_mag = distance.get_magnitude() force_mag = self.mass * other.mass / (distance_mag ** 2) force = distance.normalize() * force_mag reverse = 1 for body in self, other: acceleration = force / body.mass body.velocity += acceleration * reverse reverse = -1

`accelerate_due_to_gravity()`

is called on an object of type `SolarSystemBody`

and needs another `SolarSystemBody`

body as an argument. The parameters `self`

and `other`

represent the two bodies interacting with each other. The steps in this method are the following:

- The positions of the two bodies are used to find the distance between the two bodies. You represent this as a vector since both its magnitude and direction are important. You extract the
*x*-,*y*-, and*z*– values from the`position`

attribute using the unpacking operator`*`

and convert these into objects of type`Vector`

, which you defined earlier. Since you defined the`__sub__()`

dunder method for the`Vector`

class, you can subtract one vector from the other to get the distance between them as another vector. - You also calculate the magnitude of the distance vector using the
`get_magnitude()`

method of the`Vector`

class. - Next, you work out the magnitude of the force between the two bodies using the equation summarised above.
- However, the force has a direction as well as a magnitude. Therefore, you need to represent it as a vector. The direction of the force is the same as the direction of the vector connecting the two objects. You obtain the force vector by first normalizing the distance vector. This normalization gives a unit vector with the same direction as the vector connecting the two bodies but with a magnitude of
`1`

. Then, you multiply the unit vector by the magnitude of the force. You’re using scalar multiplication of a vector in this case which you defined when you included`__mul__()`

in the`Vector`

class. - For each of the two bodies, you work out the acceleration using the equation shown above.
`force`

is a vector. Therefore, when you divide by`body.mass`

, you’re using the scalar division you defined when you included`__truediv__()`

in the`Vector`

class.`acceleration`

is the object returned by`Vector.__truediv__()`

, which is also a`Vector`

object. - Finally, you increment the velocity using the acceleration. This method works out the values relevant for one time unit, which in this simulation is the time it takes for one iteration of the loop that will control the simulation. The
`reverse`

parameter ensures the opposite acceleration is applied to the second body since the two bodies are being pulled towards each other. The`*`

operator again calls`Vector.__mul__()`

and results in scalar multiplication.

### Calculating the interactions between all bodies in the solar system

Now that you’re able to work out the interaction between any two bodies, you can work out the interaction between all the bodies present in the solar system. You can shift your attention back to the `SolarSystem`

class for this:

# solar_system_3d.py import math import matplotlib.pyplot as plt from vectors import Vector class SolarSystem: # ... def calculate_all_body_interactions(self): bodies_copy = self.bodies.copy() for idx, first in enumerate(bodies_copy): for second in bodies_copy[idx + 1:]: first.accelerate_due_to_gravity(second) class SolarSystemBody: # ... def accelerate_due_to_gravity(self, other): distance = Vector(*other.position) - Vector(*self.position) distance_mag = distance.get_magnitude() force_mag = self.mass * other.mass / (distance_mag ** 2) force = distance.normalize() * force_mag reverse = 1 for body in self, other: acceleration = force / body.mass body.velocity += acceleration * reverse reverse = -1

The `calculate_all_body_interactions()`

method goes through all the bodies in the solar system. Each body interacts with every other body in the solar system:

- You’re using a copy of
`self.bodies`

to cater for the possibility that bodies will be removed from the solar system during the loop.*In the version you’re writing in this article, you won’t remove any bodies from the solar system. However, you may need to do so in the future if you expand this project further.* - To ensure your code doesn’t calculate the interactions between the same two bodies twice, you only work out the interactions between a body and those bodies that follow it in the list. This is why you’re using the slice
`idx + 1:`

in the second`for`

loop. - The final line calls
`accelerate_due_to_gravity()`

for the first body and includes the second body as the method’s argument.

Now, you’re ready to create a simple solar system and test the code you’ve written so far.

## Creating A Simple Solar System

In this project, you’ll focus on creating one of two types of bodies: suns and planets. You can create two classes for these bodies. The new classes inherit from `SolarSystemBody`

:

# solar_system_3d.py import itertools import math import matplotlib.pyplot as plt from vectors import Vector # class SolarSystem: # ... # class SolarSystemBody: # ... class Sun(SolarSystemBody): def __init__( self, solar_system, mass=10_000, position=(0, 0, 0), velocity=(0, 0, 0), ): super(Sun, self).__init__(solar_system, mass, position, velocity) self.colour = "yellow" class Planet(SolarSystemBody): colours = itertools.cycle([(1, 0, 0), (0, 1, 0), (0, 0, 1)]) def __init__( self, solar_system, mass=10, position=(0, 0, 0), velocity=(0, 0, 0), ): super(Planet, self).__init__(solar_system, mass, position, velocity) self.colour = next(Planet.colours)

The `Sun`

class uses a default mass of 10,000 units and sets the colour to yellow. You use the string `'yellow'`

, which is a valid colour in Matplotlib.

In the `Planet`

class, you create an `itertools.cycle`

object with three colours. In this case, the three colours are red, green, and blue. You can use any RGB colours you wish, and any number of colours, too. In this class, you define colours using a tuple with RGB values instead of a string with the colour name. This is also a valid way of defining colours in Matplotlib. You cycle through these colours using the `next()`

function each time you create a new planet.

You also set the default mass to 10 units.

Now, you can create a solar system with one sun and two planets in `simple_solar_system.py`

:

# simple_solar_system.py from solar_system_3d import SolarSystem, Sun, Planet solar_system = SolarSystem(400) sun = Sun(solar_system) planets = ( Planet( solar_system, position=(150, 50, 0), velocity=(0, 5, 5), ), Planet( solar_system, mass=20, position=(100, -50, 150), velocity=(5, 0, 0) ) ) while True: solar_system.calculate_all_body_interactions() solar_system.update_all() solar_system.draw_all()

In this script, you create a sun and two planets. You’re assigning the sun and the planets to variables called `sun`

and `planets`

, but this is not strictly required as once the `Sun`

and `Planet`

objects are created, they’re added to `solar_system`

and you don’t need to reference them directly.

You use a `while`

loop to run the simulation. The loop performs three operations in each iteration. When you run this script, you’ll get the following animation:

It works, sort of. You can see the sun anchored at the centre of this solar system and the planets being affected by the sun’s gravitational pull. In addition to the planets’ movements in the plane containing your computer screen (these are the *y*– and *z*-axes), you can also see the planets getting larger and smaller as they also move in the *x*-axis, which is perpendicular to your screen.

However, you may have noticed some peculiar behaviour of the planets. When they’re meant to be behind the sun, the planets are still displayed in front of the sun. This is not a problem with the mathematics—if you track the positions of the planets, you’ll see that their *x*-coordinates show that they actually go behind the sun, as you would expect.

### Showing bodies behind other bodies

The issue comes from the way Matplotlib draws objects on a plot. Matplotlib plots objects in layers in the order you plot them. Since you created the sun before the planets, the `Sun`

object comes first in `solar_system.bodies`

and is drawn as the bottom layer. You can verify this fact by creating the sun after the planets, and you’ll see that the planets will always appear behind the sun in this case.

You’d like Matplotlib to plot the solar system bodies in the correct order, starting with the ones that are the furthest back. To achieve this, you can sort the `SolarSystem.bodies`

list based on the value of the *x*-coordinate each time you want to refresh the 3D plot. Here’s how you can do this in the `update_all()`

method in `SolarSystem`

:

# solar_system_3d.py import itertools import math import matplotlib.pyplot as plt from vectors import Vector class SolarSystem: # ... def update_all(self): self.bodies.sort(key=lambda item: item.position[0]) for body in self.bodies: body.move() body.draw() # ... # class SolarSystemBody: # ... # class Sun(SolarSystemBody): # ... # class Planet(SolarSystemBody): # ...

You use the list method `sort`

with the `key`

parameter to define the rule you’d like to use to sort the list. The `lambda`

function sets this rule. In this case, you’re using the value of `position[0]`

of each body, which represents the *x*-coordinate. Therefore, each time you call `update_all()`

in the simulation’s `while`

loop, the list of bodies is reordered based on their position along the *x*-axis.

The result of running the `simple_solar_system.py`

script now is the following:

Now, you can visualise the orbits of the planets as they orbit the sun. The changing size shows their *x*-position, and when the planets are behind the sun, they’re hidden from sight!

Finally, you can also remove the axes and grid so that all you see in the simulation is the sun and the planets. You can do this by adding a call to the Matplotlib `axis()`

method in `SolarSystem.draw_all()`

:

# solar_system_3d.py import itertools import math import matplotlib.pyplot as plt from vectors import Vector class SolarSystem: # ... def draw_all(self): self.ax.set_xlim((-self.size / 2, self.size / 2)) self.ax.set_ylim((-self.size / 2, self.size / 2)) self.ax.set_zlim((-self.size / 2, self.size / 2)) self.ax.axis(False) plt.pause(0.001) self.ax.clear() # ... # class SolarSystemBody: # ... # class Sun(SolarSystemBody): # ... # class Planet(SolarSystemBody): # ...

And the simulation now looks like this:

The simulation of a 3D solar system in Python using Matplotlib is now complete. In the next section, you’ll add a feature that will allow you to view a 2D projection of the *xy*-plane at the bottom of the simulation. This can help with visualising the 3D dynamics of the bodies in the solar system.

## Adding a 2D Projection of The *xy*-Plane

To help visualise the motion of the bodies in the simulation of a 3D solar system in Python, you can add a 2D projection on the ‘floor’ of the animation. This 2D projection will show the position of the bodies in the *xy*-plane. To achieve this, you’ll need to add another plot to the same axes in which you’re showing the animation and only show the changes in the *x*– and *y*-coordinates. You can anchor the *z*-coordinate to the bottom of the plot so that the 2D projection is displayed on the floor of the animation.

You can start by adding a new parameter to the `__init__()`

method for the `SolarSystem`

class:

# solar_system_3d.py import itertools import math import matplotlib.pyplot as plt from vectors import Vector class SolarSystem: def __init__(self, size, projection_2d=False): self.size = size self.projection_2d = projection_2d self.bodies = [] self.fig, self.ax = plt.subplots( 1, 1, subplot_kw={"projection": "3d"}, figsize=(self.size / 50, self.size / 50), ) self.ax.view_init(0, 0) self.fig.tight_layout() # ... # class SolarSystemBody: # ... # class Sun(SolarSystemBody): # ... # class Planet(SolarSystemBody): # ...

The new parameter `projection_2d`

, which defaults to `False`

, will allow you to toggle between the two visualisation options. If `projection_2d`

is `False`

, the animation will only show the bodies moving in 3D, with no axes and grid, as in the last result you’ve seen.

Let’s start making some changes for when `projection_2d`

is `True`

:

# solar_system_3d.py import itertools import math import matplotlib.pyplot as plt from vectors import Vector class SolarSystem: def __init__(self, size, projection_2d=False): self.size = size self.projection_2d = projection_2d self.bodies = [] self.fig, self.ax = plt.subplots( 1, 1, subplot_kw={"projection": "3d"}, figsize=(self.size / 50, self.size / 50), ) self.fig.tight_layout() if self.projection_2d: self.ax.view_init(10, 0) else: self.ax.view_init(0, 0) def add_body(self, body): self.bodies.append(body) def update_all(self): self.bodies.sort(key=lambda item: item.position[0]) for body in self.bodies: body.move() body.draw() def draw_all(self): self.ax.set_xlim((-self.size / 2, self.size / 2)) self.ax.set_ylim((-self.size / 2, self.size / 2)) self.ax.set_zlim((-self.size / 2, self.size / 2)) if self.projection_2d: self.ax.xaxis.set_ticklabels([]) self.ax.yaxis.set_ticklabels([]) self.ax.zaxis.set_ticklabels([]) else: self.ax.axis(False) plt.pause(0.001) self.ax.clear() def calculate_all_body_interactions(self): bodies_copy = self.bodies.copy() for idx, first in enumerate(bodies_copy): for second in bodies_copy[idx + 1:]: first.accelerate_due_to_gravity(second) class SolarSystemBody: min_display_size = 10 display_log_base = 1.3 def __init__( self, solar_system, mass, position=(0, 0, 0), velocity=(0, 0, 0), ): self.solar_system = solar_system self.mass = mass self.position = position self.velocity = Vector(*velocity) self.display_size = max( math.log(self.mass, self.display_log_base), self.min_display_size, ) self.colour = "black" self.solar_system.add_body(self) def move(self): self.position = ( self.position[0] + self.velocity[0], self.position[1] + self.velocity[1], self.position[2] + self.velocity[2], ) def draw(self): self.solar_system.ax.plot( *self.position, marker="o", markersize=self.display_size + self.position[0] / 30, color=self.colour ) if self.solar_system.projection_2d: self.solar_system.ax.plot( self.position[0], self.position[1], -self.solar_system.size / 2, marker="o", markersize=self.display_size / 2, color=(.5, .5, .5), ) def accelerate_due_to_gravity(self, other): distance = Vector(*other.position) - Vector(*self.position) distance_mag = distance.get_magnitude() force_mag = self.mass * other.mass / (distance_mag ** 2) force = distance.normalize() * force_mag reverse = 1 for body in self, other: acceleration = force / body.mass body.velocity += acceleration * reverse reverse = -1 class Sun(SolarSystemBody): def __init__( self, solar_system, mass=10_000, position=(0, 0, 0), velocity=(0, 0, 0), ): super(Sun, self).__init__(solar_system, mass, position, velocity) self.colour = "yellow" class Planet(SolarSystemBody): colours = itertools.cycle([(1, 0, 0), (0, 1, 0), (0, 0, 1)]) def __init__( self, solar_system, mass=10, position=(0, 0, 0), velocity=(0, 0, 0), ): super(Planet, self).__init__(solar_system, mass, position, velocity) self.colour = next(Planet.colours)

The changes you’ve made are the following:

- In
`SolarSystem.__init__()`

, the 3D view is set to`view_init(0, 0)`

when the 2D projection is turned off, as before. However, the elevation is changed to 10º when the 2D projection option is turned on to allow the bottom plane to be visible. - In
`SolarSystem.draw_all()`

, the grid and axes are turned off only when there is no 2D projection. When the 2D projection is enabled, the axes and grid are displayed. However, the tick marks are replaced with blanks since the numbers on the three axes are arbitrary and are not needed. - In
`SolarSystemBody.draw()`

, a second plot is added when`projection_2d`

is`True`

. The first two arguments in`plot()`

are the bodies’*x*– and*y*-positions. However, instead of using the*z*-position as the third argument, you use the minimum value of*z*which represents the ‘floor’ of the cube containting the three axes. You then plot a grey marker half the size of the main markers in the animation.

You’ll also need to make a small change in `simple_solar_system.py`

to turn on the 2D projection:

# simple_solar_system.py from solar_system_3d import SolarSystem, Sun, Planet solar_system = SolarSystem(400, projection_2d=True) sun = Sun(solar_system) planets = ( Planet( solar_system, position=(150, 50, 0), velocity=(0, 5, 5), ), Planet( solar_system, mass=20, position=(100, -50, 150), velocity=(5, 0, 0) ) ) while True: solar_system.calculate_all_body_interactions() solar_system.update_all() solar_system.draw_all()

The simulation now looks like this:

The 2D projection of the *xy*-plane makes it easier to follow the paths of the orbiting bodies.

## Creating a Binary Star System

We’ll finish off with another simulation of a 3D solar system in Python. You’ll simulate a binary star system using the same classes you’ve already defined. Create a new file called `binary_star_system.py`

and create two suns and two planets:

# binary_star_system.py from solar_system_3d import SolarSystem, Sun, Planet solar_system = SolarSystem(400) suns = ( Sun(solar_system, position=(40, 40, 40), velocity=(6, 0, 6)), Sun(solar_system, position=(-40, -40, 40), velocity=(-6, 0, -6)), ) planets = ( Planet( solar_system, 10, position=(100, 100, 0), velocity=(0, 5.5, 5.5), ), Planet( solar_system, 20, position=(0, 0, 0), velocity=(-11, 11, 0), ), ) while True: solar_system.calculate_all_body_interactions() solar_system.update_all() solar_system.draw_all()

The simulation of this binary star system is the following:

Or you can turn on the 2D projection when creating the `SolarSystem`

object:

# binary_star_system.py from solar_system_3d import SolarSystem, Sun, Planet solar_system = SolarSystem(400, projection_2d=True) # ...

This version gives the following result:

This binary star system is not stable, and both planets are soon flung out of the system by the two suns!

If you wish, you can extend the class definitions to detect collisions between two bodies and remove a planet if it collides with a sun. The simpler, 2D version of this project, which simulates orbiting planets in 2D, includes this feature. You can look at how it was implemented in that simpler project if you’d like to add it to this project.

*The final versions of the code used in this article are also available on this GitHub repo.*

## Final Words

You can now simulate a 3D solar system in Python using Matplotlib. In this article, you’ve learned how to place objects in 3D space using vectors and the graphical capabilities of Matplotlib. You can read more about how to use Matplotlib, including making more complex animations using the `animations`

submodule in Matplotlib, in the Chapter Basics of Data Visualisation in Python Using Matplotlib of The Python Coding Book.

This completes the two-part Orbiting Planets Series. In the first post of the series, you considered only the 2D scenario and used the `turtle`

module to create the graphical animation. In the second article, the one you just finished, you looked at a 3D solar system in Python using Matplotlib for the graphical representation of the animation.

It’s now your turn to try and create simple and more complex solar systems. Can you create a stable binary star system?

I hope you enjoyed simulating a 3D solar system in Python using Matplotlib. Now you’re ready to try and create your own simulations of real-world processes.

## Further Reading

- You can read the first article in this series which simulates orbiting planets in 2D using the
`turtle`

graphics module - Read more about object-oriented programming
- You may find this article about using the 2D Fourier Transform in Python to reconstruct images from sine functions of interest, too
- Finally, if you want to get a different type of understanding of what happens behind the scenes in a Python program, try The White Room: Understanding Programming

## Get the latest blog updates

No spam promise. You’ll get an email when a new blog post is published

[…] In the second article in the series, you’ll move on to using Matplotlib to run and display the animation of a 3D solar system in Python. […]

Thank you so much for this great and comprehensive article! I felt like a roller coaster of the entire pillars of python!

Glad you enjoyed it!