Session 3: Causal Machine Learning¶

Getting started¶

First, we will load the packages we need for these exercises.

In [1]:
# From base python
import csv                                 # Read and write csv files
import os                                  # Work with file paths
import re                                  # Regular expressions 

# External
import numpy as np                         # Vectorized mathematics
import pandas as pd                        # Dataframes
# Install via `pip install shap` or `conda install -c conda-forge shap`
import doubleml as dml                     # Library for econometrics crossed with machine learning
import doubleml.datasets                   # Necessary import to grab the 401K dataset via the doubleml package

# scikit learn components
from sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier  # GBM models
from sklearn import preprocessing          # For standardizing data for LASSO and Elastic net

Double ML¶

Here we will replicate the part of the analysis from Table 1 (Estimated Mean ATE of 401(k) Eligibility on Net Financial Assets) of the Web Appendix of Chernozhukov et al. 2017 AER. This table is intended to illustrate the use of Double ML for standard economics problems where treatment is not random, and indeed may be endogenously related to the outcome.

We will replicate the "Boosting" model, which can be implemented using XGBoost, which we covered in Session 2.

In [2]:
# Grab the dataset
df = dml.datasets.fetch_401K('DataFrame')
df
Out[2]:
nifa net_tfa tw age inc fsize educ db marr twoearn e401 p401 pira hown
0 0.0 0.0 4500.0 47 6765.0 2 8 0 0 0 0 0 0 1
1 6215.0 1015.0 22390.0 36 28452.0 1 16 0 0 0 0 0 0 1
2 0.0 -2000.0 -2000.0 37 3300.0 6 12 1 0 0 0 0 0 0
3 15000.0 15000.0 155000.0 58 52590.0 2 16 0 1 1 0 0 0 1
4 0.0 0.0 58000.0 32 21804.0 1 11 0 0 0 0 0 0 1
... ... ... ... ... ... ... ... ... ... ... ... ... ... ...
9910 98498.0 98858.0 157858.0 52 73920.0 1 16 1 0 0 1 1 0 1
9911 287.0 6230.0 15730.0 41 42927.0 4 14 0 1 1 1 1 1 1
9912 99.0 6099.0 7406.0 40 23619.0 2 16 1 0 0 1 0 1 0
9913 0.0 -32.0 2468.0 47 14280.0 4 6 1 0 0 1 1 0 0
9914 4000.0 5000.0 8857.0 33 11112.0 1 14 0 0 0 1 1 0 0

9915 rows × 14 columns

In [3]:
# Alternatively, load the dataset ourselves
# Using pandas here, as polars lacks stata support without external packages like polars-readstat
df = pd.read_stata('../../Data/S3_sipp1991.dta')
df
Out[3]:
nifa net_tfa tw age inc fsize educ db marr twoearn e401 p401 pira hown
0 0.0 0.0 4500.0 47 6765.0 2 8 0 0 0 0 0 0 1
1 6215.0 1015.0 22390.0 36 28452.0 1 16 0 0 0 0 0 0 1
2 0.0 -2000.0 -2000.0 37 3300.0 6 12 1 0 0 0 0 0 0
3 15000.0 15000.0 155000.0 58 52590.0 2 16 0 1 1 0 0 0 1
4 0.0 0.0 58000.0 32 21804.0 1 11 0 0 0 0 0 0 1
... ... ... ... ... ... ... ... ... ... ... ... ... ... ...
9910 98498.0 98858.0 157858.0 52 73920.0 1 16 1 0 0 1 1 0 1
9911 287.0 6230.0 15730.0 41 42927.0 4 14 0 1 1 1 1 1 1
9912 99.0 6099.0 7406.0 40 23619.0 2 16 1 0 0 1 0 1 0
9913 0.0 -32.0 2468.0 47 14280.0 4 6 1 0 0 1 1 0 0
9914 4000.0 5000.0 8857.0 33 11112.0 1 14 0 0 0 1 1 0 0

9915 rows × 14 columns

In [4]:
y = 'net_tfa'
treat = 'e401'
controls = [x for x in df.columns.tolist() if x not in [y, treat]]

df_dml = dml.DoubleMLData(df, y_col=y, d_cols=treat, x_cols=controls)
In [5]:
print(df_dml)
================== DoubleMLData Object ==================

------------------ Data summary      ------------------
Outcome variable: net_tfa
Treatment variable(s): ['e401']
Covariates: ['nifa', 'tw', 'age', 'inc', 'fsize', 'educ', 'db', 'marr', 'twoearn', 'p401', 'pira', 'hown']
Instrument variable(s): None
No. Observations: 9915
------------------ DataFrame info    ------------------
<class 'pandas.DataFrame'>
RangeIndex: 9915 entries, 0 to 9914
Columns: 14 entries, nifa to hown
dtypes: float32(4), int8(10)
memory usage: 251.9 KB

In [6]:
#set up the nonparametric nuisance functions
g_0 = GradientBoostingRegressor(loss='squared_error',
                                learning_rate=0.01,
                                n_estimators=1000,
                                subsample=0.5,
                                max_depth=2
                               )
m_0 = GradientBoostingClassifier(loss='exponential',
                                 learning_rate=0.01,
                                 n_estimators=1000,
                                 subsample=0.5,
                                 max_depth=2
                                )
In [7]:
np.random.seed(1234)
In [8]:
dml_model_irm = dml.DoubleMLIRM(df_dml, g_0, m_0)
In [9]:
print(dml_model_irm.fit())
================== DoubleMLIRM Object ==================

------------------ Data Summary      ------------------
Outcome variable: net_tfa
Treatment variable(s): ['e401']
Covariates: ['nifa', 'tw', 'age', 'inc', 'fsize', 'educ', 'db', 'marr', 'twoearn', 'p401', 'pira', 'hown']
Instrument variable(s): None
No. Observations: 9915

------------------ Score & Algorithm ------------------
Score function: ATE

------------------ Machine Learner   ------------------
Learner ml_g: GradientBoostingRegressor(learning_rate=0.01, max_depth=2, n_estimators=1000,
                          subsample=0.5)
Learner ml_m: GradientBoostingClassifier(learning_rate=0.01, loss='exponential', max_depth=2,
                           n_estimators=1000, subsample=0.5)
Out-of-sample Performance:
Regression:
Learner ml_g0 RMSE: [[11778.09378127]]
Learner ml_g1 RMSE: [[19940.89453865]]
Classification:
Learner ml_m Log Loss: [[0.28146931]]

------------------ Resampling        ------------------
No. folds: 5
No. repeated sample splits: 1

------------------ Fit Summary       ------------------
             coef    std err         t         P>|t|        2.5 %       97.5 %
e401  3389.046741  364.77854  9.290697  1.532818e-20  2674.093941  4103.999542
In [10]:
dml_model_irm_ATTE = dml.DoubleMLIRM(df_dml, g_0, m_0, score='ATTE')
print(dml_model_irm_ATTE.fit())
================== DoubleMLIRM Object ==================

------------------ Data Summary      ------------------
Outcome variable: net_tfa
Treatment variable(s): ['e401']
Covariates: ['nifa', 'tw', 'age', 'inc', 'fsize', 'educ', 'db', 'marr', 'twoearn', 'p401', 'pira', 'hown']
Instrument variable(s): None
No. Observations: 9915

------------------ Score & Algorithm ------------------
Score function: ATTE

------------------ Machine Learner   ------------------
Learner ml_g: GradientBoostingRegressor(learning_rate=0.01, max_depth=2, n_estimators=1000,
                          subsample=0.5)
Learner ml_m: GradientBoostingClassifier(learning_rate=0.01, loss='exponential', max_depth=2,
                           n_estimators=1000, subsample=0.5)
Out-of-sample Performance:
Regression:
Learner ml_g0 RMSE: [[11179.75700392]]
Learner ml_g1 RMSE: [[20194.76926578]]
Classification:
Learner ml_m Log Loss: [[0.28201354]]

------------------ Resampling        ------------------
No. folds: 5
No. repeated sample splits: 1

------------------ Fit Summary       ------------------
              coef     std err          t          P>|t|        2.5 %  \
e401  10124.713499  402.994626  25.123694  2.740549e-139  9334.858546   

            97.5 %  
e401  10914.568453  

There are two different DML algorithms embedded in the paper. The dml1 algorithm is as described in the slides -- solve for a condition equal to zero for each model, and then average the estimators. The dml2 algorithm solves for the average of the condition being equal to zero overall. The default is to use dml2, but older versions of the doubleml packages allow dml1 to be used by adding dml_procedure='dml1' to the model function.

The library also supports tuning of parameters within the ML models. This can be done by:

  1. Creating a dictionary of the parameters to tune for each ML model
  2. Passing it to the .tune() method of the model.

For more sophisticated methods, you can also rune the models externally via scikit learn and then pass the tuned parameters to the model via the .set_nuisance_params() method of the model.

Lastly, in the spirit of Churnozhukov et al. (2017), in order to get more robust estimates, we can iteratively apply the method to our data to build out our final estimator, which is the median of the set of estimators. This is done using the n_rep= parameter in the model.

Note: This will take a long time to run, since the model is really being run 100 times in the example below.

In [11]:
dml_model_irm = dml.DoubleMLIRM(df_dml, g_0, m_0, n_rep=100)
print(dml_model_irm.fit())
================== DoubleMLIRM Object ==================

------------------ Data Summary      ------------------
Outcome variable: net_tfa
Treatment variable(s): ['e401']
Covariates: ['nifa', 'tw', 'age', 'inc', 'fsize', 'educ', 'db', 'marr', 'twoearn', 'p401', 'pira', 'hown']
Instrument variable(s): None
No. Observations: 9915

------------------ Score & Algorithm ------------------
Score function: ATE

------------------ Machine Learner   ------------------
Learner ml_g: GradientBoostingRegressor(learning_rate=0.01, max_depth=2, n_estimators=1000,
                          subsample=0.5)
Learner ml_m: GradientBoostingClassifier(learning_rate=0.01, loss='exponential', max_depth=2,
                           n_estimators=1000, subsample=0.5)
Out-of-sample Performance:
Regression:
Learner ml_g0 RMSE: [[11969.63145923]
 [11669.06604591]
 [12606.84738787]
 [11228.79002959]
 [11555.44877068]
 [11552.59008415]
 [11153.44388741]
 [11517.51961918]
 [12660.00880103]
 [11534.36766576]
 [11412.56195798]
 [12521.89420405]
 [12099.41962653]
 [11553.82376958]
 [11359.15272755]
 [11626.95602431]
 [11951.12014762]
 [11736.22361002]
 [11701.95532219]
 [11437.40117963]
 [11262.92560267]
 [11663.32602924]
 [11058.31127578]
 [11855.29034392]
 [11674.73074035]
 [11784.19462825]
 [11197.66532592]
 [11268.17931403]
 [11971.46385709]
 [11368.88872091]
 [11722.92417432]
 [11544.06041438]
 [11480.39004809]
 [12231.11109235]
 [12122.76806576]
 [11631.63259563]
 [11520.36089882]
 [11350.28174767]
 [11222.50995043]
 [11438.78305973]
 [11382.9703753 ]
 [11846.6586691 ]
 [11701.35946853]
 [11549.66722888]
 [11778.67208573]
 [11964.55811557]
 [11437.10520431]
 [11464.74646267]
 [11324.68957328]
 [11975.11571735]
 [11812.11332324]
 [11541.19753116]
 [11419.50615668]
 [11609.57454045]
 [13405.69780513]
 [11886.02147534]
 [15301.06911712]
 [11584.28491408]
 [11740.40377427]
 [11536.66371496]
 [11700.09011763]
 [11669.83393436]
 [11850.71642734]
 [11302.46525677]
 [11841.18566125]
 [11663.09642068]
 [12489.86133745]
 [11740.38554156]
 [11819.51945286]
 [11889.94474487]
 [11448.51891552]
 [11472.49609159]
 [12138.29596823]
 [11297.79464775]
 [11275.38787333]
 [13477.37015514]
 [11625.84709178]
 [11704.25385652]
 [10943.19469096]
 [11470.1943593 ]
 [11830.82857514]
 [11611.26761205]
 [12432.75740576]
 [11557.53156897]
 [11266.18659096]
 [11754.93689887]
 [11633.00414847]
 [11566.08222919]
 [11073.67649758]
 [11422.18980647]
 [12197.72067279]
 [11629.01286346]
 [11677.22715098]
 [11579.79297866]
 [11149.17261854]
 [11429.44805781]
 [11543.8112908 ]
 [12065.47781736]
 [11374.03818101]
 [11797.36247007]]
Learner ml_g1 RMSE: [[22146.51590237]
 [22105.49553042]
 [20332.08352929]
 [19870.5906871 ]
 [19577.41377038]
 [19957.83068067]
 [20201.59793539]
 [22101.04239059]
 [20002.78303274]
 [19925.23281243]
 [20166.94161671]
 [22265.36307661]
 [19884.37365623]
 [19979.3871877 ]
 [20207.03003308]
 [19728.88561711]
 [20125.90665093]
 [19918.56133921]
 [19858.77648119]
 [20510.05748008]
 [20195.62452844]
 [19779.87465155]
 [19800.25277647]
 [19979.52150831]
 [21920.48712046]
 [19861.98674962]
 [19980.60037568]
 [22213.61740996]
 [20169.09891375]
 [19934.39902764]
 [21289.73275616]
 [20233.54785034]
 [20156.9291465 ]
 [20216.10867388]
 [20768.89386679]
 [20127.85322873]
 [20415.13300629]
 [19980.49948629]
 [22626.3069747 ]
 [19988.71931773]
 [19891.30900916]
 [19929.91514579]
 [20926.0143383 ]
 [19935.27095843]
 [20049.67787545]
 [20266.361101  ]
 [20046.02721067]
 [20973.96194877]
 [19966.69843269]
 [19932.47700753]
 [20490.67610699]
 [21608.41565894]
 [20716.96471649]
 [21047.11968817]
 [20041.7362165 ]
 [19827.64987404]
 [20196.67934992]
 [20253.61382931]
 [22301.43699323]
 [21772.57965585]
 [21864.56074156]
 [20260.86939476]
 [19994.2047463 ]
 [20084.58921533]
 [21970.25308799]
 [21255.88099197]
 [22359.75870558]
 [22563.69212103]
 [20057.2470023 ]
 [22823.54320494]
 [19965.4018671 ]
 [19957.01346063]
 [20062.04999965]
 [19926.34797758]
 [21337.05924775]
 [19683.66940765]
 [22142.82044832]
 [24055.98309679]
 [20181.92504257]
 [19960.85928159]
 [22553.02499559]
 [19714.03962772]
 [20148.53839373]
 [20065.88238908]
 [20355.5301442 ]
 [19942.39253442]
 [19839.17333415]
 [19863.30713278]
 [21966.05081331]
 [22008.20913078]
 [22159.77126379]
 [19876.18426779]
 [20863.02296279]
 [19783.02597641]
 [22427.65029978]
 [21251.59741536]
 [22420.68355998]
 [20961.95176387]
 [20399.8069937 ]
 [20163.19931031]]
Classification:
Learner ml_m Log Loss: [[0.28104053]
 [0.28127035]
 [0.28139405]
 [0.28149888]
 [0.28132159]
 [0.28065756]
 [0.28118568]
 [0.28085574]
 [0.28148896]
 [0.28112525]
 [0.28135986]
 [0.28078396]
 [0.28079083]
 [0.28233891]
 [0.28055145]
 [0.28037458]
 [0.28122419]
 [0.28097508]
 [0.2807339 ]
 [0.28135843]
 [0.2813355 ]
 [0.28085714]
 [0.28146945]
 [0.28070887]
 [0.28094261]
 [0.28073838]
 [0.28060524]
 [0.28167361]
 [0.28094326]
 [0.28097812]
 [0.28092823]
 [0.28051322]
 [0.28020665]
 [0.28102623]
 [0.28103396]
 [0.28071397]
 [0.2812156 ]
 [0.28159909]
 [0.28132355]
 [0.28092987]
 [0.28083592]
 [0.2817106 ]
 [0.28093342]
 [0.28088216]
 [0.28109157]
 [0.28074962]
 [0.28080467]
 [0.28092046]
 [0.28089977]
 [0.2807708 ]
 [0.28167589]
 [0.2813748 ]
 [0.28128588]
 [0.28097445]
 [0.28135801]
 [0.28063503]
 [0.28146423]
 [0.28162482]
 [0.28109829]
 [0.28109146]
 [0.28142461]
 [0.28093473]
 [0.28095476]
 [0.28143566]
 [0.28088342]
 [0.28162472]
 [0.28110198]
 [0.28151489]
 [0.28052541]
 [0.28132601]
 [0.28068403]
 [0.28072184]
 [0.2808023 ]
 [0.28041796]
 [0.28120828]
 [0.28077007]
 [0.28061172]
 [0.28083405]
 [0.28116119]
 [0.28056192]
 [0.28102434]
 [0.28078885]
 [0.28151689]
 [0.28094951]
 [0.28074081]
 [0.28131851]
 [0.2816776 ]
 [0.28095783]
 [0.28107177]
 [0.28079649]
 [0.28092514]
 [0.28066068]
 [0.28060239]
 [0.280456  ]
 [0.28132253]
 [0.28056341]
 [0.28096625]
 [0.28127728]
 [0.2809657 ]
 [0.28128744]]

------------------ Resampling        ------------------
No. folds: 5
No. repeated sample splits: 100

------------------ Fit Summary       ------------------
             coef     std err         t         P>|t|     2.5 %       97.5 %
e401  3333.963581  410.427124  8.123156  4.542147e-16  2528.341  4138.386845

Finally, a much faster, multithreaded, but experimental implementation of GBM from scikit learn, going through the full example. The HistGradientBoosting methods are fully multithreaded and will use 100% of your CPU to solve the models, whereas the standard GBM model in scikit learn is single threaded.

In [13]:
from sklearn.ensemble import HistGradientBoostingClassifier, HistGradientBoostingRegressor

# set up the data
df = pd.read_stata('../../Data/S3_sipp1991.dta')

y = 'net_tfa'
treat = 'e401'
controls = [x for x in df.columns.tolist() if x not in [y, treat]]

df_dml3 = dml.DoubleMLData(df, y_col=y, d_cols=treat, x_cols=controls)

#set up the nonparametric nuisance functions
g_0 = HistGradientBoostingRegressor(loss='squared_error',
                                    learning_rate=0.01,
                                    max_iter=1000,
                                    max_depth=2,
                                    early_stopping=False
                                   )
m_0 = HistGradientBoostingClassifier(loss='log_loss',
                                     learning_rate=0.01,
                                     max_iter=1000,
                                     max_depth=2,
                                     early_stopping=False
                                    )


np.random.seed(1234)
dml_model_ex_irm = dml.DoubleMLIRM(df_dml, g_0, m_0, n_folds=5, n_rep=100)
print(dml_model_ex_irm.fit())
================== DoubleMLIRM Object ==================

------------------ Data Summary      ------------------
Outcome variable: net_tfa
Treatment variable(s): ['e401']
Covariates: ['nifa', 'tw', 'age', 'inc', 'fsize', 'educ', 'db', 'marr', 'twoearn', 'p401', 'pira', 'hown']
Instrument variable(s): None
No. Observations: 9915

------------------ Score & Algorithm ------------------
Score function: ATE

------------------ Machine Learner   ------------------
Learner ml_g: HistGradientBoostingRegressor(early_stopping=False, learning_rate=0.01,
                              max_depth=2, max_iter=1000)
Learner ml_m: HistGradientBoostingClassifier(early_stopping=False, learning_rate=0.01,
                               max_depth=2, max_iter=1000)
Out-of-sample Performance:
Regression:
Learner ml_g0 RMSE: [[30988.92813967]
 [30732.57461248]
 [30735.17071984]
 [29957.16012004]
 [31905.11752991]
 [30050.81347319]
 [29950.8347059 ]
 [32100.62699952]
 [31538.27387684]
 [30034.26191344]
 [31213.64052258]
 [31432.08044924]
 [29872.37706333]
 [31557.41439414]
 [31327.80451064]
 [30311.56038633]
 [30693.82347671]
 [30505.16867268]
 [31318.846411  ]
 [29802.00442859]
 [31116.14493139]
 [30498.24521657]
 [30643.05509177]
 [30073.77294785]
 [33388.77267738]
 [31165.78171728]
 [30203.3458086 ]
 [30502.58432063]
 [30186.38624931]
 [30939.07466719]
 [29408.06639156]
 [31510.14089581]
 [30467.89867211]
 [30716.90291945]
 [31215.70390277]
 [30990.97530606]
 [30644.39608014]
 [30591.36886414]
 [32961.14192986]
 [29797.72682764]
 [31586.4660577 ]
 [29547.92849315]
 [31056.09035642]
 [31165.0474639 ]
 [29931.24080209]
 [30673.28985466]
 [30580.583642  ]
 [29633.18241061]
 [31478.46914457]
 [29987.54364064]
 [33271.78593986]
 [30342.15704487]
 [30355.2101091 ]
 [31923.07133695]
 [31022.86068932]
 [31053.81515067]
 [31732.8975328 ]
 [30911.97328064]
 [32140.63754471]
 [30047.42280806]
 [32285.49804697]
 [30294.78893745]
 [30105.64676329]
 [30353.53904766]
 [31293.94756227]
 [33015.98004361]
 [30743.13287372]
 [31171.21324529]
 [30595.83916384]
 [32169.81222902]
 [31799.4269403 ]
 [32318.5842791 ]
 [31683.62561396]
 [29986.30331875]
 [32504.90383752]
 [31601.15500492]
 [31105.95715849]
 [31715.07254204]
 [34011.9370723 ]
 [31322.99080855]
 [31684.00026993]
 [31764.87421813]
 [29555.38018217]
 [30554.00488294]
 [30647.06413298]
 [31346.40995758]
 [30570.41934708]
 [32004.74014045]
 [30684.25574724]
 [30784.6306195 ]
 [31640.45395465]
 [31477.13724938]
 [31882.64343004]
 [30674.9469119 ]
 [34390.04745437]
 [31178.79912382]
 [29722.67578042]
 [29881.02445913]
 [31020.62709391]
 [30340.27718666]]
Learner ml_g1 RMSE: [[35735.48326639]
 [37370.83624368]
 [37187.1437244 ]
 [36149.90605806]
 [36423.45903783]
 [35311.25804305]
 [37490.30052102]
 [35880.25497933]
 [37642.39294745]
 [37366.07698432]
 [36646.45980808]
 [36374.14862982]
 [35767.82213595]
 [36361.33792184]
 [36127.44808837]
 [36482.90574047]
 [35735.84742998]
 [37863.58554504]
 [36066.77243748]
 [36535.54602544]
 [36328.2993226 ]
 [35788.54436935]
 [35721.63771316]
 [36411.24514958]
 [35584.63844122]
 [36317.11954195]
 [36409.83144877]
 [37552.21237907]
 [34695.42690909]
 [37503.87464349]
 [36640.4832149 ]
 [36049.49185261]
 [37083.47238385]
 [38623.54882501]
 [35645.64922523]
 [37097.54790457]
 [36323.61929632]
 [40283.17629938]
 [36704.65925588]
 [37477.99045813]
 [36440.94069001]
 [35988.09423888]
 [38122.01602081]
 [35936.26554271]
 [35475.69356526]
 [36501.20230047]
 [35169.40458274]
 [36771.55350993]
 [36362.69608734]
 [35324.00315683]
 [36952.33387844]
 [35710.4597437 ]
 [36996.93126142]
 [36965.62723585]
 [35950.70333765]
 [35448.58351516]
 [37166.33166709]
 [37138.20478368]
 [36431.28529018]
 [36506.08395882]
 [38275.2524454 ]
 [36151.8263853 ]
 [36007.35133288]
 [37541.94555221]
 [37318.08183845]
 [38009.89256284]
 [38904.14986328]
 [36291.88296399]
 [37813.11768   ]
 [35431.06385161]
 [35762.36975155]
 [35531.00422326]
 [36283.5960803 ]
 [37108.9036545 ]
 [36525.5410568 ]
 [36169.02032054]
 [36386.71279841]
 [36889.18173486]
 [36386.11427688]
 [36203.51325643]
 [36415.15979251]
 [36985.21143711]
 [37100.54813497]
 [36013.69084466]
 [37133.3735195 ]
 [35461.95083038]
 [37852.88654438]
 [36840.32914745]
 [35956.47433312]
 [36733.451148  ]
 [35641.01203089]
 [37697.47584448]
 [35447.31045441]
 [36940.25262801]
 [37243.87364485]
 [37013.43586847]
 [36164.03953156]
 [36580.82068582]
 [35919.39288366]
 [35650.26683731]]
Classification:
Learner ml_m Log Loss: [[0.28154875]
 [0.28097902]
 [0.28112478]
 [0.28155949]
 [0.2809131 ]
 [0.28161818]
 [0.28077867]
 [0.28146587]
 [0.28195925]
 [0.28109635]
 [0.28160185]
 [0.28191864]
 [0.28192459]
 [0.28223341]
 [0.281659  ]
 [0.28146221]
 [0.28213085]
 [0.28130823]
 [0.28149584]
 [0.28114933]
 [0.28084834]
 [0.28157751]
 [0.28075   ]
 [0.28128209]
 [0.281116  ]
 [0.28096792]
 [0.28127719]
 [0.28052342]
 [0.28178787]
 [0.28139011]
 [0.28139947]
 [0.28096437]
 [0.28145048]
 [0.2807697 ]
 [0.28139269]
 [0.28155078]
 [0.28165039]
 [0.28179511]
 [0.28173946]
 [0.28138954]
 [0.28146129]
 [0.28123234]
 [0.28157725]
 [0.28083057]
 [0.28134106]
 [0.28154413]
 [0.28131242]
 [0.28115745]
 [0.28169007]
 [0.28125247]
 [0.28124072]
 [0.28067901]
 [0.28208005]
 [0.28165106]
 [0.28164441]
 [0.28168723]
 [0.28081092]
 [0.28173965]
 [0.28165693]
 [0.28125744]
 [0.28187629]
 [0.28114098]
 [0.28155147]
 [0.28106444]
 [0.28170727]
 [0.28152255]
 [0.28191222]
 [0.28192741]
 [0.28157604]
 [0.28110809]
 [0.28096435]
 [0.28117794]
 [0.28124998]
 [0.28122559]
 [0.28148651]
 [0.281162  ]
 [0.28118421]
 [0.28138907]
 [0.28163966]
 [0.28083877]
 [0.2813123 ]
 [0.28168829]
 [0.28126419]
 [0.28122705]
 [0.28154337]
 [0.2810027 ]
 [0.28177674]
 [0.28129942]
 [0.28236101]
 [0.28143536]
 [0.28172892]
 [0.28133788]
 [0.28148676]
 [0.28169126]
 [0.28185015]
 [0.28176228]
 [0.28109261]
 [0.28160537]
 [0.28101752]
 [0.28185913]]

------------------ Resampling        ------------------
No. folds: 5
No. repeated sample splits: 100

------------------ Fit Summary       ------------------
             coef      std err         t     P>|t|        2.5 %       97.5 %
e401  3479.716372  1189.031254  2.926514  0.003428  1067.621409  5810.176317

Try out fully interacted version with LASSO¶

In [14]:
# Grab the dataset
df2 = dml.datasets.fetch_401K('DataFrame')

y = 'net_tfa'
treat = 'e401'
controls = [x for x in df.columns.tolist() if x not in [y, treat]]
In [15]:
# Specify the interactions
controls2 = controls.copy()
for i in controls:
    for j in controls:
        df2[i + '&' + j] = df2[i] * df2[j]
        controls2.append(i + '&' + j)
df2 = df2.copy()
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
<positron-console-cell-15>:5: PerformanceWarning: DataFrame is highly fragmented.  This is usually the result of calling `frame.insert` many times, which has poor performance.  Consider joining all columns at once using pd.concat(axis=1) instead. To get a de-fragmented frame, use `newframe = frame.copy()`
In [16]:
len(controls2)
Out[16]:
156
In [17]:
df2_dml = dml.DoubleMLData(df2, y_col=y, d_cols=treat, x_cols=controls2)
In [18]:
from sklearn.linear_model import Lasso, LogisticRegression
#set up the nonparametric nuisance functions
g2_0 = Lasso()
m2_0 = LogisticRegression(penalty='l1', solver='liblinear')
In [19]:
np.random.seed(1234)
dml2_model_irm = dml.DoubleMLIRM(df2_dml, g2_0, m2_0)
print(dml2_model_irm.fit())
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_coordinate_descent.py:840: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations, check the scale of the features or consider increasing regularisation. Duality gap: 4.325010e+11, tolerance: 1.629e+09
  model = cd_fast.enet_coordinate_descent(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_coordinate_descent.py:840: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations, check the scale of the features or consider increasing regularisation. Duality gap: 3.533061e+11, tolerance: 1.696e+09
  model = cd_fast.enet_coordinate_descent(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_coordinate_descent.py:840: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations, check the scale of the features or consider increasing regularisation. Duality gap: 3.655753e+11, tolerance: 1.519e+09
  model = cd_fast.enet_coordinate_descent(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_coordinate_descent.py:840: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations, check the scale of the features or consider increasing regularisation. Duality gap: 3.355257e+11, tolerance: 1.237e+09
  model = cd_fast.enet_coordinate_descent(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_coordinate_descent.py:840: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations, check the scale of the features or consider increasing regularisation. Duality gap: 4.252765e+11, tolerance: 1.329e+09
  model = cd_fast.enet_coordinate_descent(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_coordinate_descent.py:840: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations, check the scale of the features or consider increasing regularisation. Duality gap: 4.546084e+11, tolerance: 1.611e+09
  model = cd_fast.enet_coordinate_descent(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_coordinate_descent.py:840: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations, check the scale of the features or consider increasing regularisation. Duality gap: 4.388618e+11, tolerance: 1.843e+09
  model = cd_fast.enet_coordinate_descent(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_coordinate_descent.py:840: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations, check the scale of the features or consider increasing regularisation. Duality gap: 4.173802e+11, tolerance: 1.524e+09
  model = cd_fast.enet_coordinate_descent(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_coordinate_descent.py:840: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations, check the scale of the features or consider increasing regularisation. Duality gap: 4.696276e+11, tolerance: 1.801e+09
  model = cd_fast.enet_coordinate_descent(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_coordinate_descent.py:840: ConvergenceWarning: Objective did not converge. You might want to increase the number of iterations, check the scale of the features or consider increasing regularisation. Duality gap: 4.586217e+11, tolerance: 1.459e+09
  model = cd_fast.enet_coordinate_descent(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_logistic.py:1403: FutureWarning: 'penalty' was deprecated in version 1.8 and will be removed in 1.10. To avoid this warning, leave 'penalty' set to its default value and use 'l1_ratio' or 'C' instead. Use l1_ratio=0 instead of penalty='l2', l1_ratio=1 instead of penalty='l1', l1_ratio set to a float between 0 and 1 instead of penalty='elasticnet', and C=np.inf instead of penalty=None.
  warnings.warn(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_logistic.py:1429: UserWarning: Inconsistent values: penalty=l1 with l1_ratio=0.0. penalty is deprecated. Please use l1_ratio only.
  warnings.warn(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\svm\_base.py:1298: ConvergenceWarning: Liblinear failed to converge, increase the number of iterations.
  warnings.warn(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_logistic.py:1403: FutureWarning: 'penalty' was deprecated in version 1.8 and will be removed in 1.10. To avoid this warning, leave 'penalty' set to its default value and use 'l1_ratio' or 'C' instead. Use l1_ratio=0 instead of penalty='l2', l1_ratio=1 instead of penalty='l1', l1_ratio set to a float between 0 and 1 instead of penalty='elasticnet', and C=np.inf instead of penalty=None.
  warnings.warn(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_logistic.py:1429: UserWarning: Inconsistent values: penalty=l1 with l1_ratio=0.0. penalty is deprecated. Please use l1_ratio only.
  warnings.warn(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\svm\_base.py:1298: ConvergenceWarning: Liblinear failed to converge, increase the number of iterations.
  warnings.warn(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_logistic.py:1403: FutureWarning: 'penalty' was deprecated in version 1.8 and will be removed in 1.10. To avoid this warning, leave 'penalty' set to its default value and use 'l1_ratio' or 'C' instead. Use l1_ratio=0 instead of penalty='l2', l1_ratio=1 instead of penalty='l1', l1_ratio set to a float between 0 and 1 instead of penalty='elasticnet', and C=np.inf instead of penalty=None.
  warnings.warn(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_logistic.py:1429: UserWarning: Inconsistent values: penalty=l1 with l1_ratio=0.0. penalty is deprecated. Please use l1_ratio only.
  warnings.warn(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\svm\_base.py:1298: ConvergenceWarning: Liblinear failed to converge, increase the number of iterations.
  warnings.warn(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_logistic.py:1403: FutureWarning: 'penalty' was deprecated in version 1.8 and will be removed in 1.10. To avoid this warning, leave 'penalty' set to its default value and use 'l1_ratio' or 'C' instead. Use l1_ratio=0 instead of penalty='l2', l1_ratio=1 instead of penalty='l1', l1_ratio set to a float between 0 and 1 instead of penalty='elasticnet', and C=np.inf instead of penalty=None.
  warnings.warn(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_logistic.py:1429: UserWarning: Inconsistent values: penalty=l1 with l1_ratio=0.0. penalty is deprecated. Please use l1_ratio only.
  warnings.warn(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\svm\_base.py:1298: ConvergenceWarning: Liblinear failed to converge, increase the number of iterations.
  warnings.warn(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_logistic.py:1403: FutureWarning: 'penalty' was deprecated in version 1.8 and will be removed in 1.10. To avoid this warning, leave 'penalty' set to its default value and use 'l1_ratio' or 'C' instead. Use l1_ratio=0 instead of penalty='l2', l1_ratio=1 instead of penalty='l1', l1_ratio set to a float between 0 and 1 instead of penalty='elasticnet', and C=np.inf instead of penalty=None.
  warnings.warn(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\linear_model\_logistic.py:1429: UserWarning: Inconsistent values: penalty=l1 with l1_ratio=0.0. penalty is deprecated. Please use l1_ratio only.
  warnings.warn(
================== DoubleMLIRM Object ==================

------------------ Data Summary      ------------------
Outcome variable: net_tfa
Treatment variable(s): ['e401']
Covariates: ['nifa', 'tw', 'age', 'inc', 'fsize', 'educ', 'db', 'marr', 'twoearn', 'p401', 'pira', 'hown', 'nifa&nifa', 'nifa&tw', 'nifa&age', 'nifa&inc', 'nifa&fsize', 'nifa&educ', 'nifa&db', 'nifa&marr', 'nifa&twoearn', 'nifa&p401', 'nifa&pira', 'nifa&hown', 'tw&nifa', 'tw&tw', 'tw&age', 'tw&inc', 'tw&fsize', 'tw&educ', 'tw&db', 'tw&marr', 'tw&twoearn', 'tw&p401', 'tw&pira', 'tw&hown', 'age&nifa', 'age&tw', 'age&age', 'age&inc', 'age&fsize', 'age&educ', 'age&db', 'age&marr', 'age&twoearn', 'age&p401', 'age&pira', 'age&hown', 'inc&nifa', 'inc&tw', 'inc&age', 'inc&inc', 'inc&fsize', 'inc&educ', 'inc&db', 'inc&marr', 'inc&twoearn', 'inc&p401', 'inc&pira', 'inc&hown', 'fsize&nifa', 'fsize&tw', 'fsize&age', 'fsize&inc', 'fsize&fsize', 'fsize&educ', 'fsize&db', 'fsize&marr', 'fsize&twoearn', 'fsize&p401', 'fsize&pira', 'fsize&hown', 'educ&nifa', 'educ&tw', 'educ&age', 'educ&inc', 'educ&fsize', 'educ&educ', 'educ&db', 'educ&marr', 'educ&twoearn', 'educ&p401', 'educ&pira', 'educ&hown', 'db&nifa', 'db&tw', 'db&age', 'db&inc', 'db&fsize', 'db&educ', 'db&db', 'db&marr', 'db&twoearn', 'db&p401', 'db&pira', 'db&hown', 'marr&nifa', 'marr&tw', 'marr&age', 'marr&inc', 'marr&fsize', 'marr&educ', 'marr&db', 'marr&marr', 'marr&twoearn', 'marr&p401', 'marr&pira', 'marr&hown', 'twoearn&nifa', 'twoearn&tw', 'twoearn&age', 'twoearn&inc', 'twoearn&fsize', 'twoearn&educ', 'twoearn&db', 'twoearn&marr', 'twoearn&twoearn', 'twoearn&p401', 'twoearn&pira', 'twoearn&hown', 'p401&nifa', 'p401&tw', 'p401&age', 'p401&inc', 'p401&fsize', 'p401&educ', 'p401&db', 'p401&marr', 'p401&twoearn', 'p401&p401', 'p401&pira', 'p401&hown', 'pira&nifa', 'pira&tw', 'pira&age', 'pira&inc', 'pira&fsize', 'pira&educ', 'pira&db', 'pira&marr', 'pira&twoearn', 'pira&p401', 'pira&pira', 'pira&hown', 'hown&nifa', 'hown&tw', 'hown&age', 'hown&inc', 'hown&fsize', 'hown&educ', 'hown&db', 'hown&marr', 'hown&twoearn', 'hown&p401', 'hown&pira', 'hown&hown']
Instrument variable(s): None
No. Observations: 9915

------------------ Score & Algorithm ------------------
Score function: ATE

------------------ Machine Learner   ------------------
Learner ml_g: Lasso()
Learner ml_m: LogisticRegression(penalty='l1', solver='liblinear')
Out-of-sample Performance:
Regression:
Learner ml_g0 RMSE: [[16108.6418434]]
Learner ml_g1 RMSE: [[21007.87498748]]
Classification:
Learner ml_m Log Loss: [[0.28306736]]

------------------ Resampling        ------------------
No. folds: 5
No. repeated sample splits: 1

------------------ Fit Summary       ------------------
             coef     std err         t         P>|t|        2.5 %  \
e401  3337.207813  548.556009  6.083623  1.174969e-09  2262.057792   

           97.5 %  
e401  4412.357835  
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\sklearn\svm\_base.py:1298: ConvergenceWarning: Liblinear failed to converge, increase the number of iterations.
  warnings.warn(
M:\Python_environments\miniconda\envs\MLSS2\Lib\site-packages\doubleml\utils\propensity_score_processing.py:272: UserWarning: Propensity predictions  from learner ml_m are close to zero or one (eps=1e-12).
  warnings.warn(

Other types of DML structures¶

DiD¶

You can find details and worked out examples for the DiD implementation here: https://docs.doubleml.org/dev/guide/models.html#difference-in-differences-models-did

Robust Clustering¶

You can find details and worked out examples for the clustering specifications (one-way cluster robust DML; two-way clustering) here: https://docs.doubleml.org/stable/examples/py_double_ml_multiway_cluster.html

Causal Trees¶

Next, we shift our focus from average treatment effects to heterogeneous treatment effects.

The plan:

  • Use econml's CausalForestDML to estimate observation-level treatment effects.
  • Then fit a shallow decision tree to those estimated effects to get a simple causal-tree-style segmentation.
In [20]:
# Additional imports for the week

from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.tree import DecisionTreeRegressor, export_text

import doubleml as dml
from econml.dml import CausalForestDML
from econml.grf import CausalForest

We will use the same data frame as before, df.

In [21]:
print('Shape:', df.shape)
df.head()
Shape: (9915, 14)
Out[21]:
nifa net_tfa tw age inc fsize educ db marr twoearn e401 p401 pira hown
0 0.0 0.0 4500.0 47 6765.0 2 8 0 0 0 0 0 0 1
1 6215.0 1015.0 22390.0 36 28452.0 1 16 0 0 0 0 0 0 1
2 0.0 -2000.0 -2000.0 37 3300.0 6 12 1 0 0 0 0 0 0
3 15000.0 15000.0 155000.0 58 52590.0 2 16 0 1 1 0 0 0 1
4 0.0 0.0 58000.0 32 21804.0 1 11 0 0 0 0 0 0 1

Prep our data:

In [22]:
y_col = 'net_tfa'
t_col = 'e401'
x_cols = [c for c in df.columns if c not in [y_col, t_col]]

Y = df[y_col].to_numpy()
T = df[t_col].to_numpy()
X = df[x_cols].to_numpy()

print('Outcome:', y_col)
print('Treatment:', t_col)
print('Number of controls:', len(x_cols))
x_cols
Outcome: net_tfa
Treatment: e401
Number of controls: 12
Out[22]:
['nifa',
 'tw',
 'age',
 'inc',
 'fsize',
 'educ',
 'db',
 'marr',
 'twoearn',
 'p401',
 'pira',
 'hown']

Fit a causal forest with GRF¶

In [ ]:
np.random.seed(1234)

# CausalForest expects treatments with shape (n_samples, n_treatments).
# For a binary treatment, encode T numerically as 0/1.
T_grf = np.asarray(T).reshape(-1, 1)

cf = CausalForest(
    n_estimators=400,
    min_samples_leaf=20,
    max_depth=12,
    honest=True,
    #inference=True,
    n_jobs=-1,
    random_state=1234,
)

# Fit the model
cf.fit(X, T_grf, Y)

# predict() returns one treatment-effect coefficient per treatment.
tau_hat = cf.predict(X).ravel()
ate_hat = tau_hat.mean()

print("Estimated ATE from causal forest:", round(float(ate_hat), 3))
Estimated ATE from causal forest: 7497.59
In [24]:
results = df[[y_col, t_col] + x_cols].copy()
results["tau_hat"] = np.asarray(tau_hat).ravel()
results[["tau_hat"]].describe()
Out[24]:
tau_hat
count 9915.000000
mean 7497.589890
std 8163.195882
min -1254.292091
25% 1160.814474
50% 4486.010376
75% 11754.156008
max 37915.731410

Look at heterogeneity across subgroups¶

In [27]:
subgroup_cols = {"homeowner": "hown", "db_pension": "db", "married": "marr"}

# Summarize estimated treatment effects within each subgroup
subgroup_results = {}

for subgroup_name, subgroup in subgroup_cols.items():
    summary = (
        results.groupby(subgroup, dropna=False)["tau_hat"]
        .agg(
            n="size",
            mean_tau="mean",
            median_tau="median",
            std_tau="std",
        )
        .reset_index()
    )

    subgroup_results[subgroup] = summary

    print(f"\nTreatment-effect heterogeneity by {subgroup_name}:")
    print(summary.round(3))

    # Difference between subgroup 1 and subgroup 0
    means = results.groupby(subgroup)["tau_hat"].mean()
    if 0 in means.index and 1 in means.index:
        difference = means.loc[1] - means.loc[0]
        print(f"Difference in mean effects (1 minus 0): {difference:.3f}")
Treatment-effect heterogeneity by homeowner:
   hown     n   mean_tau  median_tau   std_tau
0     0  3617   2375.767     919.405  3805.574
1     1  6298  10439.101    7142.042  8536.744
Difference in mean effects (1 minus 0): 8063.334

Treatment-effect heterogeneity by db_pension:
   db     n  mean_tau  median_tau   std_tau
0   0  7228  7058.466    3999.788  8014.210
1   1  2687  8678.828    5614.549  8439.952
Difference in mean effects (1 minus 0): 1620.362

Treatment-effect heterogeneity by married:
   marr     n  mean_tau  median_tau  std_tau
0     0  3918  5234.574    2567.912  6704.58
1     1  5997  8976.079    5747.977  8676.59
Difference in mean effects (1 minus 0): 3741.505

Fit a causal forest with DML¶

This gives us an estimated treatment effect for each observation, rather than only one overall ATE.

In [28]:
np.random.seed(123)

cf = CausalForestDML(
    model_y=RandomForestRegressor(
        n_estimators=200,
        min_samples_leaf=20,
        random_state=123,
        n_jobs=-1,
    ),
    model_t=RandomForestClassifier(
        n_estimators=200,
        min_samples_leaf=20,
        random_state=123,
        n_jobs=-1,
    ),
    n_estimators=400,
    min_samples_leaf=20,
    max_depth=12,
    discrete_treatment=True,
    random_state=123,
)

cf.fit(Y, T, X=X)
tau_hat = cf.effect(X)
ate_hat = tau_hat.mean()

print('Estimated ATE from causal forest:', round(float(ate_hat), 3))
Estimated ATE from causal forest: -81.787
In [29]:
results = df[[y_col, t_col] + x_cols].copy()
results['tau_hat'] = tau_hat
results[['tau_hat']].describe()
Out[29]:
tau_hat
count 9915.000000
mean -81.786688
std 1976.075105
min -13341.072909
25% -649.020667
50% -91.247868
75% 234.527593
max 29348.714391

Look at heterogeneity across subgroups¶

In [30]:
subgroup_summary = pd.DataFrame({
    'homeowner': results.groupby('hown')['tau_hat'].mean(),
    'db_pension': results.groupby('db')['tau_hat'].mean(),
    'married': results.groupby('marr')['tau_hat'].mean(),
})
subgroup_summary
Out[30]:
homeowner db_pension married
0 323.141311 -93.611176 99.715942
1 -314.340606 -49.978948 -200.367196
In [31]:
results['inc_bin'] = pd.qcut(results['inc'], q=4, duplicates='drop')
results.groupby('inc_bin')['tau_hat'].agg(['mean', 'count']).sort_index()
Out[31]:
mean count
inc_bin
(-2652.001, 19413.0] -125.583017 2481
(19413.0, 31476.0] -124.120971 2477
(31476.0, 48583.5] -28.666981 2478
(48583.5, 242124.0] -48.753176 2479
In [32]:
feature_importance = pd.Series(cf.feature_importances_, index=x_cols).sort_values(ascending=False)
feature_importance.head(10)
Out[32]:
tw         0.335125
nifa       0.295516
p401       0.140880
inc        0.136740
age        0.052840
educ       0.016579
fsize      0.010145
pira       0.004316
twoearn    0.003895
db         0.002149
dtype: float64

Build a simple causal-tree-style summary¶

EconML's CausalForestDML is a forest, not a single tree. To get an interpretable tree-style partition, we fit a shallow regression tree to the estimated treatment effects.

In [33]:
tree = DecisionTreeRegressor(max_depth=3, min_samples_leaf=300, random_state=123)
tree.fit(results[x_cols], results['tau_hat'])

tree_rules = export_text(tree, feature_names=x_cols)
print(tree_rules)
|--- tw <= 315009.50
|   |--- hown <= 0.50
|   |   |--- tw <= -4393.50
|   |   |   |--- value: [2287.87]
|   |   |--- tw >  -4393.50
|   |   |   |--- value: [142.52]
|   |--- hown >  0.50
|   |   |--- tw <= 1745.00
|   |   |   |--- value: [-3108.15]
|   |   |--- tw >  1745.00
|   |   |   |--- value: [-405.92]
|--- tw >  315009.50
|   |--- value: [4254.82]

In [34]:
results['tree_leaf'] = tree.apply(results[x_cols])
leaf_summary = (
    results.groupby('tree_leaf')
    .agg(
        n=('tau_hat', 'size'),
        avg_tau_hat=('tau_hat', 'mean'),
        avg_income=('inc', 'mean'),
        avg_age=('age', 'mean'),
        share_homeowner=('hown', 'mean'),
    )
    .sort_values('avg_tau_hat', ascending=False)
)
leaf_summary
Out[34]:
n avg_tau_hat avg_income avg_age share_homeowner
tree_leaf
8 300 4254.820510 77450.101562 50.436667 0.993333
3 300 2287.872734 29732.789062 35.856667 0.000000
4 3315 142.520689 26025.126953 37.339065 0.000000
7 5700 -405.918334 42331.148438 43.153860 1.000000
6 300 -3108.148555 30428.240234 38.226667 1.000000

Interpretation guide¶

  • tau_hat is the estimated treatment effect for each household.
  • The causal forest learns where the 401(k) eligibility effect appears larger or smaller.
  • The shallow tree is only a summary tool for interpretation; the main estimator is still the causal forest.