Causal DAG analysis (and restructuring of CausalTestResult)

Author: jmafoster1

causal_testing/__main__.py

@@ -23,6 +23,7 @@ class Command(Enum):
     TEST = "test"
     GENERATE = "generate"
     DISCOVER = "discover"
+    EVALUATE = "evaluate"
 
 
 def setup_logging(level: str) -> None:
@@ -48,6 +49,15 @@ def parse_args(args: Optional[Sequence[str]] = None) -> argparse.Namespace:
         choices=["NONE", "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
         help="Set the logging level (default: WARNING).",
     )
+    main_parser.add_argument(
+        "-a",
+        "--alpha",
+        help=(
+            "The significance level of the confidence intervals used to determine causality. "
+            "This should be a value between 0 and 1. Defaults to 0.05 for 95%% confidence intervals."
+        ),
+        default=0.05,
+    )
 
     subparsers = main_parser.add_subparsers(
         help="The action you want to run - call `causal_testing {action} -h` for further details", dest="command"
@@ -108,18 +118,36 @@ def parse_args(args: Optional[Sequence[str]] = None) -> argparse.Namespace:
         default=False,
     )
 
+    # DAG evaluation
+    parser_evaluate = subparsers.add_parser(
+        Command.EVALUATE.value, help="Evaluate how well a causal DAG fits a dataset"
+    )
+    parser_evaluate.add_argument("-D", "--dag-path", help="Path to the DAG file (.dot)", required=True)
+    parser_evaluate.add_argument("-o", "--output", help="Path for output file (.csv)", required=True)
+    parser_evaluate.add_argument(
+        "-i", "--ignore-cycles", help="Ignore cycles in DAG", action="store_true", default=False
+    )
+    parser_evaluate.add_argument("-d", "--data-paths", help="Paths to data files (.csv)", nargs="+", required=True)
+    parser_evaluate.add_argument("-q", "--query", help="Query string to filter data (e.g. 'age > 18')", type=str)
+    parser_evaluate.add_argument(
+        "-b",
+        "--adequacy-bootstrap-size",
+        dest="bootstrap_size",
+        help="Number of bootstrap samples for causal test adequacy. Defaults to 100",
+        type=int,
+        default=100,
+    )
+    parser_evaluate.add_argument(
+        "-s",
+        "--silent",
+        action="store_true",
+        help="Do not crash on error. If set to true, errors are recorded as test results.",
+        default=False,
+    )
+
     # Discovery
     parser_discover = subparsers.add_parser(Command.DISCOVER.value, help="Discover causal structures from data")
     parser_discover.add_argument("-d", "--data-paths", help="Paths to data files (.csv)", nargs="+", required=True)
-    parser_discover.add_argument(
-        "-a",
-        "--alpha",
-        help=(
-            "The significance level of the confidence intervals used to determine causality. "
-            "This should be a value between 0 and 1. Defaults to 0.05 for 95%% confidence intervals."
-        ),
-        default=0.05,
-    )
     parser_discover.add_argument(
         "-t",
         "--technique",
@@ -246,6 +274,23 @@ def main() -> None:
             framework.save_results(args.output)
 
             logging.info("Causal testing completed successfully.")
+        case Command.EVALUATE:
+            # Create and setup framework
+            framework = CausalTestingFramework()
+
+            framework.setup(
+                dag_path=args.dag_path,
+                data_paths=args.data_paths,
+                test_cases_path=args.test_config,
+                query=args.query,
+                ignore_cycles=args.ignore_cycles,
+            )
+
+            logging.info("Running tests on entire dataset")
+            results = framework.evaluate_dag(alpha=args.alpha, bootstrap_size=args.bootstrap_size)
+            logging.info("Causal testing completed successfully.")
+            logging.info("Running tests on bootstrap samples")
+            results.to_csv(args.output)
 
 
 if __name__ == "__main__":

causal_testing/causal_testing_framework.py

@@ -14,6 +14,7 @@
 from causal_testing.specification.variable import Input, Output
 from causal_testing.testing.base_test_case import BaseTestCase
 from causal_testing.testing.causal_test_case import CausalTestCase
+from causal_testing.testing.causal_test_result import TestOutcome
 
 logger = logging.getLogger(__name__)
 
@@ -265,6 +266,49 @@ def run_tests(self, silent: bool = False, adequacy: bool = False, bootstrap_size
                 self.df, suppress_estimation_errors=silent, adequacy=adequacy, bootstrap_size=bootstrap_size
             )
 
+    def evaluate_dag(self, bootstrap_size: bool, alpha: float) -> pd.Series:
+        """
+        Calculate confidence intervals for how well a causal DAG fits a dataset by repeatedly resampling the dataset
+        and executing the causal tests.
+        Confidence intervals are then calcualted for how many tests pass, fail, or are inestimable.
+
+        :param bootstrap_size: The number of bootstrap samples to use when calculating causal test adequacy
+                               (defaults to 100)
+        :param alpha: The significance level to use when calculating confidence intervals.
+        """
+        self.run_tests(silent=True, adequacy=False)
+        results = {
+            test_outcome: len([test for test in self.test_cases if test.result.outcome == test_outcome])
+            for test_outcome in TestOutcome
+        }
+
+        sample_results = []
+        for sample_index in range(bootstrap_size):
+            test_outcomes = {test_outcome: 0 for test_outcome in TestOutcome}
+            for test_case in tqdm(self.test_cases):
+                effect_estimate = test_case.estimate_effect(
+                    df=self.df.sample(len(self.df), replace=True, random_state=sample_index)
+                )
+                if effect_estimate:
+                    if test_case.expected_causal_effect.apply(effect_estimate):
+                        test_outcomes[TestOutcome.PASS] += 1
+                    else:
+                        test_outcomes[TestOutcome.FAIL] += 1
+                else:
+                    test_outcomes[TestOutcome.INESTIMABLE] += 1
+            sample_results.append(test_outcomes)
+
+        sample_results = pd.DataFrame(sample_results)
+        # Calculate the confidence interval of each column
+        ci_low_inx = (alpha / 2) * bootstrap_size
+        ci_high_inx = ((1 - alpha) / 2) * bootstrap_size
+        for outcome in TestOutcome:
+            data = sorted(sample_results[outcome])
+            results[f"{outcome}_ci_low"] = data[ci_low_inx]
+            results[f"{outcome}_ci_high"] = data[ci_high_inx]
+
+        return pd.Series(results).sort_index()
+
     def save_results(self, output_path) -> list:
         """Save test results to JSON file in the expected format."""
         logger.info(f"Saving results to {output_path}")
@@ -301,15 +345,11 @@ def save_results(self, output_path) -> list:
                 result = test_case.result
                 result_index += 1
 
-                test_passed = (
-                    test_case.expected_causal_effect.apply(result) if result.effect_estimate is not None else False
-                )
-
                 output = {
                     **base_output,
                     "formula": test_case.estimator.formula if hasattr(test_case.estimator, "formula") else None,
                     "skip": False,
-                    "passed": test_passed,
+                    "passed": test_case.result.outcome == TestOutcome.PASS,
                     "result": (
                         {
                             "treatment": test_case.estimator.base_test_case.treatment_variable.name,

causal_testing/discovery/abstract_discovery.py

@@ -6,7 +6,6 @@
 import re
 import warnings
 from abc import ABC, abstractmethod
-from enum import Enum
 from itertools import permutations
 
 import networkx as nx
@@ -19,10 +18,9 @@
 from causal_testing.specification.scenario import Scenario
 from causal_testing.testing.causal_effect import Negative, Positive
 from causal_testing.testing.causal_test_case import CausalTestCase
+from causal_testing.testing.causal_test_result import TestOutcome
 from causal_testing.testing.metamorphic_relation import generate_metamorphic_relations
 
-TestResult = Enum("TestResult", [("PASS", 2), ("FAIL", 0), ("INESTIMABLE", 1)])
-
 # Ignore warnings from statsmodels when we try to evaluate test cases
 warnings.simplefilter("ignore")
 
@@ -111,9 +109,9 @@ def effect_direction(self, test_case: CausalTestCase) -> str:
         if pd.api.types.is_numeric_dtype(
             self.df[test_case.base_test_case.treatment_variable.name]
         ) and pd.api.types.is_numeric_dtype(self.df[test_case.base_test_case.outcome_variable.name]):
-            if Negative().apply(test_case.result):
+            if Negative().apply(test_case.result.effect_estimate):
                 return "negative"
-            if Positive().apply(test_case.result):
+            if Positive().apply(test_case.result.effect_estimate):
                 return "positive"
         return None
 
@@ -148,15 +146,15 @@ def write_dot(self, individual: CausalDAG, output_file: str):
 
                     print(test)
 
-                    if test["result"] == TestResult.PASS:
+                    if test["result"] == TestOutcome.PASS:
                         print("  GREEN")
                         individual[test["treatment"]][test["outcome"]]["color"] = "green"
                         individual[test["treatment"]][test["outcome"]]["fontcolor"] = "green"
-                    elif test["result"] == TestResult.INESTIMABLE:
+                    elif test["result"] == TestOutcome.INESTIMABLE:
                         print("  ORANGE")
                         individual[test["treatment"]][test["outcome"]]["color"] = "orange"
                         individual[test["treatment"]][test["outcome"]]["fontcolor"] = "orange"
-                    elif test["result"] == TestResult.FAIL:
+                    elif test["result"] == TestOutcome.FAIL:
                         print("  RED")
                         individual[test["treatment"]][test["outcome"]]["color"] = "red"
                         individual[test["treatment"]][test["outcome"]]["fontcolor"] = "red"
@@ -166,13 +164,13 @@ def write_dot(self, individual: CausalDAG, output_file: str):
                     individual.add_edge(test["treatment"], test["outcome"], ignore_cycles=True)
                     individual[test["treatment"]][test["outcome"]]["style"] = "dashed"
                     individual[test["treatment"]][test["outcome"]]["label"] = test["effect"]
-                    if test["result"] == TestResult.PASS:
+                    if test["result"] == TestOutcome.PASS:
                         individual[test["treatment"]][test["outcome"]]["style"] = "invis"
                         individual[test["treatment"]][test["outcome"]]["constraint"] = False
-                    elif test["result"] == TestResult.INESTIMABLE:
+                    elif test["result"] == TestOutcome.INESTIMABLE:
                         individual[test["treatment"]][test["outcome"]]["color"] = "orange"
                         individual[test["treatment"]][test["outcome"]]["fontcolor"] = "orange"
-                    elif test["result"] == TestResult.FAIL:
+                    elif test["result"] == TestOutcome.FAIL:
                         individual[test["treatment"]][test["outcome"]]["color"] = "red"
                         individual[test["treatment"]][test["outcome"]]["fontcolor"] = "red"
                     else:
@@ -222,9 +220,9 @@ def evaluate_tests(self, causal_dag: CausalDAG) -> pd.DataFrame:
                 results.append(
                     {
                         "result": (
-                            TestResult.PASS
-                            if test_case.expected_causal_effect.apply(test_case.result)
-                            else TestResult.FAIL
+                            TestOutcome.PASS
+                            if test_case.expected_causal_effect.apply(test_case.result.effect_estimate)
+                            else TestOutcome.FAIL
                         ),
                         "expected_effect": test_case.expected_causal_effect.__class__.__name__,
                         "treatment": test_case.base_test_case.treatment_variable.name,
@@ -235,7 +233,7 @@ def evaluate_tests(self, causal_dag: CausalDAG) -> pd.DataFrame:
             except np.linalg.LinAlgError:
                 results.append(
                     {
-                        "result": TestResult.INESTIMABLE,
+                        "result": TestOutcome.INESTIMABLE,
                         "expected_effect": test_case.expected_causal_effect.__class__.__name__,
                         "treatment": test_case.base_test_case.treatment_variable.name,
                         "outcome": test_case.base_test_case.outcome_variable.name,

causal_testing/discovery/hill_climber_discovery.py

@@ -8,8 +8,9 @@
 import numpy as np
 import pandas as pd
 
-from causal_testing.discovery.abstract_discovery import Discovery, TestResult
+from causal_testing.discovery.abstract_discovery import Discovery
 from causal_testing.specification.causal_dag import CausalDAG
+from causal_testing.testing.causal_test_result import TestOutcome
 
 
 class HillClimberDiscovery(Discovery):
@@ -48,10 +49,10 @@ def sum_test_outcomes(self, test_results: pd.DataFrame) -> dict:
             axis=1,
         )
         # Ensure every column is initialised - Test outcomes that never occurred won't be in the dataframe otherwise
-        for col in TestResult:
+        for col in TestOutcome:
             if col not in counts.columns:
                 counts[col] = 0
-        counts = counts.groupby(["treatment", "outcome"]).sum().reset_index()[list(TestResult)]
+        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
@@ -89,15 +90,15 @@ def evaluate_fitness(
         )
         problem_tests = query_df.groupby(["var1", "var2"]).filter(
             # Groups are problematic if at least one test fails or no test passes
-            lambda group: (group["result"] == TestResult.FAIL).any()
-            or ~(group["result"] == TestResult.PASS).any()
+            lambda group: (group["result"] == TestOutcome.FAIL).any()
+            or ~(group["result"] == TestOutcome.PASS).any()
         )
         problem_edges = problem_tests[["treatment", "outcome"]].apply(tuple, axis=1).tolist()
 
         fitness_values = (
-            counts.get(TestResult.PASS, 0),
-            -counts.get(TestResult.FAIL, 0),
-            -counts.get(TestResult.INESTIMABLE, 0),
+            counts.get(TestOutcome.PASS, 0),
+            -counts.get(TestOutcome.FAIL, 0),
+            -counts.get(TestOutcome.INESTIMABLE, 0),
         )
         return fitness_values, problem_edges