@thi.ng/geom-accel
Version:
n-D spatial indexing data structures with a shared ES6 Map/Set-like API
293 lines (292 loc) • 7.62 kB
JavaScript
import { equivArrayLike } from "@thi.ng/equiv";
import { assert } from "@thi.ng/errors/assert";
import { pointInCenteredBox } from "@thi.ng/geom-isec/point";
import { testCenteredBoxSphere } from "@thi.ng/geom-isec/rect-circle";
import { Heap } from "@thi.ng/heaps/heap";
import { EPS } from "@thi.ng/math/api";
import { iterate } from "@thi.ng/transducers/iterate";
import { map } from "@thi.ng/transducers/map";
import { permutations } from "@thi.ng/transducers/permutations";
import { repeat } from "@thi.ng/transducers/repeat";
import { take } from "@thi.ng/transducers/take";
import { addmN } from "@thi.ng/vectors/addmn";
import { distSq } from "@thi.ng/vectors/distsq";
import { madd } from "@thi.ng/vectors/madd";
import { mulN } from "@thi.ng/vectors/muln";
import { submN } from "@thi.ng/vectors/submn";
import { vop } from "@thi.ng/vectors/vop";
import { CMP, __addResults, __into } from "./utils.js";
class NdQtNode {
pos;
ext;
parent;
children;
numC;
k;
v;
constructor(parent, pos, ext) {
this.parent = parent;
this.pos = pos;
this.ext = ext;
this.numC = 0;
}
clear() {
delete this.children;
delete this.k;
delete this.v;
this.numC = 0;
}
set(p, val, eps, distFn) {
return (eps <= 0 || !this.queryKeys(p, eps, 1, [], distFn).length) && this.containsPoint(p) && this.setUnsafe(p, val);
}
setUnsafe(p, val) {
if (this.k) {
if (equivArrayLike(this.k, p)) {
this.v = val;
return false;
}
this.ensureChild(__childID(this.k, this.pos)).setUnsafe(
this.k,
this.v
);
delete this.k;
delete this.v;
}
if (this.children) {
return this.ensureChild(__childID(p, this.pos)).setUnsafe(p, val);
} else {
this.k = p;
this.v = val;
}
return true;
}
query(fn, p, r, max, acc, distFn) {
return __addResults(
fn,
this.doQuery(
p,
r,
max,
new Heap([[r * r]], {
compare: CMP
}),
distFn
).values,
acc
);
}
queryKeys(p, r, max, acc, distFn) {
return this.query((n) => n.k, p, r, max, acc, distFn);
}
queryValues(p, r, max, acc, distFn) {
return this.query((n) => n.v, p, r, max, acc, distFn);
}
containsPoint(p) {
return pointInCenteredBox(p, this.pos, this.ext);
}
nodeForPoint(p) {
if (this.k && equivArrayLike(this.k, p)) {
return this;
}
if (this.children) {
const child = this.children[__childID(p, this.pos)];
return child ? child.nodeForPoint(p) : void 0;
}
}
doQuery(p, r, max, acc, distFn) {
if (!testCenteredBoxSphere(this.pos, this.ext, p, r)) return acc;
if (this.k) {
const d = distFn(this.k, p);
if (d <= acc.values[0][0]) {
acc.length >= max ? acc.pushPop([d, this]) : acc.push([d, this]);
}
} else if (this.children) {
for (let i = MAX_CHILDREN[this.pos.length], j = this.numC; i-- > 0 && j > 0; ) {
if (this.children[i]) {
this.children[i].doQuery(p, r, max, acc, distFn);
j--;
}
}
}
return acc;
}
ensureChild(id) {
!this.children && (this.children = []);
let c = this.children[id];
if (!c) {
const csize = mulN([], this.ext, 0.5);
this.children[id] = c = new NdQtNode(
this,
madd([], csize, CHILD_OFFSETS[csize.length][id], this.pos),
csize
);
this.numC++;
}
return c;
}
}
class NdQuadtreeMap {
constructor(pos, ext, pairs, distanceFn = distSq) {
this.distanceFn = distanceFn;
const dim = pos.length;
assert(
dim > 0 && dim <= NdQuadtreeMap.MAX_DIM,
`illegal dimension: ${dim}`
);
assert(ext.length === dim, `pos/ext dimensions must be equal`);
__initChildOffsets(dim);
this.root = new NdQtNode(void 0, pos, ext);
this._size = 0;
pairs && this.into(pairs, -1);
}
static MAX_DIM = 16;
/**
* Returns a new point-based `NdQuadtreeMap` for nD keys in given region
* defined by `min` / `max` coordinates. The dimensionality of the tree is
* implicitly defined by the provided coordinates. Only points within that
* region can be indexed.
*
* @remarks
* Due to exponentially growing lookup tables, currently only supports up to
* 16 dimensions.
*/
static fromMinMax(min, max) {
return new NdQuadtreeMap(
addmN([], min, max, 0.5),
submN([], max, min, 0.5)
);
}
root;
_size;
get size() {
return this._size;
}
[Symbol.iterator]() {
return map((n) => [n.k, n.v], this.nodes());
}
keys() {
return map((n) => n.k, this.nodes());
}
values() {
return map((n) => n.v, this.nodes());
}
*nodes(all = false) {
let queue = [this.root];
while (queue.length) {
const n = queue.pop();
if (n) {
if (all || n.k) yield n;
if (n.children) queue = queue.concat(n.children);
}
}
}
copy() {
const tree = new NdQuadtreeMap(
this.root.pos,
this.root.ext,
this,
this.distanceFn
);
return tree;
}
clear() {
this.root.clear();
this._size = 0;
}
empty() {
return new NdQuadtreeMap(
this.root.pos,
this.root.ext,
void 0,
this.distanceFn
);
}
set(key, val, eps = EPS) {
if (this.root.set(key, val, eps, this.distanceFn)) {
this._size++;
return true;
}
return false;
}
into(pairs, eps = EPS) {
return __into(this, pairs, eps);
}
remove(p) {
let node = this.root.nodeForPoint(p);
if (!node) return false;
this._size--;
delete node.k;
delete node.v;
let doPrune = true;
while (node.parent) {
node = node.parent;
delete node.children[__childID(p, node.pos)];
doPrune = --node.numC === 0;
if (doPrune) delete node.children;
else break;
}
return true;
}
has(p, eps = EPS) {
return !!(eps <= 0 ? this.root.nodeForPoint(p) : this.root.queryKeys(p, eps, 1, [], this.distanceFn).length);
}
get(p, eps = EPS) {
if (eps <= 0) {
const node = this.root.nodeForPoint(p);
return node ? node.v : void 0;
}
return this.root.queryValues(p, eps, 1, [], this.distanceFn)[0];
}
query(p, r, max = 1, acc = []) {
return this.root.query(
(n) => [n.k, n.v],
p,
r,
max,
acc,
this.distanceFn
);
}
queryKeys(p, r, max = 1, acc = []) {
return this.root.queryKeys(p, r, max, acc, this.distanceFn);
}
queryValues(p, r, max = 1, acc = []) {
return this.root.queryValues(p, r, max, acc, this.distanceFn);
}
containsPoint(p) {
return this.root.containsPoint(p);
}
nodeForPoint(p) {
return this.root.nodeForPoint(p);
}
}
const MAX_CHILDREN = [
...take(
NdQuadtreeMap.MAX_DIM + 1,
iterate((x) => x * 2, 1)
)
];
const CHILD_OFFSETS = [];
const __initChildOffsets = (dim) => CHILD_OFFSETS[dim] || (CHILD_OFFSETS[dim] = [...permutations(...repeat([-1, 1], dim))]);
const __childID = vop(0);
__childID.add(1, (p, q) => p[0] >= q[0] ? 1 : 0);
__childID.add(2, (p, q) => (p[0] >= q[0] ? 2 : 0) | (p[1] >= q[1] ? 1 : 0));
__childID.add(
3,
(p, q) => (p[0] >= q[0] ? 4 : 0) | (p[1] >= q[1] ? 2 : 0) | (p[2] >= q[2] ? 1 : 0)
);
__childID.add(
4,
(p, q) => (p[0] >= q[0] ? 8 : 0) | (p[1] >= q[1] ? 4 : 0) | (p[2] >= q[2] ? 2 : 0) | (p[3] >= q[3] ? 1 : 0)
);
__childID.default((p, q) => {
let id = 0;
for (let i = 0, n = p.length - 1, bit = 1 << n; i <= n; i++, bit >>>= 1) {
p[i] >= q[i] && (id += bit);
}
return id;
});
export {
NdQtNode,
NdQuadtreeMap
};