Jmafoster1/visualisation

causal_testing/__main__.py

@@ -13,6 +13,7 @@
 
 from causal_testing.causal_testing_framework import CausalTestingFramework, read_dataframe
 from causal_testing.specification.causal_dag import CausalDAG
+from causal_testing.visualisation.causal_test_result_visualiser import results_dag
 
 logger = logging.getLogger(__name__)
 
@@ -237,7 +238,7 @@ def main() -> None:
                 **kwargs,
             )
             evolved_dag = discover.discover()
-            discover.write_dot(evolved_dag, args.output)
+            results_dag(test_cases=evolved_dag.test_cases, dag=evolved_dag, output_file=args.output)
             logging.info("Causal structure discovery completed successfully.")
         case Command.TEST:
             # Create and setup framework

causal_testing/discovery/abstract_discovery.py

@@ -8,15 +8,12 @@
 from abc import ABC, abstractmethod
 from itertools import permutations
 
-import networkx as nx
 import numpy as np
 import pandas as pd
 import rustworkx as rx
 
 from causal_testing.causal_testing_framework import CausalTestingFramework
 from causal_testing.specification.causal_dag import CausalDAG
-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
 
 # Ignore warnings from statsmodels when we try to evaluate test cases
@@ -97,22 +94,6 @@ def discover(self) -> CausalDAG:
         :returns: The inferred causal DAG.
         """
 
-    def effect_direction(self, test_case: CausalTestCase) -> str:
-        """
-        Check whether the estimated causal effect is negative or positive.
-
-        :param test_case: The causal test case.
-        :returns: Whether the estimated causal test is positive or negative (or no effect).
-        """
-        if pd.api.types.is_numeric_dtype(self.df[test_case.treatment_variable]) and pd.api.types.is_numeric_dtype(
-            self.df[test_case.outcome_variable]
-        ):
-            if Negative().apply(test_case.result.effect_estimate):
-                return "negative"
-            if Positive().apply(test_case.result.effect_estimate):
-                return "positive"
-        return None
-
     def remove_cycles(self, causal_dag: CausalDAG):
         """
         Remove cycles from individuals by iteratively deleting a random edge from each cycle until there are no more
@@ -130,52 +111,6 @@ def remove_cycles(self, causal_dag: CausalDAG):
             cycle = simple_cycle(causal_dag)
         causal_dag.add_nodes_from(nodes)
 
-    def write_dot(self, individual: CausalDAG, output_file: str):
-        """
-        Write the given individual to the given output file.
-
-        :param individual: The causal DAG to output.
-        :param output_file: The name of the file to write to.
-        """
-        if hasattr(individual, "test_results"):
-            for _, test in individual.test_results.iterrows():
-                if (test["treatment"], test["outcome"]) in individual.edges:
-                    individual[test["treatment"]][test["outcome"]]["label"] = test["effect"]
-
-                    print(test)
-
-                    if test["result"] == TestOutcome.PASS:
-                        print("  GREEN")
-                        individual[test["treatment"]][test["outcome"]]["color"] = "green"
-                        individual[test["treatment"]][test["outcome"]]["fontcolor"] = "green"
-                    elif test["result"] == TestOutcome.INESTIMABLE:
-                        print("  ORANGE")
-                        individual[test["treatment"]][test["outcome"]]["color"] = "orange"
-                        individual[test["treatment"]][test["outcome"]]["fontcolor"] = "orange"
-                    elif test["result"] == TestOutcome.FAIL:
-                        print("  RED")
-                        individual[test["treatment"]][test["outcome"]]["color"] = "red"
-                        individual[test["treatment"]][test["outcome"]]["fontcolor"] = "red"
-                    else:
-                        raise ValueError(f"Invalid test outcome {test['result']}")
-                else:
-                    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"] == TestOutcome.PASS:
-                        individual[test["treatment"]][test["outcome"]]["style"] = "invis"
-                        individual[test["treatment"]][test["outcome"]]["constraint"] = False
-                    elif test["result"] == TestOutcome.INESTIMABLE:
-                        individual[test["treatment"]][test["outcome"]]["color"] = "orange"
-                        individual[test["treatment"]][test["outcome"]]["fontcolor"] = "orange"
-                    elif test["result"] == TestOutcome.FAIL:
-                        individual[test["treatment"]][test["outcome"]]["color"] = "red"
-                        individual[test["treatment"]][test["outcome"]]["fontcolor"] = "red"
-                    else:
-                        raise ValueError(f"Invalid test outcome {test['result']}")
-
-        nx.drawing.nx_pydot.write_dot(individual, output_file)
-
     def evaluate_tests(self, causal_dag: CausalDAG) -> pd.DataFrame:
         """
         Generate and evaluate causal test cases from the supplied CausalDAG and return a list of edges for which the
@@ -190,6 +125,7 @@ def evaluate_tests(self, causal_dag: CausalDAG) -> pd.DataFrame:
         ctf = CausalTestingFramework(dag=causal_dag, df=self.df)
         causal_dag.datatypes = self.df.dtypes
         ctf.test_cases = causal_dag.generate_causal_tests()
+        causal_dag.test_cases = ctf.test_cases
 
         results = []
 
@@ -206,7 +142,6 @@ def evaluate_tests(self, causal_dag: CausalDAG) -> pd.DataFrame:
                         "expected_effect": test_case.expected_causal_effect.__class__.__name__,
                         "treatment": test_case.treatment_variable,
                         "outcome": test_case.outcome_variable,
-                        "effect": self.effect_direction(test_case),
                     }
                 )
             except np.linalg.LinAlgError:
@@ -219,7 +154,4 @@ def evaluate_tests(self, causal_dag: CausalDAG) -> pd.DataFrame:
                     }
                 )
 
-        causal_dag.test_results = pd.DataFrame(results)
-
-        results = pd.DataFrame(results)
         return pd.DataFrame(results)

causal_testing/discovery/hill_climber_discovery.py

@@ -75,16 +75,14 @@ def evaluate_fitness(
         :returns: Tuple of the form (X, Y), where X is a triple containing the number of passing, failing, and
                   inestimable tests respectively, and Y is a list of failing edges.
         """
