NBabel calculations accelerated

nbabel_updated.jl

"""
This is an implementation of the NBabel N-body problem.
See nbabel.org for more information.

This is 'naive & native' Julia with type and loop annotations for
fastmath and vectorization, because this is really is no effort.

Without annotations the code will be as slow as naive Python/IDL.

    include("nbabel.jl")
    NBabel("IC/input128", show=true)

Julius Donnert INAF-IRA 2017
"""

using Printf
using DelimitedFiles

tbeg = 0.0
tend = 10.0
dt = 0.001

function NBabel(fname::String; tend = tend, dt = dt, show=false)

    if show
	    println("Reading file : $fname")
    end

    ID, mass, pos, vel = read_ICs(fname)

    return NBabelCalcs(ID, mass, pos, vel, tend, dt, show)
end

function NBabelCalcs(ID, mass, pos, vel, tend = tend, dt = dt, show=false)
    acc = similar(vel)
    acc = compute_acceleration(pos, mass, acc)
    last_acc = copy(acc)

    Ekin, Epot = compute_energy(pos, vel, mass)
    Etot_ICs = Ekin + Epot

    t = 0.0
    nstep = 0

    while t < tend

        pos = update_positions(pos, vel, acc, dt)

        acc, last_acc = last_acc, acc

        acc = compute_acceleration(pos, mass, acc)

        vel = update_velocities(vel, acc, last_acc, dt)
        
        t += dt
        nstep += 1

        if show && nstep%100 == 0

            Ekin, Epot = compute_energy(pos, vel, mass)
            Etot = Ekin + Epot
            dE = (Etot - Etot_ICs)/Etot_ICs

            @printf "t = %g, Etot=%g, Ekin=%g, Epot=%g, dE=%g \n" t Etot Ekin Epot dE
        end
    end

    Ekin, Epot = compute_energy(pos, vel, mass)
    Etot = Ekin + Epot
    return (; Ekin, Epot, Etot)
    # return size(pos,1)
end

function update_positions(pos, vel, acc, dt)
    N = length(pos)

    @fastmath @inbounds @simd for i in 1:N
        pos[i] = @. (0.5 * acc[i] * dt + vel[i])*dt + pos[i]
    end

    return pos
end

function update_velocities(vel, acc, last_acc, dt)
    N = length(vel)

    @fastmath @inbounds @simd for i in 1:N
        vel[i] = @. vel[i] + 0.5 * dt * (acc[i] + last_acc[i])
    end

    return vel
end

"""
Force calculation.
"""
function compute_acceleration(pos, mass, acc)
    N = length(pos)

    @inbounds @Threads.threads for i = 1:N

        pos_i = pos[i]
        acc_i = (0.0, 0.0, 0.0)

        @fastmath @inbounds @simd for j = 1:N
            if i != j

                dr = pos_i .- pos[j]

                rinv3 = 1/sqrt(sum(dr .^ 2)) ^ 3

                acc_i = acc_i .- mass[j] .* rinv3 .* dr
            end
        end

        acc[i] = acc_i
    end

    return acc
end


"""
Kinetic and potential energy.
"""
function compute_energy(pos, vel, mass)

    N = length(vel)

    Ekin = 0.0

    @simd for i = 1:N
        @inbounds Ekin += 0.5 * mass[i] * sum(vel[i] .^ 2)
    end

    Epot = 0.0

    @inbounds for i = 1:N-1
        pos_i = pos[i]
        @fastmath @inbounds for j = i+1:N
            dr = pos_i .- pos[j]

            rinv = 1/sqrt(sum(dr .^ 2))
            Epot -= mass[i] * mass[j] * rinv
        end
    end

    return Ekin, Epot
end

function read_ICs(fname::String)

    ICs = readdlm(fname)

    N = size(ICs,1)

    id = Array{Float64}(undef, N)
    mass = Array{Float64}(undef, N)
    pos = Array{Tuple{Float64, Float64, Float64}}(undef, N)
    vel = Array{Tuple{Float64, Float64, Float64}}(undef, N)

    id .= @view ICs[:, 1]
    mass .= @view ICs[:, 2]

    for i in axes(ICs, 1)
        pos[i] = (ICs[i, 3], ICs[i, 4], ICs[i, 5])
    end

    for i in axes(ICs, 1)
        vel[i] = (ICs[i, 6], ICs[i, 7], ICs[i, 8])
    end

    return id, mass, pos, vel
end


function main(args)
    NBabel(args[1], show=true)
end
main(ARGS)