update

sdaza · Jan 1, 2025 · f273cc0 · f273cc0
1 parent 3c4efdd
commit f273cc0
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/experiment_utils/estimators.py b/experiment_utils/estimators.py
@@ -267,6 +267,7 @@ def ipw_xgboost(self, data: pd.DataFrame, covariates: List[str]) -> pd.DataFrame
         data['propensity_score'] = np.minimum(self.max_ps_score, data['propensity_score'])
         data['propensity_score'] = np.maximum(self.min_ps_score, data['propensity_score'])
         data = self.__calculate_stabilized_weights(data)
+
         return data
 
     def __calculate_stabilized_weights(self, data: pd.DataFrame) -> pd.DataFrame:
@@ -293,6 +294,7 @@ def __calculate_stabilized_weights(self, data: pd.DataFrame) -> pd.DataFrame:
         ) * (
             1 - p_treatment
         )
+
         data["tips_stabilized_weight"] = data[self.treatment_col] * p_treatment + (
             1 - data[self.treatment_col]
         ) * data["propensity_score"] / (1 - data["propensity_score"]) * (1 - p_treatment)

diff --git a/experiment_utils/experiment_analyzer.py b/experiment_utils/experiment_analyzer.py
@@ -96,6 +96,7 @@ def __init__(
                                "ATE": "ips_stabilized_weight",
                                "ATC": "cips_stabilized_weight"}
 
+
     def __check_input(self) -> None:
 
         # dataframe is empty
@@ -397,7 +398,6 @@ def get_effects(self, min_binary_count: int = 100, adjustment: Optional[str] = N
 
             # create adjustment label
             relevant_covariates = set(self.final_covariates) & set(self.regression_covariates)
-            standardized_covariates = [f"z_{cov}" for cov in relevant_covariates]
 
             adjustment_labels = {
                 'IPW': 'IPW',
@@ -415,10 +415,10 @@ def get_effects(self, min_binary_count: int = 100, adjustment: Optional[str] = N
 
             for outcome in self.outcomes:
                 if adjustment == 'IPW':
-                    output = models[adjustment](data=temp_pd, outcome_variable=outcome, covariates=standardized_covariates,
-                        weight_column='weights')
+                    output = models[adjustment](data=temp_pd, outcome_variable=outcome, covariates=relevant_covariates,
+                        weight_column=self.target_weights[self.target_ipw_effect])
                 else:
-                    output = models[adjustment](data=temp_pd, outcome_variable=outcome, covariates=standardized_covariates)
+                    output = models[adjustment](data=temp_pd, outcome_variable=outcome, covariates=relevant_covariates)
                 output['adjustment'] = adjustment_label
                 if adjustment == 'IPW':
                     output['balance'] = np.round(adjusted_balance['balance_flag'].mean(), 2)