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

enn/turbo/turbo_utils.py

@@ -24,15 +24,23 @@ def _next_power_of_2(n: int) -> int:
 
 
 @contextlib.contextmanager
-def torch_rng_context(generator: torch.Generator | Any) -> Iterator[None]:
+def torch_seed_context(
+    seed: int, device: torch.device | Any | None = None
+) -> Iterator[None]:
     import torch
 
-    old_state = torch.get_rng_state()
-    try:
-        torch.set_rng_state(generator.get_state())
+    devices: list[int] | None = None
+    if device is not None and getattr(device, "type", None) == "cuda":
+        idx = 0 if getattr(device, "index", None) is None else int(device.index)
+        devices = [idx]
+    with torch.random.fork_rng(devices=devices, enabled=True):
+        torch.manual_seed(int(seed))
+        if device is not None and getattr(device, "type", None) == "cuda":
+            torch.cuda.manual_seed_all(int(seed))
+        if device is not None and getattr(device, "type", None) == "mps":
+            if hasattr(torch, "mps") and hasattr(torch.mps, "manual_seed"):
+                torch.mps.manual_seed(int(seed))
         yield
-    finally:
-        torch.set_rng_state(old_state)
 
 
 def fit_gp(
@@ -45,8 +53,8 @@ def fit_gp(
 ) -> tuple[
     "TurboGP | TurboGPNoisy | None",
     "GaussianLikelihood | None",
-    float,
-    float,
+    float | np.ndarray,
+    float | np.ndarray,
 ]:
     import numpy as np
     import torch
@@ -60,22 +68,40 @@ def fit_gp(
     x = np.asarray(x_obs_list, dtype=float)
     y = np.asarray(y_obs_list, dtype=float)
     n = x.shape[0]
+    if y.ndim not in (1, 2):
+        raise ValueError(y.shape)
+    is_multi_output = y.ndim == 2 and y.shape[1] > 1
     if yvar_obs_list is not None:
         if len(yvar_obs_list) != len(y_obs_list):
             raise ValueError(
                 f"yvar_obs_list length {len(yvar_obs_list)} != y_obs_list length {len(y_obs_list)}"
             )
+        if is_multi_output:
+            raise ValueError("yvar_obs_list not supported for multi-output GP")
     if n == 0:
+        if is_multi_output:
+            num_outputs = int(y.shape[1])
+            return None, None, np.zeros(num_outputs), np.ones(num_outputs)
         return None, None, 0.0, 1.0
-    if n == 1:
-        gp_y_mean = float(y[0])
-        gp_y_std = 1.0
+    if n == 1 and is_multi_output:
+        gp_y_mean = y[0].copy()
+        gp_y_std = np.ones(int(y.shape[1]), dtype=float)
         return None, None, gp_y_mean, gp_y_std
-    gp_y_mean, gp_y_std = standardize_y(y)
-    y_centered = y - gp_y_mean
-    z = y_centered / gp_y_std
+
+    if is_multi_output:
+        gp_y_mean = y.mean(axis=0)
+        gp_y_std = y.std(axis=0)
+        gp_y_std = np.where(gp_y_std < 1e-6, 1.0, gp_y_std)
+        z = (y - gp_y_mean) / gp_y_std
+    else:
+        gp_y_mean, gp_y_std = standardize_y(y)
+        y_centered = y - gp_y_mean
+        z = y_centered / gp_y_std
     train_x = torch.as_tensor(x, dtype=torch.float64)
-    train_y = torch.as_tensor(z, dtype=torch.float64)
+    if is_multi_output:
+        train_y = torch.as_tensor(z.T, dtype=torch.float64)
+    else:
+        train_y = torch.as_tensor(z, dtype=torch.float64)
     lengthscale_constraint = Interval(0.005, 2.0)
     outputscale_constraint = Interval(0.05, 20.0)
     if yvar_obs_list is not None:
@@ -92,9 +118,16 @@ def fit_gp(
         likelihood = model.likelihood
     else:
         noise_constraint = Interval(5e-4, 0.2)
-        likelihood = GaussianLikelihood(noise_constraint=noise_constraint).to(
-            dtype=train_y.dtype
-        )
+        if is_multi_output:
+            num_outputs = int(y.shape[1])
+            likelihood = GaussianLikelihood(
+                noise_constraint=noise_constraint,
+                batch_shape=torch.Size([num_outputs]),
+            ).to(dtype=train_y.dtype)
+        else:
+            likelihood = GaussianLikelihood(noise_constraint=noise_constraint).to(
+                dtype=train_y.dtype
+            )
         model = TurboGP(
             train_x=train_x,
             train_y=train_y,
@@ -103,11 +136,23 @@ def fit_gp(
             outputscale_constraint=outputscale_constraint,
             ard_dims=num_dim,
         ).to(dtype=train_x.dtype)
-        likelihood.noise = torch.tensor(0.005, dtype=train_y.dtype)
-    model.covar_module.outputscale = torch.tensor(1.0, dtype=train_x.dtype)
-    model.covar_module.base_kernel.lengthscale = torch.full(
-        (num_dim,), 0.5, dtype=train_x.dtype
-    )
+        if is_multi_output:
+            likelihood.noise = torch.full(
+                (int(y.shape[1]),), 0.005, dtype=train_y.dtype
+            )
+        else:
+            likelihood.noise = torch.tensor(0.005, dtype=train_y.dtype)
+    if is_multi_output:
+        num_outputs = int(y.shape[1])
+        model.covar_module.outputscale = torch.ones(num_outputs, dtype=train_x.dtype)
+        model.covar_module.base_kernel.lengthscale = torch.full(
+            (num_outputs, 1, num_dim), 0.5, dtype=train_x.dtype
+        )
+    else:
+        model.covar_module.outputscale = torch.tensor(1.0, dtype=train_x.dtype)
+        model.covar_module.base_kernel.lengthscale = torch.full(
+            (num_dim,), 0.5, dtype=train_x.dtype
+        )
     model.train()
     likelihood.train()
     mll = ExactMarginalLogLikelihood(likelihood, model)
@@ -116,6 +161,8 @@ def fit_gp(
         optimizer.zero_grad()
         output = model(train_x)
         loss = -mll(output, train_y)
+        if loss.ndim != 0:
+            loss = loss.sum()
         loss.backward()
         optimizer.step()
     model.eval()
@@ -197,6 +244,49 @@ def raasp(
     )
 
 
+def generate_raasp_candidates(
+    center: np.ndarray | Any,
+    lb: np.ndarray | list[float] | Any,
+    ub: np.ndarray | list[float] | Any,
+    num_candidates: int,
+    *,
+    rng: Generator | Any,
+    sobol_engine: QMCEngine | Any,
+    num_pert: int = 20,
+) -> np.ndarray:
+    if num_candidates <= 0:
+        raise ValueError(num_candidates)
+    return raasp(
+        center,
+        lb,
+        ub,
+        num_candidates,
+        num_pert=num_pert,
+        rng=rng,
+        sobol_engine=sobol_engine,
+    )
+
+
+def generate_trust_region_candidates(
+    x_center: np.ndarray | Any,
+    lengthscales: np.ndarray | None,
+    num_candidates: int,
+    *,
+    compute_bounds_1d: Any,
+    rng: Generator | Any,
+    sobol_engine: QMCEngine | Any,
+) -> np.ndarray:
+    """
+    Small DRY helper for trust-region candidate generation.
+
+    `compute_bounds_1d` is typically a TR object's bound computation method.
+    """
+    lb, ub = compute_bounds_1d(x_center, lengthscales)
+    return generate_raasp_candidates(
+        x_center, lb, ub, num_candidates, rng=rng, sobol_engine=sobol_engine
+    )
+
+
 def to_unit(x: np.ndarray | Any, bounds: np.ndarray | Any) -> np.ndarray:
     import numpy as np
 
@@ -229,15 +319,15 @@ def gp_thompson_sample(
 
     x_torch = torch.as_tensor(x_cand, dtype=torch.float64)
     seed = int(rng.integers(2**31 - 1))
-    gen = torch.Generator(device=x_torch.device)
-    gen.manual_seed(seed)
     with (
         torch.no_grad(),
         gpytorch.settings.fast_pred_var(),
-        torch_rng_context(gen),
+        torch_seed_context(seed, device=x_torch.device),
     ):
         posterior = model.posterior(x_torch)
         samples = posterior.sample(sample_shape=torch.Size([1]))
+    if samples.ndim != 2:
+        raise ValueError(samples.shape)
     ts = samples[0].reshape(-1)
     scores = ts.detach().cpu().numpy().reshape(-1)
     scores = gp_y_mean + gp_y_std * scores

enn/turbo/turbo_zero_impl.py

@@ -7,9 +7,13 @@ if TYPE_CHECKING:
     from numpy.random import Generator
 
 from .base_turbo_impl import BaseTurboImpl
+from .turbo_config import TurboZeroConfig
 
 
 class TurboZeroImpl(BaseTurboImpl):
+    def __init__(self, config: TurboZeroConfig) -> None:
+        super().__init__(config)
+
     def select_candidates(
         self,
         x_cand: np.ndarray,
@@ -18,6 +22,7 @@ class TurboZeroImpl(BaseTurboImpl):
         rng: Generator,
         fallback_fn: Callable[[np.ndarray, int], np.ndarray],
         from_unit_fn: Callable[[np.ndarray], np.ndarray],
+        tr_state: object | None = None,  # noqa: ARG002
     ) -> np.ndarray:
         from .proposal import select_uniform
 

{ennbo-0.1.0.dist-info → ennbo-0.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ennbo
-Version: 0.1.0
+Version: 0.1.2
 Summary: Epistemic Nearest Neighbors
 Project-URL: Homepage, https://github.com/yubo-research/enn
 Project-URL: Source, https://github.com/yubo-research/enn
@@ -36,10 +36,10 @@ Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Scientific/Engineering :: Mathematics
 Requires-Python: >=3.11
-Requires-Dist: faiss-cpu==1.9.0
+Requires-Dist: faiss-cpu>=1.9.0
 Requires-Dist: gpytorch==1.13
 Requires-Dist: nds==0.4.3
-Requires-Dist: numpy==1.26.4
+Requires-Dist: numpy<2.0.0,>=1.26.4
 Requires-Dist: scipy==1.15.3
 Requires-Dist: torch==2.5.1
 Description-Content-Type: text/markdown
@@ -80,9 +80,10 @@ On my MacBook I can run into problems with dependencies and compatibilities.
 
 On MacOS try:
 ```
-micromamba env create -n ennbo -f conda-macos.yml
+micromamba env create -n ennbo -f admin/conda-macos.yml
 micromamba activate ennbo
 pip install --no-deps ennbo
+pytest -sv tests
 ```
 
 You may replace `micromamba` with `conda` and this will probably still work.

ennbo-0.1.2.dist-info/RECORD

@@ -0,0 +1,29 @@
+enn/__init__.py,sha256=VYIuOTCjhUFIJm78IoJv0WXtvA_IuZhY1sSMJJM3dx8,507
+enn/enn/__init__.py,sha256=K3rntg_ZkITStmXMTBcEhxeS1kel1bb7wB_C7-2WE5Y,135
+enn/enn/enn.py,sha256=HfdrK2gXoI1JvvARsh4NdGOpVpCY2qY_A2RpK4JFVZ4,9310
+enn/enn/enn_fit.py,sha256=RkyFYX4-nUteGivNNS195M2mdRWiGOrLzypVL2b_FsE,4450
+enn/enn/enn_normal.py,sha256=Lm9n-eW5WRn33nb3b9xTGv44Dfn9xAhjys5UJZm2xlc,662
+enn/enn/enn_params.py,sha256=v53qHKwUxnZFNBlcSWI5WqpgpoyFeOzvWs0BxcDqt4o,747
+enn/enn/enn_util.py,sha256=PSeYmxZHz4xLJ6pMr9n22MLxvVmmgBqmQl1ckPrLzDo,4142
+enn/turbo/__init__.py,sha256=utnD3CLZgjCvw-46AAu5Tv2M2Vbg5YXK-_TycGk5BU4,197
+enn/turbo/base_turbo_impl.py,sha256=y4SP9FDT8OaNDOBhKvx1WuTksBd0d-vD8NQxj91QAuA,4396
+enn/turbo/lhd_only_impl.py,sha256=czqLTwhb8d-pKq6jy7g28JcDWMSoopLjNrDeP1dd-3A,1254
+enn/turbo/morbo_trust_region.py,sha256=9z_DgXHEEUfgajRnZ5ieJdeoVAUt7BlLNGu-Thi5Tg4,5973
+enn/turbo/no_trust_region.py,sha256=IxZB1nvmLFgb6GjAXrNpBe1TzgzwJDPYEB8wa_ZPX3k,1813
+enn/turbo/proposal.py,sha256=obFqVyXtZ49veqwnktTJ0_F0nqCERvgFnInstgqhllM,4252
+enn/turbo/turbo_config.py,sha256=tci_GODIED3UHE63XiB9XyxPu_0J5_8R90lxBVxidOQ,2500
+enn/turbo/turbo_enn_impl.py,sha256=qcUssC4xaoMu7usIlp5oJtXrsLY1Rn-mvHUW95zlWZQ,7249
+enn/turbo/turbo_gp.py,sha256=Hi11t0nw5YEG4WM6DeoOW4X-w-M5KGG-P3Zc1sPPx1k,1069
+enn/turbo/turbo_gp_base.py,sha256=tnE5uX_eAt1Db-gemyy83ZvKpdNbMg_tsWkh6sG7zaM,638
+enn/turbo/turbo_gp_noisy.py,sha256=itTL9jUCjE566jwDODT0P36fozsfU_bXACyuKqxYMXs,1080
+enn/turbo/turbo_mode.py,sha256=JMP1jkFCRwPtOzU95MWWd04Sgze7eKF0xNkiPqtQ8SI,181
+enn/turbo/turbo_mode_impl.py,sha256=ubUkV4reOPJH3jbAh6R65cutEHOF23z7Uw1bBE3s9T0,1923
+enn/turbo/turbo_one_impl.py,sha256=PXaBNdLKCgtsLC8Q2z4pHoHYu0edPqKAXb6Tmr4Guvs,11098
+enn/turbo/turbo_optimizer.py,sha256=h6Mu3Pqb2yQRrrqEh6ODxhKIt4Nnt-C241XrNluqxtk,20274
+enn/turbo/turbo_trust_region.py,sha256=0wlN_LhsfMeLqqjhq3xhkJpXtTYIzUon0zbvUjRm2_Q,3797
+enn/turbo/turbo_utils.py,sha256=bEe1F3hBUOUVmodF1WJQv_EKRmp4X74kMxWLNnuV-z0,10733
+enn/turbo/turbo_zero_impl.py,sha256=SvPexeUTQzDPbAwPdZib5lRcIHPOmwD3-ewMxED-nlQ,832
+ennbo-0.1.2.dist-info/METADATA,sha256=-Twds_sAT4LLkcN00vzmS9_a5jUSWRsqi1bO2_RKtHw,5960
+ennbo-0.1.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+ennbo-0.1.2.dist-info/licenses/LICENSE,sha256=KTA0NjGalsl_JGrjT_x6SSq9ZYVO3gQ-hLVMEaekc5w,1070
+ennbo-0.1.2.dist-info/RECORD,,

ennbo-0.1.0.dist-info/RECORD

@@ -1,27 +0,0 @@
-enn/__init__.py,sha256=VYIuOTCjhUFIJm78IoJv0WXtvA_IuZhY1sSMJJM3dx8,507
-enn/enn/__init__.py,sha256=K3rntg_ZkITStmXMTBcEhxeS1kel1bb7wB_C7-2WE5Y,135
-enn/enn/enn.py,sha256=ZdDPivZj4SL9e87FolU1oscdPdcwUeIByIrvBLsoCfE,8060
-enn/enn/enn_fit.py,sha256=uv1BHO-nbxVXkR_tM1Ggoh6YNuR-VrjVECFxLquC7u8,4328
-enn/enn/enn_normal.py,sha256=3kOymSx2kzcBMavScXLflPm_gDDLGF9fYLBJ816I3xg,596
-enn/enn/enn_params.py,sha256=fwLZTA8ciRp4XUF5L_VAVsC3EvFuOzR85OYLVtv6TSw,184
-enn/enn/enn_util.py,sha256=ZELPVeyUl0wiHOxjHYKjxeDz88ExmKMeX3P-bQ6tCoE,3075
-enn/turbo/__init__.py,sha256=utnD3CLZgjCvw-46AAu5Tv2M2Vbg5YXK-_TycGk5BU4,197
-enn/turbo/base_turbo_impl.py,sha256=wThjwXGboRrVTamsnvzmM0WNIOZ91GNJ-BmGzjgqdhg,2699
-enn/turbo/lhd_only_impl.py,sha256=yWsOw7Oq0xfEnyXg5AXJSzZFjM7162pqNY37fHQtJQ4,1023
-enn/turbo/proposal.py,sha256=w1izo3ooiiravNRoFWK5ZK7BH-f_HWgqYP8heVtLmYs,3977
-enn/turbo/turbo_config.py,sha256=J0ww_qKDDMpbFVXdntuSbJtUTbdnXrFJyGD1svzG3RM,980
-enn/turbo/turbo_enn_impl.py,sha256=YMAS4krpPXPNtlh46RRG3VLMuGyYLFw5UkPRBU29mzA,5837
-enn/turbo/turbo_gp.py,sha256=i1bxVHima0Nv4MCLlADtlRzt1cENcnVLYk3S9vCoF4c,797
-enn/turbo/turbo_gp_base.py,sha256=tnE5uX_eAt1Db-gemyy83ZvKpdNbMg_tsWkh6sG7zaM,638
-enn/turbo/turbo_gp_noisy.py,sha256=itTL9jUCjE566jwDODT0P36fozsfU_bXACyuKqxYMXs,1080
-enn/turbo/turbo_mode.py,sha256=JMP1jkFCRwPtOzU95MWWd04Sgze7eKF0xNkiPqtQ8SI,181
-enn/turbo/turbo_mode_impl.py,sha256=3HKBjOS96Wn-R_znctQm9Ivrm3FhgZFTuBp7McNDQ88,1749
-enn/turbo/turbo_one_impl.py,sha256=nS02RdRMcEsi3II07jzcrQbsFsfWYTeahUcqoyhig4Q,5207
-enn/turbo/turbo_optimizer.py,sha256=IlofW9_ogCeQMVXa7n8xWEg5fbJBUkvAkeLKe3MoXlA,11902
-enn/turbo/turbo_trust_region.py,sha256=VHNYKWtKLt3iKHI0enL9qMMu1Bwi1nupo20L0Sv-vYY,3759
-enn/turbo/turbo_utils.py,sha256=XU9-YtW1u5-HKk3bA_M-hVNFPAuNcIYozAmej7ulVsY,7532
-enn/turbo/turbo_zero_impl.py,sha256=S4TEHYkVDowtyWSVxWO0ncd1OUIFpeV3IR-eanGr1vg,643
-ennbo-0.1.0.dist-info/METADATA,sha256=slkhtsGXaO31u8w35LNKXN2noxUJYTqHQF7bv1DZMmA,5930
-ennbo-0.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-ennbo-0.1.0.dist-info/licenses/LICENSE,sha256=KTA0NjGalsl_JGrjT_x6SSq9ZYVO3gQ-hLVMEaekc5w,1070
-ennbo-0.1.0.dist-info/RECORD,,