fix(ifbo): Transformation of categoricals and handling of unusual fidelity bounds

neps/optimizers/bayesian_optimization.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import itertools
 import math
 from collections.abc import Mapping
 from dataclasses import dataclass
@@ -89,7 +90,7 @@ def __call__(
 
         n_to_sample = 1 if n is None else n
         n_sampled = len(trials)
-        config_ids = iter(str(i + 1) for i in range(n_sampled, n_sampled + n_to_sample))
+        id_generator = iter(str(i) for i in itertools.count(n_sampled + 1))
 
         # If the amount of configs evaluated is less than the initial design
         # requirement, keep drawing from initial design
@@ -98,23 +99,39 @@ def __call__(
             for trial in trials.values()
             if trial.report is not None and trial.report.objective_to_minimize is not None
         )
+        sampled_configs: list[SampledConfig] = []
+
         if n_evaluated < self.n_initial_design:
+            # For reproducibility, we need to ensure we do the same sample of all
+            # configs each time.
             design_samples = make_initial_design(
                 parameters=parameters,
                 encoder=self.encoder,
                 sample_prior_first=self.sample_prior_first if n_sampled == 0 else False,
                 sampler=self.prior if self.prior is not None else "uniform",
                 seed=None,  # TODO: Seeding, however we need to avoid repeating configs
-                sample_size=n_to_sample,
+                sample_size=self.n_initial_design,
             )
+
+            # Then take the subset we actually need
+            design_samples = design_samples[n_evaluated:]
             for sample in design_samples:
                 sample.update(self.space.constants)
 
-            sampled_configs = [
-                SampledConfig(id=config_id, config=config)
-                for config_id, config in zip(config_ids, design_samples, strict=True)
-            ]
-            return sampled_configs[0] if n is None else sampled_configs
+            sampled_configs.extend(
+                [
+                    SampledConfig(id=config_id, config=config)
+                    for config_id, config in zip(
+                        id_generator,
+                        design_samples,
+                        # NOTE: We use a generator for the ids so no need for strict
+                        strict=False,
+                    )
+                ]
+            )
+
+            if len(sampled_configs) >= n_to_sample:
+                return sampled_configs[0] if n is None else sampled_configs
 
         # Otherwise, we encode trials and setup to fit and acquire from a GP
         data, encoder = encode_trials_for_gp(
@@ -149,6 +166,7 @@ def __call__(
             prior = None if pibo_exp_term < 1e-4 else self.prior
 
         gp = make_default_single_obj_gp(x=data.x, y=data.y, encoder=encoder)
+        n_to_acquire = n_to_sample - len(sampled_configs)
         candidates = fit_and_acquire_from_gp(
             gp=gp,
             x_train=data.x,
@@ -164,7 +182,7 @@ def __call__(
                 prune_baseline=True,
             ),
             prior=prior,
-            n_candidates_required=n_to_sample,
+            n_candidates_required=n_to_acquire,
             pibo_exp_term=pibo_exp_term,
             costs=data.cost if self.cost_aware is not False else None,
             cost_percentage_used=cost_percent,
@@ -175,8 +193,15 @@ def __call__(
         for config in configs:
             config.update(self.space.constants)
 
-        sampled_configs = [
-            SampledConfig(id=config_id, config=config)
-            for config_id, config in zip(config_ids, configs, strict=True)
-        ]
+        sampled_configs.extend(
+            [
+                SampledConfig(id=config_id, config=config)
+                for config_id, config in zip(
+                    id_generator,
+                    configs,
+                    # NOTE: We use a generator for the ids so no need for strict
+                    strict=False,
+                )
+            ]
+        )
         return sampled_configs[0] if n is None else sampled_configs

neps/optimizers/ifbo.py

@@ -23,6 +23,7 @@
 
 # NOTE: Ifbo was trained using 32 bit
 FTPFN_DTYPE = torch.float32
+BUDGET_DOMAIN_EPS = 1e-6
 
 
 def _adjust_space_to_match_stepsize(
@@ -140,20 +141,34 @@ def __call__(
 
         assert n is None, "TODO"
         ids = [int(config_id.split("_", maxsplit=1)[0]) for config_id in trials]
-        new_id = max(ids) + 1 if len(ids) > 0 else 0
+        new_id = max(ids) + 1 if len(ids) > 0 else 1
 
         # The FTPFN surrogate takes in a budget in the range [0, 1]
         # We also need to be able to map these to discrete integers
         # Hence we use the two domains below to do so.
 
         # Domain in which we should pass budgets to ifbo model
-        budget_domain = Domain.floating(lower=1 / fidelity.upper, upper=1)
+        # IFBO expects them to be in `(0 1]`, where explicitly 0 is
+        # not allowed. We map this close to `BUDGET_DOMAIN_EPS` depending
+        # on the scale of the fidelity domain.
+        budget_domain = Domain.floating(
+            lower=(fidelity.lower + BUDGET_DOMAIN_EPS)
+            / (fidelity.upper - fidelity.lower),
+            upper=1,
+        )
+
+        # However, we need to make sure we don't end up with a positive
+        # `lower=` which gauranteed with this assertion.
+        if fidelity.upper - fidelity.lower < BUDGET_DOMAIN_EPS:
+            raise ValueError(
+                f"Fidelity domain {fidelity} is too small to be used with ifBO."
+            )
 
         # Domain from which we assign an index to each budget
-        budget_index_domain = Domain.indices(self.n_fidelity_bins)
+        budget_index_domain = Domain.indices(self.n_fidelity_bins + 1)
 
         # If we havn't passed the intial design phase
-        if new_id < self.n_initial_design:
+        if new_id <= self.n_initial_design:
             init_design = make_initial_design(
                 parameters=parameters,
                 encoder=self.encoder,
@@ -163,7 +178,7 @@ def __call__(
                 sample_size=self.n_initial_design,
             )
 
-            config = init_design[new_id]
+            config = init_design[new_id - 1]
             config[fidelity_name] = fidelity.lower
             config.update(self.space.constants)
             return SampledConfig(id=f"{new_id}_0", config=config)
@@ -238,13 +253,16 @@ def _mfpi_random(samples: torch.Tensor) -> torch.Tensor:
             local_search_sample_size=256,
             local_search_confidence=0.95,
         )
+
         _id, fid, config = decode_ftpfn_data(
             best_row,
-            self.encoder,
+            encoder=self.encoder,
             budget_domain=budget_domain,
             fidelity_domain=fidelity.domain,
         )[0]
 
+        # If _id is None, that means a new config was sampled and
+        # should be evaluated one step
         if _id is None:
             config[fidelity_name] = fid
             config.update(self.space.constants)

neps/runtime.py

@@ -464,30 +464,29 @@ def _get_next_trial(self) -> Trial | Literal["break"]:
                     return this_workers_trial
                 except TrialAlreadyExistsError as e:
                     if e.trial_id in trials:
-                        logger.error(
-                            "The new sampled trial was given an id of '%s', yet this"
-                            " exists in the loaded in trials given to the optimizer. This"
-                            " indicates a bug with the optimizers allocation of ids.",
-                            e.trial_id,
-                        )
-                    else:
-                        _grace = DefaultWorker._GRACE
-                        _inc = FS_SYNC_GRACE_INC
-                        logger.warning(
-                            "The new sampled trial was given an id of '%s', which is not"
-                            " one that was loaded in by the optimizer. This is usually"
-                            " an indication that the file-system you are running on"
-                            " is not atmoic in synchoronizing file operations."
-                            " We have attempted to stabalize this but milage may vary."
-                            " We are incrementing a grace period for file-locks from"
-                            " '%s's to '%s's. You can control the initial"
-                            " grace with 'NEPS_FS_SYNC_GRACE_BASE' and the increment with"
-                            " 'NEPS_FS_SYNC_GRACE_INC'.",
-                            e.trial_id,
-                            _grace,
-                            _grace + _inc,
-                        )
-                        DefaultWorker._GRACE = _grace + FS_SYNC_GRACE_INC
+                        raise RuntimeError(
+                            f"The new sampled trial was given an id of {e.trial_id}, yet"
+                            " this exists in the loaded in trials given to the optimizer."
+                            " This is a bug with the optimizers allocation of ids."
+                        ) from e
+
+                    _grace = DefaultWorker._GRACE
+                    _inc = FS_SYNC_GRACE_INC
+                    logger.warning(
+                        "The new sampled trial was given an id of '%s', which is not"
+                        " one that was loaded in by the optimizer. This is usually"
+                        " an indication that the file-system you are running on"
+                        " is not atmoic in synchoronizing file operations."
+                        " We have attempted to stabalize this but milage may vary."
+                        " We are incrementing a grace period for file-locks from"
+                        " '%s's to '%s's. You can control the initial"
+                        " grace with 'NEPS_FS_SYNC_GRACE_BASE' and the increment with"
+                        " 'NEPS_FS_SYNC_GRACE_INC'.",
+                        e.trial_id,
+                        _grace,
+                        _grace + _inc,
+                    )
+                    DefaultWorker._GRACE = _grace + FS_SYNC_GRACE_INC
                     raise e
 
     # Forgive me lord, for I have sinned, this function is atrocious but complicated

neps/space/encoding.py

@@ -72,11 +72,10 @@ def __post_init__(self) -> None:
 
         self.domain = Domain.indices(len(self.choices), is_categorical=True)
         self._lookup = None
-        if len(self.choices) > 3:
-            try:
-                self._lookup = {c: i for i, c in enumerate(self.choices)}
-            except TypeError:
-                self._lookup = None
+        try:
+            self._lookup = {c: i for i, c in enumerate(self.choices)}
+        except TypeError:
+            self._lookup = None
 
     @override
     def encode(
@@ -117,8 +116,9 @@ class CategoricalToUnitNorm(TensorTransformer):
 
     choices: Sequence[Any]
 
-    domain: Domain = field(init=False)
-    _integer_transformer: CategoricalToIntegerTransformer = field(init=False)
+    domain: Domain[float] = field(init=False)
+    _cat_int_domain: Domain[int] = field(init=False)
+    _lookup: dict[Any, int] | None = field(init=False)
 
     def __post_init__(self) -> None:
         self.domain = Domain.floating(
@@ -127,7 +127,11 @@ def __post_init__(self) -> None:
             bins=len(self.choices),
             is_categorical=True,
         )
-        self._integer_transformer = CategoricalToIntegerTransformer(self.choices)
+        self._cat_int_domain = Domain.indices(len(self.choices), is_categorical=True)
+        try:
+            self._lookup = {c: i for i, c in enumerate(self.choices)}
+        except TypeError:
+            self._lookup = None
 
     @override
     def encode(
@@ -138,13 +142,19 @@ def encode(
         dtype: torch.dtype | None = None,
         device: torch.device | None = None,
     ) -> torch.Tensor:
-        integers = self._integer_transformer.encode(
-            x,
-            dtype=dtype if dtype is not None else torch.float64,
-            device=device,
+        if dtype is None:
+            dtype = torch.int if out is None else out.dtype
+
+        values = (
+            [self._lookup[c] for c in x]
+            if self._lookup
+            else [self.choices.index(c) for c in x]
         )
+        integers = torch.tensor(values, dtype=torch.int64, device=device)
         binned_floats = self.domain.cast(
-            integers, frm=self._integer_transformer.domain, dtype=dtype
+            integers,
+            frm=self._cat_int_domain,
+            dtype=dtype,
         )
         if out is not None:
             return out.copy_(binned_floats)
@@ -153,8 +163,8 @@ def encode(
 
     @override
     def decode(self, x: torch.Tensor) -> list[Any]:
-        x = torch.round(x * (len(self.choices) - 1)).type(torch.int64)
-        return self._integer_transformer.decode(x)
+        x = self._cat_int_domain.cast(x, frm=self.domain)
+        return [self.choices[int(i)] for i in torch.round(x).tolist()]
 
 
 # TODO: Maybe add a shift argument, could be useful to have `0` as midpoint

neps/space/parameters.py

@@ -75,6 +75,13 @@ def __post_init__(self) -> None:
         self.lower = float(self.lower)
         self.upper = float(self.upper)
 
+        if self.is_fidelity and (self.lower < 0 or self.upper < 0):
+            raise ValueError(
+                f"Float parameter: fidelity bounds error. Expected fidelity"
+                f" bounds to be >= 0, but got lower={self.lower}, "
+                f" upper={self.upper}."
+            )
+
         if np.isnan(self.lower):
             raise ValueError("Can not have lower bound that is nan")
 
@@ -149,6 +156,13 @@ def __post_init__(self) -> None:
         self.lower = lower_int
         self.upper = upper_int
 
+        if self.is_fidelity and (self.lower < 0 or self.upper < 0):
+            raise ValueError(
+                f"Integer parameter: fidelity bounds error. Expected fidelity"
+                f" bounds to be >= 0, but got lower={self.lower}, "
+                f" upper={self.upper}."
+            )
+
         if self.log and (self.lower <= 0 or self.upper <= 0):
             raise ValueError(
                 f"Integer parameter: bounds error (log scale cant have bounds <= 0). "