julia nbody

nbody.jl

# This is a translation of the C nbody benchmark from
# The Computer Language Benchmarks Game

type Body
    x :: Float64
    y :: Float64
    z :: Float64
    vx :: Float64
    vy :: Float64
    vz :: Float64
    mass :: Float64
end

typealias System Array{Body, 1}

const solarMass = 4 * pi * pi
const daysPerYear = 365.24

function offsetMomentum(b  :: Body,    px :: Float64,
                        py :: Float64, pz :: Float64)
    b.vx = -px / solarMass
    b.vy = -py / solarMass
    b.vz = -pz / solarMass
end

function newSystem(bodies :: System)
    system = deepcopy(bodies)

    px = 0.0
    py = 0.0
    pz = 0.0

    for body in system
        px = px + body.vx * body.mass
	py = py + body.vy * body.mass
	pz = pz + body.vz * body.mass
    end

    offsetMomentum(system[1], px, py, pz)

    system
end

function energy(system :: System)
    e = 0.0

    for (i, body) in enumerate(system)
        e += 0.5 * body.mass * (body.vx * body.vx +
	           body.vy * body.vy + body.vz * body.vz)

        j = i + 1
	while j <= length(system)
	    body2 = system[j]
            j = j + 1
	    dx = body.x - body2.x
	    dy = body.y - body2.y
	    dz = body.z - body2.z
            distance = sqrt(dx * dx + dy * dy + dz * dz)
	    e = e - (body.mass * body2.mass) / distance
	end
    end
    e
end

function advance(system :: System, dt :: Float64)
    const len = length(system)
    i = 1
    for body in system
        j = i + 1
        while j <= len
            body2 = system[j]

            dx = body.x - body2.x
            dy = body.y - body2.y
            dz = body.z - body2.z

            dSquared = dx * dx + dy * dy + dz * dz
            distance = sqrt(dSquared)
            mag = dt / (dSquared * distance)
            
            body.vx = body.vx - (dx * body2.mass * mag)
            body.vy = body.vy - (dy * body2.mass * mag)
            body.vz = body.vz - (dz * body2.mass * mag)

            body2.vx = body2.vx + (dx * body.mass * mag)
            body2.vy = body2.vy + (dy * body.mass * mag)
            body2.vz = body2.vz + (dz * body.mass * mag)

            j = j + 1
        end
        i = i + 1
    end

    for body in system
        body.x = body.x + (dt * body.vx)
        body.y = body.y + (dt * body.vy)
        body.z = body.z + (dt * body.vz)
    end
end

const jupiter = Body(4.84143144246472090e+00,
                     -1.16032004402742839e+00,
                     -1.03622044471123109e-01,
                     1.66007664274403694e-03 * daysPerYear,
                     7.69901118419740425e-03 * daysPerYear,
                     -6.90460016972063023e-05 * daysPerYear,
                     9.54791938424326609e-04 * solarMass,
                     )

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

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

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

const sun = Body(0, 0, 0,
                 0, 0, 0,
                 solarMass)

function main()
    system = newSystem([sun, jupiter, saturn, uranus, neptune])
    println(energy(system))
    k = 0
    const n = 50000000
    advance([sun], 1.0) # jit
    @time while k < n
        advance(system, 0.01)
        k = k + 1
    end
    println(energy(system))
end

main()