Add: initial unit tests for causal surrogate

Author: f-allian

tests/surrogate_tests/test_causal_surrogate_assisted.py

@@ -0,0 +1,42 @@
+import unittest
+from causal_testing.surrogate.causal_surrogate_assisted import SimulationResult, SearchFitnessFunction
+from causal_testing.testing.estimators import Estimator, PolynomialRegressionEstimator
+
+class TestSimulationResult(unittest.TestCase):
+
+    def setUp(self):
+
+        self.data = {'key': 'value'}
+
+    def test_inputs(self):
+
+        fault_values = [True, False]
+
+        relationship_values = ["positive", "negative", None]
+
+        for fault in fault_values:
+
+            for relationship in relationship_values:
+                with self.subTest(fault=fault, relationship=relationship):
+                    result = SimulationResult(data=self.data, fault=fault, relationship=relationship)
+
+                    self.assertIsInstance(result.data, dict)
+
+                    self.assertEqual(result.fault, fault)
+
+                    self.assertEqual(result.relationship, relationship)
+
+class TestSearchFitnessFunction(unittest.TestCase):
+
+    #TODO: complete tests for causal surrogate
+
+    def test_init_valid_values(self):
+
+        test_function = lambda x: x **2
+
+        surrogate_model = PolynomialRegressionEstimator()
+
+        search_function = SearchFitnessFunction(fitness_function=test_function, surrogate_model=surrogate_model)
+
+        self.assertIsCallable(search_function.fitness_function)
+        self.assertIsInstance(search_function.surrogate_model, PolynomialRegressionEstimator)

tests/surrogate_tests/test_surrogate_search_algorithms.py

@@ -7,6 +7,7 @@
     CausalForestEstimator,
     LogisticRegressionEstimator,
     InstrumentalVariableEstimator,
+    PolynomialRegressionEstimator
 )
 from causal_testing.specification.variable import Input
 from causal_testing.utils.validation import CausalValidator
@@ -409,6 +410,43 @@ def test_program_11_2_with_robustness_validation(self):
         self.assertEqual(round(cv.estimate_robustness(model)["treatments"], 4), 0.7353)
 
 
+class TestPolynomialRegressionEstimator(TestLinearRegressionEstimator):
+
+    @classmethod
+
+    def setUpClass(cls):
+
+        super().setUpClass()
+    def test_program_11_3_polynomial(self):
+
+        """Test whether the polynomial regression implementation produces the same results as program 11.3 (p. 162).
+        https://www.hsph.harvard.edu/miguel-hernan/wp-content/uploads/sites/1268/2023/10/hernanrobins_WhatIf_30sep23.pdf
+        """
+
+        df = self.chapter_11_df.copy()
+
+        polynomial_estimator = PolynomialRegressionEstimator(
+            "treatments", None, None, set(), "outcomes", 3, df)
+
+        model = polynomial_estimator._run_linear_regression()
+
+        ate, _ = polynomial_estimator.estimate_coefficient()
+
+        self.assertEqual(
+            round(
+                model.params["Intercept"]
+                + 90 * model.params["treatments"]
+                + 90 * 90 * model.params["np.power(treatments, 2)"],
+                1,
+            ),
+            197.1,
+        )
+        # Increasing treatments from 90 to 100 should be the same as 10 times the unit ATE
+        self.assertEqual(round(model.params["treatments"], 3), round(ate, 3))
+
+
+
+
 class TestCausalForestEstimator(unittest.TestCase):
     """Test the linear regression estimator against the programming exercises in Section 2 of Hernán and Robins [1].
 
@@ -491,3 +529,5 @@ def test_X1_effect(self):
         test_results = lr_model.estimate_ate()
         ate = test_results[0]
         self.assertAlmostEqual(ate, 2.0)
+
+