<s> import numpy as np <EOL> import time <EOL> from numba import * <EOL> flopsPerInteraction = <NUM_LIT:30> <EOL> SOFTENING_SQUARED = <NUM_LIT> <EOL> def normalize ( vector ) : <EOL> dist = np . sqrt ( ( vector * vector ) . sum ( ) ) <EOL> if dist > <NUM_LIT> : <EOL> vector /= dist <EOL> return dist <EOL> def randomize_bodies ( pos , vel , cluster_scale , velocity_scale , n ) : <EOL> np . random . seed ( <NUM_LIT> ) <EOL> scale = cluster_scale <EOL> vscale = scale * velocity_scale <EOL> inner = <NUM_LIT> * scale <EOL> outer = <NUM_LIT> * scale <EOL> i = <NUM_LIT:0> <EOL> while i < n : <EOL> point = np . random . random ( <NUM_LIT:3> ) / <NUM_LIT> <EOL> length = normalize ( point ) <EOL> if length > <NUM_LIT:1.> : <EOL> continue <EOL> pos [ i , : <NUM_LIT:3> ] = point * ( ( inner + ( outer - inner ) ) * np . random . random ( <NUM_LIT:3> ) ) <EOL> pos [ i , <NUM_LIT:3> ] = <NUM_LIT:1.0> <EOL> axis = np . array ( ( <NUM_LIT:0.> , <NUM_LIT:0.> , <NUM_LIT:1.> ) ) <EOL> normalize ( axis ) <EOL> if ( <NUM_LIT:1> - ( ( point * axis ) . sum ( ) ) ) < <NUM_LIT> : <EOL> axis [ <NUM_LIT:0> ] = point [ <NUM_LIT:1> ] <EOL> axis [ <NUM_LIT:1> ] = point [ <NUM_LIT:0> ] <EOL> normalize ( axis ) <EOL> vv = np . cross ( pos [ i , : <NUM_LIT:3> ] , axis ) <EOL> vel [ i , : <NUM_LIT:3> ] = vscale * vv <EOL> vel [ i , <NUM_LIT:3> ] = <NUM_LIT:1.0> <EOL> i += <NUM_LIT:1> <EOL> def check_correctness ( pin , pout , v , dt , n , integrate_0 , integrate_1 ) : <EOL> pin_ref = np . zeros_like ( pin ) <EOL> pout_ref = np . zeros_like ( pout ) <EOL> v_ref = np . zeros_like ( v ) <EOL> randomize_bodies ( pin_ref , v_ref , <NUM_LIT> , <NUM_LIT> , n ) <EOL> integrate_0 ( pout_ref , pin , np . copy ( v ) , dt , n ) <EOL> integrate_1 ( pout , pin , np . copy ( v ) , dt , n ) <EOL> errt = <NUM_LIT:0> <EOL> errmax = <NUM_LIT:0> <EOL> errs = np . fabs ( pout_ref - pout ) . reshape ( <NUM_LIT:4> * n ) <EOL> errt = errs . sum ( ) <EOL> errmax = errs . max ( ) <EOL> print ( "<STR_LIT>" % ( errmax , errt ) ) <EOL> def check_overflow ( x ) : <EOL> return np . isnan ( np . sum ( x ) ) <EOL> def body_body_interaction ( force , pos_mass0 , pos_mass1 ) : <EOL> r = pos_mass1 [ : <NUM_LIT:3> ] - pos_mass0 [ : <NUM_LIT:3> ] <EOL> dist_sqr = ( r * r ) . sum ( ) <EOL> dist_sqr += SOFTENING_SQUARED <EOL> inv_dist = <NUM_LIT:1.0> / np . sqrt ( dist_sqr ) <EOL> inv_dist_cube = inv_dist * inv_dist * inv_dist <EOL> s = pos_mass1 [ <NUM_LIT:3> ] * inv_dist_cube <EOL> force += r * s <EOL> def integrate ( position_out , position_in , velocity , delta_time , n ) : <EOL> for i in range ( n ) : <EOL> p = position_in [ i ] <EOL> f = np . zeros ( <NUM_LIT:3> ) <EOL> for j in range ( n ) : <EOL> body_body_interaction ( f , p , position_in [ j ] ) <EOL> inv_mass = velocity [ i , <NUM_LIT:3> ] <EOL> v = velocity [ i , : <NUM_LIT:3> ] + f * inv_mass * delta_time <EOL> position_out [ i , : <NUM_LIT:3> ] = p [ : <NUM_LIT:3> ] + v * delta_time <EOL> position_out [ i , <NUM_LIT:3> ] = position_in [ i , <NUM_LIT:3> ] <EOL> velocity [ i , : <NUM_LIT:3> ] = v <EOL> def compute_perf_stats ( milliseconds , iterations , n ) : <EOL> interactionsPerSecond = float ( n * n ) <EOL> interactionsPerSecond *= <NUM_LIT> * iterations * <NUM_LIT:1000> / milliseconds <EOL> return interactionsPerSecond * flopsPerInteraction ; <EOL> def main ( * args ) : <EOL> n = <NUM_LIT> <EOL> iterations = <NUM_LIT:10> <EOL> dt = <NUM_LIT> <EOL> if len ( args ) > <NUM_LIT:0> : <EOL> n = int ( args [ <NUM_LIT:0> ] ) <EOL> if len ( args ) > <NUM_LIT:1> : <EOL> iterations = int ( args [ <NUM_LIT:1> ] ) <EOL> pin = np . zeros ( ( n , <NUM_LIT:4> ) ) <EOL> pout = np . zeros ( ( n , <NUM_LIT:4> ) ) <EOL> v = np . zeros ( ( n , <NUM_LIT:4> ) ) <EOL> randomize_bodies ( pin , v , <NUM_LIT> , <NUM_LIT> , n ) <EOL> check_correctness ( pin , pout , v , dt , n , integrate , integrate ) <EOL> time0 = time . time ( ) <EOL> for i in range ( iterations ) : <EOL> integrate ( pout , pin , v , dt , n ) <EOL> pin , pout = pout , pin <EOL> time1 = time . time ( ) <EOL> ms = ( time1 - time0 ) * <NUM_LIT:1000> <EOL> gf = compute_perf_stats ( ms , iterations , n ) <EOL> print ( "<STR_LIT>" % iterations ) <EOL> print ( "<STR_LIT>" % ( ms , gf ) ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> import sys <EOL> main ( * sys . argv [ <NUM_LIT:1> : ] ) </s>
