성향 점수 : 실무로 통하는 인과추론 학습기(3)

import warnings
warnings.filterwarnings('ignore')

import pandas as pd
import numpy as np
import graphviz as gr
from matplotlib import style
import seaborn as sns
from matplotlib import pyplot as plt
style.use("ggplot")
import statsmodels.formula.api as smf
import matplotlib
from cycler import cycler

color=['0.2', '0.6', '1.0']
default_cycler = (cycler(color=color))
linestyle=['-', '--', ':', '-.']
marker=['o', 'v', 'd', 'p']


plt.rc('axes', prop_cycle=default_cycler)

df = pd.read_csv('data.csv')

df.describe()

#처치변수만을 넣은 회귀식
import statsmodels.formula.api as smf

formula = "Infant_Mortality_Rate ~ Intervention"

smf.ols(formula,
        data=df).fit().summary().tables[1]

#그외 교란변수를 모두 넣은 회귀식
model = smf.ols("Infant_Mortality_Rate ~ Intervention + Poverty_Level + Doctor_Ratio", data=df).fit()
model.summary().tables[1]

print("ATE:", model.params['Intervention'])
print("95% CI:", model.conf_int().loc["Intervention", :].values.T)

#성향점수 e(x) 추정
ps_model = smf.logit("Intervention ~ Poverty_Level + Doctor_Ratio", data=df).fit(disp=0)

data_ps = df.assign(
    propensity_score = ps_model.predict(df),
)
data_ps[["Intervention", "Infant_Mortality_Rate", "propensity_score"]].head()

sns.histplot(data=data_ps, x='propensity_score', hue='Intervention', element='step', stat='density', common_norm=False)
plt.title('성향점수의 분포에 따른 처치변수 개입 여부')
plt.xlabel('성향점수')
plt.ylabel('밀도')
plt.show()

model = smf.ols("Infant_Mortality_Rate ~ Intervention + propensity_score",
                data=data_ps).fit()

print("ATE:", model.params['Intervention'])

from sklearn.neighbors import KNeighborsRegressor

T = "Intervention"
X = "propensity_score"
Y = "Infant_Mortality_Rate"

treated = data_ps.query(f"{T}==1")
untreated = data_ps.query(f"{T}==0")

mt0 = KNeighborsRegressor(n_neighbors=1).fit(untreated[[X]],
                                             untreated[Y])

mt1 = KNeighborsRegressor(n_neighbors=1).fit(treated[[X]], treated[Y])

predicted = pd.concat([
    # find matches for the treated looking at the untreated knn model
    treated.assign(match=mt0.predict(treated[[X]])),

    # find matches for the untreated looking at the treated knn model
    untreated.assign(match=mt1.predict(untreated[[X]]))
])

predicted.head()

#성향점수 매칭 결과로 얻는 ATE
match_ATE = np.mean((predicted[Y] - predicted["match"])*predicted[T] 
                    + (predicted["match"] - predicted[Y])*(1-predicted[T]))

print("ATE:", match_ATE)

g_data = (data_ps
          .assign(
              weight = data_ps["Intervention"]/data_ps["propensity_score"] + (1-data_ps["Intervention"])/(1-data_ps["propensity_score"]),
              propensity_score=data_ps["propensity_score"].round(2)
          )
          .groupby(["propensity_score", "Intervention"])
          [["weight", "Infant_Mortality_Rate"]]
          .mean()
          .reset_index())

plt.figure(figsize=(10,4))
for t in [0, 1]:
    sns.scatterplot(data=g_data.query(f"Intervention=={t}"), y="Infant_Mortality_Rate", x="propensity_score", size="weight",
                    sizes=(1,100), color=color[t], legend=None, label=f"T={t}", marker=marker[t])

plt.title("Inverse Probability of Treatment Weighting")
plt.legend()

weight_t = 1/data_ps.query("Intervention==1")["propensity_score"]
weight_nt = 1/(1-data_ps.query("Intervention==0")["propensity_score"])
t1 = data_ps.query("Intervention==1")["Infant_Mortality_Rate"]
t0 = data_ps.query("Intervention==0")["Infant_Mortality_Rate"]

y1 = sum(t1*weight_t)/len(data_ps)
y0 = sum(t0*weight_nt)/len(data_ps)

print("E[Y1]:", y1)
print("E[Y0]:", y0)
print("ATE", y1 - y0)

#역확률 가중치로 얻는 ATE
IPW_ATE = np.mean(data_ps["Infant_Mortality_Rate"] 
        * (data_ps["Intervention"] - data_ps["propensity_score"]) 
        / (data_ps["propensity_score"]*(1-data_ps["propensity_score"])))

print("ATE:",IPW_ATE)

from sklearn.linear_model import LogisticRegression
from patsy import dmatrix

# define function that computes the IPW estimator
def est_ate_with_ps(df, ps_formula, T, Y):

    X = dmatrix(ps_formula, df)
    ps_model = LogisticRegression(penalty="none",
                                  max_iter=1000).fit(X, df[T])
    ps = ps_model.predict_proba(X)[:, 1]

    # compute the ATE
    return np.mean((df[T]-ps) / (ps*(1-ps)) * df[Y])

formula = """Poverty_Level + Doctor_Ratio"""
T = "Intervention"
Y = "Infant_Mortality_Rate"

est_ate_with_ps(df, formula, T, Y)

>>> - 0.3933004355383178

from joblib import Parallel, delayed # for parallel processing

def bootstrap(data, est_fn, rounds=200, seed=123, pcts=[2.5, 97.5]):
    np.random.seed(seed)

    stats = Parallel(n_jobs=4)(
        delayed(est_fn)(data.sample(frac=1, replace=True))
        for _ in range(rounds)
    )

    return np.percentile(stats, pcts)

from toolz import partial

print(f"ATE: {est_ate_with_ps(df, formula, T, Y)}")

est_fn = partial(est_ate_with_ps, ps_formula=formula, T=T, Y=Y)
print(f"95% C.I.: ", bootstrap(df, est_fn))

print("Original Sample Size", data_ps.shape[0])
print("Treated Pseudo-Population Sample Size", sum(weight_t))
print("Untreated Pseudo-Population Sample Size", sum(weight_nt))

>>> Original Sample Size 100
>>> Treated Pseudo-Population Sample Size 99.96699061739257
>>> Untreated Pseudo-Population Sample Size 100.02045315052773

p_of_t = data_ps["Intervention"].mean()

t1 = data_ps.query("Intervention==1")
t0 = data_ps.query("Intervention==0")

weight_t_stable = p_of_t/t1["propensity_score"]
weight_nt_stable = (1-p_of_t)/(1-t0["propensity_score"])

print("Treat size:", len(t1))
print("W treat", sum(weight_t_stable))

print("Control size:", len(t0))
print("W treat", sum(weight_nt_stable))

nt = len(t1)
nc = len(t0)

y1 = sum(t1["Infant_Mortality_Rate"]*weight_t_stable)/nt
y0 = sum(t0["Infant_Mortality_Rate"]*weight_nt_stable)/nc

print("ATE: ", y1 - y0)

# 그래프 설정
plt.figure(figsize=(14, 7))

# 히스토그램 및 KDE 설정
sns.histplot(data=df, x='Infant_Mortality_Rate', hue='Intervention', kde=True, element='step', stat="density", common_norm=False, palette='viridis', bins=10)

# 그래프 제목과 레이블 설정
plt.title('Distribution of Infant Mortality Rate by Intervention', fontsize=16)
plt.xlabel('Infant Mortality Rate', fontsize=14)
plt.ylabel('Density', fontsize=14)
plt.legend(title='Intervention', labels=['No Intervention', 'Intervention'])

# 그래프 출력
plt.show()

