diff --git a/experiment_utils/estimators.py b/experiment_utils/estimators.py index c7249b8..4779207 100644 --- 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 index 55b60d3..ec94a9b 100644 --- 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)