spatial/barneshut: add Barnes-Hut force approximation package and R2/R3 helpers

This commit is contained in:
Dan Kortschak
2019-05-03 06:57:12 +09:30
committed by GitHub
parent 9827ae2933
commit 4a2eb0188c
12 changed files with 1787 additions and 0 deletions

View File

@@ -0,0 +1,249 @@
// Copyright ©2019 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package barneshut
import (
"fmt"
"math"
"gonum.org/v1/gonum/spatial/r2"
)
// Particle2 is a particle in a plane.
type Particle2 interface {
Coord2() r2.Vec
Mass() float64
}
// Force2 is a force modeling function for interactions between p1 and p2,
// m1 is the mass of p1 and m2 of p2. The vector v is the vector from p1 to
// p2. The returned value is the force vector acting on p1.
//
// In models where the identity of particles must be known, p1 and p2 may be
// compared. Force2 may be passed nil for p2 when the Barnes-Hut approximation
// is being used. A nil p2 indicates that the second mass center is an
// aggregate.
type Force2 func(p1, p2 Particle2, m1, m2 float64, v r2.Vec) r2.Vec
// Gravity2 returns a vector force on m1 by m2, equal to (m1⋅m2)/‖v‖²
// in the directions of v. Gravity2 ignores the identity of the interacting
// particles and returns a zero vector when the two particles are
// coincident, but performs no other sanity checks.
func Gravity2(_, _ Particle2, m1, m2 float64, v r2.Vec) r2.Vec {
d2 := v.X*v.X + v.Y*v.Y
if d2 == 0 {
return r2.Vec{}
}
return v.Scale((m1 * m2) / (d2 * math.Sqrt(d2)))
}
// Plane implements Barnes-Hut force approximation calculations.
type Plane struct {
root tile
Particles []Particle2
}
// NewPlane returns a new Plane.
func NewPlane(p []Particle2) *Plane {
q := Plane{Particles: p}
q.Reset()
return &q
}
// Reset reconstructs the Barnes-Hut tree. Reset must be called if the
// Particles field or elements of Particles have been altered, unless
// ForceOn is called with theta=0 or no data structures have been
// previously built.
func (q *Plane) Reset() {
if len(q.Particles) == 0 {
q.root = tile{}
return
}
q.root = tile{
particle: q.Particles[0],
center: q.Particles[0].Coord2(),
mass: q.Particles[0].Mass(),
}
q.root.bounds.Min = q.root.center
q.root.bounds.Max = q.root.center
for _, e := range q.Particles[1:] {
c := e.Coord2()
if c.X < q.root.bounds.Min.X {
q.root.bounds.Min.X = c.X
}
if c.X > q.root.bounds.Max.X {
q.root.bounds.Max.X = c.X
}
if c.Y < q.root.bounds.Min.Y {
q.root.bounds.Min.Y = c.Y
}
if c.Y > q.root.bounds.Max.Y {
q.root.bounds.Max.Y = c.Y
}
}
// TODO(kortschak): Partially parallelise this by
// choosing the direction and using one of four
// goroutines to work on each root quadrant.
for _, e := range q.Particles[1:] {
q.root.insert(e)
}
q.root.summarize()
}
// ForceOn returns a force vector on p given p's mass and the force function, f,
// using the Barnes-Hut theta approximation parameter.
//
// Calls to f will include p in the p1 position and a non-nil p2 if the force
// interaction is with a non-aggregate mass center, otherwise p2 will be nil.
//
// It is safe to call ForceOn concurrently.
func (q *Plane) ForceOn(p Particle2, theta float64, f Force2) (force r2.Vec) {
var empty tile
if theta > 0 && q.root != empty {
return q.root.forceOn(p, p.Coord2(), p.Mass(), theta, f)
}
// For the degenerate case, just iterate over the
// slice of particles rather than walking the tree.
var v r2.Vec
m := p.Mass()
pv := p.Coord2()
for _, e := range q.Particles {
v = v.Add(f(p, e, m, e.Mass(), e.Coord2().Sub(pv)))
}
return v
}
// tile is a quad tree quadrant with Barnes-Hut extensions.
type tile struct {
particle Particle2
bounds r2.Box
nodes [4]*tile
center r2.Vec
mass float64
}
// insert inserts p into the subtree rooted at t.
func (t *tile) insert(p Particle2) {
if t.particle == nil {
for _, q := range t.nodes {
if q != nil {
t.passDown(p)
return
}
}
t.particle = p
t.center = p.Coord2()
t.mass = p.Mass()
return
}
t.passDown(p)
t.passDown(t.particle)
t.particle = nil
t.center = r2.Vec{}
t.mass = 0
}
func (t *tile) passDown(p Particle2) {
dir := quadrantOf(t.bounds, p)
if t.nodes[dir] == nil {
t.nodes[dir] = &tile{bounds: splitPlane(t.bounds, dir)}
}
t.nodes[dir].insert(p)
}
const (
ne = iota
se
sw
nw
)
// quadrantOf returns which quadrant of b that p should be placed in.
func quadrantOf(b r2.Box, p Particle2) int {
center := r2.Vec{
X: (b.Min.X + b.Max.X) / 2,
Y: (b.Min.Y + b.Max.Y) / 2,
}
c := p.Coord2()
if checkBounds && (c.X < b.Min.X || b.Max.X < c.X || c.Y < b.Min.Y || b.Max.Y < c.Y) {
panic(fmt.Sprintf("p out of range %+v: %#v", b, p))
}
if c.X < center.X {
if c.Y < center.Y {
return nw
} else {
return sw
}
} else {
if c.Y < center.Y {
return ne
} else {
return se
}
}
}
// splitPlane returns a quadrant subdivision of b in the given direction.
func splitPlane(b r2.Box, dir int) r2.Box {
halfX := (b.Max.X - b.Min.X) / 2
halfY := (b.Max.Y - b.Min.Y) / 2
switch dir {
case ne:
b.Min.X += halfX
b.Max.Y -= halfY
case se:
b.Min.X += halfX
b.Min.Y += halfY
case sw:
b.Max.X -= halfX
b.Min.Y += halfY
case nw:
b.Max.X -= halfX
b.Max.Y -= halfY
}
return b
}
// summarize updates node masses and centers of mass.
func (t *tile) summarize() (center r2.Vec, mass float64) {
for _, d := range &t.nodes {
if d == nil {
continue
}
c, m := d.summarize()
t.center.X += c.X * m
t.center.Y += c.Y * m
t.mass += m
}
t.center.X /= t.mass
t.center.Y /= t.mass
return t.center, t.mass
}
// forceOn returns a force vector on p given p's mass m and the force
// calculation function, using the Barnes-Hut theta approximation parameter.
func (t *tile) forceOn(p Particle2, pt r2.Vec, m, theta float64, f Force2) (vector r2.Vec) {
s := ((t.bounds.Max.X - t.bounds.Min.X) + (t.bounds.Max.Y - t.bounds.Min.Y)) / 2
d := math.Hypot(pt.X-t.center.X, pt.Y-t.center.Y)
if s/d < theta || t.particle != nil {
return f(p, t.particle, m, t.mass, t.center.Sub(pt))
}
var v r2.Vec
for _, d := range &t.nodes {
if d == nil {
continue
}
v = v.Add(d.forceOn(p, pt, m, theta, f))
}
return v
}

View File

