ennbo 0.1.2__py3-none-any.whl → 0.1.7__py3-none-any.whl

enn/turbo/turbo_utils.py

@@ -1,20 +1,35 @@
 from __future__ import annotations
-
 import contextlib
-from typing import TYPE_CHECKING, Any, Iterator
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Iterator
+import numpy as np
 
 if TYPE_CHECKING:
-    import numpy as np
     import torch
-    from gpytorch.likelihoods import GaussianLikelihood
     from numpy.random import Generator
     from scipy.stats._qmc import QMCEngine
+__all__ = [
+    "Telemetry",
+]
+
 
-    from .turbo_gp import TurboGP
-    from .turbo_gp_noisy import TurboGPNoisy
+@dataclass(frozen=True)
+class Telemetry:
+    dt_fit: float
+    dt_sel: float
+    dt_gen: float = 0.0
+    dt_tell: float = 0.0
 
 
-from enn.enn.enn_util import standardize_y
+@contextlib.contextmanager
+def record_duration(set_dt: Callable[[float], None]) -> Iterator[None]:
+    import time
+
+    t0 = time.perf_counter()
+    try:
+        yield
+    finally:
+        set_dt(time.perf_counter() - t0)
 
 
 def _next_power_of_2(n: int) -> int:
@@ -43,138 +58,27 @@ def torch_seed_context(
         yield
 
 
-def fit_gp(
-    x_obs_list: list[float] | list[list[float]],
-    y_obs_list: list[float] | list[list[float]],
-    num_dim: int,
-    *,
-    yvar_obs_list: list[float] | None = None,
-    num_steps: int = 50,
-) -> tuple[
-    "TurboGP | TurboGPNoisy | None",
-    "GaussianLikelihood | None",
-    float | np.ndarray,
-    float | np.ndarray,
-]:
-    import numpy as np
-    import torch
-    from gpytorch.constraints import Interval
-    from gpytorch.likelihoods import GaussianLikelihood
-    from gpytorch.mlls import ExactMarginalLogLikelihood
-
-    from .turbo_gp import TurboGP
-    from .turbo_gp_noisy import TurboGPNoisy
-
-    x = np.asarray(x_obs_list, dtype=float)
-    y = np.asarray(y_obs_list, dtype=float)
-    n = x.shape[0]
-    if y.ndim not in (1, 2):
-        raise ValueError(y.shape)
-    is_multi_output = y.ndim == 2 and y.shape[1] > 1
-    if yvar_obs_list is not None:
-        if len(yvar_obs_list) != len(y_obs_list):
-            raise ValueError(
-                f"yvar_obs_list length {len(yvar_obs_list)} != y_obs_list length {len(y_obs_list)}"
-            )
-        if is_multi_output:
-            raise ValueError("yvar_obs_list not supported for multi-output GP")
-    if n == 0:
-        if is_multi_output:
-            num_outputs = int(y.shape[1])
-            return None, None, np.zeros(num_outputs), np.ones(num_outputs)
-        return None, None, 0.0, 1.0
-    if n == 1 and is_multi_output:
-        gp_y_mean = y[0].copy()
-        gp_y_std = np.ones(int(y.shape[1]), dtype=float)
-        return None, None, gp_y_mean, gp_y_std
-
-    if is_multi_output:
-        gp_y_mean = y.mean(axis=0)
-        gp_y_std = y.std(axis=0)
-        gp_y_std = np.where(gp_y_std < 1e-6, 1.0, gp_y_std)
-        z = (y - gp_y_mean) / gp_y_std
-    else:
-        gp_y_mean, gp_y_std = standardize_y(y)
-        y_centered = y - gp_y_mean
-        z = y_centered / gp_y_std
-    train_x = torch.as_tensor(x, dtype=torch.float64)
-    if is_multi_output:
-        train_y = torch.as_tensor(z.T, dtype=torch.float64)
-    else:
-        train_y = torch.as_tensor(z, dtype=torch.float64)
-    lengthscale_constraint = Interval(0.005, 2.0)
-    outputscale_constraint = Interval(0.05, 20.0)
-    if yvar_obs_list is not None:
-        y_var = np.asarray(yvar_obs_list, dtype=float)
-        train_y_var = torch.as_tensor(y_var / (gp_y_std**2), dtype=torch.float64)
-        model = TurboGPNoisy(
-            train_x=train_x,
-            train_y=train_y,
-            train_y_var=train_y_var,
-            lengthscale_constraint=lengthscale_constraint,
-            outputscale_constraint=outputscale_constraint,
-            ard_dims=num_dim,
-        ).to(dtype=train_x.dtype)
-        likelihood = model.likelihood
-    else:
-        noise_constraint = Interval(5e-4, 0.2)
-        if is_multi_output:
-            num_outputs = int(y.shape[1])
-            likelihood = GaussianLikelihood(
-                noise_constraint=noise_constraint,
-                batch_shape=torch.Size([num_outputs]),
-            ).to(dtype=train_y.dtype)
-        else:
-            likelihood = GaussianLikelihood(noise_constraint=noise_constraint).to(
-                dtype=train_y.dtype
-            )
-        model = TurboGP(
-            train_x=train_x,
-            train_y=train_y,
-            likelihood=likelihood,
-            lengthscale_constraint=lengthscale_constraint,
-            outputscale_constraint=outputscale_constraint,
-            ard_dims=num_dim,
-        ).to(dtype=train_x.dtype)
-        if is_multi_output:
-            likelihood.noise = torch.full(
-                (int(y.shape[1]),), 0.005, dtype=train_y.dtype
-            )
-        else:
-            likelihood.noise = torch.tensor(0.005, dtype=train_y.dtype)
-    if is_multi_output:
-        num_outputs = int(y.shape[1])
-        model.covar_module.outputscale = torch.ones(num_outputs, dtype=train_x.dtype)
-        model.covar_module.base_kernel.lengthscale = torch.full(
-            (num_outputs, 1, num_dim), 0.5, dtype=train_x.dtype
-        )
-    else:
-        model.covar_module.outputscale = torch.tensor(1.0, dtype=train_x.dtype)
-        model.covar_module.base_kernel.lengthscale = torch.full(
-            (num_dim,), 0.5, dtype=train_x.dtype
+def get_gp_posterior_suppress_warning(model: Any, x_torch: Any) -> Any:
+    import warnings
+
+    try:
+        from gpytorch.utils.warnings import GPInputWarning
+    except Exception:
+        GPInputWarning = None
+    if GPInputWarning is None:
+        return model.posterior(x_torch)
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore",
+            message=r"The input matches the stored training data\..*",
+            category=GPInputWarning,
         )
-    model.train()
-    likelihood.train()
-    mll = ExactMarginalLogLikelihood(likelihood, model)
-    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
-    for _ in range(num_steps):
-        optimizer.zero_grad()
-        output = model(train_x)
-        loss = -mll(output, train_y)
-        if loss.ndim != 0:
-            loss = loss.sum()
-        loss.backward()
-        optimizer.step()
-    model.eval()
-    likelihood.eval()
-    return model, likelihood, gp_y_mean, gp_y_std
+        return model.posterior(x_torch)
 
 
 def latin_hypercube(
     num_points: int, num_dim: int, *, rng: Generator | Any
 ) -> np.ndarray:
-    import numpy as np
-
     x = np.zeros((num_points, num_dim))
     centers = (1.0 + 2.0 * np.arange(0.0, num_points)) / float(2 * num_points)
     for j in range(num_dim):
@@ -185,8 +89,6 @@ def latin_hypercube(
 
 
 def argmax_random_tie(values: np.ndarray | Any, *, rng: Generator | Any) -> int:
-    import numpy as np
-
     if values.ndim != 1:
         raise ValueError(values.shape)
     max_val = float(np.max(values))
@@ -208,8 +110,6 @@ def sobol_perturb_np(
     *,
     sobol_engine: QMCEngine | Any,
 ) -> np.ndarray:
-    import numpy as np
-
     n_sobol = _next_power_of_2(num_candidates)
     sobol_samples = sobol_engine.random(n_sobol)[:num_candidates]
     lb_array = np.asarray(lb)
@@ -221,29 +121,76 @@ def sobol_perturb_np(
     return candidates
 
 
-def raasp(
+def uniform_perturb_np(
     x_center: np.ndarray | Any,
     lb: np.ndarray | list[float] | Any,
     ub: np.ndarray | list[float] | Any,
     num_candidates: int,
+    mask: np.ndarray | Any,
     *,
-    num_pert: int = 20,
     rng: Generator | Any,
-    sobol_engine: QMCEngine | Any,
 ) -> np.ndarray:
-    import numpy as np
+    lb_array = np.asarray(lb)
+    ub_array = np.asarray(ub)
+    pert = lb_array + (ub_array - lb_array) * rng.uniform(
+        0.0, 1.0, size=(num_candidates, x_center.shape[-1])
+    )
+    candidates = np.tile(x_center, (num_candidates, 1))
+    if np.any(mask):
+        candidates[mask] = pert[mask]
+    return candidates
 
-    num_dim = x_center.shape[-1]
+
+def _raasp_mask(
+    *,
+    num_candidates: int,
+    num_dim: int,
+    num_pert: int,
+    rng: Generator | Any,
+) -> np.ndarray:
     prob_perturb = min(num_pert / num_dim, 1.0)
     mask = rng.random((num_candidates, num_dim)) <= prob_perturb
     ind = np.nonzero(~mask.any(axis=1))[0]
     if len(ind) > 0:
         mask[ind, rng.integers(0, num_dim, size=len(ind))] = True
+    return mask
+
+
+def raasp(
+    x_center: np.ndarray | Any,
+    lb: np.ndarray | list[float] | Any,
+    ub: np.ndarray | list[float] | Any,
+    num_candidates: int,
+    *,
+    num_pert: int = 20,
+    rng: Generator | Any,
+    sobol_engine: QMCEngine | Any,
+) -> np.ndarray:
+    num_dim = x_center.shape[-1]
+    mask = _raasp_mask(
+        num_candidates=num_candidates, num_dim=num_dim, num_pert=num_pert, rng=rng
+    )
     return sobol_perturb_np(
         x_center, lb, ub, num_candidates, mask, sobol_engine=sobol_engine
     )
 
 
+def raasp_uniform(
+    x_center: np.ndarray | Any,
+    lb: np.ndarray | list[float] | Any,
+    ub: np.ndarray | list[float] | Any,
+    num_candidates: int,
+    *,
+    num_pert: int = 20,
+    rng: Generator | Any,
+) -> np.ndarray:
+    num_dim = x_center.shape[-1]
+    mask = _raasp_mask(
+        num_candidates=num_candidates, num_dim=num_dim, num_pert=num_pert, rng=rng
+    )
+    return uniform_perturb_np(x_center, lb, ub, num_candidates, mask, rng=rng)
+
+
 def generate_raasp_candidates(
     center: np.ndarray | Any,
     lb: np.ndarray | list[float] | Any,
@@ -267,29 +214,21 @@ def generate_raasp_candidates(
     )
 
 
-def generate_trust_region_candidates(
-    x_center: np.ndarray | Any,
-    lengthscales: np.ndarray | None,
+def generate_raasp_candidates_uniform(
+    center: np.ndarray | Any,
+    lb: np.ndarray | list[float] | Any,
+    ub: np.ndarray | list[float] | Any,
     num_candidates: int,
     *,
-    compute_bounds_1d: Any,
     rng: Generator | Any,
-    sobol_engine: QMCEngine | Any,
+    num_pert: int = 20,
 ) -> np.ndarray:
-    """
-    Small DRY helper for trust-region candidate generation.
-
-    `compute_bounds_1d` is typically a TR object's bound computation method.
-    """
-    lb, ub = compute_bounds_1d(x_center, lengthscales)
-    return generate_raasp_candidates(
-        x_center, lb, ub, num_candidates, rng=rng, sobol_engine=sobol_engine
-    )
+    if num_candidates <= 0:
+        raise ValueError(num_candidates)
+    return raasp_uniform(center, lb, ub, num_candidates, num_pert=num_pert, rng=rng)
 
 
 def to_unit(x: np.ndarray | Any, bounds: np.ndarray | Any) -> np.ndarray:
-    import numpy as np
-
     lb = bounds[:, 0]
     ub = bounds[:, 1]
     if np.any(ub <= lb):
@@ -298,8 +237,6 @@ def to_unit(x: np.ndarray | Any, bounds: np.ndarray | Any) -> np.ndarray:
 
 
 def from_unit(x_unit: np.ndarray | Any, bounds: np.ndarray | Any) -> np.ndarray:
-    import numpy as np
-
     lb = np.asarray(bounds[:, 0])
     ub = np.asarray(bounds[:, 1])
     return lb + x_unit * (ub - lb)
@@ -310,11 +247,11 @@ def gp_thompson_sample(
     x_cand: np.ndarray | Any,
     num_arms: int,
     rng: Generator | Any,
+    *,
     gp_y_mean: float,
     gp_y_std: float,
 ) -> np.ndarray:
     import gpytorch
-    import numpy as np
     import torch
 
     x_torch = torch.as_tensor(x_cand, dtype=torch.float64)

enn/turbo/types/__init__.py

@@ -0,0 +1,7 @@
+from .gp_data_prep import GPDataPrep
+from .gp_fit_result import GPFitResult
+from .obs_lists import ObsLists
+from .prepare_ask_result import PrepareAskResult
+from .tell_inputs import TellInputs
+
+__all__ = ["GPDataPrep", "GPFitResult", "ObsLists", "PrepareAskResult", "TellInputs"]

enn/turbo/types/appendable_array.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import numpy as np
+
+
+class AppendableArray:
+    def __init__(self, initial_capacity: int = 100) -> None:
+        self._initial_capacity = initial_capacity
+        self._buffer: np.ndarray | None = None
+        self._size = 0
+        self._num_cols: int | None = None
+
+    @property
+    def shape(self) -> tuple[int, int]:
+        if self._num_cols is None:
+            return (0, 0)
+        return (self._size, self._num_cols)
+
+    def _initialize_buffer(self, row: np.ndarray) -> None:
+        if row.ndim == 0:
+            self._num_cols = 1
+            row = row.reshape(1, 1)
+        elif row.ndim == 1:
+            self._num_cols = row.shape[0]
+            row = row[np.newaxis, :]
+        elif row.ndim == 2:
+            if row.shape[0] != 1:
+                raise ValueError(f"Expected row shape (1, D), got {row.shape}")
+            self._num_cols = row.shape[1]
+        else:
+            raise ValueError(f"Expected 0D, 1D or 2D array, got {row.ndim}D")
+
+        self._buffer = np.empty(
+            (self._initial_capacity, self._num_cols), dtype=row.dtype
+        )
+
+    def _validate_row(self, row: np.ndarray) -> np.ndarray:
+        if row.ndim == 0:
+            if self._num_cols != 1:
+                raise ValueError(f"Expected {self._num_cols} columns, got 1 (scalar)")
+            return row.reshape(1, 1)
+        if row.ndim == 1:
+            if row.shape[0] != self._num_cols:
+                raise ValueError(
+                    f"Expected {self._num_cols} columns, got {row.shape[0]}"
+                )
+            return row[np.newaxis, :]
+        if row.ndim == 2:
+            if row.shape != (1, self._num_cols):
+                raise ValueError(
+                    f"Expected shape (1, {self._num_cols}), got {row.shape}"
+                )
+            return row
+        raise ValueError(f"Expected 0D, 1D or 2D array, got {row.ndim}D")
+
+    def append(self, row: np.ndarray) -> None:
+        row = np.asarray(row)
+
+        if self._num_cols is None:
+            self._initialize_buffer(row)
+
+        row = self._validate_row(row)
+
+        assert self._buffer is not None
+        if self._size + 1 > self._buffer.shape[0]:
+            new_capacity = self._buffer.shape[0] * 2
+            new_buffer = np.empty(
+                (new_capacity, self._num_cols), dtype=self._buffer.dtype
+            )
+            new_buffer[: self._size] = self._buffer[: self._size]
+            self._buffer = new_buffer
+
+        self._buffer[self._size] = row
+        self._size += 1
+
+    def view(self) -> np.ndarray:
+        if self._buffer is None:
+            return np.empty((0, 0))
+        return self._buffer[: self._size]
+
+    def __len__(self) -> int:
+        return self._size
+
+    def __getitem__(self, key) -> np.ndarray:
+        return self.view()[key]

enn/turbo/types/gp_data_prep.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class GPDataPrep:
+    train_x: Any
+    train_y: Any
+    is_multi: bool
+    y_mean: Any
+    y_std: Any
+    y_raw: Any

enn/turbo/types/gp_fit_result.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class GPFitResult:
+    model: Any
+    likelihood: Any
+    y_mean: Any
+    y_std: Any

enn/turbo/types/obs_lists.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+from dataclasses import dataclass
+
+
+@dataclass
+class ObsLists:
+    x_obs: list
+    y_obs: list
+    y_tr: list
+    yvar_obs: list

enn/turbo/types/prepare_ask_result.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import numpy as np
+
+
+@dataclass
+class PrepareAskResult:
+    model: Any
+    y_mean: float | None
+    y_std: float | None
+    lengthscales: np.ndarray | None

enn/turbo/types/tell_inputs.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import numpy as np
+
+
+@dataclass
+class TellInputs:
+    x: np.ndarray
+    y: np.ndarray
+    y_var: np.ndarray | None
+    num_metrics: int

{ennbo-0.1.2.dist-info → ennbo-0.1.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ennbo
-Version: 0.1.2
+Version: 0.1.7
 Summary: Epistemic Nearest Neighbors
 Project-URL: Homepage, https://github.com/yubo-research/enn
 Project-URL: Source, https://github.com/yubo-research/enn
@@ -37,11 +37,14 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Scientific/Engineering :: Mathematics
 Requires-Python: >=3.11
 Requires-Dist: faiss-cpu>=1.9.0
-Requires-Dist: gpytorch==1.13
-Requires-Dist: nds==0.4.3
+Requires-Dist: gpytorch>=1.13
 Requires-Dist: numpy<2.0.0,>=1.26.4
-Requires-Dist: scipy==1.15.3
-Requires-Dist: torch==2.5.1
+Requires-Dist: scipy>=1.15.3
+Requires-Dist: torch>=2.5.1
+Provides-Extra: dev
+Requires-Dist: nbmake>=1.5; extra == 'dev'
+Requires-Dist: pytest>=8.0; extra == 'dev'
+Requires-Dist: ruff>=0.4; extra == 'dev'
 Description-Content-Type: text/markdown
 
 # Epistemic Nearest Neighbors
@@ -50,12 +53,16 @@ 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.
+- ENN surrogate, [`EpistemicNearestNeighbors`](https://github.com/yubo-research/enn/blob/main/src/enn/enn/enn.py) [1]
+- TuRBO-ENN optimizer via [`create_optimizer`](https://github.com/yubo-research/enn/blob/main/src/enn/turbo/optimizer.py) with config factories
+	- `turbo_one_config()` - TuRBO [2], matching the reference implementation.
+	- `turbo_enn_config()` - Uses ENN instead of GP.
+	- `turbo_zero_config()` - No surrogate
+	- `lhd_only_config()` - LHD design on every `ask()`. Good for a baseline and for testing.
+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).
+
 
 [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