pytest

Author: jmafoster1

causal_testing/estimation/gp.py

@@ -85,33 +85,46 @@ def __init__(
         df: pd.DataFrame,
         features: list,
         outcome: str,
+        max_order: int = 0,
         extra_operators: list = None,
         sympy_conversions: dict = None,
         seed=0,
     ):
         # pylint: disable=too-many-arguments
         random.seed(seed)
-        np.random.seed(seed)
         self.df = df
         self.features = features
         self.outcome = outcome
+        self.max_order = max_order
         self.seed = seed
         self.pset = gp.PrimitiveSet("MAIN", len(self.features))
         self.pset.renameArguments(**{f"ARG{i}": f for i, f in enumerate(self.features)})
 
-        standard_operators = [(add, 2), (mul, 2), (reciprocal, 1)]
+        standard_operators = [(add, 2), (mul, 2)]
         if extra_operators is None:
             extra_operators = [(log, 1), (reciprocal, 1)]
-        for operator, num_args in standard_operators + extra_operators:
-            self.pset.addPrimitive(operator, num_args)
         if sympy_conversions is None:
             sympy_conversions = {}
+        for operator, num_args in standard_operators + extra_operators:
+            self.pset.addPrimitive(operator, num_args)
+
         self.sympy_conversions = {
             "mul": lambda x1, x2: f"Mul({x1},{x2})",
             "add": lambda x1, x2: f"Add({x1},{x2})",
             "reciprocal": lambda x1: f"Pow({x1},-1)",
         } | sympy_conversions
 
+        for i in range(self.max_order):
+            print("Adding in order", i)
+            name = f"power_{i}"
+            self.pset.addPrimitive(lambda x: power(x, i), 1, name=name)
+            if name in self.sympy_conversions:
+                raise ValueError(
+                    f"You have provided a function called {name}, which is reserved for raising to power"
+                    f"{i}. Please choose a different name for your function."
+                )
+            self.sympy_conversions[name] = lambda x1: f"Pow({x1},{i})"
+
         creator.create("FitnessMin", base.Fitness, weights=(-1.0,))
         creator.create("Individual", gp.PrimitiveTree, fitness=creator.FitnessMin)
 

causal_testing/estimation/linear_regression_estimator.py

@@ -63,20 +63,21 @@ def __init__(
     def gp_formula(
         self,
         ngen: int = 100,
-        mu: int = 20,
-        lambda_: int = 10,
+        pop_size: int = 20,
+        num_offspring: int = 10,
         extra_operators: list = None,
         sympy_conversions: dict = None,
+        max_order: int = 0,
         seeds: list = None,
         seed: int = 0,
     ):
-        # pylint: disable=too-many-arguments,invalid-name
+        # pylint: disable=too-many-arguments
         """
         Use Genetic Programming (GP) to infer the regression equation from the data.
 
         :param ngen: The maximum number of GP generations to run for.
         :param mu: The GP population size.
-        :param lambda_: The number of offspring per generation.
+        :param num_offspring: The number of offspring per generation.
         :param extra_operators: Additional operators for the GP (defaults are +, *, and 1/x). Operations should be of
                                 the form (fun, numArgs), e.g. (add, 2).
         :param sympy_conversions: Dictionary of conversions of extra_operators for sympy,
@@ -85,16 +86,17 @@ def gp_formula(
                       probably logarithmic, you can put that in).
         :param seed: Random seed for the GP.
         """
-        gp = GP(
+        self.gp = GP(
             df=self.df,
             features=sorted(list(self.adjustment_set.union([self.treatment]))),
             outcome=self.outcome,
             extra_operators=extra_operators,
             sympy_conversions=sympy_conversions,
             seed=seed,
+            max_order=max_order,
         )
-        formula = gp.run_gp(ngen=ngen, pop_size=mu, num_offspring=lambda_, seeds=seeds)
-        formula = gp.simplify(formula)
+        formula = self.gp.run_gp(ngen=ngen, pop_size=num_offspring, num_offspring=num_offspring, seeds=seeds)
+        formula = self.gp.simplify(formula)
         self.formula = f"{self.outcome} ~ I({formula}) - 1"
 
     def add_modelling_assumptions(self):

tests/estimation_tests/test_linear_regression_estimator.py

@@ -7,6 +7,7 @@
 from causal_testing.specification.capabilities import TreatmentSequence
 
 from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
+from causal_testing.estimation.gp import reciprocal
 
 
 def load_nhefs_df():
@@ -275,16 +276,30 @@ def test_gp(self):
         df["X"] = np.arange(10)
         df["Y"] = 1 / (df["X"] + 1)
         linear_regression_estimator = LinearRegressionEstimator("X", 0, 1, set(), "Y", df.astype(float))
-        linear_regression_estimator.gp_formula(seed=1)
-        self.assertEqual(
-            linear_regression_estimator.formula,
-            "Y ~ I((2.606801258739728e-17*X + 0.626132756132756)/(0.6261327561327561*X + 0.626132756132756)) - 1",
-        )
+        linear_regression_estimator.gp_formula(seeds=["reciprocal(add(X, 1))"])
+        print("MAPPING")
+        print(linear_regression_estimator.gp.pset.mapping)
+        self.assertEqual(linear_regression_estimator.formula, "Y ~ I(1/(X + 1)) - 1")
         ate, (ci_low, ci_high) = linear_regression_estimator.estimate_ate_calculated()
         self.assertEqual(round(ate[0], 2), 0.50)
         self.assertEqual(round(ci_low[0], 2), 0.50)
         self.assertEqual(round(ci_high[0], 2), 0.50)
 
+    def test_gp_power(self):
+        df = pd.DataFrame()
+        df["X"] = np.arange(10)
+        df["Y"] = 2 * (df["X"] ** 2)
+        linear_regression_estimator = LinearRegressionEstimator("X", 0, 1, set(), "Y", df.astype(float))
+        linear_regression_estimator.gp_formula(seed=1, max_order=0)
+        self.assertEqual(
+            linear_regression_estimator.formula,
+            "Y ~ I(2.0*X**2 + 3.8205100524608823e-31) - 1",
+        )
+        ate, (ci_low, ci_high) = linear_regression_estimator.estimate_ate_calculated()
+        self.assertEqual(round(ate[0], 2), -2.00)
+        self.assertEqual(round(ci_low[0], 2), -2.00)
+        self.assertEqual(round(ci_high[0], 2), -2.00)
+
 
 class TestLinearRegressionInteraction(unittest.TestCase):
     """Test linear regression for estimating effects involving interaction."""