Propensity score & Conditional Independence Assumption

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
import numpy as np
import pandas as pd
from plotly.subplots import make_subplots
import plotly.graph_objects as go

np.random.seed(42)

#-------------
# DGP
#-------------

def dgp(effect=[2, 8, -8]):
    N = 10000
    x1 = np.random.uniform(0, 1, N)Heterogeneityの存在
    x2 = np.random.uniform(0, 1, N)

    z1 = np.where(x1 > 0.5, 1, 0)
    z2 = np.where(x2 > 0.5, 1, 0)
    assignment_rule = z1 + z2
    assignment_prob = np.where(assignment_rule == 1, 0.3, 0.5)
    assigment = np.random.binomial(1, assignment_prob)

    potential_untreated = np.random.normal(0, 1, N)
    uplift = np.where(assignment_rule == 1, effect[0], 0.0)\
             + np.where(assignment_rule == 2, effect[1], 0.0) \
             + np.where(assignment_rule == 0, effect[2], 0.0)
    potential_treated = potential_untreated + uplift

    df = pd.DataFrame({
        'x1': x1,
        'x2': x2,
        'prob': assignment_prob,Heterogeneityの存在
        'D': assigment,
        'uplift':uplift, 
        'y0': potential_untreated,
        'y1': potential_treated
    })
    df['observed'] = (1 - df['D']) * df['y0'] + (df['D']) * df['y1']

    return df

df = dgp()
df = dgp()
df.head()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
fig = make_subplots(rows=1, cols=2, 
                    subplot_titles=("Assignment Prob distribution", 
                                    "uplift distribution"))
trace1 = go.Scatter(x=df['x1'],
                    y=df['x2'],
                    mode='markers',
                    name='prob',
                    hovertemplate="%{marker.color}%",
                    marker=dict(color=df['prob'], colorscale='bluered'),
                    showlegend=False
                   )
trace2 = go.Scatter(x=df['x1'],
                    y=df['x2'],
                    mode='markers',
                    name='uplift',
                    hovertemplate="%{marker.color}",
                    marker=dict(color=df['uplift'],colorscale='bluered'),
                    showlegend=False
                   )
fig.append_trace(trace1,1,1)
fig.append_trace(trace2,1,2)
fig.update_layout(title_text="Assignment probability and uplift distirbution conditional on X")

1
2
3
4
5
6
7
8
9
10
11
12
#--------------------
# True ATE, ATT, ATU
#--------------------
def get_true_effect(df):
    stats_df = df.groupby('D')[['y0', 'y1']].mean().reset_index(drop=True)
    stats_df = pd.concat([stats_df, 
                          pd.DataFrame(df.loc[:, ['y0', 'y1']].mean(axis=0)).T], 
                          axis=0).set_index(pd.Index(['ATU', 'ATT', 'ATE']))
    stats_df['effect'] = stats_df['y1'] - stats_df['y0']
    return stats_df

get_true_effect(df)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
def estimate_effect(df):
    y0_treated_weight = ((1 - df['D']) * df['prob']) / (1 - df['prob'])
    y1_untreated_weight = (df['D']) * (1 - df['prob']) / df['prob']
    y1_ate_weight = df['D'] / df['prob']
    y0_ate_weight = (1 - df['D'])/(1 - df['prob'])

    y0_untreated, y1_treated = df.groupby('D')['observed'].mean().values
    y0_treated = np.sum(y0_treated_weight * df['observed']) / np.sum(y0_treated_weight)
    y1_untreated = np.sum(y1_untreated_weight * df['observed']) / np.sum(y1_untreated_weight)
    
    y1_ate = np.sum(y1_ate_weight * df['observed'])/ np.sum(y1_ate_weight)
    y0_ate = np.sum(y0_ate_weight * df['observed'])/ np.sum(y0_ate_weight)

    res_df = pd.DataFrame({'y0':[y0_untreated, y0_treated, y0_ate],
                           'y1':[y1_untreated, y1_treated, y1_ate]})
    res_df = res_df.set_index(pd.Index(['ATU', 'ATT', 'ATE']))
    res_df['effect'] = res_df['y1'] - res_df['y0']

    return res_dfHeterogeneityの存在


estimate_effect(df)

1
2
3
4
5
6
7
8
9
10
res_true = []
res_estimates = []

for i in range(1000):
    tmp = dgp([2, 8, -8])
    res_true.append(get_true_effect(tmp).loc['ATE', 'effect'])
    res_estimates.append(estimate_effect(tmp).loc['ATE', 'effect'])

compare_df = pd.DataFrame({'true':res_true, 'estimate':res_estimates})
compare_df.describe()

1
2
3
hetero_df = df.groupby('prob')[['y0', 'y1']].mean()
hetero_df['effect'] = hetero_df['y1'] - hetero_df['y0']
hetero_df