mirror
/
rtred
mirror of https://github.com/tidwall/rtred.git
1
0
Fork 1
Code Issues Packages Projects Releases Wiki Activity

total recoding, kNN support

This commit is contained in:
Josh Baker 2018-01-12 16:31:15 -07:00
parent d4a8a3d30d
commit 75265604a7
11 changed files with 1697 additions and 14840 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
README.md
View File

@ -14,6 +14,7 @@ Authors
* 1995 Sphere volume fix for degeneracy problem submitted by Paul Brook
* 2004 Templated C++ port by Greg Douglas
* 2016 Go port by Josh Baker
* 2018 Added kNN and merged in some of the RBush logic by Vladimir Agafonkin
License
-------

98
base/knn.go Normal file
View File

@ -0,0 +1,98 @@
package base
import (
"github.com/tidwall/tinyqueue"
)
type queueItem struct {
node *treeNode
isItem bool
dist float64
}
func (item *queueItem) Less(b tinyqueue.Item) bool {
return item.dist < b.(*queueItem).dist
}
// KNN returns items nearest to farthest. The dist param is the "box distance".
func (tr *RTree) KNN(min, max []float64, center bool, iter func(item interface{}, dist float64) bool) bool {
var isBox bool
knnPoint := make([]float64, tr.dims)
bbox := &treeNode{min: min, max: max}
for i := 0; i < tr.dims; i++ {
knnPoint[i] = (bbox.min[i] + bbox.max[i]) / 2
if !isBox && bbox.min[i] != bbox.max[i] {
isBox = true
}
}
node := tr.data
queue := tinyqueue.New(nil)
for node != nil {
for i := 0; i < node.count; i++ {
child := node.children[i]
var dist float64
if isBox {
dist = boxDistRect(bbox, child)
} else {
dist = boxDistPoint(knnPoint, child)
}
queue.Push(&queueItem{node: child, isItem: node.leaf, dist: dist})
}
for queue.Len() > 0 && queue.Peek().(*queueItem).isItem {
item := queue.Pop().(*queueItem)
if !iter(item.node.unsafeItem().item, item.dist) {
return false
}
}
last := queue.Pop()
if last != nil {
node = (*treeNode)(last.(*queueItem).node)
} else {
node = nil
}
}
return true
}
func boxDistRect(a, b *treeNode) float64 {
var dist float64
for i := 0; i < len(a.min); i++ {
var min, max float64
if a.min[i] > b.min[i] {
min = a.min[i]
} else {
min = b.min[i]
}
if a.max[i] < b.max[i] {
max = a.max[i]
} else {
max = b.max[i]
}
squared := min - max
if squared > 0 {
dist += squared * squared
}
}
return dist
}
func boxDistPoint(point []float64, childBox *treeNode) float64 {
var dist float64
for i := 0; i < len(point); i++ {
d := axisDist(point[i], childBox.min[i], childBox.max[i])
dist += d * d
}
return dist
}
func axisDist(k, min, max float64) float64 {
if k < min {
return min - k
}
if k <= max {
return 0
}
return k - max
}

97
base/load.go Normal file
View File

@ -0,0 +1,97 @@
package base
import "math"
// Load bulk load items into the R-tree.
func (tr *RTree) Load(mins, maxs [][]float64, items []interface{}) {
if len(items) < tr.minEntries {
for i := 0; i < len(items); i++ {
tr.Insert(mins[i], maxs[i], items[i])
}
return
}
// prefill the items
fitems := make([]*treeNode, len(items))
for i := 0; i < len(items); i++ {
item := &treeItem{min: mins[i], max: maxs[i], item: items[i]}
fitems[i] = item.unsafeNode()
}
// following equations are defined in the paper describing OMT
N := len(fitems)
M := tr.maxEntries
h := int(math.Ceil(math.Log(float64(N)) / math.Log(float64(M))))
Nsubtree := int(math.Pow(float64(M), float64(h-1)))
S := int(math.Ceil(math.Sqrt(float64(N) / float64(Nsubtree))))
// sort by the initial axis
axis := 0
sortByAxis(fitems, axis)
// build the root node. it's split differently from the subtrees.
children := make([]*treeNode, 0, S)
for i := 0; i < S; i++ {
var part []*treeNode
if i == S-1 {
// last split
part = fitems[len(fitems)/S*i:]
} else {
part = fitems[len(fitems)/S*i : len(fitems)/S*(i+1)]
}
children = append(children, tr.omt(part, h-1, axis+1))
}
node := tr.createNode(children)
node.leaf = false
node.height = h
tr.calcBBox(node)
if tr.data.count == 0 {
// save as is if tree is empty
tr.data = node
} else if tr.data.height == node.height {
// split root if trees have the same height
tr.splitRoot(tr.data, node)
} else {
if tr.data.height < node.height {
// swap trees if inserted one is bigger
tr.data, node = node, tr.data
}
// insert the small tree into the large tree at appropriate level
tr.insert(node, nil, tr.data.height-node.height-1, true)
}
}
func (tr *RTree) omt(fitems []*treeNode, h, axis int) *treeNode {
if len(fitems) <= tr.maxEntries {
// reached leaf level; return leaf
children := make([]*treeNode, len(fitems))
copy(children, fitems)
node := tr.createNode(children)
node.height = h
tr.calcBBox(node)
return node
}
// sort the items on a different axis than the previous level.
sortByAxis(fitems, axis%tr.dims)
children := make([]*treeNode, 0, tr.maxEntries)
partsz := len(fitems) / tr.maxEntries
for i := 0; i < tr.maxEntries; i++ {
var part []*treeNode
if i == tr.maxEntries-1 {
// last part
part = fitems[partsz*i:]
} else {
part = fitems[partsz*i : partsz*(i+1)]
}
children = append(children, tr.omt(part, h-1, axis+1))
}
node := tr.createNode(children)
node.height = h
node.leaf = false
tr.calcBBox(node)
return node
}

673
base/rtree.go Normal file
View File

