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

enn/turbo/strategies/turbo_hybrid_strategy.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+import time
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+import numpy as np
+from ..sampling import draw_lhd
+from ..types.appendable_array import AppendableArray
+from .optimization_strategy import OptimizationStrategy
+
+if TYPE_CHECKING:
+    from ..optimizer import Optimizer
+    from ..types import TellInputs
+
+
+@dataclass
+class TurboHybridStrategy(OptimizationStrategy):
+    _bounds: np.ndarray
+    _num_dim: int
+    _rng: object
+    _num_init: int
+    _init_lhd: np.ndarray
+    _init_idx: int = 0
+
+    @classmethod
+    def create(
+        cls, *, bounds: np.ndarray, rng: object, num_init: int | None
+    ) -> TurboHybridStrategy:
+        from numpy.random import Generator
+
+        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_dim = int(bounds.shape[0])
+        if not isinstance(rng, Generator):
+            raise TypeError("rng must be a numpy.random.Generator")
+        n_init = int(num_init if num_init is not None else 2 * num_dim)
+        if n_init <= 0:
+            raise ValueError(f"num_init must be > 0, got {n_init}")
+        init_lhd = draw_lhd(bounds=bounds, num_arms=n_init, rng=rng)
+        return cls(
+            _bounds=bounds,
+            _num_dim=num_dim,
+            _rng=rng,
+            _num_init=n_init,
+            _init_lhd=init_lhd,
+        )
+
+    def _reset_init(self) -> None:
+        from numpy.random import Generator
+
+        if not isinstance(self._rng, Generator):
+            raise TypeError("rng must be a numpy.random.Generator")
+        self._init_lhd = draw_lhd(
+            bounds=self._bounds, num_arms=self._num_init, rng=self._rng
+        )
+        self._init_idx = 0
+
+    def _get_init_points(self, num_arms: int, *, fallback_fn=None) -> np.ndarray:
+        num_arms = int(num_arms)
+        num_to_return = min(num_arms, self._num_init - self._init_idx)
+        result = self._init_lhd[self._init_idx : self._init_idx + num_to_return]
+        self._init_idx += num_to_return
+        if num_to_return < num_arms:
+            extra = (
+                fallback_fn
+                or (lambda n: draw_lhd(bounds=self._bounds, num_arms=n, rng=self._rng))
+            )(num_arms - num_to_return)
+            result = np.vstack([result, extra])
+        return result
+
+    def ask(self, opt: Optimizer, num_arms: int) -> np.ndarray:
+        if opt._tr_state.needs_restart():
+            opt._tr_state.restart(opt._rng)
+            opt._restart_generation += 1
+            opt._x_obs = AppendableArray()
+            opt._y_obs = AppendableArray()
+            opt._yvar_obs = AppendableArray()
+            opt._y_tr_list = []
+            opt._incumbent_idx = None
+            opt._incumbent_x_unit = None
+            opt._incumbent_y_scalar = None
+            self._reset_init()
+            return self._get_init_points(num_arms)
+        if self._init_idx < self._num_init:
+
+            def fallback(n: int) -> np.ndarray:
+                return opt._ask_normal(n, is_fallback=True)
+
+            return self._get_init_points(
+                num_arms,
+                fallback_fn=fallback if len(opt._x_obs) > 0 else None,
+            )
+        if len(opt._x_obs) == 0:
+            return draw_lhd(bounds=self._bounds, num_arms=num_arms, rng=self._rng)
+        opt._tr_state.validate_request(int(num_arms))
+        return opt._ask_normal(int(num_arms))
+
+    def init_progress(self) -> tuple[int, int] | None:
+        return (int(self._init_idx), int(self._num_init))
+
+    def tell(
+        self, opt: Optimizer, inputs: TellInputs, *, x_unit: np.ndarray
+    ) -> np.ndarray:
+        x_all = opt._x_obs.view()
+        y_all = opt._y_obs.view()
+        y_var_all = opt._yvar_obs.view() if len(opt._yvar_obs) > 0 else None
+        t0 = time.perf_counter()
+        opt._surrogate.fit(
+            x_all, y_all, y_var_all, num_steps=opt._gp_num_steps, rng=opt._rng
+        )
+        opt._dt_fit = time.perf_counter() - t0
+        opt._y_tr_list = y_all.tolist()
+        opt._update_incumbent()
+        try:
+            new_posterior = opt._surrogate.predict(x_unit)
+            y_estimate = np.asarray(new_posterior.mu, dtype=float)
+        except RuntimeError:
+            y_estimate = np.asarray(inputs.y, dtype=float)
+        y_incumbent = opt._incumbent_y_scalar
+        y_obs = y_all
+        opt._tr_state.update(y_obs, y_incumbent)
+        if opt._trailing_obs is not None:
+            opt._trim_trailing_obs()
+        return y_estimate