@@ -0,0 +1,524 @@
// Copyright ©2019 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package barneshut
import (
"fmt"
"math"
"reflect"
"testing"
"golang.org/x/exp/rand"
"gonum.org/v1/gonum/floats"
"gonum.org/v1/gonum/spatial/r2"
)
type particle2 struct {
x, y, m float64
name string
}
func (p particle2) Coord2() r2.Vec { return r2.Vec{X: p.x, Y: p.y} }
func (p particle2) Mass() float64 { return p.m }
var planeTests = []struct {
name string
particles []particle2
want *Plane
}{
{
name: "nil",
particles: nil,
want: &Plane{},
},
{
name: "empty",
particles: []particle2{},
want: &Plane{Particles: []Particle2{}},
},
{
name: "one",
particles: []particle2{{m: 1}}, // Must have a mass to avoid vacuum decay.
want: &Plane{
root: tile{
particle: particle2{x: 0, y: 0, m: 1},
bounds: r2.Box{Min: r2.Vec{X: 0, Y: 0}, Max: r2.Vec{X: 0, Y: 0}},
center: r2.Vec{X: 0, Y: 0},
mass: 1,
},
Particles: []Particle2{
particle2{m: 1},
},
},
},
{
name: "3 corners",
particles: []particle2{
{x: 1, y: 1, m: 1},
{x: -1, y: 1, m: 1},
{x: -1, y: -1, m: 1},
},
want: &Plane{
root: tile{
bounds: r2.Box{Min: r2.Vec{X: -1, Y: -1}, Max: r2.Vec{X: 1, Y: 1}},
nodes: [4]*tile{
se: {
particle: particle2{x: 1, y: 1, m: 1},
bounds: r2.Box{Min: r2.Vec{X: 0, Y: 0}, Max: r2.Vec{X: 1, Y: 1}},
center: r2.Vec{X: 1, Y: 1}, mass: 1,
},
sw: {
particle: particle2{x: -1, y: 1, m: 1},
bounds: r2.Box{Min: r2.Vec{X: -1, Y: 0}, Max: r2.Vec{X: 0, Y: 1}},
center: r2.Vec{X: -1, Y: 1}, mass: 1,
},
nw: {
particle: particle2{x: -1, y: -1, m: 1},
bounds: r2.Box{Min: r2.Vec{X: -1, Y: -1}, Max: r2.Vec{X: 0, Y: 0}},
center: r2.Vec{X: -1, Y: -1}, mass: 1,
},
},
center: r2.Vec{X: -0.3333333333333333, Y: 0.3333333333333333},
mass: 3,
},
Particles: []Particle2{
particle2{x: 1, y: 1, m: 1},
particle2{x: -1, y: 1, m: 1},
particle2{x: -1, y: -1, m: 1},
},
},
},
{
name: "4 corners",
particles: []particle2{
{x: 1, y: 1, m: 1},
{x: -1, y: 1, m: 1},
{x: 1, y: -1, m: 1},
{x: -1, y: -1, m: 1},
},
want: &Plane{
root: tile{
bounds: r2.Box{Min: r2.Vec{X: -1, Y: -1}, Max: r2.Vec{X: 1, Y: 1}},
nodes: [4]*tile{
{
particle: particle2{x: 1, y: -1, m: 1},
bounds: r2.Box{Min: r2.Vec{X: 0, Y: -1}, Max: r2.Vec{X: 1, Y: 0}},
center: r2.Vec{X: 1, Y: -1},
mass: 1,
},
{
particle: particle2{x: 1, y: 1, m: 1},
bounds: r2.Box{Min: r2.Vec{X: 0, Y: 0}, Max: r2.Vec{X: 1, Y: 1}},
center: r2.Vec{X: 1, Y: 1},
mass: 1,
},
{
particle: particle2{x: -1, y: 1, m: 1},
bounds: r2.Box{Min: r2.Vec{X: -1, Y: 0}, Max: r2.Vec{X: 0, Y: 1}},
center: r2.Vec{X: -1, Y: 1},
mass: 1,
},
{
particle: particle2{x: -1, y: -1, m: 1},
bounds: r2.Box{Min: r2.Vec{X: -1, Y: -1}, Max: r2.Vec{X: 0, Y: 0}},
center: r2.Vec{X: -1, Y: -1},
mass: 1,
},
},
center: r2.Vec{X: 0, Y: 0},
mass: 4,
},
Particles: []Particle2{
particle2{x: 1, y: 1, m: 1},
particle2{x: -1, y: 1, m: 1},
particle2{x: 1, y: -1, m: 1},
particle2{x: -1, y: -1, m: 1},
},
},
},
{
name: "5 corners",
particles: []particle2{
{x: 1, y: 1, m: 1},
{x: -1, y: 1, m: 1},
{x: 1, y: -1, m: 1},
{x: -1, y: -1, m: 1},
{x: -1.1, y: -1, m: 1},
},
want: &Plane{
root: tile{
bounds: r2.Box{Min: r2.Vec{X: -1.1, Y: -1}, Max: r2.Vec{X: 1, Y: 1}},
nodes: [4]*tile{
{
particle: particle2{x: 1, y: -1, m: 1},
bounds: r2.Box{Min: r2.Vec{X: -0.050000000000000044, Y: -1}, Max: r2.Vec{X: 1, Y: 0}},
center: r2.Vec{X: 1, Y: -1},
mass: 1,
},
{
particle: particle2{x: 1, y: 1, m: 1},
bounds: r2.Box{Min: r2.Vec{X: -0.050000000000000044, Y: 0}, Max: r2.Vec{X: 1, Y: 1}},
center: r2.Vec{X: 1, Y: 1},
mass: 1,
},
{
particle: particle2{x: -1, y: 1, m: 1},
bounds: r2.Box{Min: r2.Vec{X: -1.1, Y: 0}, Max: r2.Vec{X: -0.050000000000000044, Y: 1}},
center: r2.Vec{X: -1, Y: 1},
mass: 1,
},
{
bounds: r2.Box{Min: r2.Vec{X: -1.1, Y: -1}, Max: r2.Vec{X: -0.050000000000000044, Y: 0}},
nodes: [4]*tile{
nw: {
bounds: r2.Box{Min: r2.Vec{X: -1.1, Y: -1}, Max: r2.Vec{X: -0.5750000000000001, Y: -0.5}},
nodes: [4]*tile{
nw: {
bounds: r2.Box{Min: r2.Vec{X: -1.1, Y: -1}, Max: r2.Vec{X: -0.8375000000000001, Y: -0.75}},
nodes: [4]*tile{
nw: {
bounds: r2.Box{Min: r2.Vec{X: -1.1, Y: -1}, Max: r2.Vec{X: -0.9687500000000001, Y: -0.875}},
nodes: [4]*tile{
ne: {
particle: particle2{x: -1, y: -1, m: 1},
bounds: r2.Box{Min: r2.Vec{X: -1.034375, Y: -1}, Max: r2.Vec{X: -0.9687500000000001, Y: -0.9375}},
center: r2.Vec{X: -1, Y: -1},
mass: 1,
},
nw: {
particle: particle2{x: -1.1, y: -1, m: 1},
bounds: r2.Box{Min: r2.Vec{X: -1.1, Y: -1}, Max: r2.Vec{X: -1.034375, Y: -0.9375}},
center: r2.Vec{X: -1.1, Y: -1},
mass: 1,
},
},
center: r2.Vec{X: -1.05, Y: -1},
mass: 2,
},
},
center: r2.Vec{X: -1.05, Y: -1},
mass: 2,
},
},
center: r2.Vec{X: -1.05, Y: -1},
mass: 2,
},
},
center: r2.Vec{X: -1.05, Y: -1},
mass: 2,
},
},
center: r2.Vec{X: -0.22000000000000003, Y: -0.2},
mass: 5,
},
Particles: []Particle2{
particle2{x: 1, y: 1, m: 1},
particle2{x: -1, y: 1, m: 1},
particle2{x: 1, y: -1, m: 1},
particle2{x: -1, y: -1, m: 1},
particle2{x: -1.1, y: -1, m: 1},
},
},
},
{
// Note that the code here subdivides the space differently to
// how it is split in the example, since Plane makes a minimum
// bounding box based on the data, while the example does not.
name: "http://arborjs.org/docs/barnes-hut example",
particles: []particle2{
{x: 64.5, y: 81.5, m: 1, name: "A"},
{x: 242, y: 34, m: 1, name: "B"},
{x: 199, y: 69, m: 1, name: "C"},
{x: 285, y: 106.5, m: 1, name: "D"},
{x: 170, y: 194.5, m: 1, name: "E"},
{x: 42.5, y: 334.5, m: 1, name: "F"},
{x: 147, y: 309, m: 1, name: "G"},
{x: 236.5, y: 324, m: 1, name: "H"},
},
want: &Plane{
root: tile{
bounds: r2.Box{Min: r2.Vec{X: 42.5, Y: 34}, Max: r2.Vec{X: 285, Y: 334.5}},
nodes: [4]*tile{
{
bounds: r2.Box{Min: r2.Vec{X: 163.75, Y: 34}, Max: r2.Vec{X: 285, Y: 184.25}},
nodes: [4]*tile{
ne: {
bounds: r2.Box{Min: r2.Vec{X: 224.375, Y: 34}, Max: r2.Vec{X: 285, Y: 109.125}},
nodes: [4]*tile{
se: {
particle: particle2{x: 285, y: 106.5, m: 1, name: "D"},
bounds: r2.Box{Min: r2.Vec{X: 254.6875, Y: 71.5625}, Max: r2.Vec{X: 285, Y: 109.125}},
center: r2.Vec{X: 285, Y: 106.5},
mass: 1,
},
nw: {
particle: particle2{x: 242, y: 34, m: 1, name: "B"},
bounds: r2.Box{Min: r2.Vec{X: 224.375, Y: 34}, Max: r2.Vec{X: 254.6875, Y: 71.5625}},
center: r2.Vec{X: 242, Y: 34},
mass: 1,
},
},
center: r2.Vec{X: 263.5, Y: 70.25},
mass: 2,
},
nw: {
particle: particle2{x: 199, y: 69, m: 1, name: "C"},
bounds: r2.Box{Min: r2.Vec{X: 163.75, Y: 34}, Max: r2.Vec{X: 224.375, Y: 109.125}},
center: r2.Vec{X: 199, Y: 69},
mass: 1,
},
},
center: r2.Vec{X: 242, Y: 69.83333333333333},
mass: 3,
},
{
bounds: r2.Box{Min: r2.Vec{X: 163.75, Y: 184.25}, Max: r2.Vec{X: 285, Y: 334.5}},
nodes: [4]*tile{
se: {
particle: particle2{x: 236.5, y: 324, m: 1, name: "H"},
bounds: r2.Box{Min: r2.Vec{X: 224.375, Y: 259.375}, Max: r2.Vec{X: 285, Y: 334.5}},
center: r2.Vec{X: 236.5, Y: 324},
mass: 1,
},
nw: {
particle: particle2{x: 170, y: 194.5, m: 1, name: "E"},
bounds: r2.Box{Min: r2.Vec{X: 163.75, Y: 184.25}, Max: r2.Vec{X: 224.375, Y: 259.375}},
center: r2.Vec{X: 170, Y: 194.5},
mass: 1,
},
},
center: r2.Vec{X: 203.25, Y: 259.25},
mass: 2,
},
{
bounds: r2.Box{Min: r2.Vec{X: 42.5, Y: 184.25}, Max: r2.Vec{X: 163.75, Y: 334.5}},
nodes: [4]*tile{
se: {
particle: particle2{x: 147, y: 309, m: 1, name: "G"},
bounds: r2.Box{Min: r2.Vec{X: 103.125, Y: 259.375}, Max: r2.Vec{X: 163.75, Y: 334.5}},
center: r2.Vec{X: 147, Y: 309},
mass: 1,
},
sw: {
particle: particle2{x: 42.5, y: 334.5, m: 1, name: "F"},
bounds: r2.Box{Min: r2.Vec{X: 42.5, Y: 259.375}, Max: r2.Vec{X: 103.125, Y: 334.5}},
center: r2.Vec{X: 42.5, Y: 334.5},
mass: 1,
},
},
center: r2.Vec{X: 94.75, Y: 321.75},
mass: 2,
},
{
particle: particle2{x: 64.5, y: 81.5, m: 1, name: "A"},
bounds: r2.Box{Min: r2.Vec{X: 42.5, Y: 34}, Max: r2.Vec{X: 163.75, Y: 184.25}},
center: r2.Vec{X: 64.5, Y: 81.5},
mass: 1,
},
},
center: r2.Vec{X: 173.3125, Y: 181.625},
mass: 8,
},
Particles: []Particle2{
particle2{x: 64.5, y: 81.5, m: 1, name: "A"},
particle2{x: 242, y: 34, m: 1, name: "B"},
particle2{x: 199, y: 69, m: 1, name: "C"},
particle2{x: 285, y: 106.5, m: 1, name: "D"},
particle2{x: 170, y: 194.5, m: 1, name: "E"},
particle2{x: 42.5, y: 334.5, m: 1, name: "F"},
particle2{x: 147, y: 309, m: 1, name: "G"},
particle2{x: 236.5, y: 324, m: 1, name: "H"},
},
},
},
}
func TestPlane(t *testing.T) {
const tol = 1e-15
for _, test := range planeTests {
var particles []Particle2
if test.particles != nil {
particles = make([]Particle2, len(test.particles))
}
for i, p := range test.particles {
particles[i] = p
}
got := NewPlane(particles)
if test.want != nil && !reflect.DeepEqual(got, test.want) {
t.Errorf("unexpected result for %q: got:%v want:%v", test.name, got, test.want)
}
// Recursively check all internal centers of mass.
walkPlane(&got.root, func(tl *tile) {
var sub []Particle2
walkPlane(tl, func(tl *tile) {
if tl.particle != nil {
sub = append(sub, tl.particle)
}
})
center, mass := centerOfMass2(sub)
if !floats.EqualWithinAbsOrRel(center.X, tl.center.X, tol, tol) || !floats.EqualWithinAbsOrRel(center.Y, tl.center.Y, tol, tol) {
t.Errorf("unexpected result for %q for center of mass: got:%f want:%f", test.name, tl.center, center)
}
if !floats.EqualWithinAbsOrRel(mass, tl.mass, tol, tol) {
t.Errorf("unexpected result for %q for total mass: got:%f want:%f", test.name, tl.mass, mass)
}
})
}
}
func centerOfMass2(particles []Particle2) (center r2.Vec, mass float64) {
for _, p := range particles {
m := p.Mass()
mass += m
c := p.Coord2()
center.X += c.X * m
center.Y += c.Y * m
}
if mass != 0 {
center.X /= mass
center.Y /= mass
}
return center, mass
}
func walkPlane(t *tile, fn func(*tile)) {
if t == nil {
return
}
fn(t)
for _, q := range t.nodes {
walkPlane(q, fn)
}
}
func TestPlaneForceOn(t *testing.T) {
const (
size = 1000
tol = 0.07
)
for _, n := range []int{3e3, 1e4, 3e4} {
rnd := rand.New(rand.NewSource(1))
particles := make([]Particle2, n)
for i := range particles {
particles[i] = particle2{x: size * rnd.Float64(), y: size * rnd.Float64(), m: 1}
}
moved := make([]r2.Vec, n)
for i, p := range particles {
var v r2.Vec
m := p.Mass()
pv := p.Coord2()
for _, e := range particles {
v = v.Add(Gravity2(p, e, m, e.Mass(), e.Coord2().Sub(pv)))
}
moved[i] = p.Coord2().Add(v)
}
plane := NewPlane(particles)
for _, theta := range []float64{0, 0.3, 0.6, 0.9} {
t.Run(fmt.Sprintf("%d-body/theta=%v", len(particles), theta), func(t *testing.T) {
var ssd, sd float64
var calls int
for i, p := range particles {
v := plane.ForceOn(p, theta, func(p1, p2 Particle2, m1, m2 float64, v r2.Vec) r2.Vec {
calls++
return Gravity2(p1, p2, m1, m2, v)
})
pos := p.Coord2().Add(v)
d := moved[i].Sub(pos)
ssd += d.X*d.X + d.Y*d.Y
sd += math.Hypot(d.X, d.Y)
}
rmsd := math.Sqrt(ssd / float64(len(particles)))
if rmsd > tol {
t.Error("RMSD for approximation too high")
}
t.Logf("rmsd=%.4v md=%.4v calls/particle=%.5v",
rmsd, sd/float64(len(particles)), float64(calls)/float64(len(particles)))
})
}
}
}
var (
fv2sink r2.Vec
planeSink *Plane
)
func BenchmarkNewPlane(b *testing.B) {
for _, n := range []int{1e3, 1e4, 1e5, 1e6} {
rnd := rand.New(rand.NewSource(1))
particles := make([]Particle2, n)
for i := range particles {
particles[i] = particle2{x: rnd.Float64(), y: rnd.Float64(), m: 1}
}
b.ResetTimer()
b.Run(fmt.Sprintf("%d-body", len(particles)), func(b *testing.B) {
for i := 0; i < b.N; i++ {
planeSink = NewPlane(particles)
}
})
}
}
func BenchmarkPlaneForceOn(b *testing.B) {
for _, n := range []int{1e3, 1e4, 1e5} {
for _, theta := range []float64{0, 0.1, 0.5, 1, 1.5} {
if n > 1e4 && theta < 0.5 {
// Don't run unreasonably long benchmarks.
continue
}
rnd := rand.New(rand.NewSource(1))
particles := make([]Particle2, n)
for i := range particles {
particles[i] = particle2{x: rnd.Float64(), y: rnd.Float64(), m: 1}
}
plane := NewPlane(particles)
b.ResetTimer()
b.Run(fmt.Sprintf("%d-body/theta=%v", len(particles), theta), func(b *testing.B) {
for i := 0; i < b.N; i++ {
for _, p := range particles {
fv2sink = plane.ForceOn(p, theta, Gravity2)
}
}
})
}
}
}
func BenchmarkPlaneFull(b *testing.B) {
for _, n := range []int{1e3, 1e4, 1e5} {
for _, theta := range []float64{0, 0.1, 0.5, 1, 1.5} {
if n > 1e4 && theta < 0.5 {
// Don't run unreasonably long benchmarks.
continue
}
rnd := rand.New(rand.NewSource(1))
particles := make([]Particle2, n)
for i := range particles {
particles[i] = particle2{x: rnd.Float64(), y: rnd.Float64(), m: 1}
}
b.ResetTimer()
b.Run(fmt.Sprintf("%d-body/theta=%v", len(particles), theta), func(b *testing.B) {
for i := 0; i < b.N; i++ {
plane := NewPlane(particles)
for _, p := range particles {
fv2sink = plane.ForceOn(p, theta, Gravity2)
}
}
})
}
}
}

