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

enn/turbo/proposal.py

@@ -1,43 +1,41 @@
 from __future__ import annotations
-
 from typing import TYPE_CHECKING, Any, Callable
 
 if TYPE_CHECKING:
     import numpy as np
     from numpy.random import Generator
-
-    from enn.enn import EpistemicNearestNeighbors
+    from enn.enn.enn_class import EpistemicNearestNeighbors
     from enn.enn.enn_params import ENNParams
-
     from .turbo_gp import TurboGP
-
-from .turbo_utils import gp_thompson_sample
+    from .config.enums import ENNIndexDriver
 
 
 def mk_enn(
     x_obs_list: list[float] | list[list[float]],
     y_obs_list: list[float] | list[list[float]],
-    *,
-    yvar_obs_list: list[float] | None = None,
     k: int,
+    yvar_obs_list: list[float] | None = None,
+    *,
     num_fit_samples: int | None = None,
     num_fit_candidates: int | None = None,
     scale_x: bool = False,
+    index_driver: ENNIndexDriver | Any | None = None,
     rng: Generator | Any | None = None,
     params_warm_start: ENNParams | Any | None = None,
 ) -> tuple[EpistemicNearestNeighbors | None, ENNParams | None]:
     import numpy as np
-
-    from enn.enn import EpistemicNearestNeighbors
+    from enn.enn.enn_class import EpistemicNearestNeighbors
     from enn.enn.enn_params import ENNParams
+    from .config.enums import ENNIndexDriver
+
+    if index_driver is None:
+        index_driver = ENNIndexDriver.FLAT
 
     if len(x_obs_list) == 0:
         return None, None
     y_obs_array = np.asarray(y_obs_list, dtype=float)
     if y_obs_array.size == 0:
         return None, None
-
-    # Preserve multi-metric shape if present, otherwise reshape to (n, 1)
     if y_obs_array.ndim == 1:
         y = y_obs_array.reshape(-1, 1)
     else:
@@ -56,10 +54,10 @@ def mk_enn(
         y,
         yvar,
         scale_x=scale_x,
+        index_driver=index_driver,
     )
     if len(enn_model) == 0:
         return None, None
-
     fitted_params: ENNParams | None = None
     if num_fit_samples is not None and rng is not None:
         from enn.enn.enn_fit import enn_fit
@@ -75,8 +73,11 @@ def mk_enn(
             params_warm_start=params_warm_start,
         )
     else:
-        fitted_params = ENNParams(k=k, epi_var_scale=1.0, ale_homoscedastic_scale=0.0)
-
+        fitted_params = ENNParams(
+            k_num_neighbors=k,
+            epistemic_variance_scale=1.0,
+            aleatoric_variance_scale=0.0,
+        )
     return enn_model, fitted_params
 
 
@@ -98,45 +99,42 @@ def select_uniform(
 def select_gp_thompson(
     x_cand: np.ndarray,
     num_arms: int,
-    x_obs_list: list[float] | list[list[float]],
-    y_obs_list: list[float] | list[list[float]],
+    x_obs_list: list,
+    y_obs_list: list,
     num_dim: int,
+    *,
     gp_num_steps: int,
     rng: Generator | Any,
-    gp_y_mean: float,
-    gp_y_std: float,
+    gp_y_stats: tuple[float, float],
     select_sobol_fn: Callable[[np.ndarray, int], np.ndarray],
     from_unit_fn: Callable[[np.ndarray], np.ndarray],
-    *,
     model: TurboGP | None = None,
-    new_gp_y_mean: float | None = None,
-    new_gp_y_std: float | None = None,
-) -> tuple[np.ndarray, float, float, TurboGP | None]:
-    from .turbo_utils import fit_gp
+) -> tuple[np.ndarray, tuple[float, float], TurboGP | None]:
+    from .turbo_gp_fit import fit_gp
 
+    gp_y_mean, gp_y_std = gp_y_stats
     if len(x_obs_list) == 0:
-        return select_sobol_fn(x_cand, num_arms), gp_y_mean, gp_y_std, None
+        return select_sobol_fn(x_cand, num_arms), (gp_y_mean, gp_y_std), None
+    fitted_mean, fitted_std = gp_y_mean, gp_y_std
     if model is None:
-        model, _likelihood, new_gp_y_mean, new_gp_y_std = fit_gp(
+        gp_result = fit_gp(
             x_obs_list,
             y_obs_list,
             num_dim,
             num_steps=gp_num_steps,
         )
+        model, fitted_mean, fitted_std = (
+            gp_result.model,
+            gp_result.y_mean,
+            gp_result.y_std,
+        )
     if model is None:
-        return select_sobol_fn(x_cand, num_arms), gp_y_mean, gp_y_std, None
-    if new_gp_y_mean is None:
-        new_gp_y_mean = gp_y_mean
-    if new_gp_y_std is None:
-        new_gp_y_std = gp_y_std
+        return select_sobol_fn(x_cand, num_arms), (gp_y_mean, gp_y_std), None
     if x_cand.shape[0] < num_arms:
         raise ValueError((x_cand.shape[0], num_arms))
+    from .turbo_utils import gp_thompson_sample
+
     idx = gp_thompson_sample(
-        model,
-        x_cand,
-        num_arms,
-        rng,
-        new_gp_y_mean,
-        new_gp_y_std,
+        model, x_cand, num_arms, rng, gp_y_mean=fitted_mean, gp_y_std=fitted_std
     )
-    return from_unit_fn(x_cand[idx]), new_gp_y_mean, new_gp_y_std, model
+    return from_unit_fn(x_cand[idx]), (fitted_mean, fitted_std), model

enn/turbo/sampling.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+import numpy as np
+from . import turbo_utils
+
+if TYPE_CHECKING:
+    from numpy.random import Generator
+
+
+def draw_lhd(*, bounds: np.ndarray, num_arms: int, rng: Generator) -> np.ndarray:
+    bounds = np.asarray(bounds, dtype=float)
+    if bounds.ndim != 2 or bounds.shape[1] != 2:
+        raise ValueError(f"bounds must be (d, 2), got {bounds.shape}")
+    num_arms = int(num_arms)
+    if num_arms <= 0:
+        raise ValueError(num_arms)
+    num_dim = int(bounds.shape[0])
+    return turbo_utils.from_unit(
+        turbo_utils.latin_hypercube(num_arms, num_dim, rng=rng),
+        bounds,
+    )

enn/turbo/strategies/__init__.py

@@ -0,0 +1,9 @@
+from .lhd_only_strategy import LHDOnlyStrategy
+from .optimization_strategy import OptimizationStrategy
+from .turbo_hybrid_strategy import TurboHybridStrategy
+
+__all__ = [
+    "LHDOnlyStrategy",
+    "OptimizationStrategy",
+    "TurboHybridStrategy",
+]

enn/turbo/strategies/lhd_only_strategy.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+import numpy as np
+from ..sampling import draw_lhd
+from .optimization_strategy import OptimizationStrategy
+
+if TYPE_CHECKING:
+    from ..optimizer import Optimizer
+    from ..types import TellInputs
+
+
+@dataclass
+class LHDOnlyStrategy(OptimizationStrategy):
+    _bounds: np.ndarray
+    _rng: object
+
+    @classmethod
+    def create(cls, *, bounds: np.ndarray, rng: object) -> LHDOnlyStrategy:
+        bounds = np.asarray(bounds, dtype=float)
+        if bounds.ndim != 2 or bounds.shape[1] != 2:
+            raise ValueError(f"bounds must be (d, 2), got {bounds.shape}")
+        return cls(_bounds=bounds, _rng=rng)
+
+    def ask(self, opt: Optimizer, num_arms: int) -> np.ndarray:
+        return draw_lhd(bounds=self._bounds, num_arms=num_arms, rng=opt._rng)
+
+    def init_progress(self) -> tuple[int, int] | None:
+        return None
+
+    def tell(
+        self, opt: Optimizer, inputs: TellInputs, *, x_unit: np.ndarray
+    ) -> np.ndarray:
+        del x_unit
+        opt._y_tr_list = inputs.y.tolist()
+        return np.asarray(inputs.y, dtype=float)

enn/turbo/strategies/optimization_strategy.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+import numpy as np
+
+if TYPE_CHECKING:
+    from ..optimizer import Optimizer
+    from ..types import TellInputs
+
+
+class OptimizationStrategy(ABC):
+    @abstractmethod
+    def ask(self, opt: Optimizer, num_arms: int) -> np.ndarray: ...
+    @abstractmethod
+    def tell(
+        self, opt: Optimizer, inputs: TellInputs, *, x_unit: np.ndarray
+    ) -> np.ndarray: ...
+    @abstractmethod
+    def init_progress(self) -> tuple[int, int] | None: ...

enn/turbo/strategies/turbo_hybrid_strategy.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+import time
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+import numpy as np
+from ..sampling import draw_lhd
+from ..types.appendable_array import AppendableArray
+from .optimization_strategy import OptimizationStrategy
+
+if TYPE_CHECKING:
+    from ..optimizer import Optimizer
+    from ..types import TellInputs
+
+
+@dataclass
+class TurboHybridStrategy(OptimizationStrategy):
+    _bounds: np.ndarray
+    _num_dim: int
+    _rng: object
+    _num_init: int
+    _init_lhd: np.ndarray
+    _init_idx: int = 0
+
+    @classmethod
+    def create(
+        cls, *, bounds: np.ndarray, rng: object, num_init: int | None
+    ) -> TurboHybridStrategy:
+        from numpy.random import Generator
+
+        bounds = np.asarray(bounds, dtype=float)
+        if bounds.ndim != 2 or bounds.shape[1] != 2:
+            raise ValueError(f"bounds must be (d, 2), got {bounds.shape}")
+        num_dim = int(bounds.shape[0])
+        if not isinstance(rng, Generator):
+            raise TypeError("rng must be a numpy.random.Generator")
+        n_init = int(num_init if num_init is not None else 2 * num_dim)
+        if n_init <= 0:
+            raise ValueError(f"num_init must be > 0, got {n_init}")
+        init_lhd = draw_lhd(bounds=bounds, num_arms=n_init, rng=rng)
+        return cls(
+            _bounds=bounds,
+            _num_dim=num_dim,
+            _rng=rng,
+            _num_init=n_init,
+            _init_lhd=init_lhd,
+        )
+
+    def _reset_init(self) -> None:
+        from numpy.random import Generator
+
+        if not isinstance(self._rng, Generator):
+            raise TypeError("rng must be a numpy.random.Generator")
+        self._init_lhd = draw_lhd(
+            bounds=self._bounds, num_arms=self._num_init, rng=self._rng
+        )
+        self._init_idx = 0
+
+    def _get_init_points(self, num_arms: int, *, fallback_fn=None) -> np.ndarray:
+        num_arms = int(num_arms)
+        num_to_return = min(num_arms, self._num_init - self._init_idx)
+        result = self._init_lhd[self._init_idx : self._init_idx + num_to_return]
+        self._init_idx += num_to_return
+        if num_to_return < num_arms:
+            extra = (
+                fallback_fn
+                or (lambda n: draw_lhd(bounds=self._bounds, num_arms=n, rng=self._rng))
+            )(num_arms - num_to_return)
+            result = np.vstack([result, extra])
+        return result
+
+    def ask(self, opt: Optimizer, num_arms: int) -> np.ndarray:
+        if opt._tr_state.needs_restart():
+            opt._tr_state.restart(opt._rng)
+            opt._restart_generation += 1
+            opt._x_obs = AppendableArray()
+            opt._y_obs = AppendableArray()
+            opt._yvar_obs = AppendableArray()
+            opt._y_tr_list = []
+            opt._incumbent_idx = None
+            opt._incumbent_x_unit = None
+            opt._incumbent_y_scalar = None
+            self._reset_init()
+            return self._get_init_points(num_arms)
+        if self._init_idx < self._num_init:
+
+            def fallback(n: int) -> np.ndarray:
+                return opt._ask_normal(n, is_fallback=True)
+
+            return self._get_init_points(
+                num_arms,
+                fallback_fn=fallback if len(opt._x_obs) > 0 else None,
+            )
+        if len(opt._x_obs) == 0:
+            return draw_lhd(bounds=self._bounds, num_arms=num_arms, rng=self._rng)
+        opt._tr_state.validate_request(int(num_arms))
+        return opt._ask_normal(int(num_arms))
+
+    def init_progress(self) -> tuple[int, int] | None:
+        return (int(self._init_idx), int(self._num_init))
+
+    def tell(
+        self, opt: Optimizer, inputs: TellInputs, *, x_unit: np.ndarray
+    ) -> np.ndarray:
+        x_all = opt._x_obs.view()
+        y_all = opt._y_obs.view()
+        y_var_all = opt._yvar_obs.view() if len(opt._yvar_obs) > 0 else None
+        t0 = time.perf_counter()
+        opt._surrogate.fit(
+            x_all, y_all, y_var_all, num_steps=opt._gp_num_steps, rng=opt._rng
+        )
+        opt._dt_fit = time.perf_counter() - t0
+        opt._y_tr_list = y_all.tolist()
+        opt._update_incumbent()
+        try:
+            new_posterior = opt._surrogate.predict(x_unit)
+            y_estimate = np.asarray(new_posterior.mu, dtype=float)
+        except RuntimeError:
+            y_estimate = np.asarray(inputs.y, dtype=float)
+        y_incumbent = opt._incumbent_y_scalar
+        y_obs = y_all
+        opt._tr_state.update(y_obs, y_incumbent)
+        if opt._trailing_obs is not None:
+            opt._trim_trailing_obs()
+        return y_estimate

enn/turbo/tr_helpers.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from .config.enums import CandidateRV, RAASPDriver
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+    from scipy.stats._qmc import QMCEngine
+
+
+def compute_full_box_bounds_1d(
+    x_center: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
+    import numpy as np
+
+    lb = np.zeros_like(x_center, dtype=float)
+    ub = np.ones_like(x_center, dtype=float)
+    return lb, ub
+
+
+def get_single_incumbent_index(
+    selector,
+    y: np.ndarray,
+    rng: Generator,
+    mu: np.ndarray | None = None,
+) -> np.ndarray:
+    import numpy as np
+
+    y = np.asarray(y, dtype=float)
+    if y.size == 0:
+        return np.array([], dtype=int)
+    best_idx = selector.select(y, mu, rng)
+    return np.array([best_idx])
+
+
+def get_incumbent_index(
+    selector,
+    y: np.ndarray,
+    rng: Generator,
+    mu: np.ndarray | None = None,
+) -> int:
+    import numpy as np
+
+    y = np.asarray(y, dtype=float)
+    if y.size == 0:
+        raise ValueError("y is empty")
+    return int(selector.select(y, mu, rng))
+
+
+def get_scalar_incumbent_value(
+    selector,
+    y_obs: np.ndarray,
+    rng: Generator,
+    *,
+    mu_obs: np.ndarray | None = None,
+) -> np.ndarray:
+    import numpy as np
+
+    y = np.asarray(y_obs, dtype=float)
+    if y.size == 0:
+        return np.array([], dtype=float)
+    idx = get_incumbent_index(selector, y, rng, mu=mu_obs)
+    use_mu = bool(getattr(selector, "noise_aware", False))
+    values = mu_obs if use_mu else y
+    if values is None:
+        raise ValueError("noise_aware incumbent selection requires mu_obs")
+    v = np.asarray(values, dtype=float)
+    if v.ndim == 2:
+        value = float(v[idx, 0])
+    elif v.ndim == 1:
+        value = float(v[idx])
+    else:
+        raise ValueError(v.shape)
+    return np.array([value], dtype=float)
+
+
+class ScalarIncumbentMixin:
+    incumbent_selector: Any
+
+    def get_incumbent_index(
+        self,
+        y: np.ndarray | Any,
+        rng: Generator,
+        mu: np.ndarray | None = None,
+    ) -> int:
+        return get_incumbent_index(self.incumbent_selector, y, rng, mu=mu)
+
+    def get_incumbent_value(
+        self,
+        y_obs: np.ndarray | Any,
+        rng: Generator,
+        mu_obs: np.ndarray | None = None,
+    ) -> np.ndarray:
+        return get_scalar_incumbent_value(
+            self.incumbent_selector, y_obs, rng, mu_obs=mu_obs
+        )
+
+
+def generate_tr_candidates_orig(
+    compute_bounds_1d: Any,
+    x_center: np.ndarray,
+    lengthscales: np.ndarray | None,
+    num_candidates: int,
+    *,
+    rng: Generator,
+    candidate_rv: CandidateRV = CandidateRV.SOBOL,
+    sobol_engine: QMCEngine | None = None,
+) -> np.ndarray:
+    from .turbo_utils import (
+        generate_raasp_candidates,
+        generate_raasp_candidates_uniform,
+    )
+
+    lb, ub = compute_bounds_1d(x_center, lengthscales)
+    if candidate_rv == CandidateRV.SOBOL:
+        if sobol_engine is None:
+            raise ValueError(
+                "sobol_engine is required when candidate_rv=CandidateRV.SOBOL"
+            )
+        return generate_raasp_candidates(
+            x_center, lb, ub, num_candidates, rng=rng, sobol_engine=sobol_engine
+        )
+    if candidate_rv == CandidateRV.UNIFORM:
+        return generate_raasp_candidates_uniform(
+            x_center, lb, ub, num_candidates, rng=rng
+        )
+    raise ValueError(candidate_rv)
+
+
+def generate_tr_candidates_fast(
+    compute_bounds_1d: Any,
+    x_center: np.ndarray,
+    lengthscales: np.ndarray | None,
+    num_candidates: int,
+    *,
+    rng: Generator,
+    candidate_rv: CandidateRV,
+    num_pert: int,
+) -> np.ndarray:
+    import numpy as np
+    from scipy.stats import qmc
+
+    lb, ub = compute_bounds_1d(x_center, lengthscales)
+    num_dim = x_center.shape[-1]
+
+    candidates = np.empty((num_candidates, num_dim), dtype=float)
+    candidates[:] = x_center
+
+    prob_perturb = min(num_pert / num_dim, 1.0)
+    ks = rng.binomial(num_dim, prob_perturb, size=num_candidates)
+    ks = np.maximum(ks, 1)
+    max_k = int(np.max(ks))
+
+    if candidate_rv == CandidateRV.SOBOL:
+        sobol = qmc.Sobol(d=max_k, scramble=True, seed=int(rng.integers(0, 2**31)))
+        samples = sobol.random(num_candidates)
+    elif candidate_rv == CandidateRV.UNIFORM:
+        samples = rng.random((num_candidates, max_k))
+    else:
+        raise ValueError(candidate_rv)
+
+    for i in range(num_candidates):
+        k = ks[i]
+        idx = rng.choice(num_dim, size=k, replace=False)
+        candidates[i, idx] = lb[idx] + (ub[idx] - lb[idx]) * samples[i, :k]
+
+    return candidates
+
+
+def generate_tr_candidates(
+    compute_bounds_1d: Any,
+    x_center: np.ndarray,
+    lengthscales: np.ndarray | None,
+    num_candidates: int,
+    *,
+    rng: Generator,
+    candidate_rv: CandidateRV,
+    sobol_engine: QMCEngine | None,
+    raasp_driver: RAASPDriver,
+    num_pert: int,
+) -> np.ndarray:
+    if raasp_driver == RAASPDriver.FAST:
+        return generate_tr_candidates_fast(
+            compute_bounds_1d,
+            x_center,
+            lengthscales,
+            num_candidates,
+            rng=rng,
+            candidate_rv=candidate_rv,
+            num_pert=num_pert,
+        )
+    return generate_tr_candidates_orig(
+        compute_bounds_1d,
+        x_center,
+        lengthscales,
+        num_candidates,
+        rng=rng,
+        candidate_rv=candidate_rv,
+        sobol_engine=sobol_engine,
+    )

enn/turbo/turbo_gp.py

@@ -1,5 +1,4 @@
 from __future__ import annotations
-
 from .turbo_gp_base import TurboGPBase
 
 

enn/turbo/turbo_gp_base.py

@@ -1,5 +1,4 @@
 from __future__ import annotations
-
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING: