diff --git a/causal_testing/__main__.py b/causal_testing/__main__.py index 76f9e6be..2ef700c7 100644 --- a/causal_testing/__main__.py +++ b/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 diff --git a/causal_testing/discovery/abstract_discovery.py b/causal_testing/discovery/abstract_discovery.py index c0b51387..0ba89871 100644 --- a/causal_testing/discovery/abstract_discovery.py +++ b/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) diff --git a/causal_testing/discovery/hill_climber_discovery.py b/causal_testing/discovery/hill_climber_discovery.py index 3faf506a..882bb1a6 100644 --- a/causal_testing/discovery/hill_climber_discovery.py +++ b/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, ) diff --git a/causal_testing/testing/causal_effect.py b/causal_testing/testing/causal_effect.py index a99a3bd1..2775faba 100644 --- a/causal_testing/testing/causal_effect.py +++ b/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): diff --git a/causal_testing/testing/causal_test_result.py b/causal_testing/testing/causal_test_result.py index 5bfb15c1..bc3363e2 100644 --- a/causal_testing/testing/causal_test_result.py +++ b/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" diff --git a/causal_testing/visualisation/causal_test_result_visualiser.py b/causal_testing/visualisation/causal_test_result_visualiser.py new file mode 100644 index 00000000..b836755d --- /dev/null +++ b/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 diff --git a/tests/discovery_tests/test_abstract_discovery.py b/tests/discovery_tests/test_abstract_discovery.py index 049c72fd..b600deca 100644 --- a/tests/discovery_tests/test_abstract_discovery.py +++ b/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) diff --git a/tests/discovery_tests/test_hill_climber_discovery.py b/tests/discovery_tests/test_hill_climber_discovery.py index e93600d3..47fe5ca6 100644 --- a/tests/discovery_tests/test_hill_climber_discovery.py +++ b/tests/discovery_tests/test_hill_climber_discovery.py @@ -6,7 +6,6 @@ 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 diff --git a/tests/testing_tests/test_causal_test_result.py b/tests/testing_tests/test_causal_test_result.py new file mode 100644 index 00000000..6490868c --- /dev/null +++ b/tests/testing_tests/test_causal_test_result.py @@ -0,0 +1,50 @@ +""" +Test the CausalTestResult class. +""" + +import unittest + +import pandas as pd + +from causal_testing.estimation.effect_estimate import EffectEstimate +from causal_testing.testing.causal_test_result import CausalTestResult + + +class TestCausalTestCase(unittest.TestCase): + + def test_effect_direction_positive(self): + 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(result.effect_direction(), "positive") + + def test_effect_direction_negative(self): + 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(result.effect_direction(), "negative") + + def test_effect_direction_none(self): + 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(result.effect_direction(), "no effect") + + def test_effect_direction_categorical(self): + result = CausalTestResult( + outcome=None, + effect_estimate=EffectEstimate( + type="ate", + value=pd.Series({"color[T.RED]": -5, "color[T.BLUE]": -4}), + ci_low=pd.Series({"color[T.RED]": -4, "color[T.BLUE]": -1}), + ci_high=pd.Series({"color[T.RED]": 5, "color[T.BLUE]": 4}), + ), + ) + self.assertEqual(result.effect_direction(), "categorical") diff --git a/tests/visualisation_tests/test_causal_test_result_visualiser.py b/tests/visualisation_tests/test_causal_test_result_visualiser.py new file mode 100644 index 00000000..0f8ab851 --- /dev/null +++ b/tests/visualisation_tests/test_causal_test_result_visualiser.py @@ -0,0 +1,40 @@ +import os +import unittest +from itertools import cycle +from tempfile import TemporaryDirectory + +import pandas as pd + +from causal_testing.estimation.effect_estimate import EffectEstimate +from causal_testing.specification.causal_dag import CausalDAG +from causal_testing.testing.causal_test_result import CausalTestResult, TestOutcome +from causal_testing.visualisation.causal_test_result_visualiser import results_dag + + +class TestVisualiser(unittest.TestCase): + def test_results_dag(self): + dag = CausalDAG() + dag.add_edges_from([("A", "B"), ("C", "D"), ("E", "F")]) + dag.datatypes = {node: float for node in dag.nodes} + test_cases = dag.generate_causal_tests() + + test_result_cycle = cycle([TestOutcome.PASS, TestOutcome.FAIL, TestOutcome.INESTIMABLE]) + effect_estimate_cycle = cycle( + [ + EffectEstimate(type="ate", value=pd.Series(5), ci_low=pd.Series(4), ci_high=pd.Series(6)), # Positive + EffectEstimate(type="ate", value=pd.Series(5), ci_low=pd.Series(-4), ci_high=pd.Series(6)), # No effect + EffectEstimate( + type="ate", value=pd.Series(-5), ci_low=pd.Series(-6), ci_high=pd.Series(-4) + ), # Negative + ] + ) + for test in test_cases: + test.result = CausalTestResult( + effect_estimate=next(effect_estimate_cycle), + outcome=next(test_result_cycle), + ) + + with TemporaryDirectory() as tmp: + results_dag(test_cases=test_cases, dag=dag, output_file=os.path.join(tmp, "dag.dot")) + dag2 = CausalDAG(os.path.join(tmp, "dag.dot"), ignore_cycles=True) + self.assertEqual(dag.nodes, dag2.nodes)