Skip to content

Instantly share code, notes, and snippets.

@robot-dreams
Last active December 6, 2018 19:50
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save robot-dreams/fb253e18fd16e5c152918679a3eaf909 to your computer and use it in GitHub Desktop.
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)
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)
}
}
@magnus-haw
Copy link

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment