Exercise & Mental Health: Addressing Selection Bias

# Propensity score estimation
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler

# Features for propensity model
features = ['age', 'income_cat', 'education', 
            'bmi', 'chronic_conditions', 'smoker',
            'employment', 'marital_status', 
            'health_insurance', 'urban_rural']

# Fit propensity model
ps_model = LogisticRegression(
    penalty='l2', 
    C=1.0,
    max_iter=1000
)
ps_model.fit(X[features], treatment)

# Get propensity scores
propensity = ps_model.predict_proba(X)[:, 1]

# Calculate IPW weights
weights = np.where(
    treatment == 1,
    1 / propensity,      # Treated
    1 / (1 - propensity) # Control
)

# Stabilize weights
weights = weights * (treatment.mean() * treatment + 
                    (1-treatment.mean()) * (1-treatment))

# Weighted outcome regression
import statsmodels.api as sm
from statsmodels.stats.weightstats import DescrStatsW

# Check covariate balance
def check_balance(df, weights, treatment):
    balanced = []
    for col in features:
        treated = DescrStatsW(
            df[treatment==1][col], 
            weights[treatment==1]
        )
        control = DescrStatsW(
            df[treatment==0][col], 
            weights[treatment==0]
        )
        smd = abs(treated.mean - control.mean) / 
              np.sqrt((treated.var + control.var)/2)
        balanced.append(smd < 0.1)
    return all(balanced)

# Estimate causal effect
model = sm.WLS(
    mental_health_days,
    sm.add_constant(treatment),
    weights=weights
).fit()

# Get robust standard errors
results = model.get_robustcov_results(cov_type='HC3')
ate = results.params[1]
ci = results.conf_int()[1]

print(f"ATE: {ate:.2f} days")
print(f"95% CI: [{ci[0]:.2f}, {ci[1]:.2f}]")