ennbo 0.1.0__py3-none-any.whl

enn/turbo/proposal.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+
+    from enn.enn import EpistemicNearestNeighbors
+    from enn.enn.enn_params import ENNParams
+
+    from .turbo_gp import TurboGP
+
+from .turbo_utils import gp_thompson_sample
+
+
+def mk_enn(
+    x_obs_list: list[float] | list[list[float]],
+    y_obs_list: list[float] | list[list[float]],
+    *,
+    yvar_obs_list: list[float] | None = None,
+    k: int,
+    num_fit_samples: int | None = None,
+    num_fit_candidates: int | None = None,
+    rng: Generator | Any | None = None,
+    params_warm_start: ENNParams | Any | None = None,
+) -> tuple[EpistemicNearestNeighbors | None, ENNParams | None]:
+    import numpy as np
+
+    from enn.enn import EpistemicNearestNeighbors
+    from enn.enn.enn_params import ENNParams
+
+    if len(x_obs_list) == 0:
+        return None, None
+    y_obs_array = np.asarray(y_obs_list, dtype=float)
+    if y_obs_array.size == 0:
+        return None, None
+
+    y = y_obs_array.reshape(-1, 1)
+    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)
+    else:
+        yvar = None
+    x_obs_array = np.asarray(x_obs_list, dtype=float)
+    enn_model = EpistemicNearestNeighbors(
+        x_obs_array,
+        y,
+        yvar,
+    )
+    if len(enn_model) == 0:
+        return None, None
+
+    fitted_params: ENNParams | None = None
+    if num_fit_samples is not None and rng is not None:
+        from enn.enn.enn_fit import enn_fit
+
+        fitted_params = enn_fit(
+            enn_model,
+            k=k,
+            num_fit_candidates=num_fit_candidates
+            if num_fit_candidates is not None
+            else 30,
+            num_fit_samples=num_fit_samples,
+            rng=rng,
+            params_warm_start=params_warm_start,
+        )
+    else:
+        fitted_params = ENNParams(k=k, epi_var_scale=1.0, ale_homoscedastic_scale=0.0)
+
+    return enn_model, fitted_params
+
+
+def select_uniform(
+    x_cand: np.ndarray,
+    num_arms: int,
+    num_dim: int,
+    rng: Generator | Any,
+    from_unit_fn: Callable[[np.ndarray], np.ndarray],
+) -> np.ndarray:
+    if x_cand.ndim != 2 or x_cand.shape[1] != num_dim:
+        raise ValueError(x_cand.shape)
+    if x_cand.shape[0] < num_arms:
+        raise ValueError((x_cand.shape[0], num_arms))
+    idx = rng.choice(x_cand.shape[0], size=num_arms, replace=False)
+    return from_unit_fn(x_cand[idx])
+
+
+def select_gp_thompson(
+    x_cand: np.ndarray,
+    num_arms: int,
+    x_obs_list: list[float] | list[list[float]],
+    y_obs_list: list[float] | list[list[float]],
+    num_dim: int,
+    gp_num_steps: int,
+    rng: Generator | Any,
+    gp_y_mean: float,
+    gp_y_std: float,
+    select_sobol_fn: Callable[[np.ndarray, int], np.ndarray],
+    from_unit_fn: Callable[[np.ndarray], np.ndarray],
+    *,
+    model: TurboGP | None = None,
+    new_gp_y_mean: float | None = None,
+    new_gp_y_std: float | None = None,
+) -> tuple[np.ndarray, float, float, TurboGP | None]:
+    from .turbo_utils import fit_gp
+
+    if len(x_obs_list) == 0:
+        return select_sobol_fn(x_cand, num_arms), gp_y_mean, gp_y_std, None
+    if model is None:
+        model, _likelihood, new_gp_y_mean, new_gp_y_std = fit_gp(
+            x_obs_list,
+            y_obs_list,
+            num_dim,
+            num_steps=gp_num_steps,
+        )
+    if model is None:
+        return select_sobol_fn(x_cand, num_arms), gp_y_mean, gp_y_std, None
+    if new_gp_y_mean is None:
+        new_gp_y_mean = gp_y_mean
+    if new_gp_y_std is None:
+        new_gp_y_std = gp_y_std
+    if x_cand.shape[0] < num_arms:
+        raise ValueError((x_cand.shape[0], num_arms))
+    idx = gp_thompson_sample(
+        model,
+        x_cand,
+        num_arms,
+        rng,
+        new_gp_y_mean,
+        new_gp_y_std,
+    )
+    return from_unit_fn(x_cand[idx]), new_gp_y_mean, new_gp_y_std, model

