Repository: fullflu/pydtr

Issue:
fix get_params
# WHY

There is a bug in the `get_params` method (`self.models` cannot be sliced) :
```
for t, m in enumerate(self.models[:-1]): <-
```
# TODO

- [ ] fix the function
- [ ] add tests of the function

Discussion:

Patch:
diff --git a/README.md b/README.md
index 0210681..76d80b4 100644
--- a/README.md
+++ b/README.md
@@ -15,7 +15,7 @@ Pydtr enables you to implement DTR methods easily by using sklearn-based interfa
 
 |                Method                 |  Single binary treatment   |  Multiple treatments  |    Multinomial treatment   |  Continuous Treatment  |
 | ---- | ---- | ---- | ---- | ---- |
-|  IqLearnReg <br> (with sklearn)      |  :white_check_mark:   |   :white_check_mark:  |   :white_check_mark: <br>(with ordinal encoded treatment)  |
+|  IqLearnReg <br> (with sklearn)      |  :white_check_mark:   |   :white_check_mark:  |   :white_check_mark: <br>(with pipeline)  |
 |  IqLearnReg <br> (with statsmodels)  |  :white_check_mark:   |   :white_check_mark:  |   :white_check_mark:       |
 | GEstimation | WIP | | WIP | WIP |
 
@@ -25,21 +25,37 @@ When a treatment variable is multinomial and you use a sklearn model as a regres
 
 G-estimation, a famous method of DTR, is now unavailable.
 
-## Requirement
+## Requirements
 
-- python (>= 3.6)
-- pip
+- python>=3.6
+- pandas>=1.1.2
+- scikit-learn>=0.23.2
+- numpy>=1.19.2
+- statsmodels>=0.12.0
 
-## Install
+## Installation
 
-`pip install pydtr`
+### From pypi
+
+```
+pip install pydtr
+```
+
+### From source
+
+```
+git clone https://github.com/fullflu/pydtr.git
+cd pydtr
+python setup.py install
+```
 
 ## Usage
 
 ### Iterative Q Learning (IqLearnReg)
 
 You need to import libraries and prepare data.
-```
+
+```python
 # import
 import numpy as np
 import pandas as pd
@@ -59,7 +75,8 @@ df["Y2"] = np.zeros(n)
 ```
 
 You can use sklearn-based models.
-```
+
+```python
 # set model info
 model_info = [
     {
@@ -89,7 +106,8 @@ opt_action_all_stages = dtr_model.predict_all_stages(df)
 ```
 
 You can also use statsmodels-based models.
-```
+
+```python
 # set model info
 model_info = [
     {
@@ -160,4 +178,4 @@ If all checkes have passed in pull-requests, I will merge and release them.
 
 ## References
 
-- Chakraborty, Bibhas. *Statistical methods for dynamic treatment regimes.* Springer, 2013.
+- Chakraborty, B, Moodie, EE. *Statistical Methods for Dynamic Treatment Regimes.* Springer, New York, 2013.
diff --git a/src/pydtr/iqlearn/base.py b/src/pydtr/iqlearn/base.py
index 52ca6e7..d08901d 100644
--- a/src/pydtr/iqlearn/base.py
+++ b/src/pydtr/iqlearn/base.py
@@ -6,7 +6,6 @@ from abc import ABCMeta, abstractmethod
 import pandas as pd
 import numpy as np
 from sklearn.utils import resample
-from sklearn.utils.estimator_checks import check_estimator
 
 
 class IqLearnBase(object):
@@ -103,10 +102,6 @@ class IqLearnBase(object):
             size_bs = df.shape[0]
         return resample(df, n_samples=size_bs)
 
-    @staticmethod
-    def _check_model_type(model):
-        assert type(model) == str or check_estimator(model)
-
     def fit(self, df: pd.DataFrame):
         """
         Fit dtr models
@@ -136,16 +131,15 @@ class IqLearnBase(object):
         # fit models using bootstrap
         for i in range(self.n_bs):
             df_i = self._sample_bs(df)
-            print("{}th bootstrap".format(i))
             for t in reversed(range(self.n_stages)):
                 # extract feature and outcome
-                X = df[self.model_info[t]["feature"]]
-                y = df[self.model_info[t]["outcome"]]
+                X = df_i[self.model_info[t]["feature"]]
+                y = df_i[self.model_info[t]["outcome"]]
                 if t == self.n_stages - 1:
                     p_outcome = y.values
                 else:
-                    X2 = df[self.model_info[t + 1]["feature"]]
-                    y2 = df[self.model_info[t + 1]["outcome"]]
+                    X2 = df_i[self.model_info[t + 1]["feature"]]
+                    y2 = df_i[self.model_info[t + 1]["outcome"]]
                     p_outcome = self._get_p_outcome(self.model_all[t + 1], X2, y2, t)
                 # fit model of stage t
                 self._fit_model(X, p_outcome, t, i)
@@ -172,9 +166,9 @@ class IqLearnBase(object):
     def get_params(self) -> pd.DataFrame:
         # get estimated parameters
         params = pd.DataFrame()
-        for t, m in enumerate(self.models[:-1]):
+        for t in reversed(range(self.n_stages)):
             if type(self.model_info[t]["model"]) == str:
-                tmp_df = pd.melt(pd.DataFrame([i.params for i in m]))
+                tmp_df = pd.melt(pd.DataFrame([i.params for i in self.models[t]]))
                 tmp_df["stage"] = t
                 params = pd.concat([params, tmp_df])
         return params
diff --git a/src/pydtr/version.py b/src/pydtr/version.py
index f102a9c..3b93d0b 100644
--- a/src/pydtr/version.py
+++ b/src/pydtr/version.py
@@ -1,1 +1,1 @@
-__version__ = "0.0.1"
+__version__ = "0.0.2"


Test patch:
diff --git a/tests/test_iqlearn_sklearn_predict.py b/tests/test_iqlearn_sklearn_predict.py
index faec86e..e8b822c 100644
--- a/tests/test_iqlearn_sklearn_predict.py
+++ b/tests/test_iqlearn_sklearn_predict.py
@@ -114,11 +114,11 @@ def test_iqlearn_regwrapper_rule():
     assert len(dtr_model.models[0]) == 2
 
 
-def test_iqlearn_regwrapper_rf():
+def test_iqlearn_rf():
     # setup params
     n = 10
     thres = int(n / 2)
-    # sample rule base models
+    # rf models
     model1 = RandomForestRegressor()
     model2 = RandomForestRegressor()
     # sample dataframe
@@ -143,7 +143,7 @@ def test_iqlearn_regwrapper_rf():
             "outcome": "Y2"
         }
     ]
-    # fit model (dummy)
+    # fit model
     dtr_model = IqLearnReg(
         n_stages=2,
         model_info=model_info
@@ -163,7 +163,7 @@ def test_iqlearn_regwrapper_rf():
     a2 = action_all.query("stage == 1")[["A2", "val"]].reset_index(drop=True)
     assert_frame_equal(action_1, a1)
     assert_frame_equal(action_2, a2)
-    # fit bootstrap model (dummy)
+    # fit bootstrap model
     dtr_model = IqLearnReg(
         n_stages=2,
         model_info=model_info,
@@ -177,7 +177,7 @@ def test_iqlearn_regwrapper_rf_multiple_actions():
     # setup params
     n = 10
     thres = int(n / 2)
-    # sample rule base models
+    # rf models
     model1 = RandomForestRegressor()
     model2 = RandomForestRegressor()
     # sample dataframe
@@ -202,7 +202,7 @@ def test_iqlearn_regwrapper_rf_multiple_actions():
             "outcome": "Y2"
         }
     ]
-    # fit model (dummy)
+    # fit model
     dtr_model = IqLearnReg(
         n_stages=2,
         model_info=model_info
@@ -222,7 +222,7 @@ def test_iqlearn_regwrapper_rf_multiple_actions():
     a2 = action_all.query("stage == 1")[["A1", "A2", "val"]].reset_index(drop=True)
     assert_frame_equal(action_1, a1)
     assert_frame_equal(action_2, a2)
-    # fit bootstrap model (dummy)
+    # fit bootstrap model
     dtr_model = IqLearnReg(
         n_stages=2,
         model_info=model_info,
@@ -236,7 +236,7 @@ def test_iqlearn_regwrapper_rf_ordinalencoder():
     # setup params
     n = 30
     thres = int(n / 2)
-    # sample rule base models
+    # rf models
     model1 = RandomForestRegressor()
     model2 = RandomForestRegressor()
     # sample dataframe
@@ -261,7 +261,7 @@ def test_iqlearn_regwrapper_rf_ordinalencoder():
             "outcome": "Y2"
         }
     ]
-    # fit model (dummy)
+    # fit model
     dtr_model = IqLearnReg(
         n_stages=2,
         model_info=model_info
@@ -281,7 +281,7 @@ def test_iqlearn_regwrapper_rf_ordinalencoder():
     a2 = action_all.query("stage == 1")[["A2", "val"]].reset_index(drop=True)
     assert_frame_equal(action_1, a1)
     assert_frame_equal(action_2, a2)
-    # fit bootstrap model (dummy)
+    # fit bootstrap model
     dtr_model = IqLearnReg(
         n_stages=2,
         model_info=model_info,
diff --git a/tests/test_iqlearn_sm_predict.py b/tests/test_iqlearn_sm_predict.py
index 448ab14..64dfc73 100644
--- a/tests/test_iqlearn_sm_predict.py
+++ b/tests/test_iqlearn_sm_predict.py
@@ -10,7 +10,7 @@ def test_iqlearn_sm():
     # setup params
     n = 10
     thres = int(n / 2)
-    # sample rule base models
+    # statsmodels
     model1 = "p_outcome ~ L1 * A1"
     model2 = "p_outcome ~ L1 + A1 + Y1 * A2"
     # sample dataframe
@@ -35,7 +35,7 @@ def test_iqlearn_sm():
             "outcome": "Y2"
         }
     ]
-    # fit model (dummy)
+    # fit model
     dtr_model = IqLearnReg(
         n_stages=2,
         model_info=model_info
@@ -55,7 +55,7 @@ def test_iqlearn_sm():
     a2 = action_all.query("stage == 1")[["A2", "val"]].reset_index(drop=True)
     assert_frame_equal(action_1, a1)
     assert_frame_equal(action_2, a2)
-    # fit bootstrap model (dummy)
+    # fit bootstrap model
     dtr_model = IqLearnReg(
         n_stages=2,
         model_info=model_info,
@@ -69,7 +69,7 @@ def test_iqlearn_sm_multiple_actions():
     # setup params
     n = 10
     thres = int(n / 2)
-    # sample rule base models
+    # statsmodels
     model1 = "p_outcome ~ L1 * A1"
     model2 = "p_outcome ~ L1 + A1 + Y1 * A2"
     # sample dataframe
@@ -94,7 +94,7 @@ def test_iqlearn_sm_multiple_actions():
             "outcome": "Y2"
         }
     ]
-    # fit model (dummy)
+    # fit model
     dtr_model = IqLearnReg(
         n_stages=2,
         model_info=model_info
@@ -114,7 +114,7 @@ def test_iqlearn_sm_multiple_actions():
     a2 = action_all.query("stage == 1")[["A1", "A2", "val"]].reset_index(drop=True)
     assert_frame_equal(action_1, a1)
     assert_frame_equal(action_2, a2)
-    # fit bootstrap model (dummy)
+    # fit bootstrap model
     dtr_model = IqLearnReg(
         n_stages=2,
         model_info=model_info,
@@ -128,7 +128,7 @@ def test_iqlearn_sm_multinomial_action():
     # setup params
     n = 30
     thres = int(n / 2)
-    # sample rule base models
+    # statsmodels
     model1 = "p_outcome ~ L1 * C(A1)"
     model2 = "p_outcome ~ L1 + A1 + Y1 * C(A2)"
     # sample dataframe
@@ -153,7 +153,7 @@ def test_iqlearn_sm_multinomial_action():
             "outcome": "Y2"
         }
     ]
-    # fit model (dummy)
+    # fit model
     dtr_model = IqLearnReg(
         n_stages=2,
         model_info=model_info
@@ -173,7 +173,7 @@ def test_iqlearn_sm_multinomial_action():
     a2 = action_all.query("stage == 1")[["A2", "val"]].reset_index(drop=True)
     assert_frame_equal(action_1, a1)
     assert_frame_equal(action_2, a2)
-    # fit bootstrap model (dummy)
+    # fit bootstrap model
     dtr_model = IqLearnReg(
         n_stages=2,
         model_info=model_info,
@@ -181,3 +181,48 @@ def test_iqlearn_sm_multinomial_action():
     )
     dtr_model.fit(df)
     assert len(dtr_model.models[0]) == 2
+
+
+def test_iqlearn_sm_get_params():
+    # setup params
+    n = 300
+    thres = int(n / 2)
+    # statsmodels
+    model1 = "p_outcome ~ L1 * A1"
+    model2 = "p_outcome ~ L1 + A1 + Y1 * A2"
+    # sample dataframe
+    df = pd.DataFrame()
+    df["L1"] = np.random.normal(0, size=n)
+    df["A1"] = [0, 1] * int(n / 2)
+    df["A2"] = [0] * int(n / 2) + [1] * int(n / 2)
+    df["Y1"] = df["L1"] * df["A1"] + np.random.normal(0, scale=5, size=n)
+    df["Y2"] = df["A1"] + df["Y1"] * df["A2"] + np.random.normal(0, scale=5, size=n)
+    # set model info
+    model_info = [
+        {
+            "model": model1,
+            "action_dict": {"A1": [0, 1]},
+            "feature": ["L1", "A1"],
+            "outcome": "Y1"
+        },
+        {
+            "model": model2,
+            "action_dict": {"A1": [0, 1], "A2": [0, 1]},
+            "feature": ["L1", "A1", "Y1", "A2"],
+            "outcome": "Y2"
+        }
+    ]
+    # fit bootstrap model
+    dtr_model = IqLearnReg(
+        n_stages=2,
+        model_info=model_info,
+        n_bs=10
+    )
+    dtr_model.fit(df)
+    # get params
+    params = dtr_model.get_params()
+    l1_unique_shape = params.query("stage == 0 & variable == 'L1'")["value"].unique().shape[0]
+    a1_unique_shape = params.query("stage == 0 & variable == 'A1'")["value"].unique().shape[0]
+    a2_unique_shape = params.query("stage == 1 & variable == 'A2'")["value"].unique().shape[0]
+    assert l1_unique_shape != 1 or a1_unique_shape != 1 or a2_unique_shape != 1
+    assert len(dtr_model.models[0]) == 10