update cov test

Author: SvenKlaassen

doubleml/irm/tests/test_iivm_unif_confset.py

@@ -5,66 +5,51 @@
 
 import doubleml as dml
 
-np.random.seed(3141)
 
-
-@pytest.fixture(scope="module")
-def true_ATE():
-    return 0.5
-
-
-@pytest.fixture(scope="module")
-def instrument_size():
-    return 0.005
-
-
-@pytest.fixture(scope="module")
-def n_samples():
-    return 1000
-
-
-@pytest.fixture(scope="module")
-def n_simulations():
-    return 100
-
-
-@pytest.fixture(scope="module")
-def weakiv_data(n_samples, instrument_size, true_ATE):
-    # Generate data
+def generate_weak_iv_data(n_samples, instrument_size, true_ATE):
     u = np.random.normal(0, 2, size=n_samples)
     X = np.random.normal(0, 1, size=n_samples)
     Z = np.random.binomial(1, 0.5, size=n_samples)
-    A = instrument_size * Z + u  # Continuous treatment A
+    A = instrument_size * Z + u
     A = np.array(A > 0, dtype=int)
-    Y = true_ATE * A + np.sign(u)  # Outcome Y
-    return dml.DoubleMLData.from_arrays(x=X, y=Y, d=A, z=Z)
-
+    Y = true_ATE * A + np.sign(u)
+    dml_data = dml.DoubleMLData.from_arrays(x=X, y=Y, d=A, z=Z)
+    return dml_data
 
-@pytest.fixture(scope="module")
-def iivm_obj(weakiv_data):
-    # Set machine learning methods for m, r & g
-    learner_g = LinearRegression()
-    classifier_m = LogisticRegression()
-    classifier_r = RandomForestClassifier(n_estimators=20, max_depth=5)
 
-    # Create DoubleMLIIVM object
-    obj_dml_data = weakiv_data
-    dml_iivm_obj = dml.DoubleMLIIVM(obj_dml_data, learner_g, classifier_m, classifier_r)
-    return dml_iivm_obj
+@pytest.mark.ci
+def test_coverage_robust_confset():
+    # Test parameters
+    true_ATE = 0.5
+    instrument_size = 0.005
+    n_samples = 1000
+    n_simulations = 100
 
-
-def test_coverage(iivm_obj, true_ATE, n_simulations):
+    np.random.seed(3141)
     coverage = []
     for _ in range(n_simulations):
-        # Fit the model
-        iivm_obj.fit()
+        data = generate_weak_iv_data(n_samples, instrument_size, true_ATE)
+
+        # Set machine learning methods
+        learner_g = LinearRegression()
+        classifier_m = LogisticRegression()
+        classifier_r = RandomForestClassifier(n_estimators=20, max_depth=5)
 
-        # Get the confidence set
-        conf_set = iivm_obj.robust_confset()
+        # Create and fit new model
+        dml_iivm_obj = dml.DoubleMLIIVM(data, learner_g, classifier_m, classifier_r)
+        dml_iivm_obj.fit()
 
-        # Check if the true ATE is in the confidence set
+        # Get confidence set
+        conf_set = dml_iivm_obj.robust_confset()
+
+        # check if conf_set is a list of tuples
+        assert isinstance(conf_set, list)
+        assert all(isinstance(x, tuple) and len(x) == 2 for x in conf_set)
+
+        # Check if true ATE is in confidence set
         ate_in_confset = any(x[0] < true_ATE < x[1] for x in conf_set)
         coverage.append(ate_in_confset)
-    # Calculate the coverage rate
+
+    # Calculate coverage rate
     coverage_rate = np.mean(coverage)
     assert coverage_rate >= 0.9, f"Coverage rate {coverage_rate} is below 0.9"