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

enn/turbo/components/gp_surrogate.py

@@ -0,0 +1,144 @@
+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
+
+
+class GPSurrogate:
+    def __init__(self) -> None:
+        self._model: Any | None = None
+        self._y_mean: float | Any = 0.0
+        self._y_std: float | Any = 1.0
+        self._lengthscales: np.ndarray | None = None
+
+    @property
+    def lengthscales(self) -> np.ndarray | None:
+        return self._lengthscales
+
+    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 ..turbo_gp_fit import fit_gp
+
+        x_obs = np.asarray(x_obs, dtype=float)
+        y_obs = np.asarray(y_obs, dtype=float)
+        num_dim = x_obs.shape[1]
+        if y_obs.ndim == 2 and y_obs.shape[1] == 1:
+            y_obs = y_obs.ravel()
+        gp_result = fit_gp(
+            x_obs.tolist(),
+            y_obs.tolist(),
+            num_dim,
+            yvar_obs_list=y_var.ravel().tolist() if y_var is not None else None,
+            num_steps=num_steps,
+        )
+        self._model = gp_result.model
+        if gp_result.y_mean is not None:
+            self._y_mean = gp_result.y_mean
+        if gp_result.y_std is not None:
+            self._y_std = gp_result.y_std
+        lengthscales = None
+        if self._model is not None:
+            lengthscale = (
+                self._model.covar_module.base_kernel.lengthscale.cpu().detach().numpy()
+            )
+            if lengthscale.ndim == 3:
+                lengthscale = lengthscale.mean(axis=0)
+            lengthscales = lengthscale.ravel()
+            lengthscales_stabilized = lengthscales / lengthscales.mean()
+            del lengthscales
+            lengthscales_geom_normed = lengthscales_stabilized / np.prod(
+                np.power(lengthscales_stabilized, 1.0 / len(lengthscales_stabilized))
+            )
+            self._lengthscales = lengthscales_geom_normed
+        return SurrogateResult(model=self._model, lengthscales=self._lengthscales)
+
+    def _as_2d(self, a: np.ndarray) -> np.ndarray:
+        a = np.asarray(a, dtype=float)
+        if a.ndim == 1:
+            return a.reshape(-1, 1)
+        if a.ndim == 2:
+            return a.T
+        raise ValueError(a.shape)
+
+    def _unstandardize(self, y_std_2d: np.ndarray) -> np.ndarray:
+        y_std_2d = np.asarray(y_std_2d, dtype=float)
+        if y_std_2d.ndim != 2:
+            raise ValueError(y_std_2d.shape)
+        y_mean = np.asarray(self._y_mean, dtype=float).reshape(-1)
+        y_std = np.asarray(self._y_std, dtype=float).reshape(-1)
+        num_metrics = y_std_2d.shape[1]
+        if y_mean.size == 1 and num_metrics != 1:
+            y_mean = np.full(num_metrics, float(y_mean[0]), dtype=float)
+        if y_std.size == 1 and num_metrics != 1:
+            y_std = np.full(num_metrics, float(y_std[0]), dtype=float)
+        return y_mean.reshape(1, -1) + y_std.reshape(1, -1) * y_std_2d
+
+    def predict(self, x: np.ndarray) -> PosteriorResult:
+        import torch
+        from ..turbo_utils import get_gp_posterior_suppress_warning
+
+        if self._model is None:
+            raise RuntimeError("GPSurrogate.predict requires a fitted model")
+        x_torch = torch.as_tensor(x, dtype=torch.float64)
+        with torch.no_grad():
+            posterior = get_gp_posterior_suppress_warning(self._model, x_torch)
+            mu_std = posterior.mean.cpu().numpy()
+            var_std = posterior.variance.cpu().numpy()
+        mu = self._unstandardize(self._as_2d(mu_std))
+        sigma_std_2d = self._as_2d(var_std**0.5)
+        y_std = np.asarray(self._y_std, dtype=float).reshape(-1)
+        if y_std.size == 1 and sigma_std_2d.shape[1] != 1:
+            y_std = np.full(sigma_std_2d.shape[1], float(y_std[0]), dtype=float)
+        sigma = y_std.reshape(1, -1) * sigma_std_2d
+        return PosteriorResult(mu=mu, sigma=sigma)
+
+    def get_incumbent_candidate_indices(self, y_obs: np.ndarray) -> np.ndarray:
+        return np.arange(len(y_obs), dtype=int)
+
+    def sample(self, x: np.ndarray, num_samples: int, rng: Generator) -> np.ndarray:
+        import gpytorch
+        import torch
+        from ..turbo_utils import torch_seed_context
+
+        if self._model is None:
+            raise RuntimeError("GPSurrogate.sample requires a fitted model")
+        x_torch = torch.as_tensor(x, dtype=torch.float64)
+        seed = int(rng.integers(2**31 - 1))
+        with (
+            torch.no_grad(),
+            gpytorch.settings.fast_pred_var(),
+            torch_seed_context(seed, device=x_torch.device),
+        ):
+            posterior = self._model.posterior(x_torch)
+            samples = posterior.sample(sample_shape=torch.Size([num_samples]))
+        samples_np = samples.detach().cpu().numpy()
+        num_candidates = len(x)
+        num_metrics = len(self._y_mean) if hasattr(self._y_mean, "__len__") else 1
+        if samples_np.ndim == 2:
+            samples_np = samples_np[:, :, np.newaxis]
+        else:
+            assert samples_np.shape == (num_samples, num_metrics, num_candidates), (
+                f"GP raw samples shape mismatch: got {samples_np.shape}, "
+                f"expected ({num_samples}, {num_metrics}, {num_candidates})"
+            )
+            samples_np = np.transpose(samples_np, (0, 2, 1))
+        assert samples_np.shape == (num_samples, num_candidates, num_metrics), (
+            f"GP samples shape after transpose: got {samples_np.shape}, "
+            f"expected ({num_samples}, {num_candidates}, {num_metrics})"
+        )
+        y_mean = np.asarray(self._y_mean, dtype=float).reshape(1, 1, -1)
+        y_std = np.asarray(self._y_std, dtype=float).reshape(1, 1, -1)
+        result = y_mean + y_std * samples_np
+        assert result.shape == (num_samples, num_candidates, num_metrics)
+        return result

