nbody physics exercise

nbody.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

"""Non-interactive command line front end to NBody implementation"""

import time
import numpy as np
import matplotlib.pyplot as plt
import itertools

#from nbodyphysics import random_galaxy, move
#from nbodyphysics_vec1 import random_galaxy, move
from nbodyphysics_vec2 import random_galaxy, move




def nbody_benchmark(bodies_list, time_step, functions=None):
    """Run benchmark simulation without visualization"""

    x_max = 500
    y_max = 500
    z_max = 500
    dt = 1.0  # One year timesteps for better accuracy

    if functions is not None:
        random_galaxy, move = functions

    times = []
    for bodies in bodies_list:
        galaxy = random_galaxy(x_max, y_max, z_max, bodies)

        start = time.time()
        for _ in range(time_step):
            move(galaxy, dt)
        stop = time.time()

        print 'Simulated ' + str(bodies) + ' bodies for ' \
            + str(time_step) + ' timesteps in ' + str(stop - start) \
            + ' seconds'

        times.append(stop-start)

    return times 


def nbody_benchmark_functions(bodies_list, time_step):
    times = {}

    from nbodyphysics import random_galaxy as random_galaxy_seq
    from nbodyphysics import move as move_seq
    print '--------- Sequential time ...'
    times['sequential'] = nbody_benchmark(bodies_list, time_step, functions=(random_galaxy_seq,move_seq))

    from nbodyphysics_vec2 import random_galaxy as random_galaxy_vec
    from nbodyphysics_vec2 import  move as move_vec
    print '--------- Vectorized time ...'
    times['vectorized'] = nbody_benchmark(bodies_list, time_step, functions=(random_galaxy_vec,move_vec))

    return times


def plot_benchmark(times, bodies_list, time_step):

    fig, axx = plt.subplots(2,sharex=True)

    fields = sorted(times.keys())

    colorcycle = itertools.cycle(['b','g','r','y'])

    n = len(bodies_list)
    margin = 0.01

    width = (1. - 2.*margin)/float(n)/float(len(fields))
    ind = np.arange(n)

    axx[0].set_yscale('log')
    for k,key in enumerate(fields):
        axx[0].bar(ind+width*k, times[key], width, label=key, log=True, color=colorcycle.next())

    axx[0].legend(loc=0)
    axx[0].set_xticks(ind+width)
    axx[0].set_xticklabels([str(i) for i in bodies_list])
    axx[0].grid(True)
    axx[0].set_ylabel('time')

    speedup = np.array(times['sequential'])/np.array(times['vectorized'])
    width = (1. - 2.*margin)/float(n)
    axx[1].bar(ind, speedup, width, align='center', color=colorcycle.next())
    axx[1].grid(True)
    axx[1].set_ylabel('speedup')

    axx[1].set_xlabel('problem size')

    plt.show()
    fig.savefig('benchmark.png',dpi=300)


if __name__ == '__main__':
    #nbody_benchmark([10, 100, 1000], 10)

    nbody_list = [10, 100, 1000]
    time_step = 10 

    times = nbody_benchmark_functions(nbody_list, time_step)
    
    plot_benchmark(times, nbody_list, time_step)

nbodyphysics_vec1.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

"""NBody in N^2 complexity
Note that we are using only Newtonian forces and do not consider relativity
Neither do we consider collisions between stars
Thus some of our stars will accelerate to speeds beyond c
This is done to keep the simulation simple enough for teaching purposes

All the work is done in the calc_force, move and random_galaxy functions.
To vectorize the code these are the functions to transform.
"""
import numpy as np

# By using the solar-mass as the mass unit and years as the standard time-unit
# the gravitational constant becomes 1

G = 1.0


def calc_force(amass, aposition, bmass, bposition, dt):
    """Calculate forces between bodies
    F = ((G m_a m_b)/r^2)*((x_b-x_a)/r)
    """
    dpos = (bposition - aposition[:,np.newaxis])

    r = np.sum((dpos ** 2),axis=0) ** 0.5
    dspeed = G * amass * bmass / r ** 2 * (dpos / r) / amass * dt
    return np.sum(dspeed,axis=1)


