Pluggable HPO backends (flaml/hyperopt/optuna) via hiertunehub

README.md

@@ -9,7 +9,7 @@ Its estimators are taken from [EconML](https://github.com/microsoft/EconML/) aug
 [DoWhy](https://github.com/microsoft/DoWhy/) wrapper.
 
 Our contribution is enabling automatic estimator tuning and selection by out-of-sample scoring of causal estimators, notably using the [energy score](https://arxiv.org/abs/2212.10076).
-We use [FLAML](https://github.com/microsoft/FLAML) for hyperparameter optimisation.
+For hyperparameter optimisation we support pluggable backends — [Optuna](https://optuna.org) (default), [Hyperopt](https://github.com/hyperopt/hyperopt), and [FLAML](https://github.com/microsoft/FLAML) — selectable via the `framework` argument to `fit()`.
 
 We perform automated hyperparameter tuning of first stage models (for the treatment and outcome models)
 as well as hyperparameter tuning and model selection for the second stage model (causal estimator).
@@ -122,17 +122,21 @@ To install from source, see [For Developers](#for-developers) section below.
 
 ### Requirements
 
-CausalTune works with Python 3.8 and 3.9.
+CausalTune works with Python 3.10, 3.11 and 3.12.
 
 It requires the following libraries to work:
 - NumPy
 - Pandas
 - EconML
 - DoWhy
 - FLAML
+- Optuna
 - Scikit-Learn
 - Dcor
 
+Hyperopt is an optional backend, installed via the `hyperopt` extra
+(`pip install causaltune[hyperopt]`).
+
 The easiest way to install the dependencies is via
 ```
 pip install -r requirements.txt
@@ -147,7 +151,7 @@ Mac/ OS users: For some machines, it can happen that the package LightGBM which
 
 2. Set Up a Conda Environment using an appropriate Python Version
 	- Ensure Anaconda or Miniconda is installed.
-	- Create a new Conda environment: `conda create -n causaltune-env python=3.9.x`
+	- Create a new Conda environment: `conda create -n causaltune-env python=3.11.x`
 	- Activate the environment: `conda activate causaltune-env`.
 
 3. Install the dependency lightgbm seperatly before attempting to install other dependencies
@@ -190,6 +194,13 @@ print(f"Best estimator: {ct.best_estimator}")
 
 ```
 
+By default `fit()` optimises with the Optuna backend. Pass `framework="hyperopt"`
+or `framework="flaml"` to switch, and `algo=` to pick a specific sampler /
+search algorithm for that backend (e.g. an Optuna sampler, a Hyperopt suggest
+function, or a FLAML search algorithm). Warm starting via `try_init_configs` and
+resuming a previous fit with `resume=True` are currently supported only with
+`framework="flaml"`.
+
 Now if ***outcome_model="auto"*** in the CausalTune constructor, we search over a simultaneous search space for the EconML estimators and for FLAML wrappers for common regressors. The old behavior is now achieved by ***outcome_model="nested"*** (Refitting AutoML for each estimator).
 
 You can also preprocess the data in the CausalityDataset using one of the popular category encoders: ***OneHot, WoE, Label, Target***.

causaltune/optimiser.py

@@ -1,14 +1,14 @@
 import copy
 import warnings
-from typing import List, Optional, Union
+from typing import Any, List, Optional, Union
 from collections import defaultdict
 import time
 
 import traceback
 import pandas as pd
 import numpy as np
 from sklearn.linear_model import _base
-from flaml import tune
+from hiertunehub import create_tuner
 
 from sklearn.dummy import DummyClassifier
 from sklearn.model_selection import train_test_split
@@ -176,6 +176,7 @@ def __init__(
         self._settings["tuner"]["resources_per_trial"] = (
             resources_per_trial if resources_per_trial is not None else {"cpu": 0.5}
         )
+        self._settings["tuner"]["algo"] = None
         self._settings["try_init_configs"] = try_init_configs
         self._settings[
             "include_experimental_estimators"
@@ -207,7 +208,7 @@ def __init__(
         self._settings["propensity_model"] = propensity_model
         self._settings["outcome_model"] = outcome_model
 
-        self.results = None
+        self.tuner = None
         self._best_estimators = defaultdict(lambda: (float("-inf"), None))
 
         self.original_estimator_list = estimator_list
@@ -216,9 +217,6 @@ def __init__(
         self.identified_estimand = None
         self.problem = None
         self.use_ray = use_ray
-        # properties that are used to resume fits (warm start)
-        self.resume_scores = []
-        self.resume_cfg = []
 
     def get_params(self, deep=False):
         return self._settings.copy()
@@ -306,6 +304,8 @@ def fit(
         encoder_type: Optional[str] = None,
         encoder_outcome: Optional[str] = None,
         use_ray: Optional[bool] = None,
+        framework: Optional[str] = "optuna",
+        algo: Any = None,
     ):
         """Performs AutoML on list of causal inference estimators
         - If estimator has a search space specified in its parameters, HPO is performed on the whole model.
@@ -325,6 +325,13 @@ def fit(
             preprocess (bool): preprocess CausalityDataset if needed.
             encoder_type (Optional[str]): Categorical Encoder for preprocessing
             encoder_outcome (Optional[str]): Categorical Encoder target for preprocessing: TargetEncoder, WOE.
+            framework (Optional[str]): HPO backend to use, one of "optuna"
+                (default), "hyperopt" or "flaml". Only "flaml" supports
+                try_init_configs warm-start and resume; the others warn/raise.
+            algo (Any): search algorithm for the chosen backend. flaml -> a
+                FLAML search_alg; hyperopt -> a suggest function (defaults to
+                hyperopt.tpe.suggest); optuna -> an optuna sampler (defaults to
+                optuna's TPESampler). None uses each backend's default.
 
         Returns:
             None
@@ -494,70 +501,123 @@ def fit(
         search_space = self.cfg.search_space(
             self.estimator_list, data_size=data.data.shape
         )
+        # init configs (warm-start points) are only wired for the FLAML backend;
+        # for hyperopt/optuna they are a best-effort no-op (warn once).
+        if self._settings["try_init_configs"] and framework != "flaml":
+            warnings.warn(
+                "try_init_configs (init config warm-start) is only applied with "
+                f"framework='flaml'; ignored for framework='{framework}'.",
+                UserWarning,
+            )
         init_cfg = (
             self.cfg.default_configs(self.estimator_list, data_size=data.data.shape)
-            if self._settings["try_init_configs"]
+            if self._settings["try_init_configs"] and framework == "flaml"
             else []
         )
 
-        if resume and self.results:
-            # pull out configs and resume_scores from previous trials:
-            for _, result in self.results.results.items():
-                self.resume_scores.append(result[self.metric])
-                self.resume_cfg.append(result["config"])
-            # append init_cfgs that have not yet been evaluated
-            for cfg in init_cfg:
-                self.resume_cfg.append(cfg) if cfg not in self.resume_cfg else None
-        try:
-            self.results = tune.run(
-                self._tune_with_config,
-                search_space,
-                metric=self.metric,
-                # use_ray=self.use_ray,
-                cost_attr="evaluation_cost",
-                points_to_evaluate=(
-                    init_cfg if len(self.resume_cfg) == 0 else self.resume_cfg
-                ),
-                evaluated_rewards=(
-                    [] if len(self.resume_scores) == 0 else self.resume_scores
-                ),
-                mode=("min" if self.metric in metrics_to_minimize() else "max"),
-                # resources_per_trial= {"cpu": 1} if self.use_ray else None,
-                low_cost_partial_config={},
-                **self._settings["tuner"],
+        if resume and framework != "flaml":
+            raise NotImplementedError(
+                "resume is currently only supported with framework='flaml', "
+                f"not framework='{framework}'."
             )
 
-            if self.results.get_best_trial() is None:
-                raise Exception(
-                    "Optimization failed! Did you set large enough time_budget and components_budget?"
+        if framework == "hyperopt" and not self._search_space_has_tunable_params():
+            raise ValueError(
+                "framework='hyperopt' needs at least one tunable hyperparameter "
+                "in the search space, but all selected estimators are "
+                "parameterless and outcome_model is not 'auto'. (hiertunehub "
+                "0.2.1's hyperopt conversion raises on a fully parameterless "
+                "search.) Use framework='optuna' or 'flaml', include a "
+                "parameterized estimator, or set outcome_model='auto'."
+            )
+
+        self._settings["tuner"]["algo"] = algo
+        mode = "min" if self.metric in metrics_to_minimize() else "max"
+        framework_params = self.cfg.parse_tuner_params(
+            self._settings["tuner"], framework
+        )
+
+        if framework == "flaml":
+            # Restore full FLAML parity: cost-aware search plus warm-start /
+            # resume seeds. Capture resume points/rewards from the PREVIOUS tuner
+            # before it is overwritten below.
+            if resume and self.tuner is not None:
+                points_to_evaluate, evaluated_rewards = self._resume_points_and_rewards(
+                    self.tuner.results, init_cfg
                 )
-        except Exception:
-            # we must have an older FLAML version that doesn't support the cost_attr parameter
-            self.results = tune.run(
-                self._tune_with_config,
-                search_space,
-                metric=self.metric,
-                points_to_evaluate=(
-                    init_cfg if len(self.resume_cfg) == 0 else self.resume_cfg
-                ),
-                evaluated_rewards=(
-                    [] if len(self.resume_scores) == 0 else self.resume_scores
-                ),
-                mode=("min" if self.metric in metrics_to_minimize() else "max"),
+            else:
+                points_to_evaluate, evaluated_rewards = init_cfg, []
+            framework_params.update(
+                cost_attr="evaluation_cost",
                 low_cost_partial_config={},
-                **self._settings["tuner"],
+                points_to_evaluate=points_to_evaluate,
+                evaluated_rewards=evaluated_rewards,
             )
-            # print("Optimization failed!\n", traceback.format_exc())
-            # raise e
+
+        self.tuner = create_tuner(
+            self._tune_with_config,
+            search_space,
+            metric=self.metric,
+            mode=mode,
+            framework=framework,
+            framework_params=framework_params,
+        )
+        self.tuner.run()
+
         self.update_summary_scores()
 
+    def _resume_points_and_rewards(self, prev_results, init_cfg):
+        """Rebuild FLAML warm-start seeds from a previous tuner's results.
+
+        Mirrors the original resume semantics: for each prior trial that carries
+        both the metric and its config, seed ``(config -> reward)``; then append
+        any init configs not already present (without a reward, so FLAML
+        evaluates them).
+
+        Args:
+            prev_results (list[dict]): ``tuner.results`` from the previous fit.
+            init_cfg (list[dict]): init configs to append if not yet evaluated.
+
+        Returns:
+            tuple[list[dict], list]: ``(points_to_evaluate, evaluated_rewards)``
+            with rewards aligned to the leading resumed configs.
+        """
+        resume_cfg = []
+        resume_scores = []
+        for result in prev_results:
+            if self.metric not in result or "config" not in result:
+                continue
+            resume_scores.append(result[self.metric])
+            resume_cfg.append(result["config"])
+        for cfg in init_cfg:
+            if cfg not in resume_cfg:
+                resume_cfg.append(cfg)
+        return resume_cfg, resume_scores
+
+    def _search_space_has_tunable_params(self) -> bool:
+        """Whether the current search space contains any tunable hyperparameter.
+
+        Sampling outcome estimators (outcome_model='auto') always adds tunable
+        component-model params; otherwise it depends on the estimators' own
+        search spaces. Used to guard the hyperopt backend, whose hiertunehub
+        converter raises on a fully parameterless search space.
+        """
+        if self.cfg.sample_outcome_estimators:
+            return True
+        configs = self.cfg._configs()
+        return any(
+            bool(configs[est].search_space)
+            for est in self.estimator_list
+            if est in configs
+        )
+
     def update_summary_scores(self):
         """Stores scores for metric of interest for each estimator
 
         Returns:
             None
         """
-        self.scores = Scorer.best_score_by_estimator(self.results.results, self.metric)
+        self.scores = Scorer.best_score_by_estimator(self.tuner.results, self.metric)
         # now inject the separately saved model objects
         for est_name in self.scores:
             # Todo: Check approximate scores for OrthoIV (possibly other IV estimators)
@@ -710,8 +770,14 @@ def _estimate_effect(self, config):
             }
         except Exception as e:
             print("Evaluation failed!\n", config, traceback.format_exc())
+            # Use the *worst* value for the metric direction as the failure
+            # sentinel, so a failed trial is never picked as best. For minimized
+            # metrics (e.g. energy_distance) that is +inf; for maximized metrics
+            # it is -inf. (A flat -inf would look optimal to a minimizing backend
+            # such as the default optuna, poisoning best_estimator selection.)
+            worst = np.inf if self.metric in metrics_to_minimize() else -np.inf
             return {
-                self.metric: -np.inf,
+                self.metric: worst,
                 "estimator_name": self.estimator_name,
                 "exception": e,
                 "traceback": traceback.format_exc(),
@@ -750,7 +816,7 @@ def best_estimator(self) -> str:
         Returns:
             None
         """
-        return self.results.best_result["estimator_name"]
+        return self.tuner.best_result["estimator_name"]
 
     @property
     def model(self):
@@ -759,7 +825,10 @@ def model(self):
         Returns:
             CausalEstimator
         """
-        return self.results.best_result["estimator"].estimator
+        # The objective pops the fitted estimator out of the result dict in the
+        # non-store_all path, so resolve it from the scores table (populated by
+        # update_summary_scores) rather than from tuner.best_result.
+        return self.scores[self.best_estimator]["estimator"].estimator
 
     def best_model_for_estimator(self, estimator_name):
         """Return the best model found for a particular estimator.
@@ -781,7 +850,7 @@ def best_config(self):
         Returns:
             (dict): the best configuration
         """
-        return self.results.best_config
+        return self.tuner.best_params
 
     @property
     def best_config_per_estimator(self):
@@ -800,7 +869,7 @@ def best_score(self):
         """
         Returns:
             (float):  the best score found."""
-        return self.results.best_result[self.metric]
+        return self.tuner.best_result[self.metric]
 
     def effect(self, df, *args, **kwargs):
         """Heterogeneous Treatment Effects for data df

causaltune/score/scoring.py

@@ -1307,43 +1307,33 @@ def make_scores(
         return out
 
     @staticmethod
-    def best_score_by_estimator(
-        scores: Dict[str, dict], metric: str
-    ) -> Dict[str, dict]:
+    def best_score_by_estimator(scores: List[dict], metric: str) -> Dict[str, dict]:
         """Obtain best score for each estimator.
 
         Args:
-            scores (Dict[str, dict]): CausalTune.scores dictionary
+            scores (List[dict]): list of per-trial result dicts (as produced by
+                the tuner), each carrying an ``estimator_name`` field
             metric (str): metric of interest
 
         Returns:
             Dict[str, dict]: dictionary containing best score by estimator
 
         """
 
-        for k, v in scores.items():
+        for v in scores:
             if "estimator_name" not in v:
                 raise ValueError(
-                    f"Malformed scores dict, 'estimator_name' field missing "
-                    f"in{k}, {v}"
+                    f"Malformed scores entry, 'estimator_name' field missing " f"in {v}"
                 )
 
         estimator_names = sorted(
-            list(
-                set(
-                    [
-                        v["estimator_name"]
-                        for v in scores.values()
-                        if "estimator_name" in v
-                    ]
-                )
-            )
+            list(set([v["estimator_name"] for v in scores if "estimator_name" in v]))
         )
         best = {}
         for name in estimator_names:
             est_scores = [
                 v
-                for v in scores.values()
+                for v in scores
                 if "estimator_name" in v and v["estimator_name"] == name
             ]
             best[name] = (