이질적 처치 효과 : 실무로 통하는 인과추론 학습기(4)

import pandas as pd
import numpy as np
from matplotlib import pyplot as plt
import seaborn as sns

import matplotlib

matplotlib.rcParams.update({'font.size': 18})

data = pd.read_csv("./data_week4.csv")

data

import statsmodels.formula.api as smf

X = ["C(month)", "C(weekday)", "is_holiday"]
regr_cate = smf.ols(f"sales ~ days_until_departure*({'+'.join(X)})",
                     data=data).fit()

ols_cate_pred = (regr_cate.predict(data.assign(days_until_departure=data["days_until_departure"]+1))-regr_cate.predict(data))
ols_cate_pred

train = data.query("day<'2024-01-01'")
test = data.query("day>='2024-01-01'")

import statsmodels.formula.api as smf
X = ["C(month)", "C(weekday)", "is_holiday"]
regr_model = smf.ols(f"sales ~ days_until_departure*({'+'.join(X)})",
                     data=train).fit()

# cate모델 생성
cate_pred = (
    regr_model.predict(test.assign(days_until_departure=test["days_until_departure"] + 1)) - regr_model.predict(test)
)

from sklearn.ensemble import GradientBoostingRegressor
# ml 모델 생성
X = ["month", "weekday", "is_holiday", "days_until_departure"]
y = "sales"

np.random.seed(1)
ml_model = GradientBoostingRegressor(n_estimators=50).fit(train[X],
                                                          train[y])

ml_pred = ml_model.predict(test[X])

# 난수모델 생성
np.random.seed(123)

test_pred = test.assign(
    ml_pred=ml_pred,
    cate_pred=cate_pred,
    rand_m_pred=np.random.uniform(-1, 1, len(test)),
)

test_pred[["flight_id", "day", "sales",
           "ml_pred", "cate_pred", "rand_m_pred"]].head()

from toolz import curry

@curry
def effect(data, y, t):
        return (np.sum((data[t] - data[t].mean())*data[y]) /
                np.sum((data[t] - data[t].mean())**2))

def effect_by_quantile(df, pred, y, t, q=10):

    # 분위수에 대한 파티션 생성
    groups = np.round(pd.IntervalIndex(pd.qcut(df[pred], q=q)).mid, 2)

    return (df
            .assign(**{f"{pred}_quantile": groups})
            .groupby(f"{pred}_quantile")
            .apply(effect(y=y, t=t)))

effect_by_quantile(test_pred, "cate_pred", y="sales", t="days_until_departure")

import warnings
warnings.filterwarnings("ignore")

fig, axs = plt.subplots(1, 3, sharey=True, figsize=(15, 4))
for m, ax in zip(["rand_m_pred", "ml_pred", "cate_pred"], axs):
    effect_by_quantile(test_pred, m, "sales", "days_until_departure").plot.bar(ax=ax)
    ax.set_ylabel("Effect")

np.set_printoptions(*linewidth*=80, *threshold*=10)

def cumulative_effect_curve(dataset, prediction, y, t,
                            ascending=False, steps=100):
    size = len(dataset)
    # prediction 열을 기준으로 오름차순으로 정렬한 다음 인덱스 재설정
    ordered_df = (dataset
                  .sort_values(prediction, ascending=ascending)
                  .reset_index(drop=True))

    # 데이터셋 크기를 기반으로 단계를 계산하여 동일한 간격으로 데이터셋 나누기
    steps = np.linspace(size/steps, size, steps).round(0)

    return np.array([effect(ordered_df.query(f"index<={row}"), t=t, y=y)
                     for row in steps])

# cumulative_effect_curve 분해해보기
size = len(test_pred)
ordered_df = (test_pred.sort_values("cate_pred", ascending=False).reset_index(drop=True))
steps = np.linspace(size/100, size, 100).round(0)
steps

cumulative_effect_curve(test_pred, "cate_pred", "sales", "days_until_departure")

plt.figure(figsize=(10,4))

for m in ["rand_m_pred", "ml_pred", "cate_pred"]:
    cumu_effect = cumulative_effect_curve(test_pred, m, "sales", "days_until_departure", steps=100)
    x = np.array(range(len(cumu_effect)))
    plt.plot(100*(x/x.max()), cumu_effect, label=m)

plt.hlines(effect(test_pred, "sales", "days_until_departure"), 0, 100, linestyles="--", color="black", label="Avg. Effect.")
plt.xlabel("Top %")
plt.ylabel("Cumulative Effect")
plt.title("Cumulative Effect Curve")
plt.legend(fontsize=14)

def cumulative_gain_curve(df, prediction, y, t,
                          ascending=False, normalize=False, steps=100):

    effect_fn = effect(t=t, y=y)
    normalizer = effect_fn(df) if normalize else 0

    size = len(df)
    ordered_df = (df
                  .sort_values(prediction, ascending=ascending)
                  .reset_index(drop=True))

    steps = np.linspace(size/steps, size, steps).round(0)
    effects = [(effect_fn(ordered_df.query(f"index<={row}"))
                -normalizer)*(row/size)
               for row in steps]

    return np.array([0] + effects)

cumulative_gain_curve(test_pred, "cate_pred", "sales", "days_until_departure")

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15,5))

for m in ["rand_m_pred", "ml_pred", "cate_pred"]:
    cumu_gain = cumulative_gain_curve(test_pred, m, "sales", "days_until_departure")
    x = np.array(range(len(cumu_gain)))
    ax1.plot(100*(x/x.max()), cumu_gain, label=m)

ax1.plot([0, 100], [0, effect(test_pred, "sales", "days_until_departure")], linestyle="--", label="Random Model", color="black")

ax1.set_xlabel("Top %")
ax1.set_ylabel("Cumulative Gain")
ax1.set_title("Cumulative Gain Curve")
ax1.legend()

for m in ["rand_m_pred", "ml_pred", "cate_pred"]:
    cumu_gain = cumulative_gain_curve(test_pred, m, "sales", "days_until_departure", normalize=True)
    x = np.array(range(len(cumu_gain)))
    ax2.plot(100*(x/x.max()), cumu_gain, label=m)

ax2.hlines(0, 0, 100, linestyle="--", label="Random Model", color="black")

ax2.set_xlabel("Top %")
ax2.set_ylabel("Cumulative Gain")
ax2.set_title("Cumulative Gain Curve (Normalized)")

for m in ["rand_m_pred", "ml_pred", "cate_pred"]:
    gain = cumulative_gain_curve(test_pred, m,
                                 "sales", "days_until_departure", normalize=True)
    print(f"AUC for {m}:", sum(gain))