#| bh-tree.scm: Barnes-Hut octtrees. 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 mzscheme (require (planet "nbody-ics.scm" ("wmfarr" "nbody-ics.plt")) (lib "42.ss" "srfi") (lib "math.ss") (lib "contract.ss")) (provide (struct cell (m q bounds sub-trees)) (struct bounds (low high))) (define-struct cell (m q bounds sub-trees) #f) (define-struct bounds (low high) #f) (define empty-tree? null?) (define (tree? obj) (or (empty-tree? obj) (body? obj) (cell? obj))) (provide/contract (empty-tree? (-> tree? boolean?)) (make-empty (-> empty-tree?)) (tree? (-> any/c boolean?)) (tree-m (-> tree? (>=/c 0.0))) (tree-q (-> tree? 3vector/c)) (tree-size-squared (-> tree? (>=/c 0.0))) (in-bounds? (-> bounds? 3vector/c boolean?)) (bodies->tree (-> nbody-system/c tree?)) (tree-fold/sc (-> (-> tree? any) (-> tree? any/c any) any/c tree? any)) (tree-fold (-> (-> tree? any/c any) any/c tree? any))) (define (make-empty) '()) (define (tree-m t) (cond ((empty-tree? t) 0.0) ((body? t) (body-m t)) ((cell? t) (cell-m t)) (else (error 'tree-m "argument not a tree: ~a" t)))) (define (tree-q t) (cond ((empty-tree? t) (make-vector 3 0.0)) ((body? t) (body-q t)) ((cell? t) (cell-q t)) (else (error 'tree-q "argument not a tree: ~a" t)))) (define (in-bounds? bds v) (every?-ec (:parallel (:vector vx v) (:vector lx (bounds-low bds)) (:vector ux (bounds-high bds))) (and (<= lx vx) (< vx ux)))) (define *epsilon* 1e-6) (define expand-bounds (let ((efactor (+ 1.0 *epsilon*))) (lambda (bds) (let ((low (bounds-low bds)) (high (bounds-high bds))) (make-bounds low (vector-of-length-ec (vector-length high) (:parallel (:vector lx low) (:vector ux high)) (let ((delta (- ux lx))) (+ lx (* delta efactor))))))))) (define (bodies->bounds bs) (let ((min (make-vector 3 +inf.0)) (max (make-vector 3 -inf.0))) (do-ec (:vector b bs) (let ((q (body-q b))) (do-ec (:parallel (:vector qx (index i) q) (:vector minx min) (:vector maxx max)) (begin (if (< qx minx) (vector-set! min i qx)) (if (> qx maxx) (vector-set! max i qx)))))) (expand-bounds (make-bounds min max)))) (define (high? i j) (> (bitwise-and (arithmetic-shift 1 j) i) 0)) (define (sub-bounds bds) (let ((low (bounds-low bds)) (high (bounds-high bds))) (let ((mids (vector-of-length-ec 3 (:parallel (:vector lx low) (:vector hx high)) (/ (+ lx hx) 2)))) (vector-of-length-ec 8 (:range i 8) (let ((sub-low (vector-of-length-ec 3 (:parallel (:vector lx (index j) low) (:vector mx mids)) (if (high? i j) mx lx))) (sub-high (vector-of-length-ec 3 (:parallel (:vector mx (index j) mids) (:vector hx high)) (if (high? i j) hx mx)))) (make-bounds sub-low sub-high)))))) (define (split-bodies bs bds) (let ((sub-bds (sub-bounds bds))) (vector-of-length-ec 8 (:vector sb sub-bds) (vector-ec (:vector b bs) (if (in-bounds? sb (body-q b))) b)))) (define (trees-total-mass ts) (sum-ec (:vector t ts) (tree-m t))) (define (trees-center-of-mass ts) (let ((M (trees-total-mass ts)) (com (make-vector 3 0.0))) (if (= M 0.0) com (begin (do-ec (:vector t ts) (let ((m (tree-m t))) (when (> m 0.0) (let ((factor (/ m M)) (q (tree-q t))) (do-ec (:parallel (:vector comx (index i) com) (:vector qx q)) (vector-set! com i (+ comx (* qx factor)))))))) com)))) (define (bodies->tree bs) (cond ((= (vector-length bs) 0) (make-empty)) ((= (vector-length bs) 1) (vector-ref bs 0)) (else (let ((bds (bodies->bounds bs))) (let ((sub-bs (split-bodies bs bds))) (let ((sub-trees (vector-of-length-ec 8 (:vector bs sub-bs) (bodies->tree bs)))) (make-cell (trees-total-mass sub-trees) (trees-center-of-mass sub-trees) bds sub-trees))))))) (define (tree-fold/sc cut? f start t) (cond ((empty-tree? t) start) ((body? t) (f t start)) ((cut? t) (f t start)) (else (let ((sub-ts (cell-sub-trees t)) (tf (lambda (t acc) (tree-fold/sc cut? f acc t)))) (fold-ec start (:vector st sub-ts) st tf))))) (define (tree-fold f start t) (tree-fold/sc (lambda (t) #f) f start t)) (define (tree-size-squared t) (cond ((empty-tree? t) 0.0) ((body? t) 0.0) (else (let ((bds (cell-bounds t))) (let ((low (bounds-low bds)) (high (bounds-high bds))) (sum-ec (:parallel (:vector lx low) (:vector hx high)) (sqr (- hx lx)))))))))