Hastebin

#lang typed/racket/base

(require racket/fixnum racket/flonum racket/math)

(provide take-energy)

(define SOLAR-MASS (* 4.0 pi pi))
(define DAYS-PER-YEAR 365.24)

(: dpe : Flonum -> Flonum)
(define (dpe v) (* v DAYS-PER-YEAR))
(: slr : Flonum -> Flonum)
(define (slr m) (* m SOLAR-MASS))

(define-syntax subn! (syntax-rules () ((_ what withs) (set! what (- what withs)))))
;(define-syntax addn! (syntax-rules () ((_ what withs) (set! what (+ what withs)))))

(: _mass : FlVector Fixnum -> Flonum)
(define (_mass v base) (flvector-ref v base))

(: _x : FlVector Fixnum -> Flonum)
(define (_x v base) (flvector-ref v (fx+ 1 base)))

(: _y : FlVector Fixnum -> Flonum)
(define (_y v base) (flvector-ref v (fx+ 2 base)))

(: _z : FlVector Fixnum -> Flonum)
(define (_z v base) (flvector-ref v (fx+ 3 base)))

(: _vx : FlVector Fixnum -> Flonum)
(define (_vx v base) (flvector-ref v (fx+ 4 base)))

(: _vy : FlVector Fixnum -> Flonum)
(define (_vy v base) (flvector-ref v (fx+ 5 base)))

(: _vz : FlVector Fixnum -> Flonum)
(define (_vz v base) (flvector-ref v (fx+ 6 base)))
;; --- setters ----

(: _x! : FlVector Fixnum Flonum -> Void)
(define (_x! v base val) (flvector-set! v (fx+ 1 base) val))

(: _y! : FlVector Fixnum Flonum -> Void)
(define (_y! v base val) (flvector-set! v (fx+ 2 base) val))

(: _z! : FlVector Fixnum Flonum -> Void)
(define (_z! v base val) (flvector-set! v (fx+ 3 base) val))

(: _vx! : FlVector Fixnum Flonum -> Void)
(define (_vx! v base val) (flvector-set! v (fx+ 4 base) val))

(: _vy! : FlVector Fixnum Flonum -> Void)
(define (_vy! v base val) (flvector-set! v (fx+ 5 base) val))

(: _vz! : FlVector Fixnum Flonum -> Void)
(define (_vz! v base val) (flvector-set! v (fx+ 6 base) val))

(: sum^2 : Flonum Flonum Flonum -> Nonnegative-Flonum)
(define (sum^2 x y z) (+ (sqr x) (sqr y) (sqr z)))

(: take-energy : Integer -> Flonum)
(define (take-energy n)
  (define bodies
    (flvector
     SOLAR-MASS 0.0 0.0 0.0 0.0 0.0 0.0  ;; SUN
     (slr 9.54791938424326609e-04) 4.84143144246472090e+00 -1.16032004402742839e+00 -1.03622044471123109e-01 (dpe 1.66007664274403694e-03) (dpe 7.69901118419740425e-03) (dpe -6.90460016972063023e-05)
     (slr 2.85885980666130812e-04) 8.34336671824457987e+00 4.12479856412430479e+00 -4.03523417114321381e-01 (dpe -2.76742510726862411e-03) (dpe 4.99852801234917238e-03) (dpe 2.30417297573763929e-05)
     (slr 4.36624404335156298e-05) 1.28943695621391310e+01 -1.51111514016986312e+01 -2.23307578892655734e-01 (dpe 2.96460137564761618e-03) (dpe 2.37847173959480950e-03) (dpe -2.96589568540237556e-05)
     (slr 5.15138902046611451e-05) 1.53796971148509165e+01 -2.59193146099879641e+01 1.79258772950371181e-01 (dpe 2.68067772490389322e-03) (dpe 1.62824170038242295e-03) (dpe -9.51592254519715870e-05)
     ))
  (define dt 0.01)
  (define bodies-len (flvector-length bodies))
  (define body-size 7)
  (for ([dummy (in-range n)])
    (do ([i : Fixnum 0 (fx+ i body-size)])
      ((= i bodies-len))
      (do ([j : Fixnum (fx+ i body-size) (fx+ j body-size)])
        ((= j bodies-len))
        (let* ([dx (- (_x bodies i) (_x bodies j))]
               [dy (- (_y bodies i) (_y bodies j))]
               [dz (- (_z bodies i) (_z bodies j))]
               [dsq (sum^2 dx dy dz)]
               [mag (/ dt (* dsq (sqrt dsq)))])
          (_vx! bodies i (- (_vx bodies i) (* dx (_mass bodies j) mag)))
          (_vy! bodies i (- (_vy bodies i) (* dy (_mass bodies j) mag)))
          (_vz! bodies i (- (_vz bodies i) (* dz (_mass bodies j) mag)))
          (_vx! bodies j (+ (_vx bodies j) (* dx (_mass bodies i) mag)))
          (_vy! bodies j (+ (_vy bodies j) (* dy (_mass bodies i) mag)))
          (_vz! bodies j (+ (_vz bodies j) (* dz (_mass bodies i) mag))))
        ))
    (do ([i : Fixnum 0 (fx+ i body-size)])
      ((= i bodies-len))
      (_x! bodies i (+ (_x bodies i) (* dt (_vx bodies i))))
      (_y! bodies i (+ (_y bodies i) (* dt (_vy bodies i))))
      (_z! bodies i (+ (_z bodies i) (* dt (_vz bodies i))))))
  (do ((i : Fixnum 0 (fx+ body-size i))
       (E 0.0 (+ E (* 0.5 (_mass bodies i) (sum^2 (_vx bodies i) (_vy bodies i) (_vz bodies i))))))
    ((= i bodies-len) E)
    (do ((j : Fixnum (fx+ i body-size) (fx+ body-size j)))
      ((= j bodies-len))
      (let ([dx (- (_x bodies i) (_x bodies j))]
            [dy (- (_y bodies i) (_y bodies j))]
            [dz (- (_z bodies i) (_z bodies j))])
        (subn! E (/ (* (_mass bodies i) (_mass bodies j))
                    (flsqrt (sum^2 dx dy dz)))))))
  )

(module* main #f
  (require racket/cmdline racket/format)
  (: ->sec : Flonum Flonum -> Flonum)
  (define (->sec from to) (* 1e-3 (- from to)))
  (: test : Integer -> Void)
  (define (test n)
    (let* ((st (current-inexact-milliseconds))
           (e (take-energy n))
           (et (current-inexact-milliseconds)))
      (printf "~a  per ~a sec ~%" (~r e #:precision 9)
              (~r (->sec et st) #:precision 3))))
  (command-line
   #:program "nbody"
   #:args args
   (cond [(pair? args) (test (assert (string->number (assert (car args) string?)) exact-integer?))]
         [else (take-energy 100000)]))
  )