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

enn/turbo/turbo_gp_fit.py

@@ -0,0 +1,187 @@
+from __future__ import annotations
+from typing import Any
+from enn.enn.enn_util import standardize_y
+from .types import GPDataPrep, GPFitResult
+
+
+def _prepare_gp_data(
+    x_obs_list: list, y_obs_list: list, yvar_obs_list: list | None
+) -> GPDataPrep:
+    import numpy as np
+    import torch
+
+    x = np.asarray(x_obs_list, dtype=float)
+    y = np.asarray(y_obs_list, dtype=float)
+    if y.ndim not in (1, 2):
+        raise ValueError(y.shape)
+    is_multi = y.ndim == 2 and y.shape[1] > 1
+    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 is_multi:
+            raise ValueError("yvar_obs_list not supported for multi-output GP")
+    if is_multi:
+        y_mean, y_std = y.mean(axis=0), y.std(axis=0)
+        y_std = np.where(y_std < 1e-6, 1.0, y_std)
+        z = (y - y_mean) / y_std
+        train_y = torch.as_tensor(z.T, dtype=torch.float64)
+    else:
+        y_mean, y_std = standardize_y(y)
+        z = (y - y_mean) / y_std
+        train_y = torch.as_tensor(z, dtype=torch.float64)
+    return GPDataPrep(
+        train_x=torch.as_tensor(x, dtype=torch.float64),
+        train_y=train_y,
+        is_multi=is_multi,
+        y_mean=y_mean,
+        y_std=y_std,
+        y_raw=y,
+    )
+
+
+def _build_gp_model(
+    train_x: Any,
+    train_y: Any,
+    is_multi: bool,
+    num_dim: int,
+    *,
+    yvar_obs_list: list | None,
+    gp_y_std: Any,
+    y: Any,
+) -> tuple[Any, Any]:
+    import numpy as np
+    import torch
+    from gpytorch.constraints import Interval
+    from gpytorch.likelihoods import GaussianLikelihood
+    from .turbo_gp import TurboGP
+    from .turbo_gp_noisy import TurboGPNoisy
+
+    ls_constr, os_constr = Interval(0.005, 2.0), 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=ls_constr,
+            outputscale_constraint=os_constr,
+            ard_dims=num_dim,
+        ).to(dtype=train_x.dtype)
+        return model, model.likelihood
+    noise_constr = Interval(5e-4, 0.2)
+    num_out = int(y.shape[1]) if is_multi else None
+    if is_multi:
+        likelihood = GaussianLikelihood(
+            noise_constraint=noise_constr, batch_shape=torch.Size([num_out])
+        ).to(dtype=train_y.dtype)
+    else:
+        likelihood = GaussianLikelihood(noise_constraint=noise_constr).to(
+            dtype=train_y.dtype
+        )
+    model = TurboGP(
+        train_x=train_x,
+        train_y=train_y,
+        likelihood=likelihood,
+        lengthscale_constraint=ls_constr,
+        outputscale_constraint=os_constr,
+        ard_dims=num_dim,
+    ).to(dtype=train_x.dtype)
+    likelihood.noise = (
+        torch.full((num_out,), 0.005, dtype=train_y.dtype)
+        if is_multi
+        else torch.tensor(0.005, dtype=train_y.dtype)
+    )
+    return model, likelihood
+
+
+def _init_gp_hyperparams(
+    model: Any, is_multi: bool, num_dim: int, num_out: int | None, dtype: Any
+) -> None:
+    import torch
+
+    if is_multi:
+        model.covar_module.outputscale = torch.ones(num_out, dtype=dtype)
+        model.covar_module.base_kernel.lengthscale = torch.full(
+            (num_out, 1, num_dim), 0.5, dtype=dtype
+        )
+    else:
+        model.covar_module.outputscale = torch.tensor(1.0, dtype=dtype)
+        model.covar_module.base_kernel.lengthscale = torch.full(
+            (num_dim,), 0.5, dtype=dtype
+        )
+
+
+def _train_gp(
+    model: Any, likelihood: Any, train_x: Any, train_y: Any, num_steps: int
+) -> None:
+    import torch
+    from gpytorch.mlls import ExactMarginalLogLikelihood
+
+    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()
+        loss = -mll(model(train_x), train_y)
+        (loss.sum() if loss.ndim != 0 else loss).backward()
+        optimizer.step()
+    model.eval()
+    likelihood.eval()
+
+
+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,
+) -> GPFitResult:
+    import numpy as np
+
+    x = np.asarray(x_obs_list, dtype=float)
+    y = np.asarray(y_obs_list, dtype=float)
+    n, is_multi = x.shape[0], y.ndim == 2 and y.shape[1] > 1
+    if n == 0:
+        return (
+            GPFitResult(
+                model=None,
+                likelihood=None,
+                y_mean=np.zeros(y.shape[1]),
+                y_std=np.ones(y.shape[1]),
+            )
+            if is_multi
+            else GPFitResult(model=None, likelihood=None, y_mean=0.0, y_std=1.0)
+        )
+    if n == 1 and is_multi:
+        return GPFitResult(
+            model=None,
+            likelihood=None,
+            y_mean=y[0].copy(),
+            y_std=np.ones(int(y.shape[1]), dtype=float),
+        )
+    gp_data = _prepare_gp_data(x_obs_list, y_obs_list, yvar_obs_list)
+    model, likelihood = _build_gp_model(
+        gp_data.train_x,
+        gp_data.train_y,
+        gp_data.is_multi,
+        num_dim,
+        yvar_obs_list=yvar_obs_list,
+        gp_y_std=gp_data.y_std,
+        y=gp_data.y_raw,
+    )
+    _init_gp_hyperparams(
+        model,
+        gp_data.is_multi,
+        num_dim,
+        int(gp_data.y_raw.shape[1]) if gp_data.is_multi else None,
+        gp_data.train_x.dtype,
+    )
+    _train_gp(model, likelihood, gp_data.train_x, gp_data.train_y, num_steps)
+    return GPFitResult(
+        model=model, likelihood=likelihood, y_mean=gp_data.y_mean, y_std=gp_data.y_std
+    )

enn/turbo/turbo_gp_noisy.py

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

enn/turbo/turbo_optimizer_utils.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+from .types import ObsLists, TellInputs
+
+if TYPE_CHECKING:
+    import numpy as np
+
+
+def sobol_seed_for_state(
+    seed_base: int, *, restart_generation: int, n_obs: int, num_arms: int
+) -> int:
+    mask64 = (1 << 64) - 1
+    x = int(seed_base) & mask64
+    x ^= (int(restart_generation) + 1) * 0xD1342543DE82EF95 & mask64
+    x ^= (int(n_obs) + 1) * 0x9E3779B97F4A7C15 & mask64
+    x ^= (int(num_arms) + 1) * 0xBF58476D1CE4E5B9 & mask64
+    x = (x + 0x9E3779B97F4A7C15) & mask64
+    z = x
+    z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9 & mask64
+    z = (z ^ (z >> 27)) * 0x94D049BB133111EB & mask64
+    z = z ^ (z >> 31)
+    return int(z & 0xFFFFFFFF)
+
+
+def reset_timing(opt: object) -> None:
+    setattr(opt, "_dt_fit", 0.0)
+    setattr(opt, "_dt_gen", 0.0)
+    setattr(opt, "_dt_sel", 0.0)
+
+
+def validate_tell_inputs(
+    x: np.ndarray, y: np.ndarray, y_var: np.ndarray | None, num_dim: int
+) -> TellInputs:
+    import numpy as np
+
+    x = np.asarray(x, dtype=float)
+    y = np.asarray(y, dtype=float)
+    if x.ndim != 2 or x.shape[1] != num_dim:
+        raise ValueError(x.shape)
+    if y.ndim == 2:
+        if y.shape[0] != x.shape[0]:
+            raise ValueError((x.shape, y.shape))
+        num_metrics = y.shape[1]
+    elif y.ndim == 1:
+        if y.shape[0] != x.shape[0]:
+            raise ValueError((x.shape, y.shape))
+        num_metrics = 1
+    else:
+        raise ValueError(y.shape)
+    if y_var is not None:
+        y_var = np.asarray(y_var, dtype=float)
+        if y_var.shape != y.shape:
+            raise ValueError((y.shape, y_var.shape))
+    return TellInputs(x=x, y=y, y_var=y_var, num_metrics=num_metrics)
+
+
+def trim_trailing_observations(
+    x_obs_list: list,
+    y_obs_list: list,
+    y_tr_list: list,
+    yvar_obs_list: list,
+    *,
+    trailing_obs: int,
+    incumbent_indices: np.ndarray,
+) -> ObsLists:
+    import numpy as np
+
+    num_total = len(x_obs_list)
+    if num_total <= trailing_obs:
+        return ObsLists(
+            x_obs=x_obs_list,
+            y_obs=y_obs_list,
+            y_tr=y_tr_list,
+            yvar_obs=yvar_obs_list,
+        )
+    start_idx = max(0, num_total - trailing_obs)
+    recent_indices = set(range(start_idx, num_total))
+    keep_indices = set(incumbent_indices.tolist()) | recent_indices
+    if len(keep_indices) > trailing_obs:
+        keep_indices = set(incumbent_indices.tolist())
+        remaining_slots = trailing_obs - len(keep_indices)
+        if remaining_slots > 0:
+            recent_non_incumbent = [
+                i for i in range(num_total - 1, -1, -1) if i not in keep_indices
+            ][:remaining_slots]
+            keep_indices.update(recent_non_incumbent)
+    indices = np.array(sorted(keep_indices), dtype=int)
+    x_array = np.asarray(x_obs_list, dtype=float)
+    y_obs_array = np.asarray(y_obs_list, dtype=float)
+    y_tr_array = np.asarray(y_tr_list, dtype=float)
+    new_x = x_array[indices].tolist()
+    new_y_obs = y_obs_array[indices].tolist()
+    new_y_tr = y_tr_array[indices].tolist() if y_tr_array.size > 0 else []
+    new_yvar = yvar_obs_list
+    if len(yvar_obs_list) == len(y_obs_array):
+        yvar_array = np.asarray(yvar_obs_list, dtype=float)
+        new_yvar = yvar_array[indices].tolist()
+    return ObsLists(x_obs=new_x, y_obs=new_y_obs, y_tr=new_y_tr, yvar_obs=new_yvar)

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,34 +111,52 @@ 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, lengthscales: np.ndarray | None = None
@@ -96,26 +166,24 @@ class TurboTrustRegion:
         if lengthscales is None:
             half_length = 0.5 * self.length
         else:
+            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,
-        lengthscales: 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 generate_trust_region_candidates
-
-        return generate_trust_region_candidates(
-            x_center,
-            lengthscales,
-            num_candidates,
-            compute_bounds_1d=self.compute_bounds_1d,
-            rng=rng,
-            sobol_engine=sobol_engine,
-        )
+        import numpy as np
+
+        return np.array([self.get_incumbent_index(y, rng, mu=mu)])