@ -0,0 +1,673 @@
package base
import (
"math"
"unsafe"
)
// precalculate infinity
var mathInfNeg = math.Inf(-1)
var mathInfPos = math.Inf(+1)
type treeNode struct {
min, max []float64
children []*treeNode
count int
height int
leaf bool
}
func (node *treeNode) unsafeItem() *treeItem {
return (*treeItem)(unsafe.Pointer(node))
}
func (tr *RTree) createNode(children []*treeNode) *treeNode {
n := &treeNode{
height: 1,
leaf: true,
children: make([]*treeNode, tr.maxEntries+1),
}
if len(children) > 0 {
n.count = len(children)
copy(n.children[:n.count], children)
}
n.min = make([]float64, tr.dims)
n.max = make([]float64, tr.dims)
for i := 0; i < tr.dims; i++ {
n.min[i] = mathInfPos
n.max[i] = mathInfNeg
}
return n
}
func (node *treeNode) extend(b *treeNode) {
for i := 0; i < len(node.min); i++ {
if b.min[i] < node.min[i] {
node.min[i] = b.min[i]
}
if b.max[i] > node.max[i] {
node.max[i] = b.max[i]
}
}
}
func (node *treeNode) area() float64 {
area := node.max[0] - node.min[0]
for i := 1; i < len(node.min); i++ {
area *= node.max[i] - node.min[i]
}
return area
}
func (node *treeNode) enlargedAreaAxis(b *treeNode, axis int) float64 {
var max, min float64
if b.max[axis] > node.max[axis] {
max = b.max[axis]
} else {
max = node.max[axis]
}
if b.min[axis] < node.min[axis] {
min = b.min[axis]
} else {
min = node.min[axis]
}
return max - min
}
func (node *treeNode) enlargedArea(b *treeNode) float64 {
area := node.enlargedAreaAxis(b, 0)
for i := 1; i < len(node.min); i++ {
area *= node.enlargedAreaAxis(b, i)
}
return area
}
func (node *treeNode) intersectionAreaAxis(b *treeNode, axis int) float64 {
var max, min float64
if node.max[axis] < b.max[axis] {
max = node.max[axis]
} else {
max = b.max[axis]
}
if node.min[axis] > b.min[axis] {
min = node.min[axis]
} else {
min = b.min[axis]
}
if max > min {
return max - min
}
return 0
}
func (node *treeNode) intersectionArea(b *treeNode) float64 {
area := node.intersectionAreaAxis(b, 0)
for i := 1; i < len(node.min); i++ {
area *= node.intersectionAreaAxis(b, i)
}
return area
}
func (node *treeNode) margin() float64 {
margin := node.max[0] - node.min[0]
for i := 1; i < len(node.min); i++ {
margin += node.max[i] - node.min[i]
}
return margin
}
type result int
const (
not result = 0
intersects result = 1
contains result = 2
)
func (node *treeNode) overlaps(b *treeNode) result {
for i := 0; i < len(node.min); i++ {
if b.min[i] > node.max[i] || b.max[i] < node.min[i] {
return not
}
if node.min[i] > b.min[i] || b.max[i] > node.max[i] {
i++
for ; i < len(node.min); i++ {
if b.min[i] > node.max[i] || b.max[i] < node.min[i] {
return not
}
}
return intersects
}
}
return contains
}
func (node *treeNode) intersects(b *treeNode) bool {
for i := 0; i < len(node.min); i++ {
if b.min[i] > node.max[i] || b.max[i] < node.min[i] {
return false
}
}
return true
}
func (node *treeNode) findItem(item interface{}) int {
for i := 0; i < node.count; i++ {
if node.children[i].unsafeItem().item == item {
return i
}
}
return -1
}
func (node *treeNode) contains(b *treeNode) bool {
for i := 0; i < len(node.min); i++ {
if node.min[i] > b.min[i] || b.max[i] > node.max[i] {
return false
}
}
return true
}
func (node *treeNode) childCount() int {
if node.leaf {
return node.count
}
var n int
for i := 0; i < node.count; i++ {
n += node.children[i].childCount()
}
return n
}
type treeItem struct {
min, max []float64
item interface{}
}
func (item *treeItem) unsafeNode() *treeNode {
return (*treeNode)(unsafe.Pointer(item))
}
// RTree is an R-tree
type RTree struct {
dims int
maxEntries int
minEntries int
data *treeNode // root node
// resusable fields, these help performance of common mutable operations.
reuse struct {
path []*treeNode // for reinsertion path
indexes []int // for remove function
stack []int // for bulk loading
}
}
// New creates a new R-tree
func New(dims, maxEntries int) *RTree {
if dims <= 0 {
panic("invalid dimensions")
}
tr := &RTree{}
tr.dims = dims
tr.maxEntries = int(math.Max(4, float64(maxEntries)))
tr.minEntries = int(math.Max(2, math.Ceil(float64(tr.maxEntries)*0.4)))
tr.data = tr.createNode(nil)
return tr
}
// Insert inserts an item
func (tr *RTree) Insert(min, max []float64, item interface{}) {
if len(min) != tr.dims || len(max) != tr.dims {
panic("invalid dimensions")
}
if item == nil {
panic("nil item")
}
bbox := treeNode{min: min, max: max}
tr.insert(&bbox, item, tr.data.height-1, false)
}
func (tr *RTree) insert(bbox *treeNode, item interface{}, level int, isNode bool) {
tr.reuse.path = tr.reuse.path[:0]
node, insertPath := tr.chooseSubtree(bbox, tr.data, level, tr.reuse.path)
if item == nil {
// item is only nil when bulk loading a node
if node.leaf {
panic("loading node into leaf")
}
node.children[node.count] = bbox
node.count++
} else {
ti := &treeItem{min: bbox.min, max: bbox.max, item: item}
node.children[node.count] = ti.unsafeNode()
node.count++
}
node.extend(bbox)
for level >= 0 {
if insertPath[level].count > tr.maxEntries {
insertPath = tr.split(insertPath, level)
level--
} else {
break
}
}
tr.adjustParentBBoxes(bbox, insertPath, level)
tr.reuse.path = insertPath
}
func (tr *RTree) adjustParentBBoxes(bbox *treeNode, path []*treeNode, level int) {
// adjust bboxes along the given tree path
for i := level; i >= 0; i-- {
path[i].extend(bbox)
}
}
func (tr *RTree) chooseSubtree(bbox, node *treeNode, level int, path []*treeNode) (*treeNode, []*treeNode) {
var targetNode *treeNode
var area, enlargement, minArea, minEnlargement float64
for {
path = append(path, node)
if node.leaf || len(path)-1 == level {
break
}
minEnlargement = mathInfPos
minArea = minEnlargement
for i := 0; i < node.count; i++ {
child := node.children[i]
area = child.area()
enlargement = bbox.enlargedArea(child) - area
if enlargement < minEnlargement {
minEnlargement = enlargement
if area < minArea {
minArea = area
}
targetNode = child
} else if enlargement == minEnlargement {
if area < minArea {
minArea = area
targetNode = child
}
}
}
if targetNode != nil {
node = targetNode
} else if node.count > 0 {
node = (*treeNode)(node.children[0])
} else {
node = nil
}
}
return node, path
}
func (tr *RTree) split(insertPath []*treeNode, level int) []*treeNode {
var node = insertPath[level]
var M = node.count
var m = tr.minEntries
tr.chooseSplitAxis(node, m, M)
splitIndex := tr.chooseSplitIndex(node, m, M)
spliced := make([]*treeNode, node.count-splitIndex)
copy(spliced, node.children[splitIndex:])
node.count = splitIndex
newNode := tr.createNode(spliced)
newNode.height = node.height
newNode.leaf = node.leaf
tr.calcBBox(node)
tr.calcBBox(newNode)
if level != 0 {
insertPath[level-1].children[insertPath[level-1].count] = newNode
insertPath[level-1].count++
} else {
tr.splitRoot(node, newNode)
}
return insertPath
}
func (tr *RTree) chooseSplitIndex(node *treeNode, m, M int) int {
var i int
var bbox1, bbox2 *treeNode
var overlap, area, minOverlap, minArea float64
var index int
minArea = mathInfPos
minOverlap = minArea
for i = m; i <= M-m; i++ {
bbox1 = tr.distBBox(node, 0, i, nil)
bbox2 = tr.distBBox(node, i, M, nil)
overlap = bbox1.intersectionArea(bbox2)
area = bbox1.area() + bbox2.area()
// choose distribution with minimum overlap
if overlap < minOverlap {
minOverlap = overlap
index = i
if area < minArea {
minArea = area
}
} else if overlap == minOverlap {
// otherwise choose distribution with minimum area
if area < minArea {
minArea = area
index = i
}
}
}
return index
}
func (tr *RTree) calcBBox(node *treeNode) {
tr.distBBox(node, 0, node.count, node)
}
func (tr *RTree) chooseSplitAxis(node *treeNode, m, M int) {
minMargin := tr.allDistMargin(node, m, M, 0)
var minAxis int
for axis := 1; axis < tr.dims; axis++ {
margin := tr.allDistMargin(node, m, M, axis)
if margin < minMargin {
minMargin = margin
minAxis = axis
}
}
if minAxis < tr.dims {
tr.sortNodes(node, minAxis)
}
}
func (tr *RTree) splitRoot(node, newNode *treeNode) {
tr.data = tr.createNode([]*treeNode{node, newNode})
tr.data.height = node.height + 1
tr.data.leaf = false
tr.calcBBox(tr.data)
}
func (tr *RTree) distBBox(node *treeNode, k, p int, destNode *treeNode) *treeNode {
if destNode == nil {
destNode = tr.createNode(nil)
} else {
for i := 0; i < tr.dims; i++ {
destNode.min[i] = mathInfPos
destNode.max[i] = mathInfNeg
}
}
for i := k; i < p; i++ {
if node.leaf {
destNode.extend(node.children[i])
} else {
destNode.extend((*treeNode)(node.children[i]))
}
}
return destNode
}
func (tr *RTree) allDistMargin(node *treeNode, m, M int, axis int) float64 {
tr.sortNodes(node, axis)
var leftBBox = tr.distBBox(node, 0, m, nil)
var rightBBox = tr.distBBox(node, M-m, M, nil)
var margin = leftBBox.margin() + rightBBox.margin()
var i int
if node.leaf {
for i = m; i < M-m; i++ {
leftBBox.extend(node.children[i])
margin += leftBBox.margin()
}
for i = M - m - 1; i >= m; i-- {
leftBBox.extend(node.children[i])
margin += rightBBox.margin()
}
} else {
for i = m; i < M-m; i++ {
child := (*treeNode)(node.children[i])
leftBBox.extend(child)
margin += leftBBox.margin()
}
for i = M - m - 1; i >= m; i-- {
child := (*treeNode)(node.children[i])
leftBBox.extend(child)
margin += rightBBox.margin()
}
}
return margin
}
func (tr *RTree) sortNodes(node *treeNode, axis int) {
sortByAxis(node.children[:node.count], axis)
}
func sortByAxis(items []*treeNode, axis int) {
if len(items) < 2 {
return
}
left, right := 0, len(items)-1
pivotIndex := len(items) / 2
items[pivotIndex], items[right] = items[right], items[pivotIndex]
for i := range items {
if items[i].min[axis] < items[right].min[axis] {
items[i], items[left] = items[left], items[i]
left++
}
}
items[left], items[right] = items[right], items[left]
sortByAxis(items[:left], axis)
sortByAxis(items[left+1:], axis)
}
// Search searches the tree for items in the input rectangle
func (tr *RTree) Search(min, max []float64, iter func(item interface{}) bool) bool {
bbox := &treeNode{min: min, max: max}
if !tr.data.intersects(bbox) {
return true
}
return tr.search(tr.data, bbox, iter)
}
func (tr *RTree) search(node, bbox *treeNode, iter func(item interface{}) bool) bool {
if node.leaf {
for i := 0; i < node.count; i++ {
if bbox.intersects(node.children[i]) {
if !iter(node.children[i].unsafeItem().item) {
return false
}
}
}
} else {
for i := 0; i < node.count; i++ {
r := bbox.overlaps(node.children[i])
if r == intersects {
if !tr.search(node.children[i], bbox, iter) {
return false
}
} else if r == contains {
if !scan(node.children[i], iter) {
return false
}
}
}
}
return true
}
func (tr *RTree) IsEmpty() bool {
empty := true
tr.Scan(func(item interface{}) bool {
empty = false
return false
})
return empty
}
// Remove removes an item from the R-tree.
func (tr *RTree) Remove(min, max []float64, item interface{}) {
bbox := &treeNode{min: min, max: max}
tr.remove(bbox, item)
}
func (tr *RTree) remove(bbox *treeNode, item interface{}) {
path := tr.reuse.path[:0]
indexes := tr.reuse.indexes[:0]
var node = tr.data
var i int
var parent *treeNode
var index int
var goingUp bool
for node != nil || len(path) != 0 {
if node == nil {
node = path[len(path)-1]
path = path[:len(path)-1]
if len(path) == 0 {
parent = nil
} else {
parent = path[len(path)-1]
}
i = indexes[len(indexes)-1]
indexes = indexes[:len(indexes)-1]
goingUp = true
}
if node.leaf {
index = node.findItem(item)
if index != -1 {
// item found, remove the item and condense tree upwards
copy(node.children[index:], node.children[index+1:])
node.children[node.count-1] = nil
node.count--
path = append(path, node)
tr.condense(path)
goto done
}
}
if !goingUp && !node.leaf && node.contains(bbox) { // go down
path = append(path, node)
indexes = append(indexes, i)
i = 0
parent = node
node = (*treeNode)(node.children[0])
} else if parent != nil { // go right
i++
if i == parent.count {
node = nil
} else {
node = (*treeNode)(parent.children[i])
}
goingUp = false
} else {
node = nil
}
}
done:
tr.reuse.path = path
tr.reuse.indexes = indexes
return
}
func (tr *RTree) condense(path []*treeNode) {
// go through the path, removing empty nodes and updating bboxes
var siblings []*treeNode
for i := len(path) - 1; i >= 0; i-- {
if path[i].count == 0 {
if i > 0 {
siblings = path[i-1].children[:path[i-1].count]
index := -1
for j := 0; j < len(siblings); j++ {
if siblings[j] == path[i] {
index = j
break
}
}
copy(siblings[index:], siblings[index+1:])
siblings[len(siblings)-1] = nil
path[i-1].count--
//siblings = siblings[:len(siblings)-1]
//path[i-1].children = siblings
} else {
tr.data = tr.createNode(nil) // clear tree
}
} else {
tr.calcBBox(path[i])
}
}
}
// Count returns the number of items in the R-tree.
func (tr *RTree) Count() int {
return tr.data.childCount()
}
// Traverse iterates over the entire R-tree and includes all nodes and items.
func (tr *RTree) Traverse(iter func(min, max []float64, level int, item interface{}) bool) bool {
return tr.traverse(tr.data, iter)
}
func (tr *RTree) traverse(node *treeNode, iter func(min, max []float64, level int, item interface{}) bool) bool {
if !iter(node.min, node.max, int(node.height), nil) {
return false
}
if node.leaf {
for i := 0; i < node.count; i++ {
child := node.children[i]
if !iter(child.min, child.max, 0, child.unsafeItem().item) {
return false
}
}
} else {
for i := 0; i < node.count; i++ {
child := node.children[i]
if !tr.traverse(child, iter) {
return false
}
}
}
return true
}
// Scan iterates over the entire R-tree
func (tr *RTree) Scan(iter func(item interface{}) bool) bool {
return scan(tr.data, iter)
}
func scan(node *treeNode, iter func(item interface{}) bool) bool {
if node.leaf {
for i := 0; i < node.count; i++ {
child := node.children[i]
if !iter(child.unsafeItem().item) {
return false
}
}
} else {
for i := 0; i < node.count; i++ {
child := node.children[i]
if !scan(child, iter) {
return false
}
}
}
return true
}
// Bounds returns the bounding box of the entire R-tree
func (tr *RTree) Bounds() (min, max []float64) {
if tr.data.count > 0 {
return tr.data.min, tr.data.max
}
return make([]float64, tr.dims), make([]float64, tr.dims)
}
// Complexity returns the complexity of the R-tree. The higher the value, the
// more complex the tree. The value of 1 is the lowest.
func (tr *RTree) Complexity() float64 {
var nodeCount int
var itemCount int
tr.Traverse(func(_, _ []float64, level int, _ interface{}) bool {
if level == 0 {
itemCount++
} else {
nodeCount++
}
return true
})
return float64(tr.maxEntries*nodeCount) / float64(itemCount)
}

