Repository: fullflu/pydtr (language: python)

Issue:
add examples
# WHY

- There is no example in this project.
- It would also be helpful to show examples using pipeline, which enables us to deal with categorical independent variables.

# TODO

- [x] add tests using pipeline
- [x] fix a table in README

add notebooks
- [ ] continuous covariate and binary action
- [ ] continuous covariate and categorical action (using pipeline)


Discussion / hints:

Fix patch:
diff --git a/Makefile b/Makefile
index 84143f5..1935119 100644
--- a/Makefile
+++ b/Makefile
@@ -1,6 +1,7 @@
 dev:
 	pip install --upgrade pip
 	pip install -e ./.
+	pip install category_encoders
 	pip install pytest
 	pip install coverage
 	pip install twine
diff --git a/README.md b/README.md
index 76d80b4..c88edd0 100644
--- a/README.md
+++ b/README.md
@@ -13,15 +13,17 @@ A DTR is a paradigm that attempts to select optimal treatments adaptively for in
 
 Pydtr enables you to implement DTR methods easily by using sklearn-based interfaces.
 
-|                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 pipeline)  |
-|  IqLearnReg <br> (with statsmodels)  |  :white_check_mark:   |   :white_check_mark:  |   :white_check_mark:       |
-| GEstimation | WIP | | WIP | WIP |
+|                Method                 |  Single binary treatment   |  Multiple treatments  |    Multinomial treatment   |  Continuous treatment  |  Modeling flexibility  |  Interpretability  |
+| ---- | ---- | ---- | ---- | ---- | ---- | ---- | 
+|  IqLearnReg <br> (with sklearn)      |  :white_check_mark:   |   :white_check_mark:  |   :white_check_mark: <br>(with pipeline)  |    |   :white_check_mark: <br>(with arbitrary regression models)  |       |
+|  IqLearnReg <br> (with statsmodels)  |  :white_check_mark:   |   :white_check_mark:  |   :white_check_mark:       |    | limited to OLS   |    :white_check_mark: <br>(with confidence intervals)  |
+| GEstimation | WIP | | WIP | WIP | WIP | WIP |
 
 `IqLearnReg` means a regression method of iterative q-learning.
 
-When a treatment variable is multinomial and you use a sklearn model as a regression function, you need to encode the treatment variable by an ordinal encoding and use a tree-based sklearn model.
+When there are categorical independent variables and you use a sklearn model as a regression function, you need to encode the categorical variables before using the model.
+
+We recommend to encode categorical variables by `category_encoders` and combine the encoders with the sklearn model by `sklearn.pipeline`.
 
 G-estimation, a famous method of DTR, is now unavailable.
 
