Add heterogeneous HPO simulation support

Author: Innixma

tabarena/tabarena/nips2025_utils/artifacts/method_metadata.py

@@ -628,15 +628,17 @@ def generate_hpo_result(
         time_limit: float | None = None,
         fixed_configs: list[str] | None = None,
         fit_order: Literal["original", "random"] = "random",
+        config_type: str | list[str] | None = None,
         holdout: bool = False,
         backend: Literal["ray", "native"] = "ray",
         seed: int = 0,
         **kwargs,
     ) -> pd.DataFrame:
         if repo is None:
             repo = self.load_processed(as_holdout=holdout)
-        assert self.config_type is not None
-        config_type = self.config_type
+        if config_type is None:
+            assert self.config_type is not None
+            config_type = self.config_type
         simulator = PaperRunTabArena(repo=repo, backend=backend)
         df_results_hpo = simulator.run_ensemble_config_type(
             config_type=config_type,
@@ -651,7 +653,7 @@ def generate_hpo_result(
         df_results_hpo = df_results_hpo.rename(columns={
             "framework": "method",
         })
-        df_results_hpo["method"] = f"HPO-N{n_configs}-{self.config_type}"
+        df_results_hpo["method"] = f"HPO-N{n_configs}-{config_type}"
         df_results_hpo["n_configs"] = n_configs
         df_results_hpo["n_iterations"] = n_iterations
         df_results_hpo["seed"] = seed
@@ -669,6 +671,8 @@ def generate_hpo_trajectories(
         time_limit: float | None = None,
         backend: Literal["ray", "native"] = "ray",
         holdout: bool = False,
+        config_type: str | list[str] | None = None,
+        repo: EvaluationRepository | None = None,
         cache: bool = False,
     ) -> pd.DataFrame:
         if n_configs == "auto":
@@ -687,7 +691,8 @@ def generate_hpo_trajectories(
             seeds = [i for i in range(seeds)]
 
         df_results_hpo_lst = []
-        repo = self.load_processed(as_holdout=holdout)
+        if repo is None:
+            repo = self.load_processed(as_holdout=holdout)
 
         # FIXME: Breaks for holdout, need to find a way to get self.config_default(holdout=True)
         # FIXME: Needed for TabPFN-2.5
@@ -719,6 +724,7 @@ def generate_hpo_trajectories(
                     time_limit=time_limit,
                     backend=backend,
                     holdout=holdout,
+                    config_type=config_type,
                 )
                 df_results_hpo["always_include_default"] = always_include_default
                 df_results_hpo_lst.append(df_results_hpo)

tabarena/tabarena/nips2025_utils/tabarena_context.py

@@ -100,7 +100,7 @@ def __init__(
 
         if extra_methods:
             for method_metadata in extra_methods:
-                assert method_metadata.method not in methods
+                assert method_metadata.method not in methods, f"{method_metadata.method} already in methods..."
                 methods.append(method_metadata.method)
                 method_metadata_lst.append(method_metadata)
 
@@ -205,6 +205,54 @@ def generate_repo_holdout(self, method: str) -> Path:
         repo.to_dir(path_processed)
         return path_processed
 
+    # FIXME: This is a hacky approach, refactor
+    def generate_hpo_trajectories(
+        self,
+        methods: list[str | MethodMetadata],
+        n_configs: list[int | None] | str = "auto",
+        seeds: int | list[int] = 20,
+        n_iterations: int = 40,
+        default_method: str = None,
+        always_include_default: bool = True,
+        fit_order: Literal["original", "random"] = "random",
+        time_limit: float | None = None,
+        backend: Literal["ray", "native"] = "ray",
+        repo: EvaluationRepository | None = None,
+        folds: list[int] | None = None,
+        ta_name: str = None,
+        ta_suite: str = None,
+        display_name: str = None,
+    ) -> pd.DataFrame:
+        methods: list[MethodMetadata] = [self.method_metadata(m) if isinstance(m, str) else m for m in methods]
+        if repo is None:
+            repo = self.load_repo(methods=methods)
+            if folds is not None:
+                repo = repo.subset(folds=folds)
+        if not default_method:
+            default_method = methods[0]
+        else:
+            for method in methods:
+                if method.method == default_method:
+                    default_method = method
+                    break
+        hpo_trajectory = default_method.generate_hpo_trajectories(
+            n_configs=n_configs,
+            repo=repo,
+            seeds=seeds,
+            n_iterations=n_iterations,
+            always_include_default=always_include_default,
+            fit_order=fit_order,
+            time_limit=time_limit,
+            backend=backend,
+            config_type=repo.config_types(),
+            cache=False,
+        )
+
+        hpo_trajectory["ta_name"] = ta_name
+        hpo_trajectory["ta_suite"] = ta_suite
+        hpo_trajectory["display_name"] = display_name
+        return hpo_trajectory
+
     def combine_hpo(
         self,
         methods: list[str],
@@ -213,6 +261,11 @@ def combine_hpo(
         ta_suite: str,
         method_default: str | None = None,
         repo: EvaluationRepository | None = None,
+        n_configs: int | None = None,
+        time_limit: float | None = None,
+        fit_order: Literal["original", "random"] = "original",
+        default_always_first: bool = True,
+        seed: int = 0,
     ) -> pd.DataFrame:
         """
         Perform HPO across multiple methods
@@ -237,16 +290,31 @@ def combine_hpo(
         else:
             default = None
 
+        if default_always_first and config_default:
+            fixed_configs = [config_default]
+        else:
+            fixed_configs = None
+
         tuned = self.run_hpo(
             method=methods,
             repo=repo,
             n_iterations=1,
+            n_configs=n_configs,
+            time_limit=time_limit,
+            fit_order=fit_order,
+            seed=seed,
+            fixed_configs=fixed_configs,
         )
 
         tuned_ens = self.run_hpo(
             method=methods,
             repo=repo,
             n_iterations=40,
+            n_configs=n_configs,
+            time_limit=time_limit,
+            fit_order=fit_order,
+            seed=seed,
+            fixed_configs=fixed_configs,
         )
 
         tuned["ta_name"] = ta_name