Python RK4 orbits

orbits.py

import math
import random
import matplotlib.pyplot as plot
from mpl_toolkits.mplot3d import Axes3D

class point:
    def __init__(self, x,y,z):
        self.x = x
        self.y = y
        self.z = z

class body:
    def __init__(self, location, mass, velocity, name = ""):
        self.location = location
        self.mass = mass
        self.velocity = velocity
        self.name = name

def partial_step(point1, point2, time_step):
    ret = point(0,0,0)
    ret.x = point1.x + point2.x * time_step
    ret.y = point1.y + point2.y * time_step
    ret.z = point1.z + point2.z * time_step
    return ret

class Euler_integrator:
    def __init__(self, time_step, bodies):
        self.time_step = time_step
        self.bodies = bodies

    def calculate_single_body_acceleration(self, body_index):
        G_const = 6.67408e-11 #m3 kg-1 s-2
        acceleration = point(0,0,0)
        target_body = self.bodies[body_index]
        for index, external_body in enumerate(bodies):
            if index != body_index:
                r = (target_body.location.x - external_body.location.x)**2 + (target_body.location.y - external_body.location.y)**2 + (target_body.location.z - external_body.location.z)**2
                r = math.sqrt(r)
                tmp = G_const * external_body.mass / r**3
                acceleration.x += tmp * (external_body.location.x - target_body.location.x)
                acceleration.y += tmp * (external_body.location.y - target_body.location.y)
                acceleration.z += tmp * (external_body.location.z - target_body.location.z)
    
        return acceleration
    
    def update_location(self):
        for target_body in self.bodies:
            target_body.location.x += target_body.velocity.x * self.time_step
            target_body.location.y += target_body.velocity.y * self.time_step
            target_body.location.z += target_body.velocity.z * self.time_step
           
    def compute_velocity(self):
        for body_index, target_body in enumerate(self.bodies):
            acceleration = self.calculate_single_body_acceleration(body_index)
            target_body.velocity.x += acceleration.x * self.time_step
            target_body.velocity.y += acceleration.y * self.time_step
            target_body.velocity.z += acceleration.z * self.time_step 
    
    
    def update_location(self):
        for target_body in self.bodies:
            target_body.location.x += target_body.velocity.x * self.time_step
            target_body.location.y += target_body.velocity.y * self.time_step
            target_body.location.z += target_body.velocity.z * self.time_step

    def compute_gravity_step(self):
        self.compute_velocity()
        self.update_location()


  

class RK4_integrator:
    def __init__(self, time_step, bodies):
        self.time_step = time_step
        self.bodies = bodies

    def calculate_single_body_acceleration(self, body_index):
        G_const = 6.67408e-11 #m3 kg-1 s-2
        acceleration = point(0,0,0)
        target_body = self.bodies[body_index]

        k1 = point (0,0,0)
        k2 = point (0,0,0)
        k3 = point (0,0,0)
        k4 = point (0,0,0)
        tmp_loc = point (0,0,0)
        tmp_vel = point (0,0,0)

        for index, external_body in enumerate(self.bodies):
            if index != body_index:
                r = (target_body.location.x - external_body.location.x)**2 + (target_body.location.y - external_body.location.y)**2 + (target_body.location.z - external_body.location.z)**2
                r = math.sqrt(r)
                tmp = G_const * external_body.mass / r**3

                k1.x = tmp * (external_body.location.x - target_body.location.x)
                k1.y = tmp * (external_body.location.y - target_body.location.y)
                k1.z = tmp * (external_body.location.z - target_body.location.z)
                
                tmp_vel = partial_step(target_body.velocity, k1, 0.5)
                tmp_loc = partial_step(target_body.location, tmp_vel, 0.5)

	    		#k2 - acceleration 0.5 timesteps in the future based on k1 acceleration value
                k2.x = (tmp_loc.x - (tmp_loc.x + tmp_vel.x * 0.5 * self.time_step)) * tmp
                k2.y = (tmp_loc.y - (tmp_loc.y + tmp_vel.y * 0.5 * self.time_step)) * tmp
                k2.z = (tmp_loc.z - (tmp_loc.z + tmp_vel.z * 0.5 * self.time_step)) * tmp

                tmp_vel = partial_step(target_body.velocity, k2, 0.5)

	    		#k3 acceleration 0.5 timesteps in the future using k2 acceleration
                k3.x = (tmp_loc.x - (tmp_loc.x + tmp_vel.x * 0.5 * self.time_step)) * tmp
                k3.y = (tmp_loc.y - (tmp_loc.y + tmp_vel.y * 0.5 * self.time_step)) * tmp
                k3.z = (tmp_loc.z - (tmp_loc.z + tmp_vel.z * 0.5 * self.time_step)) * tmp

                tmp_vel = partial_step(target_body.velocity, k3, 1)
                tmp_loc = partial_step(target_body.location, tmp_vel, 0.5)

	    		#k4 - location 1 timestep in the future using k3 acceleration
                k4.x = (external_body.location.x - (tmp_loc.x + tmp_vel.x * self.time_step)) * tmp;
                k4.y = (external_body.location.y - (tmp_loc.y + tmp_vel.y * self.time_step)) * tmp;
                k4.z = (external_body.location.z - (tmp_loc.z + tmp_vel.z * self.time_step)) * tmp;

                acceleration.x += (k1.x + k2.x * 2 + k3.x * 2 + k4.x) / 6;
                acceleration.y += (k1.y + k2.y * 2 + k3.y * 2 + k4.y) / 6;
                acceleration.z += (k1.z + k2.z * 2 + k3.z * 2 + k4.z) / 6;

        return acceleration
    
    def update_location(self):
        for target_body in self.bodies:
            target_body.location.x += target_body.velocity.x * self.time_step
            target_body.location.y += target_body.velocity.y * self.time_step
            target_body.location.z += target_body.velocity.z * self.time_step
           
    def compute_velocity(self):
        for body_index, target_body in enumerate(self.bodies):
            acceleration = self.calculate_single_body_acceleration(body_index)
            target_body.velocity.x += acceleration.x * self.time_step
            target_body.velocity.y += acceleration.y * self.time_step
            target_body.velocity.z += acceleration.z * self.time_step 
    
    
    def update_location(self):
        for target_body in self.bodies:
            target_body.location.x += target_body.velocity.x * self.time_step
            target_body.location.y += target_body.velocity.y * self.time_step
            target_body.location.z += target_body.velocity.z * self.time_step

    def compute_gravity_step(self):
        self.compute_velocity()
        self.update_location()
    



