Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def get_weighted_effect_estimate(self, assignment, df, outcome, bootstrap=False):
def estimate(df):
treated = df[df[assignment] == 1]
control = df[df[assignment] == 0]
treated_outcome = (treated[outcome]*treated['weight']).sum() / treated['weight'].sum()
control_outcome = (control[outcome]*control['weight']).sum() / control['weight'].sum()
return treated_outcome - control_outcome
if bootstrap:
return bootstrap_statistic(df, estimate)
else:
return estimate(df)