<s> import time <EOL> def combinations ( l ) : <EOL> result = [ ] <EOL> for x in xrange ( len ( l ) - <NUM_LIT:1> ) : <EOL> for j in xrange ( x + <NUM_LIT:1> , len ( l ) ) : <EOL> y = l [ j ] <EOL> result . append ( ( l [ x ] , y ) ) <EOL> return result <EOL> PI = <NUM_LIT> <EOL> SOLAR_MASS = <NUM_LIT:4> * PI * PI <EOL> DAYS_PER_YEAR = <NUM_LIT> <EOL> BODIES = { <EOL> <NUM_LIT:1> : ( [ <NUM_LIT:0.0> , <NUM_LIT:0.0> , <NUM_LIT:0.0> ] , [ <NUM_LIT:0.0> , <NUM_LIT:0.0> , <NUM_LIT:0.0> ] , SOLAR_MASS ) , <EOL> <NUM_LIT:2> : ( [ <NUM_LIT> , <EOL> - <NUM_LIT> , <EOL> - <NUM_LIT> ] , <EOL> [ <NUM_LIT> * DAYS_PER_YEAR , <EOL> <NUM_LIT> * DAYS_PER_YEAR , <EOL> - <NUM_LIT> * DAYS_PER_YEAR ] , <EOL> <NUM_LIT> * SOLAR_MASS ) , <EOL> <NUM_LIT:3> : ( [ <NUM_LIT> , <EOL> <NUM_LIT> , <EOL> - <NUM_LIT> ] , <EOL> [ - <NUM_LIT> * DAYS_PER_YEAR , <EOL> <NUM_LIT> * DAYS_PER_YEAR , <EOL> <NUM_LIT> * DAYS_PER_YEAR ] , <EOL> <NUM_LIT> * SOLAR_MASS ) , <EOL> <NUM_LIT:4> : ( [ <NUM_LIT> , <EOL> - <NUM_LIT> , <EOL> - <NUM_LIT> ] , <EOL> [ <NUM_LIT> * DAYS_PER_YEAR , <EOL> <NUM_LIT> * DAYS_PER_YEAR , <EOL> - <NUM_LIT> * DAYS_PER_YEAR ] , <EOL> <NUM_LIT> * SOLAR_MASS ) , <EOL> <NUM_LIT:5> : ( [ <NUM_LIT> , <EOL> - <NUM_LIT> , <EOL> <NUM_LIT> ] , <EOL> [ <NUM_LIT> * DAYS_PER_YEAR , <EOL> <NUM_LIT> * DAYS_PER_YEAR , <EOL> - <NUM_LIT> * DAYS_PER_YEAR ] , <EOL> <NUM_LIT> * SOLAR_MASS ) } <EOL> SYSTEM = BODIES . values ( ) <EOL> PAIRS = combinations ( SYSTEM ) <EOL> def advance ( dt , n ) : <EOL> bodies = SYSTEM <EOL> pairs = PAIRS <EOL> for i in xrange ( n ) : <EOL> for ( ( ( x1 , y1 , z1 ) , v1 , m1 ) , <EOL> ( ( x2 , y2 , z2 ) , v2 , m2 ) ) in pairs : <EOL> dx = x1 - x2 <EOL> dy = y1 - y2 <EOL> dz = z1 - z2 <EOL> mag = dt * ( ( dx * dx + dy * dy + dz * dz ) ** ( - <NUM_LIT> ) ) <EOL> b1m = m1 * mag <EOL> b2m = m2 * mag <EOL> v1 [ <NUM_LIT:0> ] = v1 [ <NUM_LIT:0> ] - dx * b2m <EOL> v1 [ <NUM_LIT:1> ] = v1 [ <NUM_LIT:1> ] - dy * b2m <EOL> v1 [ <NUM_LIT:2> ] = v1 [ <NUM_LIT:2> ] - dz * b2m <EOL> v2 [ <NUM_LIT:0> ] = v2 [ <NUM_LIT:0> ] + dx * b1m <EOL> v2 [ <NUM_LIT:1> ] = v2 [ <NUM_LIT:1> ] + dy * b1m <EOL> v2 [ <NUM_LIT:2> ] = v2 [ <NUM_LIT:2> ] + dz * b1m <EOL> for ( r , ( vx , vy , vz ) , m ) in bodies : <EOL> r [ <NUM_LIT:0> ] = r [ <NUM_LIT:0> ] + dt * vx <EOL> r [ <NUM_LIT:1> ] = r [ <NUM_LIT:1> ] + dt * vy <EOL> r [ <NUM_LIT:2> ] = r [ <NUM_LIT:2> ] + dt * vz <EOL> def report_energy ( ) : <EOL> bodies = SYSTEM <EOL> pairs = PAIRS <EOL> e = <NUM_LIT:0.0> <EOL> for ( ( ( x1 , y1 , z1 ) , v1 , m1 ) , <EOL> ( ( x2 , y2 , z2 ) , v2 , m2 ) ) in pairs : <EOL> dx = x1 - x2 <EOL> dy = y1 - y2 <EOL> dz = z1 - z2 <EOL> e = e - ( m1 * m2 ) / ( ( dx * dx + dy * dy + dz * dz ) ** <NUM_LIT:0.5> ) <EOL> for ( r , ( vx , vy , vz ) , m ) in bodies : <EOL> e = e + m * ( vx * vx + vy * vy + vz * vz ) / <NUM_LIT> <EOL> return e <EOL> def offset_momentum ( ref , bodies , px , py , pz ) : <EOL> for ( r , ( vx , vy , vz ) , m ) in bodies : <EOL> px = px - vx * m <EOL> py = py - vy * m <EOL> pz = pz - vz * m <EOL> ( r , v , m ) = ref <EOL> v [ <NUM_LIT:0> ] = px / m <EOL> v [ <NUM_LIT:1> ] = py / m <EOL> v [ <NUM_LIT:2> ] = pz / m <EOL> NUMBER_OF_ITERATIONS = <NUM_LIT> <EOL> def main ( n , ref ) : <EOL> times = [ ] <EOL> for i in range ( n ) : <EOL> t0 = time . time ( ) <EOL> offset_momentum ( BODIES [ ref ] , SYSTEM , <NUM_LIT:0.0> , <NUM_LIT:0.0> , <NUM_LIT:0.0> ) <EOL> e1 = report_energy ( ) <EOL> advance ( <NUM_LIT> , NUMBER_OF_ITERATIONS ) <EOL> print e1 - report_energy ( ) <EOL> tk = time . time ( ) <EOL> times . append ( tk - t0 ) <EOL> return times <EOL> main ( <NUM_LIT> , <NUM_LIT:2> ) </s>
