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

enn/enn/enn_like_protocol.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Any, Protocol
+
+if TYPE_CHECKING:
+    import numpy as np
+    from .enn_params import ENNParams, PosteriorFlags
+
+
+class ENNLike(Protocol):
+    _num_dim: int
+    _num_metrics: int
+    _x_scale: np.ndarray
+    _scale_x: bool
+    _enn_index: Any
+    _train_y: np.ndarray
+    _train_yvar: np.ndarray | None
+
+    def __len__(self) -> int: ...
+    def posterior(self, x: np.ndarray, *, params: ENNParams, flags: PosteriorFlags):
+        raise NotImplementedError
+
+    def _empty_posterior_internals(self, batch_size: int):
+        raise NotImplementedError
+
+    def _compute_weighted_stats(
+        self,
+        dist2s: np.ndarray,
+        y_neighbors: np.ndarray,
+        *,
+        yvar_neighbors: np.ndarray | None,
+        params: ENNParams,
+        observation_noise: bool,
+        y_scale: np.ndarray | None = None,
+    ):
+        raise NotImplementedError

enn/enn/enn_normal.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
-
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     import numpy as np
+    from numpy.random import Generator
 
 
 @dataclass
@@ -15,8 +15,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

@@ -1,10 +1,4 @@
-from __future__ import annotations
+from .enn_params_class import ENNParams
+from .posterior_flags import PosteriorFlags
 
-from dataclasses import dataclass
-
-
-@dataclass(frozen=True)
-class ENNParams:
-    k: int
-    epi_var_scale: float
-    ale_homoscedastic_scale: float
+__all__ = ["ENNParams", "PosteriorFlags"]

enn/enn/enn_params_class.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class ENNParams:
+    k_num_neighbors: int
+    epistemic_variance_scale: float
+    aleatoric_variance_scale: float
+
+    def __post_init__(self) -> None:
+        import numpy as np
+
+        k = int(self.k_num_neighbors)
+        if k <= 0:
+            raise ValueError(f"k_num_neighbors must be > 0, got {k}")
+        epi_var_scale = float(self.epistemic_variance_scale)
+        if not np.isfinite(epi_var_scale) or epi_var_scale < 0.0:
+            raise ValueError(
+                f"epistemic_variance_scale must be >= 0, got {epi_var_scale}"
+            )
+        ale_scale = float(self.aleatoric_variance_scale)
+        if not np.isfinite(ale_scale) or ale_scale < 0.0:
+            raise ValueError(f"aleatoric_variance_scale must be >= 0, got {ale_scale}")

enn/enn/enn_util.py

@@ -1,5 +1,4 @@
 from __future__ import annotations
-
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
@@ -18,9 +17,7 @@ def standardize_y(y: np.ndarray | list[float] | Any) -> tuple[float, float]:
     return center, scale
 
 
-def calculate_sobol_indices(x: np.ndarray, y: np.ndarray) -> np.ndarray:
-    import numpy as np
-
+def _validate_sobol_inputs(x, y):
     if x.ndim != 2:
         raise ValueError(f"x must be 2D, got shape {x.shape}")
     n, d = x.shape
@@ -28,16 +25,14 @@ def calculate_sobol_indices(x: np.ndarray, y: np.ndarray) -> np.ndarray:
         raise ValueError(f"x must have at least 1 dimension, got {d}")
     if y.ndim == 2 and y.shape[1] == 1:
         y = y.reshape(-1)
-    if y.ndim != 1:
-        raise ValueError(f"y must be 1D, got shape {y.shape}")
-    if y.shape[0] != n:
-        raise ValueError(f"y length {y.shape[0]} != x rows {n}")
-    if n < 9:
-        return np.ones(d, dtype=x.dtype)
-    mu = y.mean()
-    vy = y.var(ddof=0)
-    if not np.isfinite(vy) or vy <= 0:
-        return np.ones(d, dtype=x.dtype)
+    if y.ndim != 1 or y.shape[0] != n:
+        raise ValueError(f"y shape {y.shape} incompatible with x rows {n}")
+    return n, d, y
+
+
+def _compute_sobol_bins(x, y, n, d):
+    import numpy as np
+
     B = 10 if n >= 30 else 3
     order = np.argsort(x, axis=0)
     row_idx = np.arange(n).reshape(n, 1).repeat(d, axis=1)
@@ -46,17 +41,59 @@ def calculate_sobol_indices(x: np.ndarray, y: np.ndarray) -> np.ndarray:
     idx = (ranks * B) // n
     oh = np.zeros((n, d, B), dtype=x.dtype)
     oh[np.arange(n)[:, None], np.arange(d)[None, :], idx] = 1.0
-    counts = oh.sum(axis=0)
-    sums = (oh * y.reshape(n, 1, 1)).sum(axis=0)
+    counts, sums = oh.sum(axis=0), (oh * y.reshape(n, 1, 1)).sum(axis=0)
     mu_b = np.zeros_like(sums)
     mask = counts > 0
     mu_b[mask] = sums[mask] / counts[mask]
-    p_b = counts / float(n)
-    diff = mu_b - mu
-    S = (p_b * (diff * diff)).sum(axis=1) / vy
-    var_x = x.var(axis=0, ddof=0)
-    S = np.where(var_x <= 1e-12, np.zeros_like(S), S)
-    return S
+    return counts / float(n), mu_b
+
+
+def calculate_sobol_indices(x: np.ndarray, y: np.ndarray) -> np.ndarray:
+    import numpy as np
+
+    n, d, y = _validate_sobol_inputs(x, y)
+    if n < 9:
+        return np.ones(d, dtype=x.dtype)
+    mu, vy = y.mean(), y.var(ddof=0)
+    if not np.isfinite(vy) or vy <= 0:
+        return np.ones(d, dtype=x.dtype)
+    p_b, mu_b = _compute_sobol_bins(x, y, n, d)
+    S = (p_b * (mu_b - mu) ** 2).sum(axis=1) / vy
+    return np.where(x.var(axis=0, ddof=0) <= 1e-12, np.zeros_like(S), S)
+
+
+def pareto_front_2d_maximize(
+    a: np.ndarray | Any, b: np.ndarray | Any, idx: np.ndarray | Any | None = None
+) -> np.ndarray:
+    import numpy as np
+
+    a = np.asarray(a, dtype=float)
+    b = np.asarray(b, dtype=float)
+    if a.shape != b.shape or a.ndim != 1:
+        raise ValueError((a.shape, b.shape))
+    if idx is None:
+        idx = np.arange(a.size, dtype=int)
+    else:
+        idx = np.asarray(idx, dtype=int)
+        if idx.ndim != 1:
+            raise ValueError(idx.shape)
+    order = np.lexsort((-b[idx], -a[idx]))
+    sorted_idx = idx[order]
+    keep: list[int] = []
+    best_b = -float("inf")
+    last_a = float("nan")
+    last_b = float("nan")
+    for i in sorted_idx.tolist():
+        bi = float(b[i])
+        ai = float(a[i])
+        if bi > best_b:
+            keep.append(i)
+            best_b = bi
+            last_a = ai
+            last_b = bi
+        elif bi == best_b and ai == last_a and bi == last_b:
+            keep.append(i)
+    return np.asarray(keep, dtype=int)
 
 
 def arms_from_pareto_fronts(
@@ -67,7 +104,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 +111,28 @@ 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]))
-
-    combined = np.column_stack([mu, se])
-    idx_front = np.array(ndomsort.non_domin_sort(-combined, only_front_indices=True))
-
+    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")
     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/enn/neighbor_data.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import numpy as np
+
+
+@dataclass(frozen=True)
+class NeighborData:
+    dist2s: np.ndarray
+    idx: np.ndarray
+    y_neighbors: np.ndarray
+    k: int

enn/enn/neighbors.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import numpy as np
+
+
+@dataclass(frozen=True)
+class Neighbors:
+    dist2: np.ndarray
+    ids: np.ndarray
+    y: np.ndarray
+    yvar: np.ndarray | None

enn/enn/posterior_flags.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class PosteriorFlags:
+    exclude_nearest: bool = False
+    observation_noise: bool = False

enn/enn/weighted_stats.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import numpy as np
+
+
+@dataclass(frozen=True)
+class WeightedStats:
+    w_normalized: np.ndarray
+    l2: np.ndarray
+    mu: np.ndarray
+    se: np.ndarray

enn/turbo/components/__init__.py

@@ -0,0 +1,41 @@
+from .acquisition import (
+    HnRAcqOptimizer,
+    ParetoAcqOptimizer,
+    RandomAcqOptimizer,
+    ThompsonAcqOptimizer,
+    UCBAcqOptimizer,
+)
+from .incumbent_selector import (
+    ChebyshevIncumbentSelector,
+    IncumbentSelector,
+    NoIncumbentSelector,
+    ScalarIncumbentSelector,
+)
+from .protocols import (
+    AcquisitionOptimizer,
+    PosteriorResult,
+    Surrogate,
+    SurrogateResult,
+    TrustRegion,
+)
+from .surrogates import ENNSurrogate, GPSurrogate, NoSurrogate
+
+__all__ = [
+    "AcquisitionOptimizer",
+    "ChebyshevIncumbentSelector",
+    "ENNSurrogate",
+    "GPSurrogate",
+    "HnRAcqOptimizer",
+    "IncumbentSelector",
+    "NoIncumbentSelector",
+    "NoSurrogate",
+    "ParetoAcqOptimizer",
+    "PosteriorResult",
+    "RandomAcqOptimizer",
+    "ScalarIncumbentSelector",
+    "Surrogate",
+    "SurrogateResult",
+    "ThompsonAcqOptimizer",
+    "TrustRegion",
+    "UCBAcqOptimizer",
+]

enn/turbo/components/acquisition.py

@@ -0,0 +1,13 @@
+from .hnr_acq_optimizer import HnRAcqOptimizer
+from .pareto_acq_optimizer import ParetoAcqOptimizer
+from .random_acq_optimizer import RandomAcqOptimizer
+from .thompson_acq_optimizer import ThompsonAcqOptimizer
+from .ucb_acq_optimizer import UCBAcqOptimizer
+
+__all__ = [
+    "HnRAcqOptimizer",
+    "ParetoAcqOptimizer",
+    "RandomAcqOptimizer",
+    "ThompsonAcqOptimizer",
+    "UCBAcqOptimizer",
+]

enn/turbo/components/acquisition_optimizer_protocol.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Any, Protocol
+from .surrogate_protocol import Surrogate
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+
+
+class AcquisitionOptimizer(Protocol):
+    def select(
+        self,
+        x_cand: np.ndarray,
+        num_arms: int,
+        surrogate: Surrogate,
+        rng: Generator,
+        *,
+        tr_state: Any | None = None,
+    ) -> np.ndarray: ...

enn/turbo/components/builder.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+from .acquisition import HnRAcqOptimizer, ThompsonAcqOptimizer, UCBAcqOptimizer
+
+if TYPE_CHECKING:
+    from ..config.optimizer_config import OptimizerConfig
+    from .protocols import AcquisitionOptimizer, Surrogate
+
+
+def build_surrogate(config: OptimizerConfig) -> Surrogate:
+    return config.surrogate.build()
+
+
+def build_acquisition_optimizer(config: OptimizerConfig) -> AcquisitionOptimizer:
+    from ..config.acquisition import HnROptimizerConfig
+
+    base = config.acquisition.build()
+    if isinstance(config.acq_optimizer, HnROptimizerConfig):
+        if isinstance(base, (ThompsonAcqOptimizer, UCBAcqOptimizer)):
+            return HnRAcqOptimizer(base)
+        raise ValueError(f"HnR not supported with {type(base).__name__}")
+    return base

enn/turbo/components/chebyshev_incumbent_selector.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+
+
+@dataclass
+class ChebyshevIncumbentSelector:
+    num_metrics: int
+    noise_aware: bool
+    alpha: float
+    _weights: np.ndarray | None = None
+
+    def __post_init__(self) -> None:
+        if self.num_metrics < 1:
+            raise ValueError(f"num_metrics must be >= 1, got {self.num_metrics}")
+
+    @property
+    def weights(self) -> np.ndarray | None:
+        return self._weights
+
+    def _sample_weights(self, rng: Generator) -> None:
+        import numpy as np
+
+        alpha = np.ones(self.num_metrics, dtype=float)
+        self._weights = np.asarray(rng.dirichlet(alpha), dtype=float)
+
+    def reset(self, rng: Generator) -> None:
+        self._sample_weights(rng)
+
+    def select(
+        self,
+        y_obs: np.ndarray,
+        mu_obs: np.ndarray | None,
+        rng: Generator,
+    ) -> int:
+        import numpy as np
+        from ..turbo_utils import argmax_random_tie
+
+        if self._weights is None:
+            self._sample_weights(rng)
+        y = np.asarray(y_obs, dtype=float)
+        if y.ndim != 2 or y.shape[1] != self.num_metrics:
+            raise ValueError(
+                f"Expected y with {self.num_metrics} metrics, got {y.shape}"
+            )
+        if self.noise_aware:
+            if mu_obs is None:
+                raise ValueError(
+                    "noise_aware=True requires a surrogate that provides mu. "
+                    "Either use a GP/ENN surrogate or set noise_aware=False."
+                )
+            values = np.asarray(mu_obs, dtype=float)
+        else:
+            values = y
+        scores = self._scalarize(values)
+        return int(argmax_random_tie(scores, rng=rng))
+
+    def _scalarize(self, values: np.ndarray) -> np.ndarray:
+        import numpy as np
+
+        if self._weights is None:
+            raise RuntimeError("Weights not initialized; call reset() first")
+        v_min = values.min(axis=0)
+        v_max = values.max(axis=0)
+        denom = v_max - v_min
+        is_deg = denom <= 0.0
+        denom_safe = np.where(is_deg, 1.0, denom)
+        z = (values - v_min.reshape(1, -1)) / denom_safe.reshape(1, -1)
+        z = np.where(is_deg, 0.5, z)
+        z = np.clip(z, 0.0, 1.0)
+        t = z * self._weights.reshape(1, -1)
+        return np.min(t, axis=1) + self.alpha * np.sum(t, axis=1)