# df.head()에서 파악한 head의 특정 행의 Infant_Mortality_Rate 임의로 바꾸기 
# 처치를 했을 때(실험군)의 'Infant_Mortality_Rate'를 바꿔봐도 되고, 처치를 안했을 때(대조군)의 'Infant_Mortality_Rate'를 바꿔봐도 됩니다!
df.at[2, 'Infant_Mortality_Rate'] = 0.98

#처치변수만을 넣은 회귀식
import statsmodels.formula.api as smf

formula = "Infant_Mortality_Rate ~ Intervention"

smf.ols(formula,
        data=df).fit().summary().tables[1]

#그외 교란변수를 모두 넣은 회귀식
model = smf.ols("Infant_Mortality_Rate ~ Intervention + Poverty_Level + Doctor_Ratio", data=df).fit()
model.summary().tables[1]

print("ATE:", model.params['Intervention'])
print("95% CI:", model.conf_int().loc["Intervention", :].values.T) 

ps_model = smf.logit("Intervention ~ Poverty_Level + Doctor_Ratio", *data*=df).fit(*disp*=0)

data_ps = df.assign(
    propensity_score = ps_model.predict(df),
)
data_ps[["Intervention", "Infant_Mortality_Rate", "propensity_score"]].head()

model = smf.ols("Infant_Mortality_Rate ~ Intervention + propensity_score",
                data=data_ps).fit()

print("ATE:", model.params['Intervention'])

from sklearn.neighbors import KNeighborsRegressor

T = "Intervention"
X = "propensity_score"
Y = "Infant_Mortality_Rate"

treated = data_ps.query(f"{T}==1")
untreated = data_ps.query(f"{T}==0")

mt0 = KNeighborsRegressor(n_neighbors=1).fit(untreated[[X]],
                                             untreated[Y])

mt1 = KNeighborsRegressor(n_neighbors=1).fit(treated[[X]], treated[Y])

predicted = pd.concat([
    # find matches for the treated looking at the untreated knn model
    treated.assign(match=mt0.predict(treated[[X]])),

    # find matches for the untreated looking at the treated knn model
    untreated.assign(match=mt1.predict(untreated[[X]]))
])

predicted.head()

#성향점수 매칭 결과로 얻는 ATE
match_ATE = np.mean((predicted[Y] - predicted["match"])*predicted[T] 
                    + (predicted["match"] - predicted[Y])*(1-predicted[T]))

print("ATE:", match_ATE)

g_data = (data_ps
          .assign(
              weight = data_ps["Intervention"]/data_ps["propensity_score"] + (1-data_ps["Intervention"])/(1-data_ps["propensity_score"]),
              propensity_score=data_ps["propensity_score"].round(2)
          )
          .groupby(["propensity_score", "Intervention"])
          [["weight", "Infant_Mortality_Rate"]]
          .mean()
          .reset_index())

plt.figure(figsize=(10,4))
for t in [0, 1]:
    sns.scatterplot(data=g_data.query(f"Intervention=={t}"), y="Infant_Mortality_Rate", x="propensity_score", size="weight",
                    sizes=(1,100), color=color[t], legend=None, label=f"T={t}", marker=marker[t])

plt.title("Inverse Probability of Treatment Weighting")
plt.legend()

weight_t = 1/data_ps.query("Intervention==1")["propensity_score"]
weight_nt = 1/(1-data_ps.query("Intervention==0")["propensity_score"])
t1 = data_ps.query("Intervention==1")["Infant_Mortality_Rate"]
t0 = data_ps.query("Intervention==0")["Infant_Mortality_Rate"]

y1 = sum(t1*weight_t)/len(data_ps)
y0 = sum(t0*weight_nt)/len(data_ps)

print("E[Y1]:", y1)
print("E[Y0]:", y0)
print("ATE", y1 - y0)

#역확률 가중치로 얻는 ATE
IPW_ATE = np.mean(data_ps["Infant_Mortality_Rate"] 
        * (data_ps["Intervention"] - data_ps["propensity_score"]) 
        / (data_ps["propensity_score"]*(1-data_ps["propensity_score"])))

print("ATE:",IPW_ATE)