extra/benchmark/nbody-simd/nbody-simd.factor

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