rename budget to max_cost_total

neps/api.py

@@ -416,7 +416,7 @@ def _run_args(
         SearcherMapping, searcher_alg, "searcher", as_class=True
     )(
         pipeline_space=pipeline_space,
-        budget=max_cost_total,  # TODO: use max_cost_total everywhere
+        max_cost_total=max_cost_total,  # TODO: use max_cost_total everywhere
         **searcher_config,
     )
 

neps/optimizers/base_optimizer.py

@@ -89,7 +89,7 @@ def __init__(
         pipeline_space: SearchSpace,
         patience: int = 50,
         logger: logging.Logger | None = None,
-        budget: int | float | None = None,
+        max_cost_total: int | float | None = None,
         objective_to_minimize_value_on_error: float | None = None,
         cost_value_on_error: float | None = None,
         learning_curve_on_error: float | list[float] | None = None,
@@ -98,7 +98,7 @@ def __init__(
         if patience < 1:
             raise ValueError("Patience should be at least 1")
 
-        self.budget = budget
+        self.max_cost_total = max_cost_total
         self.pipeline_space = pipeline_space
         self.patience = patience
         self.logger = logger or logging.getLogger("neps")
@@ -111,13 +111,13 @@ def __init__(
     def ask(
         self,
         trials: Mapping[str, Trial],
-        budget_info: BudgetInfo | None,
+        max_cost_total_info: BudgetInfo | None,
     ) -> SampledConfig:
         """Sample a new configuration.
 
         Args:
             trials: All of the trials that are known about.
-            budget_info: information about the budget
+            max_cost_total_info: information about the max_cost_total
 
         Returns:
             SampledConfig: a sampled configuration

neps/optimizers/bayesian_optimization/acquisition_functions/cost_cooling.py

@@ -41,14 +41,14 @@ def apply_cost_cooling(
 def cost_cooled_acq(
     acq_fn: AcquisitionFunction,
     model: GPyTorchModel,
-    used_budget_percentage: float,
+    used_max_cost_total_percentage: float,
 ) -> WeightedAcquisition:
-    assert 0 <= used_budget_percentage <= 1
+    assert 0 <= used_max_cost_total_percentage <= 1
     return WeightedAcquisition(
         acq=acq_fn,
         apply_weight=partial(
             apply_cost_cooling,
             cost_model=model,
-            alpha=1 - used_budget_percentage,
+            alpha=1 - used_max_cost_total_percentage,
         ),
     )

neps/optimizers/bayesian_optimization/models/ftpfn.py

@@ -164,7 +164,7 @@ def encode_ftpfn(
         device=device,
         dtype=dtype,
     )
-    train_budgets = budget_domain.cast(
+    train_max_cost_total = budget_domain.cast(
         train_fidelities, frm=space.fidelity.domain, dtype=dtype
     )
 
@@ -192,7 +192,7 @@ def encode_ftpfn(
         )
     maximize_ys = 1 - minimize_ys
     x_train = torch.cat(
-        [ids.unsqueeze(1), train_budgets.unsqueeze(1), train_configs], dim=1
+        [ids.unsqueeze(1), train_max_cost_total.unsqueeze(1), train_configs], dim=1
     )
     return x_train, maximize_ys
 

neps/optimizers/bayesian_optimization/models/gp.py

@@ -401,7 +401,7 @@ def fit_and_acquire_from_gp(
         acquisition = cost_cooled_acq(
             acq_fn=acquisition,
             model=cost_gp,
-            used_budget_percentage=cost_percentage_used,
+            used_max_cost_total_percentage=cost_percentage_used,
         )
 
     _n = n_candidates_required if n_candidates_required is not None else 1

neps/optimizers/bayesian_optimization/optimizer.py

@@ -69,7 +69,7 @@ def __init__(
         device: torch.device | None = None,
         encoder: ConfigEncoder | None = None,
         seed: int | None = None,
-        budget: Any | None = None,  # TODO: remove
+        max_cost_total: Any | None = None,  # TODO: remove
         surrogate_model: Any | None = None,  # TODO: remove
         objective_to_minimize_value_on_error: Any | None = None,  # TODO: remove
         cost_value_on_error: Any | None = None,  # TODO: remove
@@ -134,7 +134,7 @@ def __init__(
     def ask(
         self,
         trials: Mapping[str, Trial],
-        budget_info: BudgetInfo | None = None,
+        max_cost_total_info: BudgetInfo | None = None,
     ) -> SampledConfig:
         n_sampled = len(trials)
         config_id = str(n_sampled + 1)
@@ -164,14 +164,16 @@ def ask(
 
         cost_percent = None
         if self.use_cost:
-            if budget_info is None:
+            if max_cost_total_info is None:
                 raise ValueError(
                     "Must provide a 'cost' to configurations if using cost"
                     " with BayesianOptimization."
                 )
-            if budget_info.max_cost_budget is None:
+            if max_cost_total_info.max_cost_total is None:
                 raise ValueError("Cost budget must be set if using cost")
-            cost_percent = budget_info.used_cost_budget / budget_info.max_cost_budget
+            cost_percent = (
+                max_cost_total_info.used_cost_budget / max_cost_total_info.max_cost_total
+            )
 
         # If we should use the prior, weight the acquisition function by
         # the probability of it being sampled from the prior.

neps/optimizers/grid_search/optimizer.py

@@ -97,7 +97,7 @@ def __init__(self, pipeline_space: SearchSpace, seed: int | None = None):
 
     @override
     def ask(
-        self, trials: Mapping[str, Trial], budget_info: BudgetInfo | None
+        self, trials: Mapping[str, Trial], max_cost_total_info: BudgetInfo | None
     ) -> SampledConfig:
         _num_previous_configs = len(trials)
         if _num_previous_configs > len(self.configs_list) - 1:

neps/optimizers/multi_fidelity/hyperband.py

@@ -47,7 +47,7 @@ def __init__(
         self,
         *,
         pipeline_space: SearchSpace,
-        budget: int,
+        max_cost_total: int,
         eta: int = 3,
         initial_design_type: Literal["max_budget", "unique_configs"] = "max_budget",
         use_priors: bool = False,
@@ -63,7 +63,7 @@ def __init__(
     ):
         args = {
             "pipeline_space": pipeline_space,
-            "budget": budget,
+            "max_cost_total": max_cost_total,
             "eta": eta,
             "early_stopping_rate": self.early_stopping_rate,  # HB subsumes this from SH
             "initial_design_type": initial_design_type,
@@ -125,7 +125,7 @@ def _handle_promotions(self) -> None:
     def ask(
         self,
         trials: Mapping[str, Trial],
-        budget_info: BudgetInfo | None,
+        max_cost_total_info: BudgetInfo | None,
     ) -> SampledConfig:
         completed: dict[str, ConfigResult] = {
             trial_id: trial.into_config_result(self.pipeline_space.from_dict)
@@ -273,7 +273,7 @@ def __init__(
         self,
         *,
         pipeline_space: SearchSpace,
-        budget: int,
+        max_cost_total: int,
         eta: int = 3,
         initial_design_type: Literal["max_budget", "unique_configs"] = "max_budget",
         sampling_policy: Any = FixedPriorPolicy,
@@ -288,7 +288,7 @@ def __init__(
     ):
         super().__init__(
             pipeline_space=pipeline_space,
-            budget=budget,
+            max_cost_total=max_cost_total,
             eta=eta,
             initial_design_type=initial_design_type,
             use_priors=self.use_priors,  # key change to the base HB class
@@ -311,7 +311,7 @@ def __init__(
         self,
         *,
         pipeline_space: SearchSpace,
-        budget: int,
+        max_cost_total: int,
         eta: int = 3,
         initial_design_type: Literal["max_budget", "unique_configs"] = "max_budget",
         sampling_policy: Any = EnsemblePolicy,
@@ -326,7 +326,7 @@ def __init__(
     ):
         super().__init__(
             pipeline_space=pipeline_space,
-            budget=budget,
+            max_cost_total=max_cost_total,
             eta=eta,
             initial_design_type=initial_design_type,
             sampling_policy=sampling_policy,
@@ -361,7 +361,7 @@ def __init__(
         self,
         *,
         pipeline_space: SearchSpace,
-        budget: int,
+        max_cost_total: int,
         eta: int = 3,
         initial_design_type: Literal["max_budget", "unique_configs"] = "max_budget",
         use_priors: bool = False,
@@ -377,7 +377,7 @@ def __init__(
     ):
         args = {
             "pipeline_space": pipeline_space,
-            "budget": budget,
+            "max_cost_total": max_cost_total,
             "eta": eta,
             "initial_design_type": initial_design_type,
             "use_priors": use_priors,
@@ -456,7 +456,7 @@ def __init__(
         self,
         *,
         pipeline_space: SearchSpace,
-        budget: int,
+        max_cost_total: int,
         eta: int = 3,
         initial_design_type: Literal["max_budget", "unique_configs"] = "max_budget",
         sampling_policy: Any = FixedPriorPolicy,
@@ -471,7 +471,7 @@ def __init__(
     ):
         super().__init__(
             pipeline_space=pipeline_space,
-            budget=budget,
+            max_cost_total=max_cost_total,
             eta=eta,
             initial_design_type=initial_design_type,
             use_priors=self.use_priors,  # key change to the base Async HB class
@@ -495,7 +495,7 @@ def __init__(
         self,
         *,
         pipeline_space: SearchSpace,
-        budget: int,
+        max_cost_total: int,
         eta: int = 3,
         initial_design_type: Literal["max_budget", "unique_configs"] = "max_budget",
         use_priors: bool = False,
@@ -520,7 +520,7 @@ def __init__(
     ):
         hb_args = {
             "pipeline_space": pipeline_space,
-            "budget": budget,
+            "max_cost_total": max_cost_total,
             "eta": eta,
             "initial_design_type": initial_design_type,
             "use_priors": use_priors,

neps/optimizers/multi_fidelity/ifbo.py

@@ -98,7 +98,7 @@ def __init__(
         initial_design_size: int | Literal["ndim"] = "ndim",
         n_acquisition_new_configs: int = 1_000,
         device: torch.device | None = None,
-        budget: int | float | None = None,  # TODO: Remove
+        max_cost_total: int | float | None = None,  # TODO: Remove
         objective_to_minimize_value_on_error: float | None = None,  # TODO: Remove
         cost_value_on_error: float | None = None,  # TODO: Remove
         ignore_errors: bool = False,  # TODO: Remove
@@ -171,7 +171,7 @@ def __init__(
     def ask(
         self,
         trials: Mapping[str, Trial],
-        budget_info: BudgetInfo | None = None,
+        max_cost_total_info: BudgetInfo | None = None,
     ) -> SampledConfig:
         ids = [int(config_id.split("_", maxsplit=1)[0]) for config_id in trials]
         new_id = max(ids) + 1 if len(ids) > 0 else 0

neps/optimizers/multi_fidelity/successive_halving.py

@@ -49,7 +49,7 @@ def __init__(
         self,
         *,
         pipeline_space: SearchSpace,
-        budget: int | None = None,
+        max_cost_total: int | None = None,
         eta: int = 3,
         early_stopping_rate: int = 0,
         initial_design_type: Literal["max_budget", "unique_configs"] = "max_budget",
@@ -68,7 +68,7 @@ def __init__(
 
         Args:
             pipeline_space: Space in which to search
-            budget: Maximum budget
+            max_cost_total: Maximum budget
             eta: The reduction factor used by SH
             early_stopping_rate: Determines the number of rungs in an SH bracket
                 Choosing 0 creates maximal rungs given the fidelity bounds
@@ -94,7 +94,7 @@ def __init__(
         """
         super().__init__(
             pipeline_space=pipeline_space,
-            budget=budget,
+            max_cost_total=max_cost_total,
             objective_to_minimize_value_on_error=objective_to_minimize_value_on_error,
             cost_value_on_error=cost_value_on_error,
             ignore_errors=ignore_errors,
@@ -320,7 +320,7 @@ def _fit_models(self) -> None:
     def ask(
         self,
         trials: Mapping[str, Trial],
-        budget_info: BudgetInfo | None,
+        max_cost_total_info: BudgetInfo | None,
     ) -> SampledConfig:
         """This is basically the fit method."""
         completed: dict[str, ConfigResult] = {
@@ -489,13 +489,13 @@ def _enhance_priors(self, confidence_score: dict[str, float] | None = None) -> N
 
 class SuccessiveHalving(SuccessiveHalvingBase):
     def _calc_budget_used_in_bracket(self, config_history: list[int]) -> int:
-        budget = 0
+        max_cost_total = 0
         for rung in self.config_map:
             count = sum(config_history == rung)
             # `range(min_rung, rung+1)` counts the black-box cost of promotions since
             # SH budgets assume each promotion involves evaluation from scratch
-            budget += count * sum(np.arange(self.min_rung, rung + 1))
-        return budget
+            max_cost_total += count * sum(np.arange(self.min_rung, rung + 1))
+        return max_cost_total
 
     def clear_old_brackets(self) -> None:
         """Enforces reset at each new bracket.
@@ -559,7 +559,7 @@ def __init__(
         self,
         *,
         pipeline_space: SearchSpace,
-        budget: int,
+        max_cost_total: int,
         eta: int = 3,
         early_stopping_rate: int = 0,
         initial_design_type: Literal["max_budget", "unique_configs"] = "max_budget",
@@ -575,7 +575,7 @@ def __init__(
     ):
         super().__init__(
             pipeline_space=pipeline_space,
-            budget=budget,
+            max_cost_total=max_cost_total,
             eta=eta,
             early_stopping_rate=early_stopping_rate,
             initial_design_type=initial_design_type,
@@ -599,7 +599,7 @@ def __init__(
         self,
         *,
         pipeline_space: SearchSpace,
-        budget: int,
+        max_cost_total: int,
         eta: int = 3,
         early_stopping_rate: int = 0,
         initial_design_type: Literal["max_budget", "unique_configs"] = "max_budget",
@@ -616,7 +616,7 @@ def __init__(
     ):
         super().__init__(
             pipeline_space=pipeline_space,
-            budget=budget,
+            max_cost_total=max_cost_total,
             eta=eta,
             early_stopping_rate=early_stopping_rate,
             initial_design_type=initial_design_type,
@@ -642,7 +642,7 @@ def __init__(
         self,
         *,
         pipeline_space: SearchSpace,
-        budget: int,
+        max_cost_total: int,
         eta: int = 3,
         early_stopping_rate: int = 0,
         initial_design_type: Literal["max_budget", "unique_configs"] = "max_budget",
@@ -658,7 +658,7 @@ def __init__(
     ):
         super().__init__(
             pipeline_space=pipeline_space,
-            budget=budget,
+            max_cost_total=max_cost_total,
             eta=eta,
             early_stopping_rate=early_stopping_rate,
             initial_design_type=initial_design_type,