ennbo 0.1.0__py3-none-any.whl

enn/__init__.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from .enn import EpistemicNearestNeighbors, enn_fit
+
+_LAZY_IMPORTS = ("TurboMode", "TurboOptimizer", "Turbo", "Telemetry")
+
+
+def _lazy_load(name: str):
+    from . import turbo
+
+    return getattr(turbo, name)
+
+
+def __getattr__(name: str):
+    if name in _LAZY_IMPORTS:
+        return _lazy_load(name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__: list[str] = [
+    "EpistemicNearestNeighbors",
+    "enn_fit",
+    *_LAZY_IMPORTS,
+]

enn/enn/__init__.py

@@ -0,0 +1,4 @@
+from .enn import EpistemicNearestNeighbors
+from .enn_fit import enn_fit
+
+__all__: list[str] = ["EpistemicNearestNeighbors", "enn_fit"]

enn/enn/enn.py

@@ -0,0 +1,229 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import numpy as np
+
+    from .enn_normal import ENNNormal
+    from .enn_params import ENNParams
+
+
+class EpistemicNearestNeighbors:
+    def __init__(
+        self,
+        train_x: np.ndarray | Any,
+        train_y: np.ndarray | Any,
+        train_yvar: np.ndarray | Any | None = None,
+    ) -> None:
+        import numpy as np
+
+        if train_x.ndim != 2:
+            raise ValueError(train_x.shape)
+        if train_y.ndim != 2:
+            raise ValueError(train_y.shape)
+        if train_x.shape[0] != train_y.shape[0]:
+            raise ValueError((train_x.shape, train_y.shape))
+        if train_yvar is not None:
+            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._num_obs, self._num_dim = self._train_x.shape
+        _, self._num_metrics = self._train_y.shape
+        self._eps_var = 1e-9
+        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)
+
+        self._index: Any | None = None
+        self._build_index()
+
+    @property
+    def train_x(self) -> np.ndarray:
+        return self._train_x
+
+    @property
+    def train_y(self) -> np.ndarray:
+        return self._train_y
+
+    @property
+    def train_yvar(self) -> np.ndarray | None:
+        return self._train_yvar
+
+    @property
+    def num_outputs(self) -> int:
+        return self._num_metrics
+
+    def __len__(self) -> int:
+        return self._num_obs
+
+    def _build_index(self) -> None:
+        import faiss
+        import numpy as np
+
+        if self._num_obs == 0:
+            return
+        x_f32 = self._train_x.astype(np.float32, copy=False)
+        index = faiss.IndexFlatL2(self._num_dim)
+        index.add(x_f32)
+        self._index = index
+
+    def posterior(
+        self,
+        x: np.ndarray | Any,
+        *,
+        params: ENNParams,
+        exclude_nearest: bool = False,
+        observation_noise: bool = False,
+    ) -> ENNNormal:
+        from .enn_normal import ENNNormal
+
+        post_batch = self.batch_posterior(
+            x,
+            [params],
+            exclude_nearest=exclude_nearest,
+            observation_noise=observation_noise,
+        )
+        mu = post_batch.mu[0]
+        se = post_batch.se[0]
+        return ENNNormal(mu, se)
+
+    def batch_posterior(
+        self,
+        x: np.ndarray | Any,
+        paramss: list[ENNParams],
+        *,
+        exclude_nearest: bool = False,
+        observation_noise: bool = False,
+    ) -> ENNNormal:
+        import numpy as np
+
+        from .enn_normal import ENNNormal
+
+        if x.ndim != 2:
+            raise ValueError(x.shape)
+        if x.shape[1] != self._num_dim:
+            raise ValueError(x.shape)
+        if len(paramss) == 0:
+            raise ValueError("paramss must be non-empty")
+        batch_size = x.shape[0]
+        num_params = len(paramss)
+        if len(self) == 0:
+            mu = np.zeros((num_params, batch_size, self._num_metrics), dtype=float)
+            se = np.ones((num_params, batch_size, self._num_metrics), dtype=float)
+            return ENNNormal(mu, se)
+        max_k = max(params.k for params in paramss)
+        if exclude_nearest:
+            if len(self) <= 1:
+                raise ValueError(len(self))
+            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:]
+        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
+        for i, params in enumerate(paramss):
+            k = min(params.k, available_k)
+            if k > dist2s_full.shape[1]:
+                raise RuntimeError(
+                    f"k={k} exceeds available columns={dist2s_full.shape[1]}"
+                )
+            if k == 0:
+                mu_all[i] = np.zeros((batch_size, self._num_metrics), dtype=float)
+                se_all[i] = np.ones((batch_size, self._num_metrics), dtype=float)
+                continue
+            dist2s = dist2s_full[:, :k]
+            idx = idx_full[:, :k]
+            y_neighbors = self._train_y[idx]
+
+            dist2s_expanded = dist2s[..., np.newaxis]
+            var_component = (
+                params.ale_homoscedastic_scale + params.epi_var_scale * dist2s_expanded
+            )
+            if self._train_yvar is not None:
+                yvar_neighbors = self._train_yvar[idx] / self._y_scale**2
+                var_component = var_component + yvar_neighbors
+            else:
+                yvar_neighbors = None
+
+            w = 1.0 / (self._eps_var + var_component)
+            norm = np.sum(w, axis=1)
+            mu_all[i] = np.sum(w * y_neighbors, axis=1) / norm
+            epistemic_var = 1.0 / norm
+            vvar = epistemic_var
+            if observation_noise:
+                vvar = vvar + params.ale_homoscedastic_scale
+                if yvar_neighbors is not None:
+                    ale_heteroscedastic = np.sum(w * yvar_neighbors, axis=1) / norm
+                    vvar = vvar + ale_heteroscedastic
+            vvar = np.maximum(vvar, self._eps_var)
+            se_all[i] = np.sqrt(vvar) * self._y_scale
+        return ENNNormal(mu_all, se_all)
+
+    def neighbors(
+        self,
+        x: np.ndarray | Any,
+        k: int,
+        *,
+        exclude_nearest: bool = False,
+    ) -> list[tuple[np.ndarray, np.ndarray]]:
+        import numpy as np
+
+        x = np.asarray(x, dtype=float)
+        if x.ndim == 1:
+            x = x[np.newaxis, :]
+        if x.ndim != 2:
+            raise ValueError(f"x must be 1D or 2D, got shape {x.shape}")
+        if x.shape[0] != 1:
+            raise ValueError(f"x must be a single point, got shape {x.shape}")
+        if x.shape[1] != self._num_dim:
+            raise ValueError(
+                f"x must have {self._num_dim} dimensions, got {x.shape[1]}"
+            )
+        if k < 0:
+            raise ValueError(f"k must be non-negative, got {k}")
+        if len(self) == 0:
+            return []
+        if exclude_nearest:
+            if len(self) <= 1:
+                raise ValueError(
+                    f"exclude_nearest=True requires at least 2 observations, got {len(self)}"
+                )
+            search_k = int(min(k + 1, len(self)))
+        else:
+            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:]
+        actual_k = min(k, len(idx_full[0]))
+        idx = idx_full[0, :actual_k]
+        result = []
+        for i in idx:
+            x_neighbor = self._train_x[i].copy()
+            y_neighbor = self._train_y[i].copy()
+            result.append((x_neighbor, y_neighbor))
+        return result

enn/enn/enn_fit.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+
+    from .enn import EpistemicNearestNeighbors
+    from .enn_params import ENNParams
+
+from .enn_util import standardize_y
+
+
+def subsample_loglik(
+    model: EpistemicNearestNeighbors | Any,
+    x: np.ndarray | Any,
+    y: np.ndarray | Any,
+    *,
+    paramss: list[ENNParams] | list[Any],
+    P: int = 10,
+    rng: Generator | Any,
+) -> 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))
+    if P <= 0:
+        raise ValueError(P)
+    if len(paramss) == 0:
+        raise ValueError("paramss must be non-empty")
+    n = x.shape[0]
+    if n == 0:
+        return [0.0] * len(paramss)
+    if len(model) <= 1:
+        return [0.0] * len(paramss)
+    P_actual = min(P, n)
+    if P_actual == n:
+        indices = np.arange(n, dtype=int)
+    else:
+        indices = rng.permutation(n)[:P_actual]
+    x_selected = x[indices]
+    y_selected = y[indices]
+    if not np.isfinite(y_selected).all():
+        return [0.0] * len(paramss)
+    post_batch = model.batch_posterior(
+        x_selected, paramss, exclude_nearest=True, observation_noise=True
+    )
+    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_params,
+        P_actual,
+    ):
+        raise ValueError((mu_batch.shape, se_batch.shape, (num_params, P_actual)))
+    _, y_std = standardize_y(y)
+    y_scaled = y_selected / y_std
+    mu_scaled = mu_batch / y_std
+    se_scaled = se_batch / y_std
+    result = []
+    for i in range(num_params):
+        mu_i = mu_scaled[i]
+        se_i = se_scaled[i]
+        if not np.isfinite(mu_i).all() or not np.isfinite(se_i).all():
+            result.append(0.0)
+            continue
+        if np.any(se_i <= 0.0):
+            result.append(0.0)
+            continue
+        diff = y_scaled - mu_i
+        var_scaled = se_i**2
+        log_term = np.log(2.0 * np.pi * var_scaled)
+        quad = diff**2 / var_scaled
+        loglik = -0.5 * np.sum(log_term + quad)
+        if not np.isfinite(loglik):
+            result.append(0.0)
+            continue
+        result.append(float(loglik))
+    return result
+
+
+def enn_fit(
+    model: EpistemicNearestNeighbors | Any,
+    *,
+    k: int,
+    num_fit_candidates: int,
+    num_fit_samples: int = 10,
+    rng: Generator | Any,
+    params_warm_start: ENNParams | Any | None = None,
+) -> ENNParams:
+    from .enn_params import ENNParams
+
+    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)
+    epi_var_scale_values = 10**epi_var_scale_log_values
+    ale_homoscedastic_log_values = rng.uniform(
+        log_min, log_max, size=num_fit_candidates
+    )
+    ale_homoscedastic_values = 10**ale_homoscedastic_log_values
+    paramss = [
+        ENNParams(
+            k=k,
+            epi_var_scale=float(epi_val),
+            ale_homoscedastic_scale=float(ale_val),
+        )
+        for epi_val, ale_val in zip(epi_var_scale_values, ale_homoscedastic_values)
+    ]
+    if params_warm_start is not None:
+        paramss.append(
+            ENNParams(
+                k=k,
+                epi_var_scale=params_warm_start.epi_var_scale,
+                ale_homoscedastic_scale=params_warm_start.ale_homoscedastic_scale,
+            )
+        )
+    if len(paramss) == 0:
+        return ENNParams(k=k, epi_var_scale=1.0, ale_homoscedastic_scale=0.0)
+    import numpy as np
+
+    logliks = subsample_loglik(
+        model, train_x, y, paramss=paramss, P=num_fit_samples, rng=rng
+    )
+    if len(logliks) == 0:
+        return paramss[0]
+    best_idx = int(np.argmax(logliks))
+    return paramss[best_idx]

