Improved GP fitness to use NRMSE

Author: jmafoster1

causal_testing/estimation/genetic_programming_regression_fitter.py

@@ -149,7 +149,6 @@ 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)
 
@@ -235,6 +234,8 @@ def simplify(self, expression: gp.PrimitiveTree) -> sympy.core.Expr:
 
         :return: The simplified expression as a sympy Expr object.
         """
+        if isinstance(expression, str):
+            expression = creator.Individual(gp.PrimitiveTree.from_string(expression, self.pset))
         return sympy.simplify(self._stringify_for_sympy(expression))
 
     def repair(self, expression: gp.PrimitiveTree) -> gp.PrimitiveTree:
@@ -278,16 +279,23 @@ def fitness(self, expression: gp.PrimitiveTree) -> float:
         """
         old_settings = np.seterr(all="raise")
         try:
+            if isinstance(expression, str):
+                expression = creator.Individual(gp.PrimitiveTree.from_string(expression, self.pset))
+
             # Create model, fit (run) it, give estimates from it]
             func = gp.compile(expression, self.pset)
-            y_estimates = pd.Series([func(**x) for _, x in self.df[self.features].iterrows()])
+            y_estimates = pd.Series([func(**x) for _, x in self.df[self.features].iterrows()], index=self.df.index)
 
-            # Calc errors using an improved normalised mean squared
+            # Calculate errors using root mean square error
             sqerrors = (self.df[self.outcome] - y_estimates) ** 2
-            mean_squared = sqerrors.sum() / len(self.df)
-            nmse = mean_squared / (self.df[self.outcome].sum() / len(self.df))
+            nrmse = np.sqrt(sqerrors.sum()/len(self.df))/(self.df[self.outcome].max() - self.df[self.outcome].min())
+
+
+            if pd.isnull(nrmse) or nrmse.real != nrmse:
+                return (float("inf"),)
+            assert nrmse > 0, f"NRMSE {nrmse} should be greater than zero"
 
-            return (nmse,)
+            return (nrmse,)
 
             # Fitness value of infinite if error - not return 1
         except (
@@ -321,7 +329,7 @@ def make_offspring(self, population: list, num_offspring: int) -> list:
             offspring.append(child)
         return offspring
 
-    def run_gp(self, ngen: int, pop_size: int = 20, num_offspring: int = 10, seeds: list = None) -> gp.PrimitiveTree:
+    def run_gp(self, ngen: int, pop_size: int = 20, num_offspring: int = 10, seeds: list = None, repair=True) -> gp.PrimitiveTree:
         """
         Execute Genetic Programming to find the best expression using a mu+lambda algorithm.
 
@@ -332,7 +340,9 @@ def run_gp(self, ngen: int, pop_size: int = 20, num_offspring: int = 10, seeds:
 
         :return: The best candididate expression.
         """
-        population = [self.toolbox.repair(ind) for ind in self.toolbox.population(n=pop_size)]
+        population = self.toolbox.population(n=pop_size)
+        if repair:
+            population = [self.toolbox.repair(ind) for ind in population]
         if seeds is not None:
             for seed in seeds:
                 ind = creator.Individual(gp.PrimitiveTree.from_string(seed, self.pset))
@@ -348,7 +358,8 @@ def run_gp(self, ngen: int, pop_size: int = 20, num_offspring: int = 10, seeds:
         for _ in range(1, ngen + 1):
             # Vary the population
             offspring = self.make_offspring(population, num_offspring)
-            offspring = [self.toolbox.repair(ind) for ind in offspring]
+            if repair:
+                offspring = [self.toolbox.repair(ind) for ind in offspring]
 
             # Evaluate the individuals with an invalid fitness
             for ind in offspring: