Seeding gp power

Author: jmafoster1

causal_testing/estimation/gp.py

@@ -134,7 +134,7 @@ def __init__(
             "reciprocal": lambda x1: f"Pow({x1},-1)",
         } | sympy_conversions
 
-        for i in range(self.max_order):
+        for i in range(self.max_order + 1):
             name = f"power_{i}"
             func, conversion = create_power_function(i)
             self.pset.addPrimitive(func, 1, name=name)
@@ -145,6 +145,7 @@ def __init__(
                 )
             self.sympy_conversions[name] = conversion
 
+        print(self.pset.mapping)
         creator.create("FitnessMin", base.Fitness, weights=(-1.0,))
         creator.create("Individual", gp.PrimitiveTree, fitness=creator.FitnessMin)
 

tests/estimation_tests/test_linear_regression_estimator.py

@@ -303,10 +303,10 @@ def test_gp_power(self):
         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)
+        linear_regression_estimator.gp_formula(seed=1, max_order=2, seeds=["mul(2, power_2(X))"])
         self.assertEqual(
             linear_regression_estimator.formula,
-            "Y ~ I(1.9999999999999999*X**2 - 1.0043240235058056e-116*X + 2.6645352591003757e-15) - 1",
+            "Y ~ I(2*X**2) - 1",
         )
         ate, (ci_low, ci_high) = linear_regression_estimator.estimate_ate_calculated()
         self.assertEqual(round(ate[0], 2), -2.00)