move add_covariates logic to hypothesis test

david26694 · Oct 4, 2024 · ccba301 · ccba301
1 parent 05be002
commit ccba301
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 7 deletions.
diff --git a/cluster_experiments/inference/analysis_plan.py b/cluster_experiments/inference/analysis_plan.py
@@ -108,12 +108,7 @@ def analyze(
         analysis_results = EmptyAnalysisPlanResults()
 
         for test in self.tests:
-            if test.is_cupac:
-                exp_data = test.cupac_handler.add_covariates(
-                    df=exp_data, pre_experiment_df=pre_exp_data
-                )
-
-            target_col = test.metric.target_column
+            exp_data = test.add_covariates(exp_data, pre_exp_data)
 
             for treatment_variant in self.treatment_variants:
                 for dimension in test.dimensions:
@@ -128,7 +123,7 @@ def analyze(
                             )
 
                         test._prepare_analysis_config(
-                            target_col=target_col,
+                            target_col=test.metric.target_column,
                             treatment_col=self.variant_col,
                             treatment=treatment_variant.name,
                         )

diff --git a/cluster_experiments/inference/hypothesis_test.py b/cluster_experiments/inference/hypothesis_test.py
@@ -194,3 +194,16 @@ def prepare_data(
         ).query(f"{dimension_name} == '{dimension_value}'")
 
         return prepared_df
+
+    def add_covariates(
+        self, exp_data: pd.DataFrame, pre_exp_data: pd.DataFrame
+    ) -> pd.DataFrame:
+        """
+        If the test is a cupac test, adds the covariates to the experimental data.
+        """
+        if self.is_cupac:
+            exp_data = self.cupac_handler.add_covariates(
+                df=exp_data, pre_experiment_df=pre_exp_data
+            )
+
+        return exp_data