Alioth shootout n-body benchmark Clojure implementation

nbody.clj

(ns nbody
  (:gen-class))

;;
;; Convenience aliases for unchecked arithmetic operations

(defmacro uc+ [& args] `(unchecked-add ~@args))
(defmacro uc- [& args] `(unchecked-subtract ~@args))
(defmacro uc* [& args] `(unchecked-multiply ~@args))

;;
;; Mutable three-dimensional coordinate type

(definterface ICoord
  (^nbody.ICoord clone [])
  (^double x [])
  (^double y [])
  (^double z [])
  (^double coord_add_mul_BANG_ [^nbody.ICoord c2 ^double m])
  (^double coord_mag_sq []))

(deftype Coord [^:unsynchronized-mutable ^double _x
                ^:unsynchronized-mutable ^double _y
                ^:unsynchronized-mutable ^double _z]
  ICoord
  (clone [this] (Coord. _x _y _z))
  (x [this] _x)
  (y [this] _y)
  (z [this] _z)
  (coord-add-mul! [this c2 m]
    (set! _x (uc+ _x (uc* (.x c2) m)))
    (set! _y (uc+ _y (uc* (.y c2) m)))
    (set! _z (uc+ _z (uc* (.z c2) m))))
  (coord-mag-sq [this]
    (uc+ (uc* _x _x) (uc+ (uc* _y _y) (uc* _z _z)))))

(defmacro defcoordmeth [name]
  `(defmacro ~name [this# & args#]
     (list* '. (with-meta this# {:tag '~'Coord}) '~name args#)))

(defcoordmeth coord-add-mul!)
(defcoordmeth coord-mag-sq)

(defn coord-sub [^Coord c1 ^Coord c2]
  (Coord. (uc- (.x c1) (.x c2)) (uc- (.y c1) (.y c2)) (uc- (.z c1) (.z c2))))

(defn coord-div [^Coord c1 ^double m]
  (Coord. (/ (.x c1) m) (/ (.y c1) m) (/ (.z c1) m)))

(defn coord-add-mul [^Coord c1 ^Coord c2 ^double m]
  (doto (.clone c1) (.coord-add-mul! c2 m)))

;;
;; Physics simulation bodies

(defrecord Body [pos vel ^double mass])

(defn initial-bodies [& specs]
  (let [solar-mass (* 4 Math/PI Math/PI), days-year 365.24
        body (fn [x y z vx vy vz mass]
               (let [[vx vy vz] (map #(* days-year %) [vx vy vz])]
                 (Body. (Coord. x y z) (Coord. vx vy vz) (* mass solar-mass))))
        [sun & bs] (map #(apply body %) specs)
        step (fn [p body] (coord-add-mul p (:vel body) (:mass body)))
        offset (fn [sun p] (assoc sun :vel (coord-div p (- solar-mass))))
        sun (offset sun (reduce step (Coord. 0 0 0) bs))]
    (into-array Body (cons sun bs))))

(def ^"[Lnbody.Body;" bodies
  (initial-bodies
   ;; Sun
   [0.0 0.0 0.0 0.0 0.0 0.0 1.0]

   ;; Jupiter
   [4.84143144246472090e+00 -1.16032004402742839e+00 -1.03622044471123109e-01
    1.66007664274403694e-03  7.69901118419740425e-03 -6.90460016972063023e-05
    9.54791938424326609e-04]

   ;; Saturn
   [8.34336671824457987e+00  4.12479856412430479e+00 -4.03523417114321381e-01
   -2.76742510726862411e-03  4.99852801234917238e-03  2.30417297573763929e-05
    2.85885980666130812e-04]

   ;; Uranus
   [1.28943695621391310e+01 -1.51111514016986312e+01 -2.23307578892655734e-01
    2.96460137564761618e-03  2.37847173959480950e-03 -2.96589568540237556e-05
    4.36624404335156298e-05]

   ;; Neptune
   [1.53796971148509165e+01 -2.59193146099879641e+01  1.79258772950371181e-01
    2.68067772490389322e-03  1.62824170038242295e-03 -9.51592254519715870e-05
    5.15138902046611451e-05]))

(def ^:const ^long n (alength bodies))
(def ^:const ^long n1 (dec n))

;;
;; Physics simulation operations

(defn advance [^double dt]
  (loop [i 0]
    (let [^Body ibody (aget bodies i)
          ipos (.pos ibody)
          ivel (.vel ibody)
          imass (.mass ibody)]
      (loop [j (unchecked-inc i)]
        (let [^Body jbody (aget bodies j)
              delta (coord-sub ipos (.pos jbody))
              dsq (coord-mag-sq delta)
              mag (/ dt (uc* dsq (Math/sqrt dsq)))                ]
          (coord-add-mul! ivel delta (unchecked-negate (uc* mag (.mass jbody))))
          (coord-add-mul! (.vel jbody) delta (uc* mag imass))
          (let [j (unchecked-inc j)]
            (when (> n j)
              (recur j)))))
      (coord-add-mul! ipos ivel dt)
      (let [i (unchecked-inc i)]
        (if (> n1 i)
          (recur i)
          (let [^Body ibody (aget bodies i)]
            (coord-add-mul! (.pos ibody) (.vel ibody) dt)))))))

(defn energy []
  (->> bodies (iterate rest) (take-while seq)
       (map (fn [[{:keys [pos vel mass]} & bs]]
              (reduce (fn [e b]
                        (- e (/ (* mass (:mass b))
                                (Math/sqrt (coord-mag-sq
                                            (coord-sub pos (:pos b)))))))
                      (* 0.5 mass (coord-mag-sq vel))
                      bs)))
       (reduce +)))

;;
;; Main entry point

(defn -main [& args]
  (printf "%.9f\n" (energy))
  (dotimes [_ (Long/parseLong (first args))]
    (advance 0.01))
  (printf "%.9f\n" (energy)))