-        self.evaluate_tests(individual)
-        counts = self.sum_test_outcomes(individual.test_results)
+        test_results = self.evaluate_tests(individual)
+        counts = self.sum_test_outcomes(test_results)
 
         # Add extra "var1" and "var2" columns to serve as order independent "treatment" and "outcome"
         query_df = pd.concat(
             [
-                individual.test_results,
-                pd.DataFrame(
-                    np.sort(individual.test_results[["treatment", "outcome"]], axis=1), columns=["var1", "var2"]
-                ),
+                test_results,
+                pd.DataFrame(np.sort(test_results[["treatment", "outcome"]], axis=1), columns=["var1", "var2"]),
             ],
             axis=1,
         )

causal_testing/testing/causal_effect.py

@@ -45,7 +45,7 @@ def apply(self, effect_estimate: EffectEstimate) -> bool:
         else:
             raise ValueError(f"Test Value type {effect_estimate.type} is not valid for this CausalEffect")
 
-        return (~((effect_estimate.ci_low <= value_to_check) & (value_to_check <= effect_estimate.ci_high))).all()
+        return (~((effect_estimate.ci_low <= value_to_check) & (value_to_check <= effect_estimate.ci_high))).any()
 
 
 class NoEffect(CausalEffect):

causal_testing/testing/causal_test_result.py

@@ -4,6 +4,7 @@
 from enum import Enum
 
 from causal_testing.estimation.effect_estimate import EffectEstimate
+from causal_testing.testing.causal_effect import Negative, Positive
 
 TestOutcome = Enum("TestOutcome", [("PASS", 2), ("FAIL", 0), ("INESTIMABLE", 1)])
 
@@ -50,3 +51,18 @@ def to_dict(self):
         adequacy = self.adequacy.to_dict() if self.adequacy else {}
 
         return outcome | effect_estimate | {"adequacy": adequacy}
+
+    def effect_direction(self) -> str:
+        """
+        Check whether the estimated causal effect is negative or positive.
+
+        :returns: Whether the estimated causal effect is positive or negative (or no effect).
+        """
+        if len(self.effect_estimate.value) > 1:
+            # Don't bother checking categorical estimates since they're not numeric
+            return "categorical"
+        if Negative().apply(self.effect_estimate):
+            return "negative"
+        if Positive().apply(self.effect_estimate):
+            return "positive"
+        return "no effect"

causal_testing/visualisation/causal_test_result_visualiser.py

