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

Uploaded CPython 3.12manylinux: glibc 2.31+ x86-64

ennbo-0.3.13-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.13-cp311-cp311-macosx_26_0_arm64.whl (12.2 MB view details)

Uploaded CPython 3.11macOS 26.0+ ARM64

File details

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

File metadata

  • Download URL: ennbo-0.3.13-cp312-cp312-manylinux_2_31_x86_64.whl
  • Upload date:
  • Size: 6.7 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.13-cp312-cp312-manylinux_2_31_x86_64.whl
Algorithm Hash digest
SHA256 9f83530d42cd5346ca3c75f62b0522b96cc7eecba758ec5d299e61143b7fca0d
MD5 7707d71bcf2969bfece7030d96333c8a
BLAKE2b-256 d917f8e095788cfbd7bb1d246cef029beab2e0888aa57ca42b6ae1e1cf5d4088

See more details on using hashes here.

File details

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

File metadata

  • Download URL: ennbo-0.3.13-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.13-cp312-cp312-manylinux_2_25_x86_64.whl
Algorithm Hash digest
SHA256 e9c950cf4bb71215611f4dddffc8f8af9a41f9eca47cf1c5d502d7d75add0e90
MD5 ebf405863570fb4c5c90a36333a8c5e4
BLAKE2b-256 a6c36a5d24b9d41b91594b500b412476d455fa1251c66098f0e65d2b93ece2a9

See more details on using hashes here.

File details

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

File metadata

  • Download URL: ennbo-0.3.13-cp311-cp311-macosx_26_0_arm64.whl
  • Upload date:
  • Size: 12.2 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.13-cp311-cp311-macosx_26_0_arm64.whl
Algorithm Hash digest
SHA256 4ba6a3b95739048957a7e6625df33420b6901c111a011e0d9ea81ebf5c48a108
MD5 49ce58d6008a4774993ec5570a8517b7
BLAKE2b-256 d307d117515450fe0ee86ef90e78c0cb3800256c034feb05f74c9cd834949377

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