fix: fix test case for adj estimator

CyberAgentAILab · Sep 5, 2024 · 48fdf66 · 48fdf66
1 parent 5a5f9b2
commit 48fdf66
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 3 deletions.
diff --git a/dte_adj/__init__.py b/dte_adj/__init__.py
@@ -481,7 +481,6 @@ def _compute_cumulative_distribution(
         cumulative_distribution = np.zeros(n_loc)
         superset_prediction = np.zeros((n_records, n_loc))
         treatment_mask = treatment_arms == target_treatment_arm
-        if self.is_multi_task:
         confounding_in_arm = confoundings[treatment_mask]
         n_records_in_arm = len(confounding_in_arm)
         if self.is_multi_task:
@@ -521,7 +520,6 @@ def _compute_cumulative_distribution(
                     subset_train_mask = (~superset_mask) & treatment_mask
                     subset_test_mask_inner = superset_mask[treatment_mask]
                     confounding_train = confoundings[subset_train_mask]
-                    confounding_subset_test = confoundings[subset_test_mask]
                     binominal_train = binominal[subset_train_mask]
                     if len(np.unique(binominal_train)) == 1:
                         subset_prediction[subset_test_mask_inner] = binominal_train[0]

diff --git a/tests/test_adjusted_estimator.py b/tests/test_adjusted_estimator.py
@@ -82,5 +82,5 @@ def test_compute_cumulative_distribution(self):
             self.assertAlmostEqual(cumulative_distribution[i], (i + 1) / 10, places=2)
 
         for i in range(20):
-            for j in range(1, 10):
+            for j in range(1, 8):
                 self.assertAlmostEqual(superset_prediction[i, j], 0.5, places=2)