@@ -136,6 +138,8 @@ opt_action_stage_2 = dtr_model.predict(df, 1)
 opt_action_all_stages = dtr_model.predict_all_stages(df)
 ```
 
+Please see [examples](https://github.com/fullflu/pydtr/blob/master/examples/) to get more information.
+
 ## Authors
 
 - [@fullflu](https://github.com/fullflu) 
@@ -162,6 +166,8 @@ If all checkes have passed in pull-requests, I will merge and release them.
 ├── MANIFEST.IN
 ├── Makefile
 ├── README.md
+├── examples
+│   ├── ...several notebooks...
 ├── setup.cfg
 ├── setup.py
 ├── src
diff --git a/examples/.gitkeep b/examples/.gitkeep
new file mode 100644
index 0000000..e69de29
diff --git a/setup.py b/setup.py
index 122661d..b8d7993 100644
--- a/setup.py
+++ b/setup.py
@@ -58,7 +58,7 @@ def run(self):
         'Programming Language :: Python :: 3.6',
     ],
     setup_requires=["pytest-runner"],
-    tests_require=["pytest", "pytest-cov", "coverage"],
+    tests_require=["pytest", "pytest-cov", "coverage", "category_encoders"],
     cmdclass={
         'verify': VerifyVersionCommand,
     }


Test patch:
diff --git a/tests/test_iqlearn_sklearn_predict.py b/tests/test_iqlearn_sklearn_predict.py
index e8b822c..0a4f4de 100644
--- a/tests/test_iqlearn_sklearn_predict.py
+++ b/tests/test_iqlearn_sklearn_predict.py
@@ -1,8 +1,12 @@
 import pandas as pd
 import numpy as np
+import category_encoders as ce
 from sklearn.base import BaseEstimator, RegressorMixin
 from sklearn.ensemble import RandomForestRegressor
 from pandas._testing import assert_frame_equal
+from sklearn.linear_model import Ridge
+from sklearn.preprocessing import StandardScaler
+from sklearn.pipeline import Pipeline
 
 from pydtr.iqlearn.regression import IqLearnReg
 
@@ -173,7 +177,7 @@ def test_iqlearn_rf():
     assert len(dtr_model.models[0]) == 2
 
 
-def test_iqlearn_regwrapper_rf_multiple_actions():
+def test_iqlearn_rf_multiple_actions():
     # setup params
     n = 10
     thres = int(n / 2)
@@ -232,7 +236,7 @@ def test_iqlearn_regwrapper_rf_multiple_actions():
     assert len(dtr_model.models[0]) == 2
 
 
-def test_iqlearn_regwrapper_rf_ordinalencoder():
+def test_iqlearn_rf_ordinalencoder():
     # setup params
     n = 30
     thres = int(n / 2)
@@ -289,3 +293,66 @@ def test_iqlearn_regwrapper_rf_ordinalencoder():
     )
     dtr_model.fit(df)
     assert len(dtr_model.models[0]) == 2
+
+
+def test_iqlearn_pipeline_category_encoder():
+    # setup params
+    n = 30
+    thres = int(n / 2)
+    # statsmodels
+    model1 = Pipeline(
+        [
+            ("ce0", ce.OneHotEncoder(cols=["A1"])),
+            ("scale", StandardScaler()),
+            ("model", Ridge())
+        ]
+    )
+    model2 = Pipeline(
+        [
+            ("ce0", ce.OneHotEncoder(cols=["A1", "A2"])),
+            ("scale", StandardScaler()),
+            ("model", Ridge())
+        ]
+    )
+    # sample dataframe
+    df = pd.DataFrame()
+    df["L1"] = np.arange(n)
+    df["A1"] = ["A", "B", "C"] * int(n / 3)
+    df["A2"] = ["A"] * int(n / 3) + ["C"] * int(n / 3) + ["D"] * int(n / 3)
+    df["Y1"] = np.zeros(n)
+    df["Y2"] = np.zeros(n)
+    # set model info
+    model_info = [
+        {
+            "model": model1,
+            "action_dict": {"A1": ["A", "B", "C"]},
+            "feature": ["L1", "A1"],
+            "outcome": "Y1"
+        },
+        {
+            "model": model2,
+            "action_dict": {"A2": ["A", "C", "D"]},
+            "feature": ["L1", "A1", "Y1", "A2"],
+            "outcome": "Y2"
+        }
+    ]
+    # fit model
+    dtr_model = IqLearnReg(
+        n_stages=2,
+        model_info=model_info
+    )
+    dtr_model.fit(df)
+    # predict optimal atcions
+    action_1 = dtr_model.predict(df, 0)
+    action_2 = dtr_model.predict(df, 1)
+    action_all = dtr_model.predict_all_stages(df)
+    # stage 1 test
+    assert action_1.shape[0] == df.shape[0]
+    # stage 2 test
+    assert action_2.shape[0] == df.shape[0]
+    # all stage test
+    assert action_all.shape[0] == action_1.shape[0] * 2
+    a1 = action_all.query("stage == 0")[["A1", "val"]].reset_index(drop=True)
+    a2 = action_all.query("stage == 1")[["A2", "val"]].reset_index(drop=True)
+    assert_frame_equal(action_1, a1)
+    assert_frame_equal(action_2, a2)