584
base/rtree_test.go Normal file
View File

@ -0,0 +1,584 @@
package base
import (
"fmt"
"log"
"math"
"math/rand"
"reflect"
"runtime"
"sort"
"testing"
"time"
)
const D = 2
const M = 13
type Rect struct {
min, max []float64
item interface{}
}
func (r *Rect) equals(r2 Rect) bool {
if len(r.min) != len(r2.min) || len(r.max) != len(r2.max) || r.item != r2.item {
return false
}
for i := 0; i < len(r.min); i++ {
if r.min[i] != r2.min[i] {
return false
}
}
for i := 0; i < len(r.max); i++ {
if r.max[i] != r2.max[i] {
return false
}
}
return true
}
func ptrMakePoint(vals ...float64) *Rect {
var r Rect
r.min = make([]float64, D)
r.max = make([]float64, D)
for i := 0; i < D && i < len(vals); i++ {
r.min[i] = vals[i]
r.max[i] = vals[i]
}
r.item = &r
return &r
}
func ptrMakeRect(vals ...float64) *Rect {
var r Rect
r.min = make([]float64, D)
r.max = make([]float64, D)
for i := 0; i < D && i < len(vals); i++ {
r.min[i] = vals[i]
r.max[i] = vals[i+D]
}
r.item = &r
return &r
}
func TestRTree(t *testing.T) {
tr := New(D, M)
p := ptrMakePoint(10, 10)
tr.Insert(p.min, p.max, p.item)
}
func TestPtrBasic2D(t *testing.T) {
if D != 2 {
return
}
tr := New(D, M)
p1 := ptrMakePoint(-115, 33)
p2 := ptrMakePoint(-113, 35)
tr.Insert(p1.min, p1.max, p1.item)
tr.Insert(p2.min, p2.max, p2.item)
assertEqual(t, 2, tr.Count())
var points []*Rect
bbox := ptrMakeRect(-116, 32, -114, 34)
tr.Search(bbox.min, bbox.max, func(item interface{}) bool {
points = append(points, item.(*Rect))
return true
})
assertEqual(t, 1, len(points))
tr.Remove(p1.min, p1.max, p1.item)
assertEqual(t, 1, tr.Count())
points = nil
bbox = ptrMakeRect(-116, 33, -114, 34)
tr.Search(bbox.min, bbox.max, func(item interface{}) bool {
points = append(points, item.(*Rect))
return true
})
assertEqual(t, 0, len(points))
tr.Remove(p2.min, p2.max, p2.item)
assertEqual(t, 0, tr.Count())
}
func getMemStats() runtime.MemStats {
runtime.GC()
time.Sleep(time.Millisecond)
runtime.GC()
var ms runtime.MemStats
runtime.ReadMemStats(&ms)
return ms
}
func ptrMakeRandom(what string) *Rect {
if what == "point" {
vals := make([]float64, D)
for i := 0; i < D; i++ {
if i == 0 {
vals[i] = rand.Float64()*360 - 180
} else if i == 1 {
vals[i] = rand.Float64()*180 - 90
} else {
vals[i] = rand.Float64()*100 - 50
}
}
return ptrMakePoint(vals...)
} else if what == "rect" {
vals := make([]float64, D)
for i := 0; i < D; i++ {
if i == 0 {
vals[i] = rand.Float64()*340 - 170
} else if i == 1 {
vals[i] = rand.Float64()*160 - 80
} else {
vals[i] = rand.Float64()*80 - 30
}
}
rvals := make([]float64, D*2)
for i := 0; i < D; i++ {
rvals[i] = vals[i] - rand.Float64()*10
rvals[D+i] = vals[i] + rand.Float64()*10
}
return ptrMakeRect(rvals...)
}
panic("??")
}
func TestPtrRandom(t *testing.T) {
t.Run(fmt.Sprintf("%dD", D), func(t *testing.T) {
t.Run("point", func(t *testing.T) { ptrTestRandom(t, "point", 10000) })
t.Run("rect", func(t *testing.T) { ptrTestRandom(t, "rect", 10000) })
})
}
func ptrTestRandom(t *testing.T, which string, n int) {
fmt.Println("-------------------------------------------------")
fmt.Printf("Testing Random %dD %ss\n", D, which)
fmt.Println("-------------------------------------------------")
rand.Seed(time.Now().UnixNano())
tr := New(D, M)
min, max := tr.Bounds()
assertEqual(t, make([]float64, D), min[:])
assertEqual(t, make([]float64, D), max[:])
// create random objects
m1 := getMemStats()
objs := make([]*Rect, n)
for i := 0; i < n; i++ {
objs[i] = ptrMakeRandom(which)
}
// insert the objects into tree
m2 := getMemStats()
start := time.Now()
for _, r := range objs {
tr.Insert(r.min, r.max, r.item)
}
durInsert := time.Since(start)
m3 := getMemStats()
assertEqual(t, len(objs), tr.Count())
fmt.Printf("Inserted %d random %ss in %dms -- %d ops/sec\n",
len(objs), which, int(durInsert.Seconds()*1000),
int(float64(len(objs))/durInsert.Seconds()))
fmt.Printf(" total cost is %d bytes/%s\n", int(m3.HeapAlloc-m1.HeapAlloc)/len(objs), which)
fmt.Printf(" tree cost is %d bytes/%s\n", int(m3.HeapAlloc-m2.HeapAlloc)/len(objs), which)
fmt.Printf(" tree overhead %d%%\n", int((float64(m3.HeapAlloc-m2.HeapAlloc)/float64(len(objs)))/(float64(m3.HeapAlloc-m1.HeapAlloc)/float64(len(objs)))*100))
fmt.Printf(" complexity %f\n", tr.Complexity())
start = time.Now()
// count all nodes and leaves
var nodes int
var leaves int
var maxLevel int
tr.Traverse(func(min, max []float64, level int, item interface{}) bool {
if level != 0 {
nodes++
}
if level == 1 {
leaves++
}
if level > maxLevel {
maxLevel = level
}
return true
})
fmt.Printf(" nodes: %d, leaves: %d, level: %d\n", nodes, leaves, maxLevel)
// verify mbr
for i := 0; i < D; i++ {
min[i] = math.Inf(+1)
max[i] = math.Inf(-1)
}
for _, o := range objs {
for i := 0; i < D; i++ {
if o.min[i] < min[i] {
min[i] = o.min[i]
}
if o.max[i] > max[i] {
max[i] = o.max[i]
}
}
}
minb, maxb := tr.Bounds()
assertEqual(t, min, minb)
assertEqual(t, max, maxb)
// scan
var arr []*Rect
tr.Scan(func(item interface{}) bool {
arr = append(arr, item.(*Rect))
return true
})
assertEqual(t, true, ptrTestHasSameItems(objs, arr))
// search
ptrTestSearch(t, tr, objs, 0.10, true)
ptrTestSearch(t, tr, objs, 0.50, true)
ptrTestSearch(t, tr, objs, 1.00, true)
// knn
ptrTestKNN(t, tr, objs, int(float64(len(objs))*0.01), true)
ptrTestKNN(t, tr, objs, int(float64(len(objs))*0.50), true)
ptrTestKNN(t, tr, objs, int(float64(len(objs))*1.00), true)
// remove all objects
indexes := rand.Perm(len(objs))
start = time.Now()
for _, i := range indexes {
tr.Remove(objs[i].min, objs[i].max, objs[i].item)
}
durRemove := time.Since(start)
assertEqual(t, 0, tr.Count())
fmt.Printf("Removed %d random %ss in %dms -- %d ops/sec\n",
len(objs), which, int(durRemove.Seconds()*1000),
int(float64(len(objs))/durRemove.Seconds()))
min, max = tr.Bounds()
assertEqual(t, make([]float64, D), min[:])
assertEqual(t, make([]float64, D), max[:])
}
func ptrTestHasSameItems(a1, a2 []*Rect) bool {
if len(a1) != len(a2) {
return false
}
for _, p1 := range a1 {
var found bool
for _, p2 := range a2 {
if p1.equals(*p2) {
found = true
break
}
}
if !found {
return false
}
}
return true
}
func ptrTestSearch(t *testing.T, tr *RTree, objs []*Rect, percent float64, check bool) {
var found int
var start time.Time
var stop time.Time
defer func() {
dur := stop.Sub(start)
fmt.Printf("Searched %.0f%% (%d/%d items) in %dms -- %d ops/sec\n",
percent*100, found, len(objs), int(dur.Seconds()*1000),
int(float64(1)/dur.Seconds()),
)
}()
min, max := tr.Bounds()
vals := make([]float64, D*2)
for i := 0; i < D; i++ {
vals[i] = ((max[i]+min[i])/2 - ((max[i]-min[i])*percent)/2)
vals[D+i] = ((max[i]+min[i])/2 + ((max[i]-min[i])*percent)/2)
}
var arr1 []*Rect
var box *Rect
if percent == 1 {
box = ptrMakeRect(append(append([]float64{}, min[:]...), max[:]...)...)
} else {
box = ptrMakeRect(vals...)
}
start = time.Now()
tr.Search(box.min, box.max, func(item interface{}) bool {
if check {
arr1 = append(arr1, item.(*Rect))
}
found++
return true
})
stop = time.Now()
if !check {
return
}
var arr2 []*Rect
for _, obj := range objs {
if ptrTestIntersects(obj, box) {
arr2 = append(arr2, obj)
}
}
assertEqual(t, len(arr1), len(arr2))
for _, o1 := range arr1 {
var found bool
for _, o2 := range arr2 {
if o2.equals(*o1) {
found = true
break
}
}
if !found {
t.Fatalf("not found")
}
}
}
func ptrTestKNN(t *testing.T, tr *RTree, objs []*Rect, n int, check bool) {
var start time.Time
var stop time.Time
defer func() {
dur := stop.Sub(start)
fmt.Printf("KNN %d items in %dms -- %d ops/sec\n",
n, int(dur.Seconds()*1000),
int(float64(1)/dur.Seconds()),
)
}()
min, max := tr.Bounds()
pvals := make([]float64, D)
for i := 0; i < D; i++ {
pvals[i] = (max[i] + min[i]) / 2
}
point := ptrMakePoint(pvals...)
// gather the results, make sure that is matches exactly
var arr1 []Rect
var dists1 []float64
pdist := math.Inf(-1)
start = time.Now()
tr.KNN(point.min, point.max, false, func(item interface{}, dist float64) bool {
if len(arr1) == n {
return false
}
arr1 = append(arr1, Rect{min: min, max: max, item: item})
dists1 = append(dists1, dist)
if dist < pdist {
panic("dist out of order")
}
pdist = dist
return true
})
stop = time.Now()
assertEqual(t, true, n > len(objs) || n == len(arr1))
// get the KNN for the original array
nobjs := make([]*Rect, len(objs))
copy(nobjs, objs)
sort.Slice(nobjs, func(i, j int) bool {
idist := ptrTestBoxDist(pvals, nobjs[i].min, nobjs[i].max)
jdist := ptrTestBoxDist(pvals, nobjs[j].min, nobjs[j].max)
return idist < jdist
})
arr2 := nobjs[:len(arr1)]
var dists2 []float64
for i := 0; i < len(arr2); i++ {
dist := ptrTestBoxDist(pvals, arr2[i].min, arr2[i].max)
dists2 = append(dists2, dist)
}
// only compare the distances, not the objects because rectangles with
// a dist of zero will not be ordered.
assertEqual(t, dists1, dists2)
}
func ptrTestBoxDist(point []float64, min, max []float64) float64 {
var dist float64
for i := 0; i < len(point); i++ {
d := ptrTestAxisDist(point[i], min[i], max[i])
dist += d * d
}
return dist
}
func ptrTestAxisDist(k, min, max float64) float64 {
if k < min {
return min - k
}
if k <= max {
return 0
}
return k - max
}
func ptrTestIntersects(obj, box *Rect) bool {
for i := 0; i < D; i++ {
if box.min[i] > obj.max[i] || box.max[i] < obj.min[i] {
return false
}
}
return true
}
// func TestPtrInsertFlatPNG2D(t *testing.T) {
// fmt.Println("-------------------------------------------------")
// fmt.Println("Generating Cities PNG 2D (flat-insert-2d.png)")
// fmt.Println("-------------------------------------------------")
// tr := New()
// var items []*Rect
// c := cities.Cities
// for i := 0; i < len(c); i++ {
// x := c[i].Longitude
// y := c[i].Latitude
// items = append(items, ptrMakePoint(x, y))
// }
// start := time.Now()
// for _, item := range items {
// tr.Insert(item.min, item.max, item.item)
// }
// dur := time.Since(start)
// fmt.Printf("wrote %d cities (flat) in %s (%.0f/ops)\n", len(c), dur, float64(len(c))/dur.Seconds())
// withGIF := os.Getenv("GIFOUTPUT") != ""
// if err := tr.SavePNG("ptr-flat-insert-2d.png", 1000, 1000, 1.25/360.0, 0, true, withGIF, os.Stdout); err != nil {
// t.Fatal(err)
// }
// if !withGIF {
// fmt.Println("use GIFOUTPUT=1 for animated gif")
// }
// }
// func TestPtrLoadFlatPNG2D(t *testing.T) {
// fmt.Println("-------------------------------------------------")
// fmt.Println("Generating Cities 2D PNG (flat-load-2d.png)")
// fmt.Println("-------------------------------------------------")
// tr := New()
// var items []*Rect
// c := cities.Cities
// for i := 0; i < len(c); i++ {
// x := c[i].Longitude
// y := c[i].Latitude
// items = append(items, ptrMakePoint(x, y))
// }
// var mins [][D]float64
// var maxs [][D]float64
// var ifs []interface{}
// for i := 0; i < len(items); i++ {
// mins = append(mins, items[i].min)
// maxs = append(maxs, items[i].max)
// ifs = append(ifs, items[i].item)
// }
// start := time.Now()
// tr.Load(mins, maxs, ifs)
// dur := time.Since(start)
// if true {
// var all []*Rect
// tr.Scan(func(min, max [D]float64, item interface{}) bool {
// all = append(all, &Rect{min: min, max: max, item: item})
// return true
// })
// assertEqual(t, len(all), len(items))
// for len(all) > 0 {
// item := all[0]
// var found bool
// for _, city := range items {
// if *city == *item {
// found = true
// break
// }
// }
// if !found {
// t.Fatal("item not found")
// }
// all = all[1:]
// }
// }
// fmt.Printf("wrote %d cities (flat) in %s (%.0f/ops)\n", len(c), dur, float64(len(c))/dur.Seconds())
// withGIF := os.Getenv("GIFOUTPUT") != ""
// if err := tr.SavePNG("ptr-flat-load-2d.png", 1000, 1000, 1.25/360.0, 0, true, withGIF, os.Stdout); err != nil {
// t.Fatal(err)
// }
// if !withGIF {
// fmt.Println("use GIFOUTPUT=1 for animated gif")
// }
// }
func TestBenchmarks(t *testing.T) {
var points []*Rect
for i := 0; i < 2000000; i++ {
x := rand.Float64()*360 - 180
y := rand.Float64()*180 - 90
points = append(points, ptrMakePoint(x, y))
}
tr := New(D, M)
start := time.Now()
for i := len(points) / 2; i < len(points); i++ {
tr.Insert(points[i].min, points[i].max, points[i].item)
}
dur := time.Since(start)
log.Printf("insert 1M items one by one: %.3fs", dur.Seconds())
////
rarr := rand.Perm(len(points) / 2)
start = time.Now()
for i := 0; i < len(points)/2; i++ {
a := points[rarr[i]+len(points)/2]
b := points[rarr[i]]
tr.Remove(a.min, a.max, a.item)
tr.Insert(b.min, b.max, b.item)
}
dur = time.Since(start)
log.Printf("replaced 1M items one by one: %.3fs", dur.Seconds())
points = points[:len(points)/2]
////
start = time.Now()
for i := 0; i < 1000; i++ {
tr.Remove(points[i].min, points[i].max, points[i].item)
}
dur = time.Since(start)
log.Printf("remove 100 items one by one: %.3fs", dur.Seconds())
////
bbox := ptrMakeRect(0, 0, 0+(360*0.0001), 0+(180*0.0001))
start = time.Now()
for i := 0; i < 1000; i++ {
tr.Search(bbox.min, bbox.max, func(_ interface{}) bool { return true })
}
dur = time.Since(start)
log.Printf("1000 searches of 0.01%% area: %.3fs", dur.Seconds())
////
bbox = ptrMakeRect(0, 0, 0+(360*0.01), 0+(180*0.01))
start = time.Now()
for i := 0; i < 1000; i++ {
tr.Search(bbox.min, bbox.max, func(_ interface{}) bool { return true })
}
dur = time.Since(start)
log.Printf("1000 searches of 1%% area: %.3fs", dur.Seconds())
////
bbox = ptrMakeRect(0, 0, 0+(360*0.10), 0+(180*0.10))
start = time.Now()
for i := 0; i < 1000; i++ {
tr.Search(bbox.min, bbox.max, func(_ interface{}) bool { return true })
}
dur = time.Since(start)
log.Printf("1000 searches of 10%% area: %.3fs", dur.Seconds())
///
var mins [][]float64
var maxs [][]float64
var items []interface{}
for i := 0; i < len(points); i++ {
mins = append(mins, points[i].min)
maxs = append(maxs, points[i].max)
items = append(items, points[i].item)
}
tr = New(D, M)
start = time.Now()
tr.Load(mins, maxs, items)
dur = time.Since(start)
log.Printf("bulk-insert 1M items: %.3fs", dur.Seconds())
}
func assertEqual(t *testing.T, expected, actual interface{}) {
t.Helper()
if !reflect.DeepEqual(expected, actual) {
t.Fatalf("expected '%v', got '%v'", expected, actual)
}
}

