@@ -20,52 +20,108 @@ def _fit_model(name, model, Y, T, X):
2020
2121class TestRScorer (unittest .TestCase ):
2222
23- def _get_data (self ):
23+ def _get_data (self , discrete_outcome = False ):
2424 X = np .random .normal (0 , 1 , size = (100000 , 2 ))
2525 T = np .random .binomial (1 , .5 , size = (100000 ,))
26- y = X [:, 0 ] * T + np .random .normal (size = (100000 ,))
27- return y , T , X , X [:, 0 ]
26+ if discrete_outcome :
27+ eps = np .random .normal (size = (100000 ,))
28+ log_odds = X [:, 0 ]* T + eps
29+ y_sigmoid = 1 / (1 + np .exp (- log_odds ))
30+ y = np .array ([np .random .binomial (1 , p ) for p in y_sigmoid ])
31+ # Difference in conditional probabilities P(y=1|X,T=1) - P(y=1|X,T=0)
32+ true_eff = (1 / (1 + np .exp (- (X [:, 0 ]+ eps )))) - (1 / (1 + np .exp (- eps )))
33+ else :
34+ y = X [:, 0 ] * T + np .random .normal (size = (100000 ,))
35+ true_eff = X [:, 0 ]
36+
37+ y = y .reshape (- 1 , 1 )
38+ T = T .reshape (- 1 , 1 )
39+ return y , T , X , true_eff
2840
2941 def test_comparison (self ):
42+
3043 def reg ():
3144 return LinearRegression ()
3245
3346 def clf ():
3447 return LogisticRegression ()
3548
36- y , T , X , true_eff = self ._get_data ()
37- (X_train , X_val , T_train , T_val ,
38- Y_train , Y_val , _ , true_eff_val ) = train_test_split (X , T , y , true_eff , test_size = .4 )
39-
40- models = [('ldml' , LinearDML (model_y = reg (), model_t = clf (), discrete_treatment = True , cv = 3 )),
41- ('sldml' , SparseLinearDML (model_y = reg (), model_t = clf (), discrete_treatment = True ,
42- featurizer = PolynomialFeatures (degree = 2 , include_bias = False ), cv = 3 )),
43- ('xlearner' , XLearner (models = reg (), cate_models = reg (), propensity_model = clf ())),
44- ('dalearner' , DomainAdaptationLearner (models = reg (), final_models = reg (), propensity_model = clf ())),
45- ('slearner' , SLearner (overall_model = reg ())),
46- ('tlearner' , TLearner (models = reg ())),
47- ('drlearner' , DRLearner (model_propensity = clf (), model_regression = reg (),
48- model_final = reg (), cv = 3 )),
49- ('rlearner' , NonParamDML (model_y = reg (), model_t = clf (), model_final = reg (),
50- discrete_treatment = True , cv = 3 )),
51- ('dml3dlasso' , DML (model_y = reg (), model_t = clf (), model_final = reg (), discrete_treatment = True ,
52- featurizer = PolynomialFeatures (degree = 3 ), cv = 3 ))
53- ]
54-
55- models = Parallel (n_jobs = 1 , verbose = 1 )(delayed (_fit_model )(name , mdl ,
56- Y_train , T_train , X_train )
57- for name , mdl in models )
58-
59- scorer = RScorer (model_y = reg (), model_t = clf (),
60- discrete_treatment = True , cv = 3 , mc_iters = 2 , mc_agg = 'median' )
61- scorer .fit (Y_val , T_val , X = X_val )
62- rscore = [scorer .score (mdl ) for _ , mdl in models ]
63- rootpehe_score = [np .sqrt (np .mean ((true_eff_val .flatten () - mdl .effect (X_val ).flatten ())** 2 ))
64- for _ , mdl in models ]
65- assert LinearRegression ().fit (np .array (rscore ).reshape (- 1 , 1 ), np .array (rootpehe_score )).coef_ < 0.5
66- mdl , _ = scorer .best_model ([mdl for _ , mdl in models ])
67- rootpehe_best = np .sqrt (np .mean ((true_eff_val .flatten () - mdl .effect (X_val ).flatten ())** 2 ))
68- assert rootpehe_best < 1.5 * np .min (rootpehe_score ) + 0.05
69- mdl , _ = scorer .ensemble ([mdl for _ , mdl in models ])
70- rootpehe_ensemble = np .sqrt (np .mean ((true_eff_val .flatten () - mdl .effect (X_val ).flatten ())** 2 ))
71- assert rootpehe_ensemble < 1.5 * np .min (rootpehe_score ) + 0.05
49+ test_cases = [
50+ {"name" :"continuous_outcome" , "discrete_outcome" : False },
51+ {"name" :"discrete_outcome" , "discrete_outcome" : True }
52+ ]
53+
54+ for case in test_cases :
55+ with self .subTest (case ["name" ]):
56+ discrete_outcome = case ["discrete_outcome" ]
57+
58+ if discrete_outcome :
59+ y , T , X , true_eff = self ._get_data (discrete_outcome = True )
60+
61+ models = [('ldml' , LinearDML (model_y = clf (), model_t = clf (), discrete_treatment = True ,
62+ discrete_outcome = discrete_outcome , cv = 3 )),
63+ ('sldml' , SparseLinearDML (model_y = clf (), model_t = clf (), discrete_treatment = True ,
64+ discrete_outcome = discrete_outcome ,
65+ featurizer = PolynomialFeatures (degree = 2 , include_bias = False ),
66+ cv = 3 )),
67+ ('drlearner' , DRLearner (model_propensity = clf (), model_regression = clf (), model_final = reg (),
68+ discrete_outcome = discrete_outcome , cv = 3 )),
69+ ('rlearner' , NonParamDML (model_y = clf (), model_t = clf (), model_final = reg (),
70+ discrete_treatment = True , discrete_outcome = discrete_outcome , cv = 3 )),
71+ ('dml3dlasso' , DML (model_y = clf (), model_t = clf (), model_final = reg (), discrete_treatment = True ,
72+ discrete_outcome = discrete_outcome ,
73+ featurizer = PolynomialFeatures (degree = 3 ), cv = 3 )),
74+ # SLearner as baseline for rootpehe score - not enough variation in rscore w/ above models
75+ ('slearner' , SLearner (overall_model = reg ())),
76+ ]
77+
78+ else :
79+ y , T , X , true_eff = self ._get_data ()
80+
81+ models = [('ldml' , LinearDML (model_y = reg (), model_t = clf (), discrete_treatment = True , cv = 3 )),
82+ ('sldml' , SparseLinearDML (model_y = reg (), model_t = clf (), discrete_treatment = True ,
83+ featurizer = PolynomialFeatures (degree = 2 , include_bias = False ),
84+ cv = 3 )),
85+ ('xlearner' , XLearner (models = reg (), cate_models = reg (), propensity_model = clf ())),
86+ ('dalearner' , DomainAdaptationLearner (models = reg (), final_models = reg (),
87+ propensity_model = clf ())),
88+ ('slearner' , SLearner (overall_model = reg ())),
89+ ('tlearner' , TLearner (models = reg ())),
90+ ('drlearner' , DRLearner (model_propensity = clf (), model_regression = reg (),
91+ model_final = reg (), cv = 3 )),
92+ ('rlearner' , NonParamDML (model_y = reg (), model_t = clf (), model_final = reg (),
93+ discrete_treatment = True , cv = 3 )),
94+ ('dml3dlasso' , DML (model_y = reg (), model_t = clf (), model_final = reg (),
95+ discrete_treatment = True , featurizer = PolynomialFeatures (degree = 3 ), cv = 3 ))
96+ ]
97+
98+ (X_train , X_val , T_train , T_val ,
99+ Y_train , Y_val , _ , true_eff_val ) = train_test_split (X , T , y , true_eff , test_size = .4 )
100+
101+ models = Parallel (n_jobs = 1 , verbose = 1 )(delayed (_fit_model )(name , mdl ,
102+ Y_train , T_train , X_train )
103+ for name , mdl in models )
104+
105+ if discrete_outcome :
106+ scorer = RScorer (model_y = clf (), model_t = clf (),
107+ discrete_treatment = True , discrete_outcome = discrete_outcome ,
108+ cv = 3 , mc_iters = 2 , mc_agg = 'median' )
109+ else :
110+ scorer = RScorer (model_y = reg (), model_t = clf (),
111+ discrete_treatment = True , cv = 3 ,
112+ mc_iters = 2 , mc_agg = 'median' )
113+
114+ scorer .fit (Y_val , T_val , X = X_val )
115+ rscore = [scorer .score (mdl ) for _ , mdl in models ]
116+ rootpehe_score = [np .sqrt (np .mean ((true_eff_val .flatten () - mdl .effect (X_val ).flatten ())** 2 ))
117+ for _ , mdl in models ]
118+ # Checking neg corr between rscore and rootpehe (precision in estimating heterogeneous effects)
119+ assert LinearRegression ().fit (np .array (rscore ).reshape (- 1 , 1 ), np .array (rootpehe_score )).coef_ < 0.5
120+ mdl , _ = scorer .best_model ([mdl for _ , mdl in models ])
121+ rootpehe_best = np .sqrt (np .mean ((true_eff_val .flatten () - mdl .effect (X_val ).flatten ())** 2 ))
122+ # Checking best model selection behaves as intended
123+ assert rootpehe_best < 1.5 * np .min (rootpehe_score ) + 0.05
124+ mdl , _ = scorer .ensemble ([mdl for _ , mdl in models ])
125+ rootpehe_ensemble = np .sqrt (np .mean ((true_eff_val .flatten () - mdl .effect (X_val ).flatten ())** 2 ))
126+ # Checking cate ensembling behaves as intended
127+ assert rootpehe_ensemble < 1.5 * np .min (rootpehe_score ) + 0.05
0 commit comments