def move(galaxy, dt):
    """Move the bodies
    first find forces and change velocity and then move positions
    """

    # get data from `galaxy`
    position = galaxy[['x','y','z']].view(('f8',3)).T
    speed = galaxy[['vx','vy','vz']].view(('f8',3)).T

    n = len(galaxy)

    # compute speed
    for i in xrange(n):
        # get indices of all other galaxies than the first one (`i`)
        all_others = np.hstack((np.arange(i),np.arange(i+1,n)))
        # compute the total speed contribution from all other galaxies
        speed[:,i] += calc_force(galaxy['m'][i], position[:,i], galaxy['m'][all_others], position[:,all_others], dt)
    
    # advance position
    for i in xrange(n):
        position[:,i] += speed[:,i] * dt

    # copy data back into `galaxy`
    # perhaps, this step becomes obsolete in a later version of NumPy
    for k,field in enumerate(['x','y','z']):
        galaxy[field] = position[k]
        galaxy['v'+field] = speed[k]


def random_galaxy(
    x_max,
    y_max,
    z_max,
    n,
    ):
    """Generate a galaxy of random bodies
    
    Changed here compared to the sequential version
    -----------------------------------------------
    `galaxy` is now a NumPy record array instead of a list of dictionaries

    """

    max_mass = 40.0  # Best guess of maximum known star


    # get the fields (i.e all the galaxy attributes)
    fields = ['m','x','y','z','vx','vy','vz']

    # make a custom data type for the record array
    galaxydtype = dict(names=fields,formats=['f8']*len(fields))

    # initialize the array with zeros
    galaxy = np.zeros(n,dtype=galaxydtype)

    # We let all bodies stand still initially
    # --> nothing to do here because galaxy['vx'] etc. already is 0

    # initialize mass
    galaxy['m'] = np.random.random(n) * np.random.randint(1, max_mass, n) / (4 * np.pi ** 2)

    # initialize position
    galaxy['x'] = np.random.random(n)*2*x_max-x_max
    galaxy['y'] = np.random.random(n)*2*y_max-y_max
    galaxy['z'] = np.random.random(n)*2*z_max-z_max

    return galaxy

nbodyphysics_vec2.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

"""NBody in N^2 complexity
Note that we are using only Newtonian forces and do not consider relativity
Neither do we consider collisions between stars
Thus some of our stars will accelerate to speeds beyond c
This is done to keep the simulation simple enough for teaching purposes

All the work is done in the calc_force, move and random_galaxy functions.
To vectorize the code these are the functions to transform.
"""
import numpy as np

# By using the solar-mass as the mass unit and years as the standard time-unit
# the gravitational constant becomes 1

G = 1.0


def calc_force(amass, aposition, bmass, bposition, dt):
    """Calculate forces between bodies
    F = ((G m_a m_b)/r^2)*((x_b-x_a)/r)
    """
    dpos = (bposition - aposition[:,np.newaxis])

    r = np.sum((dpos ** 2),axis=0) ** 0.5
    dspeed = G * amass * bmass / r ** 2 * (dpos / r) / amass * dt
    return np.sum(dspeed,axis=1)


def move(galaxy, dt):
    """Move the bodies
    first find forces and change velocity and then move positions
    """

    # unpack galaxy into its components
    (mass,position,speed) = galaxy

    n = len(mass)
    
    # compute speed
    for i in xrange(n):
        all_others = np.hstack((np.arange(i),np.arange(i+1,n)))
        speed[:,i] += calc_force(mass[i], position[:,i], mass[all_others], position[:,all_others], dt)

    # advance position
    for i in xrange(n):
        position[:,i] += speed[:,i] * dt
    

def random_galaxy(
    x_max,
    y_max,
    z_max,
    n,
    ):
    """Generate a galaxy of random bodies"""

    max_mass = 40.0  # Best guess of maximum known star

    # We let all bodies stand still initially

    # initialize mass
    mass = np.random.random(n) * np.random.randint(1, max_mass, n) / (4 * np.pi ** 2)

    # initialize position
    position = np.zeros((3,n))
    position[0] = np.random.random(n)*2*x_max-x_max
    position[1] = np.random.random(n)*2*y_max-y_max
    position[2] = np.random.random()*2*z_max-z_max

    # initialize speed
    speed = np.zeros((3,n))

    return (mass,position,speed)