View File

@@ -0,0 +1,300 @@
// Copyright ©2019 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package barneshut
import (
"fmt"
"math"
"gonum.org/v1/gonum/spatial/r3"
)
// Particle3 is a particle in a volume.
type Particle3 interface {
Coord3() r3.Vec
Mass() float64
}
// Force3 is a force modeling function for interactions between p1 and p2,
// m1 is the mass of p1 and m2 of p2. The vector v is the vector from p1 to
// p2. The returned value is the force vector acting on p1.
//
// In models where the identity of particles must be known, p1 and p2 may be
// compared. Force3 may be passed nil for p2 when the Barnes-Hut approximation
// is being used. A nil p2 indicates that the second mass center is an
// aggregate.
type Force3 func(p1, p2 Particle3, m1, m2 float64, v r3.Vec) r3.Vec
// Gravity3 returns a vector force on m1 by m2, equal to (m1⋅m2)/‖v‖²
// in the directions of v. Gravity3 ignores the identity of the interacting
// particles and returns a zero vector when the two particles are
// coincident, but performs no other sanity checks.
func Gravity3(_, _ Particle3, m1, m2 float64, v r3.Vec) r3.Vec {
d2 := v.X*v.X + v.Y*v.Y + v.Z*v.Z
if d2 == 0 {
return r3.Vec{}
}
return v.Scale((m1 * m2) / (d2 * math.Sqrt(d2)))
}
// Volume implements Barnes-Hut force approximation calculations.
type Volume struct {
root bucket
Particles []Particle3
}
// NewVolume returns a new Volume.
func NewVolume(p []Particle3) *Volume {
q := Volume{Particles: p}
q.Reset()
return &q
}
// Reset reconstructs the Barnes-Hut tree. Reset must be called if the
// Particles field or elements of Particles have been altered, unless
// ForceOn is called with theta=0 or no data structures have been
// previously built.
func (q *Volume) Reset() {
if len(q.Particles) == 0 {
q.root = bucket{}
return
}
q.root = bucket{
particle: q.Particles[0],
center: q.Particles[0].Coord3(),
mass: q.Particles[0].Mass(),
}
q.root.bounds.Min = q.root.center
q.root.bounds.Max = q.root.center
for _, e := range q.Particles[1:] {
c := e.Coord3()
if c.X < q.root.bounds.Min.X {
q.root.bounds.Min.X = c.X
}
if c.X > q.root.bounds.Max.X {
q.root.bounds.Max.X = c.X
}
if c.Y < q.root.bounds.Min.Y {
q.root.bounds.Min.Y = c.Y
}
if c.Y > q.root.bounds.Max.Y {
q.root.bounds.Max.Y = c.Y
}
if c.Z < q.root.bounds.Min.Z {
q.root.bounds.Min.Z = c.Z
}
if c.Z > q.root.bounds.Max.Z {
q.root.bounds.Max.Z = c.Z
}
}
// TODO(kortschak): Partially parallelise this by
// choosing the direction and using one of eight
// goroutines to work on each root octant.
for _, e := range q.Particles[1:] {
q.root.insert(e)
}
q.root.summarize()
}
// ForceOn returns a force vector on p given p's mass and the force function, f,
// using the Barnes-Hut theta approximation parameter.
//
// Calls to f will include p in the p1 position and a non-nil p2 if the force
// interaction is with a non-aggregate mass center, otherwise p2 will be nil.
//
// It is safe to call ForceOn concurrently.
func (q *Volume) ForceOn(p Particle3, theta float64, f Force3) (force r3.Vec) {
var empty bucket
if theta > 0 && q.root != empty {
return q.root.forceOn(p, p.Coord3(), p.Mass(), theta, f)
}
// For the degenerate case, just iterate over the
// slice of particles rather than walking the tree.
var v r3.Vec
m := p.Mass()
pv := p.Coord3()
for _, e := range q.Particles {
v = v.Add(f(p, e, m, e.Mass(), e.Coord3().Sub(pv)))
}
return v
}
// bucket is an oct tree octant with Barnes-Hut extensions.
type bucket struct {
particle Particle3
bounds r3.Box
nodes [8]*bucket
center r3.Vec
mass float64
}
// insert inserts p into the subtree rooted at b.
func (b *bucket) insert(p Particle3) {
if b.particle == nil {
for _, q := range b.nodes {
if q != nil {
b.passDown(p)
return
}
}
b.particle = p
b.center = p.Coord3()
b.mass = p.Mass()
return
}
b.passDown(p)
b.passDown(b.particle)
b.particle = nil
b.center = r3.Vec{}
b.mass = 0
}
func (b *bucket) passDown(p Particle3) {
dir := octantOf(b.bounds, p)
if b.nodes[dir] == nil {
b.nodes[dir] = &bucket{bounds: splitVolume(b.bounds, dir)}
}
b.nodes[dir].insert(p)
}
const (
lne = iota
lse
lsw
lnw
une
use
usw
unw
)
// octantOf returns which octant of b that p should be placed in.
func octantOf(b r3.Box, p Particle3) int {
center := r3.Vec{
X: (b.Min.X + b.Max.X) / 2,
Y: (b.Min.Y + b.Max.Y) / 2,
Z: (b.Min.Z + b.Max.Z) / 2,
}
c := p.Coord3()
if checkBounds && (c.X < b.Min.X || b.Max.X < c.X || c.Y < b.Min.Y || b.Max.Y < c.Y || c.Z < b.Min.Z || b.Max.Z < c.Z) {
panic(fmt.Sprintf("p out of range %+v: %#v", b, p))
}
if c.X < center.X {
if c.Y < center.Y {
if c.Z < center.Z {
return lnw
} else {
return unw
}
} else {
if c.Z < center.Z {
return lsw
} else {
return usw
}
}
} else {
if c.Y < center.Y {
if c.Z < center.Z {
return lne
} else {
return une
}
} else {
if c.Z < center.Z {
return lse
} else {
return use
}
}
}
}
// splitVolume returns an octant subdivision of b in the given direction.
func splitVolume(b r3.Box, dir int) r3.Box {
halfX := (b.Max.X - b.Min.X) / 2
halfY := (b.Max.Y - b.Min.Y) / 2
halfZ := (b.Max.Z - b.Min.Z) / 2
switch dir {
case lne:
b.Min.X += halfX
b.Max.Y -= halfY
b.Max.Z -= halfZ
case lse:
b.Min.X += halfX
b.Min.Y += halfY
b.Max.Z -= halfZ
case lsw:
b.Max.X -= halfX
b.Min.Y += halfY
b.Max.Z -= halfZ
case lnw:
b.Max.X -= halfX
b.Max.Y -= halfY
b.Max.Z -= halfZ
case une:
b.Min.X += halfX
b.Max.Y -= halfY
b.Min.Z += halfZ
case use:
b.Min.X += halfX
b.Min.Y += halfY
b.Min.Z += halfZ
case usw:
b.Max.X -= halfX
b.Min.Y += halfY
b.Min.Z += halfZ
case unw:
b.Max.X -= halfX
b.Max.Y -= halfY
b.Min.Z += halfZ
}
return b
}
// summarize updates node masses and centers of mass.
func (b *bucket) summarize() (center r3.Vec, mass float64) {
for _, d := range &b.nodes {
if d == nil {
continue
}
c, m := d.summarize()
b.center.X += c.X * m
b.center.Y += c.Y * m
b.center.Z += c.Z * m
b.mass += m
}
b.center.X /= b.mass
b.center.Y /= b.mass
b.center.Z /= b.mass
return b.center, b.mass
}
// forceOn returns a force vector on p given p's mass m and the force
// calculation function, using the Barnes-Hut theta approximation parameter.
func (b *bucket) forceOn(p Particle3, pt r3.Vec, m, theta float64, f Force3) (vector r3.Vec) {
s := ((b.bounds.Max.X - b.bounds.Min.X) + (b.bounds.Max.Y - b.bounds.Min.Y) + (b.bounds.Max.Z - b.bounds.Min.Z)) / 3
d := math.Hypot(math.Hypot(pt.X-b.center.X, pt.Y-b.center.Y), pt.Z-b.center.Z)
if s/d < theta || b.particle != nil {
return f(p, b.particle, m, b.mass, b.center.Sub(pt))
}
var v r3.Vec
for _, d := range &b.nodes {
if d == nil {
continue
}
v = v.Add(d.forceOn(p, pt, m, theta, f))
}
return v
}