87
gen/gen.go
View File

@ -1,87 +0,0 @@
package main
import (
"flag"
"io/ioutil"
"log"
"os"
"strconv"
"strings"
)
func main() {
var dims int
var debug bool
flag.IntVar(&dims, "dims", 4, "number of dimensions")
flag.BoolVar(&debug, "debug", false, "turn on debug tracing")
flag.Parse()
// process rtree.go
data, err := ioutil.ReadFile("src/rtree.go")
if err != nil {
log.Fatal(err)
}
data = []byte(strings.Replace(string(data), "// +build ignore", "// generated; DO NOT EDIT!", -1))
if debug {
data = []byte(strings.Replace(string(data), "TDEBUG", "true", -1))
} else {
data = []byte(strings.Replace(string(data), "TDEBUG", "false", -1))
}
var dimouts = make([]string, dims)
var output string
var recording bool
lines := strings.Split(string(data), "\n")
for _, line := range lines {
if strings.HasPrefix(strings.TrimSpace(line), "//") {
idx := strings.Index(line, "//")
switch strings.ToUpper(strings.TrimSpace(line[idx+2:])) {
case "BEGIN":
recording = true
for i := 0; i < len(dimouts); i++ {
dimouts[i] = ""
}
continue
case "END":
for _, out := range dimouts {
if out != "" {
output += out
}
}
recording = false
continue
}
}
if recording {
for i := 0; i < len(dimouts); i++ {
dimouts[i] += strings.Replace(line, "TNUMDIMS", strconv.FormatInt(int64(i+1), 10), -1) + "\n"
}
} else {
output += line + "\n"
}
}
// process rtree_base.go
if err := os.RemoveAll("../dims"); err != nil {
log.Fatal(err)
}
for i := 0; i < dims; i++ {
sdim := strconv.FormatInt(int64(i+1), 10)
data, err := ioutil.ReadFile("src/rtree_base.go")
if err != nil {
log.Fatal(err)
}
data = []byte(strings.Split(string(data), "// FILE_START")[1])
if debug {
data = []byte(strings.Replace(string(data), "TDEBUG", "true", -1))
} else {
data = []byte(strings.Replace(string(data), "TDEBUG", "false", -1))
}
data = []byte(strings.Replace(string(data), "TNUMDIMS", strconv.FormatInt(int64(i+1), 10), -1))
data = []byte(strings.Replace(string(data), "DD_", "d"+strconv.FormatInt(int64(i+1), 10), -1))
if err := os.MkdirAll("../dims/d"+sdim, 0777); err != nil {
log.Fatal(err)
}
output = string(append([]byte(output), data...))
}
if err := ioutil.WriteFile("../rtree.go", []byte(output), 0666); err != nil {
log.Fatal(err)
}
}

9
gen/gen.sh
View File

@ -1,9 +0,0 @@
#!/bin/bash
set -e
cd $(dirname "${BASH_SOURCE[0]}")
go run gen.go --dims=20 --debug=false
cd ..
go fmt

134
gen/src/rtree.go
View File

@ -1,134 +0,0 @@
// +build ignore
package rtree
import "math"
type Iterator func(item Item) bool
type Item interface {
Rect(ctx interface{}) (min []float64, max []float64)
}
type RTree struct {
ctx interface{}
// BEGIN
trTNUMDIMS *dTNUMDIMSRTree
// END
}
func New(ctx interface{}) *RTree {
return &RTree{
ctx: ctx,
// BEGIN
trTNUMDIMS: dTNUMDIMSNew(),
// END
}
}
func (tr *RTree) Insert(item Item) {
if item == nil {
panic("nil item being added to RTree")
}
min, max := item.Rect(tr.ctx)
if len(min) != len(max) {
return // just return
panic("invalid item rectangle")
}
switch len(min) {
default:
return // just return
panic("invalid dimension")
// BEGIN
case TNUMDIMS:
var amin, amax [TNUMDIMS]float64
for i := 0; i < len(min); i++ {
amin[i], amax[i] = min[i], max[i]
}
tr.trTNUMDIMS.Insert(amin, amax, item)
// END
}
}
func (tr *RTree) Remove(item Item) {
if item == nil {
panic("nil item being added to RTree")
}
min, max := item.Rect(tr.ctx)
if len(min) != len(max) {
return // just return
panic("invalid item rectangle")
}
switch len(min) {
default:
return // just return
panic("invalid dimension")
// BEGIN
case TNUMDIMS:
var amin, amax [TNUMDIMS]float64
for i := 0; i < len(min); i++ {
amin[i], amax[i] = min[i], max[i]
}
tr.trTNUMDIMS.Remove(amin, amax, item)
// END
}
}
func (tr *RTree) Reset() {
// BEGIN
tr.trTNUMDIMS = dTNUMDIMSNew()
// END
}
func (tr *RTree) Count() int {
count := 0
// BEGIN
count += tr.trTNUMDIMS.Count()
// END
return count
}
func (tr *RTree) Search(bounds Item, iter Iterator) {
if bounds == nil {
panic("nil bounds being used for search")
}
min, max := bounds.Rect(tr.ctx)
if len(min) != len(max) {
return // just return
panic("invalid item rectangle")
}
switch len(min) {
default:
return // just return
panic("invalid dimension")
// BEGIN
case TNUMDIMS:
// END
}
// BEGIN
if !tr.searchTNUMDIMS(min, max, iter) {
return
}
// END
}
// BEGIN
func (tr *RTree) searchTNUMDIMS(min, max []float64, iter Iterator) bool {
var amin, amax [TNUMDIMS]float64
for i := 0; i < TNUMDIMS; i++ {
if i < len(min) {
amin[i] = min[i]
amax[i] = max[i]
} else {
amin[i] = math.Inf(-1)
amax[i] = math.Inf(+1)
}
}
ended := false
tr.trTNUMDIMS.Search(amin, amax, func(dataID interface{}) bool {
if !iter(dataID.(Item)) {
ended = true
return false
}
return true
})
return !ended
}
// END

