This notebook teaches how to:
create groups, create immutable fields, display the table of contents, and set up complex integration instructions.

In addition to the simframe requirements, the following packages are needed for this notebook: astropy, matplotlib.

Advanced Integration

Let’s have a look at orbital integration.
Consider a planet Earth orbiting the Sun. How does its orbit look like?

The equations of motion are determined by a set of two differential equations for the position \(\vec{r}\) and velocity \(\vec{v}\) of Earth and Sun..

\(\frac{\mathrm{d}}{\mathrm{d}t} \vec{r} = \vec{v}\)

\(m\frac{\mathrm{d}}{\mathrm{d}t} \vec{v} = \vec{F}_\mathrm{G}\)

The gravitational force \(F_\mathrm{G}\) of a body of mass \(M\) and position \(\vec{R}\) acting on a body of mass \(m\) at position \(\vec{r}\) is given by

\(\vec{F}_\mathrm{G} = -GmM\frac{\vec{r}-\vec{R}}{\left| \vec{r}-\vec{R} \right|^3}\)

Creating groups

First we need to add fields for postion, velocity, and mass for Earth and Sun.
To have more structure we group the fields by Earth and Sun. This can be done with addgroup(name).
[1]:
from simframe import Frame
[2]:
sim = Frame(description="Earth-Sun system")
[3]:
sim.addgroup("Sun")
sim.addgroup("Earth")

The frame object has now two groups for Earth and Sun, that can be adressed just as fields.

[4]:
sim
[4]:
Frame (Earth-Sun system)
------------------------
    Earth        : Group
    Sun          : Group
  -----
    Integrator   : not specified
    Writer       : not specified
We will keep the problem general in three dimensions, but we will only use two of them for simplicity.
Although spherical coordinates would be more suitable for the problem, we will use Carteesian coordinates here.

We’ll use the astropy.constants and astropy.units modules for the constants and units we need here.

[5]:
import astropy.constants as c
import astropy.units as u
[6]:
AU   = (1.*u.au).si.value    # Astronomical unit [m]
day  = (1.*u.day).si.value   # Day [s]
G    = c.G.si.value          # Gravitational constant [m³/kg/s²]
year = (1.*u.year).si.value  # Year [s]

M_earth = c.M_earth.si.value # Mass of the Earth [kg]
M_sun   = c.M_sun.si.value   # Mass of the Sun [kg]

Creating immutable fields

We can now fill the groups with the variables we need, starting with the masses.
Again: the initial values must have the correct shape and data types.
[7]:
sim.Earth.addfield("M", M_earth, description="Mass [kg]")

The groups and the fields within can be easily accessed with

[8]:
sim.Earth
[8]:
Group
-----
    M            : Field (Mass [kg])
  -----
[9]:
sim.Earth.M
[9]:
5.972167867791379e+24

The mass of the Earth shall be constant throughout the simulation. We therefore set a flag so we cannot accidentally change its value.

[10]:
sim.Earth.M.constant = True
[11]:
sim.Earth
[11]:
Group
-----
    M            : Field (Mass [kg]), constant
  -----

Now we add the mass of the Sun, which we directly set to constant when adding the field.

[12]:
sim.Sun.addfield("M", M_sun, constant=True, description="Mass [kg]")
[13]:
sim.Sun
[13]:
Group
-----
    M            : Field (Mass [kg]), constant
  -----

We can do all sorts of operations with those fields, such as calculating the mass ratio.

[14]:
sim.Earth.M / sim.Sun.M
[14]:
3.0034893488507934e-06
Now we need to add fields for position and velocity. We initialize them with zeros and think later about the initial conditions.
The important thing is that the have the right shape upon initialization.
[15]:
import numpy as np
[16]:
sim.Earth.addfield("r", np.zeros(3), description="Position [m]")
sim.Earth.addfield("v", np.zeros(3), description="Velocity [m/s]")
sim.Sun.addfield("r", np.zeros(3), description="Position [m]")
sim.Sun.addfield("v", np.zeros(3), description="Velocity [m/s]")
[17]:
sim.Earth
[17]:
Group
-----
    M            : Field (Mass [kg]), constant
    r            : Field (Position [m])
    v            : Field (Velocity [m/s])
  -----
