#| bh-tree-test.scm: Tests for the bh-tree.scm module.
Copyright (C) 2007 Will M. Farr

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 bh-tree-test mzscheme
  (require "bh-tree.scm"
           (lib "math.ss")
           (planet "test.ss" ("schematics" "schemeunit.plt"))
           (planet "nbody-ics.scm" ("wmfarr" "nbody-ics.plt"))
           (lib "42.ss" "srfi"))
  
  (provide bh-tree-test-suite)
  
  (define ((close? eps) a b)
    (< (abs (- a b)) (abs eps)))
  
  (define-simple-check (check-close? eps a b)
    ((close? eps) a b))
  
  (define (distance-squared v w)
    (sum-ec (:parallel (:vector vx v)
                       (:vector wx w))
            (sqr (- vx wx))))
  
  (define (distance v w) (sqrt (distance-squared v w)))
  
  (define (body-tree-energy b t)
    (let ((mb (body-m b))
          (qb (body-q b))
          (mt (tree-m t))
          (qt (tree-q t)))
      (cond
        ((= mt 0.0) 0.0)
        ((equal? qt qb) 0.0) ;No self-energy!
        (else (let ((r (distance qt qb)))
                (- (/ (* mt mb) r)))))))
  
  (define (tree-potential-energy bs beta2)
    (let ((t (bodies->tree bs)))
      (/ (sum-ec (:vector b bs)
                 (let ((qb (body-q b)))
                   (let ((cut? (lambda (t)
                                 (let ((d2 (distance-squared (tree-q t) qb))
                                       (s2 (tree-size-squared t)))
                                   (< (/ s2 d2) beta2)))))
                     (tree-fold/sc cut?
                                   (lambda (t E)
                                     (let ((Ebt (body-tree-energy b t)))
                                       (+ E Ebt)))
                                   0
                                   t))))
         2)))
  
  (define bh-tree-test-suite
    (test-suite
     "bh-tree.scm test suite"
     (test-case
      "Total mass and com correct for large tree."
      (let ((t (bodies->tree (make-cold-spherical 1000))))
        (check-close? 1e-6 1.0 (tree-m t))
        (check-close? 1e-6 0.0 (vector-ref (tree-q t) 0))
        (check-close? 1e-6 0.0 (vector-ref (tree-q t) 1))
        (check-close? 1e-6 0.0 (vector-ref (tree-q t) 2))))
     (test-case
      "Number of bodies (computed from tree-fold ...) is OK."
      (let ((n 1000))
        (let ((t (bodies->tree (make-cold-spherical n))))
          (let ((n-fold (tree-fold (lambda (b n) (+ n 1)) 0 t)))
            (check-equal? n n-fold)))))
     (test-case
      "Compute potential energy of system"
      (let ((bs (make-cold-spherical 250)))
        (check-close? 0.1 -0.25 (tree-potential-energy bs (sqr 0.5))))))))