687
gen/src/rtree_base.go
View File

@ -1,687 +0,0 @@
// +build ignore
/*
TITLE
R-TREES: A DYNAMIC INDEX STRUCTURE FOR SPATIAL SEARCHING
DESCRIPTION
A Go version of the RTree algorithm.
AUTHORS
* 1983 Original algorithm and test code by Antonin Guttman and Michael Stonebraker, UC Berkely
* 1994 ANCI C ported from original test code by Melinda Green - melinda@superliminal.com
* 1995 Sphere volume fix for degeneracy problem submitted by Paul Brook
* 2004 Templated C++ port by Greg Douglas
* 2016 Go port by Josh Baker
LICENSE:
Entirely free for all uses. Enjoy!
*/
// Implementation of RTree, a multidimensional bounding rectangle tree.
package rtree
import "math"
// FILE_START
func DD_fmin(a, b float64) float64 {
if a < b {
return a
}
return b
}
func DD_fmax(a, b float64) float64 {
if a > b {
return a
}
return b
}
const (
DD_numDims = TNUMDIMS
DD_maxNodes = 8
DD_minNodes = DD_maxNodes / 2
DD_useSphericalVolume = true // Better split classification, may be slower on some systems
)
var DD_unitSphereVolume = []float64{
0.000000, 2.000000, 3.141593, // Dimension 0,1,2
4.188790, 4.934802, 5.263789, // Dimension 3,4,5
5.167713, 4.724766, 4.058712, // Dimension 6,7,8
3.298509, 2.550164, 1.884104, // Dimension 9,10,11
1.335263, 0.910629, 0.599265, // Dimension 12,13,14
0.381443, 0.235331, 0.140981, // Dimension 15,16,17
0.082146, 0.046622, 0.025807, // Dimension 18,19,20
}[DD_numDims]
type DD_RTree struct {
root *DD_nodeT ///< Root of tree
}
/// Minimal bounding rectangle (n-dimensional)
type DD_rectT struct {
min [DD_numDims]float64 ///< Min dimensions of bounding box
max [DD_numDims]float64 ///< Max dimensions of bounding box
}
/// May be data or may be another subtree
/// The parents level determines this.
/// If the parents level is 0, then this is data
type DD_branchT struct {
rect DD_rectT ///< Bounds
child *DD_nodeT ///< Child node
data interface{} ///< Data Id or Ptr
}
/// DD_nodeT for each branch level
type DD_nodeT struct {
count int ///< Count
level int ///< Leaf is zero, others positive
branch [DD_maxNodes]DD_branchT ///< Branch
}
func (node *DD_nodeT) isInternalNode() bool {
return (node.level > 0) // Not a leaf, but a internal node
}
func (node *DD_nodeT) isLeaf() bool {
return (node.level == 0) // A leaf, contains data
}
/// A link list of nodes for reinsertion after a delete operation
type DD_listNodeT struct {
next *DD_listNodeT ///< Next in list
node *DD_nodeT ///< Node
}
const DD_notTaken = -1 // indicates that position
/// Variables for finding a split partition
type DD_partitionVarsT struct {
partition [DD_maxNodes + 1]int
total int
minFill int
count [2]int
cover [2]DD_rectT
area [2]float64
branchBuf [DD_maxNodes + 1]DD_branchT
branchCount int
coverSplit DD_rectT
coverSplitArea float64
}
func DD_New() *DD_RTree {
// We only support machine word size simple data type eg. integer index or object pointer.
// Since we are storing as union with non data branch
return &DD_RTree{
root: &DD_nodeT{},
}
}
/// Insert entry
/// \param a_min Min of bounding rect
/// \param a_max Max of bounding rect
/// \param a_dataId Positive Id of data. Maybe zero, but negative numbers not allowed.
func (tr *DD_RTree) Insert(min, max [DD_numDims]float64, dataId interface{}) {
var branch DD_branchT
branch.data = dataId
for axis := 0; axis < DD_numDims; axis++ {
branch.rect.min[axis] = min[axis]
branch.rect.max[axis] = max[axis]
}
DD_insertRect(&branch, &tr.root, 0)
}
/// Remove entry
/// \param a_min Min of bounding rect
/// \param a_max Max of bounding rect
/// \param a_dataId Positive Id of data. Maybe zero, but negative numbers not allowed.
func (tr *DD_RTree) Remove(min, max [DD_numDims]float64, dataId interface{}) {
var rect DD_rectT
for axis := 0; axis < DD_numDims; axis++ {
rect.min[axis] = min[axis]
rect.max[axis] = max[axis]
}
DD_removeRect(&rect, dataId, &tr.root)
}
/// Find all within DD_search rectangle
/// \param a_min Min of DD_search bounding rect
/// \param a_max Max of DD_search bounding rect
/// \param a_searchResult DD_search result array. Caller should set grow size. Function will reset, not append to array.
/// \param a_resultCallback Callback function to return result. Callback should return 'true' to continue searching
/// \param a_context User context to pass as parameter to a_resultCallback
/// \return Returns the number of entries found
func (tr *DD_RTree) Search(min, max [DD_numDims]float64, resultCallback func(data interface{}) bool) int {
var rect DD_rectT
for axis := 0; axis < DD_numDims; axis++ {
rect.min[axis] = min[axis]
rect.max[axis] = max[axis]
}
foundCount, _ := DD_search(tr.root, rect, 0, resultCallback)
return foundCount
}
/// Count the data elements in this container. This is slow as no internal counter is maintained.
func (tr *DD_RTree) Count() int {
var count int
DD_countRec(tr.root, &count)
return count
}
/// Remove all entries from tree
func (tr *DD_RTree) RemoveAll() {
// Delete all existing nodes
tr.root = &DD_nodeT{}
}
func DD_countRec(node *DD_nodeT, count *int) {
if node.isInternalNode() { // not a leaf node
for index := 0; index < node.count; index++ {
DD_countRec(node.branch[index].child, count)
}
} else { // A leaf node
*count += node.count
}
}
// Inserts a new data rectangle into the index structure.
// Recursively descends tree, propagates splits back up.
// Returns 0 if node was not split. Old node updated.
// If node was split, returns 1 and sets the pointer pointed to by
// new_node to point to the new node. Old node updated to become one of two.
// The level argument specifies the number of steps up from the leaf
// level to insert; e.g. a data rectangle goes in at level = 0.
func DD_insertRectRec(branch *DD_branchT, node *DD_nodeT, newNode **DD_nodeT, level int) bool {
// recurse until we reach the correct level for the new record. data records
// will always be called with a_level == 0 (leaf)
if node.level > level {
// Still above level for insertion, go down tree recursively
var otherNode *DD_nodeT
//var newBranch DD_branchT
// find the optimal branch for this record
index := DD_pickBranch(&branch.rect, node)
// recursively insert this record into the picked branch
childWasSplit := DD_insertRectRec(branch, node.branch[index].child, &otherNode, level)
if !childWasSplit {
// Child was not split. Merge the bounding box of the new record with the
// existing bounding box
node.branch[index].rect = DD_combineRect(&branch.rect, &(node.branch[index].rect))
return false
} else {
// Child was split. The old branches are now re-partitioned to two nodes
// so we have to re-calculate the bounding boxes of each node
node.branch[index].rect = DD_nodeCover(node.branch[index].child)
var newBranch DD_branchT
newBranch.child = otherNode
newBranch.rect = DD_nodeCover(otherNode)
// The old node is already a child of a_node. Now add the newly-created
// node to a_node as well. a_node might be split because of that.
return DD_addBranch(&newBranch, node, newNode)
}
} else if node.level == level {
// We have reached level for insertion. Add rect, split if necessary
return DD_addBranch(branch, node, newNode)
} else {
// Should never occur
return false
}
}
// Insert a data rectangle into an index structure.
// DD_insertRect provides for splitting the root;
// returns 1 if root was split, 0 if it was not.
// The level argument specifies the number of steps up from the leaf
// level to insert; e.g. a data rectangle goes in at level = 0.
// InsertRect2 does the recursion.
//
func DD_insertRect(branch *DD_branchT, root **DD_nodeT, level int) bool {
var newNode *DD_nodeT
if DD_insertRectRec(branch, *root, &newNode, level) { // Root split
// Grow tree taller and new root
newRoot := &DD_nodeT{}
newRoot.level = (*root).level + 1
var newBranch DD_branchT
// add old root node as a child of the new root
newBranch.rect = DD_nodeCover(*root)
newBranch.child = *root
DD_addBranch(&newBranch, newRoot, nil)
// add the split node as a child of the new root
newBranch.rect = DD_nodeCover(newNode)
newBranch.child = newNode
DD_addBranch(&newBranch, newRoot, nil)
// set the new root as the root node
*root = newRoot
return true
}
return false
}
// Find the smallest rectangle that includes all rectangles in branches of a node.
func DD_nodeCover(node *DD_nodeT) DD_rectT {
rect := node.branch[0].rect
for index := 1; index < node.count; index++ {
rect = DD_combineRect(&rect, &(node.branch[index].rect))
}
return rect
}
// Add a branch to a node. Split the node if necessary.
// Returns 0 if node not split. Old node updated.
// Returns 1 if node split, sets *new_node to address of new node.
// Old node updated, becomes one of two.
func DD_addBranch(branch *DD_branchT, node *DD_nodeT, newNode **DD_nodeT) bool {
if node.count < DD_maxNodes { // Split won't be necessary
node.branch[node.count] = *branch
node.count++
return false
} else {
DD_splitNode(node, branch, newNode)
return true
}
}
// Disconnect a dependent node.
// Caller must return (or stop using iteration index) after this as count has changed
func DD_disconnectBranch(node *DD_nodeT, index int) {
// Remove element by swapping with the last element to prevent gaps in array
node.branch[index] = node.branch[node.count-1]
node.branch[node.count-1].data = nil
node.branch[node.count-1].child = nil
node.count--
}
// Pick a branch. Pick the one that will need the smallest increase
// in area to accomodate the new rectangle. This will result in the
// least total area for the covering rectangles in the current node.
// In case of a tie, pick the one which was smaller before, to get
// the best resolution when searching.
func DD_pickBranch(rect *DD_rectT, node *DD_nodeT) int {
var firstTime bool = true
var increase float64
var bestIncr float64 = -1
var area float64
var bestArea float64
var best int
var tempRect DD_rectT
for index := 0; index < node.count; index++ {
curRect := &node.branch[index].rect
area = DD_calcRectVolume(curRect)
tempRect = DD_combineRect(rect, curRect)
increase = DD_calcRectVolume(&tempRect) - area
if (increase < bestIncr) || firstTime {
best = index
bestArea = area
bestIncr = increase
firstTime = false
} else if (increase == bestIncr) && (area < bestArea) {
best = index
bestArea = area
bestIncr = increase
}
}
return best
}
// Combine two rectangles into larger one containing both
func DD_combineRect(rectA, rectB *DD_rectT) DD_rectT {
var newRect DD_rectT
for index := 0; index < DD_numDims; index++ {
newRect.min[index] = DD_fmin(rectA.min[index], rectB.min[index])
newRect.max[index] = DD_fmax(rectA.max[index], rectB.max[index])
}
return newRect
}
// Split a node.
// Divides the nodes branches and the extra one between two nodes.
// Old node is one of the new ones, and one really new one is created.
// Tries more than one method for choosing a partition, uses best result.
func DD_splitNode(node *DD_nodeT, branch *DD_branchT, newNode **DD_nodeT) {
// Could just use local here, but member or external is faster since it is reused
var localVars DD_partitionVarsT
parVars := &localVars
// Load all the branches into a buffer, initialize old node
DD_getBranches(node, branch, parVars)
// Find partition
DD_choosePartition(parVars, DD_minNodes)
// Create a new node to hold (about) half of the branches
*newNode = &DD_nodeT{}
(*newNode).level = node.level
// Put branches from buffer into 2 nodes according to the chosen partition
node.count = 0
DD_loadNodes(node, *newNode, parVars)
}
// Calculate the n-dimensional volume of a rectangle
func DD_rectVolume(rect *DD_rectT) float64 {
var volume float64 = 1
for index := 0; index < DD_numDims; index++ {
volume *= rect.max[index] - rect.min[index]
}
return volume
}
// The exact volume of the bounding sphere for the given DD_rectT
func DD_rectSphericalVolume(rect *DD_rectT) float64 {
var sumOfSquares float64 = 0
var radius float64
for index := 0; index < DD_numDims; index++ {
halfExtent := (rect.max[index] - rect.min[index]) * 0.5
sumOfSquares += halfExtent * halfExtent
}
radius = math.Sqrt(sumOfSquares)
// Pow maybe slow, so test for common dims just use x*x, x*x*x.
if DD_numDims == 5 {
return (radius * radius * radius * radius * radius * DD_unitSphereVolume)
} else if DD_numDims == 4 {
return (radius * radius * radius * radius * DD_unitSphereVolume)
} else if DD_numDims == 3 {
return (radius * radius * radius * DD_unitSphereVolume)
} else if DD_numDims == 2 {
return (radius * radius * DD_unitSphereVolume)
} else {
return (math.Pow(radius, DD_numDims) * DD_unitSphereVolume)
}
}
// Use one of the methods to calculate retangle volume
func DD_calcRectVolume(rect *DD_rectT) float64 {
if DD_useSphericalVolume {
return DD_rectSphericalVolume(rect) // Slower but helps certain merge cases
} else { // RTREE_USE_SPHERICAL_VOLUME
return DD_rectVolume(rect) // Faster but can cause poor merges
} // RTREE_USE_SPHERICAL_VOLUME
}
// Load branch buffer with branches from full node plus the extra branch.
func DD_getBranches(node *DD_nodeT, branch *DD_branchT, parVars *DD_partitionVarsT) {
// Load the branch buffer
for index := 0; index < DD_maxNodes; index++ {
parVars.branchBuf[index] = node.branch[index]
}
parVars.branchBuf[DD_maxNodes] = *branch
parVars.branchCount = DD_maxNodes + 1
// Calculate rect containing all in the set
parVars.coverSplit = parVars.branchBuf[0].rect
for index := 1; index < DD_maxNodes+1; index++ {
parVars.coverSplit = DD_combineRect(&parVars.coverSplit, &parVars.branchBuf[index].rect)
}
parVars.coverSplitArea = DD_calcRectVolume(&parVars.coverSplit)
}
// Method #0 for choosing a partition:
// As the seeds for the two groups, pick the two rects that would waste the
// most area if covered by a single rectangle, i.e. evidently the worst pair
// to have in the same group.
// Of the remaining, one at a time is chosen to be put in one of the two groups.
// The one chosen is the one with the greatest difference in area expansion
// depending on which group - the rect most strongly attracted to one group
// and repelled from the other.
// If one group gets too full (more would force other group to violate min
// fill requirement) then other group gets the rest.
// These last are the ones that can go in either group most easily.
func DD_choosePartition(parVars *DD_partitionVarsT, minFill int) {
var biggestDiff float64
var group, chosen, betterGroup int
DD_initParVars(parVars, parVars.branchCount, minFill)
DD_pickSeeds(parVars)
for ((parVars.count[0] + parVars.count[1]) < parVars.total) &&
(parVars.count[0] < (parVars.total - parVars.minFill)) &&
(parVars.count[1] < (parVars.total - parVars.minFill)) {
biggestDiff = -1
for index := 0; index < parVars.total; index++ {
if DD_notTaken == parVars.partition[index] {
curRect := &parVars.branchBuf[index].rect
rect0 := DD_combineRect(curRect, &parVars.cover[0])
rect1 := DD_combineRect(curRect, &parVars.cover[1])
growth0 := DD_calcRectVolume(&rect0) - parVars.area[0]
growth1 := DD_calcRectVolume(&rect1) - parVars.area[1]
diff := growth1 - growth0
if diff >= 0 {
group = 0
} else {
group = 1
diff = -diff
}
if diff > biggestDiff {
biggestDiff = diff
chosen = index
betterGroup = group
} else if (diff == biggestDiff) && (parVars.count[group] < parVars.count[betterGroup]) {
chosen = index
betterGroup = group
}
}
}
DD_classify(chosen, betterGroup, parVars)
}
// If one group too full, put remaining rects in the other
if (parVars.count[0] + parVars.count[1]) < parVars.total {
if parVars.count[0] >= parVars.total-parVars.minFill {
group = 1
} else {
group = 0
}
for index := 0; index < parVars.total; index++ {
if DD_notTaken == parVars.partition[index] {
DD_classify(index, group, parVars)
}
}
}
}
// Copy branches from the buffer into two nodes according to the partition.
func DD_loadNodes(nodeA, nodeB *DD_nodeT, parVars *DD_partitionVarsT) {
for index := 0; index < parVars.total; index++ {
targetNodeIndex := parVars.partition[index]
targetNodes := []*DD_nodeT{nodeA, nodeB}
// It is assured that DD_addBranch here will not cause a node split.
DD_addBranch(&parVars.branchBuf[index], targetNodes[targetNodeIndex], nil)
}
}
// Initialize a DD_partitionVarsT structure.
func DD_initParVars(parVars *DD_partitionVarsT, maxRects, minFill int) {
parVars.count[0] = 0
parVars.count[1] = 0
parVars.area[0] = 0
parVars.area[1] = 0
parVars.total = maxRects
parVars.minFill = minFill
for index := 0; index < maxRects; index++ {
parVars.partition[index] = DD_notTaken
}
}
func DD_pickSeeds(parVars *DD_partitionVarsT) {
var seed0, seed1 int
var worst, waste float64
var area [DD_maxNodes + 1]float64
for index := 0; index < parVars.total; index++ {
area[index] = DD_calcRectVolume(&parVars.branchBuf[index].rect)
}
worst = -parVars.coverSplitArea - 1
for indexA := 0; indexA < parVars.total-1; indexA++ {
for indexB := indexA + 1; indexB < parVars.total; indexB++ {
oneRect := DD_combineRect(&parVars.branchBuf[indexA].rect, &parVars.branchBuf[indexB].rect)
waste = DD_calcRectVolume(&oneRect) - area[indexA] - area[indexB]
if waste > worst {
worst = waste
seed0 = indexA
seed1 = indexB
}
}
}
DD_classify(seed0, 0, parVars)
DD_classify(seed1, 1, parVars)
}
// Put a branch in one of the groups.
func DD_classify(index, group int, parVars *DD_partitionVarsT) {
parVars.partition[index] = group
// Calculate combined rect
if parVars.count[group] == 0 {
parVars.cover[group] = parVars.branchBuf[index].rect
} else {
parVars.cover[group] = DD_combineRect(&parVars.branchBuf[index].rect, &parVars.cover[group])
}
// Calculate volume of combined rect
parVars.area[group] = DD_calcRectVolume(&parVars.cover[group])
parVars.count[group]++
}
// Delete a data rectangle from an index structure.
// Pass in a pointer to a DD_rectT, the tid of the record, ptr to ptr to root node.
// Returns 1 if record not found, 0 if success.
// DD_removeRect provides for eliminating the root.
func DD_removeRect(rect *DD_rectT, id interface{}, root **DD_nodeT) bool {
var reInsertList *DD_listNodeT
if !DD_removeRectRec(rect, id, *root, &reInsertList) {
// Found and deleted a data item
// Reinsert any branches from eliminated nodes
for reInsertList != nil {
tempNode := reInsertList.node
for index := 0; index < tempNode.count; index++ {
// TODO go over this code. should I use (tempNode->m_level - 1)?
DD_insertRect(&tempNode.branch[index], root, tempNode.level)
}
reInsertList = reInsertList.next
}
// Check for redundant root (not leaf, 1 child) and eliminate TODO replace
// if with while? In case there is a whole branch of redundant roots...
if (*root).count == 1 && (*root).isInternalNode() {
tempNode := (*root).branch[0].child
*root = tempNode
}
return false
} else {
return true
}
}
// Delete a rectangle from non-root part of an index structure.
// Called by DD_removeRect. Descends tree recursively,
// merges branches on the way back up.
// Returns 1 if record not found, 0 if success.
func DD_removeRectRec(rect *DD_rectT, id interface{}, node *DD_nodeT, listNode **DD_listNodeT) bool {
if node.isInternalNode() { // not a leaf node
for index := 0; index < node.count; index++ {
if DD_overlap(*rect, node.branch[index].rect) {
if !DD_removeRectRec(rect, id, node.branch[index].child, listNode) {
if node.branch[index].child.count >= DD_minNodes {
// child removed, just resize parent rect
node.branch[index].rect = DD_nodeCover(node.branch[index].child)
} else {
// child removed, not enough entries in node, eliminate node
DD_reInsert(node.branch[index].child, listNode)
DD_disconnectBranch(node, index) // Must return after this call as count has changed
}
return false
}
}
}
return true
} else { // A leaf node
for index := 0; index < node.count; index++ {
if node.branch[index].data == id {
DD_disconnectBranch(node, index) // Must return after this call as count has changed
return false
}
}
return true
}
}
// Decide whether two rectangles DD_overlap.
func DD_overlap(rectA, rectB DD_rectT) bool {
for index := 0; index < DD_numDims; index++ {
if rectA.min[index] > rectB.max[index] ||
rectB.min[index] > rectA.max[index] {
return false
}
}
return true
}
// Add a node to the reinsertion list. All its branches will later
// be reinserted into the index structure.
func DD_reInsert(node *DD_nodeT, listNode **DD_listNodeT) {
newListNode := &DD_listNodeT{}
newListNode.node = node
newListNode.next = *listNode
*listNode = newListNode
}
// DD_search in an index tree or subtree for all data retangles that DD_overlap the argument rectangle.
func DD_search(node *DD_nodeT, rect DD_rectT, foundCount int, resultCallback func(data interface{}) bool) (int, bool) {
if node.isInternalNode() {
// This is an internal node in the tree
for index := 0; index < node.count; index++ {
if DD_overlap(rect, node.branch[index].rect) {
var ok bool
foundCount, ok = DD_search(node.branch[index].child, rect, foundCount, resultCallback)
if !ok {
// The callback indicated to stop searching
return foundCount, false
}
}
}
} else {
// This is a leaf node
for index := 0; index < node.count; index++ {
if DD_overlap(rect, node.branch[index].rect) {
id := node.branch[index].data
foundCount++
if !resultCallback(id) {
return foundCount, false // Don't continue searching
}
}
}
}
return foundCount, true // Continue searching
}

