# code for loading the format for the notebook
import os
# path : store the current path to convert back to it later
path = os.getcwd()
os.chdir(os.path.join('..', '..', 'notebook_format'))
from formats import load_style
load_style(plot_style=False)
os.chdir(path)
# 1. magic for inline plot
# 2. magic to print version
# 3. magic so that the notebook will reload external python modules
# 4. magic to enable retina (high resolution) plots
# https://gist.github.com/minrk/3301035
%matplotlib inline
%load_ext watermark
%load_ext autoreload
%autoreload 2
%config InlineBackend.figure_format='retina'
import os
import time
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import statsmodels.formula.api as smf
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import OneHotEncoder
from sklearn.linear_model import LogisticRegression, LinearRegression
from sklearn.model_selection import train_test_split
# prevent scientific notations
pd.set_option('display.float_format', lambda x: '%.3f' % x)
plt.style.use('fivethirtyeight')
%watermark -a 'Ethen' -d -t -v -p numpy,pandas,sklearn,matplotlib,seaborn,statsmodels
The purpose of using propensity score is to balance the treatment/control groups on the observed covariates. The advantage with using weighting is that all individuals/records in our dataset can be leveraged. Unlike techniques such as matching, where we might be discarding large amount of data in the control group.
In this document, we'll be:
We'll be working with data coming from the The National Study of Learning Mindsets. The background behind this data is an attempt to find out if instilling students with a growth mindset will improve their overall academic performance. In this dataset, the academic performance is recorded as a standardized achievement_score
.
Besides the treated and outcome variables, the study also recorded some other features:
# we'll only work with a subset of the columns, feel free to experiment with others
cat_cols = ['ethnicity', 'gender', 'school_urbanicity']
num_cols = ['school_mindset', 'school_achievement', 'school_ethnic_minority', 'school_poverty', 'school_size']
treatment_col = 'intervention'
label_col = 'achievement_score'
use_cols = [label_col, treatment_col] + num_cols + cat_cols
df = pd.read_csv('data/learning_mindset.csv', usecols=use_cols)[use_cols]
print(df.shape)
df.head()
Our data will be fed into a logistic regression in the next section, here we one hot encode the categorical variables. As the numeric features are already standardized, we leave them as is.
one_hot_encoder = OneHotEncoder(handle_unknown='ignore', sparse=False, dtype=np.int32)
column_transformer = ColumnTransformer(
transformers = [
('one_hot', one_hot_encoder, cat_cols)
],
sparse_threshold=0,
remainder='passthrough'
)
column_transformer.fit(df)
one_hot_encoder = column_transformer.named_transformers_['one_hot']
one_hot_encoded_cols = one_hot_encoder.get_feature_names(cat_cols).tolist()
columns = one_hot_encoded_cols + [label_col, treatment_col] + num_cols
df = pd.DataFrame(column_transformer.transform(df), columns=columns)
print(df.shape)
df
It's often times useful to establish some baseline. Here, we would like to gauge what would the result look like if we don't use propensity score weighting to control for potential biases with assignment of individuals to the control and treatment group.
# change default style figure and font size
plt.rcParams['figure.figsize'] = 8, 6
plt.rcParams['font.size'] = 12
# we can check the histogram of our label between the treatment and control
plt.hist(df.loc[df[treatment_col] == 0.0, label_col], bins=20, alpha=0.3, label='control')
plt.hist(df.loc[df[treatment_col] == 1.0, label_col], bins=20, alpha=0.3, label='treatment')
plt.legend()
plt.show()
# fitting a linear regression to estimate the outcome
linear = LinearRegression()
linear.fit(df[[treatment_col]], df[label_col])
print(linear.intercept_, linear.coef_)
smf.ols(f'{label_col} ~ {treatment_col}', data=df).fit().summary().tables[1]
We can use Linear Regression from different packages to establish our baseline estimates. The one from statsmodels
will give us some additional statistical information. By blindly comparing individuals with and without the intervention, we can see that, on average, those in the treatment group achieved a achievement score 0.4723 higher than the control. Be aware that in this dataset is score was standardized, i.e. it means the treated is 0.4723 standard deviation higher than the untreated.
Upon establishing the baseline, our next task is to question these numbers.
The idea behind propensity score is we don't need to directly control for our confounders $X$ to achieve conditional independence $(Y^1,Y^0) \perp T \text{ | } X$. Hence, $Y^1$ denotes the outcome if treatment, $T$, was applied, whereas $Y^0$ measures the outcome if individual was under the control group.
Instead, it is sufficient to control for a single variable, propensity score, $P(x)$, which is the conditional probability of the treatment, $P(T|X)$. Or in notations form, with our propensity score, we now have $(Y^1,Y^0) \perp T \text{ | } P(X)$.
Here, we'll use a logistic regression to estimate our propensity score. Feel free to use other classification techniques, but keep in mind that other classification techniques might not produce well calibrated probabilities and the utmost goal of propensity score estimation is to make sure to include all the confounding variables instead of getting taken away of the different kinds of classification model that we can potentially use.
input_cols = one_hot_encoded_cols + num_cols
logistic = LogisticRegression()
logistic.fit(df[input_cols], df[treatment_col])
propensity_score = 'propensity_score'
df[propensity_score] = logistic.predict_proba(df[input_cols])[:, 1]
print(df.shape)
df.head()
After training our propensity score, it's important to check for score overlap between the treated and untreated population. Looking at the distribution plot below, we can find both treated and untreated individuals in different regions. In this case, we have a balanced dataset and our positivity assumption holds true. Remember positivity refers to the scenario where all of the individuals have at least some chance of receiving either treatment and that appears to be the case here. In summary, this would be a situation where we would feel comfortable to proceed with our propensity score matching.
Note that if we encounter a situation where the propensity score between the treatment and control does not overlap, then we should stop and better construct the prediction that controls for confounding, or in much simpler terms, see if we are missing any important features in our propensity score model. This requires domain knowledge of the data that we are working with. Failing to check for positivity and lack of balance can introduce bias in our estimation as we will be extrapolating the effect to unknown regions.
control_score = df.loc[df[treatment_col] == 0.0, propensity_score]
treatment_score = df.loc[df[treatment_col] == 1.0, propensity_score]
sns.distplot(control_score, label='control')
sns.distplot(treatment_score, label='treatment')
plt.title('Propensity Score Distribution of Control vs Treatment')
plt.ylabel('Density')
plt.xlabel('Scores')
plt.legend()
plt.tight_layout()
plt.show()
The final step in our analysis is to run our outcome model using the propensity score as weights, i.e. fit a weighted regression. In order to use our propensity score as weights, we will need to apply some transformation known as Inverse Propensity Weighting (IPW). For individuals in the treatment group, $w = \frac{1}{P(x)}$, whereas for individuals in the control group, $w = \frac{1}{1 - P(x)}$.
To understand why applying inverse propensity weighting helps us de-bias our potentially biased data. Say in the data we collected there are 200 records from group A and 2000 records from group B, to "balance" our dataset, we would like to "up-weight" the records from group A and "down-weight" the records from group B. If we were to weight each records in both groups by number of inverse records, $1 / 200$ for group A, and $1 / 2000$ for group B, we would end up with effectively 1 record on both group.
Coming back to our example, we are taking all the individuals that are in the treatment group and scaling them with the inverse propensity of being treated $w = \frac{1}{P(x)}$. What this effectively does is it makes those with a very low probability of treatment have a high weight, in other words, we have a individual in the treatment group that looks like he/she should belong to the control group, hence we will give a higher weight to this individual. What this does is create a population with the same size as the original, but where everyone is treated. We can apply the same reasoning for the control group.
treatment_weight = 1.0 / treatment_score
control_weight = 1.0 / (1.0 - control_score)
print('Original Sample Size', df.shape[0])
print('Treated Population Sample Size', treatment_weight.sum())
print('Untreated Population Sample Size', control_weight.sum())
sample_weight = 'sample_weight'
df[sample_weight] = np.where(
df[treatment_col] == 1.0,
1.0 / df[propensity_score],
1.0 / (1.0 - df[propensity_score])
)
Once the sample weight are created, we can re-estimate the outcome with a weighted Linear Regression.
smf.wls(f'{label_col} ~ {treatment_col}', data=df, weights=df[sample_weight]).fit().summary().tables[1]
linear = LinearRegression()
linear.fit(df[[treatment_col]], df[label_col], sample_weight=df[sample_weight])
print(linear.intercept_, linear.coef_)
FYI, even though scikit-learn's LinearRegression by default doesn't give us an estimated standard error, we can estimate this using bootstrapping. i.e. by sampling with replacement from the original data, and computing the average treatment effect like before. After repeating this step for lots of times, we will get a distribution of the outcome estimation. We will also use this time to organize the overall workflow into one single code cell.
def run_propensity_score_estimation(df, input_cols, treatment_col, label_col):
# df is our pre-processed data
df = df.sample(frac=1, replace=True)
# estimate the propensity score
logistic = LogisticRegression()
logistic.fit(df[input_cols], df[treatment_col])
propensity_score = logistic.predict_proba(df[input_cols])[:, 1]
# calculate the inverse propensity weight
sample_weight = np.where(
df[treatment_col] == 1.0,
1.0 / propensity_score,
1.0 / (1.0 - propensity_score)
)
# estimate the outcome using weighted regression
linear = LinearRegression()
linear.fit(df[[treatment_col]], df[label_col], sample_weight=sample_weight)
return linear.coef_[0]
from joblib import Parallel, delayed
np.random.seed(88)
# the bootstrap approach of computing standard error can be computationally expensive on large datasets.
bootstrap_sample = 1000
parallel = Parallel(n_jobs=4)
ates = parallel(delayed(run_propensity_score_estimation)(df,
input_cols,
treatment_col,
label_col)
for _ in range(bootstrap_sample))
ates = np.array(ates)
print(f"ATE: {ates.mean()}")
print(f"95% C.I.: {(np.percentile(ates, 2.5), np.percentile(ates, 97.5))}")
In this article, we look a stab at calculating average treatment effect, $E[Y|T=1] - E[Y|T=0]$, using propensity scores. We saw how using inverse propensity weighting helps us create an un-biased data from a biased data. With this method, it's important to call out identifying the features to use for the propensity score model should be treated with care. In general: