#|
nbody-ics.scm: Creates initial conditions for N-body simulations.
Copyright (C) 2007 Will M. Farr <farr@mit.edu>

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License along
with this program; if not, write to the Free Software Foundation, Inc.,
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
|#

(module nbody-ics mzscheme
  (require (lib "kw.ss")
           (lib "42.ss" "srfi")
           (lib "math.ss")
           (only (lib "43.ss" "srfi") vector-copy)
           (lib "contract.ss"))
  
  (define-struct body
    (m q p) #f)
  
  (define nbody-system/c (vectorof body?))
  (define 3vector/c (flat-named-contract "3vector/c"
                                         (lambda (v)
                                           (and (vector? v)
                                                (= (vector-length v) 3)))))
  
  (provide/contract
   (body-copy (-> body? body?))
   (copy-bodies (-> nbody-system/c nbody-system/c))
   (total-mass (-> nbody-system/c (>/c 0.0)))
   (center-of-mass (-> nbody-system/c 3vector/c))
   (total-momentum (-> nbody-system/c 3vector/c))
   (adjust-frame! (-> nbody-system/c nbody-system/c))
   (make-cold-spherical (-> natural-number/c nbody-system/c))
   (make-hot-spherical (case-> (-> natural-number/c nbody-system/c)
                               (-> natural-number/c (between/c 0.0 1.0) nbody-system/c)))
   (make-plummer (-> natural-number/c nbody-system/c)))
  
  (provide (struct body (m q p))
           nbody-system/c 3vector/c)
  
  
  (define (body-copy b)
    (make-body (body-m b) (vector-copy (body-q b)) (vector-copy (body-p b))))
  
  (define (copy-bodies bs)
    (vector-of-length-ec (vector-length bs)
      (:vector b bs)
      (body-copy b)))
  
  (define (random-between a b)
    (let ((delta (- b a)))
      (+ a (* delta (random)))))
  
  (define (square x) (* x x))
  
  ;; random-3vector generates a 3-vector with random orientation and the
  ;; given length.  We do this by choosing a longitudinal angle
  ;; uniformly in [0,2pi), and a cosine of the latitudinal angle in [-1,
  ;; 1).  Since the latitudinal angle is in [0,pi), we can compute the
  ;; sine of that angle using sin(theta) = sqrt(1-cos^2(theta)).
  (define (random-3vector r)
    (let ((phi (random-between 0.0 (* 2.0 pi)))
          (cos-theta (random-between -1.0 1.0)))
      (let ((sin-theta (sqrt (- 1.0 (square cos-theta))))) ; OK because sin-theta > 0
                                                                ; for theta in [0,pi)
        (vector (* r sin-theta (cos phi))
                (* r sin-theta (sin phi))
                (* r cos-theta)))))
  
  ;; We use the von Neumann rejection algorithm to generate random
  ;; numbers drawn from any distribution. Essentially, we're shooting
  ;; darts into the rectangle with corners (x-min, y-min) and (x-max,
  ;; y-max). This rectangle must enclose completely the distribution
  ;; we're trying to draw numbers from. If the dart hits below the
  ;; value of the distribution (i.e. (if (< y (dist x)))), then we
  ;; return the x-value; if not, we throw another dart. A bit of
  ;; thought will reveal that this returns values along the x-axis with
  ;; probability proportional to the height of the distribution function
  ;; (dist x). In principle, any values for the bounds which enclose
  ;; the distribution will work, but tighter bounds result in fewer
  ;; "misses" and therefore produce a more efficient algorithm.
  (define/kw (random-from-distribution dist #:key (x-min 0.0) (x-max 1.0) (y-min 0.0) (y-max 1.0))
    (let loop ((x (random-between x-min x-max))
               (y (random-between y-min y-max)))
      (if (< y (dist x))
          x
          (loop (random-between x-min x-max)
                (random-between y-min y-max)))))
  
  (define (vector-add! v w)
    (do-ec (:parallel (:vector velt (index i) v)
                      (:vector welt w))
           (vector-set! v i (+ velt welt)))
    v)
  
  (define (vector-sub! v w)
    (do-ec (:parallel (:vector velt (index i) v)
                      (:vector welt w))
           (vector-set! v i (- velt welt)))
    v)
  
  (define (vector-scale! v x)
    (do-ec (:vector velt (index i) v)
           (vector-set! v i (* velt x)))
    v)
  
  (define (vector-scale v x)
    (vector-of-length-ec (vector-length v)
      (:vector velt v)
      (* velt x)))
  
  (define (total-mass bs)
    (sum-ec (:vector b bs)
            (body-m b)))
  
  (define (total-momentum bs)
    (let ((p-tot (make-vector 3 0.0)))
      (do-ec (:vector b bs)
             (vector-add! p-tot (body-p b)))
      p-tot))
  
  (define (center-of-mass bs)
    (let ((M (total-mass bs))
          (mq (make-vector 3 0.0)))
      (do-ec (:vector b bs)
             (vector-add! mq (vector-scale (body-q b) (body-m b))))
      (vector-scale! mq (/ M))))
  
  (define (adjust-frame! bs)
    (let ((M (total-mass bs))
          (Q (center-of-mass bs))
          (P (total-momentum bs)))
      (do-ec (:vector b bs)
             (begin
               (vector-sub! (body-q b) Q)
               (vector-sub! (body-p b) (vector-scale P (/ (body-m b) M)))))
      bs))
  
  (define (make-cold-spherical n)
    (let ((m (/ 1.0 n)))
      (adjust-frame!
       (vector-of-length-ec n (:range i n)
         (let ((r (random-from-distribution square #:x-max 12/5 #:y-max 144/25)))
           (make-body m (random-3vector r) (make-vector 3 0.0)))))))
  
  (define/kw (make-hot-spherical n #:optional (Q 0.5))
    (let ((m (/ 1.0 n)))
      (let ((R (* 12/5 (- 1.0 Q)))
            (V (sqrt (/ (* 5/6 Q) (- 1.0 Q)))))
        (adjust-frame!
         (vector-of-length-ec n (:range i n)
           (let ((r (random-from-distribution square #:x-max R #:y-max (square R)))
                 (v (random-from-distribution square #:x-max V #:y-max (square V))))
             (make-body m (random-3vector r) (random-3vector (* m v)))))))))
  
  ;; The algorithm used here is described at
  ;; http://www.artcompsci.org/kali/vol/plummer/ch02.html ; they got it
  ;; from Sverre Aarseth, Michel Henon and Roland Wielen, _A comparison
  ;; of Numerical Methods for the Study of Star Cluster Dynamics_,
  ;; Astronomy and Astrophysics, 37, 183 (1974). Essentially, we place
  ;; each particle randomly in space according to the Plummer mass
  ;; density, and then we give that particle a velocity drawn from the
  ;; velocity distribution at that point.
  (define (make-plummer n)
    (let ((m (/ 1.0 n))
          (scalefactor (/ 16.0 (* 3.0 pi))))
      (adjust-frame!
       (vector-of-length-ec n (:range i n)
         (let ((r (/ 1.0 (sqrt (- (expt (random) -2/3) 1.0))))
               (x (random-from-distribution (lambda (x) (* (square x) (expt (- 1.0 (square x)) 3.5)))
                                            #:y-max 0.1)))
           (let ((v (* x (sqrt 2) (expt (+ 1.0 (square r)) -0.25))))
             (make-body m (random-3vector (/ r scalefactor))
                        (random-3vector (* m v (sqrt scalefactor)))))))))))