Comparing ATE Estimators

%matplotlib inline
%config InlineBackend.figure_format = 'retina'

from src.utils import *
from src.dgp import dgp_compare

df = dgp_compare().generate_data(include_beta=True)
df.head()

np.mean(df['beta'])

2.0

Simple difference

smf.ols('outcome ~ treated', data=df).fit().summary().tables[1]

Regression with control

smf.ols('outcome ~ treated + male + age + income', data=df).fit().summary().tables[1]

Matching

from causalml.match import NearestNeighborMatch, create_table_one

X = ['male', 'age', 'income']
psm = NearestNeighborMatch(replace=True, ratio=1, random_state=42)
df_matched = psm.match(data=df, 
                       treatment_col="treated",
                       score_cols=X)
smf.ols('outcome ~ treated', data=df_matched).fit().summary().tables[1]

IPW

df["pscore"] = smf.logit("np.rint(treated) ~ male + age + income", data=df).fit(disp=False).predict()
w = 1 / (df["pscore"] * df["treated"] + (1-df["pscore"]) * (1-df["treated"]))
smf.wls("outcome ~ treated", weights=w, data=df).fit().summary().tables[1]

R Learner

from causalml.inference.meta import BaseRLearner
from lightgbm import LGBMRegressor

BaseRLearner(learner=LGBMRegressor()).estimate_ate(X=df[X], treatment=df['treated'], y=df['outcome'])

(array([1.6529302]), array([1.65076231]), array([1.65509809]))

S Learner

from causalml.inference.meta import BaseSLearner

BaseSLearner(learner=LGBMRegressor()).estimate_ate(X=df[X], treatment=df['treated'], y=df['outcome'])

array([2.08090912])

T Learner

from causalml.inference.meta import BaseTLearner

BaseTLearner(learner=LGBMRegressor()).estimate_ate(X=df[X], treatment=df['treated'], y=df['outcome'])

(array([2.05252873]), array([1.86063479]), array([2.24442267]))

X Learner

from causalml.inference.meta import BaseXLearner

BaseXLearner(learner=LGBMRegressor()).estimate_ate(X=df[X], treatment=df['treated'], y=df['outcome'])

(array([2.0255193]), array([1.83456903]), array([2.21646956]))

DR Learner

from causalml.inference.meta import BaseDRLearner

BaseDRLearner(learner=LGBMRegressor()).estimate_ate(X=df[X], treatment=df['treated'], y=df['outcome'])

(array([2.08923293]), array([1.8921804]), array([2.28628547]))

Compare All

def simulate(dgp, K=100):
    
    # Initialize coefficients
    results = pd.DataFrame(columns=['k', 'Estimator', 'Estimate'])
    names = ['1. Diff ', '2. Reg  ', '3. Match', '4. IPW  ', 
             '5. R lrn', '6. S lrn', '7. T lrn', '8. X lrn', '9. DRlrn']
    
    # Compute coefficients
    for k in range(K):
        print(f"Simulation {k}/{K}", end="\r")
        temp = pd.DataFrame({'k': [k] * len(names), 
                             'Estimator': names, 
                             'Estimate': [0] * len(names)})
        
        # Draw data
        df = dgp.generate_data(seed=k)

        # Single diff
        temp['Estimate'][0] = smf.ols('outcome ~ treated', data=df).fit().params[1]
        
        # Regression with controls
        temp['Estimate'][1] = smf.ols(f'outcome ~ treated + male + age + income', data=df).fit().params[1]
        
        # Matching
        psm = NearestNeighborMatch(replace=True, ratio=1)
        df_matched = psm.match(data=df, treatment_col="treated", score_cols=X)
        temp['Estimate'][2] = smf.ols('outcome ~ treated', data=df_matched).fit().params[1]
        
        # IPW
        df["pscore"] = smf.logit("np.rint(treated) ~ male + age + income", data=df).fit(disp=False).predict()
        w = 1 / (df["pscore"] * df["treated"] + (1-df["pscore"]) * (1-df["treated"]))
        temp['Estimate'][3] = smf.wls("outcome ~ treated", weights=w, data=df).fit().params[1]
                
        # R Learner
        temp['Estimate'][4] = BaseRLearner(learner=LGBMRegressor()).estimate_ate(X=df[X], treatment=df['treated'], y=df['outcome'])[0][0]
        
        # S Learner
        temp['Estimate'][5] = BaseSLearner(learner=LGBMRegressor()).estimate_ate(X=df[X], treatment=df['treated'], y=df['outcome'])[0]
        
        # T Learner
        temp['Estimate'][6] = BaseTLearner(learner=LGBMRegressor()).estimate_ate(X=df[X], treatment=df['treated'], y=df['outcome'])[0][0]
        
        # X Learner
        temp['Estimate'][7] = BaseXLearner(learner=LGBMRegressor()).estimate_ate(X=df[X], treatment=df['treated'], y=df['outcome'])[0][0]
        
        # DR Learner
        temp['Estimate'][8] = BaseDRLearner(learner=LGBMRegressor()).estimate_ate(X=df[X], treatment=df['treated'], y=df['outcome'])[0][0]
        
        # Combine estimates
        results = pd.concat((results, temp))
    
    return results.reset_index(drop=True)

results = simulate(dgp=dgp_compare())

Simulation 99/100

Let’s plot the distribution of the estimated parameters.

p = sns.kdeplot(data=results, x="Estimate", hue="Estimator", legend=True);
sns.move_legend(p, "upper left", bbox_to_anchor=(1.05, 0.8))
plt.axvline(x=2, c='k', ls='--');
plt.title('Simulated Distributions');

We can also tabulate the simulated mean and standard deviation of each estimator.

results.groupby('Estimator').agg(mean=("Estimate", "mean"), std=("Estimate", "std"))

The only unbiased estimators seems to be the S-learner and the X-learner, followed by T-learner, DR-learner and linear regression.

he S-learner however is more efficient.