[ { "path": ".github/workflows/release.yml", "content": "name: release-upload\n\non:\n push:\n tags:\n - 'v*.*.*'\n\npermissions:\n id-token: write\n contents: read\n\njobs:\n sdist:\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v4\n - uses: actions/setup-python@v5\n with:\n python-version: '3.11'\n - name: Install build dependencies\n run: pip install build\n - name: Build sdist\n run: python -m build --sdist\n - name: Publish package distributions to PyPI\n uses: pypa/gh-action-pypi-publish@release/v1\n with:\n repository-url: https://upload.pypi.org/legacy/\n user: __token__\n password: ${{ secrets.PYPI_TOKEN }}\n" }, { "path": ".github/workflows/test.yml", "content": "name: test-upload\n\non:\n push:\n tags:\n - 'test*'\n\npermissions:\n id-token: write\n contents: read\n\njobs:\n sdist:\n runs-on: ubuntu-latest\n steps:\n - uses: actions/checkout@v4\n - uses: actions/setup-python@v5\n with:\n python-version: '3.11'\n - name: Install build dependencies\n run: pip install build\n - name: Build sdist\n run: python -m build --sdist\n - name: Publish package distributions to PyPI\n uses: pypa/gh-action-pypi-publish@release/v1\n with:\n repository-url: https://test.pypi.org/legacy/\n user: __token__\n password: ${{ secrets.TEST_PYPI_TOKEN }}\n" }, { "path": ".gitignore", "content": "/target\n\n# Byte-compiled / optimized / DLL files\n__pycache__/\n.pytest_cache/\n*.py[cod]\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\n.venv/\nenv/\nbin/\nbuild/\ndevelop-eggs/\ndist/\neggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\ninclude/\nman/\nvenv/\n*.egg-info/\n.installed.cfg\n*.egg\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\npip-selfcheck.json\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.coverage\n.cache\nnosetests.xml\ncoverage.xml\n\n# Translations\n*.mo\n\n# Mr Developer\n.mr.developer.cfg\n.project\n.pydevproject\n\n# Rope\n.ropeproject\n\n# Django stuff:\n*.log\n*.pot\n\n.DS_Store\n\n# Sphinx documentation\ndocs/_build/\n\n# PyCharm\n.idea/\n\n# VSCode\n.vscode/\n\n# Pyenv\n.python-version\n\n# tmp files\ntmp*\n*.out\n\n# tmp benchmark\n.benchmarks/\n\n# uv\nuv.lock\n" }, { "path": ".pre-commit-config.yaml", "content": "exclude: \"point-e/.*\"\nrepos:\n- repo: https://github.com/pre-commit/pre-commit-hooks\n rev: v4.4.0\n hooks:\n - id: trailing-whitespace\n - id: end-of-file-fixer\n - id: name-tests-test\n - id: requirements-txt-fixer\n- repo: https://github.com/pycqa/isort\n rev: 5.12.0\n hooks:\n - id: isort\n args: [\"--profile\", \"black\", \"--line-length=100\", \"--python-version=39\"]\n- repo: https://github.com/psf/black\n rev: 23.7.0\n hooks:\n - id: black\n args: [\"--line-length=100\", \"--target-version=py39\"]\n" }, { "path": "CMakeLists.txt", "content": "# Require CMake 3.15+ (matching scikit-build-core) Use new versions of all\n# policies up to CMake 4.0\ncmake_minimum_required(VERSION 3.15...4.0)\nset(CMAKE_EXPORT_COMPILE_COMMANDS ON)\nset(CMAKE_CXX_STANDARD 17)\nset(CMAKE_CXX_STANDARD_REQUIRED ON)\n\nproject(\n ${SKBUILD_PROJECT_NAME}\n VERSION ${SKBUILD_PROJECT_VERSION}\n LANGUAGES CXX)\n\nfind_package(pybind11 CONFIG REQUIRED)\n\npybind11_add_module(_fpsample MODULE src/lib.cpp)\ntarget_link_libraries(_fpsample PRIVATE pybind11::headers)\n\ntarget_compile_definitions(_fpsample PRIVATE VERSION_INFO=${PROJECT_VERSION})\ninstall(TARGETS _fpsample DESTINATION fpsample)\n" }, { "path": "LICENSE", "content": "MIT License\n\nCopyright (c) 2023 AyajiLin\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n" }, { "path": "README.md", "content": "# fpsample\n[![pypi package version badge](https://img.shields.io/pypi/v/fpsample)](https://pypi.org/project/fpsample/)\n![python version badge](https://img.shields.io/badge/python-%3E%3D3.7-blue)\n[![license badge](https://img.shields.io/github/license/leonardodalinky/fpsample)](https://github.com/leonardodalinky/fpsample/blob/main/LICENSE)\n[![star badge](https://img.shields.io/github/stars/leonardodalinky/fpsample?style=social)](https://github.com/leonardodalinky/fpsample)\n\nPython efficient farthest point sampling (FPS) library, 100x faster than `numpy` implementation.\n\n`fpsample` is coupled with `numpy` and built upon Rust pyo3 bindings. This library aims at achieving the best performance for FPS in single-threaded CPU environment.\n\n🎉 **PyTorch version with native multithreading, batch ops, Autograd and CUDA supports is in [pytorch_fpsample](https://github.com/leonardodalinky/pytorch_fpsample).**\n\n## Installation\n\n### Install from PyPI\n\n`numpy>=1.16.0` is required. Install `fpsample` using pip:\n\n```shell\npip install -U fpsample\n```\n\nIf you encounter any installation errors, please make an issue and try to compile from source.\n\n### Build from source\n\nSince Version 1.0.0, fpsample is built using [Scikit-build-core + pybind11](https://github.com/pybind/scikit_build_example). Therefore, you can build and install this library from source just with pip.\n\nC++ compiler must support C++17.\n\nCMake>=3.15 is required.\n\n```shell\ngit clone https://github.com/leonardodalinky/fpsample.git\ncd fpsample\npip install .\n```\n\n#### Direct porting of `QuickFPS`\n\nSee `src/warpper.hpp` and `src/_ext/` for details.\n\n## Usage\n\n```python\nimport fpsample\nimport numpy as np\n\n# Generate random point cloud\npc = np.random.rand(4096, 3)\n## sample 1024 points\n\n# Vanilla FPS\nfps_samples_idx = fpsample.fps_sampling(pc, 1024)\n\n# FPS + NPDU\nfps_npdu_samples_idx = fpsample.fps_npdu_sampling(pc, 1024)\n## or specify the windows size\nfps_npdu_samples_idx = fpsample.fps_npdu_sampling(pc, 1024, w=64)\n\n# FPS + NPDU + KDTree\nfps_npdu_kdtree_samples_idx = fpsample.fps_npdu_kdtree_sampling(pc, 1024)\n## or specify the windows size\nfps_npdu_kdtree_samples_idx = fpsample.fps_npdu_kdtree_sampling(pc, 1024, w=64)\n\n# KDTree-based FPS\nkdtree_fps_samples_idx = fpsample.bucket_fps_kdtree_sampling(pc, 1024)\n\n# NOTE: Probably the best\n# Bucket-based FPS or QuickFPS\nkdline_fps_samples_idx = fpsample.bucket_fps_kdline_sampling(pc, 1024, h=3)\n```\n\n* `FPS`: Vanilla farthest point sampling. Implemented in Rust. Achieve the same performance as `numpy`.\n* `FPS + NPDU`: Farthest point sampling with nearest-point-distance-updating (NPDU) heuristic strategy. 5x~10x faster than vanilla FPS. **Require dimensional locality and give sub-optimal answers**.\n* `FPS + NPDU + KDTree`: Farthest point sampling with NPDU heuristic strategy and KDTree. 3x~8x faster than vanilla FPS. Slightly slower than `FPS + NPDU`. But **DOES NOT** require dimensional locality.\n* `KDTree-based FPS`: A farthest point sampling algorithm based on KDTree. About 40~50x faster than vanilla FPS.\n* `Bucket-based FPS` or `QuickFPS`: A bucket-based farthest point sampling algorithm. About 80~100x faster than vanilla FPS. Require an additional hyperparameter for the height of the KDTree. In practice, `h=3` or `h=5` is recommended for small data, `h=7` is recommended for medium data, and `h=9` for extremely large data.\n\n> **NOTE**: 🔥 In most cases, `Bucket-based FPS` is the best choice, with proper hyperparameter setting.\n\n### Determinism\n\nFor deterministic results, fix the first sampled point index by passing the `start_idx` parameter.\n```python\nkdline_fps_samples_idx = fpsample.bucket_fps_kdline_sampling(pc, 1024, h=3, start_idx=0)\n```\n\n**OR** set the random seed before calling the function.\n```python\nnp.random.seed(42)\n```\n\n## Development\n\nInstall dependencies:\n```shell\nuv sync\n```\n\n## Performance\nSetup:\n - CPU: Intel(R) Xeon(R) Gold 6326 CPU @ 2.90GHz\n - RAM: 512 GiB\n - SYSTEM: Ubuntu 20.04.6 LTS\n\nRun benchmark:\n```shell\npy.test bench/ --benchmark-columns=mean,stddev --benchmark-sort=mean\n```\n\nResults:\n```\n------------------------------------------------------------------------------- benchmark '1024 of 4096': 8 tests -------------------------------------------------------------------------------\nName (time in ms) Min Max Mean StdDev Median IQR Outliers OPS Rounds Iterations\n-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------\ntest_bucket_fps_kdline_4k_h5 1.7892 (1.0) 2.2010 (1.01) 1.8301 (1.0) 0.0273 (1.20) 1.8255 (1.0) 0.0182 (1.0) 53;31 546.4295 (1.0) 491 1\ntest_bucket_fps_kdline_4k_h3 2.0452 (1.14) 2.1714 (1.0) 2.0862 (1.14) 0.0227 (1.0) 2.0817 (1.14) 0.0220 (1.21) 99;34 479.3375 (0.88) 436 1\ntest_bucket_fps_kdline_4k_h7 3.0348 (1.70) 3.2599 (1.50) 3.1082 (1.70) 0.0391 (1.72) 3.1005 (1.70) 0.0568 (3.12) 107;4 321.7322 (0.59) 306 1\ntest_fps_npdu_4k 3.9024 (2.18) 4.3511 (2.00) 3.9514 (2.16) 0.0395 (1.74) 3.9457 (2.16) 0.0283 (1.55) 29;14 253.0768 (0.46) 245 1\ntest_bucket_fps_kdtree_4k 6.5697 (3.67) 7.3863 (3.40) 6.7525 (3.69) 0.1150 (5.07) 6.7240 (3.68) 0.0926 (5.08) 26;11 148.0923 (0.27) 126 1\ntest_fps_npdu_kdtree_4k 9.4114 (5.26) 9.7609 (4.50) 9.4979 (5.19) 0.0613 (2.70) 9.4807 (5.19) 0.0362 (1.99) 14;13 105.2866 (0.19) 97 1\ntest_vanilla_fps_4k 12.3702 (6.91) 12.8378 (5.91) 12.5875 (6.88) 0.1840 (8.11) 12.5024 (6.85) 0.3712 (20.38) 40;0 79.4438 (0.15) 77 1\ntest_vanilla_fps_4k_multiple 12.4058 (6.93) 13.2153 (6.09) 12.4706 (6.81) 0.1639 (7.22) 12.4139 (6.80) 0.0268 (1.47) 7;12 80.1886 (0.15) 80 1\n-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------\n\n--------------------------------------------------------------------------------- benchmark '4096 of 50000': 7 tests --------------------------------------------------------------------------------\nName (time in ms) Min Max Mean StdDev Median IQR Outliers OPS Rounds Iterations\n-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------\ntest_bucket_fps_kdline_50k_h7 24.7285 (1.0) 25.2791 (1.0) 25.0179 (1.0) 0.1237 (1.0) 25.0139 (1.0) 0.1184 (1.0) 10;3 39.9714 (1.0) 40 1\ntest_bucket_fps_kdline_50k_h5 31.7191 (1.28) 32.9791 (1.30) 32.1800 (1.29) 0.3308 (2.68) 32.0584 (1.28) 0.3701 (3.13) 6;3 31.0752 (0.78) 31 1\ntest_bucket_fps_kdline_50k_h3 66.7806 (2.70) 67.5433 (2.67) 67.1429 (2.68) 0.2265 (1.83) 67.0898 (2.68) 0.2654 (2.24) 5;0 14.8936 (0.37) 15 1\ntest_bucket_fps_kdtree_50k 82.7767 (3.35) 85.9608 (3.40) 84.5684 (3.38) 0.9575 (7.74) 84.5108 (3.38) 1.2572 (10.62) 3;0 11.8248 (0.30) 11 1\ntest_fps_npdu_50k 179.8323 (7.27) 181.9054 (7.20) 180.5911 (7.22) 0.7326 (5.92) 180.4329 (7.21) 0.7742 (6.54) 2;0 5.5374 (0.14) 6 1\ntest_fps_npdu_kdtree_50k 253.4978 (10.25) 255.2537 (10.10) 254.2810 (10.16) 0.7408 (5.99) 253.9220 (10.15) 1.1832 (9.99) 2;0 3.9327 (0.10) 5 1\ntest_vanilla_fps_50k 592.6405 (23.97) 594.1419 (23.50) 593.2049 (23.71) 0.5970 (4.83) 593.2359 (23.72) 0.7815 (6.60) 1;0 1.6858 (0.04) 5 1\n-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------\n\n---------------------------------------------------------------------------------------- benchmark '50000 of 100000': 7 tests ---------------------------------------------------------------------------------------\nName (time in ms) Min Max Mean StdDev Median IQR Outliers OPS Rounds Iterations\n---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------\ntest_bucket_fps_kdline_100k_h7 232.5230 (1.0) 234.3285 (1.0) 233.3726 (1.0) 0.7632 (1.49) 233.5817 (1.0) 1.2734 (1.66) 2;0 4.2850 (1.0) 5 1\ntest_bucket_fps_kdtree_100k 381.3098 (1.64) 388.0433 (1.66) 384.0891 (1.65) 3.5171 (6.84) 382.9141 (1.64) 5.0501 (6.56) 1;0 2.6036 (0.61) 3 1\ntest_bucket_fps_kdline_100k_h9 409.0352 (1.76) 420.7832 (1.80) 414.3774 (1.78) 5.9458 (11.57) 413.3137 (1.77) 8.8110 (11.45) 1;0 2.4133 (0.56) 3 1\ntest_bucket_fps_kdline_100k_h5 410.1506 (1.76) 411.1764 (1.75) 410.6451 (1.76) 0.5139 (1.0) 410.6085 (1.76) 0.7693 (1.0) 1;0 2.4352 (0.57) 3 1\ntest_fps_npdu_kdtree_100k 3,134.7921 (13.48) 3,137.0544 (13.39) 3,135.7064 (13.44) 1.1919 (2.32) 3,135.2726 (13.42) 1.6967 (2.21) 1;0 0.3189 (0.07) 3 1\ntest_fps_npdu_100k 4,325.4402 (18.60) 4,341.7247 (18.53) 4,333.3156 (18.57) 8.1554 (15.87) 4,332.7821 (18.55) 12.2134 (15.87) 1;0 0.2308 (0.05) 3 1\ntest_vanilla_fps_100k 14,460.2331 (62.19) 14,664.4396 (62.58) 14,561.1574 (62.39) 102.1237 (198.73) 14,558.7995 (62.33) 153.1549 (199.07) 1;0 0.0687 (0.02) 3 1\n---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------\n\n```\n\n## Reference\nThe nearest-point-distance-updating (NPDU) heuristic strategy is proposed in the following paper:\n```bibtex\n@INPROCEEDINGS{Li2022adjust,\n author={Li, Jingtao and Zhou, Jian and Xiong, Yan and Chen, Xing and Chakrabarti, Chaitali},\n booktitle={2022 IEEE Workshop on Signal Processing Systems (SiPS)},\n title={An Adjustable Farthest Point Sampling Method for Approximately-sorted Point Cloud Data},\n year={2022},\n volume={},\n number={},\n pages={1-6},\n doi={10.1109/SiPS55645.2022.9919246}\n}\n```\n\nBucket-based farthest point sampling (QuickFPS) is proposed in the following paper. The implementation is based on the author's [Repo](https://github.com/hanm2019/bucket-based_farthest-point-sampling_CPU). To port the implementation to other C++ program, check [this](https://github.com/leonardodalinky/fpsample/tree/main/src/bucket_fps/_ext) for details.\n```bibtex\n@article{han2023quickfps,\n title={QuickFPS: Architecture and Algorithm Co-Design for Farthest Point Sampling in Large-Scale Point Clouds},\n author={Han, Meng and Wang, Liang and Xiao, Limin and Zhang, Hao and Zhang, Chenhao and Xu, Xiangrong and Zhu, Jianfeng},\n journal={IEEE Transactions on Computer-Aided Design of Integrated Circuits and Systems},\n year={2023},\n publisher={IEEE}\n}\n```\n\nThe KDTree-based NN search algorithm implemented by [nanoflann](https://github.com/jlblancoc/nanoflann)\n\nThanks to the authors for their great work.\n" }, { "path": "bench/test_bench.py", "content": "import numpy as np\nimport pytest\n\nimport fpsample\n\nTEST_SEED = 42\nTEST_BENCHMARK_SETTINGS = {\n \"4k\": {\"group\": \"1024 of 4096\", \"warmup\": False},\n \"50k\": {\"group\": \"4096 of 50000\", \"warmup\": False},\n \"100k\": {\"group\": \"50000 of 100000\", \"warmup\": False, \"min_rounds\": 3},\n}\nTEST_CASE_SETTINGS = {\n \"4k\": (4096, 1024, 3),\n \"50k\": (50000, 4096, 3),\n \"100k\": (100_000, 50_000, 3),\n}\n\n\ndef create_sample_data(n_points: int, n_dim: int = 3, seed: int = TEST_SEED):\n np.random.seed(seed)\n return np.random.rand(n_points, n_dim)\n\n\n####################\n# #\n# 4k setting #\n# #\n####################\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"4k\"])\ndef test_vanilla_fps_4k(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"4k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.fps_sampling, pc, n_samples)\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"4k\"])\ndef test_vanilla_fps_4k_multiple(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"4k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.fps_sampling, pc, n_samples, start_idx=[0, 23, 64, 128])\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"4k\"])\ndef test_fps_npdu_4k(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"4k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.fps_npdu_sampling, pc, n_samples)\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"4k\"])\ndef test_fps_npdu_kdtree_4k(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"4k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.fps_npdu_kdtree_sampling, pc, n_samples)\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"4k\"])\ndef test_bucket_fps_kdtree_4k(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"4k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.bucket_fps_kdtree_sampling, pc, n_samples)\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"4k\"])\ndef test_bucket_fps_kdline_4k_h3(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"4k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.bucket_fps_kdline_sampling, pc, n_samples, 3)\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"4k\"])\ndef test_bucket_fps_kdline_4k_h5(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"4k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.bucket_fps_kdline_sampling, pc, n_samples, 5)\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"4k\"])\ndef test_bucket_fps_kdline_4k_h7(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"4k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.bucket_fps_kdline_sampling, pc, n_samples, 7)\n\n\n#####################\n# #\n# 50k setting #\n# #\n#####################\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"50k\"])\ndef test_vanilla_fps_50k(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"50k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.fps_sampling, pc, n_samples)\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"50k\"])\ndef test_fps_npdu_50k(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"50k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.fps_npdu_sampling, pc, n_samples)\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"50k\"])\ndef test_fps_npdu_kdtree_50k(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"50k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.fps_npdu_kdtree_sampling, pc, n_samples)\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"50k\"])\ndef test_bucket_fps_kdtree_50k(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"50k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.bucket_fps_kdtree_sampling, pc, n_samples)\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"50k\"])\ndef test_bucket_fps_kdline_50k_h3(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"50k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.bucket_fps_kdline_sampling, pc, n_samples, 3)\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"50k\"])\ndef test_bucket_fps_kdline_50k_h5(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"50k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.bucket_fps_kdline_sampling, pc, n_samples, 5)\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"50k\"])\ndef test_bucket_fps_kdline_50k_h7(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"50k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.bucket_fps_kdline_sampling, pc, n_samples, 7)\n\n\n######################\n# #\n# 100k setting #\n# #\n######################\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"100k\"])\ndef test_vanilla_fps_100k(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"100k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.fps_sampling, pc, n_samples)\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"100k\"])\ndef test_fps_npdu_100k(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"100k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.fps_npdu_sampling, pc, n_samples)\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"100k\"])\ndef test_fps_npdu_kdtree_100k(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"100k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.fps_npdu_kdtree_sampling, pc, n_samples)\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"100k\"])\ndef test_bucket_fps_kdtree_100k(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"100k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.bucket_fps_kdtree_sampling, pc, n_samples)\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"100k\"])\ndef test_bucket_fps_kdline_100k_h5(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"100k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.bucket_fps_kdline_sampling, pc, n_samples, 5)\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"100k\"])\ndef test_bucket_fps_kdline_100k_h7(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"100k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.bucket_fps_kdline_sampling, pc, n_samples, 7)\n\n\n@pytest.mark.benchmark(**TEST_BENCHMARK_SETTINGS[\"100k\"])\ndef test_bucket_fps_kdline_100k_h9(benchmark):\n n_points, n_samples, n_dim = TEST_CASE_SETTINGS[\"100k\"]\n pc = create_sample_data(n_points, n_dim)\n benchmark(fpsample.bucket_fps_kdline_sampling, pc, n_samples, 9)\n" }, { "path": "pyproject.toml", "content": "[build-system]\nrequires = [\"scikit-build-core>=0.11\", \"pybind11>=2.6\"]\nbuild-backend = \"scikit_build_core.build\"\n\n[project]\nname = \"fpsample\"\nversion = \"1.0.2\"\nauthors = [\n { name = \"Leonard Lin\", email = \"leonard.keilin@gmail.com\" },\n { name = \"guch8017\", email = \"guch8017@gmail.com\" }\n]\ndescription = \"An efficient CPU implementation of farthest point sampling (FPS) for point clouds.\"\nreadme = \"README.md\"\nlicense = { file = \"LICENSE\" }\nrequires-python = \">=3.8\"\nclassifiers = [\n \"Programming Language :: Python :: Implementation :: CPython\",\n \"Programming Language :: Python :: Implementation :: PyPy\",\n \"Programming Language :: Python :: 3 :: Only\",\n]\ndependencies = [\n \"numpy>=1.16.0\",\n \"pybind11>=2.6\",\n]\n\n[project.urls]\nRepository = \"https://github.com/leonardodalinky/fpsample\"\nTracker = \"https://github.com/leonardodalinky/fpsample/issues\"\n\n[dependency-groups]\ntest = [\"pytest\"]\ndev = [\n { include-group = \"test\" },\n \"jupyterlab>=4.3.8\",\n \"pytest-benchmark>=4.0.0\",\n]\n\n\n[tool.scikit-build]\nminimum-version = \"build-system.requires\"\n" }, { "path": "src/_ext/Interval.h", "content": "//\n// Created by 韩萌 on 2022/6/14.\n// Refactored by AyajiLin on 2023/9/16.\n//\n\n#ifndef KD_TREE_BASED_FARTHEST_POINT_SAMPLING_INTERVAL_HPP\n#define KD_TREE_BASED_FARTHEST_POINT_SAMPLING_INTERVAL_HPP\n\nnamespace quickfps {\ntemplate class Interval {\n public:\n S low, high;\n Interval() : low(0), high(0){};\n Interval(S low, S high) : low(low), high(high){};\n Interval(const Interval &o) : low(o.low), high(o.high){};\n};\n} // namespace quickfps\n\n#endif // KD_TREE_BASED_FARTHEST_POINT_SAMPLING_INTERVAL_HPP\n" }, { "path": "src/_ext/KDLineTree.h", "content": "//\n// Created by hanm on 22-6-15.\n// Refactored by AyajiLin on 2023/9/16.\n//\n\n#ifndef FPS_CPU_KDLINETREE_H\n#define FPS_CPU_KDLINETREE_H\n\n#include \n#include \n\n#include \"KDTreeBase.h\"\n\nnamespace quickfps {\n\ntemplate \nclass KDLineTree : public KDTreeBase {\n public:\n using typename KDTreeBase::_Point;\n using typename KDTreeBase::_Points;\n using typename KDTreeBase::NodePtr;\n\n KDLineTree(_Points data, size_t pointSize, size_t treeHigh,\n _Points samplePoints);\n ~KDLineTree();\n\n std::vector KDNode_list;\n\n size_t high_;\n\n _Point max_point() override;\n\n void update_distance(const _Point &ref_point) override;\n\n void sample(size_t sample_num) override;\n\n bool leftNode(size_t high, size_t count) const override {\n return high == this->high_ || count == 1;\n };\n\n void addNode(NodePtr p) override;\n};\n\ntemplate \nKDLineTree::KDLineTree(_Points data, size_t pointSize,\n size_t treeHigh, _Points samplePoints)\n : KDTreeBase(data, pointSize, samplePoints), high_(treeHigh) {\n KDNode_list.clear();\n}\n\ntemplate \nKDLineTree::~KDLineTree() {\n KDNode_list.clear();\n}\n\ntemplate \ntypename KDLineTree::_Point KDLineTree::max_point() {\n _Point tmpPoint;\n S max_distance = std::numeric_limits::lowest();\n for (const auto &bucket : KDNode_list) {\n if (bucket->max_point.dis > max_distance) {\n max_distance = bucket->max_point.dis;\n tmpPoint = bucket->max_point;\n }\n }\n return tmpPoint;\n}\n\ntemplate \nvoid KDLineTree::update_distance(const _Point &ref_point) {\n for (const auto &bucket : KDNode_list) {\n bucket->send_delay_point(ref_point);\n bucket->update_distance();\n }\n}\n\ntemplate \nvoid KDLineTree::sample(size_t sample_num) {\n _Point ref_point;\n for (size_t i = 1; i < sample_num; i++) {\n ref_point = this->max_point();\n this->sample_points[i] = ref_point;\n this->update_distance(ref_point);\n }\n}\n\ntemplate \nvoid KDLineTree::addNode(NodePtr p) {\n size_t nodeIdx = KDNode_list.size();\n p->idx = nodeIdx;\n KDNode_list.push_back(p);\n}\n\n} // namespace quickfps\n\n#endif // FPS_CPU_KDLINETREE_H\n" }, { "path": "src/_ext/KDNode.h", "content": "//\n// Created by 韩萌 on 2022/6/14.\n// Refactored by AyajiLin on 2023/9/16.\n//\n\n#ifndef KD_TREE_BASED_FARTHEST_POINT_SAMPLING_KDNODE_H\n#define KD_TREE_BASED_FARTHEST_POINT_SAMPLING_KDNODE_H\n#include \n#include \n#include \n#include \n\n#include \"Interval.h\"\n#include \"Point.h\"\n\nnamespace quickfps {\n\ntemplate class KDNode {\n public:\n using _Point = Point;\n using _Points = _Point *;\n _Points points;\n size_t pointLeft, pointRight;\n size_t idx;\n\n std::array, DIM> bboxs;\n std::vector<_Point> waitpoints;\n std::vector<_Point> delaypoints;\n _Point max_point;\n KDNode *left;\n KDNode *right;\n\n KDNode();\n\n KDNode(const KDNode &a);\n\n KDNode(const std::array, DIM> &bboxs);\n\n void init(const _Point &ref);\n\n void updateMaxPoint(const _Point &lpoint, const _Point &rpoint) {\n if (lpoint.dis > rpoint.dis)\n this->max_point = lpoint;\n else\n this->max_point = rpoint;\n }\n\n S bound_distance(const _Point &ref_point) const;\n\n void send_delay_point(const _Point &point) {\n this->waitpoints.push_back(point);\n }\n\n void update_distance();\n\n void reset();\n\n size_t size() const;\n};\n\ntemplate \nKDNode::KDNode()\n : points(nullptr), pointLeft(0), pointRight(0), left(nullptr),\n right(nullptr) {}\n\ntemplate \nKDNode::KDNode(const std::array, DIM> &other_bboxs)\n : points(nullptr), pointLeft(0), pointRight(0), left(nullptr),\n right(nullptr) {\n std::copy(other_bboxs.cbegin(), other_bboxs.cend(), this->bboxs.begin());\n}\n\ntemplate \nKDNode::KDNode(const KDNode &a)\n : points(a.points), pointLeft(a.pointLeft), pointRight(a.pointRight),\n left(a.left), right(a.right), idx(a.idx) {\n std::copy(a.bboxs.cbegin(), a.bboxs.cend(), this->bboxs.begin());\n std::copy(a.waitpoints.cbegin(), a.waitpoints.cend(),\n this->waitpoints.begin());\n std::copy(a.delaypoints.cbegin(), a.delaypoints.cend(),\n this->delaypoints.begin());\n}\n\ntemplate \nvoid KDNode::init(const _Point &ref) {\n waitpoints.clear();\n delaypoints.clear();\n if (this->left && this->right) {\n this->left->init(ref);\n this->right->init(ref);\n updateMaxPoint(this->left->max_point, this->right->max_point);\n } else {\n S dis;\n S maxdis = std::numeric_limits::lowest();\n for (size_t i = pointLeft; i < pointRight; i++) {\n dis = points[i].updatedistance(ref);\n if (dis > maxdis) {\n maxdis = dis;\n max_point = points[i];\n }\n }\n }\n}\n\ntemplate \nS KDNode::bound_distance(const _Point &ref_point) const {\n S bound_dis(0);\n S dim_distance;\n for (size_t cur_dim = 0; cur_dim < DIM; cur_dim++) {\n dim_distance = 0;\n if (ref_point[cur_dim] > this->bboxs[cur_dim].high)\n dim_distance = ref_point[cur_dim] - this->bboxs[cur_dim].high;\n else if (ref_point[cur_dim] < this->bboxs[cur_dim].low)\n dim_distance = this->bboxs[cur_dim].low - ref_point[cur_dim];\n bound_dis += powi(dim_distance, 2);\n }\n return bound_dis;\n}\n\ntemplate \nvoid KDNode::update_distance() {\n for (const auto &ref_point : this->waitpoints) {\n S lastmax_distance = this->max_point.dis;\n S cur_distance = this->max_point.distance(ref_point);\n // cur_distance >\n // lastmax_distance意味着当前Node的max_point不会进行更新\n if (cur_distance > lastmax_distance) {\n S boundary_distance = bound_distance(ref_point);\n if (boundary_distance < lastmax_distance)\n this->delaypoints.push_back(ref_point);\n } else {\n if (this->right && this->left) {\n if (!delaypoints.empty()) {\n for (const auto &delay_point : delaypoints) {\n this->left->send_delay_point(delay_point);\n this->right->send_delay_point(delay_point);\n }\n delaypoints.clear();\n }\n this->left->send_delay_point(ref_point);\n this->left->update_distance();\n\n this->right->send_delay_point(ref_point);\n this->right->update_distance();\n\n updateMaxPoint(this->left->max_point, this->right->max_point);\n } else {\n S dis;\n S maxdis;\n this->delaypoints.push_back(ref_point);\n for (const auto &delay_point : delaypoints) {\n maxdis = std::numeric_limits::lowest();\n for (size_t i = pointLeft; i < pointRight; i++) {\n dis = points[i].updatedistance(delay_point);\n if (dis > maxdis) {\n maxdis = dis;\n max_point = points[i];\n }\n }\n }\n this->delaypoints.clear();\n }\n }\n }\n this->waitpoints.clear();\n}\n\ntemplate void KDNode::reset() {\n for (size_t i = pointLeft; i < pointRight; i++) {\n points[i].reset();\n }\n this->waitpoints.clear();\n this->delaypoints.clear();\n this->max_point.reset();\n if (this->left && this->right) {\n this->left->reset();\n this->right->reset();\n }\n}\n\ntemplate \nsize_t KDNode::size() const {\n if (this->left && this->right)\n return this->left->size() + this->right->size();\n return (pointRight - pointLeft);\n}\n\n} // namespace quickfps\n\n#endif // KD_TREE_BASED_FARTHEST_POINT_SAMPLING_KDNODE_H\n" }, { "path": "src/_ext/KDTree.h", "content": "//\n// Created by hanm on 22-6-15.\n// Refactored by AyajiLin on 2023/9/16.\n//\n\n#ifndef FPS_CPU_KDTREE_H\n#define FPS_CPU_KDTREE_H\n\n#include \"KDTreeBase.h\"\n\nnamespace quickfps {\n\ntemplate \nclass KDTree : public KDTreeBase {\n public:\n using typename KDTreeBase::_Point;\n using typename KDTreeBase::_Points;\n using typename KDTreeBase::NodePtr;\n explicit KDTree(_Points data, size_t pointSize, _Points samplePoints);\n\n _Point max_point() override { return this->root_->max_point; };\n\n void update_distance(const _Point &ref_point) override;\n\n void sample(size_t sample_num) override;\n\n bool leftNode(size_t, size_t count) const override { return count == 1; };\n\n void addNode(NodePtr) override{};\n};\n\ntemplate \nKDTree::KDTree(_Points data, size_t pointSize, _Points samplePoints)\n : KDTreeBase(data, pointSize, samplePoints) {}\n\ntemplate \nvoid KDTree::update_distance(const _Point &ref_point) {\n this->root_->send_delay_point(ref_point);\n this->root_->update_distance();\n}\n\ntemplate \nvoid KDTree::sample(size_t sample_num) {\n _Point ref_point;\n for (size_t i = 1; i < sample_num; i++) {\n ref_point = this->max_point();\n this->sample_points[i] = ref_point;\n this->update_distance(ref_point);\n }\n}\n\n} // namespace quickfps\n\n#endif // FPS_CPU_KDTREE_H\n" }, { "path": "src/_ext/KDTreeBase.h", "content": "//\n// Created by 韩萌 on 2022/6/14.\n// Refactored by AyajiLin on 2023/9/16.\n//\n\n#ifndef KD_TREE_BASED_FARTHEST_POINT_SAMPLING_KDTREE_H\n#define KD_TREE_BASED_FARTHEST_POINT_SAMPLING_KDTREE_H\n\n#include \"KDNode.h\"\n#include \"Point.h\"\n#include \n#include \n#include \n\nnamespace quickfps {\n\ntemplate class KDTreeBase {\n public:\n using _Point = Point;\n using _Points = _Point *;\n using NodePtr = KDNode *;\n using _Interval = Interval;\n\n size_t pointSize;\n _Points sample_points;\n NodePtr root_;\n _Points points_;\n\n public:\n KDTreeBase(_Points data, size_t pointSize, _Points samplePoints);\n\n ~KDTreeBase();\n\n void buildKDtree();\n\n NodePtr get_root() const { return this->root_; };\n\n void init(const _Point &ref);\n\n virtual _Point max_point() = 0;\n\n virtual void sample(size_t sample_num) = 0;\n\n protected:\n void deleteNode(NodePtr node_p);\n virtual void addNode(NodePtr p) = 0;\n virtual bool leftNode(size_t high, size_t count) const = 0;\n virtual void update_distance(const _Point &ref_point) = 0;\n\n NodePtr divideTree(ssize_t left, ssize_t right,\n const std::array<_Interval, DIM> &bboxs,\n size_t curr_high);\n\n size_t planeSplit(ssize_t left, ssize_t right, size_t split_dim,\n T split_val);\n\n T qSelectMedian(size_t dim, size_t left, size_t right);\n static size_t findSplitDim(const std::array<_Interval, DIM> &bboxs);\n inline std::array<_Interval, DIM> computeBoundingBox(size_t left,\n size_t right);\n};\n\ntemplate \nKDTreeBase::KDTreeBase(_Points data, size_t pointSize,\n _Points samplePoints)\n : pointSize(pointSize), sample_points(samplePoints), root_(nullptr),\n points_(data) {}\n\ntemplate \nKDTreeBase::~KDTreeBase() {\n if (root_ != nullptr)\n deleteNode(root_);\n}\n\ntemplate \nvoid KDTreeBase::deleteNode(NodePtr node_p) {\n if (node_p->left)\n deleteNode(node_p->left);\n if (node_p->right)\n deleteNode(node_p->right);\n delete node_p;\n}\n\ntemplate \nvoid KDTreeBase::buildKDtree() {\n size_t left = 0;\n size_t right = pointSize;\n std::array<_Interval, DIM> bboxs = this->computeBoundingBox(left, right);\n this->root_ = divideTree(left, right, bboxs, 0);\n}\n\ntemplate \ntypename KDTreeBase::NodePtr\nKDTreeBase::divideTree(ssize_t left, ssize_t right,\n const std::array<_Interval, DIM> &bboxs,\n size_t curr_high) {\n NodePtr node = new KDNode(bboxs);\n\n ssize_t count = right - left;\n if (this->leftNode(curr_high, count)) {\n node->pointLeft = left;\n node->pointRight = right;\n node->points = this->points_;\n this->addNode(node);\n return node;\n } else {\n size_t split_dim = this->findSplitDim(bboxs);\n T split_val = this->qSelectMedian(split_dim, left, right);\n\n size_t split_delta = planeSplit(left, right, split_dim, split_val);\n\n std::array<_Interval, DIM> bbox_cur =\n this->computeBoundingBox(left, left + split_delta);\n node->left =\n this->divideTree(left, left + split_delta, bbox_cur, curr_high + 1);\n bbox_cur = this->computeBoundingBox(left + split_delta, right);\n node->right = this->divideTree(left + split_delta, right, bbox_cur,\n curr_high + 1);\n return node;\n }\n}\n\ntemplate \nsize_t KDTreeBase::planeSplit(ssize_t left, ssize_t right,\n size_t split_dim, T split_val) {\n ssize_t start = left;\n ssize_t end = right - 1;\n\n for (;;) {\n while (start <= end && points_[start].pos[split_dim] < split_val)\n ++start;\n while (start <= end && points_[end].pos[split_dim] >= split_val)\n --end;\n\n if (start > end)\n break;\n std::swap(points_[start], points_[end]);\n ++start;\n --end;\n }\n\n ssize_t lim1 = start - left;\n if (start == left)\n lim1 = 1;\n if (start == right)\n lim1 = (right - left - 1);\n\n return static_cast(lim1);\n}\n\ntemplate \nT KDTreeBase::qSelectMedian(size_t dim, size_t left, size_t right) {\n T sum = std::accumulate(this->points_ + left, this->points_ + right, 0.0,\n [dim](const T &acc, const _Point &point) {\n return acc + point.pos[dim];\n });\n return sum / (right - left);\n}\n\ntemplate \nsize_t\nKDTreeBase::findSplitDim(const std::array<_Interval, DIM> &bboxs) {\n T min_, max_;\n T span = 0;\n size_t best_dim = 0;\n\n for (size_t cur_dim = 0; cur_dim < DIM; cur_dim++) {\n min_ = bboxs[cur_dim].low;\n max_ = bboxs[cur_dim].high;\n T cur_span = (max_ - min_);\n\n if (cur_span > span) {\n best_dim = cur_dim;\n span = cur_span;\n }\n }\n\n return best_dim;\n}\n\ntemplate \ninline std::array, DIM>\nKDTreeBase::computeBoundingBox(size_t left, size_t right) {\n T min_vals[DIM];\n T max_vals[DIM];\n std::fill(min_vals, min_vals + DIM, std::numeric_limits::max());\n std::fill(max_vals, max_vals + DIM, std::numeric_limits::lowest());\n\n for (size_t i = left; i < right; ++i) {\n const _Point &pos = points_[i];\n\n for (size_t cur_dim = 0; cur_dim < DIM; cur_dim++) {\n T val = pos[cur_dim];\n min_vals[cur_dim] = std::min(min_vals[cur_dim], val);\n max_vals[cur_dim] = std::max(max_vals[cur_dim], val);\n }\n }\n\n std::array<_Interval, DIM> bboxs;\n\n for (size_t cur_dim = 0; cur_dim < DIM; cur_dim++) {\n bboxs[cur_dim].low = min_vals[cur_dim];\n bboxs[cur_dim].high = max_vals[cur_dim];\n }\n\n return bboxs;\n}\n\ntemplate \nvoid KDTreeBase::init(const _Point &ref) {\n this->sample_points[0] = ref;\n this->root_->init(ref);\n}\n\n} // namespace quickfps\n\n#endif // KD_TREE_BASED_FARTHEST_POINT_SAMPLING_KDTREE_H\n" }, { "path": "src/_ext/Point.h", "content": "//\n// Created by 韩萌 on 2022/6/14.\n// Refactored by AyajiLin on 2023/9/16.\n//\n\n#ifndef KD_TREE_BASED_FARTHEST_POINT_SAMPLING_POINT_H\n#define KD_TREE_BASED_FARTHEST_POINT_SAMPLING_POINT_H\n\n#include \"cmath\"\n#include \"utils.h\"\n#include \n#include \n\nnamespace quickfps {\n\ntemplate class Point {\n public:\n T pos[DIM]; // x, y, z\n S dis;\n size_t id;\n\n Point();\n Point(const T pos[DIM], size_t id);\n Point(const T pos[DIM], size_t id, S dis);\n Point(const Point &obj);\n ~Point(){};\n\n bool operator<(const Point &aii) const;\n\n constexpr T operator[](size_t i) const { return pos[i]; }\n\n Point &operator=(const Point &obj) {\n std::copy(obj.pos, obj.pos + DIM, this->pos);\n this->dis = obj.dis;\n this->id = obj.id;\n return *this;\n }\n\n constexpr S distance(const Point &b) { return _distance(b, DIM); }\n\n void reset();\n\n S updatedistance(const Point &ref);\n\n S updateDistanceAndCount(const Point &ref, size_t &count);\n\n private:\n constexpr S _distance(const Point &b, size_t dim_left) {\n return (dim_left == 0)\n ? 0\n : powi((this->pos[dim_left - 1] - b.pos[dim_left - 1]), 2) +\n _distance(b, dim_left - 1);\n }\n};\n\ntemplate \nPoint::Point() : dis(std::numeric_limits::max()), id(0) {\n std::fill(pos, pos + DIM, 0);\n}\n\ntemplate \nPoint::Point(const T pos[DIM], size_t id)\n : dis(std::numeric_limits::max()), id(id) {\n std::copy(pos, pos + DIM, this->pos);\n}\n\ntemplate \nPoint::Point(const T pos[DIM], size_t id, S dis) : dis(dis), id(id) {\n std::copy(pos, pos + DIM, this->pos);\n}\n\ntemplate \nPoint::Point(const Point &obj) : dis(obj.dis), id(obj.id) {\n std::copy(obj.pos, obj.pos + DIM, this->pos);\n}\n\ntemplate \nbool Point::operator<(const Point &aii) const {\n return dis < aii.dis;\n}\n\ntemplate \nS Point::updatedistance(const Point &ref) {\n this->dis = std::min(this->dis, this->distance(ref));\n return this->dis;\n}\n\ntemplate \nS Point::updateDistanceAndCount(const Point &ref, size_t &count) {\n S tempDistance = this->distance(ref);\n if (tempDistance < this->dis) {\n this->dis = tempDistance;\n count++;\n }\n return this->dis;\n}\n\ntemplate void Point::reset() {\n this->dis = std::numeric_limits::max();\n}\n\n} // namespace quickfps\n\n#endif // KD_TREE_BASED_FARTHEST_POINT_SAMPLING_POINT_H\n" }, { "path": "src/_ext/utils.h", "content": "// Refactored by AyajiLin on 2023/9/16.\n\n/* functional */\n#ifndef KD_TREE_UTILS_HPP\n#define KD_TREE_UTILS_HPP\n#include \n#include \n\nnamespace quickfps {\nusing ssize_t = std::make_signed_t;\n\ntemplate \ninline constexpr T powi(const T base, const size_t exponent) {\n // (parentheses not required in next line)\n return (exponent == 0) ? 1 : (base * pow(base, exponent - 1));\n}\n} // namespace quickfps\n\n#endif // KD_TREE_UTILS_HPP\n" }, { "path": "src/fpsample/__init__.py", "content": "from __future__ import annotations\n\nimport warnings\nfrom typing import List, Optional, Union\n\nimport numpy as np\n\nfrom ._fpsample import (\n __doc__,\n __version__,\n _bucket_fps_kdline_sampling,\n _bucket_fps_kdtree_sampling,\n _fps_npdu_kdtree_sampling,\n _fps_npdu_sampling,\n _fps_sampling,\n)\n\n\ndef get_start_idx(n_pts: int, start_idx: Optional[Union[int, List[int]]]) -> Union[int, np.ndarray]:\n # Random pick a start or use the given start indices\n if start_idx is None:\n start_idx = np.random.randint(low=0, high=n_pts)\n elif isinstance(start_idx, int):\n # just use the int start_idx\n start_idx = start_idx\n elif isinstance(start_idx, list):\n # convert to numpy array for rust\n # TODO: maybe better performance with fortran array?\n start_idx = np.array(start_idx, dtype=np.uint64)\n else:\n raise ValueError(\"start_idx should be None, int or list\")\n return start_idx\n\n\ndef fps_sampling(\n pc: np.ndarray, n_samples: int, start_idx: Optional[Union[int, List[int]]] = None\n) -> np.ndarray:\n \"\"\"\n Vanilla FPS sampling.\n\n Args:\n pc (np.ndarray): The input point cloud of shape (n_pts, D).\n n_samples (int): Number of samples.\n start_idx (int or list[int], default=None): The starting index of sampling. If set to None, it will be randomly picked.\n If multiple start indices are given, they are used as a start set for the FPS and will all be present in the final samples.\n Returns:\n np.ndarray: The selected indices of shape (n_samples,).\n \"\"\"\n assert n_samples >= 1, \"n_samples should be >= 1\"\n assert pc.ndim == 2\n n_pts, _ = pc.shape\n assert n_pts >= n_samples, \"n_pts should be >= n_samples\"\n if isinstance(start_idx, int):\n assert (\n start_idx is None or 0 <= start_idx < n_pts\n ), \"start_idx should be None or 0 <= start_idx < n_pts\"\n if isinstance(start_idx, list):\n assert len(start_idx) <= n_samples, \"len(start_idx) should be <= n_samples\"\n for idx in start_idx:\n assert 0 <= idx < n_pts, \"start_idx should be None or 0 <= start_idx < n_pts\"\n # best performance with fortran array\n pc = np.asfortranarray(pc, dtype=np.float32)\n # Random pick a start if not given\n start_idx = get_start_idx(n_pts, start_idx)\n return _fps_sampling(pc, n_samples, start_idx)\n\n\ndef fps_npdu_sampling(\n pc: np.ndarray,\n n_samples: int,\n w: Optional[int] = None,\n start_idx: Optional[Union[int, List[int]]] = None,\n) -> np.ndarray:\n \"\"\"\n FPS sampling with nearest-point-distance-updating (NPDU) heuristic strategy.\n **Requires dimensional locality for best samples**.\n\n Args:\n pc (np.ndarray): The input point cloud of shape (n_pts, D).\n n_samples (int): Number of samples.\n w (int, default=None): Windows size of local heuristic search. If set to None, it will be set to `n_pts / n_samples * 16`.\n start_idx (int or list[int], default=None): The starting index of sampling. If set to None, it will be randomly picked.\n If multiple start indices are given, they are used as a start set for the FPS and will all be present in the final samples.\n Returns:\n np.ndarray: The selected indices of shape (n_samples,).\n \"\"\"\n assert n_samples >= 1, \"n_samples should be >= 1\"\n assert pc.ndim == 2\n n_pts, _ = pc.shape\n assert n_pts >= n_samples, \"n_pts should be >= n_samples\"\n assert (\n start_idx is None or 0 <= start_idx < n_pts\n ), \"start_idx should be None or 0 <= start_idx < n_pts\"\n if isinstance(start_idx, list):\n assert len(start_idx) <= n_samples, \"len(start_idx) should be <= n_samples\"\n pc = np.ascontiguousarray(pc, dtype=np.float32)\n w = w or int(n_pts / n_samples * 16)\n if w >= n_pts - 1:\n warnings.warn(f\"k is too large, set to {n_pts - 1}\")\n w = n_pts - 1\n # Random pick a start if not given\n start_idx = get_start_idx(n_pts, start_idx)\n return _fps_npdu_sampling(pc, n_samples, w, start_idx)\n\n\ndef fps_npdu_kdtree_sampling(\n pc: np.ndarray,\n n_samples: int,\n w: Optional[int] = None,\n start_idx: Optional[Union[int, List[int]]] = None,\n) -> np.ndarray:\n \"\"\"\n FPS sampling with nearest-point-distance-updating (NPDU) heuristic strategy.\n Using KDTree to eliminate the need of dimensional locality.\n Slower than `fps_npdu_sampling` but more robust.\n\n Args:\n pc (np.ndarray): The input point cloud of shape (n_pts, D).\n n_samples (int): Number of samples.\n w (int, default=None): Windows size of local heuristic search. If set to None, it will be set to `n_pts / n_samples * 16`.\n start_idx (int or list[int], default=None): The starting index of sampling. If set to None, it will be randomly picked.\n If multiple start indices are given, they are used as a start set for the FPS and will all be present in the final samples.\n Returns:\n np.ndarray: The selected indices of shape (n_samples,).\n \"\"\"\n assert n_samples >= 1, \"n_samples should be >= 1\"\n assert pc.ndim == 2\n n_pts, _ = pc.shape\n assert n_pts >= n_samples, \"n_pts should be >= n_samples\"\n assert (\n start_idx is None or 0 <= start_idx < n_pts\n ), \"start_idx should be None or 0 <= start_idx < n_pts\"\n if isinstance(start_idx, list):\n assert len(start_idx) <= n_samples, \"len(start_idx) should be <= n_samples\"\n pc = np.ascontiguousarray(pc, dtype=np.float32)\n w = w or int(n_pts / n_samples * 16)\n if w >= n_pts:\n warnings.warn(f\"k is too large, set to {n_pts}\")\n w = n_pts\n # Random pick a start if not given\n start_idx = get_start_idx(n_pts, start_idx)\n return _fps_npdu_kdtree_sampling(pc, n_samples, w, start_idx)\n\n\ndef bucket_fps_kdtree_sampling(\n pc: np.ndarray, n_samples: int, start_idx: Optional[Union[int, List[int]]] = None\n) -> np.ndarray:\n \"\"\"\n Bucket-based FPS sampling using KDTree. Also called \"QuickFPS\" in the paper.\n\n Args:\n pc (np.ndarray): The input point cloud of shape (n_pts, D).\n n_samples (int): Number of samples.\n start_idx (int or list[int], default=None): The starting index of sampling. If set to None, it will be randomly picked.\n If multiple start indices are given, they are used as a start set for the FPS and will all be present in the final samples.\n Returns:\n np.ndarray: The selected indices of shape (n_samples,).\n \"\"\"\n assert n_samples >= 1, \"n_samples should be >= 1\"\n assert pc.ndim == 2\n n_pts, _ = pc.shape\n assert n_pts >= n_samples, \"n_pts should be >= n_samples\"\n assert (\n start_idx is None or 0 <= start_idx < n_pts\n ), \"start_idx should be None or 0 <= start_idx < n_pts\"\n if isinstance(start_idx, list):\n assert len(start_idx) <= n_samples, \"len(start_idx) should be <= n_samples\"\n pc = np.ascontiguousarray(pc, dtype=np.float32)\n # Random pick a start if not given\n start_idx = get_start_idx(n_pts, start_idx)\n return _bucket_fps_kdtree_sampling(pc, n_samples, start_idx)\n\n\ndef bucket_fps_kdline_sampling(\n pc: np.ndarray, n_samples: int, h: int, start_idx: Optional[Union[int, List[int]]] = None\n) -> np.ndarray:\n \"\"\"\n Bucket-based FPS sampling using KDTree, with multiple points in each bucket. Also called \"QuickFPS\" in the paper.\n\n Args:\n pc (np.ndarray): The input point cloud of shape (n_pts, D).\n n_samples (int): Number of samples.\n h (int, default=None): Height of KDTree. The bucket size is `2**h`.\n According to the paper, for small workload, h=3 is enough;\n for medium workload, h=5 or 7 is enough; for large workload, h=9 is enough.\n start_idx (int or list[int], default=None): The starting index of sampling. If set to None, it will be randomly picked.\n If multiple start indices are given, they are used as a start set for the FPS and will all be present in the final samples.\n\n Returns:\n np.ndarray: The selected indices of shape (n_samples,).\n \"\"\"\n assert n_samples >= 1, \"n_samples should be >= 1\"\n assert pc.ndim == 2\n n_pts, _ = pc.shape\n assert n_pts >= n_samples, \"n_pts should be >= n_samples\"\n assert h >= 1, \"h should be >= 1\"\n assert 2**h <= n_pts, \"2**h should be <= n_pts\"\n assert (\n start_idx is None or 0 <= start_idx < n_pts\n ), \"start_idx should be None or 0 <= start_idx < n_pts\"\n if isinstance(start_idx, list):\n assert len(start_idx) <= n_samples, \"len(start_idx) should be <= n_samples\"\n pc = np.ascontiguousarray(pc, dtype=np.float32)\n # Random pick a start if not given\n start_idx = get_start_idx(n_pts, start_idx)\n return _bucket_fps_kdline_sampling(pc, n_samples, h, start_idx)\n\n\n__all__ = [\n \"__doc__\",\n \"__version__\",\n \"fps_sampling\",\n \"fps_npdu_sampling\",\n \"fps_npdu_kdtree_sampling\",\n \"bucket_fps_kdtree_sampling\",\n \"bucket_fps_kdline_sampling\",\n]\n" }, { "path": "src/lib.cpp", "content": "#include \n#include \n#include \"nanoflann.hpp\"\n#include \"wrapper.hpp\"\n\n#if defined(_MSC_VER)\n #include \n typedef SSIZE_T ssize_t;\n#endif\n\n#define STRINGIFY(x) #x\n#define MACRO_STRINGIFY(x) STRINGIFY(x)\n\nnamespace py = pybind11;\n\nclass StartIndex {\npublic:\n enum Type { SINGLE, ARRAY } type;\n size_t single_idx;\n py::array_t array_idx;\n\n StartIndex(size_t idx) : type(SINGLE), single_idx(idx) {}\n StartIndex(py::array_t arr) : type(ARRAY), array_idx(arr) {}\n};\n\nstruct PointCloud {\n size_t N, dim;\n const float* data; // row-major: N x dim\n\n inline size_t kdtree_get_point_count() const { return N; }\n\n inline float kdtree_distance(const float* a, const size_t b_idx, size_t /*dim_*/) const {\n const float* b = data + b_idx * dim;\n float dist = 0.0f;\n for (size_t i = 0; i < dim; i++) {\n float d = a[i] - b[i];\n dist += d * d;\n }\n return dist;\n }\n\n inline float kdtree_get_pt(const size_t idx, int dim_) const {\n return data[idx * dim + dim_];\n }\n\n template \n bool kdtree_get_bbox(BBOX&) const {\n return false;\n }\n};\n\nvoid check_py_input(\n py::array_t points,\n size_t n_samples,\n const StartIndex& start_idx,\n std::optional max_dim = std::nullopt\n) {\n if (points.ndim() != 2) {\n throw py::value_error(\n \"points must be a 2D array, but got shape \" + std::to_string(points.ndim())\n );\n }\n\n ssize_t P = points.shape(0);\n ssize_t C = points.shape(1);\n\n if (C == 0) {\n throw py::value_error(\"points must have at least one column\");\n }\n\n if (max_dim.has_value() && C > max_dim.value()) {\n throw py::value_error(\n \"points must have at most \" + std::to_string(max_dim.value()) +\n \" columns, but got \" + std::to_string(C)\n );\n }\n\n if (n_samples > static_cast(P)) {\n throw py::value_error(\n \"n_samples must be less than the number of points: n_samples=\" +\n std::to_string(n_samples) + \", P=\" + std::to_string(P)\n );\n }\n\n if (start_idx.type == StartIndex::SINGLE) {\n if (start_idx.single_idx >= static_cast(P)) {\n throw py::value_error(\n \"start_idx must be less than the number of points: start_idx=\" +\n std::to_string(start_idx.single_idx) + \", P=\" + std::to_string(P)\n );\n }\n } else {\n auto arr = start_idx.array_idx.unchecked<1>();\n if (arr.shape(0) > n_samples) {\n throw py::value_error(\n \"The number of start indices must be less than or equal to n_samples: \" +\n std::to_string(arr.shape(0)) + \", n_samples=\" + std::to_string(n_samples)\n );\n }\n for (ssize_t i = 0; i < arr.shape(0); ++i) {\n if (arr(i) >= static_cast(P)) {\n throw py::value_error(\n \"All indices in start_idx must be less than the number of points: \" +\n std::to_string(arr(i)) + \", P=\" + std::to_string(P)\n );\n }\n }\n }\n}\n\npy::array_t\nfps_sampling_multi_start_index(\n py::array_t points,\n size_t n_samples,\n py::array_t start_idx)\n{\n auto pts = points.unchecked<2>(); // 2D view\n auto starts = start_idx.unchecked<1>();\n\n ssize_t P = pts.shape(0);\n ssize_t C = pts.shape(1);\n\n if (P <= 0 || C <= 0) {\n throw std::runtime_error(\"points must be a 2D array\");\n }\n\n size_t res_selected_idx = 0;\n bool has_prev = false;\n\n std::vector dist_min(P, std::numeric_limits::infinity());\n std::vector selected;\n selected.reserve(n_samples);\n\n size_t start_counter = 0;\n\n while (selected.size() < n_samples) {\n if (has_prev) {\n size_t i0 = res_selected_idx;\n\n for (ssize_t i = 0; i < P; i++) {\n float dist = 0.0f;\n for (ssize_t j = 0; j < C; j++) {\n float d = pts(i, j) - pts(i0, j);\n dist += d * d;\n }\n if (dist < dist_min[i]) {\n dist_min[i] = dist;\n }\n }\n\n if (start_counter < (size_t) starts.shape(0)) {\n size_t idx = starts(start_counter);\n selected.push_back(idx);\n res_selected_idx = idx;\n start_counter++;\n } else {\n size_t max_idx = 0;\n float max_val = -1.0f;\n\n for (ssize_t i = 0; i < P; i++) {\n if (dist_min[i] >= max_val) {\n max_val = dist_min[i];\n max_idx = i;\n }\n }\n\n selected.push_back(max_idx);\n res_selected_idx = max_idx;\n }\n\n } else {\n size_t idx = starts(start_counter);\n selected.push_back(idx);\n res_selected_idx = idx;\n start_counter++;\n has_prev = true;\n }\n }\n\n py::array_t out(selected.size());\n auto out_buf = out.mutable_unchecked<1>();\n for (size_t i = 0; i < selected.size(); i++) {\n out_buf(i) = selected[i];\n }\n return out;\n}\n\npy::array_t fps_sampling(\n py::array_t points,\n size_t n_samples,\n size_t start_idx\n) {\n auto pts = points.unchecked<2>();\n ssize_t P = pts.shape(0);\n ssize_t C = pts.shape(1);\n\n if (P <= 0 || C <= 0) {\n throw py::value_error(\"points must be a 2D array with at least one column\");\n }\n if (start_idx >= static_cast(P)) {\n throw py::value_error(\"start_idx out of range\");\n }\n\n size_t res_selected_idx = start_idx;\n bool has_prev = false;\n std::vector dist_min(P, std::numeric_limits::infinity());\n\n std::vector selected;\n selected.reserve(n_samples);\n\n while (selected.size() < n_samples) {\n if (has_prev) {\n size_t i0 = res_selected_idx;\n for (ssize_t i = 0; i < P; ++i) {\n float dist = 0.0f;\n for (ssize_t j = 0; j < C; ++j) {\n float d = pts(i, j) - pts(i0, j);\n dist += d * d;\n }\n if (dist < dist_min[i]) dist_min[i] = dist;\n }\n\n size_t max_idx = 0;\n float max_val = -1.0f;\n for (ssize_t i = 0; i < P; ++i) {\n if (dist_min[i] >= max_val) {\n max_val = dist_min[i];\n max_idx = i;\n }\n }\n selected.push_back(max_idx);\n res_selected_idx = max_idx;\n } else {\n selected.push_back(start_idx);\n res_selected_idx = start_idx;\n has_prev = true;\n }\n }\n\n py::array_t out(selected.size());\n auto out_buf = out.mutable_unchecked<1>();\n for (size_t i = 0; i < selected.size(); ++i) {\n out_buf(i) = selected[i];\n }\n return out;\n}\n\n// EXPORT TO _fps_sample\npy::array_t _fps_sampling(\n py::array_t points,\n size_t n_samples,\n py::object start_idx_obj\n) {\n StartIndex start_idx = [&]() -> StartIndex {\n if (py::isinstance(start_idx_obj)) {\n return StartIndex(start_idx_obj.cast());\n } else if (py::isinstance>(start_idx_obj)) {\n return StartIndex(start_idx_obj.cast>());\n } else {\n throw py::type_error(\"start_idx must be int or 1D numpy array of size_t\");\n }\n }();\n\n check_py_input(points, n_samples, start_idx, std::nullopt);\n\n if (start_idx.type == StartIndex::SINGLE)\n return fps_sampling(points, n_samples, start_idx.single_idx);\n else\n return fps_sampling_multi_start_index(points, n_samples, start_idx.array_idx);\n}\n\npy::array_t fps_npdu_sampling(\n py::array_t points,\n size_t n_samples,\n size_t k,\n size_t start_idx\n) {\n auto pts = points.unchecked<2>();\n ssize_t P = pts.shape(0);\n ssize_t C = pts.shape(1);\n\n if (P <= 0 || C <= 0)\n throw py::value_error(\"points must be a 2D array with at least one column\");\n if (start_idx >= static_cast(P))\n throw py::value_error(\"start_idx out of range\");\n\n std::vector dist_min(P, std::numeric_limits::infinity());\n std::vector selected;\n selected.reserve(n_samples);\n\n size_t res_selected_idx = start_idx;\n bool has_prev = false;\n\n while (selected.size() < n_samples) {\n if (has_prev) {\n ssize_t hw = static_cast(k / 2);\n ssize_t start = static_cast(res_selected_idx) - hw;\n ssize_t end = static_cast(res_selected_idx) + hw;\n if (start < 0) { end -= start; start = 0; }\n if (end >= P) { start = std::max(start - (end - P + 1), ssize_t(0)); end = P - 1; }\n\n for (ssize_t i = start; i <= end; ++i) {\n float dist = 0.0f;\n for (ssize_t j = 0; j < C; ++j) {\n float d = pts(i, j) - pts(res_selected_idx, j);\n dist += d*d;\n }\n if (dist < dist_min[i]) dist_min[i] = dist;\n }\n\n size_t max_idx = 0;\n float max_val = -1.0f;\n for (ssize_t i = 0; i < P; ++i) {\n if (dist_min[i] > max_val) { max_val = dist_min[i]; max_idx = i; }\n }\n selected.push_back(max_idx);\n res_selected_idx = max_idx;\n\n } else {\n for (ssize_t i = 0; i < P; ++i) {\n float dist = 0.0f;\n for (ssize_t j = 0; j < C; ++j) {\n float d = pts(i,j) - pts(start_idx,j);\n dist += d*d;\n }\n if (dist < dist_min[i]) dist_min[i] = dist;\n }\n selected.push_back(start_idx);\n res_selected_idx = start_idx;\n has_prev = true;\n }\n }\n\n py::array_t out(selected.size());\n auto out_buf = out.mutable_unchecked<1>();\n for (size_t i = 0; i < selected.size(); ++i)\n out_buf(i) = selected[i];\n\n return out;\n}\n\n// EXPORT TO _fps_npdu_sample\npy::array_t fps_npdu_sampling_py(\n py::array_t points,\n size_t n_samples,\n size_t k,\n py::object start_idx_obj\n) {\n StartIndex start_idx = [&]() -> StartIndex {\n if (py::isinstance(start_idx_obj))\n return StartIndex(start_idx_obj.cast());\n else if (py::isinstance>(start_idx_obj))\n return StartIndex(start_idx_obj.cast>());\n else\n throw py::type_error(\"start_idx must be int or 1D numpy array of size_t\");\n }();\n\n check_py_input(points, n_samples, start_idx);\n\n if (start_idx.type == StartIndex::SINGLE)\n return fps_npdu_sampling(points, n_samples, k, start_idx.single_idx);\n else {\n PyErr_SetString(PyExc_NotImplementedError, \"Array of start indices not implemented yet\");\n throw py::error_already_set();\n }\n}\n\n// EXPORT TO _fps_npdu_kdtree_sample\npy::array_t fps_npdu_kdtree_sampling_py(\n py::array_t points,\n size_t n_samples,\n size_t k,\n py::object start_idx_obj\n) {\n StartIndex start_idx = [&]() -> StartIndex {\n if (py::isinstance(start_idx_obj))\n return StartIndex(start_idx_obj.cast());\n else if (py::isinstance>(start_idx_obj))\n return StartIndex(start_idx_obj.cast>());\n else\n throw py::type_error(\"start_idx must be int or 1D numpy array of size_t\");\n }();\n\n check_py_input(points, n_samples, start_idx);\n\n if (start_idx.type == StartIndex::ARRAY) {\n PyErr_SetString(PyExc_NotImplementedError, \"Array of start indices not implemented yet\");\n throw py::error_already_set();\n }\n\n auto pts = points.unchecked<2>();\n ssize_t P = pts.shape(0);\n ssize_t C = pts.shape(1);\n\n if (P <= 0 || C <= 0)\n throw py::value_error(\"points must be a 2D array with at least one column\");\n\n PointCloud cloud;\n cloud.N = static_cast(P);\n cloud.dim = static_cast(C);\n cloud.data = points.data();\n\n using KDTree = nanoflann::KDTreeSingleIndexAdaptor, PointCloud, -1>;\n KDTree index(static_cast(C), cloud, nanoflann::KDTreeSingleIndexAdaptorParams(10));\n index.buildIndex();\n\n std::vector dist_min(P, std::numeric_limits::infinity());\n std::vector selected;\n selected.reserve(n_samples);\n\n size_t res_selected_idx = start_idx.single_idx;\n bool has_prev = false;\n\n size_t k_use = std::min(k, static_cast(P));\n std::vector ret_indexes(k_use);\n std::vector out_dists(k_use);\n\n while (selected.size() < n_samples) {\n if (has_prev) {\n std::vector query(static_cast(C));\n for (ssize_t d = 0; d < C; ++d) query[d] = pts(res_selected_idx, d);\n\n nanoflann::KNNResultSet resultSet(static_cast(k_use));\n resultSet.init(ret_indexes.data(), out_dists.data());\n nanoflann::SearchParameters params;\n index.findNeighbors(resultSet, query.data(), params);\n\n for (size_t idx_i = 0; idx_i < k_use; ++idx_i) {\n size_t nb = ret_indexes[idx_i];\n float dist = 0.0f;\n for (ssize_t d = 0; d < C; ++d) {\n float diff = pts(nb, d) - pts(res_selected_idx, d);\n dist += diff * diff;\n }\n if (dist < dist_min[nb]) dist_min[nb] = dist;\n }\n\n size_t max_idx = 0;\n float max_val = -1.0f;\n for (ssize_t i = 0; i < P; ++i) {\n if (dist_min[i] > max_val) { max_val = dist_min[i]; max_idx = i; }\n }\n\n selected.push_back(max_idx);\n res_selected_idx = max_idx;\n } else {\n for (ssize_t i = 0; i < P; ++i) {\n float dist = 0.0f;\n for (ssize_t j = 0; j < C; ++j) {\n float d = pts(i,j) - pts(res_selected_idx,j);\n dist += d*d;\n }\n if (dist < dist_min[i]) dist_min[i] = dist;\n }\n selected.push_back(res_selected_idx);\n has_prev = true;\n }\n }\n\n py::array_t out(selected.size());\n auto out_buf = out.mutable_unchecked<1>();\n for (size_t i = 0; i < selected.size(); ++i) out_buf(i) = selected[i];\n\n return out;\n}\n\npy::array_t bucket_fps_kdtree_sampling_py(\n py::array_t points,\n size_t n_samples,\n py::object start_idx_obj\n) {\n StartIndex start_idx = [&]() -> StartIndex {\n if (py::isinstance(start_idx_obj))\n return StartIndex(start_idx_obj.cast());\n else if (py::isinstance>(start_idx_obj))\n return StartIndex(start_idx_obj.cast>());\n else\n throw py::type_error(\"start_idx must be int or 1D numpy array of size_t\");\n }();\n\n if (start_idx.type == StartIndex::ARRAY) {\n PyErr_SetString(PyExc_NotImplementedError, \"Array of start indices not implemented yet\");\n throw py::error_already_set();\n }\n\n check_py_input(points, n_samples, start_idx);\n\n if (points.ndim() != 2) {\n throw py::value_error(\"points must be a 2D float32 array\");\n }\n ssize_t P = points.shape(0);\n ssize_t C = points.shape(1);\n\n if (start_idx.single_idx >= static_cast(P)) {\n throw py::value_error(\"start_idx out of range\");\n }\n if (n_samples == 0 || n_samples > static_cast(P)) {\n throw py::value_error(\"n_samples must be in [1, num_points]\");\n }\n\n auto buf = points.unchecked<2>();\n\n py::array_t out(n_samples);\n size_t* out_ptr = static_cast(out.mutable_data());\n\n int ret = bucket_fps_kdtree(\n buf.data(0,0), // raw_data\n static_cast(P), // n_points\n static_cast(C), // dim\n n_samples, // n_samples\n start_idx.single_idx, // start_idx\n out_ptr // output buffer\n );\n\n if (ret != 0) {\n throw std::runtime_error(\"bucket_fps_kdtree failed with error code \" + std::to_string(ret));\n }\n\n return out;\n}\n\npy::array_t bucket_fps_kdline_sampling_py(\n py::array_t points,\n size_t n_samples,\n size_t height,\n py::object start_idx_obj\n) {\n StartIndex start_idx = [&]() -> StartIndex {\n if (py::isinstance(start_idx_obj))\n return StartIndex(start_idx_obj.cast());\n else if (py::isinstance>(start_idx_obj))\n return StartIndex(start_idx_obj.cast>());\n else\n throw py::type_error(\"start_idx must be int or 1D numpy array of size_t\");\n }();\n\n if (start_idx.type == StartIndex::ARRAY) {\n PyErr_SetString(PyExc_NotImplementedError, \"Array of start indices not implemented yet\");\n throw py::error_already_set();\n }\n\n if (points.ndim() != 2) {\n throw py::value_error(\"points must be a 2D float32 array\");\n }\n\n ssize_t P = points.shape(0);\n ssize_t C = points.shape(1);\n\n if (start_idx.single_idx >= static_cast(P)) {\n throw py::value_error(\"start_idx out of range\");\n }\n if (n_samples == 0 || n_samples > static_cast(P)) {\n throw py::value_error(\"n_samples must be in [1, num_points]\");\n }\n if (height == 0) {\n throw py::value_error(\"height must be >= 1\");\n }\n\n auto buf = points.unchecked<2>();\n\n py::array_t out(n_samples);\n size_t* out_ptr = static_cast(out.mutable_data());\n\n int ret = bucket_fps_kdline(\n buf.data(0,0), // raw_data\n static_cast(P), // n_points\n static_cast(C), // dim\n n_samples, // n_samples\n start_idx.single_idx, // start_idx\n height, // window height\n out_ptr // output buffer\n );\n\n if (ret != 0) {\n throw std::runtime_error(\"bucket_fps_kdline failed with error code \" + std::to_string(ret));\n }\n\n return out;\n}\n\nPYBIND11_MODULE(_fpsample, m, py::mod_gil_not_used(), py::multiple_interpreters::per_interpreter_gil()) {\n m.doc() = R\"pbdoc(\n Python efficient farthest point sampling (FPS) library\n -----------------------\n\n .. currentmodule:: fpsample\n\n .. autosummary::\n :toctree: _generate\n\n _fps_sampling\n _fps_npdu_sampling\n _fps_npdu_kdtree_sampling\n _bucket_fps_kdtree_sampling\n _bucket_fps_kdline_sampling\n )pbdoc\";\n\n m.def(\"_fps_sampling\", &_fps_sampling, R\"pbdoc(\n Farthest Point Sampling (FPS)\n Args:\n points (np.ndarray[float32, 2D]): N x C point array.\n n_samples (int): number of samples to pick.\n start_idx (int or np.ndarray[int32, 1D]): initial index or indices to start FPS.\n Returns:\n np.ndarray[int32]: sampled point indices.\n )pbdoc\");\n\n m.def(\"_fps_npdu_sampling\", &fps_npdu_sampling_py, R\"pbdoc(\n FPS with Nearest Point Distance Update\n Args:\n points (np.ndarray[float32, 2D]): N x C point array.\n n_samples (int): number of samples to pick.\n k (int): number of neighbors for local update.\n start_idx (int): initial index to start FPS.\n Returns:\n np.ndarray[int32]: sampled point indices.\n )pbdoc\");\n\n m.def(\"_fps_npdu_kdtree_sampling\", &fps_npdu_kdtree_sampling_py, R\"pbdoc(\n FPS with Nearest Point Distance Update using KD-tree acceleration\n Args:\n points (np.ndarray[float32, 2D]): N x C point array.\n n_samples (int): number of samples to pick.\n k (int): number of neighbors for local update.\n start_idx (int): initial index to start FPS.\n Returns:\n np.ndarray[int32]: sampled point indices.\n )pbdoc\");\n\n m.def(\"_bucket_fps_kdtree_sampling\",\n &bucket_fps_kdtree_sampling_py,\n R\"pbdoc(\n Bucket FPS sampling using KD-tree acceleration.\n Args:\n points (np.ndarray[float32, 2D]): N x C point array.\n n_samples (int): number of samples to pick.\n start_idx (int): initial index to start FPS.\n Returns:\n np.ndarray[int32]: sampled point indices.\n )pbdoc\");\n\nm.def(\"_bucket_fps_kdline_sampling\",\n &bucket_fps_kdline_sampling_py,\n R\"pbdoc(\n Bucket FPS sampling using KD-line local window update.\n Args:\n points (np.ndarray[float32, 2D]): N x C point array.\n n_samples (int): number of samples to pick.\n height (int): window size around selected point.\n start_idx (int): first index.\n Returns:\n np.ndarray[int32]: sampled point indices.\n )pbdoc\");\n\n#ifdef VERSION_INFO\n m.attr(\"__version__\") = MACRO_STRINGIFY(VERSION_INFO);\n#else\n m.attr(\"__version__\") = \"dev\";\n#endif\n}\n" }, { "path": "src/nanoflann.hpp", "content": "/***********************************************************************\n * Software License Agreement (BSD License)\n *\n * Copyright 2008-2009 Marius Muja (mariusm@cs.ubc.ca). All rights reserved.\n * Copyright 2008-2009 David G. Lowe (lowe@cs.ubc.ca). All rights reserved.\n * Copyright 2011-2025 Jose Luis Blanco (joseluisblancoc@gmail.com).\n * All rights reserved.\n *\n * THE BSD LICENSE\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions\n * are met:\n *\n * 1. Redistributions of source code must retain the above copyright\n * notice, this list of conditions and the following disclaimer.\n * 2. Redistributions in binary form must reproduce the above copyright\n * notice, this list of conditions and the following disclaimer in the\n * documentation and/or other materials provided with the distribution.\n *\n * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR\n * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES\n * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.\n * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,\n * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT\n * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,\n * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY\n * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF\n * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *************************************************************************/\n\n/** \\mainpage nanoflann C++ API documentation\n * nanoflann is a C++ header-only library for building KD-Trees, mostly\n * optimized for 2D or 3D point clouds.\n *\n * nanoflann does not require compiling or installing, just an\n * #include in your code.\n *\n * See:\n * - [Online README](https://github.com/jlblancoc/nanoflann)\n * - [C++ API documentation](https://jlblancoc.github.io/nanoflann/)\n */\n\n#pragma once\n\n#include \n#include \n#include \n#include \n#include // for abs()\n#include \n#include // for abs()\n#include // std::reference_wrapper\n#include \n#include \n#include // std::numeric_limits\n#include \n#include \n#include \n#include \n#include \n\n/** Library version: 0xMmP (M=Major,m=minor,P=patch) */\n#define NANOFLANN_VERSION 0x180\n\n// Avoid conflicting declaration of min/max macros in Windows headers\n#if !defined(NOMINMAX) && \\\n (defined(_WIN32) || defined(_WIN32_) || defined(WIN32) || defined(_WIN64))\n#define NOMINMAX\n#ifdef max\n#undef max\n#undef min\n#endif\n#endif\n// Avoid conflicts with X11 headers\n#ifdef None\n#undef None\n#endif\n\nnamespace nanoflann\n{\n/** @addtogroup nanoflann_grp nanoflann C++ library for KD-trees\n * @{ */\n\n/** the PI constant (required to avoid MSVC missing symbols) */\ntemplate \nT pi_const()\n{\n return static_cast(3.14159265358979323846);\n}\n\n/**\n * Traits if object is resizable and assignable (typically has a resize | assign\n * method)\n */\ntemplate \nstruct has_resize : std::false_type\n{\n};\n\ntemplate \nstruct has_resize().resize(1), 0)>\n : std::true_type\n{\n};\n\ntemplate \nstruct has_assign : std::false_type\n{\n};\n\ntemplate \nstruct has_assign().assign(1, 0), 0)>\n : std::true_type\n{\n};\n\n/**\n * Free function to resize a resizable object\n */\ntemplate \ninline typename std::enable_if::value, void>::type resize(\n Container& c, const size_t nElements)\n{\n c.resize(nElements);\n}\n\n/**\n * Free function that has no effects on non resizable containers (e.g.\n * std::array) It raises an exception if the expected size does not match\n */\ntemplate \ninline typename std::enable_if::value, void>::type\n resize(Container& c, const size_t nElements)\n{\n if (nElements != c.size())\n throw std::logic_error(\"Try to change the size of a std::array.\");\n}\n\n/**\n * Free function to assign to a container\n */\ntemplate \ninline typename std::enable_if::value, void>::type assign(\n Container& c, const size_t nElements, const T& value)\n{\n c.assign(nElements, value);\n}\n\n/**\n * Free function to assign to a std::array\n */\ntemplate \ninline typename std::enable_if::value, void>::type\n assign(Container& c, const size_t nElements, const T& value)\n{\n for (size_t i = 0; i < nElements; i++) c[i] = value;\n}\n\n/** operator \"<\" for std::sort() */\nstruct IndexDist_Sorter\n{\n /** PairType will be typically: ResultItem */\n template \n bool operator()(const PairType& p1, const PairType& p2) const\n {\n return p1.second < p2.second;\n }\n};\n\n/**\n * Each result element in RadiusResultSet. Note that distances and indices\n * are named `first` and `second` to keep backward-compatibility with the\n * `std::pair<>` type used in the past. In contrast, this structure is ensured\n * to be `std::is_standard_layout` so it can be used in wrappers to other\n * languages.\n * See: https://github.com/jlblancoc/nanoflann/issues/166\n */\ntemplate \nstruct ResultItem\n{\n ResultItem() = default;\n ResultItem(const IndexType index, const DistanceType distance)\n : first(index), second(distance)\n {\n }\n\n IndexType first; //!< Index of the sample in the dataset\n DistanceType second; //!< Distance from sample to query point\n};\n\n/** @addtogroup result_sets_grp Result set classes\n * @{ */\n\n/** Result set for KNN searches (N-closest neighbors) */\ntemplate <\n typename _DistanceType, typename _IndexType = size_t,\n typename _CountType = size_t>\nclass KNNResultSet\n{\n public:\n using DistanceType = _DistanceType;\n using IndexType = _IndexType;\n using CountType = _CountType;\n\n private:\n IndexType* indices;\n DistanceType* dists;\n CountType capacity;\n CountType count;\n\n public:\n explicit KNNResultSet(CountType capacity_)\n : indices(nullptr), dists(nullptr), capacity(capacity_), count(0)\n {\n }\n\n void init(IndexType* indices_, DistanceType* dists_)\n {\n indices = indices_;\n dists = dists_;\n count = 0;\n }\n\n CountType size() const { return count; }\n bool empty() const { return count == 0; }\n bool full() const { return count == capacity; }\n\n /**\n * Called during search to add an element matching the criteria.\n * @return true if the search should be continued, false if the results are\n * sufficient\n */\n bool addPoint(DistanceType dist, IndexType index)\n {\n CountType i;\n for (i = count; i > 0; --i)\n {\n /** If defined and two points have the same distance, the one with\n * the lowest-index will be returned first. */\n#ifdef NANOFLANN_FIRST_MATCH\n if ((dists[i - 1] > dist) ||\n ((dist == dists[i - 1]) && (indices[i - 1] > index)))\n {\n#else\n if (dists[i - 1] > dist)\n {\n#endif\n if (i < capacity)\n {\n dists[i] = dists[i - 1];\n indices[i] = indices[i - 1];\n }\n }\n else\n break;\n }\n if (i < capacity)\n {\n dists[i] = dist;\n indices[i] = index;\n }\n if (count < capacity) count++;\n\n // tell caller that the search shall continue\n return true;\n }\n\n //! Returns the worst distance among found solutions if the search result is\n //! full, or the maximum possible distance, if not full yet.\n DistanceType worstDist() const\n {\n return (count < capacity || !count)\n ? std::numeric_limits::max()\n : dists[count - 1];\n }\n\n void sort()\n {\n // already sorted\n }\n};\n\n/** Result set for RKNN searches (N-closest neighbors with a maximum radius) */\ntemplate <\n typename _DistanceType, typename _IndexType = size_t,\n typename _CountType = size_t>\nclass RKNNResultSet\n{\n public:\n using DistanceType = _DistanceType;\n using IndexType = _IndexType;\n using CountType = _CountType;\n\n private:\n IndexType* indices;\n DistanceType* dists;\n CountType capacity;\n CountType count;\n DistanceType maximumSearchDistanceSquared;\n\n public:\n explicit RKNNResultSet(\n CountType capacity_, DistanceType maximumSearchDistanceSquared_)\n : indices(nullptr),\n dists(nullptr),\n capacity(capacity_),\n count(0),\n maximumSearchDistanceSquared(maximumSearchDistanceSquared_)\n {\n }\n\n void init(IndexType* indices_, DistanceType* dists_)\n {\n indices = indices_;\n dists = dists_;\n count = 0;\n if (capacity) dists[capacity - 1] = maximumSearchDistanceSquared;\n }\n\n CountType size() const { return count; }\n bool empty() const { return count == 0; }\n bool full() const { return count == capacity; }\n\n /**\n * Called during search to add an element matching the criteria.\n * @return true if the search should be continued, false if the results are\n * sufficient\n */\n bool addPoint(DistanceType dist, IndexType index)\n {\n CountType i;\n for (i = count; i > 0; --i)\n {\n /** If defined and two points have the same distance, the one with\n * the lowest-index will be returned first. */\n#ifdef NANOFLANN_FIRST_MATCH\n if ((dists[i - 1] > dist) ||\n ((dist == dists[i - 1]) && (indices[i - 1] > index)))\n {\n#else\n if (dists[i - 1] > dist)\n {\n#endif\n if (i < capacity)\n {\n dists[i] = dists[i - 1];\n indices[i] = indices[i - 1];\n }\n }\n else\n break;\n }\n if (i < capacity)\n {\n dists[i] = dist;\n indices[i] = index;\n }\n if (count < capacity) count++;\n\n // tell caller that the search shall continue\n return true;\n }\n\n //! Returns the worst distance among found solutions if the search result is\n //! full, or the maximum possible distance, if not full yet.\n DistanceType worstDist() const\n {\n return (count < capacity || !count) ? maximumSearchDistanceSquared\n : dists[count - 1];\n }\n\n void sort()\n {\n // already sorted\n }\n};\n\n/**\n * A result-set class used when performing a radius based search.\n */\ntemplate \nclass RadiusResultSet\n{\n public:\n using DistanceType = _DistanceType;\n using IndexType = _IndexType;\n\n public:\n const DistanceType radius;\n\n std::vector>& m_indices_dists;\n\n explicit RadiusResultSet(\n DistanceType radius_,\n std::vector>& indices_dists)\n : radius(radius_), m_indices_dists(indices_dists)\n {\n init();\n }\n\n void init() { clear(); }\n void clear() { m_indices_dists.clear(); }\n\n size_t size() const { return m_indices_dists.size(); }\n size_t empty() const { return m_indices_dists.empty(); }\n\n bool full() const { return true; }\n\n /**\n * Called during search to add an element matching the criteria.\n * @return true if the search should be continued, false if the results are\n * sufficient\n */\n bool addPoint(DistanceType dist, IndexType index)\n {\n if (dist < radius) m_indices_dists.emplace_back(index, dist);\n return true;\n }\n\n DistanceType worstDist() const { return radius; }\n\n /**\n * Find the worst result (farthest neighbor) without copying or sorting\n * Pre-conditions: size() > 0\n */\n ResultItem worst_item() const\n {\n if (m_indices_dists.empty())\n throw std::runtime_error(\n \"Cannot invoke RadiusResultSet::worst_item() on \"\n \"an empty list of results.\");\n auto it = std::max_element(\n m_indices_dists.begin(), m_indices_dists.end(), IndexDist_Sorter());\n return *it;\n }\n\n void sort()\n {\n std::sort(\n m_indices_dists.begin(), m_indices_dists.end(), IndexDist_Sorter());\n }\n};\n\n/** @} */\n\n/** @addtogroup loadsave_grp Load/save auxiliary functions\n * @{ */\ntemplate \nvoid save_value(std::ostream& stream, const T& value)\n{\n stream.write(reinterpret_cast(&value), sizeof(T));\n}\n\ntemplate \nvoid save_value(std::ostream& stream, const std::vector& value)\n{\n size_t size = value.size();\n stream.write(reinterpret_cast(&size), sizeof(size_t));\n stream.write(reinterpret_cast(value.data()), sizeof(T) * size);\n}\n\ntemplate \nvoid load_value(std::istream& stream, T& value)\n{\n stream.read(reinterpret_cast(&value), sizeof(T));\n}\n\ntemplate \nvoid load_value(std::istream& stream, std::vector& value)\n{\n size_t size;\n stream.read(reinterpret_cast(&size), sizeof(size_t));\n value.resize(size);\n stream.read(reinterpret_cast(value.data()), sizeof(T) * size);\n}\n/** @} */\n\n/** @addtogroup metric_grp Metric (distance) classes\n * @{ */\n\nstruct Metric\n{\n};\n\n/** Manhattan distance functor (generic version, optimized for\n * high-dimensionality data sets). Corresponding distance traits:\n * nanoflann::metric_L1\n *\n * \\tparam T Type of the elements (e.g. double, float, uint8_t)\n * \\tparam DataSource Source of the data, i.e. where the vectors are stored\n * \\tparam _DistanceType Type of distance variables (must be signed)\n * \\tparam IndexType Type of the arguments with which the data can be\n * accessed (e.g. float, double, int64_t, T*)\n */\ntemplate <\n class T, class DataSource, typename _DistanceType = T,\n typename IndexType = size_t>\nstruct L1_Adaptor\n{\n using ElementType = T;\n using DistanceType = _DistanceType;\n\n const DataSource& data_source;\n\n L1_Adaptor(const DataSource& _data_source) : data_source(_data_source) {}\n\n DistanceType evalMetric(\n const T* a, const IndexType b_idx, size_t size,\n DistanceType worst_dist = -1) const\n {\n DistanceType result = DistanceType();\n const T* last = a + size;\n const T* lastgroup = last - 3;\n size_t d = 0;\n\n /* Process 4 items with each loop for efficiency. */\n while (a < lastgroup)\n {\n const DistanceType diff0 =\n std::abs(a[0] - data_source.kdtree_get_pt(b_idx, d++));\n const DistanceType diff1 =\n std::abs(a[1] - data_source.kdtree_get_pt(b_idx, d++));\n const DistanceType diff2 =\n std::abs(a[2] - data_source.kdtree_get_pt(b_idx, d++));\n const DistanceType diff3 =\n std::abs(a[3] - data_source.kdtree_get_pt(b_idx, d++));\n result += diff0 + diff1 + diff2 + diff3;\n a += 4;\n if ((worst_dist > 0) && (result > worst_dist)) { return result; }\n }\n /* Process last 0-3 components. Not needed for standard vector lengths.\n */\n while (a < last)\n {\n result += std::abs(*a++ - data_source.kdtree_get_pt(b_idx, d++));\n }\n return result;\n }\n\n template \n DistanceType accum_dist(const U a, const V b, const size_t) const\n {\n return std::abs(a - b);\n }\n};\n\n/** **Squared** Euclidean distance functor (generic version, optimized for\n * high-dimensionality data sets). Corresponding distance traits:\n * nanoflann::metric_L2\n *\n * \\tparam T Type of the elements (e.g. double, float, uint8_t)\n * \\tparam DataSource Source of the data, i.e. where the vectors are stored\n * \\tparam _DistanceType Type of distance variables (must be signed)\n * \\tparam IndexType Type of the arguments with which the data can be\n * accessed (e.g. float, double, int64_t, T*)\n */\ntemplate <\n class T, class DataSource, typename _DistanceType = T,\n typename IndexType = size_t>\nstruct L2_Adaptor\n{\n using ElementType = T;\n using DistanceType = _DistanceType;\n\n const DataSource& data_source;\n\n L2_Adaptor(const DataSource& _data_source) : data_source(_data_source) {}\n\n DistanceType evalMetric(\n const T* a, const IndexType b_idx, size_t size,\n DistanceType worst_dist = -1) const\n {\n DistanceType result = DistanceType();\n const T* last = a + size;\n const T* lastgroup = last - 3;\n size_t d = 0;\n\n /* Process 4 items with each loop for efficiency. */\n while (a < lastgroup)\n {\n const DistanceType diff0 =\n a[0] - data_source.kdtree_get_pt(b_idx, d++);\n const DistanceType diff1 =\n a[1] - data_source.kdtree_get_pt(b_idx, d++);\n const DistanceType diff2 =\n a[2] - data_source.kdtree_get_pt(b_idx, d++);\n const DistanceType diff3 =\n a[3] - data_source.kdtree_get_pt(b_idx, d++);\n result +=\n diff0 * diff0 + diff1 * diff1 + diff2 * diff2 + diff3 * diff3;\n a += 4;\n if ((worst_dist > 0) && (result > worst_dist)) { return result; }\n }\n /* Process last 0-3 components. Not needed for standard vector lengths.\n */\n while (a < last)\n {\n const DistanceType diff0 =\n *a++ - data_source.kdtree_get_pt(b_idx, d++);\n result += diff0 * diff0;\n }\n return result;\n }\n\n template \n DistanceType accum_dist(const U a, const V b, const size_t) const\n {\n return (a - b) * (a - b);\n }\n};\n\n/** **Squared** Euclidean (L2) distance functor (suitable for low-dimensionality\n * datasets, like 2D or 3D point clouds) Corresponding distance traits:\n * nanoflann::metric_L2_Simple\n *\n * \\tparam T Type of the elements (e.g. double, float, uint8_t)\n * \\tparam DataSource Source of the data, i.e. where the vectors are stored\n * \\tparam _DistanceType Type of distance variables (must be signed)\n * \\tparam IndexType Type of the arguments with which the data can be\n * accessed (e.g. float, double, int64_t, T*)\n */\ntemplate <\n class T, class DataSource, typename _DistanceType = T,\n typename IndexType = size_t>\nstruct L2_Simple_Adaptor\n{\n using ElementType = T;\n using DistanceType = _DistanceType;\n\n const DataSource& data_source;\n\n L2_Simple_Adaptor(const DataSource& _data_source)\n : data_source(_data_source)\n {\n }\n\n DistanceType evalMetric(\n const T* a, const IndexType b_idx, size_t size) const\n {\n DistanceType result = DistanceType();\n for (size_t i = 0; i < size; ++i)\n {\n const DistanceType diff =\n a[i] - data_source.kdtree_get_pt(b_idx, i);\n result += diff * diff;\n }\n return result;\n }\n\n template \n DistanceType accum_dist(const U a, const V b, const size_t) const\n {\n return (a - b) * (a - b);\n }\n};\n\n/** SO2 distance functor\n * Corresponding distance traits: nanoflann::metric_SO2\n *\n * \\tparam T Type of the elements (e.g. double, float, uint8_t)\n * \\tparam DataSource Source of the data, i.e. where the vectors are stored\n * \\tparam _DistanceType Type of distance variables (must be signed) (e.g.\n * float, double) orientation is constrained to be in [-pi, pi]\n * \\tparam IndexType Type of the arguments with which the data can be\n * accessed (e.g. float, double, int64_t, T*)\n */\ntemplate <\n class T, class DataSource, typename _DistanceType = T,\n typename IndexType = size_t>\nstruct SO2_Adaptor\n{\n using ElementType = T;\n using DistanceType = _DistanceType;\n\n const DataSource& data_source;\n\n SO2_Adaptor(const DataSource& _data_source) : data_source(_data_source) {}\n\n DistanceType evalMetric(\n const T* a, const IndexType b_idx, size_t size) const\n {\n return accum_dist(\n a[size - 1], data_source.kdtree_get_pt(b_idx, size - 1), size - 1);\n }\n\n /** Note: this assumes that input angles are already in the range [-pi,pi]\n */\n template \n DistanceType accum_dist(const U a, const V b, const size_t) const\n {\n DistanceType result = DistanceType();\n DistanceType PI = pi_const();\n result = b - a;\n if (result > PI)\n result -= 2 * PI;\n else if (result < -PI)\n result += 2 * PI;\n return result;\n }\n};\n\n/** SO3 distance functor (Uses L2_Simple)\n * Corresponding distance traits: nanoflann::metric_SO3\n *\n * \\tparam T Type of the elements (e.g. double, float, uint8_t)\n * \\tparam DataSource Source of the data, i.e. where the vectors are stored\n * \\tparam _DistanceType Type of distance variables (must be signed) (e.g.\n * float, double)\n * \\tparam IndexType Type of the arguments with which the data can be\n * accessed (e.g. float, double, int64_t, T*)\n */\ntemplate <\n class T, class DataSource, typename _DistanceType = T,\n typename IndexType = size_t>\nstruct SO3_Adaptor\n{\n using ElementType = T;\n using DistanceType = _DistanceType;\n\n L2_Simple_Adaptor\n distance_L2_Simple;\n\n SO3_Adaptor(const DataSource& _data_source)\n : distance_L2_Simple(_data_source)\n {\n }\n\n DistanceType evalMetric(\n const T* a, const IndexType b_idx, size_t size) const\n {\n return distance_L2_Simple.evalMetric(a, b_idx, size);\n }\n\n template \n DistanceType accum_dist(const U a, const V b, const size_t idx) const\n {\n return distance_L2_Simple.accum_dist(a, b, idx);\n }\n};\n\n/** Metaprogramming helper traits class for the L1 (Manhattan) metric */\nstruct metric_L1 : public Metric\n{\n template \n struct traits\n {\n using distance_t = L1_Adaptor;\n };\n};\n/** Metaprogramming helper traits class for the L2 (Euclidean) **squared**\n * distance metric */\nstruct metric_L2 : public Metric\n{\n template \n struct traits\n {\n using distance_t = L2_Adaptor;\n };\n};\n/** Metaprogramming helper traits class for the L2_simple (Euclidean)\n * **squared** distance metric */\nstruct metric_L2_Simple : public Metric\n{\n template \n struct traits\n {\n using distance_t = L2_Simple_Adaptor;\n };\n};\n/** Metaprogramming helper traits class for the SO3_InnerProdQuat metric */\nstruct metric_SO2 : public Metric\n{\n template \n struct traits\n {\n using distance_t = SO2_Adaptor;\n };\n};\n/** Metaprogramming helper traits class for the SO3_InnerProdQuat metric */\nstruct metric_SO3 : public Metric\n{\n template \n struct traits\n {\n using distance_t = SO3_Adaptor;\n };\n};\n\n/** @} */\n\n/** @addtogroup param_grp Parameter structs\n * @{ */\n\nenum class KDTreeSingleIndexAdaptorFlags\n{\n None = 0,\n SkipInitialBuildIndex = 1\n};\n\ninline std::underlying_type::type operator&(\n KDTreeSingleIndexAdaptorFlags lhs, KDTreeSingleIndexAdaptorFlags rhs)\n{\n using underlying =\n typename std::underlying_type::type;\n return static_cast(lhs) & static_cast(rhs);\n}\n\n/** Parameters (see README.md) */\nstruct KDTreeSingleIndexAdaptorParams\n{\n KDTreeSingleIndexAdaptorParams(\n size_t _leaf_max_size = 10,\n KDTreeSingleIndexAdaptorFlags _flags =\n KDTreeSingleIndexAdaptorFlags::None,\n unsigned int _n_thread_build = 1)\n : leaf_max_size(_leaf_max_size),\n flags(_flags),\n n_thread_build(_n_thread_build)\n {\n }\n\n size_t leaf_max_size;\n KDTreeSingleIndexAdaptorFlags flags;\n unsigned int n_thread_build;\n};\n\n/** Search options for KDTreeSingleIndexAdaptor::findNeighbors() */\nstruct SearchParameters\n{\n SearchParameters(float eps_ = 0, bool sorted_ = true)\n : eps(eps_), sorted(sorted_)\n {\n }\n\n float eps; //!< search for eps-approximate neighbours (default: 0)\n bool sorted; //!< only for radius search, require neighbours sorted by\n //!< distance (default: true)\n};\n/** @} */\n\n/** @addtogroup memalloc_grp Memory allocation\n * @{ */\n\n/**\n * Pooled storage allocator\n *\n * The following routines allow for the efficient allocation of storage in\n * small chunks from a specified pool. Rather than allowing each structure\n * to be freed individually, an entire pool of storage is freed at once.\n * This method has two advantages over just using malloc() and free(). First,\n * it is far more efficient for allocating small objects, as there is\n * no overhead for remembering all the information needed to free each\n * object or consolidating fragmented memory. Second, the decision about\n * how long to keep an object is made at the time of allocation, and there\n * is no need to track down all the objects to free them.\n *\n */\nclass PooledAllocator\n{\n static constexpr size_t WORDSIZE = 16; // WORDSIZE must >= 8\n static constexpr size_t BLOCKSIZE = 8192;\n\n /* We maintain memory alignment to word boundaries by requiring that all\n allocations be in multiples of the machine wordsize. */\n /* Size of machine word in bytes. Must be power of 2. */\n /* Minimum number of bytes requested at a time from the system. Must be\n * multiple of WORDSIZE. */\n\n using Size = size_t;\n\n Size remaining_ = 0; //!< Number of bytes left in current block of storage\n void* base_ = nullptr; //!< Pointer to base of current block of storage\n void* loc_ = nullptr; //!< Current location in block to next allocate\n\n void internal_init()\n {\n remaining_ = 0;\n base_ = nullptr;\n usedMemory = 0;\n wastedMemory = 0;\n }\n\n public:\n Size usedMemory = 0;\n Size wastedMemory = 0;\n\n /**\n Default constructor. Initializes a new pool.\n */\n PooledAllocator() { internal_init(); }\n\n /**\n * Destructor. Frees all the memory allocated in this pool.\n */\n ~PooledAllocator() { free_all(); }\n\n /** Frees all allocated memory chunks */\n void free_all()\n {\n while (base_ != nullptr)\n {\n // Get pointer to prev block\n void* prev = *(static_cast(base_));\n ::free(base_);\n base_ = prev;\n }\n internal_init();\n }\n\n /**\n * Returns a pointer to a piece of new memory of the given size in bytes\n * allocated from the pool.\n */\n void* malloc(const size_t req_size)\n {\n /* Round size up to a multiple of wordsize. The following expression\n only works for WORDSIZE that is a power of 2, by masking last bits\n of incremented size to zero.\n */\n const Size size = (req_size + (WORDSIZE - 1)) & ~(WORDSIZE - 1);\n\n /* Check whether a new block must be allocated. Note that the first\n word of a block is reserved for a pointer to the previous block.\n */\n if (size > remaining_)\n {\n wastedMemory += remaining_;\n\n /* Allocate new storage. */\n const Size blocksize =\n size > BLOCKSIZE ? size + WORDSIZE : BLOCKSIZE + WORDSIZE;\n\n // use the standard C malloc to allocate memory\n void* m = ::malloc(blocksize);\n if (!m) { throw std::bad_alloc(); }\n\n /* Fill first word of new block with pointer to previous block. */\n static_cast(m)[0] = base_;\n base_ = m;\n\n remaining_ = blocksize - WORDSIZE;\n loc_ = static_cast(m) + WORDSIZE;\n }\n void* rloc = loc_;\n loc_ = static_cast(loc_) + size;\n remaining_ -= size;\n\n usedMemory += size;\n\n return rloc;\n }\n\n /**\n * Allocates (using this pool) a generic type T.\n *\n * Params:\n * count = number of instances to allocate.\n * Returns: pointer (of type T*) to memory buffer\n */\n template \n T* allocate(const size_t count = 1)\n {\n T* mem = static_cast(this->malloc(sizeof(T) * count));\n return mem;\n }\n};\n/** @} */\n\n/** @addtogroup nanoflann_metaprog_grp Auxiliary metaprogramming stuff\n * @{ */\n\n/** Used to declare fixed-size arrays when DIM>0, dynamically-allocated vectors\n * when DIM=-1. Fixed size version for a generic DIM:\n */\ntemplate \nstruct array_or_vector\n{\n using type = std::array;\n};\n/** Dynamic size version */\ntemplate \nstruct array_or_vector<-1, T>\n{\n using type = std::vector;\n};\n\n/** @} */\n\n/** kd-tree base-class\n *\n * Contains the member functions common to the classes KDTreeSingleIndexAdaptor\n * and KDTreeSingleIndexDynamicAdaptor_.\n *\n * \\tparam Derived The name of the class which inherits this class.\n * \\tparam DatasetAdaptor The user-provided adaptor, which must be ensured to\n * have a lifetime equal or longer than the instance of this class.\n * \\tparam Distance The distance metric to use, these are all classes derived\n * from nanoflann::Metric\n * \\tparam DIM Dimensionality of data points (e.g. 3 for 3D points)\n * \\tparam IndexType Type of the arguments with which the data can be\n * accessed (e.g. float, double, int64_t, T*)\n */\ntemplate <\n class Derived, typename Distance, class DatasetAdaptor, int32_t DIM = -1,\n typename index_t = uint32_t>\nclass KDTreeBaseClass\n{\n public:\n /** Frees the previously-built index. Automatically called within\n * buildIndex(). */\n void freeIndex(Derived& obj)\n {\n obj.pool_.free_all();\n obj.root_node_ = nullptr;\n obj.size_at_index_build_ = 0;\n }\n\n using ElementType = typename Distance::ElementType;\n using DistanceType = typename Distance::DistanceType;\n using IndexType = index_t;\n\n /**\n * Array of indices to vectors in the dataset_.\n */\n std::vector vAcc_;\n\n using Offset = typename decltype(vAcc_)::size_type;\n using Size = typename decltype(vAcc_)::size_type;\n using Dimension = int32_t;\n\n /*---------------------------\n * Internal Data Structures\n * --------------------------*/\n struct Node\n {\n /** Union used because a node can be either a LEAF node or a non-leaf\n * node, so both data fields are never used simultaneously */\n union\n {\n struct leaf\n {\n Offset left, right; //!< Indices of points in leaf node\n } lr;\n struct nonleaf\n {\n Dimension divfeat; //!< Dimension used for subdivision.\n /// The values used for subdivision.\n DistanceType divlow, divhigh;\n } sub;\n } node_type;\n\n /** Child nodes (both=nullptr mean its a leaf node) */\n Node *child1 = nullptr, *child2 = nullptr;\n };\n\n using NodePtr = Node*;\n using NodeConstPtr = const Node*;\n\n struct Interval\n {\n ElementType low, high;\n };\n\n NodePtr root_node_ = nullptr;\n\n Size leaf_max_size_ = 0;\n\n /// Number of thread for concurrent tree build\n Size n_thread_build_ = 1;\n /// Number of current points in the dataset\n Size size_ = 0;\n /// Number of points in the dataset when the index was built\n Size size_at_index_build_ = 0;\n Dimension dim_ = 0; //!< Dimensionality of each data point\n\n /** Define \"BoundingBox\" as a fixed-size or variable-size container\n * depending on \"DIM\" */\n using BoundingBox = typename array_or_vector::type;\n\n /** Define \"distance_vector_t\" as a fixed-size or variable-size container\n * depending on \"DIM\" */\n using distance_vector_t = typename array_or_vector::type;\n\n /** The KD-tree used to find neighbours */\n BoundingBox root_bbox_;\n\n /**\n * Pooled memory allocator.\n *\n * Using a pooled memory allocator is more efficient\n * than allocating memory directly when there is a large\n * number small of memory allocations.\n */\n PooledAllocator pool_;\n\n /** Returns number of points in dataset */\n Size size(const Derived& obj) const { return obj.size_; }\n\n /** Returns the length of each point in the dataset */\n Size veclen(const Derived& obj) const { return DIM > 0 ? DIM : obj.dim_; }\n\n /// Helper accessor to the dataset points:\n ElementType dataset_get(\n const Derived& obj, IndexType element, Dimension component) const\n {\n return obj.dataset_.kdtree_get_pt(element, component);\n }\n\n /**\n * Computes the index memory usage\n * Returns: memory used by the index\n */\n Size usedMemory(const Derived& obj) const\n {\n return obj.pool_.usedMemory + obj.pool_.wastedMemory +\n obj.dataset_.kdtree_get_point_count() *\n sizeof(IndexType); // pool memory and vind array memory\n }\n\n /**\n * Compute the minimum and maximum element values in the specified dimension\n */\n void computeMinMax(\n const Derived& obj, Offset ind, Size count, Dimension element,\n ElementType& min_elem, ElementType& max_elem) const\n {\n min_elem = dataset_get(obj, vAcc_[ind], element);\n max_elem = min_elem;\n for (Offset i = 1; i < count; ++i)\n {\n ElementType val = dataset_get(obj, vAcc_[ind + i], element);\n if (val < min_elem) min_elem = val;\n if (val > max_elem) max_elem = val;\n }\n }\n\n /**\n * Create a tree node that subdivides the list of vecs from vind[first]\n * to vind[last]. The routine is called recursively on each sublist.\n *\n * @param left index of the first vector\n * @param right index of the last vector\n * @param bbox bounding box used as input for splitting and output for\n * parent node\n */\n NodePtr divideTree(\n Derived& obj, const Offset left, const Offset right, BoundingBox& bbox)\n {\n assert(left < obj.dataset_.kdtree_get_point_count());\n\n NodePtr node = obj.pool_.template allocate(); // allocate memory\n const auto dims = (DIM > 0 ? DIM : obj.dim_);\n\n /* If too few exemplars remain, then make this a leaf node. */\n if ((right - left) <= static_cast(obj.leaf_max_size_))\n {\n node->child1 = node->child2 = nullptr; /* Mark as leaf node. */\n node->node_type.lr.left = left;\n node->node_type.lr.right = right;\n\n // compute bounding-box of leaf points\n for (Dimension i = 0; i < dims; ++i)\n {\n bbox[i].low = dataset_get(obj, obj.vAcc_[left], i);\n bbox[i].high = dataset_get(obj, obj.vAcc_[left], i);\n }\n for (Offset k = left + 1; k < right; ++k)\n {\n for (Dimension i = 0; i < dims; ++i)\n {\n const auto val = dataset_get(obj, obj.vAcc_[k], i);\n if (bbox[i].low > val) bbox[i].low = val;\n if (bbox[i].high < val) bbox[i].high = val;\n }\n }\n }\n else\n {\n /* Determine the index, dimension and value for split plane */\n Offset idx;\n Dimension cutfeat;\n DistanceType cutval;\n middleSplit_(obj, left, right - left, idx, cutfeat, cutval, bbox);\n\n node->node_type.sub.divfeat = cutfeat;\n\n /* Recurse on left */\n BoundingBox left_bbox(bbox);\n left_bbox[cutfeat].high = cutval;\n node->child1 = this->divideTree(obj, left, left + idx, left_bbox);\n\n /* Recurse on right */\n BoundingBox right_bbox(bbox);\n right_bbox[cutfeat].low = cutval;\n node->child2 = this->divideTree(obj, left + idx, right, right_bbox);\n\n node->node_type.sub.divlow = left_bbox[cutfeat].high;\n node->node_type.sub.divhigh = right_bbox[cutfeat].low;\n\n for (Dimension i = 0; i < dims; ++i)\n {\n bbox[i].low = std::min(left_bbox[i].low, right_bbox[i].low);\n bbox[i].high = std::max(left_bbox[i].high, right_bbox[i].high);\n }\n }\n\n return node;\n }\n\n /**\n * Create a tree node that subdivides the list of vecs from vind[first] to\n * vind[last] concurrently. The routine is called recursively on each\n * sublist.\n *\n * @param left index of the first vector\n * @param right index of the last vector\n * @param bbox bounding box used as input for splitting and output for\n * parent node\n * @param thread_count count of std::async threads\n * @param mutex mutex for mempool allocation\n */\n NodePtr divideTreeConcurrent(\n Derived& obj, const Offset left, const Offset right, BoundingBox& bbox,\n std::atomic& thread_count, std::mutex& mutex)\n {\n std::unique_lock lock(mutex);\n NodePtr node = obj.pool_.template allocate(); // allocate memory\n lock.unlock();\n\n const auto dims = (DIM > 0 ? DIM : obj.dim_);\n\n /* If too few exemplars remain, then make this a leaf node. */\n if ((right - left) <= static_cast(obj.leaf_max_size_))\n {\n node->child1 = node->child2 = nullptr; /* Mark as leaf node. */\n node->node_type.lr.left = left;\n node->node_type.lr.right = right;\n\n // compute bounding-box of leaf points\n for (Dimension i = 0; i < dims; ++i)\n {\n bbox[i].low = dataset_get(obj, obj.vAcc_[left], i);\n bbox[i].high = dataset_get(obj, obj.vAcc_[left], i);\n }\n for (Offset k = left + 1; k < right; ++k)\n {\n for (Dimension i = 0; i < dims; ++i)\n {\n const auto val = dataset_get(obj, obj.vAcc_[k], i);\n if (bbox[i].low > val) bbox[i].low = val;\n if (bbox[i].high < val) bbox[i].high = val;\n }\n }\n }\n else\n {\n /* Determine the index, dimension and value for split plane */\n Offset idx;\n Dimension cutfeat;\n DistanceType cutval;\n middleSplit_(obj, left, right - left, idx, cutfeat, cutval, bbox);\n\n node->node_type.sub.divfeat = cutfeat;\n\n std::future right_future;\n\n /* Recurse on right concurrently, if possible */\n\n BoundingBox right_bbox(bbox);\n right_bbox[cutfeat].low = cutval;\n if (++thread_count < n_thread_build_)\n {\n /* Concurrent thread for right recursion */\n\n right_future = std::async(\n std::launch::async, &KDTreeBaseClass::divideTreeConcurrent,\n this, std::ref(obj), left + idx, right,\n std::ref(right_bbox), std::ref(thread_count),\n std::ref(mutex));\n }\n else { --thread_count; }\n\n /* Recurse on left in this thread */\n\n BoundingBox left_bbox(bbox);\n left_bbox[cutfeat].high = cutval;\n node->child1 = this->divideTreeConcurrent(\n obj, left, left + idx, left_bbox, thread_count, mutex);\n\n if (right_future.valid())\n {\n /* Block and wait for concurrent right from above */\n\n node->child2 = right_future.get();\n --thread_count;\n }\n else\n {\n /* Otherwise, recurse on right in this thread */\n\n node->child2 = this->divideTreeConcurrent(\n obj, left + idx, right, right_bbox, thread_count, mutex);\n }\n\n node->node_type.sub.divlow = left_bbox[cutfeat].high;\n node->node_type.sub.divhigh = right_bbox[cutfeat].low;\n\n for (Dimension i = 0; i < dims; ++i)\n {\n bbox[i].low = std::min(left_bbox[i].low, right_bbox[i].low);\n bbox[i].high = std::max(left_bbox[i].high, right_bbox[i].high);\n }\n }\n\n return node;\n }\n\n void middleSplit_(\n const Derived& obj, const Offset ind, const Size count, Offset& index,\n Dimension& cutfeat, DistanceType& cutval, const BoundingBox& bbox)\n {\n const auto dims = (DIM > 0 ? DIM : obj.dim_);\n const auto EPS = static_cast(0.00001);\n ElementType max_span = bbox[0].high - bbox[0].low;\n for (Dimension i = 1; i < dims; ++i)\n {\n ElementType span = bbox[i].high - bbox[i].low;\n if (span > max_span) { max_span = span; }\n }\n ElementType max_spread = -1;\n cutfeat = 0;\n ElementType min_elem = 0, max_elem = 0;\n for (Dimension i = 0; i < dims; ++i)\n {\n ElementType span = bbox[i].high - bbox[i].low;\n if (span >= (1 - EPS) * max_span)\n {\n ElementType min_elem_, max_elem_;\n computeMinMax(obj, ind, count, i, min_elem_, max_elem_);\n ElementType spread = max_elem_ - min_elem_;\n if (spread > max_spread)\n {\n cutfeat = i;\n max_spread = spread;\n min_elem = min_elem_;\n max_elem = max_elem_;\n }\n }\n }\n // split in the middle\n DistanceType split_val = (bbox[cutfeat].low + bbox[cutfeat].high) / 2;\n\n if (split_val < min_elem)\n cutval = min_elem;\n else if (split_val > max_elem)\n cutval = max_elem;\n else\n cutval = split_val;\n\n Offset lim1, lim2;\n planeSplit(obj, ind, count, cutfeat, cutval, lim1, lim2);\n\n if (lim1 > count / 2)\n index = lim1;\n else if (lim2 < count / 2)\n index = lim2;\n else\n index = count / 2;\n }\n\n /**\n * Subdivide the list of points by a plane perpendicular on the axis\n * corresponding to the 'cutfeat' dimension at 'cutval' position.\n *\n * On return:\n * dataset[ind[0..lim1-1]][cutfeat] < cutval\n * dataset[ind[lim1..lim2-1]][cutfeat] == cutval\n * dataset[ind[lim2..count]][cutfeat] > cutval\n */\n void planeSplit(\n const Derived& obj, const Offset ind, const Size count,\n const Dimension cutfeat, const DistanceType& cutval, Offset& lim1,\n Offset& lim2)\n {\n /* First pass.\n * Determine lim1 with all values less than cutval to the left.\n */\n Offset left = 0;\n Offset right = count - 1;\n for (;;)\n {\n while (left <= right &&\n dataset_get(obj, vAcc_[ind + left], cutfeat) < cutval)\n ++left;\n while (right && left <= right &&\n dataset_get(obj, vAcc_[ind + right], cutfeat) >= cutval)\n --right;\n if (left > right || !right)\n break; // \"!right\" was added to support unsigned Index types\n std::swap(vAcc_[ind + left], vAcc_[ind + right]);\n ++left;\n --right;\n }\n /* Second pass\n * Determine lim2 with all values greater than cutval to the right\n * The middle is used for balancing the tree\n */\n lim1 = left;\n right = count - 1;\n for (;;)\n {\n while (left <= right &&\n dataset_get(obj, vAcc_[ind + left], cutfeat) <= cutval)\n ++left;\n while (right && left <= right &&\n dataset_get(obj, vAcc_[ind + right], cutfeat) > cutval)\n --right;\n if (left > right || !right)\n break; // \"!right\" was added to support unsigned Index types\n std::swap(vAcc_[ind + left], vAcc_[ind + right]);\n ++left;\n --right;\n }\n lim2 = left;\n }\n\n DistanceType computeInitialDistances(\n const Derived& obj, const ElementType* vec,\n distance_vector_t& dists) const\n {\n assert(vec);\n DistanceType dist = DistanceType();\n\n for (Dimension i = 0; i < (DIM > 0 ? DIM : obj.dim_); ++i)\n {\n if (vec[i] < obj.root_bbox_[i].low)\n {\n dists[i] =\n obj.distance_.accum_dist(vec[i], obj.root_bbox_[i].low, i);\n dist += dists[i];\n }\n if (vec[i] > obj.root_bbox_[i].high)\n {\n dists[i] =\n obj.distance_.accum_dist(vec[i], obj.root_bbox_[i].high, i);\n dist += dists[i];\n }\n }\n return dist;\n }\n\n static void save_tree(\n const Derived& obj, std::ostream& stream, const NodeConstPtr tree)\n {\n save_value(stream, *tree);\n if (tree->child1 != nullptr) { save_tree(obj, stream, tree->child1); }\n if (tree->child2 != nullptr) { save_tree(obj, stream, tree->child2); }\n }\n\n static void load_tree(Derived& obj, std::istream& stream, NodePtr& tree)\n {\n tree = obj.pool_.template allocate();\n load_value(stream, *tree);\n if (tree->child1 != nullptr) { load_tree(obj, stream, tree->child1); }\n if (tree->child2 != nullptr) { load_tree(obj, stream, tree->child2); }\n }\n\n /** Stores the index in a binary file.\n * IMPORTANT NOTE: The set of data points is NOT stored in the file, so\n * when loading the index object it must be constructed associated to the\n * same source of data points used while building it. See the example:\n * examples/saveload_example.cpp \\sa loadIndex */\n void saveIndex(const Derived& obj, std::ostream& stream) const\n {\n save_value(stream, obj.size_);\n save_value(stream, obj.dim_);\n save_value(stream, obj.root_bbox_);\n save_value(stream, obj.leaf_max_size_);\n save_value(stream, obj.vAcc_);\n if (obj.root_node_) save_tree(obj, stream, obj.root_node_);\n }\n\n /** Loads a previous index from a binary file.\n * IMPORTANT NOTE: The set of data points is NOT stored in the file, so\n * the index object must be constructed associated to the same source of\n * data points used while building the index. See the example:\n * examples/saveload_example.cpp \\sa loadIndex */\n void loadIndex(Derived& obj, std::istream& stream)\n {\n load_value(stream, obj.size_);\n load_value(stream, obj.dim_);\n load_value(stream, obj.root_bbox_);\n load_value(stream, obj.leaf_max_size_);\n load_value(stream, obj.vAcc_);\n load_tree(obj, stream, obj.root_node_);\n }\n};\n\n/** @addtogroup kdtrees_grp KD-tree classes and adaptors\n * @{ */\n\n/** kd-tree static index\n *\n * Contains the k-d trees and other information for indexing a set of points\n * for nearest-neighbor matching.\n *\n * The class \"DatasetAdaptor\" must provide the following interface (can be\n * non-virtual, inlined methods):\n *\n * \\code\n * // Must return the number of data poins\n * size_t kdtree_get_point_count() const { ... }\n *\n *\n * // Must return the dim'th component of the idx'th point in the class:\n * T kdtree_get_pt(const size_t idx, const size_t dim) const { ... }\n *\n * // Optional bounding-box computation: return false to default to a standard\n * bbox computation loop.\n * // Return true if the BBOX was already computed by the class and returned\n * in \"bb\" so it can be avoided to redo it again.\n * // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3\n * for point clouds) template bool kdtree_get_bbox(BBOX &bb) const\n * {\n * bb[0].low = ...; bb[0].high = ...; // 0th dimension limits\n * bb[1].low = ...; bb[1].high = ...; // 1st dimension limits\n * ...\n * return true;\n * }\n *\n * \\endcode\n *\n * \\tparam DatasetAdaptor The user-provided adaptor, which must be ensured to\n * have a lifetime equal or longer than the instance of this class.\n * \\tparam Distance The distance metric to use: nanoflann::metric_L1,\n * nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc. \\tparam DIM\n * Dimensionality of data points (e.g. 3 for 3D points) \\tparam IndexType Will\n * be typically size_t or int\n */\ntemplate <\n typename Distance, class DatasetAdaptor, int32_t DIM = -1,\n typename index_t = uint32_t>\nclass KDTreeSingleIndexAdaptor\n : public KDTreeBaseClass<\n KDTreeSingleIndexAdaptor,\n Distance, DatasetAdaptor, DIM, index_t>\n{\n public:\n /** Deleted copy constructor*/\n explicit KDTreeSingleIndexAdaptor(\n const KDTreeSingleIndexAdaptor<\n Distance, DatasetAdaptor, DIM, index_t>&) = delete;\n\n /** The data source used by this index */\n const DatasetAdaptor& dataset_;\n\n const KDTreeSingleIndexAdaptorParams indexParams;\n\n Distance distance_;\n\n using Base = typename nanoflann::KDTreeBaseClass<\n nanoflann::KDTreeSingleIndexAdaptor<\n Distance, DatasetAdaptor, DIM, index_t>,\n Distance, DatasetAdaptor, DIM, index_t>;\n\n using Offset = typename Base::Offset;\n using Size = typename Base::Size;\n using Dimension = typename Base::Dimension;\n\n using ElementType = typename Base::ElementType;\n using DistanceType = typename Base::DistanceType;\n using IndexType = typename Base::IndexType;\n\n using Node = typename Base::Node;\n using NodePtr = Node*;\n\n using Interval = typename Base::Interval;\n\n /** Define \"BoundingBox\" as a fixed-size or variable-size container\n * depending on \"DIM\" */\n using BoundingBox = typename Base::BoundingBox;\n\n /** Define \"distance_vector_t\" as a fixed-size or variable-size container\n * depending on \"DIM\" */\n using distance_vector_t = typename Base::distance_vector_t;\n\n /**\n * KDTree constructor\n *\n * Refer to docs in README.md or online in\n * https://github.com/jlblancoc/nanoflann\n *\n * The KD-Tree point dimension (the length of each point in the datase, e.g.\n * 3 for 3D points) is determined by means of:\n * - The \\a DIM template parameter if >0 (highest priority)\n * - Otherwise, the \\a dimensionality parameter of this constructor.\n *\n * @param inputData Dataset with the input features. Its lifetime must be\n * equal or longer than that of the instance of this class.\n * @param params Basically, the maximum leaf node size\n *\n * Note that there is a variable number of optional additional parameters\n * which will be forwarded to the metric class constructor. Refer to example\n * `examples/pointcloud_custom_metric.cpp` for a use case.\n *\n */\n template \n explicit KDTreeSingleIndexAdaptor(\n const Dimension dimensionality, const DatasetAdaptor& inputData,\n const KDTreeSingleIndexAdaptorParams& params, Args&&... args)\n : dataset_(inputData),\n indexParams(params),\n distance_(inputData, std::forward(args)...)\n {\n init(dimensionality, params);\n }\n\n explicit KDTreeSingleIndexAdaptor(\n const Dimension dimensionality, const DatasetAdaptor& inputData,\n const KDTreeSingleIndexAdaptorParams& params = {})\n : dataset_(inputData), indexParams(params), distance_(inputData)\n {\n init(dimensionality, params);\n }\n\n private:\n void init(\n const Dimension dimensionality,\n const KDTreeSingleIndexAdaptorParams& params)\n {\n Base::size_ = dataset_.kdtree_get_point_count();\n Base::size_at_index_build_ = Base::size_;\n Base::dim_ = dimensionality;\n if (DIM > 0) Base::dim_ = DIM;\n Base::leaf_max_size_ = params.leaf_max_size;\n if (params.n_thread_build > 0)\n {\n Base::n_thread_build_ = params.n_thread_build;\n }\n else\n {\n Base::n_thread_build_ =\n std::max(std::thread::hardware_concurrency(), 1u);\n }\n\n if (!(params.flags &\n KDTreeSingleIndexAdaptorFlags::SkipInitialBuildIndex))\n {\n // Build KD-tree:\n buildIndex();\n }\n }\n\n public:\n /**\n * Builds the index\n */\n void buildIndex()\n {\n Base::size_ = dataset_.kdtree_get_point_count();\n Base::size_at_index_build_ = Base::size_;\n init_vind();\n this->freeIndex(*this);\n Base::size_at_index_build_ = Base::size_;\n if (Base::size_ == 0) return;\n computeBoundingBox(Base::root_bbox_);\n // construct the tree\n if (Base::n_thread_build_ == 1)\n {\n Base::root_node_ =\n this->divideTree(*this, 0, Base::size_, Base::root_bbox_);\n }\n else\n {\n#ifndef NANOFLANN_NO_THREADS\n std::atomic thread_count(0u);\n std::mutex mutex;\n Base::root_node_ = this->divideTreeConcurrent(\n *this, 0, Base::size_, Base::root_bbox_, thread_count, mutex);\n#else /* NANOFLANN_NO_THREADS */\n throw std::runtime_error(\"Multithreading is disabled\");\n#endif /* NANOFLANN_NO_THREADS */\n }\n }\n\n /** \\name Query methods\n * @{ */\n\n /**\n * Find set of nearest neighbors to vec[0:dim-1]. Their indices are stored\n * inside the result object.\n *\n * Params:\n * result = the result object in which the indices of the\n * nearest-neighbors are stored vec = the vector for which to search the\n * nearest neighbors\n *\n * \\tparam RESULTSET Should be any ResultSet\n * \\return True if the requested neighbors could be found.\n * \\sa knnSearch, radiusSearch\n *\n * \\note If L2 norms are used, all returned distances are actually squared\n * distances.\n */\n template \n bool findNeighbors(\n RESULTSET& result, const ElementType* vec,\n const SearchParameters& searchParams = {}) const\n {\n assert(vec);\n if (this->size(*this) == 0) return false;\n if (!Base::root_node_)\n throw std::runtime_error(\n \"[nanoflann] findNeighbors() called before building the \"\n \"index.\");\n float epsError = 1 + searchParams.eps;\n\n // fixed or variable-sized container (depending on DIM)\n distance_vector_t dists;\n // Fill it with zeros.\n auto zero = static_cast(0);\n assign(dists, (DIM > 0 ? DIM : Base::dim_), zero);\n DistanceType dist = this->computeInitialDistances(*this, vec, dists);\n searchLevel(result, vec, Base::root_node_, dist, dists, epsError);\n\n if (searchParams.sorted) result.sort();\n\n return result.full();\n }\n\n /**\n * Find all points contained within the specified bounding box. Their\n * indices are stored inside the result object.\n *\n * Params:\n * result = the result object in which the indices of the points\n * within the bounding box are stored\n * bbox = the bounding box defining the search region\n *\n * \\tparam RESULTSET Should be any ResultSet\n * \\return Number of points found within the bounding box.\n * \\sa findNeighbors, knnSearch, radiusSearch\n *\n * \\note The search is inclusive - points on the boundary are included.\n */\n template \n Size findWithinBox(RESULTSET& result, const BoundingBox& bbox) const\n {\n if (this->size(*this) == 0) return 0;\n if (!Base::root_node_)\n throw std::runtime_error(\n \"[nanoflann] findWithinBox() called before building the \"\n \"index.\");\n\n std::stack stack;\n stack.push(Base::root_node_);\n\n while (!stack.empty())\n {\n const NodePtr node = stack.top();\n stack.pop();\n\n // If this is a leaf node, then do check and return.\n if (!node->child1) // (if one node is nullptr, both are)\n {\n for (Offset i = node->node_type.lr.left;\n i < node->node_type.lr.right; ++i)\n {\n if (contains(bbox, Base::vAcc_[i]))\n {\n if (!result.addPoint(0, Base::vAcc_[i]))\n {\n // the resultset doesn't want to receive any more\n // points, we're done searching!\n return result.size();\n }\n }\n }\n }\n else\n {\n const int idx = node->node_type.sub.divfeat;\n const auto low_bound = node->node_type.sub.divlow;\n const auto high_bound = node->node_type.sub.divhigh;\n\n if (bbox[idx].low <= low_bound) stack.push(node->child1);\n if (bbox[idx].high >= high_bound) stack.push(node->child2);\n }\n }\n\n return result.size();\n }\n\n /**\n * Find the \"num_closest\" nearest neighbors to the \\a query_point[0:dim-1].\n * Their indices and distances are stored in the provided pointers to\n * array/vector.\n *\n * \\sa radiusSearch, findNeighbors\n * \\return Number `N` of valid points in the result set.\n *\n * \\note If L2 norms are used, all returned distances are actually squared\n * distances.\n *\n * \\note Only the first `N` entries in `out_indices` and `out_distances`\n * will be valid. Return is less than `num_closest` only if the\n * number of elements in the tree is less than `num_closest`.\n */\n Size knnSearch(\n const ElementType* query_point, const Size num_closest,\n IndexType* out_indices, DistanceType* out_distances) const\n {\n nanoflann::KNNResultSet resultSet(num_closest);\n resultSet.init(out_indices, out_distances);\n findNeighbors(resultSet, query_point);\n return resultSet.size();\n }\n\n /**\n * Find all the neighbors to \\a query_point[0:dim-1] within a maximum\n * radius. The output is given as a vector of pairs, of which the first\n * element is a point index and the second the corresponding distance.\n * Previous contents of \\a IndicesDists are cleared.\n *\n * If searchParams.sorted==true, the output list is sorted by ascending\n * distances.\n *\n * For a better performance, it is advisable to do a .reserve() on the\n * vector if you have any wild guess about the number of expected matches.\n *\n * \\sa knnSearch, findNeighbors, radiusSearchCustomCallback\n * \\return The number of points within the given radius (i.e. indices.size()\n * or dists.size() )\n *\n * \\note If L2 norms are used, search radius and all returned distances\n * are actually squared distances.\n */\n Size radiusSearch(\n const ElementType* query_point, const DistanceType& radius,\n std::vector>& IndicesDists,\n const SearchParameters& searchParams = {}) const\n {\n RadiusResultSet resultSet(\n radius, IndicesDists);\n const Size nFound =\n radiusSearchCustomCallback(query_point, resultSet, searchParams);\n return nFound;\n }\n\n /**\n * Just like radiusSearch() but with a custom callback class for each point\n * found in the radius of the query. See the source of RadiusResultSet<> as\n * a start point for your own classes. \\sa radiusSearch\n */\n template \n Size radiusSearchCustomCallback(\n const ElementType* query_point, SEARCH_CALLBACK& resultSet,\n const SearchParameters& searchParams = {}) const\n {\n findNeighbors(resultSet, query_point, searchParams);\n return resultSet.size();\n }\n\n /**\n * Find the first N neighbors to \\a query_point[0:dim-1] within a maximum\n * radius. The output is given as a vector of pairs, of which the first\n * element is a point index and the second the corresponding distance.\n * Previous contents of \\a IndicesDists are cleared.\n *\n * \\sa radiusSearch, findNeighbors\n * \\return Number `N` of valid points in the result set.\n *\n * \\note If L2 norms are used, all returned distances are actually squared\n * distances.\n *\n * \\note Only the first `N` entries in `out_indices` and `out_distances`\n * will be valid. Return is less than `num_closest` only if the\n * number of elements in the tree is less than `num_closest`.\n */\n Size rknnSearch(\n const ElementType* query_point, const Size num_closest,\n IndexType* out_indices, DistanceType* out_distances,\n const DistanceType& radius) const\n {\n nanoflann::RKNNResultSet resultSet(\n num_closest, radius);\n resultSet.init(out_indices, out_distances);\n findNeighbors(resultSet, query_point);\n return resultSet.size();\n }\n\n /** @} */\n\n public:\n /** Make sure the auxiliary list \\a vind has the same size than the\n * current dataset, and re-generate if size has changed. */\n void init_vind()\n {\n // Create a permutable array of indices to the input vectors.\n Base::size_ = dataset_.kdtree_get_point_count();\n if (Base::vAcc_.size() != Base::size_) Base::vAcc_.resize(Base::size_);\n for (IndexType i = 0; i < static_cast(Base::size_); i++)\n Base::vAcc_[i] = i;\n }\n\n void computeBoundingBox(BoundingBox& bbox)\n {\n const auto dims = (DIM > 0 ? DIM : Base::dim_);\n resize(bbox, dims);\n if (dataset_.kdtree_get_bbox(bbox))\n {\n // Done! It was implemented in derived class\n }\n else\n {\n const Size N = dataset_.kdtree_get_point_count();\n if (!N)\n throw std::runtime_error(\n \"[nanoflann] computeBoundingBox() called but \"\n \"no data points found.\");\n for (Dimension i = 0; i < dims; ++i)\n {\n bbox[i].low = bbox[i].high =\n this->dataset_get(*this, Base::vAcc_[0], i);\n }\n for (Offset k = 1; k < N; ++k)\n {\n for (Dimension i = 0; i < dims; ++i)\n {\n const auto val =\n this->dataset_get(*this, Base::vAcc_[k], i);\n if (val < bbox[i].low) bbox[i].low = val;\n if (val > bbox[i].high) bbox[i].high = val;\n }\n }\n }\n }\n\n bool contains(const BoundingBox& bbox, IndexType idx) const\n {\n const auto dims = (DIM > 0 ? DIM : Base::dim_);\n for (Dimension i = 0; i < dims; ++i)\n {\n const auto point = this->dataset_.kdtree_get_pt(idx, i);\n if (point < bbox[i].low || point > bbox[i].high) return false;\n }\n return true;\n }\n\n /**\n * Performs an exact search in the tree starting from a node.\n * \\tparam RESULTSET Should be any ResultSet\n * \\return true if the search should be continued, false if the results are\n * sufficient\n */\n template \n bool searchLevel(\n RESULTSET& result_set, const ElementType* vec, const NodePtr node,\n DistanceType mindist, distance_vector_t& dists,\n const float epsError) const\n {\n // If this is a leaf node, then do check and return.\n if (!node->child1) // (if one node is nullptr, both are)\n {\n for (Offset i = node->node_type.lr.left;\n i < node->node_type.lr.right; ++i)\n {\n const IndexType accessor = Base::vAcc_[i]; // reorder... : i;\n DistanceType dist = distance_.evalMetric(\n vec, accessor, (DIM > 0 ? DIM : Base::dim_));\n if (dist < result_set.worstDist())\n {\n if (!result_set.addPoint(dist, Base::vAcc_[i]))\n {\n // the resultset doesn't want to receive any more\n // points, we're done searching!\n return false;\n }\n }\n }\n return true;\n }\n\n /* Which child branch should be taken first? */\n Dimension idx = node->node_type.sub.divfeat;\n ElementType val = vec[idx];\n DistanceType diff1 = val - node->node_type.sub.divlow;\n DistanceType diff2 = val - node->node_type.sub.divhigh;\n\n NodePtr bestChild;\n NodePtr otherChild;\n DistanceType cut_dist;\n if ((diff1 + diff2) < 0)\n {\n bestChild = node->child1;\n otherChild = node->child2;\n cut_dist =\n distance_.accum_dist(val, node->node_type.sub.divhigh, idx);\n }\n else\n {\n bestChild = node->child2;\n otherChild = node->child1;\n cut_dist =\n distance_.accum_dist(val, node->node_type.sub.divlow, idx);\n }\n\n /* Call recursively to search next level down. */\n if (!searchLevel(result_set, vec, bestChild, mindist, dists, epsError))\n {\n // the resultset doesn't want to receive any more points, we're done\n // searching!\n return false;\n }\n\n DistanceType dst = dists[idx];\n mindist = mindist + cut_dist - dst;\n dists[idx] = cut_dist;\n if (mindist * epsError <= result_set.worstDist())\n {\n if (!searchLevel(\n result_set, vec, otherChild, mindist, dists, epsError))\n {\n // the resultset doesn't want to receive any more points, we're\n // done searching!\n return false;\n }\n }\n dists[idx] = dst;\n return true;\n }\n\n public:\n /** Stores the index in a binary file.\n * IMPORTANT NOTE: The set of data points is NOT stored in the file, so\n * when loading the index object it must be constructed associated to the\n * same source of data points used while building it. See the example:\n * examples/saveload_example.cpp \\sa loadIndex */\n void saveIndex(std::ostream& stream) const\n {\n Base::saveIndex(*this, stream);\n }\n\n /** Loads a previous index from a binary file.\n * IMPORTANT NOTE: The set of data points is NOT stored in the file, so\n * the index object must be constructed associated to the same source of\n * data points used while building the index. See the example:\n * examples/saveload_example.cpp \\sa loadIndex */\n void loadIndex(std::istream& stream) { Base::loadIndex(*this, stream); }\n\n}; // class KDTree\n\n/** kd-tree dynamic index\n *\n * Contains the k-d trees and other information for indexing a set of points\n * for nearest-neighbor matching.\n *\n * The class \"DatasetAdaptor\" must provide the following interface (can be\n * non-virtual, inlined methods):\n *\n * \\code\n * // Must return the number of data poins\n * size_t kdtree_get_point_count() const { ... }\n *\n * // Must return the dim'th component of the idx'th point in the class:\n * T kdtree_get_pt(const size_t idx, const size_t dim) const { ... }\n *\n * // Optional bounding-box computation: return false to default to a standard\n * bbox computation loop.\n * // Return true if the BBOX was already computed by the class and returned\n * in \"bb\" so it can be avoided to redo it again.\n * // Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3\n * for point clouds) template bool kdtree_get_bbox(BBOX &bb) const\n * {\n * bb[0].low = ...; bb[0].high = ...; // 0th dimension limits\n * bb[1].low = ...; bb[1].high = ...; // 1st dimension limits\n * ...\n * return true;\n * }\n *\n * \\endcode\n *\n * \\tparam DatasetAdaptor The user-provided adaptor (see comments above).\n * \\tparam Distance The distance metric to use: nanoflann::metric_L1,\n * nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc.\n * \\tparam DIM Dimensionality of data points (e.g. 3 for 3D points)\n * \\tparam IndexType Type of the arguments with which the data can be\n * accessed (e.g. float, double, int64_t, T*)\n */\ntemplate <\n typename Distance, class DatasetAdaptor, int32_t DIM = -1,\n typename IndexType = uint32_t>\nclass KDTreeSingleIndexDynamicAdaptor_\n : public KDTreeBaseClass<\n KDTreeSingleIndexDynamicAdaptor_<\n Distance, DatasetAdaptor, DIM, IndexType>,\n Distance, DatasetAdaptor, DIM, IndexType>\n{\n public:\n /**\n * The dataset used by this index\n */\n const DatasetAdaptor& dataset_; //!< The source of our data\n\n KDTreeSingleIndexAdaptorParams index_params_;\n\n std::vector& treeIndex_;\n\n Distance distance_;\n\n using Base = typename nanoflann::KDTreeBaseClass<\n nanoflann::KDTreeSingleIndexDynamicAdaptor_<\n Distance, DatasetAdaptor, DIM, IndexType>,\n Distance, DatasetAdaptor, DIM, IndexType>;\n\n using ElementType = typename Base::ElementType;\n using DistanceType = typename Base::DistanceType;\n\n using Offset = typename Base::Offset;\n using Size = typename Base::Size;\n using Dimension = typename Base::Dimension;\n\n using Node = typename Base::Node;\n using NodePtr = Node*;\n\n using Interval = typename Base::Interval;\n /** Define \"BoundingBox\" as a fixed-size or variable-size container\n * depending on \"DIM\" */\n using BoundingBox = typename Base::BoundingBox;\n\n /** Define \"distance_vector_t\" as a fixed-size or variable-size container\n * depending on \"DIM\" */\n using distance_vector_t = typename Base::distance_vector_t;\n\n /**\n * KDTree constructor\n *\n * Refer to docs in README.md or online in\n * https://github.com/jlblancoc/nanoflann\n *\n * The KD-Tree point dimension (the length of each point in the datase, e.g.\n * 3 for 3D points) is determined by means of:\n * - The \\a DIM template parameter if >0 (highest priority)\n * - Otherwise, the \\a dimensionality parameter of this constructor.\n *\n * @param inputData Dataset with the input features. Its lifetime must be\n * equal or longer than that of the instance of this class.\n * @param params Basically, the maximum leaf node size\n */\n KDTreeSingleIndexDynamicAdaptor_(\n const Dimension dimensionality, const DatasetAdaptor& inputData,\n std::vector& treeIndex,\n const KDTreeSingleIndexAdaptorParams& params =\n KDTreeSingleIndexAdaptorParams())\n : dataset_(inputData),\n index_params_(params),\n treeIndex_(treeIndex),\n distance_(inputData)\n {\n Base::size_ = 0;\n Base::size_at_index_build_ = 0;\n for (auto& v : Base::root_bbox_) v = {};\n Base::dim_ = dimensionality;\n if (DIM > 0) Base::dim_ = DIM;\n Base::leaf_max_size_ = params.leaf_max_size;\n if (params.n_thread_build > 0)\n {\n Base::n_thread_build_ = params.n_thread_build;\n }\n else\n {\n Base::n_thread_build_ =\n std::max(std::thread::hardware_concurrency(), 1u);\n }\n }\n\n /** Explicitly default the copy constructor */\n KDTreeSingleIndexDynamicAdaptor_(\n const KDTreeSingleIndexDynamicAdaptor_& rhs) = default;\n\n /** Assignment operator definiton */\n KDTreeSingleIndexDynamicAdaptor_ operator=(\n const KDTreeSingleIndexDynamicAdaptor_& rhs)\n {\n KDTreeSingleIndexDynamicAdaptor_ tmp(rhs);\n std::swap(Base::vAcc_, tmp.Base::vAcc_);\n std::swap(Base::leaf_max_size_, tmp.Base::leaf_max_size_);\n std::swap(index_params_, tmp.index_params_);\n std::swap(treeIndex_, tmp.treeIndex_);\n std::swap(Base::size_, tmp.Base::size_);\n std::swap(Base::size_at_index_build_, tmp.Base::size_at_index_build_);\n std::swap(Base::root_node_, tmp.Base::root_node_);\n std::swap(Base::root_bbox_, tmp.Base::root_bbox_);\n std::swap(Base::pool_, tmp.Base::pool_);\n return *this;\n }\n\n /**\n * Builds the index\n */\n void buildIndex()\n {\n Base::size_ = Base::vAcc_.size();\n this->freeIndex(*this);\n Base::size_at_index_build_ = Base::size_;\n if (Base::size_ == 0) return;\n computeBoundingBox(Base::root_bbox_);\n // construct the tree\n if (Base::n_thread_build_ == 1)\n {\n Base::root_node_ =\n this->divideTree(*this, 0, Base::size_, Base::root_bbox_);\n }\n else\n {\n#ifndef NANOFLANN_NO_THREADS\n std::atomic thread_count(0u);\n std::mutex mutex;\n Base::root_node_ = this->divideTreeConcurrent(\n *this, 0, Base::size_, Base::root_bbox_, thread_count, mutex);\n#else /* NANOFLANN_NO_THREADS */\n throw std::runtime_error(\"Multithreading is disabled\");\n#endif /* NANOFLANN_NO_THREADS */\n }\n }\n\n /** \\name Query methods\n * @{ */\n\n /**\n * Find set of nearest neighbors to vec[0:dim-1]. Their indices are stored\n * inside the result object.\n * This is the core search function, all others are wrappers around this\n * one.\n *\n * \\param result The result object in which the indices of the\n * nearest-neighbors are stored.\n * \\param vec The vector of the query point for which to search the\n * nearest neighbors.\n * \\param searchParams Optional parameters for the search.\n *\n * \\tparam RESULTSET Should be any ResultSet\n * \\return True if the requested neighbors could be found.\n *\n * \\sa knnSearch(), radiusSearch(), radiusSearchCustomCallback()\n *\n * \\note If L2 norms are used, all returned distances are actually squared\n * distances.\n */\n template \n bool findNeighbors(\n RESULTSET& result, const ElementType* vec,\n const SearchParameters& searchParams = {}) const\n {\n assert(vec);\n if (this->size(*this) == 0) return false;\n if (!Base::root_node_) return false;\n float epsError = 1 + searchParams.eps;\n\n // fixed or variable-sized container (depending on DIM)\n distance_vector_t dists;\n // Fill it with zeros.\n assign(\n dists, (DIM > 0 ? DIM : Base::dim_),\n static_cast(0));\n DistanceType dist = this->computeInitialDistances(*this, vec, dists);\n searchLevel(result, vec, Base::root_node_, dist, dists, epsError);\n return result.full();\n }\n\n /**\n * Find the \"num_closest\" nearest neighbors to the \\a query_point[0:dim-1].\n * Their indices are stored inside the result object. \\sa radiusSearch,\n * findNeighbors\n * \\return Number `N` of valid points in\n * the result set.\n *\n * \\note If L2 norms are used, all returned distances are actually squared\n * distances.\n *\n * \\note Only the first `N` entries in `out_indices` and `out_distances`\n * will be valid. Return may be less than `num_closest` only if the\n * number of elements in the tree is less than `num_closest`.\n */\n Size knnSearch(\n const ElementType* query_point, const Size num_closest,\n IndexType* out_indices, DistanceType* out_distances,\n const SearchParameters& searchParams = {}) const\n {\n nanoflann::KNNResultSet resultSet(num_closest);\n resultSet.init(out_indices, out_distances);\n findNeighbors(resultSet, query_point, searchParams);\n return resultSet.size();\n }\n\n /**\n * Find all the neighbors to \\a query_point[0:dim-1] within a maximum\n * radius. The output is given as a vector of pairs, of which the first\n * element is a point index and the second the corresponding distance.\n * Previous contents of \\a IndicesDists are cleared.\n *\n * If searchParams.sorted==true, the output list is sorted by ascending\n * distances.\n *\n * For a better performance, it is advisable to do a .reserve() on the\n * vector if you have any wild guess about the number of expected matches.\n *\n * \\sa knnSearch, findNeighbors, radiusSearchCustomCallback\n * \\return The number of points within the given radius (i.e. indices.size()\n * or dists.size() )\n *\n * \\note If L2 norms are used, search radius and all returned distances\n * are actually squared distances.\n */\n Size radiusSearch(\n const ElementType* query_point, const DistanceType& radius,\n std::vector>& IndicesDists,\n const SearchParameters& searchParams = {}) const\n {\n RadiusResultSet resultSet(\n radius, IndicesDists);\n const size_t nFound =\n radiusSearchCustomCallback(query_point, resultSet, searchParams);\n return nFound;\n }\n\n /**\n * Just like radiusSearch() but with a custom callback class for each point\n * found in the radius of the query. See the source of RadiusResultSet<> as\n * a start point for your own classes. \\sa radiusSearch\n */\n template \n Size radiusSearchCustomCallback(\n const ElementType* query_point, SEARCH_CALLBACK& resultSet,\n const SearchParameters& searchParams = {}) const\n {\n findNeighbors(resultSet, query_point, searchParams);\n return resultSet.size();\n }\n\n /** @} */\n\n public:\n void computeBoundingBox(BoundingBox& bbox)\n {\n const auto dims = (DIM > 0 ? DIM : Base::dim_);\n resize(bbox, dims);\n\n if (dataset_.kdtree_get_bbox(bbox))\n {\n // Done! It was implemented in derived class\n }\n else\n {\n const Size N = Base::size_;\n if (!N)\n throw std::runtime_error(\n \"[nanoflann] computeBoundingBox() called but \"\n \"no data points found.\");\n for (Dimension i = 0; i < dims; ++i)\n {\n bbox[i].low = bbox[i].high =\n this->dataset_get(*this, Base::vAcc_[0], i);\n }\n for (Offset k = 1; k < N; ++k)\n {\n for (Dimension i = 0; i < dims; ++i)\n {\n const auto val =\n this->dataset_get(*this, Base::vAcc_[k], i);\n if (val < bbox[i].low) bbox[i].low = val;\n if (val > bbox[i].high) bbox[i].high = val;\n }\n }\n }\n }\n\n /**\n * Performs an exact search in the tree starting from a node.\n * \\tparam RESULTSET Should be any ResultSet\n */\n template \n void searchLevel(\n RESULTSET& result_set, const ElementType* vec, const NodePtr node,\n DistanceType mindist, distance_vector_t& dists,\n const float epsError) const\n {\n // If this is a leaf node, then do check and return.\n if (!node->child1) // (if one node is nullptr, both are)\n {\n for (Offset i = node->node_type.lr.left;\n i < node->node_type.lr.right; ++i)\n {\n const IndexType index = Base::vAcc_[i]; // reorder... : i;\n if (treeIndex_[index] == -1) continue;\n DistanceType dist = distance_.evalMetric(\n vec, index, (DIM > 0 ? DIM : Base::dim_));\n if (dist < result_set.worstDist())\n {\n if (!result_set.addPoint(\n static_cast(dist),\n static_cast(\n Base::vAcc_[i])))\n {\n // the resultset doesn't want to receive any more\n // points, we're done searching!\n return; // false;\n }\n }\n }\n return;\n }\n\n /* Which child branch should be taken first? */\n Dimension idx = node->node_type.sub.divfeat;\n ElementType val = vec[idx];\n DistanceType diff1 = val - node->node_type.sub.divlow;\n DistanceType diff2 = val - node->node_type.sub.divhigh;\n\n NodePtr bestChild;\n NodePtr otherChild;\n DistanceType cut_dist;\n if ((diff1 + diff2) < 0)\n {\n bestChild = node->child1;\n otherChild = node->child2;\n cut_dist =\n distance_.accum_dist(val, node->node_type.sub.divhigh, idx);\n }\n else\n {\n bestChild = node->child2;\n otherChild = node->child1;\n cut_dist =\n distance_.accum_dist(val, node->node_type.sub.divlow, idx);\n }\n\n /* Call recursively to search next level down. */\n searchLevel(result_set, vec, bestChild, mindist, dists, epsError);\n\n DistanceType dst = dists[idx];\n mindist = mindist + cut_dist - dst;\n dists[idx] = cut_dist;\n if (mindist * epsError <= result_set.worstDist())\n {\n searchLevel(result_set, vec, otherChild, mindist, dists, epsError);\n }\n dists[idx] = dst;\n }\n\n public:\n /** Stores the index in a binary file.\n * IMPORTANT NOTE: The set of data points is NOT stored in the file, so\n * when loading the index object it must be constructed associated to the\n * same source of data points used while building it. See the example:\n * examples/saveload_example.cpp \\sa loadIndex */\n void saveIndex(std::ostream& stream) { saveIndex(*this, stream); }\n\n /** Loads a previous index from a binary file.\n * IMPORTANT NOTE: The set of data points is NOT stored in the file, so\n * the index object must be constructed associated to the same source of\n * data points used while building the index. See the example:\n * examples/saveload_example.cpp \\sa loadIndex */\n void loadIndex(std::istream& stream) { loadIndex(*this, stream); }\n};\n\n/** kd-tree dynaimic index\n *\n * class to create multiple static index and merge their results to behave as\n * single dynamic index as proposed in Logarithmic Approach.\n *\n * Example of usage:\n * examples/dynamic_pointcloud_example.cpp\n *\n * \\tparam DatasetAdaptor The user-provided adaptor (see comments above).\n * \\tparam Distance The distance metric to use: nanoflann::metric_L1,\n * nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc. \\tparam DIM\n * Dimensionality of data points (e.g. 3 for 3D points) \\tparam IndexType\n * Will be typically size_t or int\n */\ntemplate <\n typename Distance, class DatasetAdaptor, int32_t DIM = -1,\n typename IndexType = uint32_t>\nclass KDTreeSingleIndexDynamicAdaptor\n{\n public:\n using ElementType = typename Distance::ElementType;\n using DistanceType = typename Distance::DistanceType;\n\n using Offset = typename KDTreeSingleIndexDynamicAdaptor_<\n Distance, DatasetAdaptor, DIM>::Offset;\n using Size = typename KDTreeSingleIndexDynamicAdaptor_<\n Distance, DatasetAdaptor, DIM>::Size;\n using Dimension = typename KDTreeSingleIndexDynamicAdaptor_<\n Distance, DatasetAdaptor, DIM>::Dimension;\n\n protected:\n Size leaf_max_size_;\n Size treeCount_;\n Size pointCount_;\n\n /**\n * The dataset used by this index\n */\n const DatasetAdaptor& dataset_; //!< The source of our data\n\n /** treeIndex[idx] is the index of tree in which point at idx is stored.\n * treeIndex[idx]=-1 means that point has been removed. */\n std::vector treeIndex_;\n std::unordered_set removedPoints_;\n\n KDTreeSingleIndexAdaptorParams index_params_;\n\n Dimension dim_; //!< Dimensionality of each data point\n\n using index_container_t = KDTreeSingleIndexDynamicAdaptor_<\n Distance, DatasetAdaptor, DIM, IndexType>;\n std::vector index_;\n\n public:\n /** Get a const ref to the internal list of indices; the number of indices\n * is adapted dynamically as the dataset grows in size. */\n const std::vector& getAllIndices() const\n {\n return index_;\n }\n\n private:\n /** finds position of least significant unset bit */\n int First0Bit(IndexType num)\n {\n int pos = 0;\n while (num & 1)\n {\n num = num >> 1;\n pos++;\n }\n return pos;\n }\n\n /** Creates multiple empty trees to handle dynamic support */\n void init()\n {\n using my_kd_tree_t = KDTreeSingleIndexDynamicAdaptor_<\n Distance, DatasetAdaptor, DIM, IndexType>;\n std::vector index(\n treeCount_,\n my_kd_tree_t(dim_ /*dim*/, dataset_, treeIndex_, index_params_));\n index_ = index;\n }\n\n public:\n Distance distance_;\n\n /**\n * KDTree constructor\n *\n * Refer to docs in README.md or online in\n * https://github.com/jlblancoc/nanoflann\n *\n * The KD-Tree point dimension (the length of each point in the datase, e.g.\n * 3 for 3D points) is determined by means of:\n * - The \\a DIM template parameter if >0 (highest priority)\n * - Otherwise, the \\a dimensionality parameter of this constructor.\n *\n * @param inputData Dataset with the input features. Its lifetime must be\n * equal or longer than that of the instance of this class.\n * @param params Basically, the maximum leaf node size\n */\n explicit KDTreeSingleIndexDynamicAdaptor(\n const int dimensionality, const DatasetAdaptor& inputData,\n const KDTreeSingleIndexAdaptorParams& params =\n KDTreeSingleIndexAdaptorParams(),\n const size_t maximumPointCount = 1000000000U)\n : dataset_(inputData), index_params_(params), distance_(inputData)\n {\n treeCount_ = static_cast(std::log2(maximumPointCount)) + 1;\n pointCount_ = 0U;\n dim_ = dimensionality;\n treeIndex_.clear();\n if (DIM > 0) dim_ = DIM;\n leaf_max_size_ = params.leaf_max_size;\n init();\n const size_t num_initial_points = dataset_.kdtree_get_point_count();\n if (num_initial_points > 0) { addPoints(0, num_initial_points - 1); }\n }\n\n /** Deleted copy constructor*/\n explicit KDTreeSingleIndexDynamicAdaptor(\n const KDTreeSingleIndexDynamicAdaptor<\n Distance, DatasetAdaptor, DIM, IndexType>&) = delete;\n\n /** Add points to the set, Inserts all points from [start, end] */\n void addPoints(IndexType start, IndexType end)\n {\n const Size count = end - start + 1;\n int maxIndex = 0;\n treeIndex_.resize(treeIndex_.size() + count);\n for (IndexType idx = start; idx <= end; idx++)\n {\n const int pos = First0Bit(pointCount_);\n maxIndex = std::max(pos, maxIndex);\n treeIndex_[pointCount_] = pos;\n\n const auto it = removedPoints_.find(idx);\n if (it != removedPoints_.end())\n {\n removedPoints_.erase(it);\n treeIndex_[idx] = pos;\n }\n\n for (int i = 0; i < pos; i++)\n {\n for (int j = 0; j < static_cast(index_[i].vAcc_.size());\n j++)\n {\n index_[pos].vAcc_.push_back(index_[i].vAcc_[j]);\n if (treeIndex_[index_[i].vAcc_[j]] != -1)\n treeIndex_[index_[i].vAcc_[j]] = pos;\n }\n index_[i].vAcc_.clear();\n }\n index_[pos].vAcc_.push_back(idx);\n pointCount_++;\n }\n\n for (int i = 0; i <= maxIndex; ++i)\n {\n index_[i].freeIndex(index_[i]);\n if (!index_[i].vAcc_.empty()) index_[i].buildIndex();\n }\n }\n\n /** Remove a point from the set (Lazy Deletion) */\n void removePoint(size_t idx)\n {\n if (idx >= pointCount_) return;\n removedPoints_.insert(idx);\n treeIndex_[idx] = -1;\n }\n\n /**\n * Find set of nearest neighbors to vec[0:dim-1]. Their indices are stored\n * inside the result object.\n *\n * Params:\n * result = the result object in which the indices of the\n * nearest-neighbors are stored vec = the vector for which to search the\n * nearest neighbors\n *\n * \\tparam RESULTSET Should be any ResultSet\n * \\return True if the requested neighbors could be found.\n * \\sa knnSearch, radiusSearch\n *\n * \\note If L2 norms are used, all returned distances are actually squared\n * distances.\n */\n template \n bool findNeighbors(\n RESULTSET& result, const ElementType* vec,\n const SearchParameters& searchParams = {}) const\n {\n for (size_t i = 0; i < treeCount_; i++)\n {\n index_[i].findNeighbors(result, &vec[0], searchParams);\n }\n return result.full();\n }\n};\n\n/** An L2-metric KD-tree adaptor for working with data directly stored in an\n * Eigen Matrix, without duplicating the data storage. You can select whether a\n * row or column in the matrix represents a point in the state space.\n *\n * Example of usage:\n * \\code\n * Eigen::Matrix mat;\n *\n * // Fill out \"mat\"...\n * using my_kd_tree_t = nanoflann::KDTreeEigenMatrixAdaptor<\n * Eigen::Matrix>;\n *\n * const int max_leaf = 10;\n * my_kd_tree_t mat_index(mat, max_leaf);\n * mat_index.index->...\n * \\endcode\n *\n * \\tparam DIM If set to >0, it specifies a compile-time fixed dimensionality\n * for the points in the data set, allowing more compiler optimizations.\n * \\tparam Distance The distance metric to use: nanoflann::metric_L1,\n * nanoflann::metric_L2, nanoflann::metric_L2_Simple, etc.\n * \\tparam row_major If set to true the rows of the matrix are used as the\n * points, if set to false the columns of the matrix are used as the\n * points.\n */\ntemplate <\n class MatrixType, int32_t DIM = -1, class Distance = nanoflann::metric_L2,\n bool row_major = true>\nstruct KDTreeEigenMatrixAdaptor\n{\n using self_t =\n KDTreeEigenMatrixAdaptor;\n using num_t = typename MatrixType::Scalar;\n using IndexType = typename MatrixType::Index;\n using metric_t = typename Distance::template traits<\n num_t, self_t, IndexType>::distance_t;\n\n using index_t = KDTreeSingleIndexAdaptor<\n metric_t, self_t,\n row_major ? MatrixType::ColsAtCompileTime\n : MatrixType::RowsAtCompileTime,\n IndexType>;\n\n index_t* index_; //! The kd-tree index for the user to call its methods as\n //! usual with any other FLANN index.\n\n using Offset = typename index_t::Offset;\n using Size = typename index_t::Size;\n using Dimension = typename index_t::Dimension;\n\n /// Constructor: takes a const ref to the matrix object with the data points\n explicit KDTreeEigenMatrixAdaptor(\n const Dimension dimensionality,\n const std::reference_wrapper& mat,\n const int leaf_max_size = 10, const unsigned int n_thread_build = 1)\n : m_data_matrix(mat)\n {\n const auto dims = row_major ? mat.get().cols() : mat.get().rows();\n if (static_cast(dims) != dimensionality)\n throw std::runtime_error(\n \"Error: 'dimensionality' must match column count in data \"\n \"matrix\");\n if (DIM > 0 && static_cast(dims) != DIM)\n throw std::runtime_error(\n \"Data set dimensionality does not match the 'DIM' template \"\n \"argument\");\n index_ = new index_t(\n dims, *this /* adaptor */,\n nanoflann::KDTreeSingleIndexAdaptorParams(\n leaf_max_size, nanoflann::KDTreeSingleIndexAdaptorFlags::None,\n n_thread_build));\n }\n\n public:\n /** Deleted copy constructor */\n KDTreeEigenMatrixAdaptor(const self_t&) = delete;\n\n ~KDTreeEigenMatrixAdaptor() { delete index_; }\n\n const std::reference_wrapper m_data_matrix;\n\n /** Query for the \\a num_closest closest points to a given point (entered as\n * query_point[0:dim-1]). Note that this is a short-cut method for\n * index->findNeighbors(). The user can also call index->... methods as\n * desired.\n *\n * \\note If L2 norms are used, all returned distances are actually squared\n * distances.\n */\n void query(\n const num_t* query_point, const Size num_closest,\n IndexType* out_indices, num_t* out_distances) const\n {\n nanoflann::KNNResultSet resultSet(num_closest);\n resultSet.init(out_indices, out_distances);\n index_->findNeighbors(resultSet, query_point);\n }\n\n /** @name Interface expected by KDTreeSingleIndexAdaptor\n * @{ */\n\n const self_t& derived() const { return *this; }\n self_t& derived() { return *this; }\n\n // Must return the number of data points\n Size kdtree_get_point_count() const\n {\n if (row_major)\n return m_data_matrix.get().rows();\n else\n return m_data_matrix.get().cols();\n }\n\n // Returns the dim'th component of the idx'th point in the class:\n num_t kdtree_get_pt(const IndexType idx, size_t dim) const\n {\n if (row_major)\n return m_data_matrix.get().coeff(idx, IndexType(dim));\n else\n return m_data_matrix.get().coeff(IndexType(dim), idx);\n }\n\n // Optional bounding-box computation: return false to default to a standard\n // bbox computation loop.\n // Return true if the BBOX was already computed by the class and returned\n // in \"bb\" so it can be avoided to redo it again. Look at bb.size() to\n // find out the expected dimensionality (e.g. 2 or 3 for point clouds)\n template \n bool kdtree_get_bbox(BBOX& /*bb*/) const\n {\n return false;\n }\n\n /** @} */\n\n}; // end of KDTreeEigenMatrixAdaptor\n/** @} */\n\n/** @} */ // end of grouping\n} // namespace nanoflann\n" }, { "path": "src/wrapper.hpp", "content": "#include \"_ext/KDLineTree.h\"\n#include \"_ext/KDTree.h\"\n#include \n#include \n#include \n#include \n\n#ifndef BUCKET_FPS_MAX_DIM\n#define BUCKET_FPS_MAX_DIM 8\n#endif\nconstexpr size_t max_dim = BUCKET_FPS_MAX_DIM;\n\nusing quickfps::KDLineTree;\nusing quickfps::KDTree;\nusing quickfps::Point;\n\ntemplate \nstd::vector> raw_data_to_points(const float *raw_data,\n size_t n_points, size_t dim) {\n std::vector> points;\n points.reserve(n_points);\n for (size_t i = 0; i < n_points; i++) {\n const float *ptr = raw_data + i * dim;\n points.push_back(Point(ptr, i));\n }\n return points;\n}\n\ntemplate \nvoid kdtree_sample(const float *raw_data, size_t n_points, size_t dim,\n size_t n_samples, size_t start_idx,\n size_t *sampled_point_indices) {\n auto points = raw_data_to_points(raw_data, n_points, dim);\n std::unique_ptr[]> sampled_points(\n new Point[n_samples]);\n KDTree tree(points.data(), n_points, sampled_points.get());\n tree.buildKDtree();\n tree.init(points[start_idx]);\n tree.sample(n_samples);\n for (size_t i = 0; i < n_samples; i++) {\n sampled_point_indices[i] = sampled_points[i].id;\n }\n}\n\ntemplate \nvoid kdline_sample(const float *raw_data, size_t n_points, size_t dim,\n size_t n_samples, size_t start_idx, size_t height,\n size_t *sampled_point_indices) {\n auto points = raw_data_to_points(raw_data, n_points, dim);\n std::unique_ptr[]> sampled_points(\n new Point[n_samples]);\n KDLineTree tree(points.data(), n_points, height,\n sampled_points.get());\n tree.buildKDtree();\n tree.init(points[start_idx]);\n tree.sample(n_samples);\n for (size_t i = 0; i < n_samples; i++) {\n sampled_point_indices[i] = sampled_points[i].id;\n }\n}\n\n////////////////////////////////////////\n// //\n// Compile Time Function Helper //\n// //\n////////////////////////////////////////\nusing KDTreeFuncType = void (*)(const float *, size_t, size_t, size_t, size_t,\n size_t *);\nusing KDLineFuncType = void (*)(const float *, size_t, size_t, size_t, size_t,\n size_t, size_t *);\n\ntemplate \nconstexpr std::array mapIndices(M &&m, std::index_sequence) {\n std::array result{m.template operator()()...};\n return result;\n}\n\ntemplate \nconstexpr std::array map(M m) {\n return mapIndices(m, std::make_index_sequence());\n}\n\ntemplate struct kdtree_func_helper {\n template KDTreeFuncType operator()() {\n return &kdtree_sample;\n }\n};\n\ntemplate struct kdline_func_helper {\n template KDLineFuncType operator()() {\n return &kdline_sample;\n }\n};\n\n/////////////////\n// //\n// C API //\n// //\n/////////////////\n\nextern \"C\" {\nint bucket_fps_kdtree(const float *raw_data, size_t n_points, size_t dim,\n size_t n_samples, size_t start_idx,\n size_t *sampled_point_indices) {\n if (dim == 0 || dim > max_dim) {\n // only support 1 to MAX_DIM dimensions\n return 1;\n } else if (start_idx >= n_points) {\n // start_idx should be smaller than n_samples\n return 2;\n }\n auto func_arr = map(kdtree_func_helper{});\n func_arr[dim - 1](raw_data, n_points, dim, n_samples, start_idx,\n sampled_point_indices);\n return 0;\n}\n\nint bucket_fps_kdline(const float *raw_data, size_t n_points, size_t dim,\n size_t n_samples, size_t start_idx, size_t height,\n size_t *sampled_point_indices) {\n if (dim == 0 || dim > max_dim) {\n // only support 1 to MAX_DIM dimensions\n return 1;\n } else if (start_idx >= n_points) {\n // start_idx should be smaller than n_samples\n return 2;\n }\n auto func_arr = map(kdline_func_helper{});\n func_arr[dim - 1](raw_data, n_points, dim, n_samples, start_idx, height,\n sampled_point_indices);\n return 0;\n}\n}\n" } ]