Added the option for an initial individual

.github/PULL_REQUEST_TEMPLATE.md

@@ -1,3 +1,9 @@
+## Summary
+What does this PR do?
+
+## Main file changes
+Summarise changes to main files to be reviewed.
+
 ## Checklist
 Before you mark your PR as ready for review, please ensure you have completed the following.
 

causal_testing/discovery/hill_climber_discovery.py

@@ -4,6 +4,7 @@
 
 import random
 
+import networkx as nx
 import numpy as np
 import pandas as pd
 from tqdm import tqdm
@@ -26,41 +27,19 @@ def __init__(
         include_edges: str = None,
         exclude_edges: str = None,
         alpha: float = 0.05,
-        max_iterations: int = 100,
+        max_iterations: int = 30,
         max_iterations_without_improvement: int = 10,
     ):
         super().__init__(
-            df=df, random_seed=random_seed, include_edges=include_edges, exclude_edges=exclude_edges, alpha=alpha
+            df=df,
+            random_seed=random_seed,
+            include_edges=include_edges,
+            exclude_edges=exclude_edges,
+            alpha=alpha,
         )
         self.max_iterations = int(max_iterations)
         self.max_iterations_without_improvement = int(max_iterations_without_improvement)
 
-    def sum_test_outcomes(self, test_results: pd.DataFrame) -> dict:
-        """
-        Aggregate the number of passing, failing, and inestimable tests
-        :param test_results: Dataframe containing the raw pass/fail/inestimable outcome of each test case.
-        :returns: Dictionary containing the number of pass/fail/inestimable outcomes.
-        """
-        counts = pd.concat(
-            [
-                pd.DataFrame(np.sort(test_results[["treatment", "outcome"]], axis=1), columns=["treatment", "outcome"]),
-                pd.get_dummies(test_results["result"]).astype(int),
-            ],
-            axis=1,
-        )
-        # Ensure every column is initialised - Test outcomes that never occurred won't be in the dataframe otherwise
-        for col in TestOutcome:
-            if col not in counts.columns:
-                counts[col] = 0
-        counts = counts.groupby(["treatment", "outcome"]).sum().reset_index()[list(TestOutcome)]
-        # The below line normalises by the number of tests *for each edge*
-        # Independence tests X _||_ Y get two tests (X _||_ Y and Y _||_ X) because we don't know which way the
-        # causality flows. We need to normalise this (e.g. if X _||_ Y and Y _||_ X both pass, then the score should be
-        # 1 rather than 2) otherwise we end up unintentionally optimising for more independences.
-        counts = counts.apply(lambda col: col / counts.sum(axis=1))
-
-        return counts.sum(axis=0).to_dict()
-
     def evaluate_fitness(
         self,
         individual: CausalDAG,
@@ -76,7 +55,6 @@ def evaluate_fitness(
                   inestimable tests respectively, and Y is a list of failing edges.
         """
         self.evaluate_tests(individual)
-        counts = self.sum_test_outcomes(individual.test_results)
 
         # Add extra "var1" and "var2" columns to serve as order independent "treatment" and "outcome"
         query_df = pd.concat(
@@ -94,25 +72,38 @@ def evaluate_fitness(
             or ~(group["result"] == TestOutcome.PASS).any()
         )
         problem_edges = problem_tests[["treatment", "outcome"]].apply(tuple, axis=1).tolist()
-        num_tests = sum(counts.values())
+
+        counts = {key: len(group) for key, group in query_df.groupby(["result", "expected_effect"], sort=False)}
+
+        no_effect_normalisation = len(list(nx.non_edges(individual))) or 1
+        some_effect_normalisation = len(individual.edges) or 1
 
         fitness_values = (
-            counts.get(TestOutcome.PASS, 0) / num_tests,
-            -counts.get(TestOutcome.FAIL, 0) / num_tests,
-            -counts.get(TestOutcome.INESTIMABLE, 0) / num_tests,
+            (counts.get((TestOutcome.PASS, "NoEffect"), 0)) / no_effect_normalisation,
+            -(counts.get((TestOutcome.FAIL, "NoEffect"), 0)) / no_effect_normalisation,
+            (counts.get((TestOutcome.PASS, "SomeEffect"), 0)) / some_effect_normalisation,
+            -(counts.get((TestOutcome.FAIL, "SomeEffect"), 0)) / some_effect_normalisation,
+            -(counts.get((TestOutcome.INESTIMABLE, "NoEffect"), 0)) / no_effect_normalisation,
+            -(counts.get((TestOutcome.INESTIMABLE, "SomeEffect"), 0)) / some_effect_normalisation,
         )
         return fitness_values, problem_edges
 
-    def discover(self) -> CausalDAG:
+    def discover(self, individual: CausalDAG = None) -> CausalDAG:
         """
         Discover the causal DAG.
 
+        :param individual: An initial individual for the hill climber to start from
+                           (defaults to a fully connected graph).
+
         :returns: The inferred causal DAG.
         """
 
-        individual = CausalDAG(ignore_cycles=True)
-        individual.add_nodes_from(self.df.columns)
-        individual.add_edges_from(self.possible_edges)
+        if individual is None:
+            individual = CausalDAG(ignore_cycles=True)
+            individual.add_nodes_from(self.df.columns)
+            for treatment, outcome in self.possible_edges:
+                if (treatment, outcome) in self.include_edges:
+                    individual.add_edge(treatment, outcome)
         self.remove_cycles(individual)
         fitness_values, problem_edges = self.evaluate_fitness(individual)
 

causal_testing/estimation/linear_regression_estimator.py

@@ -155,7 +155,7 @@ def estimate_ate_calculated(self, df: pd.DataFrame) -> EffectEstimate:
         return EffectEstimate("ate", pd.Series(treatment_outcome["mean"] - control_outcome["mean"]), ci_low, ci_high)
 
     def _get_confidence_intervals(self, model, treatment):
-        confidence_intervals = model.conf_int(alpha=self.alpha, cols=None)
+        confidence_intervals = model.conf_int(alpha=self.alpha)
         ci_low, ci_high = (
             pd.Series(confidence_intervals[0].loc[treatment]),
             pd.Series(confidence_intervals[1].loc[treatment]),

pyproject.toml

@@ -21,7 +21,7 @@ dependencies = [
     "pandas>=2.1",
     "scikit_learn~=1.4",
     "scipy>=1.12.0,<=1.17.1",
-    "statsmodels~=0.14",
+    "statsmodels~=0.15",
     "tabulate~=0.9",
     "pydot>=2.0",
     "pygad~=3.3",

tests/discovery_tests/test_abstract_discovery.py

@@ -274,7 +274,7 @@ def test_evaluate_tests_inestimable(self):
                     "outcome": "completed",
                 },
                 {
-                    "result": TestOutcome.INESTIMABLE,
+                    "result": TestOutcome.PASS,
                     "expected_effect": "NoEffect",
                     "treatment": "color",
                     "outcome": "completed",

tests/discovery_tests/test_hill_climber_discovery.py

@@ -6,101 +6,12 @@
 
 import pandas as pd
 
-from causal_testing.discovery.abstract_discovery import simple_cycle
 from causal_testing.discovery.hill_climber_discovery import HillClimberDiscovery
 from causal_testing.specification.causal_dag import CausalDAG
-from causal_testing.testing.causal_test_result import TestOutcome
 
 
 class TestHillClimber(unittest.TestCase):
 
-    def test_sum_test_outcomes(self):
-        test_results = pd.DataFrame(
-            [
-                {
-                    "result": TestOutcome.PASS,
-                    "expected_effect": "NoEffect",
-                    "treatment": "length_in",
-                    "outcome": "large_gauge",
-                    "effect": "positive",
-                },
-                {
-                    "result": TestOutcome.INESTIMABLE,
-                    "expected_effect": "NoEffect",
-                    "treatment": "large_gauge",
-                    "outcome": "length_in",
-                    "effect": None,
-                },
-                {
-                    "result": TestOutcome.INESTIMABLE,
-                    "expected_effect": "NoEffect",
-                    "treatment": "length_in",
-                    "outcome": "color",
-                    "effect": None,
-                },
-                {
-                    "result": TestOutcome.INESTIMABLE,
-                    "expected_effect": "NoEffect",
-                    "treatment": "color",
-                    "outcome": "length_in",
-                    "effect": None,
-                },
-                {
-                    "result": TestOutcome.FAIL,
-                    "expected_effect": "SomeEffect",
-                    "treatment": "length_in",
-                    "outcome": "completed",
-                    "effect": "negative",
-                },
-                {
-                    "result": TestOutcome.INESTIMABLE,
-                    "expected_effect": "NoEffect",
-                    "treatment": "large_gauge",
-                    "outcome": "color",
-                    "effect": None,
-                },
-                {
-                    "result": TestOutcome.PASS,
-                    "expected_effect": "NoEffect",
-                    "treatment": "color",
-                    "outcome": "large_gauge",
-                    "effect": None,
-                },
-                {
-                    "result": TestOutcome.FAIL,
-                    "expected_effect": "SomeEffect",
-                    "treatment": "large_gauge",
-                    "outcome": "completed",
-                    "effect": "positive",
-                },
-                {
-                    "result": TestOutcome.PASS,
-                    "expected_effect": "NoEffect",
-                    "treatment": "color",
-                    "outcome": "completed",
-                    "effect": None,
-                },
-                {
-                    "result": TestOutcome.INESTIMABLE,
-                    "expected_effect": "NoEffect",
-                    "treatment": "completed",
-                    "outcome": "color",
-                    "effect": None,
-                },
-            ]
-        )
-        expected_results = {TestOutcome.PASS: 1.5, TestOutcome.FAIL: 2, TestOutcome.INESTIMABLE: 2.5}
-        hill_climber = HillClimberDiscovery(pd.DataFrame())
-        self.assertEqual(expected_results, hill_climber.sum_test_outcomes(test_results))
-
-    def test_sum_test_outcomes_uninitialised(self):
-        hill_climber = HillClimberDiscovery(pd.DataFrame())
-        expected_results = {TestOutcome.PASS: 0, TestOutcome.FAIL: 0, TestOutcome.INESTIMABLE: 0}
-
-        self.assertEqual(
-            expected_results, hill_climber.sum_test_outcomes(pd.DataFrame(columns=["treatment", "outcome", "result"]))
-        )
-
     def test_evaluate_fitness(self):
         scarf_df = pd.read_csv("tests/resources/data/scarf_data.csv")
         dag = CausalDAG()
@@ -109,7 +20,7 @@ def test_evaluate_fitness(self):
 
         hill_climber = HillClimberDiscovery(scarf_df)
         fitness_values, problem_edges = hill_climber.evaluate_fitness(dag)
-        expected_fitness_values = (4 / 6, -2 / 6, 0)
+        expected_fitness_values = (0.8, 0.0, 0.0, -1.0, 0.0, 0.0)
         expected_problem_edges = [
             ("length_in", "completed"),
             ("large_gauge", "completed"),

tests/main_tests/test_ctf.py

@@ -197,14 +197,14 @@ def test_ctf_evaluate_dag_inestimable(self):
         expected = pd.Series(
             {
                 "FAIL": 1,
-                "FAIL_ci_high": 2,
+                "FAIL_ci_high": 1,
                 "FAIL_ci_low": 0,
-                "INESTIMABLE": 1,
-                "INESTIMABLE_ci_high": 1,
+                "INESTIMABLE": 0,
+                "INESTIMABLE_ci_high": 0,
                 "INESTIMABLE_ci_low": 0,
-                "PASS": 4,
-                "PASS_ci_high": 4,
-                "PASS_ci_low": 0,
+                "PASS": 5,
+                "PASS_ci_high": 5,
+                "PASS_ci_low": 2,
             }
         ).sort_index()
         pd.testing.assert_series_equal(results, expected)