enn/enn/enn_normal.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import numpy as np
+
+
+@dataclass
+class ENNNormal:
+    mu: np.ndarray
+    se: np.ndarray
+
+    def sample(
+        self,
+        num_samples: int,
+        rng,
+        clip=None,
+    ) -> np.ndarray:
+        import numpy as np
+
+        size = (*self.se.shape, num_samples)
+        eps = rng.normal(size=size)
+        if clip is not None:
+            eps = np.clip(eps, a_min=-clip, a_max=clip)
+        return np.expand_dims(self.mu, -1) + np.expand_dims(self.se, -1) * eps

enn/enn/enn_params.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class ENNParams:
+    k: int
+    epi_var_scale: float
+    ale_homoscedastic_scale: float

enn/enn/enn_util.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+
+
+def standardize_y(y: np.ndarray | list[float] | Any) -> tuple[float, float]:
+    import numpy as np
+
+    y_array = np.asarray(y, dtype=float)
+    center = float(np.median(y_array))
+    scale = float(np.std(y_array))
+    if not np.isfinite(scale) or scale <= 0.0:
+        scale = 1.0
+    return center, scale
+
+
+def calculate_sobol_indices(x: np.ndarray, y: np.ndarray) -> np.ndarray:
+    import numpy as np
+
+    if x.ndim != 2:
+        raise ValueError(f"x must be 2D, got shape {x.shape}")
+    n, d = x.shape
+    if d <= 0:
+        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)
+    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)
+    ranks = np.empty_like(order)
+    ranks[order, np.arange(d)[None, :]] = row_idx
+    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)
+    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
+
+
+def arms_from_pareto_fronts(
+    x_cand: np.ndarray | Any,
+    mu: np.ndarray | Any,
+    se: np.ndarray | Any,
+    num_arms: int,
+    rng: Generator | Any,
+) -> np.ndarray:
+    import numpy as np
+    from nds import ndomsort
+
+    if x_cand.ndim != 2:
+        raise ValueError(x_cand.shape)
+    if mu.shape != se.shape or mu.ndim != 1:
+        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))
+
+    i_keep: list[int] = []
+    for n_front in range(1 + int(idx_front.max())):
+        front_indices = np.where(idx_front == n_front)[0]
+        front_indices = front_indices[np.argsort(-mu[front_indices])]
+        if len(i_keep) + len(front_indices) <= 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
+
+    i_keep = np.array(i_keep)
+    return x_cand[i_keep[np.argsort(-mu[i_keep])]]

