Adaptive RK4 Orbits

Adaptive RK4 Orbits#

import numpy as np
import matplotlib.pyplot as plt

We’ll use an implementation of the orbit problem with RK4 written as a module. The source (in both C++ and python) is here:

zingale/computational_astrophysics

import orbit_adaptive as orbit

Let’s first look at the source

%cat orbit_adaptive.py
Hide code cell output
"""
orbital motion.  We consider low mass objects orbiting the Sun.  We
work in units of AU, yr, and solar masses.  From Kepler's third law:

4 pi**2 a**3 = G M P**2

if a is in AU, P is in yr, and M is in solar masses, then

a**3 = P**2

and therefore

4 pi**2 = G

we work in coordinates with the Sun at the origin

This version implements adaptive timestepping
"""

import numpy as np
import matplotlib.pyplot as plt

# global parameters
GM = 4.0 * np.pi**2  # assuming M = 1 solar mass

class OrbitState:
    # a container to hold the star positions
    def __init__(self, x, y, u, v):
        self.x = x
        self.y = y
        self.u = u
        self.v = v

    def __add__(self, other):
        return OrbitState(self.x + other.x, self.y + other.y,
                          self.u + other.u, self.v + other.v)

    def __sub__(self, other):
        return OrbitState(self.x - other.x, self.y - other.y,
                          self.u - other.u, self.v - other.v)

    def __mul__(self, other):
        return OrbitState(other * self.x, other * self.y,
                          other * self.u, other * self.v)

    def __rmul__(self, other):
        return self.__mul__(other)

    def __str__(self):
        return f"{self.x:10.6f} {self.y:10.6f} {self.u:10.6f} {self.v:10.6f}"


class OrbitsRK4:
    """ model the evolution of a single planet around the Sun"""

    def __init__(self, a, e):
        """ a = semi-major axis (in AU),
            e = eccentricity """

        self.a = a
        self.e = e

        x0 = 0.0          # start at x = 0 by definition
        y0 = a*(1.0 - e)  # start at perihelion

        # perihelion velocity (see C&O Eq. 2.33 for ex)
        u0 = -np.sqrt( (GM/a)* (1.0 + e) / (1.0 - e) )
        v0 = 0.0

        self.history = [OrbitState(x0, y0, u0, v0)]
        self.time = [0.0]

        self.n_reset = None

    def npts(self):
        """ return the number of integration points"""
        return len(self.time)

    def get_time(self, n):
        """ return the time corresponding to step n"""
        return self.time[n]

    def kepler_period(self):
        """ return the period of the orbit in yr """
        return np.sqrt(self.a**3)

    def final_R(self):
        """ the radius at the final integration time """

        state = self.history[-1]
        return np.sqrt(state.x**2 + state.y**2)

    def displacement(self):
        """ distance between the starting and ending point """

        return np.sqrt((self.history[0].x - self.history[-1].x)**2 +
                       (self.history[0].y - self.history[-1].y)**2 )

    def energy(self, n):
        """ return the energy (per unit mass) at time n """

        state = self.history[n]
        return 0.5 * (state.u**2 + state.v**2) - \
            GM / np.sqrt(state.x**2 + state.y**2)

    def rhs(self, state):
        """ RHS of the equations of motion."""

        # current radius
        r = np.sqrt(state.x**2 + state.y**2)

        # position
        xdot = state.u
        ydot = state.v

        # velocity
        udot = -GM * state.x / r**3
        vdot = -GM * state.y / r**3

        return OrbitState(xdot, ydot, udot, vdot)

    def single_step(self, state_old, dt):
        """ take a single RK-4 timestep from t to t+dt for the system
        ydot = rhs """

        # get the RHS at several points
        ydot1 = self.rhs(state_old)

        state_tmp = state_old + 0.5 * dt * ydot1
        ydot2 = self.rhs(state_tmp)

        state_tmp = state_old + 0.5 * dt * ydot2
        ydot3 = self.rhs(state_tmp)

        state_tmp = state_old + dt * ydot3
        ydot4 = self.rhs(state_tmp)

        # advance
        state_new = state_old + (dt / 6.0) * (ydot1 + 2.0 * ydot2 +
                                              2.0 * ydot3 + ydot4)

        return state_new

    def integrate(self, dt_in, err, tmax):
        """ integrate the equations of motion using 4th order R-K method with an
            adaptive stepsize, to try to achieve the relative error err.  dt
            here is the initial timestep
            """

        S1 = 0.9
        S2 = 4.0

        # start with the old timestep
        dt_new = dt_in
        dt = dt_in

        n_reset = 0
        t = 0

        while t < tmax:

            state_old = self.history[-1]

            if err > 0.0:

                # adaptive stepping iteration loop -- keep trying to take a step
                # until we achieve our desired error

                rel_error = 1.e10

                n_try = 0
                while rel_error > err:
                    dt = dt_new
                    if t+dt > tmax:
                        dt = tmax-t

                    # take 2 half steps
                    state_tmp = self.single_step(state_old, 0.5*dt)
                    state_new = self.single_step(state_tmp, 0.5*dt)

                    # now take just a single step to cover dt
                    state_single = self.single_step(state_old, dt)

                    # state_new should be more accurate than state_single since it
                    # used smaller steps

                    # estimate the relative error now

                    rel_error = max(abs((state_new.x - state_single.x) / state_single.x),
                                    abs((state_new.y - state_single.y) / state_single.y),
                                    abs((state_new.u - state_single.u) / state_single.u),
                                    abs((state_new.v - state_single.v) / state_single.v))

                    # adaptive timestep algorithm for RK4

                    dt_est = dt * abs(err/rel_error)**0.2
                    dt_new = min(max(S1*dt_est, dt/S2), S2*dt)

                    n_try += 1

                if n_try > 1:
                    # n_try = 1 if we took only a single try at the step
                    n_reset += (n_try-1)

            else:

                # take just a single step
                if t + dt > tmax:
                    dt = tmax - t

                # take just a single step to cover dt
                state_new = self.single_step(state_old, dt)

            # successful step
            t += dt

            # store
            self.time.append(t)
            self.history.append(state_new)

    def plot(self, points=False):
        """plot the orbit"""
        fig, ax = plt.subplots()
        x = [q.x for q in self.history]
        y = [q.y for q in self.history]

        # draw the Sun
        ax.scatter([0], [0], marker=(20,1), color="y", s=250)

        if points:
            ax.scatter(x, y)
        else:
            ax.plot(x, y)
        ax.set_aspect("equal")
        return fig

