src/utils/kdtree.js
/**
* @file Kdtree
* @author Alexander Rose <alexander.rose@weirdbyte.de>
* @private
*/
import BinaryHeap from './binary-heap.js'
/**
* Kdtree
* @class
* @author Alexander Rose <alexander.rose@weirdbyte.de>, 2016
* @author Roman Bolzern <roman.bolzern@fhnw.ch>, 2013
* @author I4DS http://www.fhnw.ch/i4ds, 2013
* @license MIT License <http://www.opensource.org/licenses/mit-license.php>
* @description
* k-d Tree for typed arrays of 3d points (e.g. for Float32Array), in-place
* provides fast nearest neighbour search
*
* Based on https://github.com/ubilabs/kd-tree-javascript by Ubilabs
*
* Further information (including mathematical properties)
* http://en.wikipedia.org/wiki/Binary_tree
* http://en.wikipedia.org/wiki/K-d_tree
*
* @example
* points: [x, y, z, x, y, z, x, y, z, ...]
* metric: function(a, b){
* return Math.pow(a[0]-b[0], 2) + Math.pow(a[1]-b[1], 2) + Math.pow(a[2]-b[2], 2);
* }
*
* @param {Float32Array} points - points
* @param {Function} metric - metric
*/
function Kdtree (points, metric) {
var n = points.length / 3
var maxDepth = 0
var indices = new Uint32Array(n)
for (var i = 0; i < n; ++i) {
indices[ i ] = i
}
var nodes = new Int32Array(n * 4)
// pos, left, right, parent
var currentNode = 0
var currentDim = 0
function buildTree (depth, parent, arrBegin, arrEnd) {
if (depth > maxDepth) maxDepth = depth
var plength = arrEnd - arrBegin
if (plength === 0) {
return -1
}
var nodeIndex = currentNode * 4
currentNode += 1
if (plength === 1) {
nodes[ nodeIndex ] = arrBegin
nodes[ nodeIndex + 1 ] = -1
nodes[ nodeIndex + 2 ] = -1
nodes[ nodeIndex + 3 ] = parent
return nodeIndex
}
// if( plength <= 32 ){
// return nodeIndex;
// }
var arrMedian = arrBegin + Math.floor(plength / 2)
currentDim = depth % 3
// inlined quickselect function
var j, tmp, pivotIndex, pivotValue, storeIndex
var left = arrBegin
var right = arrEnd - 1
while (right > left) {
pivotIndex = (left + right) >> 1
pivotValue = points[ indices[ pivotIndex ] * 3 + currentDim ]
// swap( pivotIndex, right );
tmp = indices[ pivotIndex ]
indices[ pivotIndex ] = indices[ right ]
indices[ right ] = tmp
storeIndex = left
for (j = left; j < right; ++j) {
if (points[ indices[ j ] * 3 + currentDim ] < pivotValue) {
// swap( storeIndex, j );
tmp = indices[ storeIndex ]
indices[ storeIndex ] = indices[ j ]
indices[ j ] = tmp
++storeIndex
}
}
// swap( right, storeIndex );
tmp = indices[ right ]
indices[ right ] = indices[ storeIndex ]
indices[ storeIndex ] = tmp
pivotIndex = storeIndex
if (arrMedian === pivotIndex) {
break
} else if (arrMedian < pivotIndex) {
right = pivotIndex - 1
} else {
left = pivotIndex + 1
}
}
nodes[ nodeIndex ] = arrMedian
nodes[ nodeIndex + 1 ] = buildTree(depth + 1, nodeIndex, arrBegin, arrMedian)
nodes[ nodeIndex + 2 ] = buildTree(depth + 1, nodeIndex, arrMedian + 1, arrEnd)
nodes[ nodeIndex + 3 ] = parent
return nodeIndex
}
function getNodeDepth (nodeIndex) {
var parentIndex = nodes[ nodeIndex + 3 ]
if (parentIndex === -1) {
return 0
} else {
return getNodeDepth(parentIndex) + 1
}
}
// TODO
// function getNodePos( node ){}
var rootIndex = buildTree(0, -1, 0, n)
/**
* find nearest points
* @param {Array} point - array of size 3
* @param {Integer} maxNodes - max amount of nodes to return
* @param {Float} maxDistance - maximum distance of point to result nodes
* @return {Array} array of point, distance pairs
*/
function nearest (point, maxNodes, maxDistance) {
var bestNodes = new BinaryHeap(
function (e) { return -e[ 1 ] }
)
function nearestSearch (nodeIndex) {
var bestChild, otherChild
var dimension = getNodeDepth(nodeIndex) % 3
var pointIndex = indices[ nodes[ nodeIndex ] ] * 3
var ownPoint = [
points[ pointIndex + 0 ],
points[ pointIndex + 1 ],
points[ pointIndex + 2 ]
]
var ownDistance = metric(point, ownPoint)
function saveNode (nodeIndex, distance) {
bestNodes.push([ nodeIndex, distance ])
if (bestNodes.size() > maxNodes) {
bestNodes.pop()
}
}
var leftIndex = nodes[ nodeIndex + 1 ]
var rightIndex = nodes[ nodeIndex + 2 ]
// if it's a leaf
if (rightIndex === -1 && leftIndex === -1) {
if ((bestNodes.size() < maxNodes || ownDistance < bestNodes.peek()[ 1 ]) &&
ownDistance <= maxDistance
) {
saveNode(nodeIndex, ownDistance)
}
return
}
if (rightIndex === -1) {
bestChild = leftIndex
} else if (leftIndex === -1) {
bestChild = rightIndex
} else {
if (point[ dimension ] <= points[ pointIndex + dimension ]) {
bestChild = leftIndex
} else {
bestChild = rightIndex
}
}
// recursive search
nearestSearch(bestChild)
if ((bestNodes.size() < maxNodes || ownDistance < bestNodes.peek()[ 1 ]) &&
ownDistance <= maxDistance
) {
saveNode(nodeIndex, ownDistance)
}
// if there's still room or the current distance is nearer than the best distance
var linearPoint = []
for (var i = 0; i < 3; i += 1) {
if (i === dimension) {
linearPoint[ i ] = point[ i ]
} else {
linearPoint[ i ] = points[ pointIndex + i ]
}
}
var linearDistance = metric(linearPoint, ownPoint)
if ((bestNodes.size() < maxNodes || Math.abs(linearDistance) < bestNodes.peek()[ 1 ]) &&
Math.abs(linearDistance) <= maxDistance
) {
if (bestChild === leftIndex) {
otherChild = rightIndex
} else {
otherChild = leftIndex
}
if (otherChild !== -1) {
nearestSearch(otherChild)
}
}
}
nearestSearch(rootIndex)
var result = []
for (var i = 0, il = Math.min(bestNodes.size(), maxNodes); i < il; i += 1) {
result.push(bestNodes.content[ i ])
}
return result
}
function verify (nodeIndex, depth) {
var count = 1
if (nodeIndex === undefined) {
nodeIndex = rootIndex
depth = 0
}
if (nodeIndex === -1) {
throw new Error('node is null')
}
var dim = depth % 3
var leftIndex = nodes[ nodeIndex + 1 ]
var rightIndex = nodes[ nodeIndex + 2 ]
if (leftIndex !== -1) {
if (points[ indices[ nodes[ leftIndex ] ] * 3 + dim ] >
points[ indices[ nodes[ nodeIndex ] ] * 3 + dim ]
) {
throw new Error('left child is > parent!')
}
count += verify(leftIndex, depth + 1)
}
if (rightIndex !== -1) {
if (points[ indices[ nodes[ rightIndex ] ] * 3 + dim ] <
points[ indices[ nodes[ nodeIndex ] ] * 3 + dim ]
) {
throw new Error('right child is < parent!')
}
count += verify(rightIndex, depth + 1)
}
return count
}
// API
this.rootIndex = rootIndex
this.maxDepth = maxDepth
this.nearest = nearest
this.indices = indices
this.nodes = nodes
this.verify = verify
}
export default Kdtree