enn/turbo/turbo_config.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+
+@dataclass(frozen=True)
+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
+    num_fit_samples: int | None = None
+    num_fit_candidates: int | None = None
+    acq_type: Literal["thompson", "pareto", "ucb"] = "pareto"
+    local_only: bool = False
+
+    def __post_init__(self) -> None:
+        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}")

enn/turbo/turbo_enn_impl.py

@@ -0,0 +1,176 @@
+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
+from .turbo_config import TurboConfig
+
+
+class TurboENNImpl(BaseTurboImpl):
+    def __init__(self, config: TurboConfig) -> None:
+        super().__init__(config)
+        self._enn: Any | None = None
+        self._fitted_params: Any | None = None
+        self._fitted_n_obs: int = 0
+
+    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
+
+        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)
+        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)}"
+            )
+
+        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]
+
+        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
+
+    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]:
+        x_obs_list.clear()
+        y_obs_list.clear()
+        yvar_obs_list.clear()
+        return True, 0
+
+    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]:
+        from .proposal import mk_enn
+
+        k = self._config.k if self._config.k is not None else 10
+        self._enn, self._fitted_params = mk_enn(
+            x_obs_list,
+            y_obs_list,
+            yvar_obs_list=yvar_obs_list,
+            k=k,
+            num_fit_samples=self._config.num_fit_samples,
+            num_fit_candidates=self._config.num_fit_candidates,
+            rng=rng,
+            params_warm_start=self._fitted_params,
+        )
+        self._fitted_n_obs = len(x_obs_list)
+        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:
+        import numpy as np
+
+        from enn.enn.enn_params import ENNParams
+
+        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)
+
+        if self._fitted_params is not None:
+            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
+            )
+
+        posterior = self._enn.posterior(x_cand, params=params)
+        mu = posterior.mu[:, 0]
+        se = posterior.se[:, 0]
+
+        if acq_type == "pareto":
+            from enn.enn.enn_util import arms_from_pareto_fronts
+
+            x_arms = arms_from_pareto_fronts(x_cand, mu, se, num_arms, rng)
+        elif acq_type == "ucb":
+            scores = mu + se
+            shuffled_indices = rng.permutation(len(scores))
+            shuffled_scores = scores[shuffled_indices]
+            top_k_in_shuffled = np.argpartition(-shuffled_scores, num_arms - 1)[
+                :num_arms
+            ]
+            idx = shuffled_indices[top_k_in_shuffled]
+            x_arms = x_cand[idx]
+        elif acq_type == "thompson":
+            samples = posterior.sample(num_samples=1, rng=rng)
+            scores = samples[:, 0, 0]
+            shuffled_indices = rng.permutation(len(scores))
+            shuffled_scores = scores[shuffled_indices]
+            top_k_in_shuffled = np.argpartition(-shuffled_scores, num_arms - 1)[
+                :num_arms
+            ]
+            idx = shuffled_indices[top_k_in_shuffled]
+            x_arms = x_cand[idx]
+        else:
+            raise ValueError(f"Unknown acq_type: {acq_type}")
+
+        return from_unit_fn(x_arms)
+
+    def estimate_y(self, x_unit: np.ndarray, y_observed: np.ndarray) -> np.ndarray:
+        if self._enn is None or self._fitted_params is None:
+            return y_observed
+        posterior = self._enn.posterior(x_unit, params=self._fitted_params)
+        return posterior.mu[:, 0]
+
+    def get_mu_sigma(self, x_unit: np.ndarray) -> tuple[np.ndarray, np.ndarray] | None:
+        if self._enn is None:
+            return None
+        k = self._config.k if self._config.k is not None else 10
+        from enn.enn.enn_params import ENNParams
+
+        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
+            )
+        )
+        posterior = self._enn.posterior(x_unit, params=params, observation_noise=False)
+        mu = posterior.mu[:, 0]
+        sigma = posterior.se[:, 0]
+        return mu, sigma

