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

enn/turbo/turbo_trust_region.py

@@ -1,55 +1,107 @@
 from __future__ import annotations
-
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 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 .components.incumbent_selector import IncumbentSelector
+    from .config.turbo_tr_config import TurboTRConfig
 
 
 @dataclass
-class TurboTrustRegion:
+class TurboTrustRegion(ScalarIncumbentMixin):
+    config: TurboTRConfig
     num_dim: int
-    num_arms: int
-    length: float = 0.8
-    length_init: float = 0.8
-    length_min: float = 0.5**7
-    length_max: float = 1.6
+    length: float = field(init=False)
     failure_counter: int = 0
     success_counter: int = 0
     best_value: float = -float("inf")
     prev_num_obs: int = 0
+    incumbent_selector: IncumbentSelector = field(default=None, repr=False)
+    _num_arms: int | None = field(default=None, repr=False)
+    _failure_tolerance: int | None = field(default=None, repr=False)
 
     def __post_init__(self) -> None:
+        from .components.incumbent_selector import ScalarIncumbentSelector
+
+        self.length = self.config.length_init
+        self.success_tolerance = 3
+        if self.incumbent_selector is None:
+            self.incumbent_selector = ScalarIncumbentSelector(noise_aware=False)
+
+    @property
+    def length_init(self) -> float:
+        return self.config.length_init
+
+    @property
+    def length_min(self) -> float:
+        return self.config.length_min
+
+    @property
+    def length_max(self) -> float:
+        return self.config.length_max
+
+    @property
+    def num_metrics(self) -> int:
+        return 1
+
+    def _ensure_initialized(self, num_arms: int) -> None:
         import numpy as np
 
