paulb@200 | 1 | #!/usr/bin/python -OO |
paulb@200 | 2 | # The Computer Language Shootout Benchmarks |
paulb@200 | 3 | # http://shootout.alioth.debian.org/ |
paulb@200 | 4 | # |
paulb@200 | 5 | # contributed by Kevin Carson |
paulb@200 | 6 | |
paulb@200 | 7 | import sys |
paulb@200 | 8 | |
paulb@200 | 9 | pi = 3.14159265358979323 |
paulb@200 | 10 | solar_mass = 4 * pi * pi |
paulb@200 | 11 | days_per_year = 365.24 |
paulb@200 | 12 | |
paulb@200 | 13 | class body : |
paulb@200 | 14 | pass |
paulb@200 | 15 | |
paulb@200 | 16 | def advance(bodies, dt) : |
paulb@200 | 17 | for i in xrange(len(bodies)) : |
paulb@200 | 18 | b = bodies[i] |
paulb@200 | 19 | |
paulb@200 | 20 | for j in xrange(i + 1, len(bodies)) : |
paulb@200 | 21 | b2 = bodies[j] |
paulb@200 | 22 | |
paulb@200 | 23 | dx = b.x - b2.x |
paulb@200 | 24 | dy = b.y - b2.y |
paulb@200 | 25 | dz = b.z - b2.z |
paulb@200 | 26 | distance = (dx**2 + dy**2 + dz**2)**0.5 |
paulb@200 | 27 | |
paulb@200 | 28 | b_mass_x_mag = dt * b.mass / distance**3 |
paulb@200 | 29 | b2_mass_x_mag = dt * b2.mass / distance**3 |
paulb@200 | 30 | |
paulb@200 | 31 | b.vx -= dx * b2_mass_x_mag |
paulb@200 | 32 | b.vy -= dy * b2_mass_x_mag |
paulb@200 | 33 | b.vz -= dz * b2_mass_x_mag |
paulb@200 | 34 | b2.vx += dx * b_mass_x_mag |
paulb@200 | 35 | b2.vy += dy * b_mass_x_mag |
paulb@200 | 36 | b2.vz += dz * b_mass_x_mag |
paulb@200 | 37 | |
paulb@200 | 38 | for b in bodies : |
paulb@200 | 39 | b.x += dt * b.vx |
paulb@200 | 40 | b.y += dt * b.vy |
paulb@200 | 41 | b.z += dt * b.vz |
paulb@200 | 42 | |
paulb@200 | 43 | def energy(bodies) : |
paulb@200 | 44 | e = 0.0 |
paulb@200 | 45 | for i in xrange(len(bodies)) : |
paulb@200 | 46 | b = bodies[i] |
paulb@200 | 47 | e += 0.5 * b.mass * (b.vx**2 + b.vy**2 + b.vz**2) |
paulb@200 | 48 | |
paulb@200 | 49 | for j in xrange(i + 1, len(bodies)) : |
paulb@200 | 50 | b2 = bodies[j] |
paulb@200 | 51 | |
paulb@200 | 52 | dx = b.x - b2.x |
paulb@200 | 53 | dy = b.y - b2.y |
paulb@200 | 54 | dz = b.z - b2.z |
paulb@200 | 55 | distance = (dx**2 + dy**2 + dz**2)**0.5 |
paulb@200 | 56 | |
paulb@200 | 57 | e -= (b.mass * b2.mass) / distance |
paulb@200 | 58 | |
paulb@200 | 59 | return e |
paulb@200 | 60 | |
paulb@200 | 61 | def offset_momentum(bodies) : |
paulb@200 | 62 | global sun |
paulb@200 | 63 | px = py = pz = 0.0 |
paulb@200 | 64 | |
paulb@200 | 65 | for b in bodies : |
paulb@200 | 66 | px += b.vx * b.mass |
paulb@200 | 67 | py += b.vy * b.mass |
paulb@200 | 68 | pz += b.vz * b.mass |
paulb@200 | 69 | |
paulb@200 | 70 | sun.vx = - px / solar_mass |
paulb@200 | 71 | sun.vy = - py / solar_mass |
paulb@200 | 72 | sun.vz = - pz / solar_mass |
paulb@200 | 73 | |
paulb@200 | 74 | sun = body() |
paulb@200 | 75 | sun.x = sun.y = sun.z = sun.vx = sun.vy = sun.vz = 0.0 |
paulb@200 | 76 | sun.mass = solar_mass |
paulb@200 | 77 | |
paulb@200 | 78 | jupiter = body() |
paulb@200 | 79 | jupiter.x = 4.84143144246472090e+00 |
paulb@200 | 80 | jupiter.y = -1.16032004402742839e+00 |
paulb@200 | 81 | jupiter.z = -1.03622044471123109e-01 |
paulb@200 | 82 | jupiter.vx = 1.66007664274403694e-03 * days_per_year |
paulb@200 | 83 | jupiter.vy = 7.69901118419740425e-03 * days_per_year |
paulb@200 | 84 | jupiter.vz = -6.90460016972063023e-05 * days_per_year |
paulb@200 | 85 | jupiter.mass = 9.54791938424326609e-04 * solar_mass |
paulb@200 | 86 | |
paulb@200 | 87 | saturn = body() |
paulb@200 | 88 | saturn.x = 8.34336671824457987e+00 |
paulb@200 | 89 | saturn.y = 4.12479856412430479e+00 |
paulb@200 | 90 | saturn.z = -4.03523417114321381e-01 |
paulb@200 | 91 | saturn.vx = -2.76742510726862411e-03 * days_per_year |
paulb@200 | 92 | saturn.vy = 4.99852801234917238e-03 * days_per_year |
paulb@200 | 93 | saturn.vz = 2.30417297573763929e-05 * days_per_year |
paulb@200 | 94 | saturn.mass = 2.85885980666130812e-04 * solar_mass |
paulb@200 | 95 | |
paulb@200 | 96 | uranus = body() |
paulb@200 | 97 | uranus.x = 1.28943695621391310e+01 |
paulb@200 | 98 | uranus.y = -1.51111514016986312e+01 |
paulb@200 | 99 | uranus.z = -2.23307578892655734e-01 |
paulb@200 | 100 | uranus.vx = 2.96460137564761618e-03 * days_per_year |
paulb@200 | 101 | uranus.vy = 2.37847173959480950e-03 * days_per_year |
paulb@200 | 102 | uranus.vz = -2.96589568540237556e-05 * days_per_year |
paulb@200 | 103 | uranus.mass = 4.36624404335156298e-05 * solar_mass |
paulb@200 | 104 | |
paulb@200 | 105 | neptune = body() |
paulb@200 | 106 | neptune.x = 1.53796971148509165e+01 |
paulb@200 | 107 | neptune.y = -2.59193146099879641e+01 |
paulb@200 | 108 | neptune.z = 1.79258772950371181e-01 |
paulb@200 | 109 | neptune.vx = 2.68067772490389322e-03 * days_per_year |
paulb@200 | 110 | neptune.vy = 1.62824170038242295e-03 * days_per_year |
paulb@200 | 111 | neptune.vz = -9.51592254519715870e-05 * days_per_year |
paulb@200 | 112 | neptune.mass = 5.15138902046611451e-05 * solar_mass |
paulb@200 | 113 | |
paulb@200 | 114 | def main() : |
paulb@200 | 115 | try : |
paulb@200 | 116 | n = int(sys.argv[1]) |
paulb@200 | 117 | except : |
paulb@200 | 118 | print "Usage: %s <N>" % sys.argv[0] |
paulb@200 | 119 | |
paulb@200 | 120 | bodies = [sun, jupiter, saturn, uranus, neptune] |
paulb@200 | 121 | |
paulb@200 | 122 | offset_momentum(bodies) |
paulb@200 | 123 | |
paulb@200 | 124 | print "%.9f" % energy(bodies) |
paulb@200 | 125 | |
paulb@200 | 126 | for i in xrange(n) : |
paulb@200 | 127 | advance(bodies, 0.01) |
paulb@200 | 128 | |
paulb@200 | 129 | print "%.9f" % energy(bodies) |
paulb@200 | 130 | |
paulb@200 | 131 | main() |