Simulating the Solar System
Last updated on 2025-02-25 | Edit this page
Overview
Questions
- How can I work with physical units?
- How do I quickly visualize some data?
- How is dispatch used in practice?
Objectives
- Learn to work with
Unitful
- Take the first steps with
Makie
for visualisation - Work with
const
- Define a
struct
- Get familiar with some idioms in Julia
In this episode we’ll be building a simulation of the solar system. That is, we’ll only consider the Sun, Earth and Moon, but feel free to add other planets to the mix later on! To do this we need to work with the laws of gravity. We will be going on extreme tangents, exposing interesting bits of the Julia language as we go.
Introduction to Unitful
We’ll be using the Unitful
library to ensure that we
don’t mix up physical units.
Using Unitful
is without any run-time overhead. The unit
information is completely contained in the Quantity
type.
With Unitful
we can attach units to quantities using the
@u_str
syntax:
At the moment it is not important to precisely understand the syntax. Just notice that the information on the units is stored in Julia’s type system:
OUTPUT
Quantity{Float64, 𝐋, Unitful.FreeUnits{(m,), 𝐋, nothing}}
We see a Quantity
represented by a double precision
floating-point value, having dimension of length, and units of meters.
All of the unit information is captured in the type system, which only
acts at compile time. At run-time, this is an ordinary
Float64
.
Try adding incompatible quantities
For instance, add a quantity in meters to one in seconds. Be creative. Can you understand the error messages?
Next to that, we use the Vec3d
type from
GeometryBasics
to store vector particle positions.
Libraries in Julia tend to work well together, even if they’re not
principally designed to do so. We can combine Vec3d
with
Unitful
with no problems.
Generate random Vec3d
The Vec3d
type is a static 3-vector of double precision
floating point values. Read the documentation on the randn
function. Can you figure out a way to generate a random
Vec3d
?
Two bodies of mass \(M\) and \(m\) attract each other with the force
\[F = \frac{GMm}{r^2},\]
where \(r\) is the distance between those bodies, and \(G\) is the universal gravitational constant.
JULIA
#| id: gravity
const G = 6.6743e-11u"m^3*kg^-1*s^-2"
gravitational_force(m1, m2, r) = G * m1 * m2 / r^2
Verify that the force exerted between two 1 million kg bodies at a distance of 2.5 meters is about equal to holding a 1kg object on Earth. Or equivalently two one tonne objects at a distance of 2.5 mm (all of the one tonne needs to be concentrated in a space tiny enough to get two of them at that close a distance)
A better way to write the force would be,
\[\vec{F} = \hat{r} \frac{G M m}{|r|^2},\]
where \(\hat{r} = \vec{r} / |r|\), so in total we get a third power, \(|r|^3 = (r^2)^{3/2}\), where \(r^2\) can be computed with the dot-product.
JULIA
#| id: gravity
"""
gravitational_force(m1, m2, r)
Returns the gravitational force as a function of masses `m1` and `m2` and `r` the distance vector between them. The force will be in the direction of `r`.
"""
gravitational_force(m1, m2, r::AbstractVector) =
r * (G * m1 * m2 * (r ⋅ r)^(-1.5))
Callout
Not all of you will be jumping up and down for doing high-school physics. We will get to other sciences than physics later in this workshop, I promise!
The zero
function
The zero
function is overloaded to return a zero object
for any type that has a natural definition of zero.
The one
function
The return-value of zero
should be the additive identity
of the given type. So for any type T
:
There also is the one
function which returns the
multiplicative identity of a given type. Try to use the *
operator on two strings. What do you expect one(String)
to
return?
Particles
We are now ready to define the Particle
type. First we
define some constants so that our code remains readable.
JULIA
#| id: gravity
const Mass = typeof(1.0u"kg")
const MomentumVector = typeof(Vec3d(1)u"kg*m/s")
const PositionVector = typeof(Vec3d(1)u"m")
const VelocityVector = typeof(Vec3d(1)u"m/s")
JULIA
#| id: gravity
mutable struct Particle
mass::Mass
position::PositionVector
momentum::MomentumVector
end
We can write a function to generate random particles with some spread and (velocity) dispersion:
JULIA
#| id: gravity
random_particle(mass=1e6u"kg", spread=1.0u"m", dispersion=2.0u"mm/s") =
Particle(mass, randn(Vec3d) * spread, randn(Vec3d) * dispersion * mass)
function random_particles(n; seed=0, args...)
Random.seed!(seed)
[random_particle(args...) for _ in 1:n]
end
Note the ellipsis ...
. Those do variadic argument
capture. We can also define a function to obtain the velocity of a
particle:
Getters
a, Create \(N\) random particles
(say \(N=3\), but it really doesn’t
matter), and get their velocity in a one-liner. Don’t use for-loops,
just broadcasting. b. Can you do the same for obtaining the position
and/or momentum? Implement a getter for those properties. c. We can also
meaningfully extend the getter for momentum
on collections
of particles. The total momentum of a set of particles is the sum of
their individual momenta. Read the documention on the sum
function. Can you find a particularly nice way to implement
momentum(p::AbstractArray{Particle})
? d. Do the same for
the total mass of a set of particles. Can you now generalize the
definition for velocity
, so that the same function works
for particles and collections of particles?
or
JULIA
#| id: gravity
mass(p::Particle) = p.mass
position(p::Particle) = p.position
momentum(p::Particle) = p.momentum
# random_particles(3) .|> momentum
# etc...
mass(p::AbstractArray{Particle}) = sum(mass, p)
momentum(p::AbstractArray{Particle}) = sum(momentum, p)
velocity(p) = momentum(p) / mass(p)
By abstracting struct member access into functions, we could write some very elegant code to compute the total net velocity of a group of particles! In general, it is considered good practice to create accessor functions like this, so that functions that rely on member access then become more portable (i.e. they can be made to work on other types by extending the accessor methods).
It is custom to divide the computation of orbits into a kick
and drift function. (There are deep mathematical reasons for
this that we won’t get into.) We’ll first implement the
kick!
function, that updates a collection of particles,
given a certain time step.
The kick
The following function performs the kick!
operation on a
set of particles. We’ll go on a little tangent for every line in the
function.
JULIA
#| id: gravity
function kick!(particles, dt)
for i in eachindex(particles)
for j in 1:(i-1)
r = particles[j].position - particles[i].position
force = gravitational_force(particles[i].mass, particles[j].mass, r)
particles[i].momentum += dt * force
particles[j].momentum -= dt * force
end
end
return particles
end
Why the !
exclamation mark?
In Julia it is custom to have an exclamation mark at the end of names of functions that mutate their arguments.
Use eachindex
Note the way we used eachindex
. This idiom guarantees
that we can’t make out-of-bounds errors. Also this kind of indexing is
generic over other collections than vectors. However, this generality is
lost in the inner loop, where we use explicit numeric bounds. In this
case those bounds actually half the amount of work we need to do, so the
sacrifice is justified.
Broadcasting vs. specialization
We’re subtracting two vectors. We could have written that with dot-notation indicating a broadcasted function application. Generate two random numbers:
Then time the subtraction broadcasted and non-broadcasted. Which is faster? Why?
Callout
We’re iterating over the range 1:(i-1)
, and so we don’t
compute forces twice. Momentum should be conserved!
Luckily the drift!
function is much easier to implement,
and doesn’t require that we know about all particles.
JULIA
#| id: gravity
function drift!(p::Particle, dt)
p.position += dt * p.momentum / p.mass
end
function drift!(particles, dt)
for p in particles
drift!(p, dt)
end
return particles
end
Note that we defined the drift!
function twice, for
different arguments. We’re using the dispatch mechanism to write
clean/readable code.
With random particles
Let’s run a simulation with some random particles
JULIA
#| id: gravity
function run_simulation(particles, dt, n_steps)
x = deepcopy(particles)
[deepcopy(leap_frog!(x, dt)) for _ in 1:n_steps]
end
function random_orbits(n, mass; dt=1.0u"s", steps=5000, args...)
particles = random_particles(n; args...)
run_simulation(particles, dt, steps) |> collect
end
JULIA
#| classes: ["task"]
#| file: scripts/plot-two-body-drift.jl
#| creates: episodes/fig/two-body-drift.svg
#| collect: figures
module Script
using Unitful
using CairoMakie
using EfficientJulia.Gravity
function main()
fig = Figure()
ax = Axis3(fig[1, 1])
orbits = run_simulation(
random_particles(2), 1.0u"s", 5000)
for i in 1:2
orbit = position.(getindex.(orbits, i))
lines!(ax, orbit / u"m")
end
save("episodes/fig/two-body-drift.svg", fig)
end
end
Script.main()
As you see, the random particles we generated have a non-zero total momentum.
Frame of reference
We need to make sure that our entire system doesn’t have a net velocity. Otherwise it will be hard to visualize our results!
JULIA
#| id: gravity
function set_still!(particles)
v = velocity(particles)
for p in particles
p.momentum -= v * mass(p)
end
return particles
end
JULIA
#| classes: ["task"]
#| file: scripts/plot-two-body-still.jl
#| creates: episodes/fig/random-orbits.svg
#| collect: figures
module Script
using Unitful
using CairoMakie
using EfficientJulia.Gravity
function main()
fig = Figure(size=(1000, 500))
ax1 = Axis3(fig[1, 1])
orbits = run_simulation(
set_still!(random_particles(2)), 1.0u"s", 5000)
for i in 1:2
orbit = position.(getindex.(orbits, i))
lines!(ax1, orbit / u"m")
end
ax2 = Axis3(fig[1, 2])
orbits = run_simulation(
set_still!(random_particles(3, seed=3)), 1.0u"s", 5000)
for i in 1:3
orbit = position.(getindex.(orbits, i))
lines!(ax2, orbit / u"m")
scatter!(ax2, orbit[end] / u"m")
end
save("episodes/fig/random-orbits.svg", fig)
end
end
Script.main()
Solar System
JULIA
#| id: gravity
# const SUN = Particle(2e30u"kg",
# Vec3d(0.0)u"m",
# Vec3d(0.0)u"m/s")
# const EARTH = Particle(6e24u"kg",
# Vec3d(1.5e11, 0, 0)u"m",
# Vec3d(0, 3e4, 0)u"m/s")
# const MOON = Particle(7.35e22u"kg",
# EARTH.position + Vec3d(3.844e8, 0.0, 0.0)u"m",
# velocity(EARTH) + Vec3d(0, 1e3, 0)u"m/s")
Challenge
Plot the orbit of the moon around the earth. Make a
Dataframe
that contains all model data, and work from
there. Can you figure out the period of the orbit?
Key Points
- standard functions like
rand
,zero
and operators can be extended by libraries to work with new types - functions (not objects) are central to programming Julia
- don’t over-specify argument types: Julia is dynamically typed, embrace it
-
eachindex
and relatives are good ways iterate collections
JULIA
#| file: src/Gravity.jl
module Gravity
using Unitful
using GeometryBasics
using DataFrames
using LinearAlgebra
using Random
import Base: position
export random_partcle, random_particles, velocity, mass, momentum, position
export run_simulation, set_still!
<<gravity>>
end
JULIA
##| classes: ["task"]
#| creates: episodes/fig/random-orbits.png
#| collect: figures
module Script
using Unitful
using GLMakie
using DataFrames
using Random
using EfficientJulia.Gravity: random_orbits
function plot_orbits!(ax, orbits::DataFrame)
for colname in names(orbits)[2:end]
scatter!(ax, [orbits[1,colname] / u"m"])
lines!(ax, orbits[!,colname] / u"m")
end
end
function main()
Random.seed!(0)
orbs1 = random_orbits(2, 1e6u"kg")
Random.seed!(15)
orbs2 = random_orbits(3, 1e6u"kg")
fig = Figure(size=(1024, 600))
ax1 = Axis3(fig[1,1], azimuth=π/2+0.1, elevation=0.1π, title="two particles")
ax2 = Axis3(fig[1,2], azimuth=π/3, title="three particles")
plot_orbits!(ax1, orbs1)
plot_orbits!(ax2, orbs2)
save("episodes/fig/random-orbits.png", fig)
fig
end
end
Script.main()