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.pywill contain the class definitions.simple_solar_system.pywill 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
1and1to create a single set of axes in the figure. - The
subplot_kwparameter has a dictionary as its argument, which sets the projection to 3D. This means the axes created are anAxes3Dobject. figsizesets the overall size of the figure containing theAxes3Dobject.
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_systemenables you to link the body to a solar system. The argument should be of typeSolarSystem.massis 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.positionis 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.velocitydefines 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 aSolarSystemBodyis a tuple, you can convert the tuple into aVectorobject when assigning it to the attributeself.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
positionattribute using the unpacking operator*and convert these into objects of typeVector, which you defined earlier. Since you defined the__sub__()dunder method for theVectorclass, 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 theVectorclass. - 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 theVectorclass. - For each of the two bodies, you work out the acceleration using the equation shown above.
forceis a vector. Therefore, when you divide bybody.mass, you’re using the scalar division you defined when you included__truediv__()in theVectorclass.accelerationis the object returned byVector.__truediv__(), which is also aVectorobject. - 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
reverseparameter ensures the opposite acceleration is applied to the second body since the two bodies are being pulled towards each other. The*operator again callsVector.__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.bodiesto 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 secondforloop. - 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 toview_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 whenprojection_2disTrue. The first two arguments inplot()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
turtlegraphics 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
The post Simulating a 3D Solar System In Python Using Matplotlib (Orbiting Planets Series #2) appeared first on The Python Coding Book.
from Planet Python
via read more
No comments:
Post a Comment