View File

@@ -0,0 +1,521 @@
// Copyright ©2019 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package barneshut
import (
"fmt"
"math"
"reflect"
"testing"
"golang.org/x/exp/rand"
"gonum.org/v1/gonum/floats"
"gonum.org/v1/gonum/spatial/r3"
)
type particle3 struct {
x, y, z, m float64
name string
}
func (p particle3) Coord3() r3.Vec { return r3.Vec{X: p.x, Y: p.y, Z: p.z} }
func (p particle3) Mass() float64 { return p.m }
var volumeTests = []struct {
name string
particles []particle3
want *Volume
}{
{
name: "nil",
particles: nil,
want: &Volume{},
},
{
name: "empty",
particles: []particle3{},
want: &Volume{Particles: []Particle3{}},
},
{
name: "one",
particles: []particle3{{m: 1}}, // Must have a mass to avoid vacuum decay.
want: &Volume{
root: bucket{
particle: particle3{x: 0, y: 0, z: 0, m: 1},
bounds: r3.Box{Min: r3.Vec{X: 0, Y: 0, Z: 0}, Max: r3.Vec{X: 0, Y: 0, Z: 0}},
center: r3.Vec{X: 0, Y: 0, Z: 0},
mass: 1,
},
Particles: []Particle3{
particle3{x: 0, y: 0, z: 0, m: 1},
},
},
},
{
name: "3 corners",
particles: []particle3{
{x: 1, y: 1, z: 1, m: 1},
{x: -1, y: 1, z: 0, m: 1},
{x: -1, y: -1, z: -1, m: 1},
},
want: &Volume{
root: bucket{
bounds: r3.Box{Min: r3.Vec{X: -1, Y: -1, Z: -1}, Max: r3.Vec{X: 1, Y: 1, Z: 1}},
nodes: [8]*bucket{
lnw: {
particle: particle3{x: -1, y: -1, z: -1, m: 1},
bounds: r3.Box{Min: r3.Vec{X: -1, Y: -1, Z: -1}, Max: r3.Vec{X: 0, Y: 0, Z: 0}},
center: r3.Vec{X: -1, Y: -1, Z: -1},
mass: 1,
},
use: {
particle: particle3{x: 1, y: 1, z: 1, m: 1},
bounds: r3.Box{Min: r3.Vec{X: 0, Y: 0, Z: 0}, Max: r3.Vec{X: 1, Y: 1, Z: 1}},
center: r3.Vec{X: 1, Y: 1, Z: 1},
mass: 1,
},
usw: {
particle: particle3{x: -1, y: 1, z: 0, m: 1},
bounds: r3.Box{Min: r3.Vec{X: -1, Y: 0, Z: 0}, Max: r3.Vec{X: 0, Y: 1, Z: 1}},
center: r3.Vec{X: -1, Y: 1, Z: 0},
mass: 1,
},
},
center: r3.Vec{X: -0.3333333333333333, Y: 0.3333333333333333, Z: 0},
mass: 3,
},
Particles: []Particle3{
particle3{x: 1, y: 1, z: 1, m: 1},
particle3{x: -1, y: 1, z: 0, m: 1},
particle3{x: -1, y: -1, z: -1, m: 1},
},
},
},
{
name: "4 corners",
particles: []particle3{
{x: 1, y: 1, z: -1, m: 1},
{x: -1, y: 1, z: 1, m: 1},
{x: 1, y: -1, z: 1, m: 1},
{x: -1, y: -1, z: -1, m: 1},
},
want: &Volume{
root: bucket{
bounds: r3.Box{Min: r3.Vec{X: -1, Y: -1, Z: -1}, Max: r3.Vec{X: 1, Y: 1, Z: 1}},
nodes: [8]*bucket{
lse: {
particle: particle3{x: 1, y: 1, z: -1, m: 1},
bounds: r3.Box{Min: r3.Vec{X: 0, Y: 0, Z: -1}, Max: r3.Vec{X: 1, Y: 1, Z: 0}},
center: r3.Vec{X: 1, Y: 1, Z: -1},
mass: 1,
},
lnw: {
particle: particle3{x: -1, y: -1, z: -1, m: 1},
bounds: r3.Box{Min: r3.Vec{X: -1, Y: -1, Z: -1}, Max: r3.Vec{X: 0, Y: 0, Z: 0}},
center: r3.Vec{X: -1, Y: -1, Z: -1},
mass: 1,
},
une: {
particle: particle3{x: 1, y: -1, z: 1, m: 1},
bounds: r3.Box{Min: r3.Vec{X: 0, Y: -1, Z: 0}, Max: r3.Vec{X: 1, Y: 0, Z: 1}},
center: r3.Vec{X: 1, Y: -1, Z: 1},
mass: 1,
},
usw: {
particle: particle3{x: -1, y: 1, z: 1, m: 1},
bounds: r3.Box{Min: r3.Vec{X: -1, Y: 0, Z: 0}, Max: r3.Vec{X: 0, Y: 1, Z: 1}},
center: r3.Vec{X: -1, Y: 1, Z: 1},
mass: 1,
},
},
center: r3.Vec{X: 0, Y: 0, Z: 0},
mass: 4,
},
Particles: []Particle3{
particle3{x: 1, y: 1, z: -1, m: 1},
particle3{x: -1, y: 1, z: 1, m: 1},
particle3{x: 1, y: -1, z: 1, m: 1},
particle3{x: -1, y: -1, z: -1, m: 1},
},
},
},
{
name: "5 corners",
particles: []particle3{
{x: 1, y: 1, z: -1, m: 1},
{x: -1, y: 1, z: 1, m: 1},
{x: 1, y: -1, z: 1, m: 1},
{x: -1, y: -1, z: -1, m: 1},
{x: -1.1, y: -1, z: -1.1, m: 1},
},
want: &Volume{
root: bucket{
bounds: r3.Box{Min: r3.Vec{X: -1.1, Y: -1, Z: -1.1}, Max: r3.Vec{X: 1, Y: 1, Z: 1}},
nodes: [8]*bucket{
lse: {
particle: particle3{x: 1, y: 1, z: -1, m: 1},
bounds: r3.Box{Min: r3.Vec{X: -0.050000000000000044, Y: 0, Z: -1.1}, Max: r3.Vec{X: 1, Y: 1, Z: -0.050000000000000044}},
center: r3.Vec{X: 1, Y: 1, Z: -1},
mass: 1,
},
lnw: {
bounds: r3.Box{Min: r3.Vec{X: -1.1, Y: -1, Z: -1.1}, Max: r3.Vec{X: -0.050000000000000044, Y: 0, Z: -0.050000000000000044}},
nodes: [8]*bucket{
lnw: {
bounds: r3.Box{Min: r3.Vec{X: -1.1, Y: -1, Z: -1.1}, Max: r3.Vec{X: -0.5750000000000001, Y: -0.5, Z: -0.5750000000000001}},
nodes: [8]*bucket{
lnw: {
bounds: r3.Box{Min: r3.Vec{X: -1.1, Y: -1, Z: -1.1}, Max: r3.Vec{X: -0.8375000000000001, Y: -0.75, Z: -0.8375000000000001}},
nodes: [8]*bucket{
lnw: {
bounds: r3.Box{Min: r3.Vec{X: -1.1, Y: -1, Z: -1.1}, Max: r3.Vec{X: -0.9687500000000001, Y: -0.875, Z: -0.9687500000000001}},
nodes: [8]*bucket{
lnw: {
particle: particle3{x: -1.1, y: -1, z: -1.1, m: 1},
bounds: r3.Box{Min: r3.Vec{X: -1.1, Y: -1, Z: -1.1}, Max: r3.Vec{X: -1.034375, Y: -0.9375, Z: -1.034375}},
center: r3.Vec{X: -1.1, Y: -1, Z: -1.1},
mass: 1,
},
une: {
particle: particle3{x: -1, y: -1, z: -1, m: 1},
bounds: r3.Box{Min: r3.Vec{X: -1.034375, Y: -1, Z: -1.034375}, Max: r3.Vec{X: -0.9687500000000001, Y: -0.9375, Z: -0.9687500000000001}},
center: r3.Vec{X: -1, Y: -1, Z: -1},
mass: 1,
},
},
center: r3.Vec{X: -1.05, Y: -1, Z: -1.05},
mass: 2,
},
},
center: r3.Vec{X: -1.05, Y: -1, Z: -1.05},
mass: 2,
},
},
center: r3.Vec{X: -1.05, Y: -1, Z: -1.05},
mass: 2,
},
},
center: r3.Vec{X: -1.05, Y: -1, Z: -1.05},
mass: 2,
},
une: {
particle: particle3{x: 1, y: -1, z: 1, m: 1},
bounds: r3.Box{Min: r3.Vec{X: -0.050000000000000044, Y: -1, Z: -0.050000000000000044}, Max: r3.Vec{X: 1, Y: 0, Z: 1}},
center: r3.Vec{X: 1, Y: -1, Z: 1},
mass: 1,
},
usw: {
particle: particle3{x: -1, y: 1, z: 1, m: 1},
bounds: r3.Box{Min: r3.Vec{X: -1.1, Y: 0, Z: -0.050000000000000044}, Max: r3.Vec{X: -0.050000000000000044, Y: 1, Z: 1}},
center: r3.Vec{X: -1, Y: 1, Z: 1},
mass: 1,
},
},
center: r3.Vec{X: -0.22000000000000003, Y: -0.2, Z: -0.22000000000000003},
mass: 5,
},
Particles: []Particle3{
particle3{x: 1, y: 1, z: -1, m: 1},
particle3{x: -1, y: 1, z: 1, m: 1},
particle3{x: 1, y: -1, z: 1, m: 1},
particle3{x: -1, y: -1, z: -1, m: 1},
particle3{x: -1.1, y: -1, z: -1.1, m: 1},
},
},
},
{
// This case is derived from the 2D example of the same name,
// but with a monotonic increase in Z position according to name.
name: "http://arborjs.org/docs/barnes-hut example",
particles: []particle3{
{x: 64.5, y: 81.5, z: 0, m: 1, name: "A"},
{x: 242, y: 34, z: 40, m: 1, name: "B"},
{x: 199, y: 69, z: 80, m: 1, name: "C"},
{x: 285, y: 106.5, z: 120, m: 1, name: "D"},
{x: 170, y: 194.5, z: 160, m: 1, name: "E"},
{x: 42.5, y: 334.5, z: 200, m: 1, name: "F"},
{x: 147, y: 309, z: 240, m: 1, name: "G"},
{x: 236.5, y: 324, z: 280, m: 1, name: "H"},
},
want: &Volume{
root: bucket{
bounds: r3.Box{Min: r3.Vec{X: 42.5, Y: 34, Z: 0}, Max: r3.Vec{X: 285, Y: 334.5, Z: 280}},
nodes: [8]*bucket{
lne: {
bounds: r3.Box{Min: r3.Vec{X: 163.75, Y: 34, Z: 0}, Max: r3.Vec{X: 285, Y: 184.25, Z: 140}},
nodes: [8]*bucket{
lne: {
particle: particle3{x: 242, y: 34, z: 40, m: 1, name: "B"},
bounds: r3.Box{Min: r3.Vec{X: 224.375, Y: 34, Z: 0}, Max: r3.Vec{X: 285, Y: 109.125, Z: 70}},
center: r3.Vec{X: 242, Y: 34, Z: 40},
mass: 1,
},
une: {
particle: particle3{x: 285, y: 106.5, z: 120, m: 1, name: "D"},
bounds: r3.Box{Min: r3.Vec{X: 224.375, Y: 34, Z: 70}, Max: r3.Vec{X: 285, Y: 109.125, Z: 140}},
center: r3.Vec{X: 285, Y: 106.5, Z: 120},
mass: 1,
},
unw: {
particle: particle3{x: 199, y: 69, z: 80, m: 1, name: "C"},
bounds: r3.Box{Min: r3.Vec{X: 163.75, Y: 34, Z: 70}, Max: r3.Vec{X: 224.375, Y: 109.125, Z: 140}},
center: r3.Vec{X: 199, Y: 69, Z: 80},
mass: 1,
},
},
center: r3.Vec{X: 242, Y: 69.83333333333333, Z: 80},
mass: 3,
},
lnw: {
particle: particle3{x: 64.5, y: 81.5, z: 0, m: 1, name: "A"},
bounds: r3.Box{Min: r3.Vec{X: 42.5, Y: 34, Z: 0}, Max: r3.Vec{X: 163.75, Y: 184.25, Z: 140}},
center: r3.Vec{X: 64.5, Y: 81.5, Z: 0},
mass: 1,
},
(*bucket)(nil),
use: {
bounds: r3.Box{Min: r3.Vec{X: 163.75, Y: 184.25, Z: 140}, Max: r3.Vec{X: 285, Y: 334.5, Z: 280}},
nodes: [8]*bucket{
lnw: {
particle: particle3{x: 170, y: 194.5, z: 160, m: 1, name: "E"},
bounds: r3.Box{Min: r3.Vec{X: 163.75, Y: 184.25, Z: 140}, Max: r3.Vec{X: 224.375, Y: 259.375, Z: 210}},
center: r3.Vec{X: 170, Y: 194.5, Z: 160},
mass: 1,
},
use: {
particle: particle3{x: 236.5, y: 324, z: 280, m: 1, name: "H"},
bounds: r3.Box{Min: r3.Vec{X: 224.375, Y: 259.375, Z: 210}, Max: r3.Vec{X: 285, Y: 334.5, Z: 280}},
center: r3.Vec{X: 236.5, Y: 324, Z: 280},
mass: 1,
},
},
center: r3.Vec{X: 203.25, Y: 259.25, Z: 220},
mass: 2,
},
usw: {
bounds: r3.Box{Min: r3.Vec{X: 42.5, Y: 184.25, Z: 140}, Max: r3.Vec{X: 163.75, Y: 334.5, Z: 280}},
nodes: [8]*bucket{
lsw: {
particle: particle3{x: 42.5, y: 334.5, z: 200, m: 1, name: "F"},
bounds: r3.Box{Min: r3.Vec{X: 42.5, Y: 259.375, Z: 140}, Max: r3.Vec{X: 103.125, Y: 334.5, Z: 210}},
center: r3.Vec{X: 42.5, Y: 334.5, Z: 200},
mass: 1,
},
use: {
particle: particle3{x: 147, y: 309, z: 240, m: 1, name: "G"},
bounds: r3.Box{Min: r3.Vec{X: 103.125, Y: 259.375, Z: 210}, Max: r3.Vec{X: 163.75, Y: 334.5, Z: 280}},
center: r3.Vec{X: 147, Y: 309, Z: 240},
mass: 1,
},
},
center: r3.Vec{X: 94.75, Y: 321.75, Z: 220},
mass: 2,
},
},
center: r3.Vec{X: 173.3125, Y: 181.625, Z: 140},
mass: 8,
},
Particles: []Particle3{
particle3{x: 64.5, y: 81.5, z: 0, m: 1, name: "A"},
particle3{x: 242, y: 34, z: 40, m: 1, name: "B"},
particle3{x: 199, y: 69, z: 80, m: 1, name: "C"},
particle3{x: 285, y: 106.5, z: 120, m: 1, name: "D"},
particle3{x: 170, y: 194.5, z: 160, m: 1, name: "E"},
particle3{x: 42.5, y: 334.5, z: 200, m: 1, name: "F"},
particle3{x: 147, y: 309, z: 240, m: 1, name: "G"},
particle3{x: 236.5, y: 324, z: 280, m: 1, name: "H"},
},
},
},
}
func TestVolume(t *testing.T) {
const tol = 1e-15
for _, test := range volumeTests {
var particles []Particle3
if test.particles != nil {
particles = make([]Particle3, len(test.particles))
}
for i, p := range test.particles {
particles[i] = p
}
got := NewVolume(particles)
if test.want != nil && !reflect.DeepEqual(got, test.want) {
t.Errorf("unexpected result for %q: got:%v want:%v", test.name, got, test.want)
}
// Recursively check all internal centers of mass.
walkVolume(&got.root, func(b *bucket) {
var sub []Particle3
walkVolume(b, func(b *bucket) {
if b.particle != nil {
sub = append(sub, b.particle)
}
})
center, mass := centerOfMass3(sub)
if !floats.EqualWithinAbsOrRel(center.X, b.center.X, tol, tol) || !floats.EqualWithinAbsOrRel(center.Y, b.center.Y, tol, tol) || !floats.EqualWithinAbsOrRel(center.Z, b.center.Z, tol, tol) {
t.Errorf("unexpected result for %q for center of mass: got:%f want:%f", test.name, b.center, center)
}
if !floats.EqualWithinAbsOrRel(mass, b.mass, tol, tol) {
t.Errorf("unexpected result for %q for total mass: got:%f want:%f", test.name, b.mass, mass)
}
})
}
}
func centerOfMass3(particles []Particle3) (center r3.Vec, mass float64) {
for _, p := range particles {
m := p.Mass()
mass += m
c := p.Coord3()
center.X += c.X * m
center.Y += c.Y * m
center.Z += c.Z * m
}
if mass != 0 {
center.X /= mass
center.Y /= mass
center.Z /= mass
}
return center, mass
}
func walkVolume(t *bucket, fn func(*bucket)) {
if t == nil {
return
}
fn(t)
for _, q := range t.nodes {
walkVolume(q, fn)
}
}
func TestVolumeForceOn(t *testing.T) {
const (
size = 1000
tol = 1e-3
)
for _, n := range []int{3e3, 1e4, 3e4} {
rnd := rand.New(rand.NewSource(1))
particles := make([]Particle3, n)
for i := range particles {
particles[i] = particle3{x: size * rnd.Float64(), y: size * rnd.Float64(), z: size * rnd.Float64(), m: 1}
}
moved := make([]r3.Vec, n)
for i, p := range particles {
var v r3.Vec
m := p.Mass()
pv := p.Coord3()
for _, e := range particles {
v = v.Add(Gravity3(p, e, m, e.Mass(), e.Coord3().Sub(pv)))
}
moved[i] = p.Coord3().Add(v)
}
volume := NewVolume(particles)
for _, theta := range []float64{0, 0.3, 0.6, 0.9} {
t.Run(fmt.Sprintf("%d-body/theta=%v", len(particles), theta), func(t *testing.T) {
var ssd, sd float64
var calls int
for i, p := range particles {
v := volume.ForceOn(p, theta, func(p1, p2 Particle3, m1, m2 float64, v r3.Vec) r3.Vec {
calls++
return Gravity3(p1, p2, m1, m2, v)
})
pos := p.Coord3().Add(v)
d := moved[i].Sub(pos)
ssd += d.X*d.X + d.Y*d.Y + d.Z*d.Z
sd += math.Hypot(math.Hypot(d.X, d.Y), d.Z)
}
rmsd := math.Sqrt(ssd / float64(len(particles)))
if rmsd > tol {
t.Error("RMSD for approximation too high")
}
t.Logf("rmsd=%.4v md=%.4v calls/particle=%.5v",
rmsd, sd/float64(len(particles)), float64(calls)/float64(len(particles)))
})
}
}
}
var (
fv3sink r3.Vec
volumeSink *Volume
)
func BenchmarkNewVolume(b *testing.B) {
for _, n := range []int{1e3, 1e4, 1e5, 1e6} {
rnd := rand.New(rand.NewSource(1))
particles := make([]Particle3, n)
for i := range particles {
particles[i] = particle3{x: rnd.Float64(), y: rnd.Float64(), z: rnd.Float64(), m: 1}
}
b.ResetTimer()
b.Run(fmt.Sprintf("%d-body", len(particles)), func(b *testing.B) {
for i := 0; i < b.N; i++ {
volumeSink = NewVolume(particles)
}
})
}
}
func BenchmarkVolumeForceOn(b *testing.B) {
for _, n := range []int{1e3, 1e4, 1e5} {
for _, theta := range []float64{0, 0.1, 0.5, 1, 1.5} {
if n > 1e4 && theta < 0.5 {
// Don't run unreasonably long benchmarks.
continue
}
rnd := rand.New(rand.NewSource(1))
particles := make([]Particle3, n)
for i := range particles {
particles[i] = particle3{x: rnd.Float64(), y: rnd.Float64(), z: rnd.Float64(), m: 1}
}
volume := NewVolume(particles)
b.ResetTimer()
b.Run(fmt.Sprintf("%d-body/theta=%v", len(particles), theta), func(b *testing.B) {
for i := 0; i < b.N; i++ {
for _, p := range particles {
fv3sink = volume.ForceOn(p, theta, Gravity3)
}
}
})
}
}
}
func BenchmarkVolumeFull(b *testing.B) {
for _, n := range []int{1e3, 1e4, 1e5} {
for _, theta := range []float64{0, 0.1, 0.5, 1, 1.5} {
if n > 1e4 && theta < 0.5 {
// Don't run unreasonably long benchmarks.
continue
}
rnd := rand.New(rand.NewSource(1))
particles := make([]Particle3, n)
for i := range particles {
particles[i] = particle3{x: rnd.Float64(), y: rnd.Float64(), z: rnd.Float64(), m: 1}
}
b.ResetTimer()
b.Run(fmt.Sprintf("%d-body/theta=%v", len(particles), theta), func(b *testing.B) {
for i := 0; i < b.N; i++ {
volume := NewVolume(particles)
for _, p := range particles {
fv3sink = volume.ForceOn(p, theta, Gravity3)
}
}
})
}
}
}