@@ -0,0 +1,69 @@
+"""
+This module implements a set of utilities to help visualise causal test results.
+"""
+
+import networkx as nx
+import pandas as pd
+
+from causal_testing.specification.causal_dag import CausalDAG
+from causal_testing.testing.causal_test_case import CausalTestCase
+from causal_testing.testing.causal_test_result import TestOutcome
+
+
+def results_dag(
+    test_cases: list[CausalTestCase],
+    dag: CausalDAG,
+    output_file: str = None,
+    view_independences: bool = True,
+    colours: dict[TestOutcome, str] = None,
+) -> CausalDAG:
+    """
+    View causal test results as a graph.
+
+    :param test_cases: The causal tests (with results).
+    :param dag: The causal DAG that corresponds to the test results.
+    :param output_file: The name of the file to write to.
+    :param view_independences: Whether to display failed independence tests.
+    :param colours: Optional dictionary of colours to display the test outcomes.
+                    By default, pass=green, fail=red, inestimable=orange.
+    """
+    default_colours = {TestOutcome.PASS: "green", TestOutcome.INESTIMABLE: "orange", TestOutcome.FAIL: "red"}
+
+    if colours is not None:
+        colours = default_colours | colours
+    else:
+        colours = default_colours
+
+    result_dag = nx.DiGraph()
+    result_dag.add_nodes_from(dag.nodes)
+    result_dag.add_edges_from(dag.edges)
+
+    for test in test_cases:
+        if test.result:
+            effect_estimate = pd.concat(
+                [
+                    test.result.effect_estimate.ci_low,
+                    test.result.effect_estimate.value,
+                    test.result.effect_estimate.ci_high,
+                ],
+                axis=1,
+            )
+            effect_estimate.columns = ["ci_low", "estimate", "ci_high"]
+            if (test.treatment_variable, test.outcome_variable) in result_dag.edges:
+                result_dag[test.treatment_variable][test.outcome_variable]["label"] = test.result.effect_direction()
+                result_dag[test.treatment_variable][test.outcome_variable]["color"] = colours[test.result.outcome]
+                result_dag[test.treatment_variable][test.outcome_variable]["fontcolor"] = colours[test.result.outcome]
+                result_dag[test.treatment_variable][test.outcome_variable]["title"] = effect_estimate.to_html()
+
+            elif view_independences and test.result.outcome != TestOutcome.PASS:
+                result_dag.add_edge(test.treatment_variable, test.outcome_variable, ignore_cycles=True)
+                result_dag[test.treatment_variable][test.outcome_variable]["style"] = "dashed"
+                result_dag[test.treatment_variable][test.outcome_variable]["label"] = test.result.effect_direction()
+                result_dag[test.treatment_variable][test.outcome_variable]["color"] = colours[test.result.outcome]
+                result_dag[test.treatment_variable][test.outcome_variable]["fontcolor"] = colours[test.result.outcome]
+                result_dag[test.treatment_variable][test.outcome_variable]["title"] = effect_estimate.to_html()
+
+    if output_file is not None:
+        nx.drawing.nx_pydot.write_dot(result_dag, output_file)
+
+    return result_dag

tests/discovery_tests/test_abstract_discovery.py