enn/turbo/turbo_gp.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from .turbo_gp_base import TurboGPBase
+
+
+class TurboGP(TurboGPBase):
+    def __init__(
+        self,
+        train_x,
+        train_y,
+        likelihood,
+        lengthscale_constraint,
+        outputscale_constraint,
+        ard_dims: int,
+    ) -> None:
+        from gpytorch.kernels import MaternKernel, ScaleKernel
+        from gpytorch.means import ConstantMean
+
+        super().__init__(train_x, train_y, likelihood)
+        self.mean_module = ConstantMean()
+        base_kernel = MaternKernel(
+            nu=2.5,
+            ard_num_dims=ard_dims,
+            lengthscale_constraint=lengthscale_constraint,
+        )
+        self.covar_module = ScaleKernel(
+            base_kernel,
+            outputscale_constraint=outputscale_constraint,
+        )

enn/turbo/turbo_gp_base.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from gpytorch.distributions import MultivariateNormal
+
+
+def _get_exact_gp_base():
+    from gpytorch.models import ExactGP
+
+    return ExactGP
+
+
+class TurboGPBase(_get_exact_gp_base()):
+    mean_module: Any
+    covar_module: Any
+
+    def forward(self, x) -> MultivariateNormal:
+        from gpytorch.distributions import MultivariateNormal
+
+        mean_x = self.mean_module(x)
+        covar_x = self.covar_module(x)
+        return MultivariateNormal(mean_x, covar_x)
+
+    def posterior(self, x) -> MultivariateNormal:
+        return self(x)

enn/turbo/turbo_gp_noisy.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from .turbo_gp_base import TurboGPBase
+
+
+class TurboGPNoisy(TurboGPBase):
+    def __init__(
+        self,
+        train_x,
+        train_y,
+        train_y_var,
+        lengthscale_constraint,
+        outputscale_constraint,
+        ard_dims: int,
+        *,
+        learn_additional_noise: bool = True,
+    ) -> None:
+        from gpytorch.kernels import MaternKernel, ScaleKernel
+        from gpytorch.likelihoods import FixedNoiseGaussianLikelihood
+        from gpytorch.means import ConstantMean
+
+        likelihood = FixedNoiseGaussianLikelihood(
+            noise=train_y_var,
+            learn_additional_noise=learn_additional_noise,
+        )
+        super().__init__(train_x, train_y, likelihood)
+        self.mean_module = ConstantMean()
+        base_kernel = MaternKernel(
+            nu=2.5,
+            ard_num_dims=ard_dims,
+            lengthscale_constraint=lengthscale_constraint,
+        )
+        self.covar_module = ScaleKernel(
+            base_kernel,
+            outputscale_constraint=outputscale_constraint,
+        )

enn/turbo/turbo_mode.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from enum import Enum, auto
+
+
+class TurboMode(Enum):
+    TURBO_ONE = auto()
+    TURBO_ZERO = auto()
+    TURBO_ENN = auto()
+    LHD_ONLY = auto()