View File

@@ -0,0 +1,9 @@
// Copyright ©2019 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// +build bounds
package barneshut
const checkBounds = true

10
spatial/barneshut/doc.go Normal file
View File

@@ -0,0 +1,10 @@
// Copyright ©2019 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package barneshut provides routines for calculating n-body force approximations
// using the Barnes-Hut algorithm.
//
// See https://en.wikipedia.org/wiki/BarnesHut_simulation, http://arborjs.org/docs/barnes-hut
// and https://jheer.github.io/barnes-hut/ for details.
package barneshut // import "gonum.org/v1/gonum/spatial/barneshut"

View File

@@ -0,0 +1,78 @@
// Copyright ©2019 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package barneshut_test
import (
"golang.org/x/exp/rand"
"gonum.org/v1/gonum/spatial/barneshut"
"gonum.org/v1/gonum/spatial/r2"
)
type mass struct {
d r2.Vec
v r2.Vec
m float64
}
func (m *mass) Coord2() r2.Vec { return m.d }
func (m *mass) Mass() float64 { return m.m }
func (m *mass) move(f r2.Vec) {
m.v = m.v.Add(f.Scale(1 / m.m))
m.d = m.d.Add(m.v)
}
func Example_galaxy() {
rnd := rand.New(rand.NewSource(1))
// Make 1000 stars in random locations.
stars := make([]*mass, 1000)
p := make([]barneshut.Particle2, len(stars))
for i := range stars {
s := &mass{
d: r2.Vec{
X: 100 * rnd.Float64(),
Y: 100 * rnd.Float64(),
},
v: r2.Vec{
X: rnd.NormFloat64(),
Y: rnd.NormFloat64(),
},
m: 10 * rnd.Float64(),
}
stars[i] = s
p[i] = s
}
vectors := make([]r2.Vec, len(stars))
// Make a plane to calculate approximate forces
plane := barneshut.Plane{Particles: p}
// Run a simulation for 100 updates.
for i := 0; i < 1000; i++ {
// Build the data structure. For small systems
// this step may be omitted and ForceOn will
// perform the naive quadratic calculation
// without building the data structure.
plane.Reset()
// Calculate the force vectors using the theta
// parameter...
const theta = 0.5
// and an imaginary gravitational constant.
const G = 10
for j, s := range stars {
vectors[j] = plane.ForceOn(s, theta, barneshut.Gravity2).Scale(G)
}
// Update positions.
for j, s := range stars {
s.move(vectors[j])
}
// Rendering stars is left as an exercise for
// the reader.
}
}