[18]:
sim.Sun
[18]:
Group
-----
    M            : Field (Mass [kg]), constant
    r            : Field (Position [m])
    v            : Field (Velocity [m/s])
  -----

Setting the integration variable

For our simulation we need an integration variable. In our case this is the time.

[19]:
sim.addintegrationvariable("t", 0., description="Time [s]")

We’ll set the step size to a constant value of one day.

[20]:
dt = 1.*day
[21]:
def f_dt(frame):
    return dt
[22]:
sim.t.updater = f_dt

We want to integrate for two years and want to have a snapshot every ten days.

[23]:
snapwidth = 10.*day
tmax = 2.*year
[24]:
sim.t.snapshots = np.arange(snapwidth, tmax, snapwidth)

Note: If the initial value of the integration variable is smaller the first snapshot, simframe automatically writes an output with initial conditions, if a writer is set.

[25]:
sim
[25]:
Frame (Earth-Sun system)
------------------------
    Earth        : Group
    Sun          : Group
  -----
    t            : IntVar (Time [s]), Integration variable
  -----
    Integrator   : not specified
    Writer       : not specified

Setting the writer

As a writer we use the namespacewriter, which does not write the data into files, but stores them within a buffer in the writer.

[26]:
from simframe import writers
[27]:
sim.writer = writers.namespacewriter()

The namespacewriter is by default not writing dump files.

[28]:
sim.writer
[28]:
Writer (Temporary namespace writer)
-----------------------------------
    Data directory : data
    Dumping        : False
    Verbosity      : 1

Adding differential equations

As a next step we’ll add differential equations to the quantities. The differential equations for the positions are simple. We simply return the velocities, which we can adress with the frame argument.

[29]:
def dr_Earth(frame, x, Y):
    return frame.Earth.v

def dr_Sun(frame, x, Y):
    return frame.Sun.v

For the differential equations of the velocities we’ll write a little helper function that computes the gravitational accelleration.

[30]:
# Gravitational acceleration
def ag(M, r, R):
    direction = r-R
    distance = np.linalg.norm(direction)
    return -G * M * direction / distance**3
[31]:
def dv_Earth(frame, x, Y):
    return ag(frame.Sun.M, frame.Earth.r, frame.Sun.r)

def dv_Sun(frame, x, Y):
    return ag(frame.Earth.M, frame.Sun.r, frame.Earth.r)

Now we need to add the differential equations to their fields.

[32]:
sim.Earth.v.differentiator = dv_Earth
sim.Earth.r.differentiator = dr_Earth
sim.Sun.v.differentiator   = dv_Sun
sim.Sun.r.differentiator   = dr_Sun

Displaying the table of contents

You can also display the complete tree structure of your frame.

[33]:
sim.toc
Frame (Earth-Sun system)
    - Earth: Group
        - M: Field (Mass [kg]), constant
        - r: Field (Position [m])
        - v: Field (Velocity [m/s])
    - Sun: Group
        - M: Field (Mass [kg]), constant
        - r: Field (Position [m])
        - v: Field (Velocity [m/s])
    - t: IntVar (Time [s]), Integration variable

Setting up complex integration instructions

Next we need to set up the integrator. We integrate all quantities with the explicit Euler 1st-order scheme as in the previous tutorial, but this time he have to integrate four fields. This can be easily achieved by adding four instructions.

[34]:
from simframe import Integrator
from simframe.integration import Instruction
from simframe import schemes
[35]:
sim.integrator = Integrator(sim.t, description="Euler 1st-order")
[36]:
instructions_euler = [Instruction(schemes.expl_1_euler, sim.Earth.r),
                      Instruction(schemes.expl_1_euler, sim.Earth.v),
                      Instruction(schemes.expl_1_euler, sim.Sun.r  ),
                      Instruction(schemes.expl_1_euler, sim.Sun.v  ),
                     ]
