From 4a2eb0188cbc0e070b369cf7197a3ecb1034fb46 Mon Sep 17 00:00:00 2001 From: Dan Kortschak Date: Fri, 3 May 2019 06:57:12 +0930 Subject: [PATCH] spatial/barneshut: add Barnes-Hut force approximation package and R2/R3 helpers --- spatial/barneshut/barneshut2.go | 249 +++++++++++ spatial/barneshut/barneshut2_test.go | 524 +++++++++++++++++++++++ spatial/barneshut/barneshut3.go | 300 +++++++++++++ spatial/barneshut/barneshut3_test.go | 521 ++++++++++++++++++++++ spatial/barneshut/bounds.go | 9 + spatial/barneshut/doc.go | 10 + spatial/barneshut/galaxy_example_test.go | 78 ++++ spatial/barneshut/no_bounds.go | 9 + spatial/r2/doc.go | 6 + spatial/r2/vector.go | 36 ++ spatial/r3/doc.go | 6 + spatial/r3/vector.go | 39 ++ 12 files changed, 1787 insertions(+) create mode 100644 spatial/barneshut/barneshut2.go create mode 100644 spatial/barneshut/barneshut2_test.go create mode 100644 spatial/barneshut/barneshut3.go create mode 100644 spatial/barneshut/barneshut3_test.go create mode 100644 spatial/barneshut/bounds.go create mode 100644 spatial/barneshut/doc.go create mode 100644 spatial/barneshut/galaxy_example_test.go create mode 100644 spatial/barneshut/no_bounds.go create mode 100644 spatial/r2/doc.go create mode 100644 spatial/r2/vector.go create mode 100644 spatial/r3/doc.go create mode 100644 spatial/r3/vector.go diff --git a/spatial/barneshut/barneshut2.go b/spatial/barneshut/barneshut2.go new file mode 100644 index 00000000..e229b654 --- /dev/null +++ b/spatial/barneshut/barneshut2.go @@ -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 +} diff --git a/spatial/barneshut/barneshut2_test.go b/spatial/barneshut/barneshut2_test.go new file mode 100644 index 00000000..29f5d654 --- /dev/null +++ b/spatial/barneshut/barneshut2_test.go @@ -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) + } + } + }) + } + } +} diff --git a/spatial/barneshut/barneshut3.go b/spatial/barneshut/barneshut3.go new file mode 100644 index 00000000..c19cf773 --- /dev/null +++ b/spatial/barneshut/barneshut3.go @@ -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 +} diff --git a/spatial/barneshut/barneshut3_test.go b/spatial/barneshut/barneshut3_test.go new file mode 100644 index 00000000..a9d099e4 --- /dev/null +++ b/spatial/barneshut/barneshut3_test.go @@ -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) + } + } + }) + } + } +} diff --git a/spatial/barneshut/bounds.go b/spatial/barneshut/bounds.go new file mode 100644 index 00000000..841807d4 --- /dev/null +++ b/spatial/barneshut/bounds.go @@ -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 diff --git a/spatial/barneshut/doc.go b/spatial/barneshut/doc.go new file mode 100644 index 00000000..14aab74b --- /dev/null +++ b/spatial/barneshut/doc.go @@ -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/Barnes–Hut_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" diff --git a/spatial/barneshut/galaxy_example_test.go b/spatial/barneshut/galaxy_example_test.go new file mode 100644 index 00000000..cc6a6bf9 --- /dev/null +++ b/spatial/barneshut/galaxy_example_test.go @@ -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. + } +} diff --git a/spatial/barneshut/no_bounds.go b/spatial/barneshut/no_bounds.go new file mode 100644 index 00000000..03925dfb --- /dev/null +++ b/spatial/barneshut/no_bounds.go @@ -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 diff --git a/spatial/r2/doc.go b/spatial/r2/doc.go new file mode 100644 index 00000000..c21a5290 --- /dev/null +++ b/spatial/r2/doc.go @@ -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" diff --git a/spatial/r2/vector.go b/spatial/r2/vector.go new file mode 100644 index 00000000..b7b9cd6d --- /dev/null +++ b/spatial/r2/vector.go @@ -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 +} diff --git a/spatial/r3/doc.go b/spatial/r3/doc.go new file mode 100644 index 00000000..e67b5bd7 --- /dev/null +++ b/spatial/r3/doc.go @@ -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" diff --git a/spatial/r3/vector.go b/spatial/r3/vector.go new file mode 100644 index 00000000..39c14c1a --- /dev/null +++ b/spatial/r3/vector.go @@ -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 +}