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

enn/turbo/morbo_trust_region.py

@@ -0,0 +1,189 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+    from scipy.stats._qmc import QMCEngine
+
+from .turbo_trust_region import TurboTrustRegion
+
+
+class MorboTrustRegion:
+    def __init__(
+        self,
+        num_dim: int,
+        num_arms: int,
+        num_metrics: int,
+        *,
+        rng: Generator,
+    ) -> None:
+        import numpy as np
+
+        self._tr = TurboTrustRegion(num_dim=num_dim, num_arms=num_arms)
+        self._num_dim = int(num_dim)
+        self._num_arms = int(num_arms)
+        self._num_metrics = int(num_metrics)
+        if self._num_metrics <= 0:
+            raise ValueError(self._num_metrics)
+
+        alpha = np.ones(self._num_metrics, dtype=float)
+        self._weights = np.asarray(rng.dirichlet(alpha), dtype=float)
+        self._alpha = 0.05
+
+        self._y_min: np.ndarray | Any | None = None
+        self._y_max: np.ndarray | Any | None = None
+
+    @property
+    def num_dim(self) -> int:
+        return self._num_dim
+
+    @property
+    def num_arms(self) -> int:
+        return self._num_arms
+
+    @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)
+
+    def update(self, values: np.ndarray | Any) -> None:
+        raise NotImplementedError(
+            "Use update_xy(x_obs, y_obs) with multi-objective observations."
+        )
+
+    def update_xy(
+        self, x_obs: np.ndarray | Any, y_obs: np.ndarray | Any, *, k: Any = None
+    ) -> None:  # noqa: ARG002
+        import numpy as np
+
+        x_obs = np.asarray(x_obs, dtype=float)
+        y_obs = np.asarray(y_obs, dtype=float)
+
+        if x_obs.ndim != 2 or x_obs.shape[1] != self._num_dim:
+            raise ValueError(x_obs.shape)
+        if y_obs.ndim != 2 or y_obs.shape[0] != x_obs.shape[0]:
+            raise ValueError((x_obs.shape, y_obs.shape))
+        if y_obs.shape[1] != self._num_metrics:
+            raise ValueError((y_obs.shape, self._num_metrics))
+
+        n = int(x_obs.shape[0])
+        if n == 0:
+            self._y_min = None
+            self._y_max = None
+            self._tr.restart()
+            return
+
+        prev_n = int(self._tr.prev_num_obs)
+        if n < prev_n:
+            raise ValueError((n, prev_n))
+
+        y_min_all = y_obs.min(axis=0)
+        y_max_all = y_obs.max(axis=0)
+        y_min_prev = y_obs[:prev_n].min(axis=0) if prev_n > 0 else y_min_all
+        y_max_prev = y_obs[:prev_n].max(axis=0) if prev_n > 0 else y_max_all
+
+        self._y_min = y_min_all
+        self._y_max = y_max_all
+
+        if prev_n == 0:
+            values = np.asarray(self.scalarize(y_obs, clip=True), dtype=float)
+            if values.shape != (n,):
+                raise RuntimeError((values.shape, n))
+            self._tr.update(values)
+            return
+
+        if not np.isfinite(self._tr.best_value):
+            raise RuntimeError(self._tr.best_value)
+
+        values_old = self._scalarize_with_ranges(
+            y_obs, y_min=y_min_prev, y_max=y_max_prev, clip=True
+        )
+        values_old = np.asarray(values_old, dtype=float)
+        if values_old.shape != (n,):
+            raise RuntimeError((values_old.shape, n))
+
+        incumbent_old = float(np.max(values_old[:prev_n]))
+        self._tr.best_value = incumbent_old
+        if prev_n == n:
+            return
+
+        self._tr.update(values_old)
+
+    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) -> None:
+        self._y_min = None
+        self._y_max = None
+        self._tr.restart()
+
+    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,
+    ) -> np.ndarray:
+        return self._tr.generate_candidates(
+            x_center, lengthscales, num_candidates, rng, sobol_engine
+        )

enn/turbo/no_trust_region.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+    from scipy.stats._qmc import QMCEngine
+
+
+@dataclass
+class NoTrustRegion:
+    num_dim: int
+    num_arms: int
+    length: float = 1.0
+
+    def update(self, values: np.ndarray | Any) -> None:
+        return
+
+    def needs_restart(self) -> bool:
+        return False
+
+    def restart(self) -> None:
+        return
+
+    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}"
+                )
+
+    def compute_bounds_1d(
+        self, x_center: np.ndarray | Any, lengthscales: np.ndarray | None = None
+    ) -> 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 generate_candidates(
+        self,
+        x_center: np.ndarray,
+        lengthscales: np.ndarray | None,
+        num_candidates: int,
+        rng: Generator,
+        sobol_engine: QMCEngine,
+    ) -> 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,
+        )

enn/turbo/proposal.py

@@ -22,6 +22,7 @@ def mk_enn(
     k: int,
     num_fit_samples: int | None = None,
     num_fit_candidates: int | None = None,
+    scale_x: bool = False,
     rng: Generator | Any | None = None,
     params_warm_start: ENNParams | Any | None = None,
 ) -> tuple[EpistemicNearestNeighbors | None, ENNParams | None]:
@@ -36,10 +37,17 @@ def mk_enn(
     if y_obs_array.size == 0:
         return None, None
 
-    y = y_obs_array.reshape(-1, 1)
+    # Preserve multi-metric shape if present, otherwise reshape to (n, 1)
+    if y_obs_array.ndim == 1:
+        y = y_obs_array.reshape(-1, 1)
+    else:
+        y = y_obs_array
     if yvar_obs_list is not None and len(yvar_obs_list) > 0:
         yvar_array = np.asarray(yvar_obs_list, dtype=float)
-        yvar = yvar_array.reshape(-1, 1)
+        if yvar_array.ndim == 1:
+            yvar = yvar_array.reshape(-1, 1)
+        else:
+            yvar = yvar_array
     else:
         yvar = None
     x_obs_array = np.asarray(x_obs_list, dtype=float)
@@ -47,6 +55,7 @@ def mk_enn(
         x_obs_array,
         y,
         yvar,
+        scale_x=scale_x,
     )
     if len(enn_model) == 0:
         return None, None

enn/turbo/turbo_config.py

@@ -9,20 +9,64 @@ class TurboConfig:
     k: int | None = None
     num_candidates: int | None = None
     num_init: int | None = None
-    var_scale: float = 1.0
 
     # Experimental
     trailing_obs: int | None = None
+    tr_type: Literal["turbo", "morbo", "none"] = "turbo"
+    num_metrics: int | None = None
+
+    def __post_init__(self) -> None:
+        if self.tr_type not in ["turbo", "morbo", "none"]:
+            raise ValueError(
+                f"tr_type must be 'turbo', 'morbo', or 'none', got {self.tr_type!r}"
+            )
+        if self.num_metrics is not None and self.num_metrics < 1:
+            raise ValueError(f"num_metrics must be >= 1, got {self.num_metrics}")
+        if self.tr_type == "turbo":
+            if self.num_metrics is not None and self.num_metrics != 1:
+                raise ValueError(
+                    f"num_metrics must be 1 for tr_type='turbo', got {self.num_metrics}"
+                )
+        if self.tr_type == "none":
+            if self.num_metrics is not None and self.num_metrics != 1:
+                raise ValueError(
+                    f"num_metrics must be 1 for tr_type='none', got {self.num_metrics}"
+                )
+
+
+@dataclass(frozen=True)
+class TurboOneConfig(TurboConfig):
+    pass
+
+
+@dataclass(frozen=True)
+class TurboZeroConfig(TurboConfig):
+    pass
+
+
+@dataclass(frozen=True)
+class LHDOnlyConfig(TurboConfig):
+    pass
+
+
+@dataclass(frozen=True)
+class TurboENNConfig(TurboConfig):
+    acq_type: Literal["thompson", "pareto", "ucb"] = "pareto"
     num_fit_samples: int | None = None
     num_fit_candidates: int | None = None
-    acq_type: Literal["thompson", "pareto", "ucb"] = "pareto"
-    local_only: bool = False
+    scale_x: bool = False
 
     def __post_init__(self) -> None:
+        super().__post_init__()
         if self.acq_type not in ["thompson", "pareto", "ucb"]:
             raise ValueError(
                 f"acq_type must be 'thompson', 'pareto', or 'ucb', got {self.acq_type!r}"
             )
-        # Pareto acquisition is the only type that works well without hyperparameter fitting
         if self.num_fit_samples is None and self.acq_type != "pareto":
             raise ValueError(f"num_fit_samples required for acq_type={self.acq_type!r}")
+        if self.num_fit_samples is not None and int(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 int(self.num_fit_candidates) <= 0:
+            raise ValueError(
+                f"num_fit_candidates must be > 0, got {self.num_fit_candidates}"
+            )

enn/turbo/turbo_enn_impl.py

@@ -7,11 +7,11 @@ if TYPE_CHECKING:
     from numpy.random import Generator
 
 from .base_turbo_impl import BaseTurboImpl
-from .turbo_config import TurboConfig
+from .turbo_config import TurboENNConfig
 
 
 class TurboENNImpl(BaseTurboImpl):
-    def __init__(self, config: TurboConfig) -> None:
+    def __init__(self, config: TurboENNConfig) -> None:
         super().__init__(config)
         self._enn: Any | None = None
         self._fitted_params: Any | None = None
@@ -22,6 +22,7 @@ class TurboENNImpl(BaseTurboImpl):
         x_obs_list: list,
         y_obs_list: list,
         rng: Generator,
+        tr_state: Any = None,
     ) -> np.ndarray | None:
         import numpy as np
 
@@ -30,7 +31,7 @@ class TurboENNImpl(BaseTurboImpl):
         if len(y_obs_list) == 0:
             return None
         if self._enn is None or self._fitted_params is None:
-            return super().get_x_center(x_obs_list, y_obs_list, rng)
+            return super().get_x_center(x_obs_list, y_obs_list, rng, tr_state)
         if self._fitted_n_obs != len(x_obs_list):
             raise RuntimeError(
                 f"ENN fitted on {self._fitted_n_obs} obs but get_x_center called with {len(x_obs_list)}"
@@ -38,17 +39,40 @@ class TurboENNImpl(BaseTurboImpl):
 
         y_array = np.asarray(y_obs_list, dtype=float)
         x_array = np.asarray(x_obs_list, dtype=float)
-
         k = self._config.k if self._config.k is not None else 10
-        num_top = min(k, len(y_array))
-        top_indices = np.argpartition(-y_array, num_top - 1)[:num_top]
 
-        x_top = x_array[top_indices]
-        posterior = self._enn.posterior(x_top, params=self._fitted_params)
-        mu = posterior.mu[:, 0]
+        # For morbo: top-k per metric → union → scalarize mu
+        if self._config.tr_type == "morbo" and tr_state is not None:
+            if y_array.ndim == 1:
+                y_array = y_array.reshape(-1, tr_state.num_metrics)
+            num_metrics = y_array.shape[1]
+
+            # Find top-k indices for each metric and take union
+            union_indices = set()
+            for m in range(num_metrics):
+                num_top = min(k, len(y_array))
+                top_m = np.argpartition(-y_array[:, m], num_top - 1)[:num_top]
+                union_indices.update(top_m.tolist())
+            union_indices = np.array(sorted(union_indices), dtype=int)
+
+            x_union = x_array[union_indices]
+            posterior = self._enn.posterior(x_union, params=self._fitted_params)
+            mu = posterior.mu  # (len(union), num_metrics)
+
+            scalarized = tr_state.scalarize(mu, clip=False)
+            best_idx_in_union = argmax_random_tie(scalarized, rng=rng)
+            return x_union[best_idx_in_union]
+        else:
+            # Single-objective: original logic
+            num_top = min(k, len(y_array))
+            top_indices = np.argpartition(-y_array, num_top - 1)[:num_top]
+
+            x_top = x_array[top_indices]
+            posterior = self._enn.posterior(x_top, params=self._fitted_params)
+            mu = posterior.mu[:, 0]
 
-        best_idx_in_top = argmax_random_tie(mu, rng=rng)
-        return x_top[best_idx_in_top]
+            best_idx_in_top = argmax_random_tie(mu, rng=rng)
+            return x_top[best_idx_in_top]
 
     def needs_tr_list(self) -> bool:
         return True
@@ -85,6 +109,7 @@ class TurboENNImpl(BaseTurboImpl):
             k=k,
             num_fit_samples=self._config.num_fit_samples,
             num_fit_candidates=self._config.num_fit_candidates,
+            scale_x=self._config.scale_x,
             rng=rng,
             params_warm_start=self._fitted_params,
         )
@@ -99,6 +124,7 @@ class TurboENNImpl(BaseTurboImpl):
         rng: Generator,
         fallback_fn: Callable[[np.ndarray, int], np.ndarray],
         from_unit_fn: Callable[[np.ndarray], np.ndarray],
+        tr_state: Any = None,  # noqa: ARG002
     ) -> np.ndarray:
         import numpy as np
 
@@ -106,7 +132,6 @@ class TurboENNImpl(BaseTurboImpl):
 
         acq_type = self._config.acq_type
         k = self._config.k
-        var_scale = self._config.var_scale
 
         if self._enn is None:
             return fallback_fn(x_cand, num_arms)
@@ -115,9 +140,7 @@ class TurboENNImpl(BaseTurboImpl):
             params = self._fitted_params
         else:
             k_val = k if k is not None else 10
-            params = ENNParams(
-                k=k_val, epi_var_scale=var_scale, ale_homoscedastic_scale=0.0
-            )
+            params = ENNParams(k=k_val, epi_var_scale=1.0, ale_homoscedastic_scale=0.0)
 
         posterior = self._enn.posterior(x_cand, params=params)
         mu = posterior.mu[:, 0]
@@ -155,6 +178,9 @@ class TurboENNImpl(BaseTurboImpl):
         if self._enn is None or self._fitted_params is None:
             return y_observed
         posterior = self._enn.posterior(x_unit, params=self._fitted_params)
+        # For multi-metric (morbo), return full mu; for single-metric, return 1D
+        if posterior.mu.shape[1] > 1:
+            return posterior.mu
         return posterior.mu[:, 0]
 
     def get_mu_sigma(self, x_unit: np.ndarray) -> tuple[np.ndarray, np.ndarray] | None:
@@ -166,11 +192,10 @@ class TurboENNImpl(BaseTurboImpl):
         params = (
             self._fitted_params
             if self._fitted_params is not None
-            else ENNParams(
-                k=k, epi_var_scale=self._config.var_scale, ale_homoscedastic_scale=0.0
-            )
+            else ENNParams(k=k, epi_var_scale=1.0, ale_homoscedastic_scale=0.0)
         )
         posterior = self._enn.posterior(x_unit, params=params, observation_noise=False)
-        mu = posterior.mu[:, 0]
-        sigma = posterior.se[:, 0]
-        return mu, sigma
+        # For multi-metric (morbo), return full mu/sigma; for single-metric, return 1D
+        if posterior.mu.shape[1] > 1:
+            return posterior.mu, posterior.se
+        return posterior.mu[:, 0], posterior.se[:, 0]

enn/turbo/turbo_gp.py

@@ -13,17 +13,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_mode_impl.py

@@ -13,18 +13,25 @@ class TurboModeImpl(Protocol):
         x_obs_list: list,
         y_obs_list: list,
         rng: Generator,
+        tr_state: Any = None,
     ) -> np.ndarray | None: ...
 
     def needs_tr_list(self) -> bool: ...
 
-    def create_trust_region(self, num_dim: int, num_arms: int) -> Any: ...
+    def create_trust_region(
+        self,
+        num_dim: int,
+        num_arms: int,
+        rng: Generator,
+        num_metrics: int | None = None,
+    ) -> Any: ...
 
     def try_early_ask(
         self,
         num_arms: int,
         x_obs_list: list,
         draw_initial_fn: Callable[[int], np.ndarray],
-        get_init_lhd_points_fn: Callable[[int], np.ndarray | None],
+        get_init_lhd_points_fn: Callable[[int], np.ndarray],
     ) -> np.ndarray | None: ...
 
     def handle_restart(
@@ -54,11 +61,13 @@ class TurboModeImpl(Protocol):
         rng: Generator,
         fallback_fn: Callable[[np.ndarray, int], np.ndarray],
         from_unit_fn: Callable[[np.ndarray], np.ndarray],
+        tr_state: Any = None,
     ) -> np.ndarray: ...
 
     def update_trust_region(
         self,
         tr_state: Any,
+        x_obs_list: list,
         y_obs_list: list,
         x_center: np.ndarray | None = None,
         k: int | None = None,