@@ -2,19 +2,14 @@
 This module tests common causal discovery functionality provided within the abstract_discovery module.
 """
 
-import os
 import unittest
-from tempfile import TemporaryDirectory
 
 import pandas as pd
-from numpy import nan
 
 from causal_testing.discovery.abstract_discovery import Discovery, simple_cycle
-from causal_testing.estimation.effect_estimate import EffectEstimate
 from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator
 from causal_testing.specification.causal_dag import CausalDAG
-from causal_testing.testing.causal_test_case import CausalTestCase
-from causal_testing.testing.causal_test_result import CausalTestResult, TestOutcome
+from causal_testing.testing.causal_test_result import TestOutcome
 
 
 class AbstractDiscovery(Discovery):
@@ -53,46 +48,6 @@ def test_simple_cycle_no_cycles(self):
         dag.add_edges_from([("A", "B"), ("B", "C")])
         self.assertEqual(simple_cycle(dag), [])
 
-    def test_effect_direction_positive(self):
-        causal_test_case = CausalTestCase(
-            estimator=LinearRegressionEstimator(treatment_variable="A", outcome_variable="B", adjustment_set=set()),
-            effect_measure="ate",
-            expected_causal_effect=None,
-        )
-        causal_test_case.result = CausalTestResult(
-            outcome=None,
-            effect_estimate=EffectEstimate(
-                type="ate", value=pd.Series(5.05), ci_low=pd.Series(5), ci_high=pd.Series(6)
-            ),
-        )
-        self.assertEqual(self.abstract_discovery.effect_direction(causal_test_case), "positive")
-
-    def test_effect_direction_negative(self):
-        causal_test_case = CausalTestCase(
-            estimator=LinearRegressionEstimator(treatment_variable="A", outcome_variable="B", adjustment_set=set()),
-            expected_causal_effect=None,
-            effect_measure="ate",
-        )
-        causal_test_case.result = CausalTestResult(
-            outcome=None,
-            effect_estimate=EffectEstimate(
-                type="ate", value=pd.Series(-5.05), ci_low=pd.Series(-6), ci_high=pd.Series(-5)
-            ),
-        )
-        self.assertEqual(self.abstract_discovery.effect_direction(causal_test_case), "negative")
-
-    def test_effect_direction_none(self):
-        causal_test_case = CausalTestCase(
-            estimator=LinearRegressionEstimator(treatment_variable="A", outcome_variable="B", adjustment_set=set()),
-            effect_measure="ate",
-            expected_causal_effect=None,
-        )
-        causal_test_case.result = CausalTestResult(
-            outcome=None,
-            effect_estimate=EffectEstimate(type="ate", value=pd.Series(0), ci_low=pd.Series(-1), ci_high=pd.Series(1)),
-        )
-        self.assertEqual(self.abstract_discovery.effect_direction(causal_test_case), None)
-
     def test_include_edge_wildcard(self):
         abstract_discovery = AbstractDiscovery(
             df=pd.DataFrame(columns=["x_1", "x_2", "x_3", "y_1", "y_2", "y_3", "z_1", "z_2"]),
@@ -155,50 +110,6 @@ def test_remove_cycles_multiple_cycles(self):
         self.assertTrue(dag.has_edge("A", "B") or dag.has_edge("B", "A"))
         self.assertTrue(dag.has_edge("C", "D") or dag.has_edge("D", "C"))
 
-    def test_write_dot(self):
-        dag = CausalDAG()
-        dag.add_edges_from([("A", "B"), ("C", "D"), ("E", "F")])
-        dag.test_results = pd.DataFrame(
-            [  # Edges
-                {"treatment": "A", "outcome": "B", "effect": "positive", "result": TestOutcome.PASS},
-                {"treatment": "C", "outcome": "D", "effect": "positive", "result": TestOutcome.FAIL},
-                {"treatment": "E", "outcome": "F", "effect": "None", "result": TestOutcome.INESTIMABLE},
-                # Independences
-                {"treatment": "A", "outcome": "C", "effect": None, "result": TestOutcome.PASS},
-                {"treatment": "A", "outcome": "D", "effect": "negative", "result": TestOutcome.FAIL},
-                {"treatment": "A", "outcome": "E", "effect": None, "result": TestOutcome.INESTIMABLE},
-            ]
-        )
-        abstract_discovery = AbstractDiscovery(pd.DataFrame())
-        with TemporaryDirectory() as tmp:
-            abstract_discovery.write_dot(dag, os.path.join(tmp, "dag.dot"))
-            dag2 = CausalDAG(os.path.join(tmp, "dag.dot"))
-            self.assertEqual(dag.nodes, dag2.nodes)
-
-    def test_write_dot_invalid_edge_outcome(self):
-        dag = CausalDAG()
-        dag.add_edges_from([("A", "B"), ("C", "D"), ("E", "F")])
-        dag.test_results = pd.DataFrame(
-            [  # Edges
-                {"treatment": "A", "outcome": "B", "effect": None, "result": None},
-            ]
-        )
-        abstract_discovery = AbstractDiscovery(pd.DataFrame())
-        with self.assertRaises(ValueError):
-            abstract_discovery.write_dot(dag, "dag.dot")
-
-    def test_write_dot_invalid_independence_outcome(self):
-        dag = CausalDAG()
-        dag.add_edges_from([("A", "B"), ("C", "D"), ("E", "F")])
-        dag.test_results = pd.DataFrame(
-            [  # Edges
-                {"treatment": "A", "outcome": "C", "effect": None, "result": None},
-            ]
-        )
-        abstract_discovery = AbstractDiscovery(pd.DataFrame())
-        with self.assertRaises(ValueError):
-            abstract_discovery.write_dot(dag, "dag.dot")
-
     def test_evaluate_tests_invalid_datatype(self):
         scarf_df = pd.read_csv("tests/resources/data/scarf_data.csv")
         scarf_df["completed"] = pd.to_datetime(["2026-01-01" for _ in range(len(scarf_df))], format="%Y-%m-%d")
@@ -286,7 +197,6 @@ def test_evaluate_tests_inestimable(self):
                 },
             ]
         )
-        expected_results["effect"] = nan
         pd.testing.assert_frame_equal(test_results, expected_results)
 
     def test_evaluate_tests(self):
@@ -362,5 +272,5 @@ def test_evaluate_tests(self):
                 },
             ]
         )
-        expected_results["effect"] = None
+        print(test_results)
         pd.testing.assert_frame_equal(test_results, expected_results)