ennbo 0.0.7__tar.gz → 0.1.1__tar.gz

ennbo-0.1.1/.gitignore

@@ -0,0 +1,6 @@
+*~
+dist/
+style.md
+requirements.md
+.cursorrules
+__pycache__/

{ennbo-0.0.7 → ennbo-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ennbo
-Version: 0.0.7
+Version: 0.1.1
 Summary: Epistemic Nearest Neighbors
 Project-URL: Homepage, https://github.com/yubo-research/enn
 Project-URL: Source, https://github.com/yubo-research/enn
@@ -36,8 +36,9 @@ Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Scientific/Engineering :: Mathematics
 Requires-Python: >=3.11
-Requires-Dist: faiss-cpu==1.9.0
+Requires-Dist: faiss-cpu>=1.9.0
 Requires-Dist: gpytorch==1.13
+Requires-Dist: nds==0.4.3
 Requires-Dist: numpy==1.26.4
 Requires-Dist: scipy==1.15.3
 Requires-Dist: torch==2.5.1
@@ -49,27 +50,60 @@ 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 model, [`EpistemicNearestNeighbors`](https://github.com/yubo-research/enn/blob/main/src/enn/core.py) [1]
-- TuRBO-ENN optimizer, class [`TurboOptimizer`](https://github.com/yubo-research/enn/blob/main/src/enn/turbo_optimizer.py) has four modes
+- ENN model, [`EpistemicNearestNeighbors`](https://github.com/yubo-research/enn/blob/main/src/enn/enn/enn.py) [1]
+- TuRBO-ENN optimizer, class [`TurboOptimizer`](https://github.com/yubo-research/enn/blob/main/src/enn/turbo/turbo_optimizer.py) has four modes
 	- `TURBO_ONE` - A clone of the TuRBO [2] reference [code](https://github.com/uber-research/TuRBO), reworked to have an `ask()`/`tell()` interface.
 	- `TURBO_ENN` - Same as TURBO_ONE, except uses ENN instead of GP and Pareto(mu, se) instead of Thompson sampling.
 	- `TURBO_ZERO` - Same as TURBO_ONE, except randomly-chosen RAASP [3] candidates are picked to be proposals. There is no surrogate.
 	- `LHD_ONLY` - Just creates an LHD design for every `ask()`. Good for a baseline and for testing.
 
-[1] **Sweet, D., & Jadhav, S. A. (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  
+[1] **Sweet, D., & Jadhav, S. A. (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`
 
 ## Demonstration
-[`demo_enn.ipynb`](https://github.com/yubo-research/enn/tree/main/examples) - Shows how to use [`EpistemicNearestNeighbors`](https://github.com/yubo-research/enn/blob/main/src/enn/core.py) to build and query an ENN model.  
-[`demo_turbo_enn.ipynb`](https://github.com/yubo-research/enn/tree/main/examples) - Shows how to use [`TurboOptimizer`](https://github.com/yubo-research/enn/blob/main/src/enn/turbo_optimizer.py) to optimize the Ackley function.
+[`demo_enn.ipynb`](https://github.com/yubo-research/enn/tree/main/examples/demo_enn.ipynb) - Shows how to use [`EpistemicNearestNeighbors`](https://github.com/yubo-research/enn/blob/main/src/enn/enn/enn.py) to build and query an ENN model.
+[`demo_turbo_enn.ipynb`](https://github.com/yubo-research/enn/tree/main/examples/demo_turbo_enn.ipynb) - Shows how to use [`TurboOptimizer`](https://github.com/yubo-research/enn/blob/main/src/enn/turbo/turbo_optimizer.py) 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 conda-macos.yml
+micromamba activate ennbo
+pip install --no-deps ennbo
+```
+
+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](https://github.com/pytorch/pytorch/issues/44282) [faiss issue](https://github.com/faiss-wheels/faiss-wheels/issues/40).
+- You use old enough versions of NumPy and PyTorch to be compatible with faiss [faiss issue](https://github.com/faiss-wheels/faiss-wheels/issues/104).
+- 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 your, try this [hack](https://discuss.pytorch.org/t/ran-into-this-issue-while-executing/101460):
+```
+export KMP_DUPLICATE_LIB_OK=TRUE
+export OMP_NUM_THREADS=1
+```
+I don't recommend this, however, as it will slow things down.

ennbo-0.1.1/README.md

@@ -0,0 +1,63 @@
+# 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 model, [`EpistemicNearestNeighbors`](https://github.com/yubo-research/enn/blob/main/src/enn/enn/enn.py) [1]
+- TuRBO-ENN optimizer, class [`TurboOptimizer`](https://github.com/yubo-research/enn/blob/main/src/enn/turbo/turbo_optimizer.py) has four modes
+	- `TURBO_ONE` - A clone of the TuRBO [2] reference [code](https://github.com/uber-research/TuRBO), reworked to have an `ask()`/`tell()` interface.
+	- `TURBO_ENN` - Same as TURBO_ONE, except uses ENN instead of GP and Pareto(mu, se) instead of Thompson sampling.
+	- `TURBO_ZERO` - Same as TURBO_ONE, except randomly-chosen RAASP [3] candidates are picked to be proposals. There is no surrogate.
+	- `LHD_ONLY` - Just creates an LHD design for every `ask()`. Good for a baseline and for testing.
+
+[1] **Sweet, D., & Jadhav, S. A. (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`
+
+## Demonstration
+[`demo_enn.ipynb`](https://github.com/yubo-research/enn/tree/main/examples/demo_enn.ipynb) - Shows how to use [`EpistemicNearestNeighbors`](https://github.com/yubo-research/enn/blob/main/src/enn/enn/enn.py) to build and query an ENN model.
+[`demo_turbo_enn.ipynb`](https://github.com/yubo-research/enn/tree/main/examples/demo_turbo_enn.ipynb) - Shows how to use [`TurboOptimizer`](https://github.com/yubo-research/enn/blob/main/src/enn/turbo/turbo_optimizer.py) 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 conda-macos.yml
+micromamba activate ennbo
+pip install --no-deps ennbo
+```
+
+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](https://github.com/pytorch/pytorch/issues/44282) [faiss issue](https://github.com/faiss-wheels/faiss-wheels/issues/40).
+- You use old enough versions of NumPy and PyTorch to be compatible with faiss [faiss issue](https://github.com/faiss-wheels/faiss-wheels/issues/104).
+- 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 your, try this [hack](https://discuss.pytorch.org/t/ran-into-this-issue-while-executing/101460):
+```
+export KMP_DUPLICATE_LIB_OK=TRUE
+export OMP_NUM_THREADS=1
+```
+I don't recommend this, however, as it will slow things down.

ennbo-0.1.1/admin/conda-macos.yml

@@ -0,0 +1,16 @@
+channels:
+  - conda-forge
+  - pkgs/main
+dependencies:
+  - python=3.11.14=h18782d2_2_cpython
+  - faiss-cpu=1.9.0=hd610c6a_0
+  - nomkl=1.0=h5ca1d4c_0
+  - llvm-openmp=20.1.2=hdb05f8b_1
+  - gpytorch=1.13=pyh101cb37_1
+  - pytorch=2.5.1=gpu_mps_py311h80af30b_202
+  - scipy=1.15.3=py311hb81b3df_0
+  - pytest=8.3.4=pyhd8ed1ab_1
+  - ipykernel=6.29.5=pyh57ce528_0
+  - pip:
+    - matplotlib==3.10.8
+    - nds==0.4.3

ennbo-0.1.1/examples/demo_enn.ipynb

@@ -0,0 +1,127 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "534cb992",
+   "metadata": {},
+   "source": [
+    "# Epistemic Nearest Neighbors (ENN)\n",
+    "\n",
+    "ENN is a non-parametric surrogate with $O(N)$ computation-time scaling, where $N$ is the number of observations in the data set. ENN can be used in Bayesian optimization as a scalable alternative to a GP (which scales as $O(N^2)$.)\n",
+    "\n",
+    "**Sweet, D., & Jadhav, S. A. (2025).** Taking the GP Out of the Loop. *arXiv preprint arXiv:2506.12818*.  \n",
+    "   https://arxiv.org/abs/2506.12818\n",
+    "\n",
+    "   ---"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8792c830",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "from enn import EpistemicNearestNeighbors, enn_fit\n",
+    "\n",
+    "\n",
+    "def plot_enn_demo(ax, num_samples: int, k: int, noise: float, m: int = 1) -> None:\n",
+    "    x = np.sort(np.random.rand(num_samples + 4))\n",
+    "    x[-3] = x[-4]\n",
+    "    x[-2] = x[-4]\n",
+    "    x[-1] = x[-4]\n",
+    "    x[1] = x[0] + 0.03\n",
+    "    eps = np.random.randn(num_samples + 4)\n",
+    "    y = np.sin(2 * m * np.pi * x) + noise * eps\n",
+    "    yvar = (noise**2) * np.ones_like(y)\n",
+    "    train_x = x[:, None]\n",
+    "    train_y = y[:, None]\n",
+    "    train_yvar = yvar[:, None]\n",
+    "    model = EpistemicNearestNeighbors(\n",
+    "        train_x,\n",
+    "        train_y,\n",
+    "        train_yvar,\n",
+    "    )\n",
+    "    rng = np.random.default_rng(0)\n",
+    "    result = enn_fit(\n",
+    "        model,\n",
+    "        k=k,\n",
+    "        num_fit_candidates=100,\n",
+    "        num_fit_samples=min(10, num_samples),\n",
+    "        rng=rng,\n",
+    "    )\n",
+    "    print(k, noise, result)\n",
+    "    params = result\n",
+    "    x_hat = np.linspace(0.0, 1.0, 30)\n",
+    "    x_hat_2d = x_hat[:, None]\n",
+    "    posterior = model.posterior(x_hat_2d, params=params, exclude_nearest=False)\n",
+    "    mu = posterior.mu[:, 0]\n",
+    "    se = posterior.se[:, 0]\n",
+    "    marker_size = 3 if num_samples >= 100 else 15\n",
+    "    ax.scatter(x, y, s=marker_size, color=\"black\", alpha=0.5)\n",
+    "    ax.plot(x_hat, mu, linestyle=\"--\", color=\"tab:blue\", alpha=0.7)\n",
+    "    ax.fill_between(x_hat, mu - 2 * se, mu + 2 * se, color=\"tab:blue\", alpha=0.2)\n",
+    "    ax.set_ylim(-5, 5)\n",
+    "    ax.set_title(f\"n={num_samples}, noise={noise}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "992d16f9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "# Blue area is the epistemic uncertainty only\n",
+    "\n",
+    "k = 5\n",
+    "fig, axes = plt.subplots(2, 3, figsize=(9, 6), sharex=True, sharey=True)\n",
+    "num_samples_list = [5, 10]\n",
+    "noise_list = [0.0, 0.1, 0.3]\n",
+    "for row_idx, num_samples in enumerate(num_samples_list):\n",
+    "    for col_idx, noise in enumerate(noise_list):\n",
+    "        ax = axes[row_idx, col_idx]\n",
+    "        np.random.seed(4)\n",
+    "        plot_enn_demo(ax, num_samples=num_samples, k=k, noise=noise)\n",
+    "for ax in axes[-1, :]:\n",
+    "    ax.set_xlabel(\"x\")\n",
+    "for ax in axes[:, 0]:\n",
+    "    ax.set_ylabel(\"y\")\n",
+    "fig.tight_layout()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "84968071",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import time\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "np.random.seed(1)\n",
+    "fig, ax = plt.subplots(figsize=(5, 3))\n",
+    "t_0 = time.time()\n",
+    "plot_enn_demo(ax, num_samples=1_000_000, k=5, noise=0.3, m=3)\n",
+    "t_1 = time.time()\n",
+    "print(f\"Time taken: {t_1 - t_0:.2f} seconds\")\n",
+    "ax.set_xlabel(\"x\")\n",
+    "ax.set_ylabel(\"y\")\n",
+    "fig.tight_layout()"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python",
+   "pygments_lexer": "ipython3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}