enn/turbo/tr_helpers.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from .config.enums import CandidateRV, RAASPDriver
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+    from scipy.stats._qmc import QMCEngine
+
+
+def compute_full_box_bounds_1d(
+    x_center: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
+    import numpy as np
+
+    lb = np.zeros_like(x_center, dtype=float)
+    ub = np.ones_like(x_center, dtype=float)
+    return lb, ub
+
+
+def get_single_incumbent_index(
+    selector,
+    y: np.ndarray,
+    rng: Generator,
+    mu: np.ndarray | None = None,
+) -> np.ndarray:
+    import numpy as np
+
+    y = np.asarray(y, dtype=float)
+    if y.size == 0:
+        return np.array([], dtype=int)
+    best_idx = selector.select(y, mu, rng)
+    return np.array([best_idx])
+
+
+def get_incumbent_index(
+    selector,
+    y: np.ndarray,
+    rng: Generator,
+    mu: np.ndarray | None = None,
+) -> int:
+    import numpy as np
+
+    y = np.asarray(y, dtype=float)
+    if y.size == 0:
+        raise ValueError("y is empty")
+    return int(selector.select(y, mu, rng))
+
+
+def get_scalar_incumbent_value(
+    selector,
+    y_obs: np.ndarray,
+    rng: Generator,
+    *,
+    mu_obs: np.ndarray | None = None,
+) -> np.ndarray:
+    import numpy as np
+
+    y = np.asarray(y_obs, dtype=float)
+    if y.size == 0:
+        return np.array([], dtype=float)
+    idx = get_incumbent_index(selector, y, rng, mu=mu_obs)
+    use_mu = bool(getattr(selector, "noise_aware", False))
+    values = mu_obs if use_mu else y
+    if values is None:
+        raise ValueError("noise_aware incumbent selection requires mu_obs")
+    v = np.asarray(values, dtype=float)
+    if v.ndim == 2:
+        value = float(v[idx, 0])
+    elif v.ndim == 1:
+        value = float(v[idx])
+    else:
+        raise ValueError(v.shape)
+    return np.array([value], dtype=float)
+
+
+class ScalarIncumbentMixin:
+    incumbent_selector: Any
+
+    def get_incumbent_index(
+        self,
+        y: np.ndarray | Any,
+        rng: Generator,
+        mu: np.ndarray | None = None,
+    ) -> int:
+        return get_incumbent_index(self.incumbent_selector, y, rng, mu=mu)
+
+    def get_incumbent_value(
+        self,
+        y_obs: np.ndarray | Any,
+        rng: Generator,
+        mu_obs: np.ndarray | None = None,
+    ) -> np.ndarray:
+        return get_scalar_incumbent_value(
+            self.incumbent_selector, y_obs, rng, mu_obs=mu_obs
+        )
+
+
+def generate_tr_candidates_orig(
+    compute_bounds_1d: Any,
+    x_center: np.ndarray,
+    lengthscales: np.ndarray | None,
+    num_candidates: int,
+    *,
+    rng: Generator,
+    candidate_rv: CandidateRV = CandidateRV.SOBOL,
+    sobol_engine: QMCEngine | None = None,
+) -> np.ndarray:
+    from .turbo_utils import (
+        generate_raasp_candidates,
+        generate_raasp_candidates_uniform,
+    )
+
+    lb, ub = compute_bounds_1d(x_center, lengthscales)
+    if candidate_rv == CandidateRV.SOBOL:
+        if sobol_engine is None:
+            raise ValueError(
+                "sobol_engine is required when candidate_rv=CandidateRV.SOBOL"
+            )
+        return generate_raasp_candidates(
+            x_center, lb, ub, num_candidates, rng=rng, sobol_engine=sobol_engine
+        )
+    if candidate_rv == CandidateRV.UNIFORM:
+        return generate_raasp_candidates_uniform(
+            x_center, lb, ub, num_candidates, rng=rng
+        )
+    raise ValueError(candidate_rv)
+
+
+def generate_tr_candidates_fast(
+    compute_bounds_1d: Any,
+    x_center: np.ndarray,
+    lengthscales: np.ndarray | None,
+    num_candidates: int,
+    *,
+    rng: Generator,
+    candidate_rv: CandidateRV,
+    num_pert: int,
+) -> np.ndarray:
+    import numpy as np
+    from scipy.stats import qmc
+
+    lb, ub = compute_bounds_1d(x_center, lengthscales)
+    num_dim = x_center.shape[-1]
+
+    candidates = np.empty((num_candidates, num_dim), dtype=float)
+    candidates[:] = x_center
+
+    prob_perturb = min(num_pert / num_dim, 1.0)
+    ks = rng.binomial(num_dim, prob_perturb, size=num_candidates)
+    ks = np.maximum(ks, 1)
+    max_k = int(np.max(ks))
+
+    if candidate_rv == CandidateRV.SOBOL:
+        sobol = qmc.Sobol(d=max_k, scramble=True, seed=int(rng.integers(0, 2**31)))
+        samples = sobol.random(num_candidates)
+    elif candidate_rv == CandidateRV.UNIFORM:
+        samples = rng.random((num_candidates, max_k))
+    else:
+        raise ValueError(candidate_rv)
+
+    for i in range(num_candidates):
+        k = ks[i]
+        idx = rng.choice(num_dim, size=k, replace=False)
+        candidates[i, idx] = lb[idx] + (ub[idx] - lb[idx]) * samples[i, :k]
+
+    return candidates
+
+
+def generate_tr_candidates(
+    compute_bounds_1d: Any,
+    x_center: np.ndarray,
+    lengthscales: np.ndarray | None,
+    num_candidates: int,
+    *,
+    rng: Generator,
+    candidate_rv: CandidateRV,
+    sobol_engine: QMCEngine | None,
+    raasp_driver: RAASPDriver,
+    num_pert: int,
+) -> np.ndarray:
+    if raasp_driver == RAASPDriver.FAST:
+        return generate_tr_candidates_fast(
+            compute_bounds_1d,
+            x_center,
+            lengthscales,
+            num_candidates,
+            rng=rng,
+            candidate_rv=candidate_rv,
+            num_pert=num_pert,
+        )
+    return generate_tr_candidates_orig(
+        compute_bounds_1d,
+        x_center,
+        lengthscales,
+        num_candidates,
+        rng=rng,
+        candidate_rv=candidate_rv,
+        sobol_engine=sobol_engine,
+    )

enn/turbo/turbo_gp.py

@@ -1,5 +1,4 @@
 from __future__ import annotations
-
 from .turbo_gp_base import TurboGPBase
 
 
@@ -13,17 +12,25 @@ class TurboGP(TurboGPBase):
         outputscale_constraint,
         ard_dims: int,
     ) -> None:
+        import torch
         from gpytorch.kernels import MaternKernel, ScaleKernel
         from gpytorch.means import ConstantMean
 
         super().__init__(train_x, train_y, likelihood)
-        self.mean_module = ConstantMean()
+        batch_shape = (
+            torch.Size(train_y.shape[:-1])
+            if getattr(train_y, "ndim", 0) > 1
+            else torch.Size()
+        )
+        self.mean_module = ConstantMean(batch_shape=batch_shape)
         base_kernel = MaternKernel(
             nu=2.5,
             ard_num_dims=ard_dims,
+            batch_shape=batch_shape,
             lengthscale_constraint=lengthscale_constraint,
         )
         self.covar_module = ScaleKernel(
             base_kernel,
+            batch_shape=batch_shape,
             outputscale_constraint=outputscale_constraint,
         )

enn/turbo/turbo_gp_base.py

@@ -1,5 +1,4 @@
 from __future__ import annotations
-
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:

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)