if __name__ == "__main__":

    o = OrbitsRK4(1.0, 0.95)
    o.integrate(0.05, 1.e-8, 10.0);

    print("number of points = ", o.npts())

    # plot the orbit

    fig = o.plot()
    fig.tight_layout()
    fig.savefig("adaptive_rk4_orbit.png", bbox_inches="tight")


    E = []
    for n in range(o.npts()):
        E.append(o.energy(n))

    fig, ax = plt.subplots()
    ax.plot(o.time, E/E[0])
    ax.set_xlabel("time [yr]")
    ax.set_ylabel("E/E(t=0)")
    fig.savefig("adaptive_rk4_energy.png")

Highly eccentric orbit#

Let’s create an orbit with a tolerance of \(10^{-5}\), and the a very high eccentricity (\(a = 1\), \(e = 0.95\))

o = orbit.OrbitsRK4(1.0, 0.95)
o.integrate(0.05, 1.e-5, o.kepler_period())

How many points did it take?

o.npts()
92

Let’s plot it

fig = o.plot(points=True)
fig.set_size_inches(5, 12)
../_images/2502cac03d8651b5abba8d413ec4fe33c19cf91c1c4482ac7c1d1a72917fba01.png

Clearly we see that the size of the timestep is much larger at aphelion as compared to perihelion.

Let’s create a version with a fixed timestep, with a timestep of \(0.001~\mathrm{yr}\). We can do this by passing in a negative error.

o_fixed = orbit.OrbitsRK4(1.0, 0.95)
o_fixed.integrate(0.001, -1, o_fixed.kepler_period())
fig = o_fixed.plot()
fig.set_size_inches(5, 12)
../_images/8dd0c4ebb94d6b4c2bf0301f5f9dee044aaf911e40d1cb26796797063ff821cb.png
print(o_fixed.npts())
1001

We need a really small fixed step size to integrate this at all reasonably

Let’s look at the adaptive version some more. First, let’s look at the timestep evolution.

ts = np.array(o.time)

dt = ts[1:] - ts[:-1]
tc = 0.5 * (ts[1:] + ts[:-1])

fig, ax = plt.subplots()
ax.plot(tc, dt)
ax.set_xlabel("t [yr]")
ax.set_ylabel("dt [yr]")
ax.set_yscale("log")
../_images/29389e491948abaec1e61c829939834ab3c9bad306c1f855545aa9b864edc364.png

Notice that the timestep changes by ~ 3 orders of magnitude over the evolution.

What about energy conservation?

e = []
ts = o.time
for n in range(o.npts()):
    e.append(o.energy(n))
    
fig, ax = plt.subplots()
ax.plot(ts, e/e[0])
ax.set_xlabel("time [yr]")
ax.set_ylabel("E/E(t=0)")
ax.ticklabel_format(useOffset=False)
../_images/a416ce75a68fe4583f6d6172581b9a85520615941de846de198217cd83c6ce7a.png

Now let’s do 10 orbits

o = orbit.OrbitsRK4(1.0, 0.95)
o.integrate(0.05, 1.e-5, 20*o.kepler_period())
e = []
ts = o.time
for n in range(o.npts()):
    e.append(o.energy(n))
    
fig, ax = plt.subplots()
ax.plot(ts, e/e[0])
ax.set_xlabel("time [yr]")
ax.set_ylabel("E/E(t=0)")
ax.ticklabel_format(useOffset=False)
../_images/eb7e5ebcc49c9f3be8c59233f72c8a7018917f257020abb98e0138c5b99d1abe.png

We can see that RK4 does not conserve energy! There is a steady decrease in the total energy over 20 orbits. If we wanted to evolve for millions of years, this would certainly be a problem. Making the tolerance tighter certainly will help, but it will also make things a lot more expensive.

The problem is that 4th order Runge-Kutta does not know anything about energy conservation.