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.12-cp312-cp312-manylinux_2_31_x86_64.whl (6.8 MB view details)

Uploaded CPython 3.12manylinux: glibc 2.31+ x86-64

ennbo-0.3.12-cp312-cp312-manylinux_2_25_x86_64.whl (6.6 MB view details)

Uploaded CPython 3.12manylinux: glibc 2.25+ x86-64

ennbo-0.3.12-cp311-cp311-macosx_26_0_arm64.whl (12.3 MB view details)

Uploaded CPython 3.11macOS 26.0+ ARM64

File details

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

File metadata

  • Download URL: ennbo-0.3.12-cp312-cp312-manylinux_2_31_x86_64.whl
  • Upload date:
  • Size: 6.8 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.12-cp312-cp312-manylinux_2_31_x86_64.whl
Algorithm Hash digest
SHA256 519e0094f8b1c1a218368d6ea831806c5175fd04046b5f939c18793cb205abc3
MD5 0dd3c318bc743a4ee3299baf81b4edf4
BLAKE2b-256 72825d134ad7f440a2f72759a9ea1a74b779e8d8bd24e6ac053b494d6de92bf6

See more details on using hashes here.

File details

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

File metadata

  • Download URL: ennbo-0.3.12-cp312-cp312-manylinux_2_25_x86_64.whl
  • Upload date:
  • Size: 6.6 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.12-cp312-cp312-manylinux_2_25_x86_64.whl
Algorithm Hash digest
SHA256 8c85660723f22dda0741f0886f1a1fb0e73049eccd82f052a2acf62ef1a9f2b8
MD5 ed2e1f968fe95faebd0d71fe7a421f8b
BLAKE2b-256 1011135bdb6ae45958d0c2fb51d78d9d7f2452fea8fedceb3c8e9dbac8c81c3c

See more details on using hashes here.

File details

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

File metadata

  • Download URL: ennbo-0.3.12-cp311-cp311-macosx_26_0_arm64.whl
  • Upload date:
  • Size: 12.3 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.12-cp311-cp311-macosx_26_0_arm64.whl
Algorithm Hash digest
SHA256 1aaaea0d1bec708928e19f947674fdbe27cf4b68db3a9d8d876fe980d957bbd6
MD5 a734838307220e414f60ef8e6d547471
BLAKE2b-256 090b68d23dd541780577b1249dd26332f68e3f2d51b317ef7ff4f73823ee56b7

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