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

enn/enn/enn.py

@@ -12,12 +12,16 @@ if TYPE_CHECKING:
 class EpistemicNearestNeighbors:
     def __init__(
         self,
-        train_x: np.ndarray | Any,
-        train_y: np.ndarray | Any,
-        train_yvar: np.ndarray | Any | None = None,
+        train_x: np.ndarray,
+        train_y: np.ndarray,
+        train_yvar: np.ndarray | None = None,
+        *,
+        scale_x: bool = False,
     ) -> None:
         import numpy as np
 
+        train_x = np.asarray(train_x, dtype=float)
+        train_y = np.asarray(train_y, dtype=float)
         if train_x.ndim != 2:
             raise ValueError(train_x.shape)
         if train_y.ndim != 2:
@@ -25,22 +29,41 @@ class EpistemicNearestNeighbors:
         if train_x.shape[0] != train_y.shape[0]:
             raise ValueError((train_x.shape, train_y.shape))
         if train_yvar is not None:
+            train_yvar = np.asarray(train_yvar, dtype=float)
             if train_yvar.ndim != 2:
                 raise ValueError(train_yvar.shape)
             if train_y.shape != train_yvar.shape:
                 raise ValueError((train_y.shape, train_yvar.shape))
-        self._train_x = np.asarray(train_x, dtype=float)
-        self._train_y = np.asarray(train_y, dtype=float)
-        self._train_yvar = (
-            np.asarray(train_yvar, dtype=float) if train_yvar is not None else None
-        )
+
+        self._train_x = train_x
+        self._train_y = train_y
+        self._train_yvar = train_yvar
         self._num_obs, self._num_dim = self._train_x.shape
         _, self._num_metrics = self._train_y.shape
         self._eps_var = 1e-9
+        self._scale_x = bool(scale_x)
+        if self._scale_x:
+            if len(self._train_x) < 2:
+                x_scale = np.ones((1, self._num_dim), dtype=float)
+            else:
+                x_scale = np.std(self._train_x, axis=0, keepdims=True).astype(float)
+                x_scale = np.where(
+                    np.isfinite(x_scale) & (x_scale > 1e-12),
+                    x_scale,
+                    1.0,
+                )
+            self._x_scale = x_scale
+            self._train_x_scaled = self._train_x / self._x_scale
+        else:
+            self._x_scale = np.ones((1, self._num_dim), dtype=float)
+            self._train_x_scaled = self._train_x
         if len(self._train_y) < 2:
             self._y_scale = np.ones(shape=(1, self._num_metrics), dtype=float)
         else:
-            self._y_scale = np.std(self._train_y, axis=0, keepdims=True).astype(float)
+            y_scale = np.std(self._train_y, axis=0, keepdims=True).astype(float)
+            self._y_scale = np.where(
+                np.isfinite(y_scale) & (y_scale > 0.0), y_scale, 1.0
+            )
 
         self._index: Any | None = None
         self._build_index()
@@ -70,14 +93,42 @@ class EpistemicNearestNeighbors:
 
         if self._num_obs == 0:
             return
-        x_f32 = self._train_x.astype(np.float32, copy=False)
+        x_f32 = self._train_x_scaled.astype(np.float32, copy=False)
         index = faiss.IndexFlatL2(self._num_dim)
         index.add(x_f32)
         self._index = index
 
+    def _search_index(
+        self,
+        x: np.ndarray,
+        *,
+        search_k: int,
+        exclude_nearest: bool,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        import numpy as np
+
+        search_k = int(search_k)
+        if search_k <= 0:
+            raise ValueError(search_k)
+        x = np.asarray(x, dtype=float)
+        if x.ndim != 2 or x.shape[1] != self._num_dim:
+            raise ValueError(x.shape)
+        if self._index is None:
+            raise RuntimeError("index is not initialized")
+
+        x_scaled = x / self._x_scale if self._scale_x else x
+        x_f32 = x_scaled.astype(np.float32, copy=False)
+        dist2s_full, idx_full = self._index.search(x_f32, search_k)
+        dist2s_full = dist2s_full.astype(float)
+        idx_full = idx_full.astype(int)
+        if exclude_nearest:
+            dist2s_full = dist2s_full[:, 1:]
+            idx_full = idx_full[:, 1:]
+        return dist2s_full, idx_full
+
     def posterior(
         self,
-        x: np.ndarray | Any,
+        x: np.ndarray,
         *,
         params: ENNParams,
         exclude_nearest: bool = False,
@@ -97,7 +148,7 @@ class EpistemicNearestNeighbors:
 
     def batch_posterior(
         self,
-        x: np.ndarray | Any,
+        x: np.ndarray,
         paramss: list[ENNParams],
         *,
         exclude_nearest: bool = False,
@@ -107,6 +158,7 @@ class EpistemicNearestNeighbors:
 
         from .enn_normal import ENNNormal
 
+        x = np.asarray(x, dtype=float)
         if x.ndim != 2:
             raise ValueError(x.shape)
         if x.shape[1] != self._num_dim:
@@ -126,15 +178,9 @@ class EpistemicNearestNeighbors:
             search_k = int(min(max_k + 1, len(self)))
         else:
             search_k = int(min(max_k, len(self)))
-        x_f32 = x.astype(np.float32, copy=False)
-        if self._index is None:
-            raise RuntimeError("index is not initialized")
-        dist2s_full, idx_full = self._index.search(x_f32, search_k)
-        dist2s_full = dist2s_full.astype(float)
-        idx_full = idx_full.astype(int)
-        if exclude_nearest:
-            dist2s_full = dist2s_full[:, 1:]
-            idx_full = idx_full[:, 1:]
+        dist2s_full, idx_full = self._search_index(
+            x, search_k=search_k, exclude_nearest=exclude_nearest
+        )
         mu_all = np.zeros((num_params, batch_size, self._num_metrics), dtype=float)
         se_all = np.zeros((num_params, batch_size, self._num_metrics), dtype=float)
         available_k = search_k - 1 if exclude_nearest else search_k
@@ -178,7 +224,7 @@ class EpistemicNearestNeighbors:
 
     def neighbors(
         self,
-        x: np.ndarray | Any,
+        x: np.ndarray,
         k: int,
         *,
         exclude_nearest: bool = False,
@@ -210,15 +256,9 @@ class EpistemicNearestNeighbors:
             search_k = int(min(k, len(self)))
         if search_k == 0:
             return []
-        x_f32 = x.astype(np.float32, copy=False)
-        if self._index is None:
-            raise RuntimeError("index is not initialized")
-        dist2s_full, idx_full = self._index.search(x_f32, search_k)
-        dist2s_full = dist2s_full.astype(float)
-        idx_full = idx_full.astype(int)
-        if exclude_nearest:
-            dist2s_full = dist2s_full[:, 1:]
-            idx_full = idx_full[:, 1:]
+        dist2s_full, idx_full = self._search_index(
+            x, search_k=search_k, exclude_nearest=exclude_nearest
+        )
         actual_k = min(k, len(idx_full[0]))
         idx = idx_full[0, :actual_k]
         result = []

enn/enn/enn_fit.py

@@ -9,8 +9,6 @@ if TYPE_CHECKING:
     from .enn import EpistemicNearestNeighbors
     from .enn_params import ENNParams
 
-from .enn_util import standardize_y
-
 
 def subsample_loglik(
     model: EpistemicNearestNeighbors | Any,
@@ -23,17 +21,21 @@ def subsample_loglik(
 ) -> list[float]:
     import numpy as np
 
-    if x.ndim != 2:
-        raise ValueError(x.shape)
-    if y.ndim != 1:
-        raise ValueError(y.shape)
-    if x.shape[0] != y.shape[0]:
-        raise ValueError((x.shape, y.shape))
+    x_array = np.asarray(x, dtype=float)
+    if x_array.ndim != 2:
+        raise ValueError(x_array.shape)
+    y_array = np.asarray(y, dtype=float)
+    if y_array.ndim == 1:
+        y_array = y_array.reshape(-1, 1)
+    if y_array.ndim != 2:
+        raise ValueError(y_array.shape)
+    if x_array.shape[0] != y_array.shape[0]:
+        raise ValueError((x_array.shape, y_array.shape))
     if P <= 0:
         raise ValueError(P)
     if len(paramss) == 0:
         raise ValueError("paramss must be non-empty")
-    n = x.shape[0]
+    n = x_array.shape[0]
     if n == 0:
         return [0.0] * len(paramss)
     if len(model) <= 1:
@@ -43,8 +45,8 @@ def subsample_loglik(
         indices = np.arange(n, dtype=int)
     else:
         indices = rng.permutation(n)[:P_actual]
-    x_selected = x[indices]
-    y_selected = y[indices]
+    x_selected = x_array[indices]
+    y_selected = y_array[indices]
     if not np.isfinite(y_selected).all():
         return [0.0] * len(paramss)
     post_batch = model.batch_posterior(
@@ -52,16 +54,22 @@ def subsample_loglik(
     )
     mu_batch = post_batch.mu
     se_batch = post_batch.se
-    if mu_batch.shape[2] == 1:
-        mu_batch = mu_batch[:, :, 0]
-        se_batch = se_batch[:, :, 0]
     num_params = len(paramss)
-    if mu_batch.shape != (num_params, P_actual) or se_batch.shape != (
+    num_outputs = y_selected.shape[1]
+    if mu_batch.shape != (num_params, P_actual, num_outputs) or se_batch.shape != (
         num_params,
         P_actual,
+        num_outputs,
     ):
-        raise ValueError((mu_batch.shape, se_batch.shape, (num_params, P_actual)))
-    _, y_std = standardize_y(y)
+        raise ValueError(
+            (
+                mu_batch.shape,
+                se_batch.shape,
+                (num_params, P_actual, num_outputs),
+            )
+        )
+    y_std = np.std(y_array, axis=0, keepdims=True).astype(float)
+    y_std = np.where(np.isfinite(y_std) & (y_std > 0.0), y_std, 1.0)
     y_scaled = y_selected / y_std
     mu_scaled = mu_batch / y_std
     se_scaled = se_batch / y_std
@@ -100,12 +108,6 @@ def enn_fit(
 
     train_x = model.train_x
     train_y = model.train_y
-    train_yvar = model.train_yvar
-    if train_y.shape[1] != 1:
-        raise ValueError(train_y.shape)
-    if train_yvar is not None and train_yvar.shape[1] != 1:
-        raise ValueError(train_yvar.shape)
-    y = train_y[:, 0]
     log_min = -3.0
     log_max = 3.0
     epi_var_scale_log_values = rng.uniform(log_min, log_max, size=num_fit_candidates)
@@ -135,7 +137,7 @@ def enn_fit(
     import numpy as np
 
     logliks = subsample_loglik(
-        model, train_x, y, paramss=paramss, P=num_fit_samples, rng=rng
+        model, train_x, train_y, paramss=paramss, P=num_fit_samples, rng=rng
     )
     if len(logliks) == 0:
         return paramss[0]

enn/enn/enn_normal.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     import numpy as np
+    from numpy.random import Generator
 
 
 @dataclass
@@ -15,8 +16,8 @@ class ENNNormal:
     def sample(
         self,
         num_samples: int,
-        rng,
-        clip=None,
+        rng: Generator,
+        clip: float | None = None,
     ) -> np.ndarray:
         import numpy as np
 

enn/enn/enn_params.py

@@ -8,3 +8,16 @@ class ENNParams:
     k: int
     epi_var_scale: float
     ale_homoscedastic_scale: float
+
+    def __post_init__(self) -> None:
+        import numpy as np
+
+        k = int(self.k)
+        if k <= 0:
+            raise ValueError(f"k must be > 0, got {k}")
+        epi_var_scale = float(self.epi_var_scale)
+        if not np.isfinite(epi_var_scale) or epi_var_scale < 0.0:
+            raise ValueError(f"epi_var_scale must be >= 0, got {epi_var_scale}")
+        ale_scale = float(self.ale_homoscedastic_scale)
+        if not np.isfinite(ale_scale) or ale_scale < 0.0:
+            raise ValueError(f"ale_homoscedastic_scale must be >= 0, got {ale_scale}")

enn/enn/enn_util.py

@@ -67,7 +67,6 @@ def arms_from_pareto_fronts(
     rng: Generator | Any,
 ) -> np.ndarray:
     import numpy as np
-    from nds import ndomsort
 
     if x_cand.ndim != 2:
         raise ValueError(x_cand.shape)
@@ -75,22 +74,51 @@ def arms_from_pareto_fronts(
         raise ValueError((mu.shape, se.shape))
     if mu.size != x_cand.shape[0]:
         raise ValueError((mu.size, x_cand.shape[0]))
+    num_arms = int(num_arms)
+    if num_arms <= 0:
+        raise ValueError(num_arms)
+    if not np.all(np.isfinite(mu)) or not np.all(np.isfinite(se)):
+        raise ValueError("mu and se must be finite")
 
-    combined = np.column_stack([mu, se])
-    idx_front = np.array(ndomsort.non_domin_sort(-combined, only_front_indices=True))
+    def _pareto_front_2d_maximize(
+        mu_: np.ndarray, se_: np.ndarray, idx: np.ndarray
+    ) -> np.ndarray:
+        order = np.lexsort((-se_[idx], -mu_[idx]))
+        sorted_idx = idx[order]
+        keep: list[int] = []
+        best_se = -float("inf")
+        last_mu = float("nan")
+        last_se = float("nan")
+        for i in sorted_idx.tolist():
+            s = float(se_[i])
+            m = float(mu_[i])
+            if s > best_se:
+                keep.append(i)
+                best_se = s
+                last_mu = m
+                last_se = s
+            elif s == best_se and m == last_mu and s == last_se:
+                keep.append(i)
+        return np.asarray(keep, dtype=int)
 
     i_keep: list[int] = []
-    for n_front in range(1 + int(idx_front.max())):
-        front_indices = np.where(idx_front == n_front)[0]
+    remaining = np.arange(mu.size, dtype=int)
+    while remaining.size > 0 and len(i_keep) < num_arms:
+        front_indices = _pareto_front_2d_maximize(mu, se, remaining)
+        if front_indices.size == 0:
+            raise RuntimeError("pareto front extraction failed")
         front_indices = front_indices[np.argsort(-mu[front_indices])]
-        if len(i_keep) + len(front_indices) <= num_arms:
+        if len(i_keep) + int(front_indices.size) <= num_arms:
             i_keep.extend(front_indices.tolist())
-        else:
-            remaining = num_arms - len(i_keep)
-            i_keep.extend(
-                rng.choice(front_indices, size=remaining, replace=False).tolist()
-            )
-            break
+            is_front = np.zeros(mu.size, dtype=bool)
+            is_front[front_indices] = True
+            remaining = remaining[~is_front[remaining]]
+            continue
+        remaining_arms = num_arms - len(i_keep)
+        i_keep.extend(
+            rng.choice(front_indices, size=remaining_arms, replace=False).tolist()
+        )
+        break
 
     i_keep = np.array(i_keep)
     return x_cand[i_keep[np.argsort(-mu[i_keep])]]

enn/turbo/base_turbo_impl.py

@@ -18,6 +18,7 @@ class BaseTurboImpl:
         x_obs_list: list,
         y_obs_list: list,
         rng: Generator,
+        tr_state: Any = None,
     ) -> np.ndarray | None:
         import numpy as np
 
@@ -26,24 +27,58 @@ class BaseTurboImpl:
         y_array = np.asarray(y_obs_list, dtype=float)
         if y_array.size == 0:
             return None
-        idx = argmax_random_tie(y_array, rng=rng)
         x_array = np.asarray(x_obs_list, dtype=float)
+
+        # For morbo: scalarize raw y observations
+        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)
+            scalarized = tr_state.scalarize(y_array, clip=True)
+            idx = argmax_random_tie(scalarized, rng=rng)
+        else:
+            idx = argmax_random_tie(y_array, rng=rng)
+
         return x_array[idx]
 
     def needs_tr_list(self) -> bool:
         return False
 
-    def create_trust_region(self, num_dim: int, num_arms: int) -> Any:
-        from .turbo_trust_region import TurboTrustRegion
-
-        return TurboTrustRegion(num_dim=num_dim, num_arms=num_arms)
+    def create_trust_region(
+        self,
+        num_dim: int,
+        num_arms: int,
+        rng: Generator,
+        num_metrics: int | None = None,
+    ) -> Any:
+        if self._config.tr_type == "none":
+            from .no_trust_region import NoTrustRegion
+
+            return NoTrustRegion(num_dim=num_dim, num_arms=num_arms)
+        elif self._config.tr_type == "turbo":
+            from .turbo_trust_region import TurboTrustRegion
+
+            return TurboTrustRegion(num_dim=num_dim, num_arms=num_arms)
+        elif self._config.tr_type == "morbo":
+            from .morbo_trust_region import MorboTrustRegion
+
+            effective_num_metrics = num_metrics or self._config.num_metrics
+            if effective_num_metrics is None:
+                raise ValueError("num_metrics required for tr_type='morbo'")
+            return MorboTrustRegion(
+                num_dim=num_dim,
+                num_arms=num_arms,
+                num_metrics=effective_num_metrics,
+                rng=rng,
+            )
+        else:
+            raise ValueError(f"Unknown tr_type: {self._config.tr_type!r}")
 
     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:
         return None
 
@@ -55,6 +90,11 @@ class BaseTurboImpl:
         init_idx: int,
         num_init: int,
     ) -> tuple[bool, int]:
+        if self._config.tr_type == "morbo":
+            x_obs_list.clear()
+            y_obs_list.clear()
+            yvar_obs_list.clear()
+            return True, 0
         return False, init_idx
 
     def prepare_ask(
@@ -76,20 +116,26 @@ class BaseTurboImpl:
         rng: Generator,
         fallback_fn: Callable[[np.ndarray, int], np.ndarray],
         from_unit_fn: Callable[[np.ndarray], np.ndarray],
+        tr_state: Any = None,
     ) -> np.ndarray:
         raise NotImplementedError("Subclasses must implement select_candidates")
 
     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,
     ) -> None:
         import numpy as np
 
+        x_obs_array = np.asarray(x_obs_list, dtype=float)
         y_obs_array = np.asarray(y_obs_list, dtype=float)
-        tr_state.update(y_obs_array)
+        if hasattr(tr_state, "update_xy"):
+            tr_state.update_xy(x_obs_array, y_obs_array, k=k)
+        else:
+            tr_state.update(y_obs_array)
 
     def estimate_y(self, x_unit: np.ndarray, y_observed: np.ndarray) -> np.ndarray:
         return y_observed

enn/turbo/lhd_only_impl.py

@@ -7,14 +7,19 @@ if TYPE_CHECKING:
     from numpy.random import Generator
 
 from .base_turbo_impl import BaseTurboImpl
+from .turbo_config import LHDOnlyConfig
 
 
 class LHDOnlyImpl(BaseTurboImpl):
+    def __init__(self, config: LHDOnlyConfig) -> None:
+        super().__init__(config)
+
     def get_x_center(
         self,
         x_obs_list: list,
         y_obs_list: list,
         rng: Generator,
+        tr_state: Any = None,
     ) -> np.ndarray | None:
         return None
 
@@ -26,6 +31,7 @@ class LHDOnlyImpl(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:
         from .turbo_utils import latin_hypercube
 
@@ -35,6 +41,7 @@ class LHDOnlyImpl(BaseTurboImpl):
     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,