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

enn/turbo/config/optimizer_config.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+from .candidate_gen_config import CandidateGenConfig
+from .init_config import InitConfig
+from .surrogate import NoSurrogateConfig, SurrogateConfig
+from .trust_region import TrustRegionConfig, TurboTRConfig
+
+if TYPE_CHECKING:
+    from .acquisition import AcqOptimizerConfig, AcquisitionConfig
+    from .enums import CandidateRV, RAASPDriver
+
+
+@dataclass(frozen=True)
+class ObservationHistoryConfig:
+    trailing_obs: int | None = None
+
+    def __post_init__(self) -> None:
+        if self.trailing_obs is not None and self.trailing_obs <= 0:
+            raise ValueError(f"trailing_obs must be > 0, got {self.trailing_obs}")
+
+
+def _default_acquisition():
+    from .acquisition import RandomAcquisitionConfig
+
+    return RandomAcquisitionConfig()
+
+
+def _default_acq_optimizer():
+    from .acquisition import RAASPOptimizerConfig
+
+    return RAASPOptimizerConfig()
+
+
+@dataclass(frozen=True)
+class OptimizerConfig:
+    trust_region: TrustRegionConfig = TurboTRConfig()
+    candidates: CandidateGenConfig = CandidateGenConfig()
+    init: InitConfig = InitConfig()
+    surrogate: SurrogateConfig = NoSurrogateConfig()
+    acquisition: AcquisitionConfig = field(default_factory=_default_acquisition)
+    acq_optimizer: AcqOptimizerConfig = field(default_factory=_default_acq_optimizer)
+    observation_history: ObservationHistoryConfig = ObservationHistoryConfig()
+
+    def __post_init__(self) -> None:
+        from .validation import validate_optimizer_config
+
+        validate_optimizer_config(self)
+
+    @property
+    def num_metrics(self) -> int | None:
+        from .morbo_tr_config import MorboTRConfig
+
+        if isinstance(self.trust_region, MorboTRConfig):
+            return self.trust_region.num_metrics
+        return None
+
+    @property
+    def candidate_rv(self) -> CandidateRV:
+        return self.candidates.candidate_rv
+
+    @property
+    def raasp_driver(self) -> RAASPDriver:
+        return self.candidates.raasp_driver
+
+    @property
+    def num_candidates(self):
+        return self.candidates.num_candidates
+
+    @property
+    def trailing_obs(self) -> int | None:
+        return self.observation_history.trailing_obs

enn/turbo/config/pareto_acquisition_config.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..components.protocols import AcquisitionOptimizer
+
+
+@dataclass(frozen=True)
+class ParetoAcquisitionConfig:
+    def build(self) -> AcquisitionOptimizer:
+        from ..components.acquisition import ParetoAcqOptimizer
+
+        return ParetoAcqOptimizer()

enn/turbo/config/raasp_driver.py

@@ -0,0 +1,6 @@
+from enum import Enum, auto
+
+
+class RAASPDriver(Enum):
+    ORIG = auto()
+    FAST = auto()

enn/turbo/config/raasp_optimizer_config.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class RAASPOptimizerConfig:
+    pass

enn/turbo/config/random_acquisition_config.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..components.protocols import AcquisitionOptimizer
+
+
+@dataclass(frozen=True)
+class RandomAcquisitionConfig:
+    def build(self) -> AcquisitionOptimizer:
+        from ..components.acquisition import RandomAcqOptimizer
+
+        return RandomAcqOptimizer()

enn/turbo/config/rescalarize.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+from enum import Enum
+
+
+class Rescalarize(Enum):
+    ON_RESTART = "on_restart"
+    ON_PROPOSE = "on_propose"

enn/turbo/config/surrogate.py

@@ -0,0 +1,12 @@
+from .enn_surrogate_config import ENNFitConfig, ENNSurrogateConfig
+from .gp_surrogate_config import GPSurrogateConfig
+from .no_surrogate_config import NoSurrogateConfig
+
+SurrogateConfig = NoSurrogateConfig | GPSurrogateConfig | ENNSurrogateConfig
+__all__ = [
+    "ENNFitConfig",
+    "ENNSurrogateConfig",
+    "GPSurrogateConfig",
+    "NoSurrogateConfig",
+    "SurrogateConfig",
+]

