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

enn/turbo/config/__init__.py

@@ -0,0 +1,87 @@
+from .acquisition import (
+    AcqOptimizerConfig,
+    AcquisitionConfig,
+    DrawAcquisitionConfig,
+    HnROptimizerConfig,
+    NDSOptimizerConfig,
+    ParetoAcquisitionConfig,
+    RAASPOptimizerConfig,
+    RandomAcquisitionConfig,
+    UCBAcquisitionConfig,
+)
+from .base import (
+    CandidateGenConfig,
+    InitConfig,
+)
+from .enums import (
+    AcqType,
+    CandidateRV,
+)
+from .init_strategies import HybridInit, InitStrategy, LHDOnlyInit
+from .optimizer_config import OptimizerConfig
+from .surrogate import (
+    ENNFitConfig,
+    ENNSurrogateConfig,
+    GPSurrogateConfig,
+    NoSurrogateConfig,
+    SurrogateConfig,
+)
+from .trust_region import (
+    MorboTRConfig,
+    MultiObjectiveConfig,
+    NoTRConfig,
+    RescalePolicyConfig,
+    TRLengthConfig,
+    TrustRegionConfig,
+    TurboTRConfig,
+)
+
+
+def __getattr__(name: str) -> object:
+    if name in (
+        "lhd_only_config",
+        "turbo_enn_config",
+        "turbo_one_config",
+        "turbo_zero_config",
+    ):
+        from . import factory
+
+        return getattr(factory, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "AcqOptimizerConfig",
+    "AcqType",
+    "AcquisitionConfig",
+    "CandidateGenConfig",
+    "CandidateRV",
+    "DrawAcquisitionConfig",
+    "ENNFitConfig",
+    "ENNSurrogateConfig",
+    "GPSurrogateConfig",
+    "HnROptimizerConfig",
+    "InitConfig",
+    "InitStrategy",
+    "HybridInit",
+    "LHDOnlyInit",
+    "lhd_only_config",
+    "MorboTRConfig",
+    "MultiObjectiveConfig",
+    "NDSOptimizerConfig",
+    "NoSurrogateConfig",
+    "NoTRConfig",
+    "OptimizerConfig",
+    "ParetoAcquisitionConfig",
+    "RAASPOptimizerConfig",
+    "RandomAcquisitionConfig",
+    "RescalePolicyConfig",
+    "SurrogateConfig",
+    "TRLengthConfig",
+    "TrustRegionConfig",
+    "turbo_enn_config",
+    "turbo_one_config",
+    "TurboTRConfig",
+    "turbo_zero_config",
+    "UCBAcquisitionConfig",
+]

enn/turbo/config/acq_type.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+from enum import Enum
+
+
+class AcqType(Enum):
+    THOMPSON = "thompson"
+    PARETO = "pareto"
+    UCB = "ucb"

enn/turbo/config/acquisition.py

@@ -0,0 +1,26 @@
+from .draw_acquisition_config import DrawAcquisitionConfig
+from .hnr_optimizer_config import HnROptimizerConfig
+from .nds_optimizer_config import NDSOptimizerConfig
+from .pareto_acquisition_config import ParetoAcquisitionConfig
+from .raasp_optimizer_config import RAASPOptimizerConfig
+from .random_acquisition_config import RandomAcquisitionConfig
+from .ucb_acquisition_config import UCBAcquisitionConfig
+
+AcquisitionConfig = (
+    UCBAcquisitionConfig
+    | DrawAcquisitionConfig
+    | ParetoAcquisitionConfig
+    | RandomAcquisitionConfig
+)
+AcqOptimizerConfig = RAASPOptimizerConfig | HnROptimizerConfig | NDSOptimizerConfig
+__all__ = [
+    "AcqOptimizerConfig",
+    "AcquisitionConfig",
+    "DrawAcquisitionConfig",
+    "HnROptimizerConfig",
+    "NDSOptimizerConfig",
+    "ParetoAcquisitionConfig",
+    "RAASPOptimizerConfig",
+    "RandomAcquisitionConfig",
+    "UCBAcquisitionConfig",
+]

enn/turbo/config/base.py

@@ -0,0 +1,4 @@
+from .candidate_gen_config import CandidateGenConfig
+from .init_config import InitConfig
+
+__all__ = ["CandidateGenConfig", "InitConfig"]

enn/turbo/config/candidate_gen_config.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+from .candidate_rv import CandidateRV
+from .raasp_driver import RAASPDriver
+
+if TYPE_CHECKING:
+
+    class NumCandidatesFn:
+        def __call__(self, *, num_dim: int, num_arms: int) -> int: ...
+else:
+    NumCandidatesFn = Any
+
+
+def default_num_candidates(*, num_dim: int, num_arms: int) -> int:
+    return min(5000, 100 * int(num_dim))
+
+
+def const_num_candidates(n: int) -> NumCandidatesFn:
+    n = int(n)
+    if n <= 0:
+        raise ValueError(f"num_candidates must be > 0, got {n}")
+
+    def fn(*, num_dim: int, num_arms: int) -> int:
+        return n
+
+    return fn
+
+
+@dataclass(frozen=True)
+class CandidateGenConfig:
+    candidate_rv: CandidateRV = CandidateRV.SOBOL
+    num_candidates: NumCandidatesFn = field(
+        default_factory=lambda: default_num_candidates
+    )
+    raasp_driver: RAASPDriver = RAASPDriver.ORIG
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.candidate_rv, CandidateRV):
+            raise ValueError(
+                f"candidate_rv must be a CandidateRV enum, got {self.candidate_rv!r}"
+            )
+        if not callable(self.num_candidates):
+            raise ValueError(
+                f"num_candidates must be callable, got {type(self.num_candidates)!r}"
+            )
+        test_n = int(self.num_candidates(num_dim=1, num_arms=1))
+        if test_n <= 0:
+            raise ValueError(f"num_candidates must be > 0, got {test_n}")

enn/turbo/config/candidate_rv.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+from enum import Enum
+
+
+class CandidateRV(Enum):
+    SOBOL = "sobol"
+    UNIFORM = "uniform"

enn/turbo/config/draw_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 DrawAcquisitionConfig:
+    def build(self) -> AcquisitionOptimizer:
+        from ..components.acquisition import ThompsonAcqOptimizer
+
+        return ThompsonAcqOptimizer()

enn/turbo/config/enn_index_driver.py

@@ -0,0 +1,6 @@
+from enum import Enum, auto
+
+
+class ENNIndexDriver(Enum):
+    FLAT = auto()
+    HNSW = auto()

enn/turbo/config/enn_surrogate_config.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+from .enums import ENNIndexDriver
+
+if TYPE_CHECKING:
+    from ..components.protocols import Surrogate
+
+
+@dataclass(frozen=True)
+class ENNFitConfig:
+    num_fit_samples: int | None = None
+    num_fit_candidates: int | None = None
+
+    def __post_init__(self) -> None:
+        if self.num_fit_samples is not None and self.num_fit_samples <= 0:
+            raise ValueError(f"num_fit_samples must be > 0, got {self.num_fit_samples}")
+        if self.num_fit_candidates is not None and self.num_fit_candidates <= 0:
+            raise ValueError(
+                f"num_fit_candidates must be > 0, got {self.num_fit_candidates}"
+            )
+
+
+@dataclass(frozen=True)
+class ENNSurrogateConfig:
+    k: int | None = None
+    fit: ENNFitConfig = ENNFitConfig()
+    scale_x: bool = False
+    index_driver: ENNIndexDriver = ENNIndexDriver.FLAT
+
+    @property
+    def num_fit_samples(self) -> int | None:
+        return self.fit.num_fit_samples
+
+    @property
+    def num_fit_candidates(self) -> int | None:
+        return self.fit.num_fit_candidates
+
+    def build(self) -> Surrogate:
+        from ..components.surrogates import ENNSurrogate
+
+        return ENNSurrogate(self)

enn/turbo/config/enums.py

@@ -0,0 +1,7 @@
+from .acq_type import AcqType
+from .candidate_rv import CandidateRV
+from .enn_index_driver import ENNIndexDriver
+from .raasp_driver import RAASPDriver
+from .rescalarize import Rescalarize
+
+__all__ = ["AcqType", "CandidateRV", "ENNIndexDriver", "RAASPDriver", "Rescalarize"]

enn/turbo/config/factory.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+from . import acquisition as acq
+from . import surrogate as sur
+from . import trust_region as tr
+from .candidate_gen_config import (
+    CandidateGenConfig,
+    NumCandidatesFn,
+    const_num_candidates,
+)
+from .enums import AcqType, CandidateRV
+from .init_config import InitConfig
+from .optimizer_config import ObservationHistoryConfig, OptimizerConfig
+
+
+def _make_candidate_gen_config(
+    candidate_rv: CandidateRV,
+    num_candidates: NumCandidatesFn | int | None,
+) -> CandidateGenConfig:
+    if num_candidates is None:
+        return CandidateGenConfig(candidate_rv=candidate_rv)
+    if isinstance(num_candidates, int):
+        num_candidates = const_num_candidates(num_candidates)
+    return CandidateGenConfig(candidate_rv=candidate_rv, num_candidates=num_candidates)
+
+
+def turbo_one_config(
+    *,
+    num_candidates: int | None = None,
+    num_init: int | None = None,
+    trailing_obs: int | None = None,
+    trust_region: tr.TrustRegionConfig | None = None,
+    candidate_rv: CandidateRV = CandidateRV.SOBOL,
+) -> OptimizerConfig:
+    return OptimizerConfig(
+        trust_region=trust_region or tr.TurboTRConfig(),
+        candidates=_make_candidate_gen_config(candidate_rv, num_candidates),
+        init=InitConfig(num_init=num_init),
+        surrogate=sur.GPSurrogateConfig(),
+        acquisition=acq.DrawAcquisitionConfig(),
+        acq_optimizer=acq.RAASPOptimizerConfig(),
+        observation_history=ObservationHistoryConfig(trailing_obs=trailing_obs),
+    )
+
+
+def turbo_zero_config(
+    *,
+    num_candidates: int | None = None,
+    num_init: int | None = None,
+    trailing_obs: int | None = None,
+    trust_region: tr.TrustRegionConfig | None = None,
+    candidate_rv: CandidateRV = CandidateRV.SOBOL,
+) -> OptimizerConfig:
+    return OptimizerConfig(
+        trust_region=trust_region or tr.TurboTRConfig(),
+        candidates=_make_candidate_gen_config(candidate_rv, num_candidates),
+        init=InitConfig(num_init=num_init),
+        surrogate=sur.NoSurrogateConfig(),
+        acquisition=acq.RandomAcquisitionConfig(),
+        acq_optimizer=acq.RAASPOptimizerConfig(),
+        observation_history=ObservationHistoryConfig(trailing_obs=trailing_obs),
+    )
+
+
+def turbo_enn_config(
+    *,
+    enn: sur.ENNSurrogateConfig | None = None,
+    trust_region: tr.TrustRegionConfig | None = None,
+    candidates: CandidateGenConfig | None = None,
+    num_init: int | None = None,
+    trailing_obs: int | None = None,
+    acq_type: AcqType = AcqType.PARETO,
+) -> OptimizerConfig:
+    if acq_type == AcqType.PARETO:
+        acquisition = acq.ParetoAcquisitionConfig()
+        acq_optimizer = acq.NDSOptimizerConfig()
+    elif acq_type == AcqType.UCB:
+        acquisition = acq.UCBAcquisitionConfig()
+        acq_optimizer = acq.RAASPOptimizerConfig()
+    elif acq_type == AcqType.THOMPSON:
+        acquisition = acq.DrawAcquisitionConfig()
+        acq_optimizer = acq.RAASPOptimizerConfig()
+    else:
+        raise ValueError(
+            f"acq_type must be AcqType.THOMPSON, AcqType.PARETO, or AcqType.UCB, got {acq_type!r}"
+        )
+    surrogate = enn if enn is not None else sur.ENNSurrogateConfig()
+    if surrogate.num_fit_samples is None and acq_type != AcqType.PARETO:
+        raise ValueError(f"enn.num_fit_samples required for acq_type={acq_type!r}")
+    return OptimizerConfig(
+        trust_region=trust_region or tr.TurboTRConfig(),
+        candidates=candidates or CandidateGenConfig(),
+        init=InitConfig(num_init=num_init),
+        surrogate=surrogate,
+        acquisition=acquisition,
+        acq_optimizer=acq_optimizer,
+        observation_history=ObservationHistoryConfig(trailing_obs=trailing_obs),
+    )
+
+
+def lhd_only_config(
+    *,
+    num_candidates: int | None = None,
+    num_init: int | None = None,
+    trailing_obs: int | None = None,
+    trust_region: tr.TrustRegionConfig | None = None,
+    candidate_rv: CandidateRV = CandidateRV.SOBOL,
+) -> OptimizerConfig:
+    from .init_strategies import LHDOnlyInit
+
+    return OptimizerConfig(
+        trust_region=trust_region or tr.NoTRConfig(),
+        candidates=_make_candidate_gen_config(candidate_rv, num_candidates),
+        init=InitConfig(init_strategy=LHDOnlyInit(), num_init=num_init),
+        surrogate=sur.NoSurrogateConfig(),
+        acquisition=acq.RandomAcquisitionConfig(),
+        acq_optimizer=acq.RAASPOptimizerConfig(),
+        observation_history=ObservationHistoryConfig(trailing_obs=trailing_obs),
+    )

enn/turbo/config/gp_surrogate_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 Surrogate
+
+
+@dataclass(frozen=True)
+class GPSurrogateConfig:
+    def build(self) -> Surrogate:
+        from ..components.surrogates import GPSurrogate
+
+        return GPSurrogate()

enn/turbo/config/hnr_optimizer_config.py

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

enn/turbo/config/init_config.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from .init_strategies import HybridInit, InitStrategy
+
+
+@dataclass(frozen=True)
+class InitConfig:
+    init_strategy: InitStrategy = HybridInit()
+    num_init: int | None = None
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.init_strategy, InitStrategy):
+            raise ValueError(
+                f"init_strategy must be an InitStrategy, got {self.init_strategy!r}"
+            )
+        if self.num_init is not None and self.num_init <= 0:
+            raise ValueError(f"num_init must be > 0, got {self.num_init}")

enn/turbo/config/init_strategies/__init__.py

@@ -0,0 +1,9 @@
+from .hybrid_init import HybridInit
+from .init_strategy import InitStrategy
+from .lhd_only_init import LHDOnlyInit
+
+__all__ = [
+    "HybridInit",
+    "InitStrategy",
+    "LHDOnlyInit",
+]

enn/turbo/config/init_strategies/hybrid_init.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+from .init_strategy import InitStrategy
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+    from ...strategies import OptimizationStrategy
+
+
+@dataclass(frozen=True)
+class HybridInit(InitStrategy):
+    def create_runtime_strategy(
+        self,
+        *,
+        bounds: np.ndarray,
+        rng: Generator,
+        num_init: int | None,
+    ) -> OptimizationStrategy:
+        from ...strategies import TurboHybridStrategy
+
+        return TurboHybridStrategy.create(bounds=bounds, rng=rng, num_init=num_init)

enn/turbo/config/init_strategies/init_strategy.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+    from ...strategies import OptimizationStrategy
+
+
+class InitStrategy(ABC):
+    @abstractmethod
+    def create_runtime_strategy(
+        self,
+        *,
+        bounds: np.ndarray,
+        rng: Generator,
+        num_init: int | None,
+    ) -> OptimizationStrategy: ...

enn/turbo/config/init_strategies/lhd_only_init.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+from .init_strategy import InitStrategy
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+    from ...strategies import OptimizationStrategy
+
+
+@dataclass(frozen=True)
+class LHDOnlyInit(InitStrategy):
+    def create_runtime_strategy(
+        self,
+        *,
+        bounds: np.ndarray,
+        rng: Generator,
+        num_init: int | None,
+    ) -> OptimizationStrategy:
+        from ...strategies import LHDOnlyStrategy
+
+        del num_init
+        return LHDOnlyStrategy.create(bounds=bounds, rng=rng)

enn/turbo/config/morbo_tr_config.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from .rescalarize import Rescalarize
+from .turbo_tr_config import TRLengthConfig
+
+if TYPE_CHECKING:
+    from numpy.random import Generator
+
+    from ..components.protocols import TrustRegion
+    from .enums import CandidateRV
+
+from .enums import CandidateRV
+
+
+@dataclass(frozen=True)
+class MultiObjectiveConfig:
+    num_metrics: int
+    alpha: float = 0.05
+
+    def __post_init__(self) -> None:
+        if self.num_metrics < 2:
+            raise ValueError(
+                f"num_metrics must be >= 2 for MORBO, got {self.num_metrics}"
+            )
+        if self.alpha <= 0:
+            raise ValueError(f"alpha must be > 0, got {self.alpha}")
+
+
+@dataclass(frozen=True)
+class RescalePolicyConfig:
+    rescalarize: Rescalarize = Rescalarize.ON_PROPOSE
+
+
+@dataclass(frozen=True)
+class MorboTRConfig:
+    multi_objective: MultiObjectiveConfig
+    length: TRLengthConfig = TRLengthConfig()
+    rescale_policy: RescalePolicyConfig = RescalePolicyConfig()
+    noise_aware: bool = False
+
+    @property
+    def rescalarize(self) -> Rescalarize:
+        return self.rescale_policy.rescalarize
+
+    @property
+    def num_metrics(self) -> int:
+        return self.multi_objective.num_metrics
+
+    @property
+    def alpha(self) -> float:
+        return self.multi_objective.alpha
+
+    @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 = CandidateRV.SOBOL,
+    ) -> TrustRegion:
+        from ..morbo_trust_region import MorboTrustRegion
+
+        return MorboTrustRegion(
+            config=self,
+            num_dim=num_dim,
+            rng=rng,
+            candidate_rv=candidate_rv,
+        )

enn/turbo/config/nds_optimizer_config.py

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

enn/turbo/config/no_surrogate_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 Surrogate
+
+
+@dataclass(frozen=True)
+class NoSurrogateConfig:
+    def build(self) -> Surrogate:
+        from ..components.surrogates import NoSurrogate
+
+        return NoSurrogate()

enn/turbo/config/no_tr_config.py

@@ -0,0 +1,31 @@
+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 NoTRConfig:
+    noise_aware: bool = False
+
+    def build(
+        self,
+        *,
+        num_dim: int,
+        rng: Generator,
+        candidate_rv: CandidateRV | None = None,
+    ) -> TrustRegion:
+        from ..components.incumbent_selector import ScalarIncumbentSelector
+        from ..no_trust_region import NoTrustRegion
+
+        return NoTrustRegion(
+            config=self,
+            num_dim=num_dim,
+            incumbent_selector=ScalarIncumbentSelector(noise_aware=self.noise_aware),
+        )