Last active
December 6, 2018 19:50
-
-
Save robot-dreams/fb253e18fd16e5c152918679a3eaf909 to your computer and use it in GitHub Desktop.
Fast approximation of n-body simulation via Barnes-Hut algorithm, O(n log n)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package main | |
import ( | |
"math" | |
) | |
// We represent a region of space, containing some number of particles, as an | |
// OctTree. Leaves of the tree contain at most 1 particle (but are also allowed | |
// to represent empty regions). | |
// | |
// For the time being, we assume all particles have unit mass, but it is a | |
// straightforward extension to allow particles with different mass. | |
type octNode struct { | |
// Sum of all particle positions | |
sx, sy, sz float32 | |
// Number of particles in the region | |
n int | |
// Side length of region | |
l float32 | |
// Minimum coordinates of the region | |
x, y, z float32 | |
// If the region contains more than 1 particle, then it's recursively | |
// subdivided into octants. | |
children []*octNode | |
} | |
// theta is a threshold that determines when a region can be approximated as a | |
// single particle located at the region's center of mass. | |
const theta = 0.5 | |
// Approximates gravitational force on a unit-mass particle with given position. | |
func (o *octNode) approxF(px, py, pz float32) (float32, float32, float32) { | |
// An empty region exerts no gravitational force. | |
if o.n == 0 { | |
return 0, 0, 0 | |
} | |
// Region's center of mass | |
cx, cy, cz := o.sx/float32(o.n), o.sy/float32(o.n), o.sz/float32(o.n) | |
// Displacement from particle to center of mass | |
dx, dy, dz := cx-px, cy-py, cz-pz | |
r := float32(math.Sqrt(float64(dx*dx + dy*dy + dz*dz))) | |
// A particle doesn't exert force on itself (we use 1e-9 to account for | |
// floating point precision issues) | |
if r < 1e-9 { | |
return 0, 0, 0 | |
} | |
// We can consider the region as a single particle located at the center | |
// of mass in these cases: | |
// | |
// - The region only has a single particle | |
// | |
// - The displacement from the particle to the center of mass is | |
// sufficiently larger than the side length of the region | |
// | |
// Otherwise, we add up the individual contributions of each sub-region. | |
if o.n == 1 || o.l/r < theta { | |
// Since we're assuming that all particles have unit mass, the total | |
// mass of the region is just the number of particles. | |
m := float32(o.n) | |
r2 := r * r | |
return m * dx / r2, m * dy / r2, m * dz / r2 | |
} else { | |
var fx, fy, fz float32 | |
for _, child := range o.children { | |
dfx, dfy, dfz := child.approxF(px, py, pz) | |
fx += dfx | |
fy += dfy | |
fz += dfz | |
} | |
return fx, fy, fz | |
} | |
} | |
func (o *octNode) inRegion(px, py, pz float32) bool { | |
return px >= o.x && py >= o.y && pz >= o.z && | |
px < o.x+o.l && py < o.y+o.l && pz < o.z+o.l | |
} | |
// Adds a particle to the OctTree rooted at the receiver node; particles outside | |
// the OctTree's region will be ignored. | |
func (o *octNode) addParticle(px, py, pz float32) { | |
if !o.inRegion(px, py, pz) { | |
return | |
} | |
// Since leaf nodes can have at most 1 particle, adding a second particle | |
// requires us to subdivide the OctTree's region. | |
if o.n == 1 { | |
o.subdivide() | |
} | |
o.sx += px | |
o.sy += py | |
o.sz += pz | |
o.n++ | |
for _, child := range o.children { | |
child.addParticle(px, py, pz) | |
} | |
} | |
// Precondition: The region represented by this OctTree contains exactly one | |
// particle. | |
func (o *octNode) subdivide() { | |
for i := 0; i < 8; i++ { | |
child := &octNode{ | |
l: o.l / 2, | |
x: o.x, | |
y: o.y, | |
z: o.z, | |
} | |
// We set the minimum coordinates for the i-th region based on the | |
// binary representation of i. | |
if i&1 != 0 { | |
child.x += child.l | |
} | |
if i&2 != 0 { | |
child.y += child.l | |
} | |
if i&4 != 0 { | |
child.z += child.l | |
} | |
o.children = append(o.children, child) | |
} | |
for _, child := range o.children { | |
// Since there's only one particle, the "sum of all particle positions" | |
// is just the position of that particle. | |
child.addParticle(o.sx, o.sy, o.sz) | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Aha! You are intimately familiar with this algorithm.
Some other examples I found:
https://benedikt-bitterli.me/vorton-fluid.html
http://on-demand.gputechconf.com/gtc/2012/presentations/S0111-Efficient-CUDA-Implementation-of-Tree-Based-N-Body-Algorithm.pdf