enn/turbo/components/hnr_acq_optimizer.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Any
+import numpy as np
+from .thompson_acq_optimizer import ThompsonAcqOptimizer
+from .ucb_acq_optimizer import UCBAcqOptimizer
+
+if TYPE_CHECKING:
+    from numpy.random import Generator
+    from .protocols import Surrogate
+
+
+class HnRAcqOptimizer:
+    def __init__(
+        self,
+        base_optimizer: ThompsonAcqOptimizer | UCBAcqOptimizer,
+        num_iterations: int = 100,
+    ) -> None:
+        self._base = base_optimizer
+        self._num_iterations = num_iterations
+
+    def _score_fn_ucb(
+        self, x_pt: np.ndarray, surrogate: Surrogate, beta: float = 1.0
+    ) -> float:
+        posterior = surrogate.predict(x_pt.reshape(1, -1))
+        mu = float(posterior.mu[0, 0])
+        sigma = float(posterior.sigma[0, 0]) if posterior.sigma is not None else 0.0
+        return mu + beta * sigma
+
+    def _score_fn_thompson(
+        self, x_pt: np.ndarray, surrogate: Surrogate, seed: int
+    ) -> float:
+        fixed_rng = np.random.default_rng(seed)
+        samples = surrogate.sample(x_pt.reshape(1, -1), 1, fixed_rng)
+        return float(samples[0, 0, 0])
+
+    def _optimize_one_arm(
+        self,
+        x_start: np.ndarray,
+        num_dim: int,
+        rng: Generator,
+        score_fn,
+    ) -> np.ndarray:
+        x_current = x_start.copy()
+        current_score = score_fn(x_current)
+        for _ in range(self._num_iterations):
+            direction = rng.standard_normal(num_dim)
+            direction = direction / np.linalg.norm(direction)
+            step_size = rng.uniform(0.01, 0.1)
+            x_proposed = np.clip(x_current + step_size * direction, 0.0, 1.0)
+            proposed_score = score_fn(x_proposed)
+            if proposed_score > current_score:
+                x_current = x_proposed
+                current_score = proposed_score
+        return x_current
+
+    def select(
+        self,
+        x_cand: np.ndarray,
+        num_arms: int,
+        surrogate: Surrogate,
+        rng: Generator,
+        *,
+        tr_state: Any | None = None,
+    ) -> np.ndarray:
+        num_dim = x_cand.shape[1]
+        x_arms = np.zeros((num_arms, num_dim), dtype=float)
+        is_ucb = isinstance(self._base, UCBAcqOptimizer)
+        for arm_idx in range(num_arms):
+            start_idx = rng.integers(0, len(x_cand))
+            x_start = x_cand[start_idx]
+            if is_ucb:
+                beta = getattr(self._base, "_beta", 1.0)
+
+                def score_fn(x_pt):
+                    return self._score_fn_ucb(x_pt, surrogate, beta)
+            else:
+                seed = int(rng.integers(0, 2**31))
+
+                def score_fn(x_pt, s=seed):
+                    return self._score_fn_thompson(x_pt, surrogate, s)
+
+            x_arms[arm_idx] = self._optimize_one_arm(x_start, num_dim, rng, score_fn)
+        return x_arms

enn/turbo/components/incumbent_selector.py

@@ -0,0 +1,11 @@
+from .chebyshev_incumbent_selector import ChebyshevIncumbentSelector
+from .incumbent_selector_protocol import IncumbentSelector
+from .no_incumbent_selector import NoIncumbentSelector
+from .scalar_incumbent_selector import ScalarIncumbentSelector
+
+__all__ = [
+    "ChebyshevIncumbentSelector",
+    "IncumbentSelector",
+    "NoIncumbentSelector",
+    "ScalarIncumbentSelector",
+]

enn/turbo/components/incumbent_selector_protocol.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Protocol
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+
+
+class IncumbentSelector(Protocol):
+    def select(
+        self,
+        y_obs: np.ndarray,
+        mu_obs: np.ndarray | None,
+        rng: Generator,
+    ) -> int: ...
+    def reset(self, rng: Generator) -> None: ...

enn/turbo/components/no_incumbent_selector.py

@@ -0,0 +1,21 @@
+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 NoIncumbentSelector:
+    def select(
+        self,
+        y_obs: np.ndarray,
+        mu_obs: np.ndarray | None,
+        rng: Generator,
+    ) -> int:
+        return 0
+
+    def reset(self, rng: Generator) -> None:
+        pass

enn/turbo/components/no_surrogate.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+import numpy as np
+from .posterior_result import PosteriorResult
+from .surrogate_result import SurrogateResult
+
+if TYPE_CHECKING:
+    from numpy.random import Generator
+
+
+class NoSurrogate:
+    def __init__(self) -> None:
+        self._x_obs: np.ndarray | None = None
+        self._y_obs: np.ndarray | None = None
+
+    @property
+    def lengthscales(self) -> np.ndarray | None:
+        return getattr(self, "_lengthscales", 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:
+        self._x_obs = np.asarray(x_obs, dtype=float)
+        self._y_obs = np.asarray(y_obs, dtype=float)
+        if self._y_obs.ndim == 1:
+            self._y_obs = self._y_obs.reshape(-1, 1)
+        return SurrogateResult(model=None, lengthscales=None)
+
+    def predict(self, x: np.ndarray) -> PosteriorResult:
+        if self._x_obs is None or self._y_obs is None:
+            raise RuntimeError("NoSurrogate.predict requires fit() first")
+        x = np.asarray(x, dtype=float)
+        if np.array_equal(x, self._x_obs):
+            return PosteriorResult(mu=self._y_obs.copy(), sigma=None)
+        raise RuntimeError("NoSurrogate.predict only works for training points")
+
+    def get_incumbent_candidate_indices(self, y_obs: np.ndarray) -> np.ndarray:
+        return np.arange(len(y_obs), dtype=int)
+
+    def sample(self, x: np.ndarray, num_samples: int, rng: Generator) -> np.ndarray:
+        n = len(x)
+        num_metrics = self._y_obs.shape[1] if hasattr(self, "_y_obs") else 1
+        return rng.standard_normal((num_samples, n, num_metrics))

enn/turbo/components/pareto_acq_optimizer.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Any
+import numpy as np
+
+if TYPE_CHECKING:
+    from numpy.random import Generator
+    from .protocols import Surrogate
+
+
+class ParetoAcqOptimizer:
+    def select(
+        self,
+        x_cand: np.ndarray,
+        num_arms: int,
+        surrogate: Surrogate,
+        rng: Generator,
+        *,
+        tr_state: Any | None = None,
+    ) -> np.ndarray:
+        from enn.enn.enn_util import arms_from_pareto_fronts
+
+        posterior = surrogate.predict(x_cand)
+        mu = posterior.mu
+        se = posterior.sigma if posterior.sigma is not None else np.zeros_like(mu)
+        if mu.ndim == 2 and mu.shape[1] > 1:
+            from nds import ndomsort
+
+            n = mu.shape[0]
+            i_keep: list[int] = []
+            remaining_mask = np.ones(n, dtype=bool)
+            while len(i_keep) < num_arms and np.any(remaining_mask):
+                remaining_indices = np.where(remaining_mask)[0]
+                fronts = ndomsort.non_domin_sort(
+                    -mu[remaining_indices], only_front_indices=True
+                )
+                front_indices = remaining_indices[np.where(fronts == 0)[0]]
+                if len(i_keep) + len(front_indices) <= num_arms:
+                    i_keep.extend(front_indices.tolist())
+                    remaining_mask[front_indices] = False
+                else:
+                    needed = num_arms - len(i_keep)
+                    selected = rng.choice(front_indices, size=needed, replace=False)
+                    i_keep.extend(selected.tolist())
+                    break
+            return x_cand[np.array(i_keep, dtype=int)]
+        else:
+            mu_1d = mu[:, 0] if mu.ndim == 2 else mu
+            se_1d = se[:, 0] if se.ndim == 2 else se
+            return arms_from_pareto_fronts(x_cand, mu_1d, se_1d, num_arms, rng)

enn/turbo/components/posterior_result.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import numpy as np
+
+
+@dataclass
+class PosteriorResult:
+    mu: np.ndarray
+    sigma: np.ndarray | None = None

enn/turbo/components/protocols.py

@@ -0,0 +1,13 @@
+from .acquisition_optimizer_protocol import AcquisitionOptimizer
+from .posterior_result import PosteriorResult
+from .surrogate_protocol import Surrogate
+from .surrogate_result import SurrogateResult
+from .trust_region_protocol import TrustRegion
+
+__all__ = [
+    "AcquisitionOptimizer",
+    "PosteriorResult",
+    "Surrogate",
+    "SurrogateResult",
+    "TrustRegion",
+]

enn/turbo/components/random_acq_optimizer.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Any
+import numpy as np
+
+if TYPE_CHECKING:
+    from numpy.random import Generator
+    from .protocols import Surrogate
+
+
+class RandomAcqOptimizer:
+    def select(
+        self,
+        x_cand: np.ndarray,
+        num_arms: int,
+        surrogate: Surrogate,
+        rng: Generator,
+        *,
+        tr_state: Any | None = None,
+    ) -> np.ndarray:
+        idx = rng.choice(len(x_cand), size=num_arms, replace=False)
+        return x_cand[idx]

enn/turbo/components/scalar_incumbent_selector.py

@@ -0,0 +1,39 @@
+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 ScalarIncumbentSelector:
+    noise_aware: bool
+
+    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
+
+        y = np.asarray(y_obs, dtype=float)
+        if y.ndim == 2:
+            y = y[:, 0]
+        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."
+                )
+            mu = np.asarray(mu_obs, dtype=float)
+            if mu.ndim == 2:
+                mu = mu[:, 0]
+            return int(argmax_random_tie(mu, rng=rng))
+        return int(argmax_random_tie(y, rng=rng))
+
+    def reset(self, rng: Generator) -> None:
+        pass

enn/turbo/components/surrogate_protocol.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Any, Protocol
+from .posterior_result import PosteriorResult
+from .surrogate_result import SurrogateResult
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+
+
+class Surrogate(Protocol):
+    @property
+    def lengthscales(self) -> np.ndarray | 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: ...
+    def predict(self, x: np.ndarray) -> PosteriorResult: ...
+    def sample(self, x: np.ndarray, num_samples: int, rng: Generator) -> np.ndarray: ...
+    def find_x_center(
+        self,
+        x_obs: np.ndarray,
+        y_obs: np.ndarray,
+        tr_state: Any,
+        rng: Generator,
+    ) -> np.ndarray | None: ...
+    def get_incumbent_candidate_indices(self, y_obs: np.ndarray) -> np.ndarray: ...

enn/turbo/components/surrogate_result.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import numpy as np
+
+
+@dataclass
+class SurrogateResult:
+    model: Any
+    lengthscales: np.ndarray | None = None

enn/turbo/components/surrogates.py

@@ -0,0 +1,5 @@
+from .enn_surrogate import ENNSurrogate
+from .gp_surrogate import GPSurrogate
+from .no_surrogate import NoSurrogate
+
+__all__ = ["ENNSurrogate", "GPSurrogate", "NoSurrogate"]