14111
rtree.go
View File

File diff suppressed because it is too large Load Diff

56
rtree_test.go
View File

@ -167,3 +167,59 @@ func BenchmarkInsert(t *testing.B) {
t.StartTimer()
}
func TestKNN(t *testing.T) {
n := 25000
tr := New(nil)
var points []*tPoint
rand.Seed(1)
for i := 0; i < n; i++ {
r := tRandPoint()
points = append(points, r)
tr.Insert(r)
}
if tr.Count() != n {
t.Fatalf("expecting %v, got %v", n, tr.Count())
}
var count int
tr.Search(&tRect{-100, -100, -100, -100, 100, 100, 100, 100}, func(item Item) bool {
count++
return true
})
var pdist float64
var i int
center := []float64{50, 50}
centerRect := &tRect{center[0], center[1], center[0], center[1]}
tr.KNN(centerRect, true, func(item Item, dist float64) bool {
dist2 := boxDistPoint(center, item)
if i > 0 && dist2 < pdist {
t.Fatal("out of order")
}
pdist = dist
i++
return true
})
if i != n {
t.Fatal("mismatch")
}
}
func boxDistPoint(point []float64, item Item) float64 {
var dist float64
min, max := item.Rect(nil)
for i := 0; i < len(point); i++ {
d := axisDist(point[i], min[i], max[i])
dist += d * d
}
return dist
}
func axisDist(k, min, max float64) float64 {
if k < min {
return min - k
}
if k <= max {
return 0
}
return k - max
}