ennbo 0.1.1__tar.gz → 0.1.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ennbo
-Version: 0.1.1
+Version: 0.1.2
 Summary: Epistemic Nearest Neighbors
 Project-URL: Homepage, https://github.com/yubo-research/enn
 Project-URL: Source, https://github.com/yubo-research/enn
@@ -39,7 +39,7 @@ Requires-Python: >=3.11
 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: numpy<2.0.0,>=1.26.4
 Requires-Dist: scipy==1.15.3
 Requires-Dist: torch==2.5.1
 Description-Content-Type: text/markdown
@@ -80,9 +80,10 @@ 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 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.

{ennbo-0.1.1 → ennbo-0.1.2}/README.md

@@ -34,9 +34,10 @@ 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 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.

{ennbo-0.1.1 → ennbo-0.1.2}/examples/demo_enn.ipynb

@@ -114,6 +114,14 @@
     "ax.set_ylabel(\"y\")\n",
     "fig.tight_layout()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7abb3c1a",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

ennbo-0.1.2/examples/demo_morbo_enn.ipynb

@@ -0,0 +1,253 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# MORBO-ENN: Multi-Objective Optimization\n",
+    "\n",
+    "This notebook demonstrates multi-objective Bayesian optimization using the MORBO trust region with ENN surrogate."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "---"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "\n",
+    "def ackley_core(\n",
+    "    x: np.ndarray, a: float = 20.0, b: float = 0.2, c: float = 2 * np.pi\n",
+    ") -> np.ndarray:\n",
+    "    if x.ndim == 1:\n",
+    "        x = x[None, :]\n",
+    "    x = x - 1\n",
+    "    term1 = -a * np.exp(-b * np.sqrt((x**2).mean(axis=1)))\n",
+    "    term2 = -np.exp(np.cos(c * x).mean(axis=1))\n",
+    "    return term1 + term2 + a + np.e\n",
+    "\n",
+    "\n",
+    "class DoubleAckley:\n",
+    "    def __init__(self, noise: float, rng):\n",
+    "        self.noise = noise\n",
+    "        self.rng = rng\n",
+    "        self.bounds = [-32.768, 32.768]\n",
+    "\n",
+    "    def __call__(self, x: np.ndarray) -> np.ndarray:\n",
+    "        x = np.asarray(x, dtype=float)\n",
+    "        if x.ndim == 1:\n",
+    "            x = x[None, :]\n",
+    "        n, d = x.shape\n",
+    "        assert d % 2 == 0, \"num_dim must be even for DoubleAckley\"\n",
+    "\n",
+    "        mid = d // 2\n",
+    "        x1 = x[:, :mid]\n",
+    "        x2 = x[:, mid:]\n",
+    "\n",
+    "        y1 = -ackley_core(x1) + self.noise * self.rng.normal(size=n)\n",
+    "        y2 = -ackley_core(x2) + self.noise * self.rng.normal(size=n)\n",
+    "\n",
+    "        return np.stack([y1, y2], axis=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from botorch.utils.multi_objective.hypervolume import Hypervolume\n",
+    "from nds import ndomsort\n",
+    "\n",
+    "\n",
+    "def get_pareto_front_indices(y: np.ndarray) -> np.ndarray:\n",
+    "    y = np.asarray(y, dtype=float)\n",
+    "    if y.ndim != 2:\n",
+    "        raise ValueError(y.shape)\n",
+    "    front_ids = ndomsort.non_domin_sort(-y, only_front_indices=True)\n",
+    "    return np.array([i for i, f in enumerate(front_ids) if f == 0], dtype=int)\n",
+    "\n",
+    "\n",
+    "def compute_hypervolume(y: np.ndarray, ref_point: np.ndarray) -> float:\n",
+    "    if len(y) == 0:\n",
+    "        return 0.0\n",
+    "\n",
+    "    pareto_idx = get_pareto_front_indices(y)\n",
+    "    if len(pareto_idx) == 0:\n",
+    "        return 0.0\n",
+    "\n",
+    "    pareto_y = y[pareto_idx]\n",
+    "    pareto_y_tensor = torch.tensor(pareto_y, dtype=torch.float64)\n",
+    "    ref_point_tensor = torch.tensor(ref_point, dtype=torch.float64)\n",
+    "\n",
+    "    hv = Hypervolume(ref_point=ref_point_tensor)\n",
+    "    return hv.compute(pareto_y_tensor)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import time\n",
+    "\n",
+    "from enn import Turbo, TurboMode\n",
+    "from enn.turbo.turbo_config import (\n",
+    "    LHDOnlyConfig,\n",
+    "    TurboENNConfig,\n",
+    "    TurboOneConfig,\n",
+    "    TurboZeroConfig,\n",
+    ")\n",
+    "\n",
+    "\n",
+    "def run_optimization(turbo_mode: TurboMode):\n",
+    "    num_dim = 30\n",
+    "    num_iterations = 100\n",
+    "    num_arms = 10\n",
+    "    noise = 0.1\n",
+    "    ref_point = np.array([-25.0, -25.0])\n",
+    "\n",
+    "    rng = np.random.default_rng(42)\n",
+    "    objective = DoubleAckley(noise=noise, rng=rng)\n",
+    "    bounds = np.array([objective.bounds] * num_dim, dtype=float)\n",
+    "\n",
+    "    common_cfg = dict(k=10, tr_type=\"morbo\", num_metrics=2)\n",
+    "    if turbo_mode == TurboMode.TURBO_ENN:\n",
+    "        config = TurboENNConfig(**common_cfg, num_fit_samples=100, acq_type=\"ucb\")\n",
+    "    elif turbo_mode == TurboMode.TURBO_ONE:\n",
+    "        config = TurboOneConfig(**common_cfg)\n",
+    "    elif turbo_mode == TurboMode.TURBO_ZERO:\n",
+    "        config = TurboZeroConfig(**common_cfg)\n",
+    "    elif turbo_mode == TurboMode.LHD_ONLY:\n",
+    "        config = LHDOnlyConfig(**common_cfg)\n",
+    "    else:\n",
+    "        raise ValueError(turbo_mode)\n",
+    "\n",
+    "    optimizer = Turbo(bounds=bounds, mode=turbo_mode, rng=rng, config=config)\n",
+    "\n",
+    "    all_y = []\n",
+    "    hypervolume_history = []\n",
+    "    proposal_times = []\n",
+    "\n",
+    "    for iteration in range(num_iterations):\n",
+    "        t0 = time.time()\n",
+    "        x_arms = optimizer.ask(num_arms=num_arms)\n",
+    "        proposal_times.append(time.time() - t0)\n",
+    "\n",
+    "        y_obs = objective(x_arms)\n",
+    "        optimizer.tell(x_arms, y_obs)\n",
+    "\n",
+    "        all_y.append(y_obs)\n",
+    "        hv = compute_hypervolume(np.vstack(all_y), ref_point)\n",
+    "        hypervolume_history.append(hv)\n",
+    "\n",
+    "        if iteration % 10 == 0:\n",
+    "            print(f\"{turbo_mode.name} iter {iteration:3d}: HV = {hv:.2f}\")\n",
+    "\n",
+    "    evals = num_arms * np.arange(1, len(hypervolume_history) + 1)\n",
+    "    return hypervolume_history, proposal_times, evals"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "RUN_TURBO_ONE = True\n",
+    "\n",
+    "modes = [TurboMode.TURBO_ENN, TurboMode.LHD_ONLY, TurboMode.TURBO_ZERO]\n",
+    "if RUN_TURBO_ONE:\n",
+    "    modes.append(TurboMode.TURBO_ONE)\n",
+    "\n",
+    "results = {}\n",
+    "for mode in modes:\n",
+    "    hv, times, evals = run_optimization(mode)\n",
+    "    results[mode] = {\"hv\": hv, \"times\": times, \"evals\": evals}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Hypervolume vs Function Evaluations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "plt.figure(figsize=(10, 6))\n",
+    "for mode, data in results.items():\n",
+    "    plt.plot(data[\"evals\"], data[\"hv\"], linewidth=2, label=mode.name)\n",
+    "plt.xlabel(\"Function Evaluations\")\n",
+    "plt.ylabel(\"Hypervolume\")\n",
+    "plt.legend()\n",
+    "plt.grid(True, alpha=0.3)\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Proposal Time vs Function Evaluations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.figure(figsize=(10, 6))\n",
+    "for mode, data in results.items():\n",
+    "    plt.plot(data[\"evals\"], data[\"times\"], linewidth=2, label=mode.name)\n",
+    "plt.xlabel(\"Function Evaluations\")\n",
+    "plt.ylabel(\"Proposal Time (seconds)\")\n",
+    "plt.legend()\n",
+    "plt.grid(True, alpha=0.3)\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python",
+   "pygments_lexer": "ipython3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

ennbo-0.1.2/examples/demo_turbo_enn.ipynb

@@ -0,0 +1,246 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# TuRBO-ENN\n",
+    "\n",
+    "This code implements TuRBO [1], a SOTA Bayesian optimization algorithm.\n",
+    "\n",
+    "The optimization class, `Turbo`, supports four modes of operation.\n",
+    "\n",
+    "**LHD_ONLY**  \n",
+    "Generate a Latin Hypercube Design (LHD) for every batch of arms. This is included as a simple baseline.\n",
+    "\n",
+    "**TURBO_ZERO**  \n",
+    "Initialze with LHD. Afterward, sample near the best-so-far x value, x_best. Samples are \"near\" x_best in two senses: (i) They are in a trust region, an adaptively-sized box around x_best, and (ii) They perturb only a small number of dimensions using RAASP sampling [2]. Other dimensions take the same value as in x_best. The num_arms proposals are chosen randomly from RAASP candidates inside the trust region.\n",
+    "\n",
+    "This is included to help differentiate the impact of the trust region from the impact of the surrogate. Notice (below) that the trust region has high impact.\n",
+    "\n",
+    "**TURBO_ONE**  \n",
+    "This adds a GP surrogate to TURBO_ZERO. The num_arms proposals are chosen via Thompson sampling from RAASP candidates inside the trust region. Occasionally, the trust region adapter resets and (i) discards all observations, and (ii) begins anew with and LHD design.\n",
+    "\n",
+    "This is the standard SOTA method. It should match the TuRBO reference [implementation](https://github.com/uber-research/TuRBO). \n",
+    "\n",
+    "**TURBO_ENN**  \n",
+    "This replaces the GP surrogate with a simpler, more scalable surrogate called Epistemic Nearest Neighbors (ENN). ENN's proposal time scales as $O(N)$ rather than the $O(N^2)$ of a GP surrogate. [3]\n",
+    "\n",
+    "\n",
+    "## References\n",
+    "\n",
+    "1. **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*.  \n",
+    "   https://arxiv.org/abs/1910.01739\n",
+    "\n",
+    "2. **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.  \n",
+    "   https://proceedings.mlr.press/v238/rashidi24a.html\n",
+    "\n",
+    "3. **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"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "---"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "\n",
+    "class Ackley:\n",
+    "    def __init__(self, noise: float, rng):\n",
+    "        self.noise = noise\n",
+    "        self.rng = rng\n",
+    "        self.a = 20.0\n",
+    "        self.b = 0.2\n",
+    "        self.c = 2 * np.pi\n",
+    "        self.bounds = [-32.768, 32.768]\n",
+    "\n",
+    "    def __call__(self, x):\n",
+    "        x = np.asarray(x, dtype=float)\n",
+    "        if x.ndim == 1:\n",
+    "            x = x[None, :]\n",
+    "        x = x - 1\n",
+    "        y = (\n",
+    "            -self.a * np.exp(-self.b * np.sqrt((x**2).mean(axis=1)))\n",
+    "            - np.exp(np.cos(self.c * x).mean(axis=1))\n",
+    "            + self.a\n",
+    "            + np.e\n",
+    "        )\n",
+    "        result = -y + self.noise * self.rng.normal(size=(x.shape[0],))\n",
+    "        return result if result.ndim > 0 else float(result)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import time\n",
+    "import torch\n",
+    "\n",
+    "from enn import Turbo, TurboMode\n",
+    "from enn.turbo.turbo_config import (\n",
+    "    LHDOnlyConfig,\n",
+    "    TurboENNConfig,\n",
+    "    TurboOneConfig,\n",
+    "    TurboZeroConfig,\n",
+    ")\n",
+    "\n",
+    "\n",
+    "def run_optimization(turbo_mode: TurboMode):\n",
+    "    num_dim = 30\n",
+    "    num_iterations = 100\n",
+    "    num_arms = 10\n",
+    "    noise = 0.1\n",
+    "    use_y_hat = False\n",
+    "\n",
+    "    rng = np.random.default_rng(17)\n",
+    "    torch.manual_seed(19)\n",
+    "    objective = Ackley(noise=noise, rng=rng)\n",
+    "    bounds = np.array([objective.bounds] * num_dim, dtype=float)\n",
+    "\n",
+    "    common_cfg = dict(k=10)\n",
+    "    if turbo_mode == TurboMode.TURBO_ENN:\n",
+    "        config = TurboENNConfig(**common_cfg, num_fit_samples=100, acq_type=\"ucb\")\n",
+    "    elif turbo_mode == TurboMode.TURBO_ONE:\n",
+    "        config = TurboOneConfig(**common_cfg)\n",
+    "    elif turbo_mode == TurboMode.TURBO_ZERO:\n",
+    "        config = TurboZeroConfig(**common_cfg)\n",
+    "    elif turbo_mode == TurboMode.LHD_ONLY:\n",
+    "        config = LHDOnlyConfig(**common_cfg)\n",
+    "    else:\n",
+    "        raise ValueError(turbo_mode)\n",
+    "\n",
+    "    optimizer = Turbo(\n",
+    "        bounds=bounds,\n",
+    "        mode=turbo_mode,\n",
+    "        rng=rng,\n",
+    "        config=config,\n",
+    "    )\n",
+    "\n",
+    "    best_values = []\n",
+    "    proposal_times = []\n",
+    "    y_hat_best = -np.inf\n",
+    "    y_best = -np.inf\n",
+    "\n",
+    "    for iteration in range(num_iterations):\n",
+    "        t_0 = time.time()\n",
+    "        x_arms = optimizer.ask(num_arms=num_arms)\n",
+    "        t_1 = time.time()\n",
+    "        proposal_times.append(t_1 - t_0)\n",
+    "\n",
+    "        y_obs = objective(x_arms)\n",
+    "\n",
+    "        y_hat = optimizer.tell(x_arms, y_obs, y_var=noise**2 * np.ones_like(y_obs))\n",
+    "        if use_y_hat:\n",
+    "            y_decide = y_hat\n",
+    "        else:\n",
+    "            y_decide = y_obs\n",
+    "        i = np.where(y_decide == np.max(y_decide))[0][0]\n",
+    "        y_hat_batch_best = y_decide[i]\n",
+    "\n",
+    "        if y_hat_batch_best > y_hat_best:\n",
+    "            y_hat_best = y_hat_batch_best\n",
+    "            y_best = objective(x_arms[i])[0]\n",
+    "\n",
+    "        best_values.append(y_best)\n",
+    "        if iteration % 10 == 0:\n",
+    "            print(\n",
+    "                f\"{turbo_mode} iteration = {iteration} y_best = {y_best} y_hat_best = {y_hat_best}\"\n",
+    "            )\n",
+    "\n",
+    "    evals = num_arms * np.arange(len(best_values))\n",
+    "    return best_values, proposal_times, evals"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "\n",
+    "RUN_TURBO_ONE = True\n",
+    "\n",
+    "best_values_zero, proposal_times_zero, evals_zero = run_optimization(\n",
+    "    TurboMode.TURBO_ZERO\n",
+    ")\n",
+    "if RUN_TURBO_ONE:\n",
+    "    best_values_one, proposal_times_one, evals_one = run_optimization(\n",
+    "        TurboMode.TURBO_ONE\n",
+    "    )\n",
+    "best_values_enn, proposal_times_enn, evals_enn = run_optimization(TurboMode.TURBO_ENN)\n",
+    "best_values_lhd, proposal_times_lhd, evals_lhd = run_optimization(TurboMode.LHD_ONLY)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.figure(figsize=(10, 6))\n",
+    "plt.plot(evals_zero, best_values_zero, linewidth=2, label=\"TURBO_ZERO\")\n",
+    "plt.plot(evals_enn, best_values_enn, linewidth=2, label=\"TURBO_ENN\")\n",
+    "plt.plot(evals_lhd, best_values_lhd, linewidth=2, label=\"LHD_ONLY\")\n",
+    "if RUN_TURBO_ONE:\n",
+    "    plt.plot(evals_one, best_values_one, linewidth=2, label=\"TURBO_ONE\")\n",
+    "plt.xlabel(\"Function Evaluations\")\n",
+    "plt.ylabel(\"Best Function Value\")\n",
+    "plt.legend()\n",
+    "plt.grid(True, alpha=0.3)\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.figure(figsize=(10, 6))\n",
+    "plt.plot(evals_zero, proposal_times_zero, linewidth=2, label=\"TURBO_ZERO\")\n",
+    "plt.plot(evals_enn, proposal_times_enn, linewidth=2, label=\"TURBO_ENN\")\n",
+    "plt.plot(evals_lhd, proposal_times_lhd, linewidth=2, label=\"LHD_ONLY\")\n",
+    "if RUN_TURBO_ONE:\n",
+    "    plt.plot(evals_one, proposal_times_one, linewidth=2, label=\"TURBO_ONE\")\n",
+    "plt.xlabel(\"Function Evaluations\")\n",
+    "plt.ylabel(\"Proposal Time (seconds)\")\n",
+    "# c = plt.axis()\n",
+    "# plt.axis([c[0], c[1], 0, 5])\n",
+    "plt.legend()\n",
+    "plt.grid(True, alpha=0.3)\n",
+    "plt.tight_layout()\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python",
+   "pygments_lexer": "ipython3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

{ennbo-0.1.1 → ennbo-0.1.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "ennbo"
-version = "0.1.1"
+version = "0.1.2"
 description = "Epistemic Nearest Neighbors"
 readme = "README.md"
 requires-python = ">=3.11"
@@ -13,7 +13,7 @@ authors = [
   { name = "YUBO Lab", email = "david.sweet@yu.edu" },
 ]
 dependencies = [
-    "numpy==1.26.4",
+    "numpy>=1.26.4,<2.0.0",
     "torch==2.5.1",
     "gpytorch==1.13",
     "faiss-cpu>=1.9.0",