Fix the test_accuracy function by modifying the assertion logic (#867)

Fix the test_accuracy function by modifying the assertion logic

Signed-off-by: kgao <[email protected]>

Author: kgao

econml/tests/test_discrete_outcome.py

@@ -27,9 +27,7 @@ def test_accuracy(self):
         discrete_outcome = True
         discrete_treatment = True
         true_ate = 0.3
-        W = np.random.uniform(-1, 1, size=(n, 1))
-        D = np.random.binomial(1, .5 + .1 * W[:, 0], size=(n,))
-        Y = np.random.binomial(1, .5 + true_ate * D + .1 * W[:, 0], size=(n,))
+        num_iterations = 10
 
         ests = [
             LinearDML(discrete_outcome=discrete_outcome, discrete_treatment=discrete_treatment),
@@ -39,34 +37,42 @@ def test_accuracy(self):
 
         for est in ests:
 
-            if isinstance(est, CausalForestDML):
-                est.fit(Y, D, X=W)
-                ate = est.ate(X=W)
-                ate_lb, ate_ub = est.ate_interval(X=W)
+            count_within_interval = 0
 
-            else:
-                est.fit(Y, D, W=W)
-                ate = est.ate()
-                ate_lb, ate_ub = est.ate_interval()
+            for _ in range(num_iterations):
 
-            if isinstance(est, LinearDRLearner):
-                est.summary(T=1)
-            else:
-                est.summary()
+                W = np.random.uniform(-1, 1, size=(n, 1))
+                D = np.random.binomial(1, .5 + .1 * W[:, 0], size=(n,))
+                Y = np.random.binomial(1, .5 + true_ate * D + .1 * W[:, 0], size=(n,))
+
+                if isinstance(est, CausalForestDML):
+                    est.fit(Y, D, X=W)
+                    ate_lb, ate_ub = est.ate_interval(X=W)
+
+                else:
+                    est.fit(Y, D, W=W)
+                    ate_lb, ate_ub = est.ate_interval()
+
+                if isinstance(est, LinearDRLearner):
+                    est.summary(T=1)
+                else:
+                    est.summary()
+
+                if ate_lb <= true_ate <= ate_ub:
+                    count_within_interval += 1
 
-            proportion_in_interval = ((ate_lb < true_ate) & (true_ate < ate_ub)).mean()
-            np.testing.assert_array_less(0.50, proportion_in_interval)
+            assert count_within_interval >= 7, (
+                f"{est.__class__.__name__}: True ATE falls within the interval bounds "
+                f"only {count_within_interval} times out of {num_iterations}"
+            )
 
     # accuracy test, DML
     def test_accuracy_iv(self):
-        n = 10000
+        n = 1000
         discrete_outcome = True
         discrete_treatment = True
         true_ate = 0.3
-        W = np.random.uniform(-1, 1, size=(n, 1))
-        Z = np.random.uniform(-1, 1, size=(n, 1))
-        D = np.random.binomial(1, .5 + .1 * W[:, 0] + .1 * Z[:, 0], size=(n,))
-        Y = np.random.binomial(1, .5 + true_ate * D + .1 * W[:, 0], size=(n,))
+        num_iterations = 10
 
         ests = [
             OrthoIV(discrete_outcome=discrete_outcome, discrete_treatment=discrete_treatment),
@@ -75,14 +81,26 @@ def test_accuracy_iv(self):
 
         for est in ests:
 
-            est.fit(Y, D, W=W, Z=Z)
-            ate = est.ate()
-            ate_lb, ate_ub = est.ate_interval()
+            count_within_interval = 0
+
+            for _ in range(num_iterations):
+
+                W = np.random.uniform(-1, 1, size=(n, 1))
+                Z = np.random.uniform(-1, 1, size=(n, 1))
+                D = np.random.binomial(1, .5 + .1 * W[:, 0] + .1 * Z[:, 0], size=(n,))
+                Y = np.random.binomial(1, .5 + true_ate * D + .1 * W[:, 0], size=(n,))
+
+                est.fit(Y, D, W=W, Z=Z)
+                ate_lb, ate_ub = est.ate_interval()
+                est.summary()
 
-            est.summary()
+                if ate_lb <= true_ate <= ate_ub:
+                    count_within_interval += 1
 
-            proportion_in_interval = ((ate_lb < true_ate) & (true_ate < ate_ub)).mean()
-            np.testing.assert_array_less(0.50, proportion_in_interval)
+            assert count_within_interval >= 7, (
+                f"{est.__class__.__name__}: True ATE falls within the interval bounds "
+                f"only {count_within_interval} times out of {num_iterations}"
+            )
 
     def test_string_outcome(self):
         n = 100