-        self.failure_tolerance = int(
-            np.ceil(
-                max(
-                    4.0 / float(self.num_arms),
-                    float(self.num_dim) / float(self.num_arms),
+        if self._num_arms is None:
+            self._num_arms = num_arms
+            self._failure_tolerance = int(
+                np.ceil(
+                    max(
+                        4.0 / float(num_arms),
+                        float(self.num_dim) / float(num_arms),
+                    )
                 )
             )
-        )
-        self.success_tolerance = 3
+        elif num_arms != self._num_arms:
+            raise ValueError(
+                f"num_arms changed from {self._num_arms} to {num_arms}; "
+                "must be consistent across ask() calls"
+            )
+
+    @property
+    def failure_tolerance(self) -> int:
+        if self._failure_tolerance is None:
+            raise RuntimeError("failure_tolerance not initialized; call ask() first")
+        return self._failure_tolerance
 
-    def update(self, values: np.ndarray | Any) -> None:
+    def _coerce_y_obs_1d(self, y_obs: np.ndarray | Any) -> np.ndarray:
         import numpy as np
 
-        if values.ndim != 1:
-            raise ValueError(values.shape)
-        if values.size == 0:
-            return
-        new_values = values[self.prev_num_obs :]
-        if new_values.size == 0:
-            return
-        if not np.isfinite(self.best_value):
-            self.best_value = float(np.max(new_values))
-            self.prev_num_obs = values.size
-            return
-        improved = np.max(new_values) > self.best_value + 1e-3 * np.abs(self.best_value)
+        y_obs = np.asarray(y_obs, dtype=float)
+        if y_obs.ndim == 2:
+            if y_obs.shape[1] != 1:
+                raise ValueError(f"TurboTrustRegion expects m=1, got {y_obs.shape}")
+            return y_obs[:, 0]
+        if y_obs.ndim != 1:
+            raise ValueError(y_obs.shape)
+        return y_obs
+
+    def _coerce_y_incumbent_value(self, y_incumbent: np.ndarray | Any) -> float:
+        import numpy as np
+
+        y_incumbent = np.asarray(y_incumbent, dtype=float).reshape(-1)
+        if y_incumbent.shape != (self.num_metrics,):
+            raise ValueError(
+                f"y_incumbent must have shape ({self.num_metrics},), got {y_incumbent.shape}"
+            )
+        return float(y_incumbent[0])
+
+    def _improvement_scale(self, prev_values: np.ndarray) -> float:
+        import numpy as np
+
+        if prev_values.size == 0:
+            return 0.0
+        return float(np.max(prev_values) - np.min(prev_values))
+
+    def _update_counters_and_length(self, *, improved: bool) -> None:
         if improved:
             self.success_counter += 1
             self.failure_counter = 0
@@ -59,65 +111,79 @@ class TurboTrustRegion:
         if self.success_counter >= self.success_tolerance:
             self.length = min(2.0 * self.length, self.length_max)
             self.success_counter = 0
-        elif self.failure_counter >= self.failure_tolerance:
+        elif (
+            self._failure_tolerance is not None
+            and self.failure_counter >= self._failure_tolerance
+        ):
             self.length = 0.5 * self.length
             self.failure_counter = 0
 
-        self.best_value = max(self.best_value, float(np.max(new_values)))
-        self.prev_num_obs = values.size
+    def update(self, y_obs: np.ndarray | Any, y_incumbent: np.ndarray | Any) -> None:
+        if self._failure_tolerance is None:
+            return
+        y_obs = self._coerce_y_obs_1d(y_obs)
+        n = int(y_obs.size)
+        if n <= 0:
+            return
+        if n < self.prev_num_obs:
+            raise ValueError((n, self.prev_num_obs))
+        if n == self.prev_num_obs:
+            return
+        y_incumbent_value = self._coerce_y_incumbent_value(y_incumbent)
+        import math
+
+        if not math.isfinite(self.best_value):
+            self.best_value = y_incumbent_value
+            self.prev_num_obs = n
+            return
+        prev_values = y_obs[: self.prev_num_obs]
+        scale = self._improvement_scale(prev_values)
+        improved = y_incumbent_value > self.best_value + 1e-3 * scale
+        self._update_counters_and_length(improved=improved)
+        self.best_value = max(self.best_value, y_incumbent_value)
+        self.prev_num_obs = n
 
     def needs_restart(self) -> bool:
         return self.length < self.length_min
 
-    def restart(self) -> None:
+    def restart(self, rng: Any | None = None) -> None:
         self.length = self.length_init
         self.failure_counter = 0
         self.success_counter = 0
         self.best_value = -float("inf")
         self.prev_num_obs = 0
+        self._num_arms = None
+        self._failure_tolerance = None
 
     def validate_request(self, num_arms: int, *, is_fallback: bool = False) -> None:
-        if is_fallback:
-            if num_arms > self.num_arms:
-                raise ValueError(
-                    f"num_arms {num_arms} > configured num_arms {self.num_arms}"
-                )
-        else:
-            if num_arms != self.num_arms:
-                raise ValueError(
-                    f"num_arms {num_arms} != configured num_arms {self.num_arms}"
-                )
+        self._ensure_initialized(num_arms)
 
     def compute_bounds_1d(
-        self, x_center: np.ndarray | Any, weights: np.ndarray | None = None
+        self, x_center: np.ndarray | Any, lengthscales: np.ndarray | None = None
     ) -> tuple[np.ndarray, np.ndarray]:
         import numpy as np
 
-        if weights is None:
+        if lengthscales is None:
             half_length = 0.5 * self.length
         else:
-            half_length = weights * self.length / 2.0
+            lengthscales = np.asarray(lengthscales, dtype=float).reshape(-1)
+            if lengthscales.shape != (self.num_dim,):
+                raise ValueError(
+                    f"lengthscales must have shape ({self.num_dim},), got {lengthscales.shape}"
+                )
+            if not np.all(np.isfinite(lengthscales)):
+                raise ValueError("lengthscales must be finite")
+            half_length = lengthscales * self.length / 2.0
         lb = np.clip(x_center - half_length, 0.0, 1.0)
         ub = np.clip(x_center + half_length, 0.0, 1.0)
         return lb, ub
 
-    def generate_candidates(
+    def get_incumbent_indices(
         self,
-        x_center: np.ndarray,
-        weights: np.ndarray | None,
-        num_candidates: int,
+        y: np.ndarray | Any,
         rng: Generator,
-        sobol_engine: QMCEngine,
+        mu: np.ndarray | None = None,
     ) -> np.ndarray:
-        from .turbo_utils import raasp
-
-        lb, ub = self.compute_bounds_1d(x_center, weights)
-        return raasp(
-            x_center,
-            lb,
-            ub,
-            num_candidates,
-            num_pert=20,
-            rng=rng,
-            sobol_engine=sobol_engine,
-        )
+        import numpy as np
+
+        return np.array([self.get_incumbent_index(y, rng, mu=mu)])

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:
@@ -24,110 +39,46 @@ def _next_power_of_2(n: int) -> int:
 
 
 @contextlib.contextmanager
-def torch_rng_context(generator: torch.Generator | Any) -> Iterator[None]:
+def torch_seed_context(
+    seed: int, device: torch.device | Any | None = None
+) -> Iterator[None]:
     import torch
 
-    old_state = torch.get_rng_state()
-    try:
-        torch.set_rng_state(generator.get_state())
+    devices: list[int] | None = None
+    if device is not None and getattr(device, "type", None) == "cuda":
+        idx = 0 if getattr(device, "index", None) is None else int(device.index)
+        devices = [idx]
+    with torch.random.fork_rng(devices=devices, enabled=True):
+        torch.manual_seed(int(seed))
+        if device is not None and getattr(device, "type", None) == "cuda":
+            torch.cuda.manual_seed_all(int(seed))
+        if device is not None and getattr(device, "type", None) == "mps":
+            if hasattr(torch, "mps") and hasattr(torch.mps, "manual_seed"):
+                torch.mps.manual_seed(int(seed))
         yield
-    finally:
-        torch.set_rng_state(old_state)
 
 
-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,
-    float,
-]:
-    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 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 n == 0:
-        return None, None, 0.0, 1.0
-    if n == 1:
-        gp_y_mean = float(y[0])
-        gp_y_std = 1.0
-        return None, None, gp_y_mean, gp_y_std
-    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)
-    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)
-        likelihood = GaussianLikelihood(noise_constraint=noise_constraint).to(
-            dtype=train_y.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 = 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)
-        likelihood.noise = torch.tensor(0.005, dtype=train_y.dtype)
-    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
-    )
-    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)
-        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):
@@ -138,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))
@@ -161,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)
@@ -174,32 +121,114 @@ 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 to_unit(x: np.ndarray | Any, bounds: np.ndarray | Any) -> np.ndarray:
-    import numpy as np
+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,
+    ub: np.ndarray | list[float] | Any,
+    num_candidates: int,
+    *,
+    rng: Generator | Any,
+    sobol_engine: QMCEngine | Any,
+    num_pert: int = 20,
+) -> np.ndarray:
+    if num_candidates <= 0:
+        raise ValueError(num_candidates)
+    return raasp(
+        center,
+        lb,
+        ub,
+        num_candidates,
+        num_pert=num_pert,
+        rng=rng,
+        sobol_engine=sobol_engine,
+    )
+
 
+def generate_raasp_candidates_uniform(
+    center: np.ndarray | Any,
+    lb: np.ndarray | list[float] | Any,
+    ub: np.ndarray | list[float] | Any,
+    num_candidates: int,
+    *,
+    rng: Generator | Any,
+    num_pert: int = 20,
+) -> np.ndarray:
+    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:
     lb = bounds[:, 0]
     ub = bounds[:, 1]
     if np.any(ub <= lb):
@@ -208,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)
@@ -220,24 +247,24 @@ 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)
     seed = int(rng.integers(2**31 - 1))
-    gen = torch.Generator(device=x_torch.device)
-    gen.manual_seed(seed)
     with (
         torch.no_grad(),
         gpytorch.settings.fast_pred_var(),
-        torch_rng_context(gen),
+        torch_seed_context(seed, device=x_torch.device),
     ):
         posterior = model.posterior(x_torch)
         samples = posterior.sample(sample_shape=torch.Size([1]))
+    if samples.ndim != 2:
+        raise ValueError(samples.shape)
     ts = samples[0].reshape(-1)
     scores = ts.detach().cpu().numpy().reshape(-1)
     scores = gp_y_mean + gp_y_std * scores

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"]