[37]:
sim.integrator.instructions = instructions_euler

Initial conditions

Before we can start the simulation, we have to think about initial conditions.

If we simply set the Sun at rest in the center of our simulation and only set the position and velocity of the Earth, then the center of mass would have a non-zero momentum and would slowly drift away from its initial position.

So what we do first: we set the Sun’s positon to zero (i.e., don’t do anything) and the Earth’s position to a distance of 1 AU in positive x-direction. Then we’ll offset their positions to center the system on the center of mass instead onto the Sun.

[38]:
r_Earth_ini = np.array([AU, 0., 0.])
r_Sun_ini = np.zeros(3)
# Center of mass
COM_ini = (M_earth*r_Earth_ini + M_sun*r_Sun_ini) / (M_earth+M_sun)
# Offset both positions
r_Earth_ini -= COM_ini
r_Sun_ini -= COM_ini

We save them in a separate variable instead of assigning them directly for later use.

The inital orbital velocities of the Earth shall be in positive y-direction, the velocity of the Sun in negative y-direction.
For calculating the value we use the reduced mass \(mu\) of the system.
[39]:
mu = M_earth*M_sun / (M_earth+M_sun)
[40]:
v_Earth_ini = np.array([0., np.sqrt(G*M_sun/M_earth*mu/AU), 0.])
v_Sun_ini   = np.array([0., -np.sqrt(G*M_earth/M_sun*mu/AU), 0.])

Now we assign them to their fields.

[41]:
sim.Earth.r = r_Earth_ini
sim.Earth.v = v_Earth_ini
sim.Sun.r   = r_Sun_ini
sim.Sun.v   = v_Sun_ini

Starting the simulation

[42]:
sim.run()
Saving frame 0000
Saving frame 0001
Saving frame 0002
Saving frame 0003
Saving frame 0004
Saving frame 0005
Saving frame 0006
Saving frame 0007
Saving frame 0008
Saving frame 0009
Saving frame 0010
Saving frame 0011
Saving frame 0012
Saving frame 0013
Saving frame 0014
Saving frame 0015
Saving frame 0016
Saving frame 0017
Saving frame 0018
Saving frame 0019
Saving frame 0020
Saving frame 0021
Saving frame 0022
Saving frame 0023
Saving frame 0024
Saving frame 0025
Saving frame 0026
Saving frame 0027
Saving frame 0028
Saving frame 0029
Saving frame 0030
Saving frame 0031
Saving frame 0032
Saving frame 0033
Saving frame 0034
Saving frame 0035
Saving frame 0036
Saving frame 0037
Saving frame 0038
Saving frame 0039
Saving frame 0040
Saving frame 0041
Saving frame 0042
Saving frame 0043
Saving frame 0044
Saving frame 0045
Saving frame 0046
Saving frame 0047
Saving frame 0048
Saving frame 0049
Saving frame 0050
Saving frame 0051
Saving frame 0052
Saving frame 0053
Saving frame 0054
Saving frame 0055
Saving frame 0056
Saving frame 0057
Saving frame 0058
Saving frame 0059
Saving frame 0060
Saving frame 0061
Saving frame 0062
Saving frame 0063
Saving frame 0064
Saving frame 0065
Saving frame 0066
Saving frame 0067
Saving frame 0068
Saving frame 0069
Saving frame 0070
Saving frame 0071
Saving frame 0072
Saving frame 0073
Execution time: 0:00:02

Reading and plotting

Reading data from the namespace writer works identical to the writer discussed earlier.

[43]:
data_euler = sim.writer.read.all()
[44]:
import matplotlib.pyplot as plt

