From 39a9e75edfa59ec57c9f42e16f07151f836a4dde Mon Sep 17 00:00:00 2001 From: Yikun Jiang Date: Tue, 4 Mar 2025 20:56:01 +0000 Subject: [PATCH 1/4] support hyperopt and optuna --- causaltune/optimiser.py | 108 ++++++++++++++++++++---------------- causaltune/score/scoring.py | 8 +-- causaltune/search/params.py | 21 +++++++ 3 files changed, 86 insertions(+), 51 deletions(-) diff --git a/causaltune/optimiser.py b/causaltune/optimiser.py index ff5f037e..16d90ef1 100644 --- a/causaltune/optimiser.py +++ b/causaltune/optimiser.py @@ -8,7 +8,7 @@ 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 @@ -200,7 +200,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 @@ -288,6 +288,7 @@ def fit( encoder_type: Optional[str] = None, encoder_outcome: Optional[str] = None, use_ray: Optional[bool] = None, + framework: Optional[str] = "flaml", ): """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. @@ -474,47 +475,60 @@ def fit( 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"], - ) + # TODO: intergrate resume + # 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 + + self.tuner = create_tuner( + self._tune_with_config, + search_space, + metric=self.metric, + mode=("min" if self.metric in metrics_to_minimize() else "max"), + framework=framework, + **self.cfg.parse_tuner_params(self._settings["tuner"], framework), + ) - if self.results.get_best_trial() is None: - raise Exception( - "Optimization failed! Did you set large enough time_budget and components_budget?" - ) - 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"), - low_cost_partial_config={}, - **self._settings["tuner"], - ) - # print("Optimization failed!\n", traceback.format_exc()) - # raise e + self.tuner.run() + + # 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 self.results.get_best_trial() is None: + # raise Exception( + # "Optimization failed! Did you set large enough time_budget and components_budget?" + # ) + # 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"), + # low_cost_partial_config={}, + # **self._settings["tuner"], + # ) + # # print("Optimization failed!\n", traceback.format_exc()) + # # raise e self.update_summary_scores() def update_summary_scores(self): @@ -523,7 +537,7 @@ def update_summary_scores(self): 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) @@ -699,7 +713,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): @@ -708,7 +722,7 @@ def model(self): Returns: CausalEstimator """ - return self.results.best_result["estimator"].estimator + return self.tuner.best_result["estimator"].estimator def best_model_for_estimator(self, estimator_name): """Return the best model found for a particular estimator. @@ -730,7 +744,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): @@ -749,7 +763,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 diff --git a/causaltune/score/scoring.py b/causaltune/score/scoring.py index 9aa5cd1f..6499e6be 100644 --- a/causaltune/score/scoring.py +++ b/causaltune/score/scoring.py @@ -1321,7 +1321,7 @@ 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: @@ -1333,19 +1333,19 @@ def best_score_by_estimator(scores: Dict[str, dict], metric: str) -> Dict[str, d """ - 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}" ) 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() if "estimator_name" in v and v["estimator_name"] == name + v for v in scores if "estimator_name" in v and v["estimator_name"] == name ] best[name] = ( min(est_scores, key=lambda x: x[metric]) diff --git a/causaltune/search/params.py b/causaltune/search/params.py index c75f8993..5c8d011e 100644 --- a/causaltune/search/params.py +++ b/causaltune/search/params.py @@ -7,6 +7,7 @@ import warnings from econml.inference import BootstrapInference # noqa F401 from sklearn import linear_model +from hiertunehub import SearchSpace from causaltune.utils import clean_config from causaltune.search.component import model_from_cfg, joint_config @@ -161,8 +162,28 @@ def search_space( data_size, outcome_estimator_list ) + out = SearchSpace.from_flaml(out, name="estimator_name") + return out + @staticmethod + def parse_tuner_params(params: dict, framework: str) -> dict: + if framework == "flaml": + return params + elif framework == "hyperopt": + return { + "max_evals": params["num_samples"] if params["num_samples"] != -1 else None, + "timeout": params["time_budget_s"], + "verbose": params["verbose"], + } + elif framework == "optuna": + return { + "n_trials": params["num_samples"] if params["num_samples"] != -1 else None, + "timeout": params["time_budget_s"], + } + else: + raise ValueError(f"Framework {framework} not supported") + def default_configs( self, estimator_list: Iterable[str], From 3cd34ed86b1d8cef2b1ccfcc1146eae749913d6c Mon Sep 17 00:00:00 2001 From: Yikun Jiang Date: Wed, 5 Mar 2025 15:09:20 +0000 Subject: [PATCH 2/4] pass in kwarg --- causaltune/optimiser.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/causaltune/optimiser.py b/causaltune/optimiser.py index 16d90ef1..9e980339 100644 --- a/causaltune/optimiser.py +++ b/causaltune/optimiser.py @@ -289,6 +289,7 @@ def fit( encoder_outcome: Optional[str] = None, use_ray: Optional[bool] = None, framework: Optional[str] = "flaml", + **kwargs, ): """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. @@ -492,6 +493,7 @@ def fit( mode=("min" if self.metric in metrics_to_minimize() else "max"), framework=framework, **self.cfg.parse_tuner_params(self._settings["tuner"], framework), + **kwargs, ) self.tuner.run() From 79d9c1fbd75eb97ce571c2a473e8a78c0f1d629c Mon Sep 17 00:00:00 2001 From: Yikun Jiang Date: Wed, 5 Mar 2025 15:26:51 +0000 Subject: [PATCH 3/4] can pass in algo parameter --- causaltune/optimiser.py | 5 +++-- causaltune/search/params.py | 10 +++++++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/causaltune/optimiser.py b/causaltune/optimiser.py index 9e980339..14bb08e4 100644 --- a/causaltune/optimiser.py +++ b/causaltune/optimiser.py @@ -174,6 +174,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"] = include_experimental_estimators @@ -289,7 +290,7 @@ def fit( encoder_outcome: Optional[str] = None, use_ray: Optional[bool] = None, framework: Optional[str] = "flaml", - **kwargs, + algo = 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. @@ -438,6 +439,7 @@ def fit( self._settings["tuner"]["time_budget_s"] = ( 2.5 * len(self.estimator_list) * self._settings["component_models"]["time_budget"] ) + self._settings["tuner"]["algo"] = algo cmtb = self._settings["component_models"]["time_budget"] @@ -493,7 +495,6 @@ def fit( mode=("min" if self.metric in metrics_to_minimize() else "max"), framework=framework, **self.cfg.parse_tuner_params(self._settings["tuner"], framework), - **kwargs, ) self.tuner.run() diff --git a/causaltune/search/params.py b/causaltune/search/params.py index 5c8d011e..5e3326d4 100644 --- a/causaltune/search/params.py +++ b/causaltune/search/params.py @@ -169,17 +169,25 @@ def search_space( @staticmethod def parse_tuner_params(params: dict, framework: str) -> dict: if framework == "flaml": - return params + return { + "num_samples": params["num_samples"], + "time_budget_s": params["time_budget_s"], + "verbose": params["verbose"], + "resources_per_trial": params["resources_per_trial"], + "search_alg": params["algo"], + } elif framework == "hyperopt": return { "max_evals": params["num_samples"] if params["num_samples"] != -1 else None, "timeout": params["time_budget_s"], "verbose": params["verbose"], + "algo": params["algo"], } elif framework == "optuna": return { "n_trials": params["num_samples"] if params["num_samples"] != -1 else None, "timeout": params["time_budget_s"], + "sampler": params["algo"], } else: raise ValueError(f"Framework {framework} not supported") From 6a97bfd14e26a04cf4385aa799b67872d7b70dc7 Mon Sep 17 00:00:00 2001 From: Yikun Jiang Date: Tue, 17 Jun 2025 13:01:24 +0100 Subject: [PATCH 4/4] complete docs --- causaltune/optimiser.py | 8 +++++--- setup.py | 3 +++ 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/causaltune/optimiser.py b/causaltune/optimiser.py index 14bb08e4..53936348 100644 --- a/causaltune/optimiser.py +++ b/causaltune/optimiser.py @@ -1,6 +1,6 @@ import copy import warnings -from typing import List, Optional, Union +from typing import List, Optional, Union, Any from collections import defaultdict import time @@ -290,7 +290,7 @@ def fit( encoder_outcome: Optional[str] = None, use_ray: Optional[bool] = None, framework: Optional[str] = "flaml", - algo = None, + 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. @@ -310,6 +310,8 @@ 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]): framework to use for HPO, choices are "flaml", "hyperopt" and "optuna". + algo (Optional[str]): algorithm to use for HPO, each framework has its own set of algorithms to choose from. Returns: None @@ -478,7 +480,7 @@ def fit( else [] ) - # TODO: intergrate resume + # TODO: intergrate resume and init cfg # if resume and self.results: # # pull out configs and resume_scores from previous trials: # for _, result in self.results.results.items(): diff --git a/setup.py b/setup.py index cbef4079..7b306f8b 100644 --- a/setup.py +++ b/setup.py @@ -34,6 +34,9 @@ "wise-pizza", "seaborn", "category_encoders==2.6.3", + "hiertunehub", + "hyperopt", + "optuna" ], extras_require={ "test": [