diff --git a/causaltune/dataset_processor.py b/causaltune/dataset_processor.py index 08b0b31c..11ef5313 100644 --- a/causaltune/dataset_processor.py +++ b/causaltune/dataset_processor.py @@ -14,7 +14,10 @@ def __init__(self): self.encoder = None def fit( - self, cd: CausalityDataset, encoder_type: Optional[str] = "onehot", outcome: str = None + self, + cd: CausalityDataset, + encoder_type: Optional[str] = "onehot", + outcome: str = None, ): cd = copy.deepcopy(cd) self.preprocess_dataset( diff --git a/causaltune/models/regression.py b/causaltune/models/regression.py new file mode 100644 index 00000000..28aa53a8 --- /dev/null +++ b/causaltune/models/regression.py @@ -0,0 +1,86 @@ +from sklearn.linear_model import ElasticNet, LassoLars + + +from flaml.automl.model import SKLearnEstimator +from flaml import tune + +# These models are for some reason not in the deployed version of flaml 2.2.0, +# but in the source code they are there +# So keep this file in the project for now + + +class ElasticNetEstimator(SKLearnEstimator): + """The class for tuning Elastic Net regression model.""" + + """Reference: https://scikit-learn.org/stable/modules/generated/sklearn.linear_model.ElasticNet.html""" + + ITER_HP = "max_iter" + + @classmethod + def search_space(cls, data_size, task="regresssion", **params): + return { + "alpha": { + "domain": tune.loguniform(lower=0.0001, upper=1.0), + "init_value": 0.1, + }, + "l1_ratio": { + "domain": tune.uniform(lower=0.0, upper=1.0), + "init_value": 0.5, + }, + "selection": { + "domain": tune.choice(["cyclic", "random"]), + "init_value": "cyclic", + }, + } + + def config2params(self, config: dict) -> dict: + params = super().config2params(config) + params["tol"] = params.get("tol", 0.0001) + if "n_jobs" in params: + params.pop("n_jobs") + return params + + def __init__(self, task="regression", **config): + super().__init__(task, **config) + assert self._task.is_regression(), "ElasticNet for regression task only" + self.estimator_class = ElasticNet + + +class LassoLarsEstimator(SKLearnEstimator): + """The class for tuning Lasso model fit with Least Angle Regression a.k.a. Lars.""" + + """Reference: https://scikit-learn.org/stable/modules/generated/sklearn.linear_model.LassoLars.html""" + + ITER_HP = "max_iter" + + @classmethod + def search_space(cls, task=None, **params): + return { + "alpha": { + "domain": tune.loguniform(lower=1e-4, upper=1.0), + "init_value": 0.1, + }, + "fit_intercept": { + "domain": tune.choice([True, False]), + "init_value": True, + }, + "eps": { + "domain": tune.loguniform(lower=1e-16, upper=1e-4), + "init_value": 2.220446049250313e-16, + }, + } + + def config2params(self, config: dict) -> dict: + params = super().config2params(config) + if "n_jobs" in params: + params.pop("n_jobs") + return params + + def __init__(self, task="regression", **config): + super().__init__(task, **config) + assert self._task.is_regression(), "LassoLars for regression task only" + self.estimator_class = LassoLars + + def predict(self, X, **kwargs): + X = self._preprocess(X) + return self._model.predict(X, **kwargs) diff --git a/causaltune/optimiser.py b/causaltune/optimiser.py index 08b921ae..9331295c 100644 --- a/causaltune/optimiser.py +++ b/causaltune/optimiser.py @@ -1,5 +1,4 @@ import copy -from copy import deepcopy import warnings from typing import List, Optional, Union from collections import defaultdict @@ -19,10 +18,9 @@ from joblib import Parallel, delayed -from causaltune.params import SimpleParamService +from causaltune.search.params import SimpleParamService from causaltune.scoring import Scorer -from causaltune.r_score import RScoreWrapper -from causaltune.utils import clean_config, treatment_is_multivalue +from causaltune.utils import treatment_is_multivalue from causaltune.models.monkey_patches import ( AutoML, apply_multitreatment, @@ -96,7 +94,7 @@ def __init__( test_size=None, num_samples=-1, propensity_model="dummy", - outcome_model=None, + outcome_model="nested", components_task="regression", components_verbose=0, components_pred_time_limit=10 / 1e6, @@ -181,9 +179,9 @@ def __init__( resources_per_trial if resources_per_trial is not None else {"cpu": 0.5} ) self._settings["try_init_configs"] = try_init_configs - self._settings["include_experimental_estimators"] = ( - include_experimental_estimators - ) + self._settings[ + "include_experimental_estimators" + ] = include_experimental_estimators # params for FLAML on component models: self._settings["component_models"] = {} @@ -246,33 +244,44 @@ def init_propensity_model(self, propensity_model: str): ) def init_outcome_model(self, outcome_model): + # TODO: implement filtering like below, when there are propensity-only features + # feature_filter below acts on classes not instances + # to preserve all the extra methods through inheritance # if we are only supplying certain features to the propensity function, # make them invisible to the outcome component model # This is a workaround for the DoWhy/EconML data model which doesn't # support that out of the box - if outcome_model is not None: - # TODO: implement filtering like below, when there are propensity-only features - # feature_filter below acts on classes not instances - # to preserve all the extra methods through inheritance - self.outcome_model = outcome_model + + if hasattr(outcome_model, "fit") and hasattr(outcome_model, "predict"): + return outcome_model + elif outcome_model == "auto": + # Will be dynamically chosen at optimization time + return outcome_model + elif outcome_model == "nested": + # The current default behavior + return self.auto_outcome_model() else: - data = self.data - propensity_only_cols = [ - p - for p in data.propensity_modifiers - if p not in data.common_causes + data.effect_modifiers - ] - - if len(propensity_only_cols): - outcome_model_class = feature_filter( - AutoML, data.effect_modifiers + data.common_causes, first_cols=True - ) - else: - outcome_model_class = AutoML + raise ValueError( + 'outcome_model valid values are None, "auto", or an estimator object' + ) - self.outcome_model = outcome_model_class( - **self._settings["component_models"] + def auto_outcome_model(self): + data = self.data + propensity_only_cols = [ + p + for p in data.propensity_modifiers + if p not in data.common_causes + data.effect_modifiers + ] + + if len(propensity_only_cols): + # TODO: implement feature_filter for arbitrary outcome models + outcome_model_class = feature_filter( + AutoML, data.effect_modifiers + data.common_causes, first_cols=True ) + else: + outcome_model_class = AutoML + + return outcome_model_class(**self._settings["component_models"]) def fit( self, @@ -363,7 +372,6 @@ def fit( ) self.init_propensity_model(self._settings["propensity_model"]) - self.init_outcome_model(self._settings["outcome_model"]) self.identified_estimand: IdentifiedEstimand = ( self.causal_model.identify_effect(proceed_when_unidentifiable=True) @@ -406,11 +414,10 @@ def fit( # config with method-specific params self.cfg = SimpleParamService( - self.propensity_model, - self.outcome_model, n_jobs=self._settings["component_models"]["n_jobs"], include_experimental=self._settings["include_experimental_estimators"], multivalue=treatment_is_multivalue(self._treatment_values), + sample_outcome_estimators=self._settings["outcome_model"] == "auto", ) self.estimator_list = self.cfg.estimator_names_from_patterns( @@ -454,24 +461,30 @@ def fit( if self._settings["test_size"] is not None: self.test_df = self.test_df.sample(self._settings["test_size"]) - self.r_scorer = ( - None - if "r_scorer" not in self.metrics_to_report - else RScoreWrapper( - self.outcome_model, - self.propensity_model, - self.train_df, - self.test_df, - outcome, - treatment, - common_causes, - effect_modifiers, + if "r_scorer" in self.metrics_to_report: + raise NotImplementedError( + "R-squared scorer no longer suported, please raise an issue if you want it back" ) + # self.r_scorer = ( + # None + # if "r_scorer" not in self.metrics_to_report + # else RScoreWrapper( + # self.outcome_model, + # self.propensity_model, + # self.train_df, + # self.test_df, + # outcome, + # treatment, + # common_causes, + # effect_modifiers, + # ) + # ) + + search_space = self.cfg.search_space( + self.estimator_list, data_size=data.data.shape ) - - search_space = self.cfg.search_space(self.estimator_list) init_cfg = ( - self.cfg.default_configs(self.estimator_list) + self.cfg.default_configs(self.estimator_list, data_size=data.data.shape) if self._settings["try_init_configs"] else [] ) @@ -543,7 +556,7 @@ def _tune_with_config(self, config: dict) -> dict: # to spawn a separate process to prevent cross-talk between tuner and automl on component models: estimates = Parallel(n_jobs=2, backend="threading")( - delayed(self._estimate_effect)(config["estimator"]) for i in range(1) + delayed(self._estimate_effect)(config) for i in range(1) )[0] # estimates = self._estimate_effect(config["estimator"]) @@ -582,19 +595,15 @@ def _est_effect_stub(self, method_params): def _estimate_effect(self, config): """estimates effect with chosen estimator""" - # add params that are tuned by flaml: - config = clean_config(copy.copy(config)) - self.estimator_name = config.pop("estimator_name") - # params_to_tune = { - # k: v for k, v in config.items() if (not k == "estimator_name") - # } - cfg = self.cfg.method_params(self.estimator_name) - method_params = { - "init_params": {**deepcopy(config), **cfg.init_params}, - "fit_params": {}, - } + # Do we need an boject property for this, instead of a local var? + self.estimator_name = config["estimator"]["estimator_name"] + outcome_model = self.init_outcome_model(self._settings["outcome_model"]) + method_params = self.cfg.method_params( + config, outcome_model, self.propensity_model + ) + try: # - # if True: # + # This calls the causal model's estimate_effect method estimate = self._est_effect_stub(method_params) scores = { "estimator_name": self.estimator_name, @@ -611,8 +620,9 @@ def _estimate_effect(self, config): return { self.metric: scores["validation"][self.metric], "estimator": estimate, - "estimator_name": scores.pop("estimator_name"), + "estimator_name": self.estimator_name, "scores": scores, + # TODO: return full config! "config": config, } except Exception as e: @@ -641,7 +651,14 @@ def score_dataset(self, df: pd.DataFrame, dataset_name: str): None. """ for scr in self.scores.values(): - scr["scores"][dataset_name] = self._compute_metrics(scr["estimator"], df) + if scr["estimator"] is None: + warnings.warn( + "Skipping scoring for estimator %s", scr["estimator_name"] + ) + else: + scr["scores"][dataset_name] = self._compute_metrics( + scr["estimator"], df + ) @property def best_estimator(self) -> str: @@ -788,21 +805,20 @@ def effect_stderr(self, df, n_bootstrap_samples=5, n_jobs=1, *args, **kwargs): if "Econml" in str(type(self.model)): # Get a list of "Inference" objects from EconML, one per treatment self.model.__class__.effect_stderr = effect_stderr - cfg = self.cfg.method_params(self.best_estimator) + outcome_model = self.init_outcome_model(self._settings["outcome_model"]) + method_params = self.cfg.method_params( + self.best_config, outcome_model, self.propensity_model + ) - if cfg.inference == "bootstrap": + if self.cfg.full_config(self.best_estimator).inference == "bootstrap": # TODO: before bootstrapping, check whether that's already been done bootstrap = BootstrapInference( n_bootstrap_samples=n_bootstrap_samples, n_jobs=n_jobs ) - - best_cfg = { - k: v for k, v in self.best_config.items() if k not in ["estimator"] - } - method_params = { - "init_params": {**best_cfg, **cfg.init_params}, - "fit_params": {"inference": bootstrap}, - } + method_params["fit_params"]["inference"] = bootstrap + self.estimator_name = ( + self.best_estimator + ) # needed for _est_effect_stub, just in case self.bootstrapped_estimate = self._est_effect_stub(method_params) est = self.bootstrapped_estimate.estimator else: diff --git a/causaltune/search/__init__.py b/causaltune/search/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/causaltune/search/component.py b/causaltune/search/component.py new file mode 100644 index 00000000..f3920aab --- /dev/null +++ b/causaltune/search/component.py @@ -0,0 +1,147 @@ +import warnings +from typing import Tuple +import copy + +import numpy as np +import pandas as pd + +from flaml import tune +from flaml.automl.model import ( + KNeighborsEstimator, + XGBoostSklearnEstimator, + XGBoostLimitDepthEstimator, + RandomForestEstimator, + LGBMEstimator, + CatBoostEstimator, + ExtraTreesEstimator, +) +from flaml.automl.task.factory import task_factory +import flaml + +from causaltune.models.regression import ElasticNetEstimator, LassoLarsEstimator + + +def flaml_config_to_tune_config(flaml_config: dict) -> Tuple[dict, dict, dict]: + cfg = {} + init_params = {} + low_cost_init_params = {} + for key, value in flaml_config.items(): + if isinstance(value["domain"], dict): + raise NotImplementedError("Nested dictionaries are not supported yet") + cfg[key] = value["domain"] + if "init_value" in value: + init_params[key] = value["init_value"] + if "low_cost_init_value" in value: + low_cost_init_params[key] = value["low_cost_init_value"] + + return cfg, init_params, low_cost_init_params + + +estimators = { + "elastic_net": ElasticNetEstimator, + "lasso_lars": LassoLarsEstimator, + "knn": KNeighborsEstimator, + "xgboost": XGBoostSklearnEstimator, + "xgboost_limit_depth": XGBoostLimitDepthEstimator, + "random_forest": RandomForestEstimator, + "lgbm": LGBMEstimator, + "catboost": CatBoostEstimator, + "extra_trees": ExtraTreesEstimator, +} + + +def joint_config(data_size: Tuple[int, int], estimator_list=None): + joint_cfg = [] + joint_init_params = [] + joint_low_cost_init_params = {} + for name, cls in estimators.items(): + if estimator_list is not None and name not in estimator_list: + continue + task = task_factory("regression") + cfg, init_params, low_cost_init_params = flaml_config_to_tune_config( + cls.search_space(data_size=data_size, task=task) + ) + + # Test if the estimator instantiates fine + try: + cls(task=task, **init_params) + cfg["estimator_name"] = name + joint_cfg.append(cfg) + init_params["estimator_name"] = name + joint_init_params.append(init_params) + joint_low_cost_init_params[name] = low_cost_init_params + except ImportError as e: + print(f"Error instantiating {name}: {e}") + + return tune.choice(joint_cfg), joint_init_params, joint_low_cost_init_params + + +def model_from_cfg(cfg: dict): + cfg = copy.deepcopy(cfg) + model_name = cfg.pop("estimator_name") + estimator_class = estimators[model_name] + + # Some Econml estimators pass a weights vector as an unnamed third argument, + # which is not supported by flaml. We need to wrap the estimator to ignore + # TODO: expose better estimator wrappers that support weights + class FlamlEstimatorWrapper(estimator_class): + wrapped_class = estimator_class + + def fit(self, X, y, *args, **kwargs): + if len(kwargs): + warnings.warn(f"Extra args {args} {kwargs} are being ignored") + return self.wrapped_class.fit(self, X, y) + + out = FlamlEstimatorWrapper(task=task_factory("regression"), **cfg) + return out + + +def config2score(cfg: dict, X, y): + model = model_from_cfg(cfg["estimator"]) + model.fit(X, y) + ypred = model.predict(X) + err = y - ypred + return {"score": np.mean(err**2)} + + +def make_fake_data(): + + # Set random seed for reproducibility + np.random.seed(42) + + # Parameters for the DataFrame + num_samples = 1000 # Number of rows (samples) + num_features = 5 # Number of features (columns) + + # Generate random float features + X = np.random.rand(num_samples, num_features) + + # Define the coefficients for each feature to generate the target variable + coefficients = np.random.rand(num_features) + + # Generate the target variable y as a linear combination of the features plus some noise + noise = np.random.normal(0, 0.1, num_samples) # Add some Gaussian noise + y = np.dot(X, coefficients) + noise + + # Create a DataFrame + column_names = [f"feature_{i + 1}" for i in range(num_features)] + df = pd.DataFrame(X, columns=column_names) + + return df, y + + +if __name__ == "__main__": + + # Create fake data + X, y = make_fake_data() + cfg, init_params, low_cost_init_params = joint_config(data_size=X.shape) + flaml.tune.run( + evaluation_function=lambda cfgs: config2score(cfgs, X, y), + metric="score", + mode="min", + config={"estimator": cfg}, + points_to_evaluate=init_params, + num_samples=10, + ) + + print("yay!") diff --git a/causaltune/params.py b/causaltune/search/params.py similarity index 78% rename from causaltune/params.py rename to causaltune/search/params.py index 36403d0f..c75f8993 100644 --- a/causaltune/params.py +++ b/causaltune/search/params.py @@ -1,15 +1,22 @@ +import numpy as np from flaml import tune from copy import deepcopy -from typing import Optional, Sequence, Union, Iterable, Dict +from typing import Optional, Sequence, Union, Iterable, Dict, Any, Tuple from dataclasses import dataclass, field import warnings from econml.inference import BootstrapInference # noqa F401 from sklearn import linear_model +from causaltune.utils import clean_config +from causaltune.search.component import model_from_cfg, joint_config + @dataclass class EstimatorConfig: + outcome_model_name: str = None + final_model_name: str = None + propensity_model_name: str = None init_params: dict = field(default_factory=dict) fit_params: dict = field(default_factory=dict) search_space: dict = field(default_factory=dict) @@ -22,21 +29,17 @@ class EstimatorConfig: class SimpleParamService: def __init__( self, - propensity_model, - outcome_model, multivalue: bool, - final_model=None, n_bootstrap_samples: Optional[int] = None, n_jobs: Optional[int] = None, include_experimental=False, + sample_outcome_estimators: bool = False, ): - self.propensity_model = propensity_model - self.outcome_model = outcome_model - self.final_model = final_model self.n_jobs = n_jobs self.include_experimental = include_experimental self.n_bootstrap_samples = n_bootstrap_samples self.multivalue = multivalue + self.sample_outcome_estimators = sample_outcome_estimators def estimator_names_from_patterns( self, @@ -107,7 +110,7 @@ def problem_match(est_name: str, problem: str) -> bool: assert isinstance(p, str) except Exception: raise ValueError( - "Invalid estimator list, must be 'auto', 'all', or a list of strings" + "Invalid estimator list, must be 'auto', 'all', 'cheap_inference' or a list of strings" ) out = [ @@ -129,7 +132,12 @@ def estimator_names(self): else: return [est for est, cfg in cfgs.items() if not cfg.experimental] - def search_space(self, estimator_list: Iterable[str]): + def search_space( + self, + estimator_list: Iterable[str], + data_size: Tuple[int, int], + outcome_estimator_list: Iterable[str] = None, + ): """Constructs search space with estimators and their respective configs Args: @@ -147,9 +155,21 @@ def search_space(self, estimator_list: Iterable[str]): if est in estimator_list ] - return {"estimator": tune.choice(search_space)} + out = {"estimator": tune.choice(search_space)} + if self.sample_outcome_estimators: + out["outcome_estimator"], _, _ = joint_config( + data_size, outcome_estimator_list + ) - def default_configs(self, estimator_list: Iterable[str]): + return out + + def default_configs( + self, + estimator_list: Iterable[str], + data_size: Tuple[int, int], + outcome_estimator_list: Iterable[str] = None, + num_outcome_samples: int = 3, + ): """Creates list with initial configs to try before moving on to hierarchical HPO. The list has been identified by evaluating performance of all @@ -164,24 +184,85 @@ def default_configs(self, estimator_list: Iterable[str]): Returns: list: list of dicts with promising initial configs """ - points = [ + pre_points = [ {"estimator": {"estimator_name": est, **est_params.defaults}} for est, est_params in self._configs().items() if est in estimator_list ] - print("Initial configs:", points) + cfgs = self._configs() + + if self.sample_outcome_estimators: + points = [] + _, init_params, _ = joint_config(data_size, outcome_estimator_list) + for p in pre_points: + if cfgs[p["estimator"]["estimator_name"]].outcome_model_name is None: + this_p = deepcopy(p) + # this won't have any effect, so pick any valid config to mitigate sampling bias + this_p["outcome_estimator"] = np.random.choice(init_params) + points.append(p) + continue + else: # Sample different outcome functions for first pass + for outcome_est in np.random.choice( + init_params, size=num_outcome_samples, replace=False + ): + this_p = deepcopy(p) + this_p["outcome_estimator"] = outcome_est + points.append(this_p) + else: + points = pre_points + return points def method_params( self, - estimator: str, + config: dict, + outcome_model: Any, + propensity_model: Any, + final_model: Any = None, ): - return self._configs()[estimator] + est_config = clean_config(deepcopy(config["estimator"])) + estimator_name = est_config.pop("estimator_name") + + cfg = self._configs()[estimator_name] + + if outcome_model == "auto" and cfg.outcome_model_name is not None: + # Spawn the outcome model dynamically + outcome_model = model_from_cfg(config["outcome_estimator"]) + + if ( + cfg.outcome_model_name is not None + and cfg.outcome_model_name not in cfg.init_params + ): + cfg.init_params[cfg.outcome_model_name] = deepcopy(outcome_model) + + if ( + cfg.propensity_model_name is not None + and cfg.propensity_model_name not in cfg.init_params + ): + cfg.init_params[cfg.propensity_model_name] = deepcopy(propensity_model) + + if ( + cfg.final_model_name is not None + and cfg.final_model_name not in cfg.init_params + ): + cfg.init_params[cfg.final_model_name] = ( + deepcopy(final_model) + if final_model is not None + else deepcopy(outcome_model) + ) + + method_params = { + "init_params": {**deepcopy(est_config), **cfg.init_params}, + "fit_params": {}, + } + return method_params + + def full_config(self, estimator_name: str): + cfg = self._configs()[estimator_name] + return cfg def _configs(self) -> Dict[str, EstimatorConfig]: - propensity_model = deepcopy(self.propensity_model) - outcome_model = deepcopy(self.outcome_model) if self.n_bootstrap_samples is not None: # TODO Egor please look into this # bootstrap is causing recursion errors (see notes below) @@ -190,48 +271,39 @@ def _configs(self) -> Dict[str, EstimatorConfig]: # ) pass - if self.final_model is None: - final_model = deepcopy(self.outcome_model) - else: - final_model = deepcopy(self.final_model) - configs: dict[str:EstimatorConfig] = { "backdoor.causaltune.models.NaiveDummy": EstimatorConfig(), "backdoor.causaltune.models.Dummy": EstimatorConfig( - init_params={"propensity_model": propensity_model}, + propensity_model_name="propensity_model", experimental=False, ), "backdoor.propensity_score_weighting": EstimatorConfig( - init_params={"propensity_model": propensity_model}, + propensity_model_name="propensity_model", experimental=True, ), "backdoor.econml.metalearners.SLearner": EstimatorConfig( - init_params={"overall_model": outcome_model}, + outcome_model_name="overall_model", supports_multivalue=True, ), "backdoor.econml.metalearners.TLearner": EstimatorConfig( - init_params={"models": outcome_model}, + outcome_model_name="models", supports_multivalue=True, ), "backdoor.econml.metalearners.XLearner": EstimatorConfig( - init_params={ - "propensity_model": propensity_model, - "models": outcome_model, - }, + outcome_model_name="models", + propensity_model_name="propensity_model", supports_multivalue=True, ), "backdoor.econml.metalearners.DomainAdaptationLearner": EstimatorConfig( - init_params={ - "propensity_model": propensity_model, - "models": outcome_model, - "final_models": final_model, - }, + outcome_model_name="models", + propensity_model_name="propensity_model", + final_model_name="final_models", supports_multivalue=True, ), "backdoor.econml.dr.ForestDRLearner": EstimatorConfig( + outcome_model_name="model_regression", + propensity_model_name="model_propensity", init_params={ - "model_propensity": propensity_model, - "model_regression": outcome_model, # putting these here for now, until default values can be reconciled with search space "mc_iters": None, "max_depth": None, @@ -268,9 +340,9 @@ def _configs(self) -> Dict[str, EstimatorConfig]: inference="blb", ), "backdoor.econml.dr.LinearDRLearner": EstimatorConfig( + outcome_model_name="model_regression", + propensity_model_name="model_propensity", init_params={ - "model_propensity": propensity_model, - "model_regression": outcome_model, "mc_iters": None, }, search_space={ @@ -286,9 +358,9 @@ def _configs(self) -> Dict[str, EstimatorConfig]: inference="auto", ), "backdoor.econml.dr.SparseLinearDRLearner": EstimatorConfig( + outcome_model_name="model_regression", + propensity_model_name="model_propensity", init_params={ - "model_propensity": propensity_model, - "model_regression": outcome_model, "mc_iters": None, }, search_space={ @@ -314,9 +386,9 @@ def _configs(self) -> Dict[str, EstimatorConfig]: inference="auto", ), "backdoor.econml.dml.LinearDML": EstimatorConfig( + outcome_model_name="model_y", + propensity_model_name="model_t", init_params={ - "model_t": propensity_model, - "model_y": outcome_model, "discrete_treatment": True, # it runs out of memory fast if the below is not set "linear_first_stages": False, @@ -335,9 +407,9 @@ def _configs(self) -> Dict[str, EstimatorConfig]: inference="statsmodels", ), "backdoor.econml.dml.SparseLinearDML": EstimatorConfig( + outcome_model_name="model_y", + propensity_model_name="model_t", init_params={ - "model_t": propensity_model, - "model_y": outcome_model, "discrete_treatment": True, # it runs out of memory fast if the below is not set "linear_first_stages": False, @@ -364,9 +436,9 @@ def _configs(self) -> Dict[str, EstimatorConfig]: inference="auto", ), "backdoor.econml.dml.CausalForestDML": EstimatorConfig( + outcome_model_name="model_y", + propensity_model_name="model_t", init_params={ - "model_t": propensity_model, - "model_y": outcome_model, # "max_depth": self.max_depth, # "n_estimators": self.n_estimators, "discrete_treatment": True, @@ -416,10 +488,8 @@ def _configs(self) -> Dict[str, EstimatorConfig]: inference="auto", ), "backdoor.causaltune.models.TransformedOutcome": EstimatorConfig( - init_params={ - "propensity_model": propensity_model, - "outcome_model": outcome_model, - }, + outcome_model_name="outcome_model", + propensity_model_name="propensity_model", ), # leaving out DML and NonParamDML as they're base classes for the 3 # above @@ -434,8 +504,8 @@ def _configs(self) -> Dict[str, EstimatorConfig]: # "fit_params": {}, # }, "backdoor.econml.orf.DROrthoForest": EstimatorConfig( + propensity_model_name="propensity_model", init_params={ - "propensity_model": propensity_model, "model_Y": linear_model.Ridge( alpha=0.01 ), # WeightedLasso(alpha=0.01), # @@ -465,8 +535,8 @@ def _configs(self) -> Dict[str, EstimatorConfig]: inference="blb", ), "backdoor.econml.orf.DMLOrthoForest": EstimatorConfig( + propensity_model_name="model_T", init_params={ - "model_T": propensity_model, "model_Y": linear_model.Ridge( alpha=0.01 ), # WeightedLasso(alpha=0.01), # @@ -499,20 +569,16 @@ def _configs(self) -> Dict[str, EstimatorConfig]: inference="blb", ), "iv.econml.iv.dr.LinearDRIV": EstimatorConfig( - init_params={ - "model_y_xw": outcome_model, - "model_t_xw": propensity_model, - }, + outcome_model_name="model_y_xw", + propensity_model_name="model_t_xw", search_space={ "projection": tune.choice([0, 1]), }, defaults={"projection": True}, ), "iv.econml.iv.dml.OrthoIV": EstimatorConfig( - init_params={ - "model_y_xw": outcome_model, - "model_t_xw": propensity_model, - }, + outcome_model_name="model_y_xw", + propensity_model_name="model_t_xw", search_space={ "mc_agg": tune.choice(["mean", "median"]), }, @@ -521,10 +587,8 @@ def _configs(self) -> Dict[str, EstimatorConfig]: }, ), "iv.econml.iv.dml.DMLIV": EstimatorConfig( - init_params={ - "model_y_xw": outcome_model, - "model_t_xw": propensity_model, - }, + outcome_model_name="model_y_xw", + propensity_model_name="model_t_xw", search_space={ "mc_agg": tune.choice(["mean", "median"]), }, @@ -533,10 +597,8 @@ def _configs(self) -> Dict[str, EstimatorConfig]: }, ), "iv.econml.iv.dr.SparseLinearDRIV": EstimatorConfig( - init_params={ - "model_y_xw": outcome_model, - "model_t_xw": propensity_model, - }, + outcome_model_name="model_y_xw", + propensity_model_name="model_t_xw", search_space={ "projection": tune.choice([0, 1]), "opt_reweighted": tune.choice([0, 1]), @@ -549,9 +611,7 @@ def _configs(self) -> Dict[str, EstimatorConfig]: }, ), "iv.econml.iv.dr.LinearIntentToTreatDRIV": EstimatorConfig( - init_params={ - "model_y_xw": outcome_model, - }, + outcome_model_name="model_y_xw", search_space={ "cov_clip": tune.quniform(0.08, 0.2, 0.01), "opt_reweighted": tune.choice([0, 1]), diff --git a/causaltune/utils.py b/causaltune/utils.py index df21f89b..6af2097b 100644 --- a/causaltune/utils.py +++ b/causaltune/utils.py @@ -8,6 +8,7 @@ def clean_config(params: dict): + # TODO: move this to formal constraints in tune? if "subforest_size" in params and "n_estimators" in params: params["n_estimators"] = params["subforest_size"] * math.ceil( params["n_estimators"] / params["subforest_size"] diff --git a/setup.py b/setup.py index 08299e37..0ad67cc4 100644 --- a/setup.py +++ b/setup.py @@ -21,7 +21,7 @@ install_requires=[ "dowhy==0.9.1", "econml==0.14.1", - "FLAML==1.0.14", + "FLAML==2.2.0", "xgboost==1.7.6", "numpy==1.23.5", "pandas", diff --git a/tests/causaltune/test_custom_outcome_model.py b/tests/causaltune/test_custom_outcome_model.py index 138d8d83..898832e9 100644 --- a/tests/causaltune/test_custom_outcome_model.py +++ b/tests/causaltune/test_custom_outcome_model.py @@ -5,7 +5,7 @@ from causaltune import CausalTune from causaltune.datasets import linear_multi_dataset, generate_synthetic_data -from causaltune.params import SimpleParamService +from causaltune.search.params import SimpleParamService warnings.filterwarnings("ignore") # suppress sklearn deprecation warnings for now.. @@ -20,8 +20,6 @@ def test_custom_outcome_model(self): data.preprocess_dataset() cfg = SimpleParamService( - propensity_model=None, - outcome_model=None, n_jobs=-1, include_experimental=False, multivalue=False, @@ -56,8 +54,6 @@ def test_custom_outcome_model(self): def test_custom_outcome_model_multivalue(self): data = linear_multi_dataset(10000) cfg = SimpleParamService( - propensity_model=None, - outcome_model=None, n_jobs=-1, include_experimental=False, multivalue=True, diff --git a/tests/causaltune/test_endtoend.py b/tests/causaltune/test_endtoend.py index 741cc035..ccf2ebec 100644 --- a/tests/causaltune/test_endtoend.py +++ b/tests/causaltune/test_endtoend.py @@ -1,10 +1,9 @@ -import pytest import warnings from causaltune import CausalTune from causaltune.datasets import generate_non_random_dataset, linear_multi_dataset -from causaltune.params import SimpleParamService +from causaltune.search.params import SimpleParamService warnings.filterwarnings("ignore") # suppress sklearn deprecation warnings for now.. @@ -42,8 +41,6 @@ def test_endtoend_cate(self): data.preprocess_dataset() cfg = SimpleParamService( - propensity_model=None, - outcome_model=None, n_jobs=-1, include_experimental=False, multivalue=False, @@ -78,8 +75,6 @@ def test_endtoend_cate(self): def test_endtoend_multivalue(self): data = linear_multi_dataset(5000) cfg = SimpleParamService( - propensity_model=None, - outcome_model=None, n_jobs=-1, include_experimental=False, multivalue=True, @@ -101,5 +96,6 @@ def test_endtoend_multivalue(self): if __name__ == "__main__": - pytest.main([__file__]) + TestEndToEnd().test_endtoend_cate() + # pytest.main([__file__]) # TestEndToEnd().test_endtoend_iv() diff --git a/tests/causaltune/test_endtoend_automl_propensity.py b/tests/causaltune/test_endtoend_automl_propensity.py index f5b2f893..c6a6d4f1 100644 --- a/tests/causaltune/test_endtoend_automl_propensity.py +++ b/tests/causaltune/test_endtoend_automl_propensity.py @@ -3,7 +3,7 @@ from causaltune import CausalTune from causaltune.datasets import linear_multi_dataset -from causaltune.params import SimpleParamService +from causaltune.search.params import SimpleParamService warnings.filterwarnings("ignore") # suppress sklearn deprecation warnings for now.. @@ -53,8 +53,6 @@ def test_endtoend_multivalue_propensity(self): data = linear_multi_dataset(10000) cfg = SimpleParamService( - propensity_model=None, - outcome_model=None, n_jobs=-1, include_experimental=False, multivalue=True, diff --git a/tests/causaltune/test_endtoend_flat_search.py b/tests/causaltune/test_endtoend_flat_search.py new file mode 100644 index 00000000..7a20436d --- /dev/null +++ b/tests/causaltune/test_endtoend_flat_search.py @@ -0,0 +1,102 @@ +import warnings + + +from causaltune import CausalTune +from causaltune.datasets import generate_non_random_dataset, linear_multi_dataset +from causaltune.search.params import SimpleParamService + +warnings.filterwarnings("ignore") # suppress sklearn deprecation warnings for now.. + + +class TestEndToEnd(object): + """tests causaltune model end-to-end + 1/ import causaltune object + 2/ preprocess data + 3/ init causaltune object + 4/ run causaltune on data + """ + + def test_imports(self): + """tests if CausalTune can be imported""" + + from causaltune import CausalTune # noqa F401 + + def test_data_preprocessing(self): + """tests data preprocessing routines""" + data = generate_non_random_dataset() # noqa F484 + + def test_init_causaltune(self): + """tests if causaltune object can be instantiated without errors""" + + from causaltune import CausalTune # noqa F401 + + ct = CausalTune(time_budget=0) # noqa F484 + + def test_endtoend_cate(self): + """tests if CATE model can be instantiated and fit to data""" + + from causaltune.shap import shap_values # noqa F401 + + data = generate_non_random_dataset() + data.preprocess_dataset() + + cfg = SimpleParamService( + n_jobs=-1, + include_experimental=False, + multivalue=False, + ) + estimator_list = cfg.estimator_names_from_patterns("backdoor", "all", 1) + # outcome = targets[0] + ct = CausalTune( + components_time_budget=10, + num_samples=len(estimator_list) * 4, + estimator_list=estimator_list, # "all", # + outcome_model="auto", + use_ray=False, + verbose=3, + components_verbose=2, + resources_per_trial={"cpu": 0.5}, + ) + + ct.fit(data) + # ct.fit(data, resume=True) + ct.effect(data.data) + ct.score_dataset(data.data, "test") + + # now let's test Shapley ct calculation + for est_name, scores in ct.scores.items(): + # Dummy model doesn't support Shapley values + # Orthoforest shapley calc is VERY slow + if "Dummy" not in est_name and "Ortho" not in est_name: + print("Calculating Shapley values for", est_name) + shap_values(scores["estimator"], data.data[:10]) + + print(f"Best estimator: {ct.best_estimator}") + + def test_endtoend_multivalue(self): + data = linear_multi_dataset(5000) + cfg = SimpleParamService( + n_jobs=-1, + include_experimental=False, + multivalue=True, + ) + estimator_list = cfg.estimator_names_from_patterns( + "backdoor", "all", data_rows=len(data) + ) + + ct = CausalTune( + estimator_list="all", + num_samples=len(estimator_list), + components_time_budget=10, + ) + ct.fit(data) + # ct.fit(data, resume=True) + + # TODO add an effect() call and an effect_tt call + print("yay!") + + +if __name__ == "__main__": + TestEndToEnd().test_endtoend_cate() + # pytest.main([__file__]) + # TestEndToEnd().test_endtoend_iv() diff --git a/tests/causaltune/test_endtoend_inference.py b/tests/causaltune/test_endtoend_inference.py index 154fe748..659fee7f 100644 --- a/tests/causaltune/test_endtoend_inference.py +++ b/tests/causaltune/test_endtoend_inference.py @@ -5,7 +5,7 @@ from causaltune import CausalTune from causaltune.datasets import linear_multi_dataset -from causaltune.params import SimpleParamService +from causaltune.search.params import SimpleParamService warnings.filterwarnings("ignore") # suppress sklearn deprecation warnings for now.. @@ -21,8 +21,6 @@ def test_endtoend_inference_nobootstrap(self): data.preprocess_dataset() cfg = SimpleParamService( - propensity_model=None, - outcome_model=None, n_jobs=-1, include_experimental=False, multivalue=False, @@ -71,8 +69,6 @@ def test_endtoend_inference_bootstrap(self): def test_endtoend_multivalue_nobootstrap(self): data = linear_multi_dataset(1000) cfg = SimpleParamService( - propensity_model=None, - outcome_model=None, n_jobs=-1, include_experimental=False, multivalue=True, @@ -111,6 +107,7 @@ def test_endtoend_multivalue_bootstrap(self): estimator_list=[e], use_ray=False, verbose=3, + outcome_model="auto", components_verbose=2, resources_per_trial={"cpu": 0.5}, ) @@ -124,4 +121,4 @@ def test_endtoend_multivalue_bootstrap(self): if __name__ == "__main__": pytest.main([__file__]) - # TestEndToEnd().test_endtoend_iv() + # TestEndToEndInference().test_endtoend_multivalue_bootstrap() diff --git a/tests/causaltune/test_estimator_list.py b/tests/causaltune/test_estimator_list.py index 8db6b30a..2b28c726 100644 --- a/tests/causaltune/test_estimator_list.py +++ b/tests/causaltune/test_estimator_list.py @@ -2,7 +2,7 @@ import pandas as pd from causaltune import CausalTune -from causaltune.params import SimpleParamService +from causaltune.search.params import SimpleParamService class TestEstimatorListGenerator: @@ -10,9 +10,7 @@ class TestEstimatorListGenerator: def test_auto_list(self): """tests if "auto" setting yields all available estimators""" - cfg = SimpleParamService( - propensity_model=None, outcome_model=None, multivalue=False - ) + cfg = SimpleParamService(multivalue=False) auto_estimators_iv = cfg.estimator_names_from_patterns("iv", "auto") auto_estimators_backdoor = cfg.estimator_names_from_patterns("backdoor", "auto") # verify that returned estimator list includes all available estimators @@ -40,9 +38,7 @@ def test_auto_list(self): def test_substring_group(self): """tests if substring match to group of estimators works""" - cfg = SimpleParamService( - propensity_model=None, outcome_model=None, multivalue=False - ) + cfg = SimpleParamService(multivalue=False) estimator_list = cfg.estimator_names_from_patterns("backdoor", ["dml"]) available_estimators = [e for e in cfg._configs().keys() if "dml" in e] @@ -57,9 +53,7 @@ def test_substring_group(self): def test_substring_single(self): """tests if substring match to single estimators works""" - cfg = SimpleParamService( - propensity_model=None, outcome_model=None, multivalue=False - ) + cfg = SimpleParamService(multivalue=False) estimator_list = cfg.estimator_names_from_patterns( "backdoor", ["DomainAdaptationLearner"] ) @@ -69,9 +63,7 @@ def test_substring_single(self): def test_checkduplicates(self): """tests if duplicates are removed""" - cfg = SimpleParamService( - propensity_model=None, outcome_model=None, multivalue=False - ) + cfg = SimpleParamService(multivalue=False) estimator_list = cfg.estimator_names_from_patterns( "backdoor", [ @@ -87,9 +79,7 @@ def test_invalid_choice(self): # this should raise a ValueError # unavailable estimators: - cfg = SimpleParamService( - propensity_model=None, outcome_model=None, multivalue=False - ) + cfg = SimpleParamService(multivalue=False) with pytest.raises(ValueError): cfg.estimator_names_from_patterns( diff --git a/tests/causaltune/test_sklearn_propensity_model.py b/tests/causaltune/test_sklearn_propensity_model.py index 7a1ce204..94ef157e 100644 --- a/tests/causaltune/test_sklearn_propensity_model.py +++ b/tests/causaltune/test_sklearn_propensity_model.py @@ -9,7 +9,7 @@ generate_synth_data_with_categories, linear_multi_dataset, ) -from causaltune.params import SimpleParamService +from causaltune.search.params import SimpleParamService warnings.filterwarnings("ignore") # suppress sklearn deprecation warnings for now.. @@ -24,8 +24,6 @@ def test_sklearn_propensity_model(self): data.preprocess_dataset() cfg = SimpleParamService( - propensity_model=None, - outcome_model=None, n_jobs=-1, include_experimental=False, multivalue=False, @@ -60,8 +58,6 @@ def test_sklearn_propensity_model(self): def test_sklearn_propensity_model_multivalue(self): data = linear_multi_dataset(5000) cfg = SimpleParamService( - propensity_model=None, - outcome_model=None, n_jobs=-1, include_experimental=False, multivalue=True,