enn/turbo/components/thompson_acq_optimizer.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Any
+import numpy as np
+
+if TYPE_CHECKING:
+    from numpy.random import Generator
+    from .protocols import Surrogate
+
+
+class ThompsonAcqOptimizer:
+    def select(
+        self,
+        x_cand: np.ndarray,
+        num_arms: int,
+        surrogate: Surrogate,
+        rng: Generator,
+        *,
+        tr_state: Any | None = None,
+    ) -> np.ndarray:
+        num_candidates = len(x_cand)
+        samples = surrogate.sample(x_cand, num_arms, rng)
+        assert samples.ndim == 3, f"samples.ndim={samples.ndim}, expected 3"
+        assert (
+            samples.shape[0] == num_arms
+        ), f"samples.shape[0]={samples.shape[0]}, expected num_arms={num_arms}"
+        assert (
+            samples.shape[1] == num_candidates
+        ), f"samples.shape[1]={samples.shape[1]}, expected num_candidates={num_candidates}"
+        num_metrics = samples.shape[2]
+        if tr_state is not None and hasattr(tr_state, "scalarize"):
+            indices = []
+            for i in range(num_arms):
+                sample_i = samples[i]
+                assert sample_i.shape == (num_candidates, num_metrics), (
+                    f"sample_i.shape={sample_i.shape}, "
+                    f"expected ({num_candidates}, {num_metrics})"
+                )
+                scores = tr_state.scalarize(sample_i, clip=False)
+                assert scores.shape == (
+                    num_candidates,
+                ), f"scores.shape={scores.shape}, expected ({num_candidates},)"
+                for prev_idx in indices:
+                    scores[prev_idx] = -np.inf
+                idx = np.argmax(scores)
+                indices.append(idx)
+            return x_cand[indices]
+        else:
+            arm_indices = np.argmax(samples[:, :, 0], axis=1)
+            return x_cand[arm_indices]

enn/turbo/components/trust_region_protocol.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Protocol
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.random import Generator
+
+
+class TrustRegion(Protocol):
+    @property
+    def length(self) -> float: ...
+    def compute_bounds(
+        self, x_center: np.ndarray, lengthscales: np.ndarray | None = None
+    ) -> tuple[np.ndarray, np.ndarray]: ...
+    def update(self, y_obs: np.ndarray, y_incumbent: np.ndarray) -> None: ...
+    def needs_restart(self) -> bool: ...
+    def restart(self) -> None: ...
+    def get_incumbent_indices(self, y: np.ndarray, rng: Generator) -> np.ndarray: ...
+    def get_incumbent_index(
+        self, y: np.ndarray, rng: Generator, mu: np.ndarray | None = None
+    ) -> int: ...
+    def get_incumbent_value(
+        self, y_obs: np.ndarray, rng: Generator, mu_obs: np.ndarray | None = None
+    ) -> np.ndarray: ...

enn/turbo/components/ucb_acq_optimizer.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Any
+import numpy as np
+
+if TYPE_CHECKING:
+    from numpy.random import Generator
+    from .protocols import Surrogate
+
+
+class UCBAcqOptimizer:
+    def __init__(self, beta: float = 1.0) -> None:
+        self._beta = beta
+
+    def select(
+        self,
+        x_cand: np.ndarray,
+        num_arms: int,
+        surrogate: Surrogate,
+        rng: Generator,
+        *,
+        tr_state: Any | None = None,
+    ) -> np.ndarray:
+        num_candidates = len(x_cand)
+        posterior = surrogate.predict(x_cand)
+        mu = posterior.mu
+        sigma = posterior.sigma if posterior.sigma is not None else np.zeros_like(mu)
+        assert mu.ndim == 2, f"mu.ndim={mu.ndim}, expected 2"
+        assert (
+            mu.shape[0] == num_candidates
+        ), f"mu.shape[0]={mu.shape[0]}, expected {num_candidates}"
+        num_metrics = mu.shape[1]
+        if tr_state is not None and hasattr(tr_state, "scalarize"):
+            ucb = mu + self._beta * sigma
+            assert ucb.shape == (
+                num_candidates,
+                num_metrics,
+            ), f"ucb.shape={ucb.shape}, expected ({num_candidates}, {num_metrics})"
+            scores = tr_state.scalarize(ucb, clip=False)
+            assert scores.shape == (
+                num_candidates,
+            ), f"scores.shape={scores.shape}, expected ({num_candidates},)"
+        else:
+            scores = mu[:, 0] + self._beta * sigma[:, 0]
+            assert scores.shape == (num_candidates,)
+        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]
+        return x_cand[idx]