extra/benchmark/nbody/nbody.factor

   1 ! Copyright (C) 2008, 2009 Slava Pestov.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: accessors specialized-arrays fry kernel
   4 locals math math.constants math.functions math.vectors
   5 prettyprint combinators.smart sequences hints arrays ;
   6 FROM: alien.c-types => double ;
   7 SPECIALIZED-ARRAY: double
   8 IN: benchmark.nbody
   9
  10 : solar-mass ( -- x ) 4 pi sq * ; inline
  11 CONSTANT: days-per-year 365.24
  12
  13 TUPLE: body
  14 { location double-array }
  15 { velocity double-array }
  16 { mass float read-only } ;
  17
  18 : <body> ( location velocity mass -- body )
  19     [ days-per-year v*n ] [ solar-mass * ] bi* body boa ; inline
  20
  21 : <jupiter> ( -- body )
  22     double-array{ 4.84143144246472090e00 -1.16032004402742839e00 -1.03622044471123109e-01 }
  23     double-array{ 1.66007664274403694e-03 7.69901118419740425e-03 -6.90460016972063023e-05 }
  24     9.54791938424326609e-04
  25     <body> ;
  26
  27 : <saturn> ( -- body )
  28     double-array{ 8.34336671824457987e00 4.12479856412430479e00 -4.03523417114321381e-01 }
  29     double-array{ -2.76742510726862411e-03 4.99852801234917238e-03 2.30417297573763929e-05 }
  30     2.85885980666130812e-04
  31     <body> ;
  32
  33 : <uranus> ( -- body )
  34     double-array{ 1.28943695621391310e01 -1.51111514016986312e01 -2.23307578892655734e-01 }
  35     double-array{ 2.96460137564761618e-03 2.37847173959480950e-03 -2.96589568540237556e-05 }
  36     4.36624404335156298e-05
  37     <body> ;
  38
  39 : <neptune> ( -- body )
  40     double-array{ 1.53796971148509165e01 -2.59193146099879641e01 1.79258772950371181e-01 }
  41     double-array{ 2.68067772490389322e-03 1.62824170038242295e-03 -9.51592254519715870e-05 }
  42     5.15138902046611451e-05
  43     <body> ;
  44
  45 : <sun> ( -- body )
  46     double-array{ 0 0 0 } double-array{ 0 0 0 } 1 <body> ;
  47
  48 : offset-momentum ( body offset -- body )
  49     vneg solar-mass v/n >>velocity ; inline
  50
  51 TUPLE: nbody-system { bodies array read-only } ;
  52
  53 : init-bodies ( bodies -- )
  54     [ first ] [ double-array{ 0 0 0 } [ [ velocity>> ] [ mass>> ] bi v*n v+ ] reduce ] bi
  55     offset-momentum drop ; inline
  56
  57 : <nbody-system> ( -- system )
  58     [ <sun> <jupiter> <saturn> <uranus> <neptune> ] output>array nbody-system boa
  59     dup bodies>> init-bodies ; inline
  60
  61 :: each-pair ( ... bodies pair-quot: ( ... other-body body -- ... ) each-quot: ( ... body -- ... ) -- ... )
  62     bodies [| body i |
  63         body each-quot call
  64         bodies i 1 + tail-slice [
  65             body pair-quot call
  66         ] each
  67     ] each-index ; inline
  68
  69 : update-position ( body dt -- )
  70     [ dup velocity>> ] dip '[ _ _ v*n v+ ] change-location drop ;
  71
  72 : mag ( dt body other-body -- mag d )
  73     [ location>> ] bi@ v- [ norm-sq dup sqrt * / ] keep ; inline
  74
  75 :: update-velocity ( other-body body dt -- )
  76     dt body other-body mag
  77     [ [ body ] 2dip '[ other-body mass>> _ * _ n*v v- ] change-velocity drop ]
  78     [ [ other-body ] 2dip '[ body mass>> _ * _ n*v v+ ] change-velocity drop ] 2bi ;
  79
  80 : advance ( system dt -- )
  81     [ bodies>> ] dip
  82     [ '[ _ update-velocity ] [ drop ] each-pair ]
  83     [ '[ _ update-position ] each ]
  84     2bi ; inline
  85
  86 : inertia ( body -- e )
  87     [ mass>> ] [ velocity>> norm-sq ] bi * 0.5 * ;
  88
  89 : newton's-law ( other-body body -- e )
  90     [ [ mass>> ] bi@ * ] [ [ location>> ] bi@ distance ] 2bi / ;
  91
  92 : energy ( system -- x )
  93     [ 0.0 ] dip bodies>> [ newton's-law - ] [ inertia + ] each-pair ; inline
  94
  95 : nbody ( n -- )
  96     <nbody-system>
  97     [ energy . ] [ '[ _ 0.01 advance ] times ] [ energy . ] tri ;
  98
  99 HINTS: update-position body float ;
 100 HINTS: update-velocity body body float ;
 101 HINTS: inertia body ;
 102 HINTS: newton's-law body body ;
 103 HINTS: nbody fixnum ;
 104
 105 : nbody-benchmark ( -- ) 1000000 nbody ;
 106
 107 MAIN: nbody-benchmark