src/test/bench/shootout-nbody.rs

   1 // Copyright 2014 The Rust Project Developers. See the COPYRIGHT
   2 // file at the top-level directory of this distribution and at
   3 // http://rust-lang.org/COPYRIGHT.
   4 //
   5 // Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
   6 // http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
   7 // <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
   8 // option. This file may not be copied, modified, or distributed
   9 // except according to those terms.
  10
  11 use std::os;
  12
  13 static PI: f64 = 3.141592653589793;
  14 static SOLAR_MASS: f64 = 4.0 * PI * PI;
  15 static YEAR: f64 = 365.24;
  16 static N_BODIES: uint = 5;
  17
  18 static BODIES: [Planet, ..N_BODIES] = [
  19     // Sun
  20     Planet {
  21         x: [ 0.0, 0.0, 0.0 ],
  22         v: [ 0.0, 0.0, 0.0 ],
  23         mass: SOLAR_MASS,
  24     },
  25     // Jupiter
  26     Planet {
  27         x: [
  28             4.84143144246472090e+00,
  29             -1.16032004402742839e+00,
  30             -1.03622044471123109e-01,
  31         ],
  32         v: [
  33             1.66007664274403694e-03 * YEAR,
  34             7.69901118419740425e-03 * YEAR,
  35             -6.90460016972063023e-05 * YEAR,
  36         ],
  37         mass: 9.54791938424326609e-04 * SOLAR_MASS,
  38     },
  39     // Saturn
  40     Planet {
  41         x: [
  42             8.34336671824457987e+00,
  43             4.12479856412430479e+00,
  44             -4.03523417114321381e-01,
  45         ],
  46         v: [
  47             -2.76742510726862411e-03 * YEAR,
  48             4.99852801234917238e-03 * YEAR,
  49             2.30417297573763929e-05 * YEAR,
  50         ],
  51         mass: 2.85885980666130812e-04 * SOLAR_MASS,
  52     },
  53     // Uranus
  54     Planet {
  55         x: [
  56             1.28943695621391310e+01,
  57             -1.51111514016986312e+01,
  58             -2.23307578892655734e-01,
  59         ],
  60         v: [
  61             2.96460137564761618e-03 * YEAR,
  62             2.37847173959480950e-03 * YEAR,
  63             -2.96589568540237556e-05 * YEAR,
  64         ],
  65         mass: 4.36624404335156298e-05 * SOLAR_MASS,
  66     },
  67     // Neptune
  68     Planet {
  69         x: [
  70             1.53796971148509165e+01,
  71             -2.59193146099879641e+01,
  72             1.79258772950371181e-01,
  73         ],
  74         v: [
  75             2.68067772490389322e-03 * YEAR,
  76             1.62824170038242295e-03 * YEAR,
  77             -9.51592254519715870e-05 * YEAR,
  78         ],
  79         mass: 5.15138902046611451e-05 * SOLAR_MASS,
  80     },
  81 ];
  82
  83 struct Planet {
  84     x: [f64, ..3],
  85     v: [f64, ..3],
  86     mass: f64,
  87 }
  88
  89 fn advance(bodies: &mut [Planet, ..N_BODIES], dt: f64, steps: i32) {
  90     let mut d = [ 0.0, ..3 ];
  91     for _ in range(0, steps) {
  92         for i in range(0u, N_BODIES) {
  93             for j in range(i + 1, N_BODIES) {
  94                 d[0] = bodies[i].x[0] - bodies[j].x[0];
  95                 d[1] = bodies[i].x[1] - bodies[j].x[1];
  96                 d[2] = bodies[i].x[2] - bodies[j].x[2];
  97
  98                 let d2 = d[0]*d[0] + d[1]*d[1] + d[2]*d[2];
  99                 let mag = dt / (d2 * d2.sqrt());
 100
 101                 let a_mass = bodies[i].mass;
 102                 let b_mass = bodies[j].mass;
 103                 bodies[i].v[0] -= d[0] * b_mass * mag;
 104                 bodies[i].v[1] -= d[1] * b_mass * mag;
 105                 bodies[i].v[2] -= d[2] * b_mass * mag;
 106
 107                 bodies[j].v[0] += d[0] * a_mass * mag;
 108                 bodies[j].v[1] += d[1] * a_mass * mag;
 109                 bodies[j].v[2] += d[2] * a_mass * mag;
 110             }
 111         }
 112
 113         for a in bodies.mut_iter() {
 114             a.x[0] += dt * a.v[0];
 115             a.x[1] += dt * a.v[1];
 116             a.x[2] += dt * a.v[2];
 117         }
 118     }
 119 }
 120
 121 fn energy(bodies: &[Planet, ..N_BODIES]) -> f64 {
 122     let mut e = 0.0;
 123     let mut d = [ 0.0, ..3 ];
 124     for i in range(0u, N_BODIES) {
 125         for k in range(0u, 3) {
 126             e += bodies[i].mass * bodies[i].v[k] * bodies[i].v[k] / 2.0;
 127         }
 128
 129         for j in range(i + 1, N_BODIES) {
 130             for k in range(0u, 3) {
 131                 d[k] = bodies[i].x[k] - bodies[j].x[k];
 132             }
 133             let dist = (d[0]*d[0] + d[1]*d[1] + d[2]*d[2]).sqrt();
 134             e -= bodies[i].mass * bodies[j].mass / dist;
 135         }
 136     }
 137     e
 138 }
 139
 140 fn offset_momentum(bodies: &mut [Planet, ..N_BODIES]) {
 141     for i in range(0u, N_BODIES) {
 142         for k in range(0u, 3) {
 143             bodies[0].v[k] -= bodies[i].v[k] * bodies[i].mass / SOLAR_MASS;
 144         }
 145     }
 146 }
 147
 148 fn main() {
 149     let args = os::args();
 150     let args = if os::getenv("RUST_BENCH").is_some() {
 151         vec!(~"", ~"1000")
 152     } else if args.len() <= 1u {
 153         vec!(~"", ~"1000")
 154     } else {
 155         args.move_iter().collect()
 156     };
 157
 158     let n: i32 = from_str::<i32>(*args.get(1)).unwrap();
 159     let mut bodies = BODIES;
 160
 161     offset_momentum(&mut bodies);
 162     println!("{:.9f}", energy(&bodies) as f64);
 163
 164     advance(&mut bodies, 0.01, n);
 165
 166     println!("{:.9f}", energy(&bodies) as f64);
 167 }
 168