nbody.py

from math import pi, sqrt

solar_mass = 4 * pi * pi
days_per_year = 365.24

class Body:
    def __init__(self, x, y, z, vx, vy, vz, mass):
        self.x = x
        self.y = y
        self.z = z
        self.vx = vx
        self.vy = vy
        self.vz = vz
        self.mass = mass

    def __str__(self):
        return "<%lf %lf %lf>, <%lf %lf %lf>, %lf" % \
                (self.x, self.y, self.z, self.vx, self.vy, self.vz, self.mass)

jupiter = Body( 4.84143144246472090e+00,                   # x
               -1.16032004402742839e+00,                   # y
               -1.03622044471123109e-01,                   # z
                1.66007664274403694e-03 * days_per_year,   # vx
                7.69901118419740425e-03 * days_per_year,   # vy
               -6.90460016972063023e-05 * days_per_year,   # vz
                9.54791938424326609e-04 * solar_mass)      # mass

saturn = Body( 8.34336671824457987e+00,
               4.12479856412430479e+00,
              -4.03523417114321381e-01,
              -2.76742510726862411e-03 * days_per_year,
               4.99852801234917238e-03 * days_per_year,
               2.30417297573763929e-05 * days_per_year,
               2.85885980666130812e-04 * solar_mass)

uranus = Body( 1.28943695621391310e+01,
              -1.51111514016986312e+01,
              -2.23307578892655734e-01,
               2.96460137564761618e-03 * days_per_year,
               2.37847173959480950e-03 * days_per_year,
              -2.96589568540237556e-05 * days_per_year,
               4.36624404335156298e-05 * solar_mass)

neptune = Body( 1.53796971148509165e+01,
               -2.59193146099879641e+01,
                1.79258772950371181e-01,
                2.68067772490389322e-03 * days_per_year,
                1.62824170038242295e-03 * days_per_year,
               -9.51592254519715870e-05 * days_per_year,
                5.15138902046611451e-05 * solar_mass)

sun = Body(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, solar_mass)


def init_sun(sun, bodies):
    px = 0.0
    py = 0.0
    pz = 0.0
    for b in bodies:
        px += b.vx * b.mass
        py += b.vy * b.mass
        pz += b.vz * b.mass
    sun.vx = -px / solar_mass
    sun.vy = -py / solar_mass
    sun.vz = -pz / solar_mass

import numpy as np
import hope
hope.config.optimize = True
hope.config.verbose = False
hope.config.keeptemp = True

def advance(N, dt, bodies):
    size = len(bodies)
    bodies_ndarray = np.array([(b.x, b.y, b.z, b.vx, b.vy, b.vz, b.mass) for b in bodies])
    advance_hope(N, dt, bodies_ndarray, size)

@hope.jit
def advance_hope(N, dt, bodies, size):
    for _ in xrange(N):
        for i in xrange(size):
            for j in xrange(i+1, size):
                dx = bodies[i,0] - bodies[j,0]
                dy = bodies[i,1] - bodies[j,1]
                dz = bodies[i,2] - bodies[j,2]
                dsq = dx**2 + dy**2 + dz**2
                distance = np.sqrt(dsq)
                mag = dt / (dsq * distance)

                bodies[i,3] -= dx * bodies[j,6] * mag
                bodies[i,4] -= dy * bodies[j,6] * mag
                bodies[i,5] -= dz * bodies[j,6] * mag

                bodies[j,3] += dx * bodies[i,6] * mag
                bodies[j,4] += dy * bodies[i,6] * mag
                bodies[j,5] += dz * bodies[i,6] * mag

        for k in xrange(size):
            bodies[k,0] += dt * bodies[k,3]
            bodies[k,1] += dt * bodies[k,4]
            bodies[k,2] += dt * bodies[k,5]

def energy(bodies):
    e = 0.0
    for i, body1 in enumerate(bodies):
        e += 0.5 * body1.mass * (body1.vx**2 + body1.vy**2 + body1.vz**2)
        for body2 in bodies[i+1:]:
            dx = body1.x - body2.x
            dy = body1.y - body2.y
            dz = body1.z - body2.z
            distance = sqrt(dx**2 + dy**2 + dz**2)
            e -= (body1.mass * body2.mass) / distance
    return e

def nbody(N):

    bodies = [sun, jupiter, saturn, uranus, neptune]
    init_sun(bodies[0], bodies)

    print energy(bodies)
    advance(N, 0.01, bodies)
    print energy(bodies)


if __name__ == '__main__':
    import sys
    nbody(int(sys.argv[1]))

@grakic I just stumbled across your blog entry about evaluating the HOPE package. Using HOPE to run a planet-n-body simulation is pretty cool! 👍

I’m currently writing a set of examples to show the various features of the package. I therefore wanted to ask you if I could add your code to the set (or you could just send me a PR).