CM4/python-dml-401k-IV.ipynb

{"cells":[{"cell_type":"markdown","id":"narrative-sailing","metadata":{"papermill":{"duration":0.040767,"end_time":"2021-03-24T14:13:51.082088","exception":false,"start_time":"2021-03-24T14:13:51.041321","status":"completed"},"tags":[],"id":"narrative-sailing"},"source":["# Inference on Predictive and Causal Effects in High-Dimensional Nonlinear Models"]},{"cell_type":"markdown","id":"ready-appearance","metadata":{"papermill":{"duration":0.036878,"end_time":"2021-03-24T14:13:51.155414","exception":false,"start_time":"2021-03-24T14:13:51.118536","status":"completed"},"tags":[],"id":"ready-appearance"},"source":["## Impact of 401(k) on  Financial Wealth\n","\n","As a practical illustration of the methods developed in this lecture, we consider estimation of the effect of 401(k) eligibility and participation\n","on accumulated assets. 401(k) plans are pension accounts sponsored by employers. The key problem in determining the effect of participation in 401(k) plans on accumulated assets is saver heterogeneity coupled with the fact that the decision to enroll in a 401(k) is non-random. It is generally recognized that some people have a higher preference for saving than others. It also seems likely that those individuals with high unobserved preference for saving would be most likely to choose to participate in tax-advantaged retirement savings plans and would tend to have otherwise high amounts of accumulated assets. The presence of unobserved savings preferences with these properties then implies that conventional estimates that do not account for saver heterogeneity and endogeneity of participation will be biased upward, tending to overstate the savings effects of 401(k) participation.\n","\n","One can argue that eligibility for enrolling in a 401(k) plan in this data can be taken as exogenous after conditioning on a few observables of which the most important for their argument is income. The basic idea is that, at least around the time 401(k)’s initially became available, people were unlikely to be basing their employment decisions on whether an employer offered a 401(k) but would instead focus on income and other aspects of the job."]},{"cell_type":"markdown","id":"divine-phoenix","metadata":{"papermill":{"duration":0.036497,"end_time":"2021-03-24T14:13:51.229370","exception":false,"start_time":"2021-03-24T14:13:51.192873","status":"completed"},"tags":[],"id":"divine-phoenix"},"source":["### Data\n","\n","The data set can be downloaded from the github repo\n"]},{"cell_type":"code","source":["!pip install wget"],"metadata":{"id":"teq-s8pb3a82"},"id":"teq-s8pb3a82","execution_count":null,"outputs":[]},{"cell_type":"code","execution_count":null,"id":"75662e18","metadata":{"id":"75662e18"},"outputs":[],"source":["# Import relevant packages\n","import pandas as pd\n","import numpy as np\n","import matplotlib.pyplot as plt\n","from sklearn.model_selection import train_test_split\n","from sklearn.model_selection import cross_val_score, cross_val_predict, KFold\n","from sklearn.pipeline import make_pipeline\n","from sklearn.linear_model import LassoCV, RidgeCV, ElasticNetCV, LinearRegression, Ridge, Lasso, LogisticRegressionCV\n","from sklearn.preprocessing import StandardScaler\n","from sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier\n","from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier\n","from sklearn.tree import DecisionTreeRegressor, DecisionTreeClassifier\n","import patsy\n","import warnings\n","from sklearn.base import BaseEstimator, clone\n","import statsmodels.api as sm\n","from IPython.display import Markdown\n","import os\n","import wget\n","import seaborn as sns\n","warnings.simplefilter('ignore')\n","np.random.seed(1234)"]},{"cell_type":"code","execution_count":null,"id":"6e66e77f","metadata":{"id":"6e66e77f"},"outputs":[],"source":["file = \"https://raw.githubusercontent.com/CausalAIBook/MetricsMLNotebooks/main/data/401k.csv\"\n","data = pd.read_csv(file)"]},{"cell_type":"code","execution_count":null,"id":"65a2d086","metadata":{"id":"65a2d086"},"outputs":[],"source":["data.describe()"]},{"cell_type":"code","execution_count":null,"id":"00884061","metadata":{"id":"00884061"},"outputs":[],"source":["data.head()"]},{"cell_type":"code","execution_count":null,"id":"1d3dc85c","metadata":{"id":"1d3dc85c"},"outputs":[],"source":["readme = \"https://raw.githubusercontent.com/CausalAIBook/MetricsMLNotebooks/main/data/401k.md\"\n","filename = wget.download(readme)\n","display(Markdown(open(filename, 'r').read()))"]},{"cell_type":"markdown","id":"looking-invention","metadata":{"papermill":{"duration":0.037509,"end_time":"2021-03-24T14:13:52.250583","exception":false,"start_time":"2021-03-24T14:13:52.213074","status":"completed"},"tags":[],"id":"looking-invention"},"source":["The data consist of 9,915 observations at the household level drawn from the 1991 Survey of Income and Program Participation (SIPP).  All the variables are referred to 1990. We use net financial assets (*net\\_tfa*) as the outcome variable, $Y$,  in our analysis. The net financial assets are computed as the sum of IRA balances, 401(k) balances, checking accounts, saving bonds, other interest-earning accounts, other interest-earning assets, stocks, and mutual funds less non mortgage debts."]},{"cell_type":"markdown","id":"received-nutrition","metadata":{"papermill":{"duration":0.038843,"end_time":"2021-03-24T14:13:52.327055","exception":false,"start_time":"2021-03-24T14:13:52.288212","status":"completed"},"tags":[],"id":"received-nutrition"},"source":["Among the $9915$ individuals, $3682$ are eligible to participate in the program. The variable *e401* indicates eligibility and *p401* indicates participation, respectively."]},{"cell_type":"code","execution_count":null,"id":"0d26d552","metadata":{"id":"0d26d552"},"outputs":[],"source":["sns.countplot(data['e401'],)\n","plt.show()"]},{"cell_type":"markdown","id":"material-sending","metadata":{"papermill":{"duration":0.039406,"end_time":"2021-03-24T14:13:53.435306","exception":false,"start_time":"2021-03-24T14:13:53.395900","status":"completed"},"tags":[],"id":"material-sending"},"source":["Eligibility is highly associated with financial wealth:"]},{"cell_type":"code","execution_count":null,"id":"a7741bb6","metadata":{"id":"a7741bb6"},"outputs":[],"source":["sns.displot(data=data, x='net_tfa', kind='kde', col='e401', hue='e401', fill=True)\n","plt.show()"]},{"cell_type":"markdown","id":"awful-antigua","metadata":{"papermill":{"duration":0.040645,"end_time":"2021-03-24T14:13:54.245822","exception":false,"start_time":"2021-03-24T14:13:54.205177","status":"completed"},"tags":[],"id":"awful-antigua"},"source":["The unconditional APE of e401 is about $19559$:"]},{"cell_type":"code","execution_count":null,"id":"ef070d79","metadata":{"id":"ef070d79"},"outputs":[],"source":["e1 = data[data['e401'] == 1]['net_tfa']\n","e0 = data[data['e401'] == 0]['net_tfa']\n","print(f'{np.mean(e1) - np.mean(e0):.0f}')"]},{"cell_type":"markdown","id":"cross-priority","metadata":{"papermill":{"duration":0.041187,"end_time":"2021-03-24T14:13:54.437695","exception":false,"start_time":"2021-03-24T14:13:54.396508","status":"completed"},"tags":[],"id":"cross-priority"},"source":["Among the $3682$ individuals that  are eligible, $2594$ decided to participate in the program. The unconditional APE of p401 is about $27372$:"]},{"cell_type":"code","execution_count":null,"id":"33cd014e","metadata":{"id":"33cd014e"},"outputs":[],"source":["e1 = data[data['p401'] == 1]['net_tfa']\n","e0 = data[data['p401'] == 0]['net_tfa']\n","print(f'{np.mean(e1) - np.mean(e0):.0f}')"]},{"cell_type":"markdown","id":"suitable-vulnerability","metadata":{"papermill":{"duration":0.041678,"end_time":"2021-03-24T14:13:54.640532","exception":false,"start_time":"2021-03-24T14:13:54.598854","status":"completed"},"tags":[],"id":"suitable-vulnerability"},"source":["As discussed, these estimates are biased since they do not account for saver heterogeneity and endogeneity of participation."]},{"cell_type":"code","execution_count":null,"id":"a9354414","metadata":{"id":"a9354414"},"outputs":[],"source":["y = data['net_tfa'].values\n","Z = data['e401'].values\n","D = data['p401'].values\n","X = data.drop(['e401', 'p401', 'a401', 'tw', 'tfa', 'net_tfa', 'tfa_he',\n","               'hval', 'hmort', 'hequity',\n","               'nifa', 'net_nifa', 'net_n401', 'ira',\n","               'dum91', 'icat', 'ecat', 'zhat',\n","               'i1', 'i2', 'i3', 'i4', 'i5', 'i6', 'i7',\n","               'a1', 'a2', 'a3', 'a4', 'a5'], axis=1)\n","X.columns"]},{"cell_type":"markdown","id":"c4c3e489","metadata":{"id":"c4c3e489"},"source":["### We define a transformer that constructs the engineered features for controls"]},{"cell_type":"code","execution_count":null,"id":"a2cfdf4c","metadata":{"id":"a2cfdf4c"},"outputs":[],"source":["!pip install formulaic"]},{"cell_type":"code","execution_count":null,"id":"53b1283d","metadata":{"id":"53b1283d"},"outputs":[],"source":["from sklearn.base import TransformerMixin, BaseEstimator\n","from formulaic import Formula\n","\n","class FormulaTransformer(TransformerMixin, BaseEstimator):\n","\n","    def __init__(self, formula, array=False):\n","        self.formula = formula\n","        self.array = array\n","\n","    def fit(self, X, y=None):\n","        return self\n","\n","    def transform(self, X, y=None):\n","        df = Formula(self.formula).get_model_matrix(X)\n","        if self.array:\n","            return df.values\n","        return df"]},{"cell_type":"code","execution_count":null,"id":"824ee320","metadata":{"id":"824ee320"},"outputs":[],"source":["transformer = FormulaTransformer(\"0 + poly(age, degree=6, raw=True) + poly(inc, degree=8, raw=True) \"\n","                                 \"+ poly(educ, degree=4, raw=True) + poly(fsize, degree=2, raw=True) \"\n","                                 \"+ male + marr + twoearn + db + pira + hown\")"]},{"cell_type":"code","execution_count":null,"id":"b35bbd34","metadata":{"id":"b35bbd34"},"outputs":[],"source":["transformer.fit_transform(X).head()"]},{"cell_type":"code","execution_count":null,"id":"8c5e20e7","metadata":{"id":"8c5e20e7"},"outputs":[],"source":["transformer = FormulaTransformer(\"0 + poly(age, degree=6, raw=True) + poly(inc, degree=8, raw=True) \"\n","                                 \"+ poly(educ, degree=4, raw=True) + poly(fsize, degree=2, raw=True) \"\n","                                 \"+ male + marr + twoearn + db + pira + hown\", array=True)"]},{"cell_type":"markdown","id":"f751665e","metadata":{"id":"f751665e"},"source":["# Effect of Eligibility on Financial Assets"]},{"cell_type":"code","execution_count":null,"id":"051af0ae","metadata":{"id":"051af0ae"},"outputs":[],"source":["modely = make_pipeline(transformer, StandardScaler(), LassoCV())\n","modelz = make_pipeline(transformer, StandardScaler(), LassoCV())"]},{"cell_type":"code","execution_count":null,"id":"53bc1c83","metadata":{"id":"53bc1c83"},"outputs":[],"source":["resy = y - modely.fit(X, y).predict(X)\n","resZ = Z - modelz.fit(X, Z).predict(X)"]},{"cell_type":"code","execution_count":null,"id":"51244ab4","metadata":{"id":"51244ab4"},"outputs":[],"source":["np.mean(resy * resZ) / np.mean(resZ**2)"]},{"cell_type":"markdown","source":["Recall if we want to do inference, we need to either use the theoretically driven penalty paramter for Lasso or perform cross-fitting."],"metadata":{"id":"bWR5j12xZQHJ"},"id":"bWR5j12xZQHJ"},{"cell_type":"markdown","id":"de8e13e7","metadata":{"id":"de8e13e7"},"source":["# Instrumental Variables: Effect of 401k Participation on Financial Assets"]},{"cell_type":"markdown","id":"61e02e77","metadata":{"id":"61e02e77"},"source":["## Double ML IV under Partial Linearity"]},{"cell_type":"markdown","id":"83869c99","metadata":{"id":"83869c99"},"source":["Now, we consider estimation of average treatment effects of participation in 401k, i.e. `p401`, with the binary instrument being eligibility in 401k, i.e. `e401`. As before, $Y$ denotes the outcome `net_tfa`, and $X$ is the vector of covariates. We consider a partially linear structural equation model:\n","\\begin{eqnarray*}\n","Y & := & g_Y(\\epsilon_Y) D + f_Y(A, X, \\epsilon_Y),  \\\\\n","D & := & f_D(Z, X, A, \\epsilon_D), \\\\\n","Z & := & f_Z(X, \\epsilon_Z),\\\\\n","A & : =  & f_A(X, \\epsilon_A), \\\\\n","X & := &  \\epsilon_X,\n","\\end{eqnarray*}\n","where $A$ is a vector of un-observed confounders.\n","\n","Under this structural equation model, the average treatment effect:\n","\\begin{align}\n","\\alpha = E[Y(1) - Y(0)]\n","\\end{align}\n","can be identified by the moment restriction:\n","\\begin{align}\n","E[(\\tilde{Y} - \\alpha \\tilde{D}) \\tilde{Z}] = 0\n","\\end{align}\n","where for any variable $V$, we denote with $\\tilde{V} = V - E[V|X]$."]},{"cell_type":"code","execution_count":null,"id":"0d743ebe","metadata":{"id":"0d743ebe"},"outputs":[],"source":["modely = make_pipeline(transformer, StandardScaler(), LassoCV())\n","modeld = make_pipeline(transformer, StandardScaler(), LassoCV())\n","modelz = make_pipeline(transformer, StandardScaler(), LassoCV())"]},{"cell_type":"code","execution_count":null,"id":"e7bd792d","metadata":{"id":"e7bd792d"},"outputs":[],"source":["resy = y - modely.fit(X, y).predict(X)\n","resZ = Z - modelz.fit(X, Z).predict(X) # instrument is e401k (eligibility)\n","resD = D - modeld.fit(X, D).predict(X) # treatment is p401k (participation)"]},{"cell_type":"code","execution_count":null,"id":"7ebedcb0","metadata":{"id":"7ebedcb0"},"outputs":[],"source":["np.mean(resy * resZ) / np.mean(resD * resZ)"]},{"cell_type":"markdown","source":["Recall if we want to do inference, we need to either use the theoretically driven penalty paramter for Lasso or perform cross-fitting."],"metadata":{"id":"cH3YkvuSZfP0"},"id":"cH3YkvuSZfP0"},{"cell_type":"markdown","id":"b218c25e","metadata":{"id":"b218c25e"},"source":["### DML with Non-Linear ML Models and Cross-fitting"]},{"cell_type":"code","execution_count":null,"id":"333c19ed","metadata":{"id":"333c19ed"},"outputs":[],"source":["def dml(X, Z, D, y, modely, modeld, modelz, *, nfolds, classifier=False):\n","    '''\n","    DML for the Partially Linear Model setting with cross-fitting\n","\n","    Input\n","    -----\n","    X: the controls\n","    Z: the instrument\n","    D: the treatment\n","    y: the outcome\n","    modely: the ML model for predicting the outcome y\n","    modeld: the ML model for predicting the treatment D\n","    modelz: the ML model for predicting the instrument Z\n","    nfolds: the number of folds in cross-fitting\n","    classifier: bool, whether the modeld is a classifier or a regressor\n","\n","    Output\n","    ------\n","    point: the point estimate of the treatment effect of D on y\n","    stderr: the standard error of the treatment effect\n","    yhat: the cross-fitted predictions for the outcome y\n","    Dhat: the cross-fitted predictions for the treatment D\n","    Zhat: the cross-fitted predictions for the instrument Z\n","    resy: the outcome residuals\n","    resD: the treatment residuals\n","    resZ: the instrument residuals\n","    epsilon: the final residual-on-residual OLS regression residual\n","    '''\n","    cv = KFold(n_splits=nfolds, shuffle=True, random_state=123) # shuffled k-folds\n","    yhat = cross_val_predict(modely, X, y, cv=cv, n_jobs=-1) # out-of-fold predictions for y\n","    # out-of-fold predictions for D\n","    # use predict or predict_proba dependent on classifier or regressor for D\n","    if classifier:\n","        Dhat = cross_val_predict(modeld, X, D, cv=cv, method='predict_proba', n_jobs=-1)[:, 1]\n","        Zhat = cross_val_predict(modelz, X, Z, cv=cv, method='predict_proba', n_jobs=-1)[:, 1]\n","    else:\n","        Dhat = cross_val_predict(modeld, X, D, cv=cv, n_jobs=-1)\n","        Zhat = cross_val_predict(modelz, X, Z, cv=cv, n_jobs=-1)\n","    # calculate outcome and treatment residuals\n","    resy = y - yhat\n","    resD = D - Dhat\n","    resZ = Z - Zhat\n","    # final stage ols based point estimate and standard error\n","    point = np.mean(resy * resZ) / np.mean(resD*resZ)\n","    epsilon = resy - point * resD\n","    var = np.mean(epsilon**2 * resZ**2) / np.mean(resD*resZ)**2\n","    stderr = np.sqrt(var / X.shape[0])\n","    return point, stderr, yhat, Dhat, Zhat, resy, resD, resZ, epsilon"]},{"cell_type":"code","execution_count":null,"id":"7d01b969","metadata":{"id":"7d01b969"},"outputs":[],"source":["def summary(point, stderr, yhat, Dhat, Zhat, resy, resD, resZ, epsilon, X, Z, D, y, *, name):\n","    '''\n","    Convenience summary function that takes the results of the DML function\n","    and summarizes several estimation quantities and performance metrics.\n","    '''\n","    return pd.DataFrame({'estimate': point, # point estimate\n","                         'stderr': stderr, # standard error\n","                         'lower': point - 1.96*stderr, # lower end of 95% confidence interval\n","                         'upper': point + 1.96*stderr, # upper end of 95% confidence interval\n","                         'rmse y': np.sqrt(np.mean(resy**2)), # RMSE of model that predicts outcome y\n","                         'rmse D': np.sqrt(np.mean(resD**2)), # RMSE of model that predicts treatment D\n","                         'rmse Z': np.sqrt(np.mean(resZ**2)), # RMSE of model that predicts treatment D\n","                         'accuracy D': np.mean(np.abs(resD) < .5), # binary classification accuracy of model for D\n","                         'accuracy Z': np.mean(np.abs(resZ) < .5), # binary classification accuracy of model for Z\n","                         }, index=[name])"]},{"cell_type":"markdown","id":"50821e2a","metadata":{"id":"50821e2a"},"source":["#### Double Lasso with Cross-Fitting"]},{"cell_type":"code","execution_count":null,"id":"3245907f","metadata":{"id":"3245907f"},"outputs":[],"source":["cv = KFold(n_splits=5, shuffle=True, random_state=123)\n","lassoy = make_pipeline(transformer, StandardScaler(), LassoCV(cv=cv))\n","lassod = make_pipeline(transformer, StandardScaler(), LassoCV(cv=cv))\n","lassoz = make_pipeline(transformer, StandardScaler(), LassoCV(cv=cv))\n","result = dml(X, Z, D, y, lassoy, lassod, lassoz, nfolds=3)"]},{"cell_type":"code","execution_count":null,"id":"cb8e6883","metadata":{"id":"cb8e6883"},"outputs":[],"source":["table = summary(*result, X, Z, D, y, name='double lasso')\n","table"]},{"cell_type":"markdown","id":"30317504","metadata":{"id":"30317504"},"source":["#### Using a Penalized Logistic Regression for D"]},{"cell_type":"code","execution_count":null,"id":"24237b0c","metadata":{"id":"24237b0c"},"outputs":[],"source":["cv = KFold(n_splits=5, shuffle=True, random_state=123)\n","lassoy = make_pipeline(transformer, StandardScaler(), LassoCV(cv=cv))\n","lgrd = make_pipeline(transformer, StandardScaler(), LogisticRegressionCV(cv=cv))\n","lgrz = make_pipeline(transformer, StandardScaler(), LogisticRegressionCV(cv=cv))\n","result = dml(X, Z, D, y, lassoy, lgrd, lgrz, nfolds=3, classifier=True)"]},{"cell_type":"code","execution_count":null,"id":"de68f123","metadata":{"id":"de68f123"},"outputs":[],"source":["table = table.append(summary(*result,  X, Z, D, y, name='lasso/logistic'))\n","table"]},{"cell_type":"markdown","id":"1b45a835","metadata":{"id":"1b45a835"},"source":["### Random Forests"]},{"cell_type":"code","execution_count":null,"id":"4007ff4e","metadata":{"id":"4007ff4e"},"outputs":[],"source":["rfy = make_pipeline(transformer, RandomForestRegressor(n_estimators=100, min_samples_leaf=10, ccp_alpha=.001))\n","rfd = make_pipeline(transformer, RandomForestClassifier(n_estimators=100, min_samples_leaf=10, ccp_alpha=.001))\n","rfz = make_pipeline(transformer, RandomForestClassifier(n_estimators=100, min_samples_leaf=10, ccp_alpha=.001))\n","result = dml(X, Z, D, y, rfy, rfd, rfz, nfolds=3, classifier=True)"]},{"cell_type":"code","execution_count":null,"id":"801cf22e","metadata":{"id":"801cf22e"},"outputs":[],"source":["table = table.append(summary(*result,  X, Z, D, y, name='random forest'))\n","table"]},{"cell_type":"markdown","id":"f831c007","metadata":{"id":"f831c007"},"source":["### Decision Trees"]},{"cell_type":"code","execution_count":null,"id":"36d6e07a","metadata":{"id":"36d6e07a"},"outputs":[],"source":["dtry = make_pipeline(transformer, DecisionTreeRegressor(min_samples_leaf=10, ccp_alpha=.001))\n","dtrd = make_pipeline(transformer, DecisionTreeClassifier(min_samples_leaf=10, ccp_alpha=.001))\n","dtrz = make_pipeline(transformer, DecisionTreeClassifier(min_samples_leaf=10, ccp_alpha=.001))\n","result = dml(X, Z, D, y, dtry, dtrd, dtrz, nfolds=3, classifier=True)"]},{"cell_type":"code","execution_count":null,"id":"2cdb226d","metadata":{"id":"2cdb226d"},"outputs":[],"source":["table = table.append(summary(*result,  X, Z, D, y, name='decision tree'))\n","table"]},{"cell_type":"markdown","id":"58ccc781","metadata":{"id":"58ccc781"},"source":["### Boosted Trees"]},{"cell_type":"code","source":["gbfy = make_pipeline(transformer, GradientBoostingRegressor(max_depth=2, n_iter_no_change=5))\n","gbfd = make_pipeline(transformer, GradientBoostingClassifier(max_depth=2, n_iter_no_change=5))\n","gbfz = make_pipeline(transformer, GradientBoostingClassifier(max_depth=2, n_iter_no_change=5))\n","result = dml(X, Z, D, y, gbfy, gbfd, gbfz, nfolds=3, classifier=True)"],"metadata":{"id":"cedPovUlUcJt"},"id":"cedPovUlUcJt","execution_count":null,"outputs":[]},{"cell_type":"code","execution_count":null,"id":"76fd3e93","metadata":{"id":"76fd3e93"},"outputs":[],"source":["table = table.append(summary(*result,  X, Z, D, y, name='boosted forest'))\n","table"]},{"cell_type":"markdown","id":"911a6099","metadata":{"id":"911a6099"},"source":["## Semi-Crossfitting and AutoML"]},{"cell_type":"code","source":["!pip install flaml"],"metadata":{"id":"EEaB50ly5FhE"},"id":"EEaB50ly5FhE","execution_count":null,"outputs":[]},{"cell_type":"code","execution_count":null,"id":"32bfd161","metadata":{"id":"32bfd161"},"outputs":[],"source":["from flaml import AutoML\n","\n","flamly = make_pipeline(transformer, AutoML(time_budget=100, task='regression', early_stop=True,\n","                                    eval_method='cv', n_splits=3, metric='r2', verbose=0))\n","flamld = make_pipeline(transformer, AutoML(time_budget=100, task='classification', early_stop=True,\n","                                           eval_method='cv', n_splits=3, metric='r2', verbose=0))\n","flamlz = make_pipeline(transformer, AutoML(time_budget=100, task='classification', early_stop=True,\n","                                           eval_method='cv', n_splits=3, metric='r2', verbose=0))"]},{"cell_type":"code","execution_count":null,"id":"e6bf41b6","metadata":{"id":"e6bf41b6"},"outputs":[],"source":["flamly.fit(X, y)\n","besty = make_pipeline(transformer, clone(flamly[-1].best_model_for_estimator(flamly[-1].best_estimator)))"]},{"cell_type":"code","execution_count":null,"id":"0841262d","metadata":{"id":"0841262d"},"outputs":[],"source":["flamld.fit(X, D)\n","bestd = make_pipeline(transformer, clone(flamld[-1].best_model_for_estimator(flamld[-1].best_estimator)))"]},{"cell_type":"code","execution_count":null,"id":"d0a3a5fd","metadata":{"id":"d0a3a5fd"},"outputs":[],"source":["flamlz.fit(X, Z)\n","bestz = make_pipeline(transformer, clone(flamlz[-1].best_model_for_estimator(flamlz[-1].best_estimator)))"]},{"cell_type":"code","execution_count":null,"id":"87e10995","metadata":{"id":"87e10995"},"outputs":[],"source":["result = dml(X, Z, D, y, besty, bestd, bestz, nfolds=3, classifier=True)"]},{"cell_type":"code","execution_count":null,"id":"dd2fd52a","metadata":{"id":"dd2fd52a"},"outputs":[],"source":["table = table.append(summary(*result,  X, Z, D, y, name='automl (semi-cfit)'))\n","table"]},{"cell_type":"markdown","id":"a14f6004","metadata":{"id":"a14f6004"},"source":["## Inference Robust to Weak Identification"]},{"cell_type":"code","execution_count":null,"id":"3622cc4d","metadata":{"id":"3622cc4d"},"outputs":[],"source":["import scipy.stats\n","\n","def robust_inference(point, stderr, yhat, Dhat, Zhat, resy, resD, resZ, epsilon, X, Z, D, y, *, grid, alpha=0.05):\n","    '''\n","    Inference in the partially linear IV model that is robust to weak identification.\n","    grid: grid of theta values to search over when trying to identify the confidence region\n","    alpha: confidence level\n","    '''\n","    n = X.shape[0]\n","    thr = scipy.stats.chi2.ppf(1 - alpha, df=1)\n","    accept = []d\n","    for theta in grid:\n","        moment = (resy - theta * resD) * resZ\n","        test = n * np.mean(moment)**2 / np.var(moment)\n","        if test <= thr:\n","            accept.append(theta)\n","    return accept"]},{"cell_type":"code","execution_count":null,"id":"9f668209","metadata":{"id":"9f668209"},"outputs":[],"source":["region = robust_inference(*result, X, Z, D, y, grid=np.linspace(0, 20000, 10000))"]},{"cell_type":"code","execution_count":null,"id":"f6eb57a1","metadata":{"id":"f6eb57a1"},"outputs":[],"source":["np.min(region), np.max(region)"]},{"cell_type":"markdown","id":"8326cd85","metadata":{"id":"8326cd85"},"source":["We find that the robust inference confidence region is almost identical to the normal based inference. We are most probably in the strong instrument regime. We can check the t-statistic for the effect of the instrument on the treatment, to verify that"]},{"cell_type":"code","execution_count":null,"id":"e0e7e6d9","metadata":{"id":"e0e7e6d9"},"outputs":[],"source":["beta = np.mean(resZ * resD) / np.mean(resZ**2)\n","var_beta = np.mean((resD - beta * resZ)**2 * resZ**2) / np.mean(resZ**2)**2\n","se_beta = np.sqrt(var_beta / resD.shape[0])\n","print(np.abs(beta) / se_beta)"]},{"cell_type":"markdown","id":"74e607de","metadata":{"id":"74e607de"},"source":["Since the $t$-statistic is very large (much larger than the rule of thumb of $4$), the normal based approximation and confidence intervals should be fine. We can also get this $t$-statistic by simply using the statsmodels package"]},{"cell_type":"code","execution_count":null,"id":"7be6a610","metadata":{"id":"7be6a610"},"outputs":[],"source":["from statsmodels.api import OLS\n","OLS(endog=resD, exog=resZ, hasconst=False).fit(cov_type='HC0').summary()"]},{"cell_type":"markdown","id":"db39aa34","metadata":{"id":"db39aa34"},"source":["# Interactive IV Model and LATE"]},{"cell_type":"markdown","id":"7255c0a8","metadata":{"id":"7255c0a8"},"source":["Now, we consider estimation of local average treatment effects (LATE) of participation `p401`, with the binary instrument `e401`. As before, $Y$ denotes the outcome `net_tfa`, and $X$ is the vector of covariates.  Here the structural equation model is:\n","\\begin{eqnarray}\n","Y &:=&  f_Y (D, X, A, \\epsilon_Y) \\\\\n","D &:= & f_D(Z, X, A, \\epsilon_D) \\in \\{0,1\\},  \\\\\n","Z  &:= & f_Z(X,\\epsilon_Z) \\in \\{0,1\\},  \\\\\n","X &:=&  \\epsilon_X, \\quad A = \\epsilon_A,\n","\\end{eqnarray}\n","where $\\epsilon$'s are all exogenous and independent,\n","and\n","$$\n","z \\mapsto f_D(z , A, X, \\epsilon_D) \\text{ is weakly increasing (weakly monotone)}.\n","$$\n","and $A$ is a vector of unobserved confounders. Note that in our setting monotonicity is satisfied, since participation is only feasible when it is eligible. Thus we have that $D=0$ whenever $Z=0$. Thus it can only be that $f_D(1, A, X, \\epsilon_D) \\geq 0 = f_D(0, A, X, \\epsilon_D)$.\n","\n","In this case, we can estimate the local average treatment effect (LATE):\n","$$\n","\\alpha = E[Y(1) - Y(0) | D(1) > D(0)]\n","$$\n","This can be identified using the Neyman orthogonal moment equation:\n","\\begin{align}\n","E\\left[g(1, X) - g(0, X) + H(Z) (Y - g(Z, X)) - \\alpha \\cdot  (m(1, X) - m(0, X) + H(Z) (D - m(Z, X))\\right] = 0\n","\\end{align}\n","where\n","\\begin{align}\n","g(Z,X) =~& E[Y|Z,X],\\\\\n","m(Z,X) =~& E[D|Z,X],\\\\\n","H(Z) =~& \\frac{Z}{Pr(Z=1|X)} - \\frac{1 - Z}{1 - Pr(Z=1|X)}\n","\\end{align}"]},{"cell_type":"code","execution_count":null,"id":"efc66d4b","metadata":{"id":"efc66d4b"},"outputs":[],"source":["def iiv(X, Z, D, y, modely0, modely1, modeld1, modeld0, modelz, *, trimming=0.01, nfolds):\n","    '''\n","    DML for the Interactive IV Model setting with cross-fitting\n","\n","    Input\n","    -----\n","    X: the controls\n","    D: the treatment\n","    y: the outcome\n","    modely0: the ML model for predicting the outcome y in the Z=0 population\n","    modely1: the ML model for predicting the outcome y in the Z=1 population\n","    modeld0: the ML model for predicting the treatment D in the Z=0 population\n","    modeld1: the ML model for predicting the treatment D in the Z=1 population\n","    modelz: the ML model for predicting the instrument Z\n","    trimming: threshold below which to trim propensities\n","    nfolds: the number of folds in cross-fitting\n","\n","    Output\n","    ------\n","    point: the point estimate of the treatment effect of D on y\n","    stderr: the standard error of the treatment effect\n","    yhat: the cross-fitted predictions for the outcome y\n","    Dhat: the cross-fitted predictions for the outcome D\n","    resy: the outcome residuals\n","    resD: the treatment residuals\n","    drhat: the doubly robust quantity for each sample\n","    '''\n","    cv = KFold(n_splits=nfolds, shuffle=True, random_state=123)\n","\n","    yhat0, yhat1 = np.zeros(y.shape), np.zeros(y.shape)\n","    Dhat0, Dhat1 = np.zeros(D.shape), np.zeros(D.shape)\n","\n","    # we will fit a model E[Y| D, X] by fitting a separate model for D==0\n","    # and a separate model for D==1.\n","    for train, test in cv.split(X, y):\n","        # train an outcome model on training data that received Z=0 and predict outcome on all data in test set\n","        yhat0[test] = clone(modely0).fit(X.iloc[train][Z[train]==0], y[train][Z[train]==0]).predict(X.iloc[test])\n","        # train an outcome model on training data that received Z=1 and predict outcome on all data in test set\n","        yhat1[test] = clone(modely1).fit(X.iloc[train][Z[train]==1], y[train][Z[train]==1]).predict(X.iloc[test])\n","        # train a treatment model on training data that received Z=0 and predict treatment on all data in test set\n","        if np.mean(D[train][Z[train]==0]) > 0: # it could be that D=0, whenever Z=0 deterministically\n","            modeld0_ = clone(modeld0).fit(X.iloc[train][Z[train]==0], D[train][Z[train]==0])\n","            Dhat0[test] = modeld0_.predict_proba(X.iloc[test])[:, 1]\n","        # train a treamtent model on training data that received Z=1 and predict treatment on all data in test set\n","        if np.mean(D[train][Z[train]==1]) < 1: # it could be that D=1, whenever Z=1 deterministically\n","            modeld1_ = clone(modeld1).fit(X.iloc[train][Z[train]==1], D[train][Z[train]==1])\n","            Dhat1[test] = modeld1_.predict_proba(X.iloc[test])[:, 1]\n","        else:\n","            Dhat1[test] = 1\n","\n","    # prediction of treatment and outcome for observed instrument\n","    yhat = yhat0 * (1 - Z) + yhat1 * Z\n","    Dhat = Dhat0 * (1 - Z) + Dhat1 * Z\n","    # propensity scores\n","    Zhat = cross_val_predict(modelz, X, Z, cv=cv, method='predict_proba', n_jobs=-1)[:, 1]\n","    Zhat = np.clip(Zhat, trimming, 1 - trimming)\n","    # doubly robust quantity for every sample\n","    HZ = Z/Zhat - (1 - Z)/(1 - Zhat)\n","    drZ = yhat1 - yhat0 + (y - yhat) * HZ\n","    drD = Dhat1 - Dhat0 + (D - Dhat) * HZ\n","    point = np.mean(drZ) / np.mean(drD)\n","    psi = drZ - point * drD\n","    Jhat = np.mean(drD)\n","    var = np.mean(psi**2) / Jhat**2\n","    stderr = np.sqrt(var / X.shape[0])\n","    return point, stderr, yhat, Dhat, Zhat, y - yhat, D - Dhat, Z - Zhat, drZ, drD"]},{"cell_type":"code","execution_count":null,"id":"7d65020c","metadata":{"id":"7d65020c"},"outputs":[],"source":["def summary(point, stderr, yhat, Dhat, Zhat, resy, resD, resZ, drZ, drD, X, Z, D, y, *, name):\n","    '''\n","    Convenience summary function that takes the results of the DML function\n","    and summarizes several estimation quantities and performance metrics.\n","    '''\n","    return pd.DataFrame({'estimate': point, # point estimate\n","                         'stderr': stderr, # standard error\n","                         'lower': point - 1.96*stderr, # lower end of 95% confidence interval\n","                         'upper': point + 1.96*stderr, # upper end of 95% confidence interval\n","                         'rmse y': np.sqrt(np.mean(resy**2)), # RMSE of model that predicts outcome y\n","                         'rmse D': np.sqrt(np.mean(resD**2)), # RMSE of model that predicts treatment D\n","                         'rmse Z': np.sqrt(np.mean(resZ**2)), # RMSE of model that predicts treatment D\n","                         'accuracy D': np.mean(np.abs(resD) < .5), # binary classification accuracy of model for D\n","                         'accuracy Z': np.mean(np.abs(resZ) < .5), # binary classification accuracy of model for Z\n","                         }, index=[name])"]},{"cell_type":"code","execution_count":null,"id":"7d894f58","metadata":{"id":"7d894f58"},"outputs":[],"source":["cv = KFold(n_splits=5, shuffle=True, random_state=123)\n","lassoy = make_pipeline(transformer, StandardScaler(), LassoCV(cv=cv))\n","lgrd = make_pipeline(transformer, StandardScaler(), LogisticRegressionCV(cv=cv))\n","lgrz = make_pipeline(transformer, StandardScaler(), LogisticRegressionCV(cv=cv))\n","result = iiv(X, Z, D, y, lassoy, lassoy, lgrd, lgrd, lgrz, nfolds=3)"]},{"cell_type":"code","execution_count":null,"id":"c663299a","metadata":{"id":"c663299a"},"outputs":[],"source":["tableiiv = summary(*result, X, Z, D, y, name='lasso/logistic')\n","tableiiv"]},{"cell_type":"code","execution_count":null,"id":"8617137d","metadata":{"id":"8617137d"},"outputs":[],"source":["rfy = make_pipeline(transformer, RandomForestRegressor(n_estimators=100, min_samples_leaf=10, ccp_alpha=.001))\n","rfd = make_pipeline(transformer, RandomForestClassifier(n_estimators=100, min_samples_leaf=10, ccp_alpha=.001))\n","rfz = make_pipeline(transformer, RandomForestClassifier(n_estimators=100, min_samples_leaf=10, ccp_alpha=.001))\n","result = iiv(X, Z, D, y, rfy, rfy, rfd, rfd, rfz, nfolds=3)"]},{"cell_type":"code","execution_count":null,"id":"365ff0f4","metadata":{"id":"365ff0f4"},"outputs":[],"source":["tableiiv = tableiiv.append(summary(*result, X, Z, D, y, name='random forest'))\n","tableiiv"]},{"cell_type":"code","execution_count":null,"id":"571112a4","metadata":{"id":"571112a4"},"outputs":[],"source":["dtry = make_pipeline(transformer, DecisionTreeRegressor(min_samples_leaf=10, ccp_alpha=.001))\n","dtrd = make_pipeline(transformer, DecisionTreeClassifier(min_samples_leaf=10, ccp_alpha=.001))\n","dtrz = make_pipeline(transformer, DecisionTreeClassifier(min_samples_leaf=10, ccp_alpha=.001))\n","result = iiv(X, Z, D, y, dtry, dtry, dtrd, dtrd, dtrz, nfolds=3)"]},{"cell_type":"code","execution_count":null,"id":"04228398","metadata":{"id":"04228398"},"outputs":[],"source":["tableiiv = tableiiv.append(summary(*result, X, Z, D, y, name='decision trees'))\n","tableiiv"]},{"cell_type":"code","execution_count":null,"id":"cba4c9e9","metadata":{"id":"cba4c9e9"},"outputs":[],"source":["gbfy = make_pipeline(transformer, GradientBoostingRegressor(max_depth=2, n_iter_no_change=5))\n","gbfd = make_pipeline(transformer, GradientBoostingClassifier(max_depth=2, n_iter_no_change=5))\n","gbfz = make_pipeline(transformer, GradientBoostingClassifier(max_depth=2, n_iter_no_change=5))\n","result = iiv(X, Z, D, y, gbfy, gbfy, gbfd, gbfd, gbfz, nfolds=3)"]},{"cell_type":"code","execution_count":null,"id":"b221dba2","metadata":{"id":"b221dba2"},"outputs":[],"source":["tableiiv = tableiiv.append(summary(*result, X, Z, D, y, name='boosted trees'))\n","tableiiv"]},{"cell_type":"markdown","id":"9c96e854","metadata":{"id":"9c96e854"},"source":["## Semi-Crossfitting"]},{"cell_type":"code","execution_count":null,"id":"457b513f","metadata":{"id":"457b513f"},"outputs":[],"source":["from flaml import AutoML\n","flamly0 = make_pipeline(transformer, AutoML(time_budget=60, task='regression', early_stop=True,\n","                                     eval_method='cv', n_splits=3, metric='r2', verbose=0))\n","flamly1 = make_pipeline(transformer, AutoML(time_budget=60, task='regression', early_stop=True,\n","                                     eval_method='cv', n_splits=3, metric='r2', verbose=0))\n","flamld1 = make_pipeline(transformer, AutoML(time_budget=60, task='classification', early_stop=True,\n","                                           eval_method='cv', n_splits=3, metric='r2', verbose=0))\n","flamlz = make_pipeline(transformer, AutoML(time_budget=60, task='classification', early_stop=True,\n","                                           eval_method='cv', n_splits=3, metric='r2', verbose=0))"]},{"cell_type":"code","execution_count":null,"id":"621d086a","metadata":{"id":"621d086a"},"outputs":[],"source":["flamly0.fit(X[Z==0], y[Z==0])\n","besty0 = make_pipeline(transformer, clone(flamly0[-1].best_model_for_estimator(flamly0[-1].best_estimator)))"]},{"cell_type":"code","execution_count":null,"id":"a348199a","metadata":{"id":"a348199a"},"outputs":[],"source":["flamly1.fit(X[Z==1], y[Z==1])\n","besty1 = make_pipeline(transformer, clone(flamly1[-1].best_model_for_estimator(flamly1[-1].best_estimator)))"]},{"cell_type":"code","execution_count":null,"id":"4b8c3648","metadata":{"id":"4b8c3648"},"outputs":[],"source":["from sklearn.dummy import DummyClassifier\n","bestd0 = DummyClassifier() # since D=0 whenever Z=0"]},{"cell_type":"code","execution_count":null,"id":"e7326607","metadata":{"id":"e7326607"},"outputs":[],"source":["flamld1.fit(X[Z==1], D[Z==1])\n","bestd1 = make_pipeline(transformer, clone(flamld1[-1].best_model_for_estimator(flamld1[-1].best_estimator)))"]},{"cell_type":"code","execution_count":null,"id":"0cad169d","metadata":{"id":"0cad169d"},"outputs":[],"source":["flamlz.fit(X, Z)\n","bestz = make_pipeline(transformer, clone(flamlz[-1].best_model_for_estimator(flamlz[-1].best_estimator)))"]},{"cell_type":"code","execution_count":null,"id":"50a9df7f","metadata":{"id":"50a9df7f"},"outputs":[],"source":["result = iiv(X, Z, D, y, besty0, besty1, bestd0, bestd1, bestz, nfolds=3)"]},{"cell_type":"code","execution_count":null,"id":"788a551b","metadata":{"id":"788a551b"},"outputs":[],"source":["tableiiv = tableiiv.append(summary(*result, X, Z, D, y, name='automl (semi-cfit)'))\n","tableiiv"]},{"cell_type":"markdown","id":"062e99e8","metadata":{"id":"062e99e8"},"source":["Comparing with the PLR model"]},{"cell_type":"code","execution_count":null,"id":"7f36700c","metadata":{"id":"7f36700c"},"outputs":[],"source":["table"]},{"cell_type":"markdown","id":"a1ecb257","metadata":{"id":"a1ecb257"},"source":["We find that the PLR model overestimates the effect by around 1k; though both sets of results have overlapping confidence intervals."]},{"cell_type":"code","execution_count":null,"id":"e2b38fb1","metadata":{"id":"e2b38fb1"},"outputs":[],"source":["import scipy.stats\n","\n","def iivm_robust_inference(point, stderr, yhat, Dhat, Zhat, resy, resD, resZ, drZ, drD, X, Z, D, y, *, grid, alpha=0.05):\n","    '''\n","    Inference in the partially linear IV model that is robust to weak identification.\n","    grid: grid of theta values to search over when trying to identify the confidence region\n","    alpha: confidence level\n","    '''\n","    n = X.shape[0]\n","    thr = scipy.stats.chi2.ppf(1 - alpha, df=1)\n","    accept = []\n","    for theta in grid:\n","        moment = drZ - theta * drD\n","        test = n * np.mean(moment)**2 / np.var(moment)\n","        if test <= thr:\n","            accept.append(theta)\n","    return accept"]},{"cell_type":"code","execution_count":null,"id":"60ace0a2","metadata":{"id":"60ace0a2"},"outputs":[],"source":["region = iivm_robust_inference(*result, X, Z, D, y, grid=np.linspace(0, 20000, 10000))"]},{"cell_type":"code","execution_count":null,"id":"f7bb488d","metadata":{"id":"f7bb488d"},"outputs":[],"source":["np.min(region), np.max(region)"]},{"cell_type":"markdown","id":"225cb858","metadata":{"id":"225cb858"},"source":["We find again that the robust inference confidence region is almost identical to the normal based inference. We are most probably in the strong instrument regime. We can check the t-statistic for the effect of the instrument on the treatment to verify this."]},{"cell_type":"markdown","id":"8cda6326","metadata":{"id":"8cda6326"},"source":["# Using EconML\n","\n","There exist nice packages out there that can help us do our estimation with the simple call of a function. Such packages include EconML (Python) and DoubleML (Python and R). We run through PLR IV with EconML and IIVM with DoubleML to illustrate."]},{"cell_type":"code","execution_count":null,"id":"257b9330","metadata":{"id":"257b9330"},"outputs":[],"source":["!pip install econml"]},{"cell_type":"code","execution_count":null,"id":"6108d89d","metadata":{"id":"6108d89d"},"outputs":[],"source":["W = StandardScaler().fit_transform(transformer.fit_transform(X))"]},{"cell_type":"code","execution_count":null,"id":"49587174","metadata":{"id":"49587174"},"outputs":[],"source":["from econml.iv.dml import OrthoIV\n","\n","cv = KFold(n_splits=5, shuffle=True, random_state=123)\n","plriv = OrthoIV(model_y_xw=LassoCV(cv=cv),\n","                model_t_xw=LogisticRegressionCV(cv=cv),\n","                model_z_xw=LogisticRegressionCV(cv=cv),\n","                cv=3, discrete_treatment=True, discrete_instrument=True, random_state=123)"]},{"cell_type":"code","execution_count":null,"id":"54524cc0","metadata":{"id":"54524cc0"},"outputs":[],"source":["plriv.fit(y, D, Z=Z, W=W)"]},{"cell_type":"code","execution_count":null,"id":"fa3f1497","metadata":{"id":"fa3f1497"},"outputs":[],"source":["plriv.summary()"]},{"cell_type":"markdown","id":"216ed55d","metadata":{"id":"216ed55d"},"source":["EconML does not yet support LATE estimation under the fully interactive IV model.\n","\n","It does support however a more general IV model where the variables `X` are allowed to alter the effect heterogeneity and the compliance heterogeneity, but the un-observed confounder is not allowed to jointly alter both. In other words, it assumes a structural equation model of the form:\n","\\begin{eqnarray*}\n","Y & := & g_Y(X, \\epsilon_Y) D + f_Y(A, X, \\epsilon_Y),  \\\\\n","D & := & f_D(Z, X, A, \\epsilon_D), \\\\\n","Z & := & f_Z(X, \\epsilon_Z),\\\\\n","A & : =  & f_A(X, \\epsilon_A), \\\\\n","X & := &  \\epsilon_X,\n","\\end{eqnarray*}\n","where $A$ is a vector of un-observed confounders. Under these assumptions the average treatment effect is identifiable (not just the local average treatment effect). In particular, the average treatment effect can be identified as:\n","\\begin{align}\n","\\alpha := E[Y(1) - Y(0)] = E\\left[ \\frac{Cov(Y, Z\\mid X)}{Cov(D, Z\\mid X)} \\right] = E\\left[ \\frac{E[\\tilde{Y} \\tilde{Z}\\mid X]}{E[\\tilde{Z} \\tilde{D}\\mid X]} \\right]\n","\\end{align}\n","where for any variable $V$ we have $\\tilde{V}=V-E[V|X]$.\n","\n","However, the variance of this method can be quite larger than the LATE method, since it uses a local compliance measure, i.e. $E[\\tilde{Z} \\tilde{D}\\mid X]$, which can be small for some regions of $X$. This extra variance stems from the fact that we are going after a more challenging causal quantity which is the average treatment effect, instead of the local average treatment effect, and hence we need to re-weight the data based on compliance levels, conditional on observable covariates $X$."]},{"cell_type":"code","execution_count":null,"id":"122ab658","metadata":{"id":"122ab658"},"outputs":[],"source":["from econml.iv.dr import LinearDRIV\n","\n","driv = LinearDRIV(model_y_xw=LassoCV(cv=cv),\n","                  model_t_xw=LogisticRegressionCV(cv=cv),\n","                  model_t_xwz=LogisticRegressionCV(cv=cv),\n","                  model_tz_xw=LassoCV(cv=cv),\n","                  flexible_model_effect=LassoCV(cv=cv),\n","                  projection=True,\n","                  discrete_instrument=True, discrete_treatment=True, cv=3, cov_clip=0.01, random_state=123)"]},{"cell_type":"code","execution_count":null,"id":"7e878a86","metadata":{"id":"7e878a86"},"outputs":[],"source":["driv.fit(y, D, Z=Z, W=W)"]},{"cell_type":"code","execution_count":null,"id":"dd8e664a","metadata":{"id":"dd8e664a"},"outputs":[],"source":["driv.summary()"]},{"cell_type":"markdown","id":"a629b091","metadata":{"id":"a629b091"},"source":["# Using the DoubleML Package\n","\n","We run through IIVM using `DoubleML` below to illustrate. The `DoubleML` package internally builds on `mlr3`. We use the meta package `mlr3` to generate predictions with machine learning methods. A comprehensive introduction and description of the `mlr3` package is provided in the [mlr3book](https://mlr3book.mlr-org.com/). A list of all learners that you can use in `mlr3` can be found [here](https://mlr3extralearners.mlr-org.com/articles/learners/list_learners.html). The entry in the columns *mlr3 Package* and *Packages* indicate which packages must be installed/loaded in your R session.\n","\n","You find additional information about `DoubleML` on the package on the package website https://docs.doubleml.org/ and the R documentation page https://docs.doubleml.org/r/stable/."]},{"cell_type":"code","execution_count":null,"id":"4f29d8e7","metadata":{"id":"4f29d8e7"},"outputs":[],"source":["!pip install doubleml"]},{"cell_type":"code","execution_count":null,"id":"ff0bca33","metadata":{"id":"ff0bca33"},"outputs":[],"source":["from doubleml import DoubleMLData\n","dml_data = DoubleMLData.from_arrays(W, y, D, z=Z)\n","print(dml_data)"]},{"cell_type":"code","execution_count":null,"id":"47f735fb","metadata":{"id":"47f735fb"},"outputs":[],"source":["import doubleml as dml\n","\n","class RegWrapper(BaseEstimator):\n","\n","    def __init__(self, clf):\n","        self.clf = clf\n","\n","    def fit(self, X, y):\n","        self.clf_ = clone(self.clf).fit(X, y)\n","        return self\n","\n","    def predict(self, X):\n","        return self.clf_.predict_proba(X)[:, 1]\n","\n","dml_plr_obj = dml.DoubleMLPLIV(dml_data,\n","                               LassoCV(cv=cv),\n","                               RegWrapper(LogisticRegressionCV(cv=cv)),\n","                               RegWrapper(LogisticRegressionCV(cv=cv)),\n","                               n_folds=3)\n","print(dml_plr_obj.fit())"]},{"cell_type":"code","execution_count":null,"id":"bdd0ac4d","metadata":{"id":"bdd0ac4d"},"outputs":[],"source":["import doubleml as dml\n","\n","class Wrapper(BaseEstimator):\n","\n","    def __init__(self, clf):\n","        self.clf = clf\n","\n","    def fit(self, X, y):\n","        if np.mean(y) == 0 or np.mean(y) == 1:\n","            self.clf_ = np.mean(y)\n","        else:\n","            self.clf_ = clone(self.clf).fit(X, y)\n","        self.classes_ = np.array([0, 1])\n","        return self\n","\n","    def predict_proba(self, X):\n","        probs = np.zeros((X.shape[0], 2))\n","        if self.clf_ == 0:\n","            probs[:, 0] = 1\n","            return probs\n","        if self.clf_ == 1:\n","            probs[:, 1] = 1\n","            return probs\n","        return self.clf_.predict_proba(X)\n","\n","    def predict(self, X):\n","        return self.predict_proba(X)[:, 1] >= .5\n","\n","dml_plr_obj = dml.DoubleMLIIVM(dml_data,\n","                               LassoCV(cv=cv),\n","                               Wrapper(LogisticRegressionCV(cv=cv)),\n","                               Wrapper(LogisticRegressionCV(cv=cv)),\n","                               n_folds=3)\n","print(dml_plr_obj.fit())"]}],"metadata":{"kernelspec":{"display_name":"Python 3 (ipykernel)","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.9.13"},"papermill":{"default_parameters":{},"duration":537.547458,"end_time":"2021-03-24T14:22:44.931595","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2021-03-24T14:13:47.384137","version":"2.3.2"},"colab":{"provenance":[{"file_id":"https://github.com/CausalAIBook/MetricsMLNotebooks/blob/main/CM4/python-dml-401k-IV.ipynb","timestamp":1702247540262}]}},"nbformat":4,"nbformat_minor":5}