Last active
December 12, 2024 19:08
-
-
Save nicklanng/0142556a805839a053eecf0a6dc22af8 to your computer and use it in GitHub Desktop.
My first attempt at implementing John Ratcliffe's Sphere Trees in go
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package main | |
import ( | |
"fmt" | |
"github.com/faiface/pixel" | |
"github.com/faiface/pixel/imdraw" | |
"github.com/faiface/pixel/pixelgl" | |
"github.com/faiface/pixel/text" | |
"golang.org/x/image/colornames" | |
"golang.org/x/image/font/basicfont" | |
"image/color" | |
"math" | |
"math/rand" | |
"time" | |
) | |
var ZeroVector3 = Vector3{} | |
// Zero is a utility function that returns the zero value for any type | |
func Zero[T any]() T { | |
return *new(T) | |
} | |
// FreeList is a structure that holds any kind of object. | |
// When an entry is erased, it marks that slot as free so that new entries can be filled in to | |
// the existing allocated memory. | |
type FreeList[T any] struct { | |
data []freeListEntry[T] | |
FirstFree int | |
} | |
type freeListEntry[T any] struct { | |
element T | |
nextFree int | |
} | |
func NewFreeList[T any]() FreeList[T] { | |
return FreeList[T]{ | |
FirstFree: -1, | |
} | |
} | |
func (f *FreeList[T]) Insert(element T) int { | |
// if there is a free entry | |
if f.FirstFree != -1 { | |
// get the index of the free entry | |
index := f.FirstFree | |
// set the next free entry from this index | |
f.FirstFree = f.data[index].nextFree | |
f.data[index].nextFree = 0 | |
// set the data at current free index | |
f.data[index].element = element | |
return index | |
} | |
// insert at the end of the list | |
f.data = append(f.data, freeListEntry[T]{ | |
element: element, | |
nextFree: 0, | |
}) | |
return len(f.data) - 1 | |
} | |
func (f *FreeList[T]) Set(n int, element T) { | |
f.data[n].element = element | |
} | |
func (f *FreeList[T]) Erase(n int) { | |
// set the old next free index to this node | |
f.data[n].nextFree = f.FirstFree | |
// set the first free index to this nodes next free Raw | |
f.FirstFree = n | |
} | |
func (f *FreeList[T]) Clear() { | |
f.data = f.data[:0] | |
f.FirstFree = -1 | |
} | |
func (f *FreeList[T]) Get(n int) T { | |
return f.data[n].element | |
} | |
func (f *FreeList[T]) Len() int { | |
return len(f.data) | |
} | |
// Queue is your standard First In First Out queue. | |
type Queue[T any] struct { | |
data []T | |
} | |
func (q *Queue[T]) Push(e T) { | |
q.data = append(q.data, e) | |
} | |
func (q *Queue[T]) Pop() (T, bool) { | |
if len(q.data) == 0 { | |
return Zero[T](), false | |
} | |
e := q.data[0] | |
q.data = q.data[1:] | |
return e, true | |
} | |
func (q *Queue[T]) Peek(i int) T { | |
return q.data[i] | |
} | |
func (q *Queue[T]) Len() int { | |
return len(q.data) | |
} | |
// Vector3 represents a point or a line in 3d space | |
type Vector3 struct { | |
X float64 | |
Y float64 | |
Z float64 | |
} | |
func (v Vector3) Add(v2 Vector3) Vector3 { | |
v.X += v2.X | |
v.Y += v2.Y | |
v.Z += v2.Z | |
return v | |
} | |
func (v Vector3) Sub(v2 Vector3) Vector3 { | |
v.X -= v2.X | |
v.Y -= v2.Y | |
v.Z -= v2.Z | |
return v | |
} | |
func (v Vector3) Multiply(scalar float64) Vector3 { | |
v.X *= scalar | |
v.Y *= scalar | |
v.Z *= scalar | |
return v | |
} | |
func (v Vector3) Magnitude() float64 { | |
return math.Sqrt(v.X*v.X + v.Y*v.Y + v.Z*v.Z) | |
} | |
func (v Vector3) Normalize() Vector3 { | |
div := v.Magnitude() | |
v.X /= div | |
v.Y /= div | |
v.Z /= div | |
return v | |
} | |
func (v Vector3) Distance(v2 Vector3) float64 { | |
return math.Sqrt((v2.X-v.X)*(v2.X-v.X) + (v2.Y-v.Y)*(v2.Y-v.Y) + (v2.Z-v.Z)*(v2.Z-v.Z)) | |
} | |
// Sphere is literally a sphere in 3d space | |
type Sphere struct { | |
center Vector3 | |
radius float64 | |
} | |
func NewSphere(center Vector3, radius float64) Sphere { | |
return Sphere{ | |
center: center, | |
radius: radius, | |
} | |
} | |
func (s Sphere) IntersectsSphere(s2 Sphere) bool { | |
return s.center.Distance(s2.center) < s.radius+s2.radius | |
} | |
func (s Sphere) ContainsSphere(s2 Sphere) bool { | |
return s.radius >= s.center.Distance(s2.center)+s2.radius | |
} | |
// SphereTree is a space partitioning data structure that contains spheres inside larger super sphere | |
type SphereTree struct { | |
spheres FreeList[SphereEntry] | |
integrateFIFO Queue[int] | |
recomputeFIFO Queue[int] | |
maxSize float64 | |
gravy float64 | |
} | |
type SphereEntry struct { | |
ID uint64 | |
sphere Sphere | |
parent int | |
firstChild int | |
next int | |
} | |
func NewSphereTree(rootSphere Sphere, maxSize, gravy float64) *SphereTree { | |
st := &SphereTree{ | |
spheres: FreeList[SphereEntry]{FirstFree: -1}, | |
maxSize: maxSize, | |
gravy: gravy, | |
} | |
st.spheres.Insert(SphereEntry{sphere: rootSphere, firstChild: -1, next: -1}) | |
return st | |
} | |
func (st *SphereTree) Insert(id uint64, sphere Sphere) int { | |
entry := SphereEntry{ID: id, sphere: sphere, firstChild: -2, next: -1} | |
// insert entry into list | |
i := st.spheres.Insert(entry) | |
// set the entry to be integrated into the tree, don't do it now | |
st.QueueIntegrate(i) | |
return i | |
} | |
func (st *SphereTree) Remove(entryID int) { | |
entry := st.spheres.Get(entryID) | |
st.spheres.Erase(entryID) | |
// remove this entry from the parent's linked list | |
st.removeChild(entry.parent, entryID) | |
} | |
func (st *SphereTree) Move(entryID int, sphere Sphere) { | |
entry := st.spheres.Get(entryID) | |
entry.sphere = sphere | |
st.spheres.Set(entryID, entry) | |
parentID := entry.parent | |
parent := st.spheres.Get(parentID) | |
// parent spheres still contains | |
if parent.sphere.ContainsSphere(entry.sphere) { | |
//st.QueueRecompute(parentID) // TODO: maybe dont | |
return | |
} | |
// entry has broken out of sphere | |
// so detach it and queue it for integration | |
// and recompute its old parent | |
st.removeChild(parentID, entryID) | |
st.QueueIntegrate(entryID) | |
} | |
func (st *SphereTree) Integrate() { | |
for { | |
integrateCandidateID, ok := st.integrateFIFO.Pop() | |
if !ok { | |
break | |
} | |
st.integrate(integrateCandidateID) | |
} | |
} | |
func (st *SphereTree) integrate(entryID int) { | |
integrateCandidate := st.spheres.Get(entryID) | |
containsUs := -1 | |
nearest := -1 | |
nearestDist := math.MaxFloat64 | |
// look through all super spheres for candidate parent | |
// look for super spheres that fully contain the candidate or find the closet one | |
rootSphere := st.spheres.Get(0) // root | |
superSphereIndex := rootSphere.firstChild | |
for { | |
if superSphereIndex < 0 { | |
break | |
} | |
superSphere := st.spheres.Get(superSphereIndex) | |
if superSphere.sphere.ContainsSphere(integrateCandidate.sphere) { | |
containsUs = superSphereIndex | |
break | |
} | |
// TODO: Verify this | |
dist := superSphere.sphere.center.Distance(integrateCandidate.sphere.center) + integrateCandidate.sphere.radius - superSphere.sphere.radius | |
if dist < nearestDist { | |
nearest = superSphereIndex | |
nearestDist = dist | |
} | |
superSphereIndex = superSphere.next | |
} | |
// if a super sphere contains it, just insert it! | |
if containsUs != -1 { | |
st.addChild(containsUs, entryID) | |
return | |
} | |
// check to see if the nearest sphere can grow to contain us | |
if nearest != -1 { | |
parent := st.spheres.Get(nearest) | |
newSize := nearestDist + parent.sphere.radius | |
if newSize <= st.maxSize { | |
parent.sphere.radius = newSize + st.gravy | |
st.spheres.Set(nearest, parent) | |
st.addChild(nearest, entryID) | |
return | |
} | |
} | |
// we'll have to make a new super sphere | |
parent := SphereEntry{ | |
sphere: integrateCandidate.sphere, | |
firstChild: -1, | |
parent: 0, | |
} | |
parent.sphere.radius += st.gravy | |
rootSphere = st.spheres.Get(0) // root | |
parent.next = rootSphere.firstChild | |
parentID := st.spheres.Insert(parent) | |
rootSphere.firstChild = parentID | |
st.spheres.Set(0, rootSphere) | |
st.addChild(parentID, entryID) | |
} | |
func (st *SphereTree) Recompute() { | |
for { | |
recomputeCandidateID, ok := st.recomputeFIFO.Pop() | |
if !ok { | |
break | |
} | |
st.recompute(recomputeCandidateID) | |
} | |
} | |
func (st *SphereTree) recompute(superSphereID int) { | |
superSphere := st.spheres.Get(superSphereID) | |
if superSphere.firstChild == -1 { | |
st.removeChild(0, superSphereID) | |
st.spheres.Erase(superSphereID) | |
return | |
} | |
var childCount int | |
var total Vector3 | |
childIndex := superSphere.firstChild | |
for { | |
if childIndex == -1 { | |
break | |
} | |
childCount++ | |
child := st.spheres.Get(childIndex) | |
total = total.Add(child.sphere.center) | |
childIndex = child.next | |
} | |
recip := 1.0 / float64(childCount) | |
oldCenter := superSphere.sphere.center | |
superSphere.sphere.center = total.Multiply(recip) | |
newRadius := 0.0 | |
childIndex = superSphere.firstChild | |
for { | |
if childIndex == -1 { | |
break | |
} | |
child := st.spheres.Get(childIndex) | |
radius := superSphere.sphere.center.Distance(child.sphere.center) + child.sphere.radius | |
if radius > newRadius { | |
newRadius = radius | |
if newRadius+st.gravy > superSphere.sphere.radius { | |
superSphere.sphere.center = oldCenter | |
return | |
} | |
} | |
childIndex = child.next | |
} | |
superSphere.sphere.radius = newRadius + st.gravy | |
st.spheres.Set(superSphereID, superSphere) | |
// RINSE: look for circle that can own this one | |
root := st.spheres.Get(0) | |
possibleParentID := root.firstChild | |
for { | |
if possibleParentID == 0 { | |
panic("no") | |
} | |
if possibleParentID == -1 { | |
break | |
} | |
if possibleParentID == superSphereID { | |
possibleParentID = superSphere.next | |
continue | |
} | |
possibleParent := st.spheres.Get(possibleParentID) | |
if possibleParent.sphere.ContainsSphere(superSphere.sphere) { | |
// find last child of possible parent | |
childIndex := superSphere.firstChild | |
for { | |
if childIndex == -1 { | |
break | |
} | |
child := st.spheres.Get(childIndex) | |
next := child.next | |
child.next = possibleParent.firstChild | |
child.parent = possibleParentID | |
possibleParent.firstChild = childIndex | |
st.spheres.Set(childIndex, child) | |
childIndex = next | |
} | |
superSphere.firstChild = -1 | |
st.spheres.Set(superSphereID, superSphere) | |
st.spheres.Set(possibleParentID, possibleParent) | |
st.QueueRecompute(superSphereID) | |
st.QueueRecompute(possibleParentID) | |
break | |
} | |
possibleParentID = possibleParent.next | |
} | |
} | |
func (st *SphereTree) Walk(f func(s Sphere, isSuper bool)) { | |
root := st.spheres.Get(0) | |
childIndex := root.firstChild | |
for { | |
if childIndex == -1 { | |
break | |
} | |
child := st.spheres.Get(childIndex) | |
entryIndex := child.firstChild | |
for { | |
if entryIndex == -1 { | |
break | |
} | |
entry := st.spheres.Get(entryIndex) | |
f(entry.sphere, false) | |
entryIndex = entry.next | |
} | |
childIndex = child.next | |
} | |
childIndex = root.firstChild | |
for { | |
if childIndex == -1 { | |
break | |
} | |
child := st.spheres.Get(childIndex) | |
f(child.sphere, true) | |
childIndex = child.next | |
} | |
} | |
func (st *SphereTree) QueueIntegrate(entryID int) { | |
for i := 0; i < st.integrateFIFO.Len(); i++ { | |
if st.integrateFIFO.Peek(i) == entryID { | |
return | |
} | |
} | |
st.integrateFIFO.Push(entryID) | |
} | |
func (st *SphereTree) QueueRecompute(parentID int) { | |
if parentID == 0 { | |
return | |
} | |
for i := 0; i < st.recomputeFIFO.Len(); i++ { | |
if st.recomputeFIFO.Peek(i) == parentID { | |
return | |
} | |
} | |
st.recomputeFIFO.Push(parentID) | |
} | |
func (st *SphereTree) GetEntries(selectionSphere Sphere) []SphereEntry { | |
var results []SphereEntry | |
root := st.spheres.Get(0) | |
sphereIndex := root.firstChild | |
for { | |
if sphereIndex == -1 { | |
break | |
} | |
sphere := st.spheres.Get(sphereIndex) | |
if selectionSphere.IntersectsSphere(sphere.sphere) { | |
childIndex := sphere.firstChild | |
for { | |
if childIndex == -1 { | |
break | |
} | |
child := st.spheres.Get(childIndex) | |
if selectionSphere.IntersectsSphere(child.sphere) { | |
results = append(results, child) | |
} | |
childIndex = child.next | |
} | |
} | |
sphereIndex = sphere.next | |
} | |
return results | |
} | |
func (st *SphereTree) addChild(parentID, sphereID int) { | |
parent := st.spheres.Get(parentID) | |
entry := st.spheres.Get(sphereID) | |
entry.parent = parentID | |
entry.next = parent.firstChild | |
parent.firstChild = sphereID | |
st.spheres.Set(parentID, parent) | |
st.spheres.Set(sphereID, entry) | |
st.QueueRecompute(parentID) | |
} | |
func (st *SphereTree) removeChild(parentID, entryID int) { | |
parent := st.spheres.Get(parentID) | |
entry := st.spheres.Get(entryID) | |
childIndex := parent.firstChild | |
if childIndex == entryID { | |
parent.firstChild = entry.next | |
} else { | |
for { | |
if childIndex == -1 { | |
break | |
} | |
child := st.spheres.Get(childIndex) | |
if child.next == entryID { | |
child.next = entry.next | |
st.spheres.Set(childIndex, child) | |
break | |
} | |
childIndex = child.next | |
} | |
} | |
st.spheres.Set(parentID, parent) | |
st.spheres.Set(entryID, entry) | |
st.QueueRecompute(parentID) | |
} | |
type actor struct { | |
id uint64 | |
sphere Sphere | |
entryID int | |
attractor Vector3 | |
speed float64 | |
color color.Color | |
} | |
func main() { | |
pixelgl.Run(run) | |
} | |
func run() { | |
rand.Seed(time.Now().UnixMilli()) | |
cfg := pixelgl.WindowConfig{ | |
Title: "Pixel Rocks!", | |
Bounds: pixel.R(0, 0, 1920, 1080), | |
VSync: true, | |
Resizable: true, | |
Maximized: false, | |
} | |
win, err := pixelgl.NewWindow(cfg) | |
if err != nil { | |
panic(err) | |
} | |
atlas := text.NewAtlas(basicfont.Face7x13, text.ASCII) | |
moustTxt := text.New(pixel.V(50, win.Bounds().H()-50), atlas) | |
var actors []actor | |
var attractors []Vector3 | |
for i := 0; i < 19; i++ { | |
attractors = append(attractors, Vector3{X: rand.Float64() * win.Bounds().W(), Y: rand.Float64() * win.Bounds().H()}) | |
} | |
st := NewSphereTree(Sphere{center: Vector3{X: win.Bounds().W(), Y: win.Bounds().H() / 2.0}, radius: 500}, 100, 10) | |
imd := imdraw.New(nil) | |
for !win.Closed() { | |
// updates | |
if len(actors) < 4000 { | |
for i := 0; i < 100; i++ { | |
sphere := Sphere{center: Vector3{X: rand.Float64() * win.Bounds().W(), Y: rand.Float64() * win.Bounds().H()}, radius: 1.0} | |
entryID := st.Insert(uint64(len(actors)), sphere) | |
var attractor Vector3 | |
c := colornames.Deepskyblue | |
if len(actors) > 2000 { | |
c = colornames.Red | |
if i%3 == 0 { | |
c = colornames.Purple | |
} | |
if i%3 == 1 { | |
c = colornames.Hotpink | |
} | |
attractor = attractors[rand.Intn(len(attractors))] | |
} | |
actors = append(actors, actor{ | |
id: uint64(len(actors)), | |
sphere: sphere, | |
entryID: entryID, | |
attractor: attractor, | |
speed: 1, | |
color: c, | |
}) | |
} | |
} else { | |
for i := 0; i < len(actors); i++ { | |
if actors[i].attractor == ZeroVector3 { | |
continue | |
} | |
delta := actors[i].attractor.Sub(actors[i].sphere.center) | |
if delta.Magnitude() < 1.0 { | |
actors[i].sphere.center = Vector3{X: rand.Float64() * win.Bounds().W(), Y: rand.Float64() * win.Bounds().H()} | |
delta = actors[i].attractor.Sub(actors[i].sphere.center) | |
} | |
actors[i].sphere.center = actors[i].sphere.center.Add(delta.Normalize().Multiply(actors[i].speed)) | |
st.Move(actors[i].entryID, actors[i].sphere) | |
} | |
} | |
st.Integrate() | |
st.Recompute() | |
// renders | |
imd.Clear() | |
for i := range actors { | |
imd.Color = actors[i].color | |
imd.Push(pixel.V(actors[i].sphere.center.X, actors[i].sphere.center.Y)) | |
imd.Circle(actors[i].sphere.radius, 1) | |
} | |
imd.Color = colornames.Navy | |
st.Walk(func(s Sphere, isSuper bool) { | |
if !isSuper { | |
return | |
} | |
imd.Push(pixel.V(s.center.X, s.center.Y)) | |
imd.Circle(s.radius, 1) | |
}) | |
for _, a := range attractors { | |
imd.Color = colornames.Limegreen | |
imd.Push(pixel.V(a.X, a.Y)) | |
imd.Circle(4, 2) | |
} | |
imd.Color = colornames.Yellowgreen | |
//now := time.Now() | |
entries := st.GetEntries(Sphere{ | |
center: Vector3{X: win.MousePosition().X, Y: win.MousePosition().Y}, | |
radius: 40, | |
}) | |
//fmt.Println(time.Since(now)) | |
for i := range entries { | |
imd.Push(pixel.V(entries[i].sphere.center.X, entries[i].sphere.center.Y)) | |
imd.Circle(entries[i].sphere.radius, 1) | |
} | |
imd.Push(win.MousePosition()) | |
imd.Circle(40, 1) | |
moustTxt.Clear() | |
fmt.Fprintf(moustTxt, "Selected: %d", len(entries)) | |
win.Clear(color.Black) | |
imd.Draw(win) | |
moustTxt.Draw(win, pixel.IM) | |
win.Update() | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment