Merge remote-tracking branch 'origin/surrogateassisted' into surrogateassisted

# Conflicts:
#	tests/surrogate_tests/test_causal_surrogate_assisted.py
#	tests/testing_tests/test_estimators.py

Author: f-allian

causal_testing/specification/causal_dag.py

@@ -531,6 +531,9 @@ def identification(self, base_test_case: BaseTestCase, scenario: Scenario = None
         if scenario is not None:
             minimal_adjustment_sets = self.remove_hidden_adjustment_sets(minimal_adjustment_sets, scenario)
 
+        if len(minimal_adjustment_sets) == 0:
+            return set()
+
         minimal_adjustment_set = min(minimal_adjustment_sets, key=len)
         return minimal_adjustment_set
 

causal_testing/surrogate/causal_surrogate_assisted.py

@@ -1,46 +1,61 @@
+"""Module containing classes to define and run causal surrogate assisted test cases"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Callable
+
 from causal_testing.data_collection.data_collector import ObservationalDataCollector
 from causal_testing.specification.causal_specification import CausalSpecification
 from causal_testing.testing.base_test_case import BaseTestCase
-from causal_testing.testing.estimators import Estimator, PolynomialRegressionEstimator
-
-from dataclasses import dataclass
-from typing import Callable, Any
-from abc import ABC
+from causal_testing.testing.estimators import CubicSplineRegressionEstimator
 
 
 @dataclass
-class SimulationResult(ABC):
+class SimulationResult:
+    """Data class holding the data and result metadata of a simulation"""
+
     data: dict
     fault: bool
     relationship: str
 
 
-@dataclass
-class SearchFitnessFunction(ABC):
-    fitness_function: Any
-    surrogate_model: PolynomialRegressionEstimator
-
+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"""
 
-class SearchAlgorithm:
-    def generate_fitness_functions(self, surrogate_models: list[Estimator]) -> list[SearchFitnessFunction]:
-        pass
+    @abstractmethod
+    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 surrogate_models: The surrogate models to be searched
+        :param specification:  The Causal Specification (combination of Scenario and Causal Dag)"""
 
-    def search(self, fitness_functions: list[SearchFitnessFunction], specification: CausalSpecification) -> list:
-        pass
 
+class Simulator(ABC):
+    """Class to be inherited with Simulator specific functions to start, shutdown and run the simulation with the give
+    config file"""
 
-class Simulator:
+    @abstractmethod
     def startup(self, **kwargs):
-        pass
+        """Function that when run, initialises and opens the Simulator"""
 
+    @abstractmethod
     def shutdown(self, **kwargs):
-        pass
+        """Function to safely exit and shutdown the Simulator"""
 
-    def run_with_config(self, configuration) -> SimulationResult:
-        pass
+    @abstractmethod
+    def run_with_config(self, configuration: dict) -> SimulationResult:
+        """Run the simulator with the given configuration and return the results in the structure of a
+        SimulationResult
+        :param configuration: The configuration required to initialise the Simulation
+        :return: Simulation results in the structure of the SimulationResult data class"""
 
 
 class CausalSurrogateAssistedTestCase:
+    """A class representing a single causal surrogate assisted test case."""
+
     def __init__(
         self,
         specification: CausalSpecification,
@@ -57,27 +72,33 @@ def execute(
         max_executions: int = 200,
         custom_data_aggregator: Callable[[dict, dict], dict] = None,
     ):
+        """For this specific test case, a search algorithm is used to find the most contradictory point in the input
+        space which is, therefore, most likely to indicate incorrect behaviour. This cadidate test case is run against
+        the simulator, checked for faults and the result returned with collected data
+        :param data_collector: An ObservationalDataCollector which gathers data relevant to the specified scenario
+        :param max_executions: Maximum number of simulator executions before exiting the search
+        :param custom_data_aggregator:
+        :return: tuple containing SimulationResult or str, execution number and collected data"""
         data_collector.collect_data()
 
         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, _fitness, 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)
             self.simulator.shutdown()
 
             if custom_data_aggregator is not None:
-                data_collector.data = custom_data_aggregator(data_collector.data, test_result.data)
+                if data_collector.data is not None:
+                    data_collector.data = custom_data_aggregator(data_collector.data, test_result.data)
             else:
                 data_collector.data = data_collector.data.append(test_result.data, ignore_index=True)
 
             if test_result.fault:
                 print(
-                    f"Fault found between {surrogate.treatment} causing {surrogate.outcome}. Contradiction with expected {surrogate.expected_relationship}."
+                    f"Fault found between {surrogate.treatment} causing {surrogate.outcome}. Contradiction with "
+                    f"expected {surrogate.expected_relationship}."
                 )
                 test_result.relationship = (
                     f"{surrogate.treatment} -> {surrogate.outcome} expected {surrogate.expected_relationship}"
@@ -89,7 +110,12 @@ 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
+        :return: A list of surrogate models
+        """
         surrogate_models = []
 
         for u, v in specification.causal_dag.graph.edges:
@@ -101,7 +127,7 @@ def generate_surrogates(
 
                 minimal_adjustment_set = specification.causal_dag.identification(base_test_case, specification.scenario)
 
-                surrogate = PolynomialRegressionEstimator(
+                surrogate = CubicSplineRegressionEstimator(
                     u,
                     0,
                     0,

causal_testing/surrogate/surrogate_search_algorithms.py

@@ -1,12 +1,18 @@
-from causal_testing.specification.causal_specification import CausalSpecification
-from causal_testing.testing.estimators import Estimator, PolynomialRegressionEstimator
-from causal_testing.surrogate.causal_surrogate_assisted import SearchAlgorithm, SearchFitnessFunction
+"""Module containing implementation of search algorithm for surrogate search """
+# pylint: disable=cell-var-from-loop
+# Fitness functions are required to be iteratively defined, including all variables within.
 
-from pygad import GA
 from operator import itemgetter
+from pygad import GA
+
+from causal_testing.specification.causal_specification import CausalSpecification
+from causal_testing.testing.estimators import CubicSplineRegressionEstimator
+from causal_testing.surrogate.causal_surrogate_assisted import SearchAlgorithm
 
 
 class GeneticSearchAlgorithm(SearchAlgorithm):
+    """Implementation of SearchAlgorithm class. Implements genetic search algorithm for surrogate models."""
+
     def __init__(self, delta=0.05, config: dict = None) -> None:
         super().__init__()
 
@@ -15,91 +21,93 @@ def __init__(self, delta=0.05, config: dict = None) -> None:
         self.contradiction_functions = {
             "positive": lambda x: -1 * x,
             "negative": lambda x: x,
-            "no_effect": lambda x: abs(x),
+            "no_effect": abs,
             "some_effect": lambda x: abs(1 / x),
         }
 
-    def generate_fitness_functions(
-        self, surrogate_models: list[PolynomialRegressionEstimator]
-    ) -> 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
-            def fitness_function(_ga, solution, idx):
+            # 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
                 surrogate.treatment_value = solution[0] + self.delta
 
-                adjustment_dict = dict()
+                adjustment_dict = {}
                 for i, adjustment in enumerate(surrogate.adjustment_set):
                     adjustment_dict[adjustment] = solution[i + 1]
 
                 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:
-            var_space = dict()
-            var_space[fitness_function.surrogate_model.treatment] = dict()
-            for adj in fitness_function.surrogate_model.adjustment_set:
-                var_space[adj] = dict()
-
-            for relationship in list(specification.scenario.constraints):
-                rel_split = str(relationship).split(" ")
-
-                if rel_split[1] == ">=":
-                    var_space[rel_split[0]]["low"] = int(rel_split[2])
-                elif rel_split[1] == "<=":
-                    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[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(adj).datatype)
+            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,
             )
 
             if self.config is not None:
                 for k, v in self.config.items():
                     if k == "gene_space":
-                        raise Exception(
-                            "Gene space should not be set through config. This is generated from the causal specification"
+                        raise ValueError(
+                            "Gene space should not be set through config. This is generated from the causal "
+                            "specification"
                         )
                     setattr(ga, k, v)
 
             ga.run()
-            solution, fitness, _idx = ga.best_solution()
+            solution, fitness, _ = ga.best_solution()
 
-            solution_dict = dict()
-            solution_dict[fitness_function.surrogate_model.treatment] = solution[0]
-            for idx, adj in enumerate(fitness_function.surrogate_model.adjustment_set):
+            solution_dict = {}
+            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(
+        surrogate_model: CubicSplineRegressionEstimator, specification: CausalSpecification
+    ) -> tuple[list, list]:
+        """Generate the gene_types and gene_space for a given fitness function and specification
+        :param surrogate_model: Instance of a CubicSplineRegressionEstimator
+        :param specification: The Causal Specification (combination of Scenario and Causal Dag)"""
+
+        var_space = {}
+        var_space[surrogate_model.treatment] = {}
+        for adj in surrogate_model.adjustment_set:
+            var_space[adj] = {}
+
+        for relationship in list(specification.scenario.constraints):
+            rel_split = str(relationship).split(" ")
+
+            if rel_split[0] in var_space:
+                if rel_split[1] == ">=":
+                    var_space[rel_split[0]]["low"] = int(rel_split[2])
+                elif rel_split[1] == "<=":
+                    var_space[rel_split[0]]["high"] = int(rel_split[2])
+
+        gene_space = []
+        gene_space.append(var_space[surrogate_model.treatment])
+        for adj in surrogate_model.adjustment_set:
+            gene_space.append(var_space[adj])
 
-        return max(
-            solutions, key=itemgetter(1)
-        )  # TODO This can be done better with fitness normalisation between edges
+        gene_types = []
+        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