diff --git a/causality/estimation/parametric.py b/causality/estimation/parametric.py index 997e440..104a781 100644 --- a/causality/estimation/parametric.py +++ b/causality/estimation/parametric.py @@ -32,22 +32,22 @@ def average_treatment_effect(self, X, start='Start', end='End', assignment='Assi control_final = control['End'] del test, control - df = pd.DataFrame({'y' : test_initial, - 'assignment' : [1. for i in test_initial], + df = pd.DataFrame({'y' : test_initial, + 'assignment' : [1. for i in test_initial], 't' :[0. for i in test_initial] }) - df = df.append(pd.DataFrame({'y' : test_final, - 'assignment' : [1. for i in test_final], + df = df.append(pd.DataFrame({'y' : test_final, + 'assignment' : [1. for i in test_final], 't' :[1. for i in test_final] })) - df = df.append(pd.DataFrame({'y' : control_initial, - 'assignment' : [0. for i in control_initial], + df = df.append(pd.DataFrame({'y' : control_initial, + 'assignment' : [0. for i in control_initial], 't' :[0. for i in control_initial] })) - df = df.append(pd.DataFrame({'y' : control_final, - 'assignment' : [0. for i in control_final], + df = df.append(pd.DataFrame({'y' : control_final, + 'assignment' : [0. for i in control_final], 't' :[1. for i in control_final] })) del test_initial, test_final, control_initial, control_final - df['did'] = df['t'] * df['assignment'] + df['did'] = df['t'] * df['assignment'] df['intercept'] = 1. model = self.model(df['y'], df[['t', 'assignment','did', 'intercept']]) @@ -55,7 +55,7 @@ def average_treatment_effect(self, X, start='Start', end='End', assignment='Assi conf_int = result.conf_int().ix['did'] expected = result.params['did'] return conf_int[0], expected, conf_int[1] - + def test_parallel_trend(self, X, start='Start', end='End', assignment='Assignment'): """ This will find the average treatment effect on @@ -64,7 +64,7 @@ def test_parallel_trend(self, X, start='Start', end='End', assignment='Assignmen that the average treatment effect between the test and control groups when neither is treated is 0. - The format for this dataset is the same as that + The format for this dataset is the same as that for the real estimation task, except that the start time is some time before the experiment is run, and the end time is the starting point for the experiment. @@ -145,4 +145,9 @@ def estimate_ATC(self, X, assignment, outcome, confounder_types, n_neighbors=5): def estimate_ATE(self, X, assignment, outcome, confounder_types, n_neighbors=5): att = estimate_ATT(self, X, assignment, outcome, confounder_types, n_neighbors=n_neighbors) atc = estimate_ATC(self, X, assignment, outcome, confounder_types, n_neighbors=n_neighbors) - return (atc+att)/2. + return (atc+att)/2. + + +class CausalImpact(object): + def __init__(self): + pass