Fixed bug in fitness function scope

Author: rsomers1998

causal_testing/surrogate/causal_surrogate_assisted.py

@@ -19,29 +19,15 @@ class SimulationResult:
     relationship: str
 
 
-@dataclass
-class SearchFitnessFunction:
-    """Data class containing the Fitness function and related model"""
-
-    fitness_function: Any
-    surrogate_model: CubicSplineRegressionEstimator
-
-
 class SearchAlgorithm(ABC):
     """Class to be inherited with the search algorithm consisting of a search function and the fitness function of the
     space to be searched"""
 
     @abstractmethod
-    def generate_fitness_functions(self, surrogate_models: list[Estimator]) -> list[SearchFitnessFunction]:
-        """Generates the fitness function of the search space
-        :param surrogate_models: A list of CubicSplineRegressionEstimator generated for each edge of the DAG
-        :return: A list of fitness functions mapping to each of the surrogate models in the input"""
-
-    @abstractmethod
-    def search(self, fitness_functions: list[SearchFitnessFunction], specification: CausalSpecification) -> list:
+    def search(self, surrogate_models: list[CubicSplineRegressionEstimator], specification: CausalSpecification) -> list:
         """Function which implements a search routine which searches for the optimal fitness value for the specified
         scenario
-        :param fitness_functions: The fitness function to be optimised
+        :param surrogate_models: The surrogate models to be searched
         :param specification:  The Causal Specification (combination of Scenario and Causal Dag)"""
 
 
@@ -95,8 +81,7 @@ def execute(
 
         for i in range(max_executions):
             surrogate_models = self.generate_surrogates(self.specification, data_collector)
-            fitness_functions = self.search_algorithm.generate_fitness_functions(surrogate_models)
-            candidate_test_case, _, surrogate = self.search_algorithm.search(fitness_functions, self.specification)
+            candidate_test_case, _, surrogate = self.search_algorithm.search(surrogate_models, self.specification)
 
             self.simulator.startup()
             test_result = self.simulator.run_with_config(candidate_test_case)
@@ -123,7 +108,7 @@ def execute(
 
     def generate_surrogates(
         self, specification: CausalSpecification, data_collector: ObservationalDataCollector
-    ) -> list[SearchFitnessFunction]:
+    ) -> list[CubicSplineRegressionEstimator]:
         """Generate a surrogate model for each edge of the dag that specifies it is included in the DAG metadata.
         :param specification: The Causal Specification (combination of Scenario and Causal Dag)
         :param data_collector: An ObservationalDataCollector which gathers data relevant to the specified scenario

causal_testing/surrogate/surrogate_search_algorithms.py

@@ -7,7 +7,7 @@
 
 from causal_testing.specification.causal_specification import CausalSpecification
 from causal_testing.testing.estimators import CubicSplineRegressionEstimator
-from causal_testing.surrogate.causal_surrogate_assisted import SearchAlgorithm, SearchFitnessFunction
+from causal_testing.surrogate.causal_surrogate_assisted import SearchAlgorithm
 
 
 class GeneticSearchAlgorithm(SearchAlgorithm):
@@ -25,15 +25,15 @@ def __init__(self, delta=0.05, config: dict = None) -> None:
             "some_effect": lambda x: abs(1 / x),
         }
 
-    def generate_fitness_functions(
-        self, surrogate_models: list[CubicSplineRegressionEstimator]
-    ) -> list[SearchFitnessFunction]:
-        fitness_functions = []
+    def search(
+            self, surrogate_models: list[CubicSplineRegressionEstimator], specification: CausalSpecification
+    ) -> list:
+        solutions = []
 
         for surrogate in surrogate_models:
             contradiction_function = self.contradiction_functions[surrogate.expected_relationship]
 
-            # The returned fitness function after including required variables into the function's scope
+            # The GA fitness function after including required variables into the function's scope
             # Unused arguments are required for pygad's fitness function signature
             def fitness_function(ga, solution, idx): # pylint: disable=unused-argument
                 surrogate.control_value = solution[0] - self.delta
@@ -46,25 +46,15 @@ def fitness_function(ga, solution, idx): # pylint: disable=unused-argument
                 ate = surrogate.estimate_ate_calculated(adjustment_dict)
 
                 return contradiction_function(ate)
-
-            search_fitness_function = SearchFitnessFunction(fitness_function, surrogate)
-
-            fitness_functions.append(search_fitness_function)
-
-        return fitness_functions
-
-    def search(self, fitness_functions: list[SearchFitnessFunction], specification: CausalSpecification) -> list:
-        solutions = []
-
-        for fitness_function in fitness_functions:
-            gene_types, gene_space = self.create_gene_types(fitness_function, specification)
+            
+            gene_types, gene_space = self.create_gene_types(surrogate, specification)
 
             ga = GA(
                 num_generations=200,
                 num_parents_mating=4,
-                fitness_func=fitness_function.fitness_function,
+                fitness_func=fitness_function,
                 sol_per_pop=10,
-                num_genes=1 + len(fitness_function.surrogate_model.adjustment_set),
+                num_genes=1 + len(surrogate.adjustment_set),
                 gene_space=gene_space,
                 gene_type=gene_types,
             )
@@ -82,24 +72,24 @@ def search(self, fitness_functions: list[SearchFitnessFunction], specification:
             solution, fitness, _ = ga.best_solution()
 
             solution_dict = {}
-            solution_dict[fitness_function.surrogate_model.treatment] = solution[0]
-            for idx, adj in enumerate(fitness_function.surrogate_model.adjustment_set):
+            solution_dict[surrogate.treatment] = solution[0]
+            for idx, adj in enumerate(surrogate.adjustment_set):
                 solution_dict[adj] = solution[idx + 1]
-            solutions.append((solution_dict, fitness, fitness_function.surrogate_model))
+            solutions.append((solution_dict, fitness, surrogate))
 
         return max(solutions, key=itemgetter(1))  # This can be done better with fitness normalisation between edges
 
     @staticmethod
     def create_gene_types(
-        fitness_function: SearchFitnessFunction, specification: CausalSpecification
+        surrogate_model: CubicSplineRegressionEstimator, specification: CausalSpecification
     ) -> tuple[list, list]:
         """Generate the gene_types and gene_space for a given fitness function and specification
-        :param fitness_function: Instance of SearchFitnessFunction
+        :param surrogate_model: Instance of a CubicSplineRegressionEstimator
         :param specification: The Causal Specification (combination of Scenario and Causal Dag)"""
 
         var_space = {}
-        var_space[fitness_function.surrogate_model.treatment] = {}
-        for adj in fitness_function.surrogate_model.adjustment_set:
+        var_space[surrogate_model.treatment] = {}
+        for adj in surrogate_model.adjustment_set:
             var_space[adj] = {}
 
         for relationship in list(specification.scenario.constraints):
@@ -112,12 +102,12 @@ def create_gene_types(
                     var_space[rel_split[0]]["high"] = int(rel_split[2])
 
         gene_space = []
-        gene_space.append(var_space[fitness_function.surrogate_model.treatment])
-        for adj in fitness_function.surrogate_model.adjustment_set:
+        gene_space.append(var_space[surrogate_model.treatment])
+        for adj in surrogate_model.adjustment_set:
             gene_space.append(var_space[adj])
 
         gene_types = []
-        gene_types.append(specification.scenario.variables.get(fitness_function.surrogate_model.treatment).datatype)
-        for adj in fitness_function.surrogate_model.adjustment_set:
+        gene_types.append(specification.scenario.variables.get(surrogate_model.treatment).datatype)
+        for adj in surrogate_model.adjustment_set:
             gene_types.append(specification.scenario.variables.get(adj).datatype)
         return gene_types, gene_space

tests/surrogate_tests/test_causal_surrogate_assisted.py

@@ -4,7 +4,7 @@
 from causal_testing.specification.causal_specification import CausalSpecification
 from causal_testing.specification.scenario import Scenario
 from causal_testing.specification.variable import Input
-from causal_testing.surrogate.causal_surrogate_assisted import SearchAlgorithm, SimulationResult, SearchFitnessFunction, CausalSurrogateAssistedTestCase, Simulator
+from causal_testing.surrogate.causal_surrogate_assisted import SimulationResult, CausalSurrogateAssistedTestCase, Simulator
 from causal_testing.surrogate.surrogate_search_algorithms import GeneticSearchAlgorithm
 from causal_testing.testing.estimators import CubicSplineRegressionEstimator
 from tests.test_helpers import create_temp_dir_if_non_existent, remove_temp_dir_if_existent
@@ -36,7 +36,7 @@ def test_inputs(self):
 
                     self.assertEqual(result.relationship, relationship)
 
-class TestSearchFitnessFunction(unittest.TestCase):
+class TestCausalSurrogate(unittest.TestCase):
 
     @classmethod
     def setUpClass(cls) -> None:
@@ -49,17 +49,6 @@ def setUp(self):
         with open(self.dag_dot_path, "w") as f:
             f.write(dag_dot)
 
-    def test_init_valid_values(self):
-
-        test_function = lambda x: x **2
-
-        surrogate_model = CubicSplineRegressionEstimator("", 0, 0, set(), "", 4)
-
-        search_function = SearchFitnessFunction(fitness_function=test_function, surrogate_model=surrogate_model)
-
-        self.assertTrue(callable(search_function.fitness_function))
-        self.assertIsInstance(search_function.surrogate_model, CubicSplineRegressionEstimator)
-
     def test_surrogate_model_generation(self):
         c_s_a_test_case = CausalSurrogateAssistedTestCase(None, None, None)