enn/turbo/turbo_mode_impl.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable, Protocol
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+
+
+class TurboModeImpl(Protocol):
+    def get_x_center(
+        self,
+        x_obs_list: list,
+        y_obs_list: list,
+        rng: Generator,
+    ) -> np.ndarray | None: ...
+
+    def needs_tr_list(self) -> bool: ...
+
+    def create_trust_region(self, num_dim: int, num_arms: int) -> 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],
+    ) -> np.ndarray | 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]: ...
+
+    def prepare_ask(
+        self,
+        x_obs_list: list,
+        y_obs_list: list,
+        yvar_obs_list: list,
+        num_dim: int,
+        gp_num_steps: int,
+        rng: Generator | Any | None = None,
+    ) -> tuple[Any, float | None, float | None, np.ndarray | 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: ...
+
+    def update_trust_region(
+        self,
+        tr_state: Any,
+        y_obs_list: list,
+        x_center: np.ndarray | None = None,
+        k: int | None = None,
+    ) -> None: ...
+
+    def estimate_y(self, x_unit: np.ndarray, y_observed: np.ndarray) -> np.ndarray: ...

enn/turbo/turbo_one_impl.py

@@ -0,0 +1,163 @@
+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
+from .turbo_config import TurboConfig
+from .turbo_utils import gp_thompson_sample
+
+
+class TurboOneImpl(BaseTurboImpl):
+    def __init__(self, config: TurboConfig) -> None:
+        super().__init__(config)
+        self._gp_model: Any | None = None
+        self._gp_y_mean: float = 0.0
+        self._gp_y_std: float = 1.0
+        self._fitted_n_obs: int = 0
+
+    def get_x_center(
+        self,
+        x_obs_list: list,
+        y_obs_list: list,
+        rng: Generator,
+    ) -> np.ndarray | None:
+        import numpy as np
+        import torch
+
+        from .turbo_utils import argmax_random_tie
+
+        if len(y_obs_list) == 0:
+            return None
+        if self._gp_model is None:
+            return super().get_x_center(x_obs_list, y_obs_list, rng)
+        if self._fitted_n_obs != len(x_obs_list):
+            raise RuntimeError(
+                f"GP fitted on {self._fitted_n_obs} obs but get_x_center called with {len(x_obs_list)}"
+            )
+
+        x_array = np.asarray(x_obs_list, dtype=float)
+        x_torch = torch.as_tensor(x_array, dtype=torch.float64)
+        with torch.no_grad():
+            posterior = self._gp_model.posterior(x_torch)
+            mu = posterior.mean.cpu().numpy().ravel()
+
+        best_idx = argmax_random_tie(mu, rng=rng)
+        return x_array[best_idx]
+
+    def needs_tr_list(self) -> bool:
+        return True
+
+    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:
+        if len(x_obs_list) == 0:
+            return get_init_lhd_points_fn(num_arms)
+        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]:
+        x_obs_list.clear()
+        y_obs_list.clear()
+        yvar_obs_list.clear()
+        return True, 0
+
+    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]:
+        import numpy as np
+
+        from .turbo_utils import fit_gp
+
+        if len(x_obs_list) == 0:
+            return None, None, None, None
+        self._gp_model, _likelihood, gp_y_mean_fitted, gp_y_std_fitted = fit_gp(
+            x_obs_list,
+            y_obs_list,
+            num_dim,
+            yvar_obs_list=yvar_obs_list if yvar_obs_list else None,
+            num_steps=gp_num_steps,
+        )
+        self._fitted_n_obs = len(x_obs_list)
+        if gp_y_mean_fitted is not None:
+            self._gp_y_mean = gp_y_mean_fitted
+        if gp_y_std_fitted is not None:
+            self._gp_y_std = gp_y_std_fitted
+        weights = None
+        if self._gp_model is not None:
+            weights = (
+                self._gp_model.covar_module.base_kernel.lengthscale.cpu()
+                .detach()
+                .numpy()
+                .ravel()
+            )
+            # First line helps stabilize second line.
+            weights = weights / weights.mean()
+            weights = weights / np.prod(np.power(weights, 1.0 / len(weights)))
+        return self._gp_model, gp_y_mean_fitted, gp_y_std_fitted, weights
+
+    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:
+        if self._gp_model is None:
+            return fallback_fn(x_cand, num_arms)
+
+        idx = gp_thompson_sample(
+            self._gp_model,
+            x_cand,
+            num_arms,
+            rng,
+            self._gp_y_mean,
+            self._gp_y_std,
+        )
+        return from_unit_fn(x_cand[idx])
+
+    def estimate_y(self, x_unit: np.ndarray, y_observed: np.ndarray) -> np.ndarray:
+        import torch
+
+        if self._gp_model is None:
+            return y_observed
+        x_torch = torch.as_tensor(x_unit, dtype=torch.float64)
+        with torch.no_grad():
+            posterior = self._gp_model.posterior(x_torch)
+            mu = posterior.mean.cpu().numpy().ravel()
+        return self._gp_y_mean + self._gp_y_std * mu
+
+    def get_mu_sigma(self, x_unit: np.ndarray) -> tuple[np.ndarray, np.ndarray] | None:
+        import torch
+
+        if self._gp_model is None:
+            return None
+        x_torch = torch.as_tensor(x_unit, dtype=torch.float64)
+        with torch.no_grad():
+            posterior = self._gp_model.posterior(x_torch)
+            mu_std = posterior.mean.cpu().numpy().ravel()
+            sigma_std = posterior.variance.cpu().numpy().ravel() ** 0.5
+        mu = self._gp_y_mean + self._gp_y_std * mu_std
+        sigma = self._gp_y_std * sigma_std
+        return mu, sigma