Full Code of mourner/rbush-knn for AI

Repository: mourner/rbush-knn
Branch: main
Commit: 9e06d0fcd23b
Files: 10
Total size: 10.3 KB

Directory structure:
gitextract__3u9_1mp/

├── .github/
│   └── workflow/
│       └── node.yml
├── .gitignore
├── LICENSE
├── README.md
├── bench.js
├── eslint.config.js
├── index.js
├── package.json
├── rollup.config.js
└── test.js

================================================
FILE CONTENTS
================================================

================================================
FILE: .github/workflow/node.yml
================================================
name: Node
on: [push, pull_request]
jobs:
  test:
    runs-on: ubuntu-latest
    steps:
    - name: Checkout
      uses: actions/checkout@v4

    - name: Setup Node
      uses: actions/setup-node@v4
      with:
        node-version: 20

    - name: Install dependencies
      run: npm ci

    - name: Build the bundle
      run: npm run build

    - name: Run tests
      run: npm test


================================================
FILE: .gitignore
================================================
node_modules
rbush-knn.js
rbush-knn.min.js


================================================
FILE: LICENSE
================================================
Copyright (c) 2016, Vladimir Agafonkin

Permission to use, copy, modify, and/or distribute this software for any purpose
with or without fee is hereby granted, provided that the above copyright notice
and this permission notice appear in all copies.

THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND
FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER
TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF
THIS SOFTWARE.


================================================
FILE: README.md
================================================
## rbush-knn

_k_-nearest neighbors search for [RBush](https://github.com/mourner/rbush).
Implements a simple depth-first kNN search algorithm using a priority queue.

```js
import RBush from 'rbush';
import knn from 'rbush-knn';

const tree = new RBush(); // create RBush tree
tree.load(data); // bulk insert

const neighbors = knn(tree, 40, 40, 10); // return 10 nearest items around point [40, 40]
```

You can optionally pass a filter function to find k neighbors that satisfy a certain condition:

```js
const neighbors = knn(tree, 40, 40, 10, function (item) {
    return item.foo === 'bar';
});
```

### API

**knn(tree, x, y, [k, filterFn, maxDistance])**

- `tree`: an RBush tree
- `x`, `y`: query coordinates
- `k`: number of neighbors to search for (`Infinity` by default)
- `filterFn`: optional filter function; `k` nearest items where `filterFn(item) === true` will be returned.
- `maxDistance` (optional): maximum distance between neighbors and the query coordinates (`Infinity` by default)


================================================
FILE: bench.js
================================================
import RBush from 'rbush';
import knn from './index.js';

const N = 200000,
    M = 20000,
    K = 5;

const points = [];
const queries = [];
for (let i = 0; i < N; i++) {
    points.push(randPoint());
    queries.push(randPoint());
}

console.time(`load ${N} points`);
const tree = new RBush().load(points);
console.timeEnd(`load ${N} points`);

console.time(`knn query ${K} neighbors x ${M}`);
for (let i = 0; i < M; i++) {
    knn(tree, queries[i].minX, queries[i].minY, K);
}
console.timeEnd(`knn query ${K} neighbors x ${M}`);


console.time(`bbox query x ${  M}`);
for (let i = 0; i < M; i++) {
    tree.search(queries[i]);
}
console.timeEnd(`bbox query x ${M}`);

function randPoint() {
    const x = Math.floor(Math.random() * 100000),
        y = Math.floor(Math.random() * 100000);
    return {
        minX: x,
        minY: y,
        maxX: x,
        maxY: y
    };
}


================================================
FILE: eslint.config.js
================================================
export {default} from 'eslint-config-mourner';


================================================
FILE: index.js
================================================
import Queue from 'tinyqueue';

export default function knn(tree, x, y, n, predicate, maxDistance) {
    let node = tree.data;
    const result = [];
    const toBBox = tree.toBBox;

    const queue = new Queue(undefined, compareDist);

    while (node) {
        for (let i = 0; i < node.children.length; i++) {
            const child = node.children[i];
            const dist = boxDist(x, y, node.leaf ? toBBox(child) : child);
            if (!maxDistance || dist <= maxDistance * maxDistance) {
                queue.push({
                    node: child,
                    isItem: node.leaf,
                    dist
                });
            }
        }

        while (queue.length && queue.peek().isItem) {
            const candidate = queue.pop().node;
            if (!predicate || predicate(candidate))
                result.push(candidate);
            if (n && result.length === n) return result;
        }

        node = queue.pop();
        if (node) node = node.node;
    }

    return result;
}

function compareDist(a, b) {
    return a.dist - b.dist;
}

function boxDist(x, y, box) {
    const dx = axisDist(x, box.minX, box.maxX),
        dy = axisDist(y, box.minY, box.maxY);
    return dx * dx + dy * dy;
}

function axisDist(k, min, max) {
    return k < min ? min - k : k <= max ? 0 : k - max;
}


================================================
FILE: package.json
================================================
{
  "name": "rbush-knn",
  "version": "4.0.0",
  "description": "k-neareset neighbors search for RBush",
  "main": "rbush-knn.js",
  "module": "index.js",
  "browser": "rbush-knn.min.js",
  "jsdelivr": "rbush-knn.min.js",
  "unpkg": "rbush-knn.min.js",
  "exports": "./index.js",
  "scripts": {
    "bench": "node bench.js",
    "lint": "eslint index.js test.js bench.js",
    "pretest": "npm run lint",
    "test": "node --test",
    "build": "rollup -c",
    "prepare": "npm run build"
  },
  "keywords": [
    "rbush",
    "knn",
    "k-nearest neighbors",
    "data structure",
    "query"
  ],
  "author": "Vladimir Agafonkin",
  "license": "ISC",
  "type": "module",
  "devDependencies": {
    "@rollup/plugin-node-resolve": "^15.2.3",
    "@rollup/plugin-terser": "^0.4.4",
    "eslint": "^9.6.0",
    "eslint-config-mourner": "^4.0.1",
    "rbush": "^4.0.0",
    "rollup": "^4.18.0"
  },
  "dependencies": {
    "tinyqueue": "^2.0.3"
  },
  "files": [
    "index.js",
    "rbush-knn.js",
    "rbush-knn.min.js"
  ],
  "repository": "mourner/rbush-knn"
}


================================================
FILE: rollup.config.js
================================================
import terser from '@rollup/plugin-terser';
import resolve from '@rollup/plugin-node-resolve';

const output = (file, plugins) => ({
    input: 'index.js',
    output: {
        name: 'rbush-knn',
        format: 'umd',
        file
    },
    plugins
});

export default [
    output('rbush-knn.js', [resolve()]),
    output('rbush-knn.min.js', [resolve(), terser()])
];


================================================
FILE: test.js
================================================
import RBush from 'rbush';
import test from 'node:test';
import assert from 'node:assert/strict';

import knn from './index.js';

function rbush() {
    return new RBush();
}

/*eslint  @stylistic/js/comma-spacing: 0 */

function arrToBox(arr) {
    return {
        minX: arr[0],
        minY: arr[1],
        maxX: arr[2],
        maxY: arr[3]
    };
}

const data = [[87,55,87,56],[38,13,39,16],[7,47,8,47],[89,9,91,12],[4,58,5,60],[0,11,1,12],[0,5,0,6],[69,78,73,78],
    [56,77,57,81],[23,7,24,9],[68,24,70,26],[31,47,33,50],[11,13,14,15],[1,80,1,80],[72,90,72,91],[59,79,61,83],
    [98,77,101,77],[11,55,14,56],[98,4,100,6],[21,54,23,58],[44,74,48,74],[70,57,70,61],[32,9,33,12],[43,87,44,91],
    [38,60,38,60],[62,48,66,50],[16,87,19,91],[5,98,9,99],[9,89,10,90],[89,2,92,6],[41,95,45,98],[57,36,61,40],
    [50,1,52,1],[93,87,96,88],[29,42,33,42],[34,43,36,44],[41,64,42,65],[87,3,88,4],[56,50,56,52],[32,13,35,15],
    [3,8,5,11],[16,33,18,33],[35,39,38,40],[74,54,78,56],[92,87,95,90],[12,97,16,98],[76,39,78,40],[16,93,18,95],
    [62,40,64,42],[71,87,71,88],[60,85,63,86],[39,52,39,56],[15,18,19,18],[91,62,94,63],[10,16,10,18],[5,86,8,87],
    [85,85,88,86],[44,84,44,88],[3,94,3,97],[79,74,81,78],[21,63,24,66],[16,22,16,22],[68,97,72,97],[39,65,42,65],
    [51,68,52,69],[61,38,61,42],[31,65,31,65],[16,6,19,6],[66,39,66,41],[57,32,59,35],[54,80,58,84],[5,67,7,71],
    [49,96,51,98],[29,45,31,47],[31,72,33,74],[94,25,95,26],[14,7,18,8],[29,0,31,1],[48,38,48,40],[34,29,34,32],
    [99,21,100,25],[79,3,79,4],[87,1,87,5],[9,77,9,81],[23,25,25,29],[83,48,86,51],[79,94,79,95],[33,95,33,99],
    [1,14,1,14],[33,77,34,77],[94,56,98,59],[75,25,78,26],[17,73,20,74],[11,3,12,4],[45,12,47,12],[38,39,39,39],
    [99,3,103,5],[41,92,44,96],[79,40,79,41],[29,2,29,4]].map(arrToBox);

test('finds n neighbours', () => {
    const tree = rbush().load(data);
    const result = knn(tree, 40, 40, 10);
    assert.deepEqual(result, [[38,39,39,39],[35,39,38,40],[34,43,36,44],[29,42,33,42],[48,38,48,40],[31,47,33,50],[34,29,34,32],
        [29,45,31,47],[39,52,39,56],[57,36,61,40]].map(arrToBox));
});

test('does not throw if requesting too many items', () => {
    const tree = rbush().load(data);
    assert.doesNotThrow(() => {
        const result = knn(tree, 40, 40, 1000);
        assert.equal(result.length, data.length);
    });
});

test('finds all neighbors for maxDistance', () => {
    const tree = rbush().load(data);
    const result = knn(tree, 40, 40, 0, null, 10);
    assert.deepEqual(result, [[38,39,39,39],[35,39,38,40],[34,43,36,44],[29,42,33,42],[48,38,48,40],[31,47,33,50],[34,29,34,32]].map(arrToBox));
});

test('finds n neighbors for maxDistance', () => {
    const tree = rbush().load(data);
    const result = knn(tree, 40, 40, 1, null, 10);
    assert.deepEqual(result, [[38,39,39,39]].map(arrToBox));
});

test('does not throw if requesting too many items for maxDistance', () => {
    const tree = rbush().load(data);
    assert.doesNotThrow(() => {
        const result = knn(tree, 40, 40, 1000, null, 10);
        assert.deepEqual(result, [[38,39,39,39],[35,39,38,40],[34,43,36,44],[29,42,33,42],[48,38,48,40],[31,47,33,50],[34,29,34,32]].map(arrToBox));
    });
});

const pythData = [[0,0,0,0],[9,9,9,9],[12,12,12,12],[13,14,19,11]].map(arrToBox);

test('verify maxDistance excludes items too far away, in order to adhere to pythagoras theorem a^2+b^2=c^2', () => {
    const tree = rbush().load(pythData);
    // sqrt(9^2+9^2)~=12.727
    const result = knn(tree, 0, 0, 1000, null, 12.6);
    assert.deepEqual(result, [[0,0,0,0]].map(arrToBox));
});

test('verify maxDistance includes all items within range, in order to adhere to pythagoras theorem a^2+b^2=c^2', () => {
    const tree = rbush().load(pythData);
    // sqrt(9^2+9^2)~=12.727
    const result = knn(tree, 0, 0, 1000, null, 12.8);
    assert.deepEqual(result, [[0,0,0,0],[9,9,9,9]].map(arrToBox));
});

const richData = [[1,2,1,2],[3,3,3,3],[5,5,5,5],[4,2,4,2],[2,4,2,4],[5,3,5,3]].map((a, i) => {
    const item = arrToBox(a);
    item.version = i + 1;
    return item;
});

test('find n neighbours that do satisfy a given predicate', () => {
    const tree = rbush().load(richData);
    const result = knn(tree, 2, 4, 1, item => item.version < 5);
    if (result.length === 1) {
        const item = result[0];
        if (item.minX === 3 && item.minY === 3 && item.maxX === 3 && item.maxY === 3 && item.version === 2) {
            // OK
        } else {
            assert.fail(`Could not find the correct item, found ${  JSON.stringify(item)}`);
        }
    } else {
        assert.fail('Could not find the correct item');
    }
});