enn/turbo/config/trust_region.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol
+
+from .morbo_tr_config import MorboTRConfig, MultiObjectiveConfig, RescalePolicyConfig
+from .no_tr_config import NoTRConfig
+from .turbo_tr_config import TRLengthConfig, TurboTRConfig
+
+if TYPE_CHECKING:
+    from numpy.random import Generator
+
+    from ..components.protocols import TrustRegion
+    from .enums import CandidateRV
+
+
+class TrustRegionConfig(Protocol):
+    def build(
+        self,
+        *,
+        num_dim: int,
+        rng: Generator,
+        candidate_rv: CandidateRV,
+    ) -> TrustRegion: ...
+
+
+__all__ = [
+    "MorboTRConfig",
+    "MultiObjectiveConfig",
+    "NoTRConfig",
+    "RescalePolicyConfig",
+    "TRLengthConfig",
+    "TrustRegionConfig",
+    "TurboTRConfig",
+]

enn/turbo/config/turbo_tr_config.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from numpy.random import Generator
+
+    from ..components.protocols import TrustRegion
+    from .enums import CandidateRV
+
+
+@dataclass(frozen=True)
+class TRLengthConfig:
+    length_init: float = 0.8
+    length_min: float = 0.5**7
+    length_max: float = 1.6
+
+    def __post_init__(self) -> None:
+        if self.length_init <= 0:
+            raise ValueError(f"length_init must be > 0, got {self.length_init}")
+        if self.length_min <= 0:
+            raise ValueError(f"length_min must be > 0, got {self.length_min}")
+        if self.length_max <= 0:
+            raise ValueError(f"length_max must be > 0, got {self.length_max}")
+        if self.length_min >= self.length_max:
+            raise ValueError(
+                f"length_min must be < length_max, got {self.length_min} >= {self.length_max}"
+            )
+        if self.length_init > self.length_max:
+            raise ValueError(
+                f"length_init must be <= length_max, got {self.length_init} > {self.length_max}"
+            )
+        if self.length_min > self.length_init:
+            raise ValueError(
+                f"length_min must be <= length_init, got {self.length_min} > {self.length_init}"
+            )
+
+
+@dataclass(frozen=True)
+class TurboTRConfig:
+    length: TRLengthConfig = TRLengthConfig()
+    noise_aware: bool = False
+
+    @property
+    def length_init(self) -> float:
+        return self.length.length_init
+
+    @property
+    def length_min(self) -> float:
+        return self.length.length_min
+
+    @property
+    def length_max(self) -> float:
+        return self.length.length_max
+
+    def build(
+        self,
+        *,
+        num_dim: int,
+        rng: Generator,
+        candidate_rv: CandidateRV | None = None,
+    ) -> TrustRegion:
+        from ..components.incumbent_selector import ScalarIncumbentSelector
+        from ..turbo_trust_region import TurboTrustRegion
+
+        return TurboTrustRegion(
+            config=self,
+            num_dim=num_dim,
+            incumbent_selector=ScalarIncumbentSelector(noise_aware=self.noise_aware),
+        )

enn/turbo/config/ucb_acquisition_config.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..components.protocols import AcquisitionOptimizer
+
+
+@dataclass(frozen=True)
+class UCBAcquisitionConfig:
+    def build(self) -> AcquisitionOptimizer:
+        from ..components.acquisition import UCBAcqOptimizer
+
+        return UCBAcqOptimizer()

enn/turbo/config/validation.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+from typing import Any
+
+
+def validate_optimizer_config(cfg: Any) -> None:
+    from .acquisition import (
+        DrawAcquisitionConfig,
+        HnROptimizerConfig,
+        NDSOptimizerConfig,
+        ParetoAcquisitionConfig,
+        UCBAcquisitionConfig,
+    )
+    from .init_strategies import LHDOnlyInit
+    from .surrogate import GPSurrogateConfig, NoSurrogateConfig
+
+    if isinstance(cfg.init.init_strategy, LHDOnlyInit):
+        if not isinstance(cfg.surrogate, NoSurrogateConfig):
+            raise ValueError(
+                "init_strategy='lhd_only' requires NoSurrogateConfig surrogate"
+            )
+    if isinstance(cfg.surrogate, NoSurrogateConfig):
+        if isinstance(cfg.acquisition, DrawAcquisitionConfig):
+            raise ValueError(
+                "DrawAcquisitionConfig (Thompson sampling) requires a surrogate. "
+                "NoSurrogateConfig is not compatible with DrawAcquisitionConfig."
+            )
+        if isinstance(cfg.acquisition, UCBAcquisitionConfig):
+            raise ValueError(
+                "UCBAcquisitionConfig requires a surrogate. "
+                "NoSurrogateConfig is not compatible with UCBAcquisitionConfig."
+            )
+    if isinstance(cfg.acquisition, ParetoAcquisitionConfig):
+        if not isinstance(cfg.acq_optimizer, NDSOptimizerConfig):
+            raise ValueError("ParetoAcquisitionConfig requires NDSOptimizerConfig")
+    if isinstance(cfg.acq_optimizer, HnROptimizerConfig):
+        if isinstance(cfg.acquisition, ParetoAcquisitionConfig):
+            raise ValueError(
+                "HnROptimizerConfig is not compatible with ParetoAcquisitionConfig"
+            )
+        if isinstance(cfg.surrogate, GPSurrogateConfig) and isinstance(
+            cfg.acquisition, DrawAcquisitionConfig
+        ):
+            raise NotImplementedError(
+                "GP surrogate with DrawAcquisitionConfig and HnROptimizerConfig is not yet implemented"
+            )

enn/turbo/hypervolume.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+import numpy as np
+
+
+def hypervolume_2d_max(y: np.ndarray, ref_point: np.ndarray) -> float:
+    y = np.asarray(y, dtype=float)
+    ref_point = np.asarray(ref_point, dtype=float)
+    if y.size == 0:
+        return 0.0
+    if y.ndim != 2:
+        raise ValueError(y.shape)
+    if y.shape[1] != 2:
+        raise ValueError(y.shape)
+    if ref_point.shape != (2,):
+        raise ValueError(ref_point.shape)
+    mask = (y[:, 0] > ref_point[0]) & (y[:, 1] > ref_point[1])
+    y = y[mask]
+    if y.size == 0:
+        return 0.0
+    order = np.argsort(y[:, 0], kind="mergesort")[::-1]
+    y = y[order]
+    hv = 0.0
+    best_y1 = ref_point[1]
+    for i in range(len(y)):
+        x0, y1 = y[i]
+        if y1 > best_y1:
+            best_y1 = y1
+        x_next = y[i + 1, 0] if i + 1 < len(y) else ref_point[0]
+        hv += (x0 - x_next) * (best_y1 - ref_point[1])
+    return float(hv)

enn/turbo/impl_helpers.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+
+
+def get_x_center_fallback(
+    config: Any,
+    x_obs_list: list,
+    y_obs_list: list,
+    rng: Generator,
+    tr_state: Any = None,
+) -> np.ndarray | None:
+    import numpy as np
+    from .components.incumbent_selector import ScalarIncumbentSelector
+
+    y_array = np.asarray(y_obs_list, dtype=float)
+    if y_array.size == 0:
+        return None
+    x_array = np.asarray(x_obs_list, dtype=float)
+    if tr_state is not None and hasattr(tr_state, "incumbent_selector"):
+        selector = tr_state.incumbent_selector
+    else:
+        selector = ScalarIncumbentSelector(noise_aware=False)
+    idx = selector.select(y_array, None, rng)
+    return x_array[idx]
+
+
+def handle_restart_clear_always(
+    x_obs_list: list,
+    y_obs_list: list,
+    yvar_obs_list: list,
+) -> tuple[bool, int]:
+    x_obs_list.clear()
+    y_obs_list.clear()
+    yvar_obs_list.clear()
+    return True, 0
+
+
+def handle_restart_check_multi_objective(
+    tr_state: Any,
+    x_obs_list: list,
+    y_obs_list: list,
+    yvar_obs_list: list,
+    init_idx: int,
+) -> tuple[bool, int]:
+    is_multi = (
+        tr_state is not None
+        and hasattr(tr_state, "num_metrics")
+        and tr_state.num_metrics > 1
+    )
+    if is_multi:
+        x_obs_list.clear()
+        y_obs_list.clear()
+        yvar_obs_list.clear()
+        return True, 0
+    return False, init_idx
+
+
+def estimate_y_passthrough(y_observed: np.ndarray) -> np.ndarray:
+    import numpy as np
+
+    y = np.asarray(y_observed, dtype=float)
+    if y.ndim == 1:
+        return y.reshape(-1, 1)
+    return y

enn/turbo/morbo_trust_region.py

