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

enn/turbo/no_trust_region.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+from .tr_helpers import ScalarIncumbentMixin
+
+if TYPE_CHECKING:
+    import numpy as np
+    from .components.incumbent_selector import IncumbentSelector
+    from .config.no_tr_config import NoTRConfig
+
+
+@dataclass
+class NoTrustRegion(ScalarIncumbentMixin):
+    config: NoTRConfig
+    num_dim: int
+    incumbent_selector: IncumbentSelector = field(default=None, repr=False)
+    length: float = field(default=1.0, init=False)
+
+    def __post_init__(self) -> None:
+        from .components.incumbent_selector import ScalarIncumbentSelector
+
+        if self.incumbent_selector is None:
+            self.incumbent_selector = ScalarIncumbentSelector(noise_aware=False)
+
+    @property
+    def num_metrics(self) -> int:
+        return 1
+
+    def update(self, y_obs: np.ndarray | Any, y_incumbent: np.ndarray | Any) -> None:
+        return
+
+    def needs_restart(self) -> bool:
+        return False
+
+    def restart(self, rng=None) -> None:
+        return
+
+    def validate_request(self, num_arms: int, *, is_fallback: bool = False) -> None:
+        pass
+
+    def compute_bounds_1d(
+        self,
+        x_center: np.ndarray | Any,
+        lengthscales: np.ndarray | None = None,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        from .tr_helpers import compute_full_box_bounds_1d
+
+        return compute_full_box_bounds_1d(x_center)
+
+    def get_incumbent_indices(
+        self,
+        y: np.ndarray | Any,
+        rng,
+        mu: np.ndarray | None = None,
+    ) -> np.ndarray:
+        import numpy as np
+
+        return np.array([self.get_incumbent_index(y, rng, mu=mu)])

enn/turbo/optimizer.py

@@ -0,0 +1,300 @@
+from __future__ import annotations
+
+import time
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from . import turbo_optimizer_utils, turbo_utils
+from .components import AcquisitionOptimizer, Surrogate
+from .config.enums import CandidateRV
+from .strategies import OptimizationStrategy
+from .types.appendable_array import AppendableArray
+
+if TYPE_CHECKING:
+    from numpy.random import Generator
+
+    from .config.optimizer_config import OptimizerConfig
+
+
+class Optimizer:
+    def __init__(
+        self,
+        *,
+        bounds: np.ndarray,
+        config: OptimizerConfig,
+        rng: Generator,
+        surrogate: Surrogate,
+        acquisition_optimizer: AcquisitionOptimizer,
+        strategy: OptimizationStrategy | None = None,
+    ) -> None:
+        self._config = config
+        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}")
+        self._bounds = bounds
+        self._num_dim = bounds.shape[0]
+        self._rng = rng
+        self._surrogate = surrogate
+        self._acq_optimizer = acquisition_optimizer
+        self._strategy = (
+            strategy
+            if strategy is not None
+            else config.init.init_strategy.create_runtime_strategy(
+                bounds=self._bounds, rng=self._rng, num_init=config.init.num_init
+            )
+        )
+        self._tr_state = config.trust_region.build(
+            num_dim=self._num_dim,
+            rng=rng,
+            candidate_rv=config.candidate_rv,
+        )
+        self._trailing_obs = (
+            None if config.trailing_obs is None else int(config.trailing_obs)
+        )
+        self._gp_num_steps = 50
+        if self._trailing_obs is not None and self._trailing_obs <= 0:
+            raise ValueError(f"trailing_obs must be > 0, got {self._trailing_obs}")
+        self._x_obs = AppendableArray()
+        self._y_obs = AppendableArray()
+        self._yvar_obs = AppendableArray()
+        self._y_tr_list: list[float] | list[list[float]] = []
+        self._expects_yvar: bool | None = None
+        self._dt_fit = 0.0
+        self._dt_gen = 0.0
+        self._dt_sel = 0.0
+        self._dt_tell = 0.0
+        self._sobol_seed_base = int(rng.integers(2**31 - 1))
+        self._restart_generation = 0
+        self._incumbent_idx: int | None = None
+        self._incumbent_x_unit: np.ndarray | None = None
+        self._incumbent_y_scalar: np.ndarray | None = None
+
+    @property
+    def tr_obs_count(self) -> int:
+        return len(self._y_obs)
+
+    @property
+    def tr_length(self) -> float:
+        return float(self._tr_state.length)
+
+    def telemetry(self) -> turbo_utils.Telemetry:
+        return turbo_utils.Telemetry(
+            dt_fit=self._dt_fit,
+            dt_gen=self._dt_gen,
+            dt_sel=self._dt_sel,
+            dt_tell=self._dt_tell,
+        )
+
+    @property
+    def init_progress(self) -> tuple[int, int] | None:
+        return self._strategy.init_progress()
+
+    def ask(self, num_arms: int) -> np.ndarray:
+        num_arms = int(num_arms)
+        if num_arms <= 0:
+            raise ValueError(num_arms)
+        turbo_optimizer_utils.reset_timing(self)
+        return self._strategy.ask(self, num_arms)
+
+    def _ask_normal(self, num_arms: int, *, is_fallback: bool = False) -> np.ndarray:
+        self._tr_state.validate_request(num_arms, is_fallback=is_fallback)
+        self._maybe_resample_weights()
+        x_center = self._incumbent_x_unit
+        if x_center is None:
+            if len(self._y_obs) == 0:
+                raise RuntimeError("no observations")
+            x_center = np.full(self._num_dim, 0.5)
+        t0 = time.perf_counter()
+        lengthscales = self._surrogate.lengthscales
+        x_cand = self._generate_candidates(x_center, lengthscales, num_arms=num_arms)
+        self._dt_gen = time.perf_counter() - t0
+        t0 = time.perf_counter()
+        selected = self._acq_optimizer.select(
+            x_cand,
+            num_arms,
+            self._surrogate,
+            self._rng,
+            tr_state=self._tr_state,
+        )
+        self._dt_sel = time.perf_counter() - t0
+        return turbo_utils.from_unit(selected, self._bounds)
+
+    def _find_x_center(self, x_obs: np.ndarray, y_obs: np.ndarray) -> np.ndarray | None:
+        return self._incumbent_x_unit
+
+    def _maybe_resample_weights(self) -> None:
+        from .config.rescalarize import Rescalarize
+
+        if hasattr(self._tr_state, "rescalarize"):
+            if self._tr_state.rescalarize == Rescalarize.ON_PROPOSE:
+                self._tr_state.resample_weights(self._rng)
+
+    def _generate_candidates(
+        self,
+        x_center: np.ndarray,
+        lengthscales: np.ndarray | None,
+        *,
+        num_arms: int,
+    ) -> np.ndarray:
+        from . import tr_helpers
+
+        if lengthscales is not None:
+            lengthscales = np.asarray(lengthscales, dtype=float).reshape(-1)
+            if not np.all(np.isfinite(lengthscales)):
+                raise ValueError("lengthscales must be finite")
+        num_candidates = int(
+            self._config.candidates.num_candidates(
+                num_dim=self._num_dim, num_arms=num_arms
+            )
+        )
+        if num_candidates <= 0:
+            raise ValueError(num_candidates)
+        candidate_rv = self._config.candidate_rv
+        if candidate_rv == CandidateRV.SOBOL:
+            from scipy.stats import qmc
+
+            sobol_seed = turbo_optimizer_utils.sobol_seed_for_state(
+                self._sobol_seed_base,
+                restart_generation=self._restart_generation,
+                n_obs=len(self._x_obs),
+                num_arms=num_arms,
+            )
+            sobol_engine = qmc.Sobol(d=self._num_dim, scramble=True, seed=sobol_seed)
+        else:
+            sobol_engine = None
+        if getattr(self._tr_state, "uses_custom_candidate_gen", False):
+            return self._tr_state.generate_candidates(
+                x_center,
+                lengthscales,
+                num_candidates,
+                rng=self._rng,
+                sobol_engine=sobol_engine,
+                raasp_driver=self._config.raasp_driver,
+                num_pert=20,
+            )
+        return tr_helpers.generate_tr_candidates(
+            self._tr_state.compute_bounds_1d,
+            x_center,
+            lengthscales,
+            num_candidates,
+            rng=self._rng,
+            candidate_rv=candidate_rv,
+            sobol_engine=sobol_engine,
+            raasp_driver=self._config.raasp_driver,
+            num_pert=20,
+        )
+
+    def _validate_tell_inputs(
+        self, x: np.ndarray, y: np.ndarray, y_var: np.ndarray | None
+    ) -> turbo_optimizer_utils.TellInputs:
+        inputs = turbo_optimizer_utils.validate_tell_inputs(x, y, y_var, self._num_dim)
+        tr_num_metrics = getattr(self._tr_state, "num_metrics", 1)
+        if inputs.num_metrics != tr_num_metrics:
+            raise ValueError(
+                f"y has {inputs.num_metrics} metrics but trust region expects {tr_num_metrics}"
+            )
+        if self._expects_yvar is None:
+            self._expects_yvar = inputs.y_var is not None
+        if (inputs.y_var is not None) != bool(self._expects_yvar):
+            raise ValueError(
+                f"y_var must be {'provided' if self._expects_yvar else 'omitted'} on every tell()"
+            )
+        return inputs
+
+    def _update_incumbent(self) -> None:
+        if len(self._y_obs) == 0:
+            self._incumbent_idx = None
+            self._incumbent_x_unit = None
+            self._incumbent_y_scalar = None
+            return
+        x_obs = self._x_obs.view()
+        y_obs = self._y_obs.view()
+        candidate_indices = self._surrogate.get_incumbent_candidate_indices(y_obs)
+        x_cand = x_obs[candidate_indices]
+        y_cand = y_obs[candidate_indices]
+        mu_cand = None
+        noise_aware = False
+        if hasattr(self._tr_state, "incumbent_selector"):
+            noise_aware = getattr(
+                self._tr_state.incumbent_selector, "noise_aware", False
+            )
+        elif hasattr(self._tr_state, "config"):
+            noise_aware = getattr(self._tr_state.config, "noise_aware", False)
+        if noise_aware:
+            try:
+                mu_cand = self._surrogate.predict(x_cand).mu
+            except RuntimeError:
+                mu_cand = None
+        idx_in_cand = self._tr_state.get_incumbent_index(y_cand, self._rng, mu=mu_cand)
+        self._incumbent_idx = int(candidate_indices[idx_in_cand])
+        self._incumbent_x_unit = x_obs[self._incumbent_idx]
+        if noise_aware and mu_cand is not None:
+            self._incumbent_y_scalar = mu_cand[idx_in_cand : idx_in_cand + 1].copy()
+        else:
+            self._incumbent_y_scalar = y_cand[idx_in_cand : idx_in_cand + 1].copy()
+
+    def _trim_trailing_obs(self) -> None:
+        incumbent_indices = np.array([self._incumbent_idx], dtype=int)
+        obs = turbo_optimizer_utils.trim_trailing_observations(
+            self._x_obs.view().tolist(),
+            self._y_obs.view().tolist(),
+            self._y_tr_list,
+            self._yvar_obs.view().tolist() if len(self._yvar_obs) > 0 else [],
+            trailing_obs=self._trailing_obs,
+            incumbent_indices=incumbent_indices,
+        )
+        self._x_obs = AppendableArray()
+        for x in obs.x_obs:
+            self._x_obs.append(np.array(x))
+        self._y_obs = AppendableArray()
+        for y in obs.y_obs:
+            self._y_obs.append(np.array(y))
+        self._yvar_obs = AppendableArray()
+        if obs.yvar_obs:
+            for yvar in obs.yvar_obs:
+                self._yvar_obs.append(np.array(yvar))
+        self._y_tr_list = obs.y_tr
+
+    def _update_best_value_if_needed(self) -> None:
+        pass
+
+    def tell(
+        self, x: np.ndarray, y: np.ndarray, y_var: np.ndarray | None = None
+    ) -> np.ndarray:
+        with turbo_utils.record_duration(
+            lambda dt: setattr(self, "_dt_tell", float(dt))
+        ):
+            inputs = self._validate_tell_inputs(x, y, y_var)
+            if inputs.x.shape[0] == 0:
+                return (
+                    np.array([], dtype=float)
+                    if inputs.num_metrics == 1
+                    else np.empty((0, inputs.num_metrics), dtype=float)
+                )
+            x_unit = turbo_utils.to_unit(inputs.x, self._bounds)
+            for i in range(inputs.x.shape[0]):
+                self._x_obs.append(x_unit[i])
+                self._y_obs.append(inputs.y[i])
+                if inputs.y_var is not None:
+                    self._yvar_obs.append(inputs.y_var[i])
+            return self._strategy.tell(self, inputs, x_unit=x_unit)
+
+
+def create_optimizer(
+    *,
+    bounds: np.ndarray,
+    config: OptimizerConfig,
+    rng: Generator,
+) -> Optimizer:
+    from .components.builder import build_acquisition_optimizer, build_surrogate
+
+    surrogate = build_surrogate(config)
+    acq_optimizer = build_acquisition_optimizer(config)
+    return Optimizer(
+        bounds=bounds,
+        config=config,
+        rng=rng,
+        surrogate=surrogate,
+        acquisition_optimizer=acq_optimizer,
+    )

enn/turbo/optimizer_config.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+from . import config as _config
+
+__all__ = list(_config.__all__)
+
+
+def __getattr__(name: str) -> object:
+    return getattr(_config, name)

enn/turbo/proposal.py

@@ -1,45 +1,51 @@
 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
-
-    y = y_obs_array.reshape(-1, 1)
+    if y_obs_array.ndim == 1:
+        y = y_obs_array.reshape(-1, 1)
+    else:
+        y = y_obs_array
     if yvar_obs_list is not None and len(yvar_obs_list) > 0:
         yvar_array = np.asarray(yvar_obs_list, dtype=float)
-        yvar = yvar_array.reshape(-1, 1)
+        if yvar_array.ndim == 1:
+            yvar = yvar_array.reshape(-1, 1)
+        else:
+            yvar = yvar_array
     else:
         yvar = None
     x_obs_array = np.asarray(x_obs_list, dtype=float)
@@ -47,10 +53,11 @@ def mk_enn(
         x_obs_array,
         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
@@ -66,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
 
 
@@ -89,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: ...