enn/turbo/components/enn_surrogate.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Any
+import numpy as np
+from .posterior_result import PosteriorResult
+from .surrogate_result import SurrogateResult
+
+if TYPE_CHECKING:
+    from numpy.random import Generator
+    from ..config.surrogate import ENNSurrogateConfig
+
+
+class ENNSurrogate:
+    def __init__(self, config: ENNSurrogateConfig) -> None:
+        self._config = config
+        self._enn: Any | None = None
+        self._params: Any | None = None
+
+    @property
+    def lengthscales(self) -> np.ndarray | None:
+        return None
+
+    def fit(
+        self,
+        x_obs: np.ndarray,
+        y_obs: np.ndarray,
+        y_var: np.ndarray | None = None,
+        *,
+        num_steps: int = 0,
+        rng: Generator | None = None,
+    ) -> SurrogateResult:
+        from ..proposal import mk_enn
+        from ..config.enums import ENNIndexDriver
+
+        k = self._config.k if self._config.k is not None else 10
+        if (
+            self._config.index_driver == ENNIndexDriver.HNSW
+            and self._enn is not None
+            and len(x_obs) > len(self._enn)
+        ):
+            n_old = len(self._enn)
+            new_x = x_obs[n_old:]
+            new_y = y_obs[n_old:]
+            new_yvar = y_var[n_old:] if y_var is not None else None
+            self._enn.add(new_x, new_y, new_yvar)
+            if self._config.num_fit_samples is not None and rng is not None:
+                from ...enn.enn_fit import enn_fit
+
+                self._params = enn_fit(
+                    self._enn,
+                    k=k,
+                    num_fit_candidates=self._config.num_fit_candidates
+                    if self._config.num_fit_candidates is not None
+                    else 30,
+                    num_fit_samples=self._config.num_fit_samples,
+                    rng=rng,
+                    params_warm_start=self._params,
+                )
+        else:
+            self._enn, self._params = mk_enn(
+                list(x_obs),
+                list(y_obs),
+                k,
+                list(y_var) if y_var is not None else [],
+                num_fit_samples=self._config.num_fit_samples,
+                num_fit_candidates=self._config.num_fit_candidates,
+                scale_x=self._config.scale_x,
+                index_driver=self._config.index_driver,
+                rng=rng,
+                params_warm_start=self._params,
+            )
+        return SurrogateResult(model=self._enn, lengthscales=None)
+
+    def get_incumbent_candidate_indices(self, y_obs: np.ndarray) -> np.ndarray:
+        y_array = np.asarray(y_obs, dtype=float)
+        k = self._config.k
+        if k is None:
+            num_fit_candidates = (
+                self._config.num_fit_candidates
+                if self._config.num_fit_candidates is not None
+                else 100
+            )
+            k = min(len(y_array), max(10, 2 * num_fit_candidates))
+        if y_array.ndim == 2 and y_array.shape[1] > 1:
+            num_top = min(k, len(y_array))
+            union_indices: set[int] = set()
+            for m in range(y_array.shape[1]):
+                top_m = np.argpartition(-y_array[:, m], num_top - 1)[:num_top]
+                union_indices.update(top_m.tolist())
+            return np.array(sorted(union_indices), dtype=int)
+        else:
+            y_flat = y_array[:, 0] if y_array.ndim == 2 else y_array
+            num_top = min(k, len(y_flat))
+            return np.argpartition(-y_flat, num_top - 1)[:num_top]
+
+    def predict(self, x: np.ndarray) -> PosteriorResult:
+        if self._enn is None or self._params is None:
+            raise RuntimeError("ENNSurrogate.predict requires a fitted model")
+        posterior = self._enn.posterior(x, params=self._params)
+        return PosteriorResult(mu=posterior.mu, sigma=posterior.se)
+
+    def sample(self, x: np.ndarray, num_samples: int, rng: Generator) -> np.ndarray:
+        if self._enn is None or self._params is None:
+            raise RuntimeError("ENNSurrogate.sample requires a fitted model")
+        num_candidates = len(x)
+        num_metrics = self._enn.num_outputs
+        base_seed = rng.integers(0, 2**31)
+        function_seeds = np.arange(base_seed, base_seed + num_samples, dtype=np.int64)
+        samples = self._enn.posterior_function_draw(
+            x, self._params, function_seeds=function_seeds
+        )
+        assert samples.shape == (num_samples, num_candidates, num_metrics), (
+            f"ENN samples shape mismatch: got {samples.shape}, "
+            f"expected ({num_samples}, {num_candidates}, {num_metrics})"
+        )
+        return samples