@@ -0,0 +1,250 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from .tr_helpers import ScalarIncumbentMixin
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+    from scipy.stats._qmc import QMCEngine
+
+    from .config.morbo_tr_config import MorboTRConfig
+    from .config.rescalarize import Rescalarize
+
+from .config.enums import CandidateRV, RAASPDriver
+
+
+class MorboTrustRegion(ScalarIncumbentMixin):
+    def __init__(
+        self,
+        config: MorboTRConfig,
+        num_dim: int,
+        *,
+        rng: Generator,
+        candidate_rv: CandidateRV = CandidateRV.SOBOL,
+    ) -> None:
+        from .components.incumbent_selector import ChebyshevIncumbentSelector
+        from .config.turbo_tr_config import TurboTRConfig
+        from .turbo_trust_region import TurboTrustRegion
+
+        self._config = config
+        self._candidate_rv = candidate_rv
+        inner_config = TurboTRConfig(length=config.length)
+        self._tr = TurboTrustRegion(
+            config=inner_config,
+            num_dim=num_dim,
+        )
+        self._num_dim = int(num_dim)
+        self._num_metrics = int(config.num_metrics)
+        if self._num_metrics <= 0:
+            raise ValueError(self._num_metrics)
+        self._alpha = float(config.alpha)
+        self._rescalarize = config.rescalarize
+        self.incumbent_selector = ChebyshevIncumbentSelector(
+            num_metrics=self._num_metrics,
+            alpha=self._alpha,
+            noise_aware=config.noise_aware,
+        )
+        self.incumbent_selector.reset(rng)
+        self._weights = self.incumbent_selector.weights
+        self._y_min: np.ndarray | Any | None = None
+        self._y_max: np.ndarray | Any | None = None
+        self._incumbent_y_raw: np.ndarray | None = None
+
+    @property
+    def num_dim(self) -> int:
+        return self._num_dim
+
+    @property
+    def num_metrics(self) -> int:
+        return self._num_metrics
+
+    @property
+    def weights(self) -> np.ndarray:
+        return self._weights
+
+    @property
+    def length(self) -> float:
+        return float(self._tr.length)
+
+    @property
+    def rescalarize(self) -> Rescalarize:
+        return self._rescalarize
+
+    def resample_weights(self, rng: Generator) -> None:
+        self.incumbent_selector.reset(rng)
+        self._weights = self.incumbent_selector.weights
+
+    def _update_ranges(self, y_obs):
+        self._y_min, self._y_max = y_obs.min(axis=0), y_obs.max(axis=0)
+
+    def update(self, y_obs: np.ndarray | Any, y_incumbent: np.ndarray | Any) -> None:
+        import numpy as np
+
+        y_obs = np.asarray(y_obs, dtype=float)
+        if y_obs.ndim != 2 or y_obs.shape[1] != self._num_metrics:
+            raise ValueError((y_obs.shape, self._num_metrics))
+        n = int(y_obs.shape[0])
+        if n == 0:
+            self._y_min, self._y_max = None, None
+            self._incumbent_y_raw = None
+            self._tr.restart()
+            return
+        prev_n = int(self._tr.prev_num_obs)
+        if n < prev_n:
+            raise ValueError((n, prev_n))
+        self._y_min, self._y_max = y_obs.min(axis=0), y_obs.max(axis=0)
+        y_incumbent = np.asarray(y_incumbent, dtype=float).reshape(1, -1)
+        if y_incumbent.shape != (1, self._num_metrics):
+            raise ValueError(
+                f"y_incumbent must have shape (1, {self._num_metrics}), got {y_incumbent.shape}"
+            )
+        if prev_n == 0:
+            self._handle_initial_update(y_incumbent, n)
+            return
+        if self._incumbent_y_raw is None:
+            self._handle_initial_update(y_incumbent, n)
+            return
+        scores = self.scalarize(
+            np.vstack([self._incumbent_y_raw, y_incumbent]), clip=True
+        )
+        old_score = float(scores[0])
+        new_score = float(scores[1])
+        self._tr.best_value = old_score
+        dummy_y_obs = np.zeros((n, 1))
+        self._tr.update(dummy_y_obs, np.array([new_score]))
+        if new_score > old_score:
+            self._incumbent_y_raw = y_incumbent.copy()
+
+    def _handle_initial_update(self, y_incumbent: np.ndarray, n: int) -> None:
+        import numpy as np
+
+        self._incumbent_y_raw = y_incumbent.copy()
+        score = self.scalarize(y_incumbent, clip=True)
+        dummy_y_obs = np.zeros((n, 1))
+        self._tr.update(dummy_y_obs, score)
+
+    def scalarize(self, y: np.ndarray | Any, *, clip: bool) -> np.ndarray:
+        import numpy as np
+
+        y = np.asarray(y, dtype=float)
+        if y.ndim != 2 or y.shape[1] != self._num_metrics:
+            raise ValueError(y.shape)
+        if self._y_min is None or self._y_max is None:
+            raise RuntimeError("scalarize called before any observations")
+        return self._scalarize_with_ranges(
+            y, y_min=self._y_min, y_max=self._y_max, clip=clip
+        )
+
+    def _scalarize_with_ranges(
+        self,
+        y: np.ndarray | Any,
+        *,
+        y_min: np.ndarray,
+        y_max: np.ndarray,
+        clip: bool,
+    ) -> np.ndarray:
+        import numpy as np
+
+        y = np.asarray(y, dtype=float)
+        if y.ndim != 2 or y.shape[1] != self._num_metrics:
+            raise ValueError(y.shape)
+        y_min = np.asarray(y_min, dtype=float).reshape(-1)
+        y_max = np.asarray(y_max, dtype=float).reshape(-1)
+        if y_min.shape != (self._num_metrics,) or y_max.shape != (self._num_metrics,):
+            raise ValueError((y_min.shape, y_max.shape, self._num_metrics))
+        denom = y_max - y_min
+        is_deg = denom <= 0.0
+        denom_safe = np.where(is_deg, 1.0, denom)
+        z = (y - y_min.reshape(1, -1)) / denom_safe.reshape(1, -1)
+        z = np.where(is_deg, 0.5, z)
+        if clip:
+            z = np.clip(z, 0.0, 1.0)
+        t = z * self._weights.reshape(1, -1)
+        scores = np.min(t, axis=1) + self._alpha * np.sum(t, axis=1)
+        return scores
+
+    def needs_restart(self) -> bool:
+        return self._tr.needs_restart()
+
+    def restart(self, rng: Generator | None = None) -> None:
+        from .config.rescalarize import Rescalarize
+
+        self._y_min = None
+        self._y_max = None
+        self._incumbent_y_raw = None
+        self._tr.restart()
+        if rng is not None and self._rescalarize == Rescalarize.ON_RESTART:
+            self.resample_weights(rng)
+
+    def validate_request(self, num_arms: int, *, is_fallback: bool = False) -> None:
+        return self._tr.validate_request(num_arms, is_fallback=is_fallback)
+
+    def compute_bounds_1d(
+        self, x_center: np.ndarray | Any, lengthscales: np.ndarray | None = None
+    ) -> tuple[np.ndarray, np.ndarray]:
+        return self._tr.compute_bounds_1d(x_center, lengthscales)
+
+    def generate_candidates(
+        self,
+        x_center: np.ndarray,
+        lengthscales: np.ndarray | None,
+        num_candidates: int,
+        rng: Generator,
+        sobol_engine: QMCEngine,
+        raasp_driver: RAASPDriver = RAASPDriver.ORIG,
+        num_pert: int = 20,
+    ) -> np.ndarray:
+        from .tr_helpers import generate_tr_candidates
+
+        return generate_tr_candidates(
+            self._tr.compute_bounds_1d,
+            x_center,
+            lengthscales,
+            num_candidates,
+            rng=rng,
+            candidate_rv=self._candidate_rv,
+            sobol_engine=sobol_engine,
+            raasp_driver=raasp_driver,
+            num_pert=num_pert,
+        )
+
+    def get_incumbent_indices(
+        self,
+        y: np.ndarray | Any,
+        rng: Generator,
+    ) -> np.ndarray:
+        import numpy as np
+
+        y = np.asarray(y, dtype=float)
+        if y.ndim != 2:
+            raise ValueError(y.shape)
+        n = y.shape[0]
+        if n == 0:
+            return np.array([], dtype=int)
+        from nds import ndomsort
+
+        idx_front = np.array(ndomsort.non_domin_sort(-y, only_front_indices=True))
+        return np.where(idx_front == 0)[0]
+
+    def get_incumbent_value(
+        self,
+        y_obs: np.ndarray | Any,
+        rng: Generator,
+        mu_obs: np.ndarray | None = None,
+    ) -> np.ndarray:
+        import numpy as np
+
+        y_obs = np.asarray(y_obs, dtype=float)
+        if y_obs.ndim != 2 or y_obs.shape[1] != self._num_metrics:
+            raise ValueError((y_obs.shape, self._num_metrics))
+        n = int(y_obs.shape[0])
+        if n == 0:
+            return np.array([], dtype=float)
+        idx = self.get_incumbent_index(y_obs, rng, mu=mu_obs)
+        use_mu = bool(getattr(self.incumbent_selector, "noise_aware", False))
+        values = np.asarray(mu_obs if use_mu else y_obs, dtype=float)
+        if values.ndim != 2 or values.shape[1] != self._num_metrics:
+            raise ValueError((values.shape, self._num_metrics))
+        return values[idx : idx + 1].copy()