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.visualisation_dashboard import Dashboard
 
 logger = logging.getLogger(__name__)
 
@@ -26,6 +27,7 @@ class Command(Enum):
     GENERATE = "generate"
     DISCOVER = "discover"
     EVALUATE = "evaluate"
+    VISUALISE = "visualise"
 
 
 def setup_logging(level: str) -> None:
@@ -83,6 +85,13 @@ def parse_args(args: Optional[Sequence[str]] = None) -> argparse.Namespace:
         default=False,
     )
 
+    # Visualisation
+    parser_visualise = subparsers.add_parser(Command.VISUALISE.value, help="Visualise causal test results")
+    parser_visualise.add_argument("-D", "--dag-path", help="Path to the DAG file (.dot)", required=True)
+    parser_visualise.add_argument(
+        "-t", "--result-config", help="Path to causal test result file (.json)", required=True
+    )
+
     # DAG evaluation
     parser_evaluate = subparsers.add_parser(
         Command.EVALUATE.value, help="Evaluate how well a causal DAG fits a dataset"
@@ -139,7 +148,7 @@ def parse_args(args: Optional[Sequence[str]] = None) -> argparse.Namespace:
         default=[],
     )
 
-    for parser in [parser_generate, parser_discover, parser_test, parser_evaluate]:
+    for parser in [parser_generate, parser_discover, parser_test, parser_evaluate, parser_visualise]:
         parser.add_argument(
             "-l",
             "--log_level",
@@ -148,6 +157,7 @@ def parse_args(args: Optional[Sequence[str]] = None) -> argparse.Namespace:
             choices=["NONE", "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
             help="Set the logging level (default: WARNING).",
         )
+    for parser in [parser_generate, parser_discover, parser_test, parser_evaluate]:
         parser.add_argument(
             "-a",
             "--alpha",
@@ -194,7 +204,7 @@ def main() -> None:
                 skip=False,
             )
             with open(args.output, "w", encoding="utf-8") as f:
-                json.dump({"tests": [test.to_dict() for test in causal_tests]}, f)
+                json.dump([test.to_dict() for test in causal_tests], f)
             logging.info("Causal test generation completed successfully.")
 
         case Command.DISCOVER:
@@ -237,7 +247,8 @@ def main() -> None:
                 **kwargs,
             )
             evolved_dag = discover.discover()
-            discover.write_dot(evolved_dag, args.output)
+            if args.output is not None:
+                nx.drawing.nx_pydot.write_dot(evolved_dag, args.output)
             logging.info("Causal structure discovery completed successfully.")
         case Command.TEST:
             # Create and setup framework
@@ -256,6 +267,11 @@ def main() -> None:
             framework.save_results(args.output)
 
             logging.info("Causal testing completed successfully.")
+        case Command.VISUALISE:
+            framework = CausalTestingFramework()
+            framework.setup(dag_path=args.dag_path, test_cases_path=args.result_config)
+            dashboard = Dashboard(framework)
+            dashboard.serve()
         case Command.EVALUATE:
             # Create and setup framework
             framework = CausalTestingFramework()

causal_testing/causal_testing_framework.py

@@ -11,40 +11,44 @@
 import pandas as pd
 from tqdm import tqdm
 
+from causal_testing.estimation.effect_estimate import EffectEstimate
 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
+from causal_testing.testing.causal_test_result import CausalTestResult, TestOutcome
+from causal_testing.testing.data_adequacy import DataAdequacy
 
 logger = logging.getLogger(__name__)
 
-
-def read_dataframe(file_path: str, **kwargs: dict) -> pd.DataFrame:
+data_readers = {
+    ".csv": pd.read_csv,
+    ".xlsx": pd.read_excel,
+    ".xls": pd.read_excel,
+    ".html": pd.read_html,
+    ".xml": pd.read_xml,
+    ".feather": pd.read_feather,
+    ".parquet": pd.read_parquet,
+    ".pq": pd.read_parquet,
+    ".pqt": pd.read_parquet,
+    ".json": pd.read_json,
+    ".stata": pd.read_stata,
+}
+
+
+def read_dataframe(file_path: str, content: bytes = None, **kwargs: dict) -> pd.DataFrame:
     """
     Read data into a dataframe.
 
     :param file_path: The path to the data.
+    :param content: The bytes content of the file.
     :param kwargs: Keyword arguments to be passed to the `read_` function.
 
     :returns: The read-in DataFrame.
     """
-    readers = {
-        ".csv": pd.read_csv,
-        ".xlsx": pd.read_excel,
-        ".xls": pd.read_excel,
-        ".html": pd.read_html,
-        ".xml": pd.read_xml,
-        ".feather": pd.read_feather,
-        ".parquet": pd.read_parquet,
-        ".pq": pd.read_parquet,
-        ".pqt": pd.read_parquet,
-        ".json": pd.read_json,
-        ".stata": pd.read_stata,
-    }
 
     suffix = Path(file_path).suffix.lower()
 
-    if suffix in readers:
-        return readers[suffix](file_path, **kwargs)
+    if suffix in data_readers:
+        return data_readers[suffix](content if content is not None else file_path, **kwargs)
     raise ValueError(f"Unsupported file extension: '{suffix}'")
 
 
@@ -61,9 +65,9 @@ def __init__(self, dag: CausalDAG = None, test_cases: list[CausalTestCase] = Non
 
     def setup(
         self,
-        dag_path: str,
-        data_paths: list[str],
-        test_cases_path: str,
+        dag_path: str = None,
+        data_paths: list[str] = None,
+        test_cases_path: str = None,
         ignore_cycles: bool = False,
         query: str = None,
         **kwargs: dict,
@@ -78,9 +82,12 @@ def setup(
         :param query: Optional pandas query string to filter the loaded data
         :param kwargs: Keyword arguments to be passed to the `read_` function.
         """
-        self.load_dag(dag_path, ignore_cycles)
-        self.load_data(data_paths, query, **kwargs)
-        self.load_test_cases_from_json(test_cases_path)
+        if dag_path is not None:
+            self.load_dag(dag_path, ignore_cycles)
+        if data_paths is not None:
+            self.load_data(data_paths, query, **kwargs)
+        if test_cases_path is not None:
+            self.load_test_cases_from_json(test_cases_path)
 
     def load_dag(self, dag_path: str, ignore_cycles: bool = False):
         """
@@ -120,21 +127,13 @@ def load_test_cases_from_json(self, test_cases_path: str):
         """
         logger.info(f"Loading test configurations from {test_cases_path}")
 
-        if self.dag is None or self.df is None:
-            raise ValueError("Please load DAG and data before attempting to load tests.")
+        if self.dag is None:
+            raise ValueError("Please load DAG before attempting to load tests.")
 
         with open(test_cases_path, "r", encoding="utf-8") as f:
             test_configs = json.load(f)
 
-        test_cases = []
-
-        for test in test_configs.get("tests", []):
-
-            # Create causal test case
-            causal_test = self.create_causal_test(test)
-            test_cases.append(causal_test)
-
-        self.test_cases = test_cases
+        self.test_cases = [self.create_causal_test(test) for test in test_configs]
 
     def create_causal_test(self, test: dict) -> CausalTestCase:
         """
@@ -145,54 +144,56 @@ def create_causal_test(self, test: dict) -> CausalTestCase:
         :return: CausalTestCase object
         :raises: ValueError if invalid estimator or configuration is provided
         """
+        # Create the estimator with correct parameters
         estimator_map = {ff.name: ff for ff in entry_points(group="estimators")}
-        effect_map = {ff.name: ff for ff in entry_points(group="causal_effects")}
 
         if "estimator" not in test:
             raise ValueError("Test configuration must specify an estimator")
 
-        if test["estimator"] not in estimator_map:
+        estimator_class = test["estimator"].pop("name")
+        if estimator_class not in estimator_map:
             raise ValueError(
-                f"Unsupported estimator {test['estimator']}. Supported: {sorted(estimator_map)}. "
+                f"Unsupported estimator {estimator_class}. Supported: {sorted(estimator_map)}. "
                 "If you have implemented a custom estimator, you will need to add this to your entrypoints via your "
                 "pyproject.toml file."
             )
 
-        # Create the estimator with correct parameters
-        treatment_variable = test.get("treatment_variable")
-        outcome_variable = test.get("outcome_variable")
-        estimator_class = estimator_map.get(test["estimator"]).load()
-        estimator_kwargs = test.get("estimator_kwargs", {})
-        effect_type = test.get("expected_effect", {}).get("effect_type", "direct")
-
-        estimator = estimator_class(
-            treatment_variable=treatment_variable,
-            outcome_variable=outcome_variable,
-            treatment_value=test.get("treatment_value"),
-            control_value=test.get("control_value"),
-            alpha=test.get("alpha", 0.05),
-            **estimator_kwargs,
-        )
-
-        # Get effect type and create expected effect
-        expected_effect = test["expected_effect"]
-        effect_type = expected_effect.pop("name")
-        if effect_type not in effect_map:
+        estimator_class = estimator_map.get(estimator_class).load()
+        test["estimator"] = estimator_class(**test["estimator"])
+
+        # Create the expected effect with correct parameters
+        effect_map = {ff.name: ff for ff in entry_points(group="causal_effects")}
+
+        if "expected_causal_effect" not in test:
+            raise ValueError("Test configuration must specify an expected causal effect.")
+
+        effect_class = test["expected_causal_effect"].pop("name")
+        if effect_class not in effect_map:
             raise ValueError(
-                f"Unsupported causal effect {effect_type}. Supported: {sorted(effect_map)}. "
+                f"Unsupported causal effect {effect_class}. Supported: {sorted(effect_map)}. "
                 "If you have implemented a custom causal effect, you will need to add this to your entrypoints via "
                 "your pyproject.toml file."
             )
-        expected_effect = effect_map[effect_type].load()(**expected_effect)
+        effect_class = effect_map.get(effect_class).load()
+        test["expected_causal_effect"] = effect_class(**test["expected_causal_effect"])
+
+        if "result" in test:
+            outcome = getattr(TestOutcome, test["result"]["outcome"]) if "outcome" in test["result"] else None
+            effect_estimate = (
+                EffectEstimate(**test["result"]["effect_estimate"]) if "effect_estimate" in test["result"] else None
+            )
+            adequacy = DataAdequacy(**test["result"]["adequacy"]) if "adequacy" in test["result"] else None
+
+            test["result"] = CausalTestResult(outcome=outcome, effect_estimate=effect_estimate, adequacy=adequacy)
+
+        return CausalTestCase(**test)
 
-        return CausalTestCase(
-            name=test.get("name"),
-            effect_measure=test.get("effect_measure"),
-            query=test.get("query"),
-            expected_causal_effect=expected_effect,
-            estimator=estimator,
-            skip=test.get("skip", False),
-        )
+    def ready_to_run(self) -> bool:
+        """
+        Test whether framework is ready to run test cases.
+        :returns: True if the DAG, data, and test cases are defined.
+        """
+        return all(x is not None for x in (self.test_cases, self.dag, self.df)) and bool(self.test_cases)
 
     def run_tests(self, silent: bool = False, adequacy: bool = False, bootstrap_size: int = 100):
         """
@@ -279,3 +280,10 @@ def save_results(self, output_path) -> list:
             json.dump([test.to_dict() for test in self.test_cases], f, indent=2)
 
         logger.info("Results saved successfully")
+
+    def test_dataframe(self) -> pd.DataFrame:
+        """
+        :returns: The causal test cases as a dataframe. Nested objects such as results are indexed as, e.g.
+        `result.outcome`.
+        """
+        return pd.json_normalize(map(lambda t: t.to_dict(), self.test_cases))