Propensity score

%load_ext autoreload
%autoreload 2

import logging

import pandas as pd
import statsmodels.formula.api as smf
from IPython.display import display

import helpers.hnotebook as hnotebo
import helpers.hpandas_display as hpandisp
import helpers.hpandas_stats as hpanstat
import helpers.htutorial as ut
import L08_04_05_propensity_score_utils as mtl0psu

ut.config_notebook()

# Initialize logger.
logging.basicConfig(level=logging.INFO)
_LOG = logging.getLogger(__name__)
# Note: the original code called a non-existent `set_all_loggers_to_print()`
# top-level function; `init_loggers(..., set_all_loggers_to_print=True)` is
# the actual public API for what this was trying to do.
hnotebo.init_loggers(_LOG, set_all_loggers_to_print=True)

dir_name = "L08_data"
!ls $dir_name

out_dir_name = "figures/"

df = pd.read_csv(f"{dir_name}/management_training.csv")
hpandisp.display_df(df)

show_distributions = True
show_correlations = True

hpanstat.explore_dataframe(
    df,
    show_distributions=show_distributions,
    show_correlations=show_correlations,
)

model = smf.ols("engagement_score ~ intervention", data=df).fit()
print("ATE:", model.params["intervention"])
print("95% CI:", model.conf_int().loc["intervention", :].values.T)

display(model.summary().tables[1])

mtl0psu.plot_engagement_vs_intervention(df)

mtl0psu.plot_engagement_vs_intervention_by_department(df)

mtl0psu.plot_all_variables_density_by_intervention(df)

# To reduce this bias, you can adjust for the covariates you have in your data.
model_adj = smf.ols(
    """
    engagement_score ~ intervention
        + tenure + last_engagement_score + department_score
        + n_of_reports + C(gender) + C(role)""",
    data=df,
).fit()

print("ATE (adjusted):", model_adj.params["intervention"])
print("95% CI (adjusted):", model_adj.conf_int().loc["intervention", :].values.T)
print("ATE (naive):", model.params["intervention"])
print("95% CI (naive):", model.conf_int().loc["intervention", :].values.T)

ps_model = smf.logit(
    """
    intervention ~
        tenure + last_engagement_score + department_score
        + C(n_of_reports) + C(gender) + C(role)""",
    data=df,
).fit(disp=0)

data_ps = df.copy()
data_ps["propensity_score"] = ps_model.predict(df)

display(data_ps[["intervention", "engagement_score", "propensity_score"]].head())

# Estimate using propensity score as confounder / covariate.
model_ps = smf.ols(
    "engagement_score ~ intervention + propensity_score",
    data=data_ps,
).fit()
print("ATE (propensity-score covariate):", model_ps.params["intervention"])

# Perform 1-nearest neighbor propensity score matching.
predicted = mtl0psu.propensity_score_matching(data_ps)
display(predicted.head())

# Calculate average treatment effect from propensity score matching.
hat_ATE = mtl0psu.calculate_psm_ate(predicted)
print(f"ATE (propensity score matching): {hat_ATE:.4f}")

# Plot inverse probability of treatment weighting results.
mtl0psu.plot_iptw(data_ps)

# Estimate ATE using IPTW.
weighted_e_y1, weighted_e_y0, hat_ATE = mtl0psu.estimate_ate_iptw(data_ps)

print("E[Y1]:", weighted_e_y1)
print("E[Y0]:", weighted_e_y0)
print("ATE:", hat_ATE)

# Prepare formula and variables for IPW estimation.
formula = """
tenure + last_engagement_score + department_score
    + C(n_of_reports) + C(gender) + C(role)
"""
treatment_col = "intervention"
outcome_col = "engagement_score"

# Estimate ATE using IPW estimator.
ate_ipw = mtl0psu.estimate_ate_with_ps(
    df, formula, treatment_col=treatment_col, outcome_col=outcome_col
)
print(f"ATE (IPW): {ate_ipw:.4f}")

# Compute bootstrap 95% confidence interval for ATE using IPW.
est_fn = lambda data: mtl0psu.estimate_ate_with_ps(
    data,
    ps_formula=formula,
    treatment_col=treatment_col,
    outcome_col=outcome_col,
)

ci = mtl0psu.estimate_confidence_interval_bootstrap(
    df, est_fn, rounds=200, seed=123, n_jobs=4, pcts=[2.5, 97.5]
)
print(f"ATE: {ate_ipw:.4f}")
print(f"95% confidence interval: {ci}")

# Show sample sizes for original and pseudo-population.
print("Original sample size:", data_ps.shape[0])

# Compute sample sizes after IPTW weighting.
treated = data_ps.query("intervention==1")
control = data_ps.query("intervention==0")

weight_t = 1 / treated["propensity_score"]
weight_nt = 1 / (1 - control["propensity_score"])

print("Treated pseudo-population sample size:", sum(weight_t))
print("Untreated pseudo-population sample size:", sum(weight_nt))

# Estimate ATE using stabilized propensity weights.
ate_stabilized = mtl0psu.estimate_ate_stabilized_weights(data_ps)
print(f"ATE (stabilized weights): {ate_stabilized:.4f}")

# Plot propensity score distributions before and after weighting.
mtl0psu.plot_propensity_distributions(data_ps)