pylint

Author: jmafoster1

causal_testing/estimation/gp.py

@@ -75,6 +75,16 @@ def mut_insert(expression: gp.PrimitiveTree, pset: gp.PrimitiveSet):
     return (expression,)
 
 
+def create_power_function(power):
+    def power_func(x):
+        return power(x, power)
+
+    def sympy_conversion(x1):
+        return f"Pow({x1},{i})"
+
+    return power_func, sympy_conversion
+
+
 class GP:
     """
     Object to perform genetic programming.
@@ -90,7 +100,7 @@ def __init__(
         sympy_conversions: dict = None,
         seed=0,
     ):
-        # pylint: disable=too-many-arguments
+        # pylint: disable=too-many-arguments,too-many-instance-attributes
         random.seed(seed)
         self.df = df
         self.features = features
@@ -115,15 +125,15 @@ def __init__(
         } | 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)
+            func, conversion = create_power_function(i)
+            self.pset.addPrimitive(func, 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})"
+            self.sympy_conversions[name] = conversion
 
         creator.create("FitnessMin", base.Fitness, weights=(-1.0,))
         creator.create("Individual", gp.PrimitiveTree, fitness=creator.FitnessMin)

causal_testing/estimation/linear_regression_estimator.py

@@ -65,9 +65,9 @@ def gp_formula(
         ngen: int = 100,
         pop_size: int = 20,
         num_offspring: int = 10,
+        max_order: int = 0,
         extra_operators: list = None,
         sympy_conversions: dict = None,
-        max_order: int = 0,
         seeds: list = None,
         seed: int = 0,
     ):
@@ -76,17 +76,18 @@ def gp_formula(
         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 pop_size: The GP population size.
         :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
+        :param max_order: The maximum polynomial order to use, e.g. `max_order=2` will give polynomials of the form `ax^2 + bx + c`.
+        :param extra_operators: Additional operators for the GP (defaults are +, *, log(x), 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,
                                   e.g. `"mul": lambda *args_: "Mul({},{})".format(*args_)`.
         :param seeds: Seed individuals for the population (e.g. if you think that the relationship between X and Y is
                       probably logarithmic, you can put that in).
         :param seed: Random seed for the GP.
         """
-        self.gp = GP(
+        gp = GP(
             df=self.df,
             features=sorted(list(self.adjustment_set.union([self.treatment]))),
             outcome=self.outcome,
@@ -95,8 +96,8 @@ def gp_formula(
             seed=seed,
             max_order=max_order,
         )
-        formula = self.gp.run_gp(ngen=ngen, pop_size=num_offspring, num_offspring=num_offspring, seeds=seeds)
-        formula = self.gp.simplify(formula)
+        formula = gp.run_gp(ngen=ngen, pop_size=pop_size, num_offspring=num_offspring, seeds=seeds)
+        formula = gp.simplify(formula)
         self.formula = f"{self.outcome} ~ I({formula}) - 1"
 
     def add_modelling_assumptions(self):

tests/estimation_tests/test_linear_regression_estimator.py

@@ -277,8 +277,6 @@ def test_gp(self):
         df["Y"] = 1 / (df["X"] + 1)
         linear_regression_estimator = LinearRegressionEstimator("X", 0, 1, set(), "Y", df.astype(float))
         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)