From be83130087a2e5b85a073d0d891f8a06584b6030 Mon Sep 17 00:00:00 2001
From: EgorKraevTransferwise
 <62890791+EgorKraevTransferwise@users.noreply.github.com>
Date: Wed, 15 Jun 2022 08:39:08 +0100
Subject: [PATCH 1/2] Add FLAML to setup.cfg and bump dowhy version

---
 setup.cfg | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/setup.cfg b/setup.cfg
index dc2ffa146..5be3c4798 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -39,8 +39,9 @@ install_requires =
     statsmodels >= 0.10
     pandas
     shap >= 0.38.1, < 0.41.0
-    dowhy < 0.8
+    dowhy <= 0.8
     lightgbm
+    FLAML >= 1.0.1
 test_suite = econml.tests
 tests_require =
     pytest

From 5f401104dbc880ffc2d95b01833e3847db797098 Mon Sep 17 00:00:00 2001
From: EgorKraevTransferwise
 <62890791+EgorKraevTransferwise@users.noreply.github.com>
Date: Wed, 15 Jun 2022 08:43:23 +0100
Subject: [PATCH 2/2] Add search space functions to ForestDRLearner

---
 econml/dr/_drlearner.py | 50 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/econml/dr/_drlearner.py b/econml/dr/_drlearner.py
index 1cc88eda8..cddd9e937 100644
--- a/econml/dr/_drlearner.py
+++ b/econml/dr/_drlearner.py
@@ -1433,6 +1433,56 @@ def _gen_model_final(self):
 
     def _gen_ortho_learner_model_final(self):
         return _ModelFinal(self._gen_model_final(), self._gen_featurizer(), False)
+    
+    @classmethod
+    def search_space(cls) -> dict:
+        """
+        Specify the hyperparameter search space for FLAML
+
+        Returns
+        -------
+        A FLAML-compatible search space for __init__ args we want to search over
+        """
+        return {
+            "min_propensity": tune.loguniform(1e-6, 1e-1),
+            # "mc_iters": tune.randint(0, 10), # is this worth searching over?
+            "n_estimators": tune.randint(2, 200),
+            # "max_depth": is tune.choice([None, tune.randint(2, 1000)]) the right syntax?
+            "min_samples_split": tune.randint(2, 20),
+            "min_samples_leaf": tune.randint(1, 25),
+            "min_weight_fraction_leaf": tune.uniform(0, 0.5),
+            "max_features": tune.choice(["auto", "sqrt", "log2", None]),
+            "min_impurity_decrease": tune.uniform(0, 10),
+            "max_samples": tune.uniform(0, 0.5),
+            "min_balancedness_tol": tune.uniform(0, 0.5),
+            "honest": tune.choice([0, 1]),
+            "subforest_size": tune.randint(2, 10),
+        }
+
+    @classmethod
+    def search_start(cls) -> dict:
+        """
+        Specify the default start point for hyperparameter search
+
+        Returns
+        -------
+        A dict of hyperparameter values consistent with the output from search_space()
+        """
+
+        return {
+            "min_propensity": 1e-6,
+            "n_estimators": 100,  # default is 1000 but that seems excessive
+            "min_samples_split": 5,
+            "min_samples_leaf": 5,
+            "min_weight_fraction_leaf": 0.0,
+            "max_features": "auto",
+            "min_impurity_decrease": 0.0,
+            "max_samples": 0.45,
+            "min_balancedness_tol": 0.45,
+            "honest": True,
+            "subforest_size": 4,
+        }
+
 
     def fit(self, Y, T, *, X=None, W=None, sample_weight=None, groups=None,
             cache_values=False, inference='auto'):