def plot_orbits(data):
    fig, ax = plt.subplots(dpi=150)
    ax.set_aspect(1)
    ax.axis("off")
    fig.set_facecolor("#000000")
    imax = data.t.shape[0]
    for i in range(imax):
        alpha = np.maximum(i/imax-0.1, 0.5)
        ax.plot(data.Sun.r[i, 0], data.Sun.r[i, 1], "o", c="#FFFF00", markersize=4, alpha=alpha)
        ax.plot(data.Earth.r[i, 0], data.Earth.r[i, 1], "o", c="#0000FF", markersize=1, alpha=alpha)
    ax.plot(data.Sun.r[-1, 0], data.Sun.r[-1, 1], "o", c="#FFFF00", markersize=16)
    ax.plot(data.Earth.r[-1, 0], data.Earth.r[-1, 1], "o", c="#0000FF", markersize=4)
    ax.set_xlim(-1.5*AU, 1.5*AU)
    ax.set_ylim(-1.5*AU, 1.5*AU)
    fig.tight_layout()
    plt.show()
[45]:
plot_orbits(data_euler)
_images/2_advanced_integration_72_0.png

Leapfrog integration

As you can see, the Earth is not on a circular orbit. Its orbital distance is increasing and because of that the Earth could not achieve two full orbital cycles.

The problem is the simple integration scheme used here.
Every integration scheme has numerical error. Some have larger ones than others. Euler’s 1st order method is simply not suited for orbital integration.
But there is a way out: Symplectic integration
Symplectic integrators conserve the energy of the system.
One of these methods is the Leapfrog method.
Leapfrogging means the velocity and the position are not updated synchronous, but in between each other. They are leapfroggin each other.

In our case we first update the velocities by using half of the step size. Then we update the positions for a full time step. And finally, we update the velocities for another semi time step. We can do this by using the fstep keyword argument for the instructions, which is the fraction of the time for which this instruction is applied.

But there is one caveat: usually the fields that are integrated in an instructions set are only updated after all instructions have been executed, such that the second instruction still sees the original value of the first variable and not already the new value.

But for leapfrog integration, the instruction for the positions already need to know the new values of the velocities. And the second set of velocity instructions already need to know the new positions. We therefore have to add update instructions in between.

We do not have to add update instructions after the second set of velocity operations, because at the end of an instruction set all fields contained in the instruction set will be updated.

The new leapfrog instruction set now looks like this.

[46]:
instructions_leapfrog = [Instruction(schemes.expl_1_euler, sim.Sun.v,   fstep=0.5),
                         Instruction(schemes.expl_1_euler, sim.Earth.v, fstep=0.5),
                         Instruction(schemes.update,       sim.Sun.v             ),
                         Instruction(schemes.update,       sim.Earth.v           ),
                         Instruction(schemes.expl_1_euler, sim.Sun.r,   fstep=1.0),
                         Instruction(schemes.expl_1_euler, sim.Earth.r, fstep=1.0),
                         Instruction(schemes.update,       sim.Sun.r             ),
                         Instruction(schemes.update,       sim.Earth.r           ),
                         Instruction(schemes.expl_1_euler, sim.Sun.v,   fstep=0.5),
                         Instruction(schemes.expl_1_euler, sim.Earth.v, fstep=0.5),
                        ]

We can assign this new instruction set to our frame. To rerun the simulation we have to reset the initial conditions, that we’ve saved earlier.

[47]:
sim.integrator.instructions = instructions_leapfrog
sim.integrator.description = "Leapfrog integrator"
sim.Earth.r = r_Earth_ini
sim.Earth.v = v_Earth_ini
sim.Sun.r   = r_Sun_ini
sim.Sun.v   = v_Sun_ini
sim.t       = 0.

Additionally we have to reset the buffer of the namespace writer. Otherwise, we would simply add snapshots to the old dataset. Furthermore, we descrease the verbosity of the writer to prevent it from writing information on screen.

[48]:
sim.writer.reset()
sim.writer.verbosity = 0
[49]:
sim.run()
Execution time: 0:00:01
[50]:
data_leapfrog = sim.writer.read.all()
[51]:
plot_orbits(data_leapfrog)
_images/2_advanced_integration_82_0.png