enn/turbo/__init__.py

@@ -0,0 +1,11 @@
+from .turbo_mode import TurboMode
+from .turbo_optimizer import Telemetry, TurboOptimizer
+
+Turbo = TurboOptimizer
+
+__all__ = [
+    "TurboMode",
+    "TurboOptimizer",
+    "Turbo",
+    "Telemetry",
+]

enn/turbo/base_turbo_impl.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+
+from .turbo_config import TurboConfig
+
+
+class BaseTurboImpl:
+    def __init__(self, config: TurboConfig) -> None:
+        self._config = config
+
+    def get_x_center(
+        self,
+        x_obs_list: list,
+        y_obs_list: list,
+        rng: Generator,
+    ) -> np.ndarray | None:
+        import numpy as np
+
+        from .turbo_utils import argmax_random_tie
+
+        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)
+        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 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],
+    ) -> np.ndarray | None:
+        return None
+
+    def handle_restart(
+        self,
+        x_obs_list: list,
+        y_obs_list: list,
+        yvar_obs_list: list,
+        init_idx: int,
+        num_init: int,
+    ) -> tuple[bool, int]:
+        return False, init_idx
+
+    def prepare_ask(
+        self,
+        x_obs_list: list,
+        y_obs_list: list,
+        yvar_obs_list: list,
+        num_dim: int,
+        gp_num_steps: int,
+        rng: Any | None = None,
+    ) -> tuple[Any, float | None, float | None, np.ndarray | None]:
+        return None, None, None, None
+
+    def select_candidates(
+        self,
+        x_cand: np.ndarray,
+        num_arms: int,
+        num_dim: int,
+        rng: Generator,
+        fallback_fn: Callable[[np.ndarray, int], np.ndarray],
+        from_unit_fn: Callable[[np.ndarray], np.ndarray],
+    ) -> np.ndarray:
+        raise NotImplementedError("Subclasses must implement select_candidates")
+
+    def update_trust_region(
+        self,
+        tr_state: Any,
+        y_obs_list: list,
+        x_center: np.ndarray | None = None,
+        k: int | None = None,
+    ) -> None:
+        import numpy as np
+
+        y_obs_array = np.asarray(y_obs_list, dtype=float)
+        tr_state.update(y_obs_array)
+
+    def estimate_y(self, x_unit: np.ndarray, y_observed: np.ndarray) -> np.ndarray:
+        return y_observed
+
+    def get_mu_sigma(self, x_unit: np.ndarray) -> tuple[np.ndarray, np.ndarray] | None:
+        return None

enn/turbo/lhd_only_impl.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+
+from .base_turbo_impl import BaseTurboImpl
+
+
+class LHDOnlyImpl(BaseTurboImpl):
+    def get_x_center(
+        self,
+        x_obs_list: list,
+        y_obs_list: list,
+        rng: Generator,
+    ) -> np.ndarray | None:
+        return None
+
+    def select_candidates(
+        self,
+        x_cand: np.ndarray,
+        num_arms: int,
+        num_dim: int,
+        rng: Generator,
+        fallback_fn: Callable[[np.ndarray, int], np.ndarray],
+        from_unit_fn: Callable[[np.ndarray], np.ndarray],
+    ) -> np.ndarray:
+        from .turbo_utils import latin_hypercube
+
+        unit = latin_hypercube(num_arms, num_dim, rng=rng)
+        return from_unit_fn(unit)
+
+    def update_trust_region(
+        self,
+        tr_state: Any,
+        y_obs_list: list,
+        x_center: np.ndarray | None = None,
+        k: int | None = None,
+    ) -> None:
+        pass