n-body-code

nb.py

import argparse

import matplotlib.pyplot as plt
import numpy as np
from scipy.constants import G

day = 24 * 60 * 60
year = day * 365

parser = argparse.ArgumentParser(description="This is a simple N-body code.")
parser.add_argument("-t", dest="timestep", action="store",
                    default=1, type=float,
                    help="length of timestep for numerical integration (in days)")
parser.add_argument("-m", dest="max_time", action="store", default=1, type=int,
                    help="maximum time to compute (in years)")
parser.add_argument("-p", dest="plot_interval", action="store", default=10, type=int,
                    help="only plot every n calculated positions")
parser.add_argument("-i", dest="inputfile", action="store", default="3Jupiters.input",
                    help="file containing data on all bodies (one body per line, 5 columns: mass x-vector v-vector")

args = parser.parse_args()
print(args)

ts = args.timestep * day
max_time = args.max_time * year

bodies = []
time = 0
with open(args.inputfile) as inputfile:
    body_number = 0
    for bodyline in inputfile:
        data = [float(n) for n in bodyline.strip().split()]
        mass = data[0]
        location = np.array(data[1:3])
        velocity = np.array(data[3:5])
        body = (body_number, mass, location, velocity)
        print(body)
        bodies.append(body)
        body_number += 1

while time <= max_time:
    for body in bodies:
        (body_number, mass, location, velocity) = body
        for other_body in bodies:
            (other_body_number, other_mass, other_location, other_velocity) = other_body
            if body_number == other_body_number:  # don't try to calculate distance to itself (division by 0)
                continue
            velocity += -G * other_mass / np.linalg.norm(location - other_location) ** 3 * (
                    location - other_location) * ts
        location += velocity * ts
        if time % (ts * 10) == 0:  # only plot every n-th result (speedup!)
            plt.plot(location[0], location[1], "o", c="C" + str(body_number))
    time += ts
    print(time / max_time * 100)
plt.show()