View File

@@ -0,0 +1,9 @@
// Copyright ©2019 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// +build !bounds
package barneshut
const checkBounds = false

6
spatial/r2/doc.go Normal file
View File

@@ -0,0 +1,6 @@
// Copyright ©2019 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package r2 provides 2D vectors and boxes and operations on them.
package r2 // import "gonum.org/v1/gonum/spatial/r2"

36
spatial/r2/vector.go Normal file
View File

@@ -0,0 +1,36 @@
// Copyright ©2019 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package r2
// Vec is a 2D vector.
type Vec struct {
X, Y float64
}
// Add returns the vector sum of p and q.
func (p Vec) Add(q Vec) Vec {
p.X += q.X
p.Y += q.Y
return p
}
// Sub returns the vector sum of p and -q.
func (p Vec) Sub(q Vec) Vec {
p.X -= q.X
p.Y -= q.Y
return p
}
// Scale returns the vector p scaled by f.
func (p Vec) Scale(f float64) Vec {
p.X *= f
p.Y *= f
return p
}
// Box is a 2D bounding box.
type Box struct {
Min, Max Vec
}

6
spatial/r3/doc.go Normal file
View File

@@ -0,0 +1,6 @@
// Copyright ©2019 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package r3 provides 3D vectors and boxes and operations on them.
package r3 // import "gonum.org/v1/gonum/spatial/r3"

39
spatial/r3/vector.go Normal file
View File

@@ -0,0 +1,39 @@
// Copyright ©2019 The Gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package r3
// Vec is a 3D vector.
type Vec struct {
X, Y, Z float64
}
// Add returns the vector sum of p and q.
func (p Vec) Add(q Vec) Vec {
p.X += q.X
p.Y += q.Y
p.Z += q.Z
return p
}
// Sub returns the vector sum of p and -q.
func (p Vec) Sub(q Vec) Vec {
p.X -= q.X
p.Y -= q.Y
p.Z -= q.Z
return p
}
// Scale returns the vector p scaled by f.
func (p Vec) Scale(f float64) Vec {
p.X *= f
p.Y *= f
p.Z *= f
return p
}
// Box is a 3D bounding box.
type Box struct {
Min, Max Vec
}