Skip to main content

Epistemic Nearest Neighbors

Project description

Epistemic Nearest Neighbors

A fast, alternative surrogate for Bayesian optimization

ENN estimates a function's value and associated epistemic uncertainty using a K-Nearest Neighbors model. Queries take $O(N lnK)$ time, where $N$ is the number of observations available for KNN lookups. Compare to an exact GP, which takes $O(N^2)$ time. Additionally, measured running times are very small compared to GPs and other alternative surrogates. [1]

Contents

  • ENN surrogate, EpistemicNearestNeighbors [1]

  • TuRBO optimizer via create_optimizer with config factories

    • turbo_enn_config() - TuRBO-ENN (Rust-backed by default)
    • turbo_zero_config() - TuRBO-zero (Rust-backed)
    • lhd_only_config() - LHD design on every ask() (Rust-backed)
    • turbo_one_config() - TuRBO with GP surrogate (Python fallback until GP is ported) The optimizer has an ask()/tell() interface. All turbo_*() methods follow TuRBO:
    • Generate candidates with RAASP [3] sampling.
    • Select a candidate with Thompson sampling (TuRBO-one), UCB (TuRBO-ENN), or randomly (TURBO-zero).
  • Overview of algorithms: algos.pdf

[1] M. Bafna, Jadhav, S. a., & Sweet, D., (2025). Taking the GP Out of the Loop. arXiv preprint arXiv:2506.12818. https://arxiv.org/abs/2506.12818 [2] Eriksson, D., Pearce, M., Gardner, J. R., Turner, R., & Poloczek, M. (2020). Scalable Global Optimization via Local Bayesian Optimization. Advances in Neural Information Processing Systems, 32. https://arxiv.org/abs/1910.01739 [3] Rashidi, B., Johnstonbaugh, K., & Gao, C. (2024). Cylindrical Thompson Sampling for High-Dimensional Bayesian Optimization. Proceedings of The 27th International Conference on Artificial Intelligence and Statistics (pp. 3502–3510). PMLR. https://proceedings.mlr.press/v238/rashidi24a.html

Installation

pip install ennbo[with-deps] or cargo add ennbo

PyPI wheels are platform-specific (they include the enn.enn_rust native extension). If pip install ennbo gives an import error about enn.enn_rust, install a matching wheel (same OS/arch/Python) or build from source (Rust + linkable Faiss C API; see below).

Demonstration

demo_enn.ipynb - Shows how to use EpistemicNearestNeighbors to build and query an ENN model. demo_turbo_enn.ipynb - Shows how to use TurboOptimizer to optimize the Ackley function.

Installation, MacOS

On my MacBook I can run into problems with dependencies and compatibilities.

On MacOS try:

micromamba env create -n ennbo -f admin/conda-macos.yml
micromamba activate ennbo
pip install --no-deps ennbo
pytest -sv tests

You may replace micromamba with conda and this will probably still work.

The commands above make sure

  • You use the MacOS-specific PyTorch (with mps).
  • You avoid having multiple, competing OpenMPs installed PyTorch issue faiss issue.
  • You use old enough versions of NumPy and PyTorch to be compatible with faiss faiss issue.
  • Prevent matplotlib's installation from upgrading your NumPy to an incompatible version.
  • ennbo's listed dependencies do not undo any of the above (which is fine b/c the above commands set the up correctly).

Run tests with

pytest -x -sv tests

and they should all pass fairly quickly (~10s-30s).

If your code still crashes or hangs, try this hack:

export KMP_DUPLICATE_LIB_OK=TRUE
export OMP_NUM_THREADS=1

I don't recommend this, however, as it will slow things down.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distributions

No source distribution files available for this release.See tutorial on generating distribution archives.

Built Distributions

If you're not sure about the file name format, learn more about wheel file names.

ennbo-0.3.10-cp312-cp312-manylinux_2_31_x86_64.whl (6.5 MB view details)

Uploaded CPython 3.12manylinux: glibc 2.31+ x86-64

ennbo-0.3.10-cp312-cp312-manylinux_2_25_x86_64.whl (6.4 MB view details)

Uploaded CPython 3.12manylinux: glibc 2.25+ x86-64

ennbo-0.3.10-cp311-cp311-macosx_26_0_arm64.whl (12.0 MB view details)

Uploaded CPython 3.11macOS 26.0+ ARM64

File details

Details for the file ennbo-0.3.10-cp312-cp312-manylinux_2_31_x86_64.whl.

File metadata

  • Download URL: ennbo-0.3.10-cp312-cp312-manylinux_2_31_x86_64.whl
  • Upload date:
  • Size: 6.5 MB
  • Tags: CPython 3.12, manylinux: glibc 2.31+ x86-64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.9.21 {"installer":{"name":"uv","version":"0.9.21","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"macOS","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}

File hashes

Hashes for ennbo-0.3.10-cp312-cp312-manylinux_2_31_x86_64.whl
Algorithm Hash digest
SHA256 ab222b5a2dcefb1570b04c5bc5371af084b4826b6795e7544cec77ec258ea9f8
MD5 58179af8832f266b5f3d2decf2e361ef
BLAKE2b-256 1e93182d0ffdf623f4f80d9fd714dfa7e84b32600ad3791b6aeb138bfbd28e35

See more details on using hashes here.

File details

Details for the file ennbo-0.3.10-cp312-cp312-manylinux_2_25_x86_64.whl.

File metadata

  • Download URL: ennbo-0.3.10-cp312-cp312-manylinux_2_25_x86_64.whl
  • Upload date:
  • Size: 6.4 MB
  • Tags: CPython 3.12, manylinux: glibc 2.25+ x86-64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.9.21 {"installer":{"name":"uv","version":"0.9.21","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"macOS","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}

File hashes

Hashes for ennbo-0.3.10-cp312-cp312-manylinux_2_25_x86_64.whl
Algorithm Hash digest
SHA256 64a0154b0f1b388dbe89672e61e6e608d821821e3e3415dc97355e198f9f5317
MD5 42d0541ca596fc6f3c69f325becf7136
BLAKE2b-256 72a1a28d3fe8015fc4b8f4646c239080d0e2874612d0707563c769ec351d396e

See more details on using hashes here.

File details

Details for the file ennbo-0.3.10-cp311-cp311-macosx_26_0_arm64.whl.

File metadata

  • Download URL: ennbo-0.3.10-cp311-cp311-macosx_26_0_arm64.whl
  • Upload date:
  • Size: 12.0 MB
  • Tags: CPython 3.11, macOS 26.0+ ARM64
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.9.21 {"installer":{"name":"uv","version":"0.9.21","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"macOS","version":null,"id":null,"libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}

File hashes

Hashes for ennbo-0.3.10-cp311-cp311-macosx_26_0_arm64.whl
Algorithm Hash digest
SHA256 572bb936ff0cd7d460bb57f92ef0657710616845bba7bb6bf4db847ac04f7c60
MD5 6de20fe8b2cfb17b8b0bfc9efb18b3cd
BLAKE2b-256 87fd1c0ffd9e2f2500c3e5abb757522da524d6af2720e7df6d312e641d181941

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page