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

Uploaded CPython 3.12manylinux: glibc 2.31+ x86-64

ennbo-0.3.11-cp312-cp312-manylinux_2_25_x86_64.whl (6.5 MB view details)

Uploaded CPython 3.12manylinux: glibc 2.25+ x86-64

ennbo-0.3.11-cp311-cp311-macosx_26_0_arm64.whl (12.1 MB view details)

Uploaded CPython 3.11macOS 26.0+ ARM64

File details

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

File metadata

  • Download URL: ennbo-0.3.11-cp312-cp312-manylinux_2_31_x86_64.whl
  • Upload date:
  • Size: 6.6 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.11-cp312-cp312-manylinux_2_31_x86_64.whl
Algorithm Hash digest
SHA256 5961741b99b4781f45a99b5920e2c32ed53e2cb338ac55d60f38c732ed0f5e53
MD5 4a2e98240b468f417ae541bdf52399b0
BLAKE2b-256 f071a5b201f85882509bc318cee2cf6e0e209caa4c98f071a6ff609a1eaf5a1e

See more details on using hashes here.

File details

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

File metadata

  • Download URL: ennbo-0.3.11-cp312-cp312-manylinux_2_25_x86_64.whl
  • Upload date:
  • Size: 6.5 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.11-cp312-cp312-manylinux_2_25_x86_64.whl
Algorithm Hash digest
SHA256 783dc88eacd514f0fd5efffd457f593dba1c68ce02d67ec51d130d59acf7150b
MD5 acd6399aa5ccbbf8ef47149c43750970
BLAKE2b-256 57bf1ecf9465275697a030cbbffdd79a9fdb21dda8d8a283ca195108f7eb8bff

See more details on using hashes here.

File details

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

File metadata

  • Download URL: ennbo-0.3.11-cp311-cp311-macosx_26_0_arm64.whl
  • Upload date:
  • Size: 12.1 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.11-cp311-cp311-macosx_26_0_arm64.whl
Algorithm Hash digest
SHA256 ca82a7681d2811ee29760f532705455c76714487cd1ac28b0c9889264500c5dc
MD5 12dd422ad1a6e6b666f5649f48359a27
BLAKE2b-256 f0906891b6c9f795bb7e534873126b7582123140b7746009036600dbfad7382e

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