def plot_output(bodies, outfile = None):
    fig = plot.figure()
    colours = ['r','b','g','y','m','c']
    ax = fig.add_subplot(1,1,1, projection='3d')
    max_range = 0
    for current_body in bodies: 
        max_dim = max(max(current_body["x"]),max(current_body["y"]),max(current_body["z"]))
        if max_dim > max_range:
            max_range = max_dim
        ax.plot(current_body["x"], current_body["y"], current_body["z"], c = random.choice(colours), label = current_body["name"])        
    
    ax.set_xlim([-max_range,max_range])    
    ax.set_ylim([-max_range,max_range])
    ax.set_zlim([-max_range,max_range])
    ax.legend()        

    if outfile:
        plot.savefig(outfile)
    else:
        plot.show()

def run_simulation(integrator, names = None, number_of_steps = 10000, report_freq = 100):

    #create output container for each body
    body_locations_hist = []
    for current_body in bodies:
        body_locations_hist.append({"x":[], "y":[], "z":[], "name":current_body.name})
        
    for i in range(1, int(number_of_steps)):
        integrator.compute_gravity_step()            
        
        if i % report_freq == 0:
            for index, body_location in enumerate(body_locations_hist):
                body_location["x"].append(bodies[index].location.x)
                body_location["y"].append(bodies[index].location.y)           
                body_location["z"].append(bodies[index].location.z)       

    return body_locations_hist        
            
#planet data (location (m), mass (kg), velocity (m/s)
sun = {"location":point(0,0,0), "mass":2e30, "velocity":point(0,0,0)}
mercury = {"location":point(0,5.0e10,0), "mass":3.285e23, "velocity":point(47000,0,0)}
venus = {"location":point(0,1.1e11,0), "mass":4.8e24, "velocity":point(35000,0,0)}
earth = {"location":point(0,1.5e11,1e10), "mass":6e24, "velocity":point(30000,0,0)}
mars = {"location":point(0,2.2e11,0), "mass":6.4e23, "velocity":point(24000,0,0)}
jupiter = {"location":point(0,7.7e11,0), "mass":1e28, "velocity":point(13000,0,0)}
saturn = {"location":point(0,1.4e12,0), "mass":5.7e26, "velocity":point(9000,0,0)}
uranus = {"location":point(0,2.8e12,0), "mass":8.7e25, "velocity":point(6835,0,0)}
neptune = {"location":point(0,4.5e12,0), "mass":1e26, "velocity":point(5477,0,0)}
pluto = {"location":point(0,3.7e12,0), "mass":1.3e22, "velocity":point(4748,0,0)}

if __name__ == "__main__":

    #build list of planets in the simulation, or create your own
    bodies = [
        body( location = sun["location"], mass = sun["mass"], velocity = sun["velocity"], name = "sun"),
        body( location = earth["location"], mass = earth["mass"], velocity = earth["velocity"], name = "earth"),
        body ( location = jupiter["location"], mass = jupiter["mass"], velocity = jupiter["velocity"], name = "jupiter")
        ]
    
    integrator = RK4_integrator(time_step = 100, bodies = bodies)
    motions = run_simulation(integrator, number_of_steps = 1e4, report_freq = 1e2)
    plot_output(motions)

Hi @benrules2, for RK2 is just a matter of ignoring lines 108 to 121, and calculating accelerations as acceleration.x += k2.x etc.?

Hello @oy321, no, it is not that easy. The present RK4 code is wrong in at least two places.

• First it uses a first-order update for the positions, which destroys the order of the method.
• Then the whole body of the RK4 stages is rather strange, the force computation is not repeated for the temporary position of the stage, stages 2 and 3 should look at least like stage 4, but they would have to also use the updated versions of the state data of external_body, including the updated radius and thus force factor.

There are two approaches to repair this in an intelligible way

1. Provide the bodies with an acceleration field and the solar system object with a function that computes the accelerations at the current state. Also add a constructor with arguments (system1, system2, dt) that computes the states from the states of system1, the derivatives of system 2 and the time step dt. Then the integrator would construct temporary solar system objects corresponding to the stages and in a final sweep collect all the derivatives data to update the state of the original system. Or keep everything in one solar system object and have each body contain a list/stack of states (pos, vel, acc). This might be easier to coordinate.
2. Provide the bodies with an index range for their state data in a flat array and the solar system object with functionality to read and export positions, velocities and (after their computation) accelerations to such a flat array. Then RK4 can be implemented with simple vector arithmetic in the fashion of https://stackoverflow.com/a/60414658/3088138 or with some insight into the scheme https://scicomp.stackexchange.com/a/34330/6839. The solar system object is then only used for the physics part of the computation and for evaluation and visualization of the state data, but not for the time propagation. It would also be easy to write an ODE function that can be passed to a standard solver.