diff --git a/causal_testing/__main__.py b/causal_testing/__main__.py index 3ecd443a..845098e1 100644 --- a/causal_testing/__main__.py +++ b/causal_testing/__main__.py @@ -1,18 +1,163 @@ """This module contains the main entrypoint functionality to the Causal Testing Framework.""" -import json +import argparse import logging -import os -import tempfile +from enum import Enum from importlib.metadata import entry_points -from pathlib import Path +from typing import Optional, Sequence import networkx as nx import pandas as pd +from causal_testing.causal_testing_framework import CausalTestingFramework, read_dataframe from causal_testing.testing.metamorphic_relation import generate_causal_tests -from .main import CausalTestingFramework, CausalTestingPaths, Command, parse_args, setup_logging +logger = logging.getLogger(__name__) + + +class Command(Enum): + """ + Enum for supported CTF commands. + """ + + TEST = "test" + GENERATE = "generate" + DISCOVER = "discover" + + +def setup_logging(level: str) -> None: + """Set up logging configuration.""" + logging.basicConfig( + level=level, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S" + ) + + +def parse_args(args: Optional[Sequence[str]] = None) -> argparse.Namespace: + """Parse command line arguments.""" + main_parser = argparse.ArgumentParser( + add_help=True, + description="Causal Testing Framework - " + "A causal inference-driven framework for functional black-box testing of complex software.", + ) + + main_parser.add_argument( + "-l", + "--log_level", + default="WARNING", + type=str.upper, + choices=["NONE", "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], + help="Set the logging level (default: WARNING).", + ) + + subparsers = main_parser.add_subparsers( + help="The action you want to run - call `causal_testing {action} -h` for further details", dest="command" + ) + + # Generation + parser_generate = subparsers.add_parser(Command.GENERATE.value, help="Generate causal tests from a DAG") + parser_generate.add_argument("-D", "--dag-path", help="Path to the DAG file (.dot)", required=True) + parser_generate.add_argument("-o", "--output", help="Path for output file (.json)", required=True) + parser_generate.add_argument( + "-e", + "--estimator", + help="The name of the estimator class to use when evaluating tests (defaults to LinearRegressionEstimator)", + default="LinearRegressionEstimator", + ) + parser_generate.add_argument( + "-T", + "--effect-type", + help="The effect type to estimate {direct, total}", + default="direct", + ) + parser_generate.add_argument( + "-E", + "--estimate-type", + help="The estimate type to use when evaluating tests (defaults to coefficient)", + default="coefficient", + ) + parser_generate.add_argument( + "-i", "--ignore-cycles", help="Ignore cycles in DAG", action="store_true", default=False + ) + parser_generate.add_argument( + "--threads", "-t", type=int, help="The number of parallel threads to use.", required=False, default=0 + ) + + # Testing + parser_test = subparsers.add_parser(Command.TEST.value, help="Run causal tests") + parser_test.add_argument("-D", "--dag-path", help="Path to the DAG file (.dot)", required=True) + parser_test.add_argument("-o", "--output", help="Path for output file (.json)", required=True) + parser_test.add_argument("-i", "--ignore-cycles", help="Ignore cycles in DAG", action="store_true", default=False) + parser_test.add_argument("-d", "--data-paths", help="Paths to data files (.csv)", nargs="+", required=True) + parser_test.add_argument("-t", "--test-config", help="Path to test configuration file (.json)", required=True) + parser_test.add_argument("-q", "--query", help="Query string to filter data (e.g. 'age > 18')", type=str) + parser_test.add_argument( + "-a", "--adequacy", help="Calculate causal test adequacy for each test case", action="store_true", default=False + ) + parser_test.add_argument( + "-b", + "--adequacy-bootstrap-size", + dest="bootstrap_size", + help="Number of bootstrap samples for causal test adequacy. Defaults to 100", + type=int, + ) + parser_test.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", + help="The name of the technique to use. Currently supported are 'HillClimberDiscovery' and 'NSGADiscovery'", + required=True, + ) + parser_discover.add_argument( + "-V", + "--variables", + help="The subset of variables from the data to consider. Defaults to all.", + nargs="*", + default=[], + ) + parser_discover.add_argument("-o", "--output", help="Path for output DAG file (.dot)", required=True) + parser_discover.add_argument( + "-i", "--include-edges", help="Path to file containing edges to include", required=False + ) + parser_discover.add_argument( + "-e", "--exclude-edges", help="Path to file containing edges to exclude", required=False + ) + parser_discover.add_argument( + "--technique-kwargs", + help="Keywords for the discovery technique. These should be specified as `arg1=value1 arg2=value2...`.", + nargs="*", + default=[], + ) + + args = main_parser.parse_args(args) + + # Assume the user wants test adequacy if they're setting bootstrap_size + if getattr(args, "bootstrap_size", None) is not None: + args.adequacy = True + if getattr(args, "adequacy", False) and getattr(args, "bootstrap_size", None) is None: + # Need this here rather than a default value because otherwise the above always sets adequacy to True + args.bootstrap_size = 100 + + args.command = Command(args.command) + return args def main() -> None: @@ -60,9 +205,14 @@ def main() -> None: logging.info("Discovering causal structure") # Need to reset index to allow for multiple files having the same index (i.e. starting at zero). # Otherwise you end up with duplicate indices, which causes problems further down the line - df = pd.concat([pd.read_csv(path) for path in args.data_paths]).reset_index() + df = pd.concat([read_dataframe(path) for path in args.data_paths]).reset_index() if args.variables: df = df[args.variables] + # Drop unnamed columns + unnamed_columns = [c for c in df.columns if c.startswith("Unnamed: ")] + if unnamed_columns: + logger.warning(f"Dropping unnamed columns: {unnamed_columns}") + df = df.drop(unnamed_columns) discover_class = discover_map[args.technique].load() discover = discover_class( @@ -80,55 +230,20 @@ def main() -> None: discover.write_dot(evolved_dag, args.output) logging.info("Causal structure discovery completed successfully.") case Command.TEST: - # Create paths object - paths = CausalTestingPaths( + # Create and setup framework + framework = CausalTestingFramework() + + framework.setup( dag_path=args.dag_path, data_paths=args.data_paths, - test_config_path=args.test_config, - output_path=args.output, + test_cases_path=args.test_config, + query=args.query, + ignore_cycles=args.ignore_cycles, ) - # Create and setup framework - framework = CausalTestingFramework(paths, ignore_cycles=args.ignore_cycles, query=args.query) - framework.setup() - - # Load and run tests - framework.load_tests() - - if args.batch_size > 0: - logging.info(f"Running tests in batches of size {args.batch_size}") - with tempfile.TemporaryDirectory() as tmpdir: - output_files = [] - for i, results in enumerate( - framework.run_tests_in_batches( - batch_size=args.batch_size, - silent=args.silent, - adequacy=args.adequacy, - bootstrap_size=args.bootstrap_size, - ) - ): - temp_file_path = os.path.join(tmpdir, f"output_{i}.json") - framework.save_results(results, temp_file_path) - output_files.append(temp_file_path) - del results - - # Now stitch the results together from the temporary files - all_results = [] - for file_path in output_files: - with open(file_path, "r", encoding="utf-8") as f: - all_results.extend(json.load(f)) - - output_path = Path(args.output) - output_path.parent.mkdir(parents=True, exist_ok=True) - - with open(args.output, "w", encoding="utf-8") as f: - json.dump(all_results, f, indent=4) - else: - logging.info("Running tests in regular mode") - results = framework.run_tests( - silent=args.silent, adequacy=args.adequacy, bootstrap_size=args.bootstrap_size - ) - framework.save_results(results) + logging.info("Running tests") + framework.run_tests(silent=args.silent, adequacy=args.adequacy, bootstrap_size=args.bootstrap_size) + framework.save_results(args.output) logging.info("Causal testing completed successfully.") diff --git a/causal_testing/causal_testing_framework.py b/causal_testing/causal_testing_framework.py new file mode 100644 index 00000000..a364036b --- /dev/null +++ b/causal_testing/causal_testing_framework.py @@ -0,0 +1,335 @@ +""" +This module implements the CausalTestingFramework class, which is the main interaction point for causal testing. +""" + +import json +import logging +from importlib.metadata import entry_points +from pathlib import Path + +import pandas as pd +from tqdm import tqdm + +from causal_testing.specification.causal_dag import CausalDAG +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 + +logger = logging.getLogger(__name__) + + +def read_dataframe(file_path: str, **kwargs: dict) -> pd.DataFrame: + """ + Read data into a dataframe. + + :param file_path: The path to the data. + :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) + raise ValueError(f"Unsupported file extension: '{suffix}'") + + +class CausalTestingFramework: + """ + Main class for running causal tests. + """ + + def __init__(self, dag: CausalDAG = None, test_cases: list[CausalTestCase] = None, df: pd.DataFrame = None): + self.dag = dag + self.test_cases = test_cases + self.df = df + self.variables = {"inputs": {}, "outputs": {}} + if self.dag is not None and self.df is not None: + self.create_variables() + + def create_variables(self) -> None: + """ + Create variable objects from DAG nodes based on their connectivity. + """ + for node_name, node_data in self.dag.nodes(data=True): + if node_name not in self.df.columns and not node_data.get("hidden", False): + raise ValueError(f"Node {node_name} missing from data. Should it be marked as hidden?") + + dtype = self.df.dtypes.get(node_name) + + # If node has no incoming edges, it's an input + if self.dag.in_degree(node_name) == 0: + self.variables["inputs"][node_name] = Input(name=node_name, datatype=dtype) + + # Otherwise it's an output + if self.dag.in_degree(node_name) > 0: + self.variables["outputs"][node_name] = Output(name=node_name, datatype=dtype) + + def setup( + self, + dag_path: str, + data_paths: list[str], + test_cases_path: str, + ignore_cycles: bool = False, + query: str = None, + **kwargs: dict, + ): + """ + Shortcut for loading in the DAG, data, and test cases. + :param dag_path: Path to the DAG definition file. + :param data_paths: List of paths to input data files. + :param test_cases_path: Path to the test configuration file + :param ignore_cycles: Whether to ignore cycles in the causal graph. + NOTE: Setting this to True severely limits the testing that can be performed. + :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) + + def load_dag(self, dag_path: str, ignore_cycles: bool = False): + """ + Load the causal DAG from the specified file path. + + :param dag_path: Path to the DAG definition file. + :param ignore_cycles: Whether to ignore cycles in the causal graph. + NOTE: Setting this to True severely limits the testing that can be performed. + """ + logger.info(f"Loading DAG from {dag_path}") + self.dag = CausalDAG(dag_path, ignore_cycles=ignore_cycles) + logger.info(f"DAG loaded with {len(self.dag.nodes)} nodes and {len(self.dag.edges)} edges") + + def load_data(self, data_paths: list[str], query: str = None, **kwargs: dict): + """Load and combine all data sources with optional filtering. + + :param data_paths: List of paths to input data files. + :param query: Optional pandas query string to filter the loaded data + :param kwargs: Keyword arguments to be passed to the `read_` function. + """ + logger.info(f"Loading data from {len(data_paths)} source(s)") + + data = pd.concat([read_dataframe(data_path, **kwargs) for data_path in data_paths], axis=0, ignore_index=True) + logger.info(f"Initial data shape: {data.shape}") + + if query: + logger.info(f"Attempting to apply query: '{query}'") + data = data.query(query) + + self.df = data + + def load_test_cases_from_json(self, test_cases_path: str): + """ + Load and prepare test configurations from JSON file. + + :param test_cases_path: Path to the test configuration file + """ + 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.") + + self.create_variables() + + 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 + + def create_base_test(self, test: dict) -> BaseTestCase: + """ + Create base test case from test configuration. + + :param test: Dictionary containing test configuration parameters + + :return: BaseTestCase object + :raises: KeyError if required variables are not found in inputs or outputs + """ + treatment_name = test["treatment_variable"] + outcome_name = next(iter(test["expected_effect"].keys())) + + # Look for treatment variable in both inputs and outputs + treatment_var = self.variables["inputs"].get(treatment_name) or self.variables["outputs"].get(treatment_name) + if not treatment_var: + raise KeyError(f"Treatment variable '{treatment_name}' not found in inputs or outputs") + + # Look for outcome variable in both inputs and outputs + outcome_var = self.variables["inputs"].get(outcome_name) or self.variables["outputs"].get(outcome_name) + if not outcome_var: + raise KeyError(f"Outcome variable '{outcome_name}' not found in inputs or outputs") + + return BaseTestCase( + treatment_variable=treatment_var, outcome_variable=outcome_var, effect=test.get("effect", "total") + ) + + def create_causal_test(self, test: dict) -> CausalTestCase: + """ + Create causal test case from test configuration and base test. + + :param test: Dictionary containing test configuration parameters + + :return: CausalTestCase object + :raises: ValueError if invalid estimator or configuration is provided + """ + 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")} + + base_test = self.create_base_test(test) + + if "estimator" not in test: + raise ValueError("Test configuration must specify an estimator") + + if test["estimator"] not in estimator_map: + raise ValueError( + f"Unsupported estimator {test['estimator']}. 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 + estimator_class = estimator_map.get(test["estimator"]).load() + estimator_kwargs = test.get("estimator_kwargs", {}) + estimator = estimator_class( + base_test_case=base_test, + treatment_value=test.get("treatment_value"), + control_value=test.get("control_value"), + adjustment_set=test.get( + "adjustment_set", + self.dag.identification(base_test), + ), + alpha=test.get("alpha", 0.05), + **estimator_kwargs, + ) + + # Get effect type and create expected effect + effect_type = test["expected_effect"][base_test.outcome_variable.name] + if effect_type not in effect_map: + raise ValueError( + f"Unsupported causal effect {effect_type}. 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()(**test.get("effect_kwargs", {})) + + return CausalTestCase( + name=test.get("name"), + query=test.get("query"), + base_test_case=base_test, + expected_causal_effect=expected_effect, + estimate_type=test.get("estimate_type", "ate"), + estimator=estimator, + skip=test.get("skip", False), + ) + + def run_tests(self, silent: bool = False, adequacy: bool = False, bootstrap_size: int = 100): + """ + Run all test cases and return their results. + + :param silent: Whether to suppress errors + :param adequacy: Whether to calculate causal test adequacy (defaults to False) + :param bootstrap_size: The number of bootstrap samples to use when calculating causal test adequacy + (defaults to 100) + + :raises: ValueError if no tests are loaded + :raises: Exception if test execution fails + """ + logger.info("Running causal tests...") + + if not self.test_cases: + raise ValueError("No tests to run.") + + for test_case in tqdm(self.test_cases): + test_case.execute_test( + self.df, suppress_estimation_errors=silent, adequacy=adequacy, bootstrap_size=bootstrap_size + ) + + 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}") + + # Create parent directory if it doesn't exist + Path(output_path).parent.mkdir(parents=True, exist_ok=True) + + json_results = [] + result_index = 0 + + for test_case in self.test_cases: + # Create a base output first of common entries + base_output = { + "name": test_case.name, + "estimate_type": test_case.estimate_type, + "effect": test_case.base_test_case.effect, + "treatment_variable": test_case.base_test_case.treatment_variable.name, + "expected_effect": test_case.expected_causal_effect.__class__.__name__, + "alpha": test_case.estimator.alpha, + } + if test_case.skip: + # Include those skipped test entry without execution results + output = { + **base_output, + "formula": test_case.estimator.formula if hasattr(test_case.estimator, "formula") else None, + "skip": True, + "passed": None, + "result": { + "status": "skipped", + "reason": "Test marked as skip:true in the causal test config file.", + }, + } + else: + 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, + "result": ( + { + "treatment": test_case.estimator.base_test_case.treatment_variable.name, + "outcome": test_case.estimator.base_test_case.outcome_variable.name, + "adjustment_set": ( + list(test_case.estimator.adjustment_set) if test_case.estimator.adjustment_set else [] + ), + } + | result.effect_estimate.to_dict() + | (result.adequacy.to_dict() if result.adequacy else {}) + if result.effect_estimate + else {"status": "error", "reason": result.error_message} + ), + } + + json_results.append(output) + + # Save to file + with open(output_path, "w", encoding="utf-8") as f: + json.dump(json_results, f, indent=2) + + logger.info("Results saved successfully") + return json_results diff --git a/causal_testing/discovery/abstract_discovery.py b/causal_testing/discovery/abstract_discovery.py index cb0c049e..ee599ad9 100644 --- a/causal_testing/discovery/abstract_discovery.py +++ b/causal_testing/discovery/abstract_discovery.py @@ -14,11 +14,11 @@ import pandas as pd import rustworkx as rx -from causal_testing.main import CausalTestingFramework +from causal_testing.causal_testing_framework import CausalTestingFramework from causal_testing.specification.causal_dag import CausalDAG from causal_testing.specification.scenario import Scenario from causal_testing.testing.causal_effect import Negative, Positive -from causal_testing.testing.causal_test_result import CausalTestResult +from causal_testing.testing.causal_test_case import CausalTestCase from causal_testing.testing.metamorphic_relation import generate_metamorphic_relations TestResult = Enum("TestResult", [("PASS", 2), ("FAIL", 0), ("INESTIMABLE", 1)]) @@ -38,23 +38,6 @@ def simple_cycle(causal_dag: CausalDAG): return [(rx_graph[i], rx_graph[j]) for i, j in rx.digraph_find_cycle(rx_graph)] -def effect_direction(result: CausalTestResult) -> str: - """ - Check whether the estimated causal effect is negative or positive. - - :param result: The causal test result object. - :returns: Whether the estimated causal test is positive or negative (or no effect). - """ - if pd.api.types.is_numeric_dtype( - result.estimator.df[result.estimator.base_test_case.treatment_variable.name] - ) and pd.api.types.is_numeric_dtype(result.estimator.df[result.estimator.base_test_case.outcome_variable.name]): - if Negative().apply(result): - return "negative" - if Positive().apply(result): - return "positive" - return None - - def is_match(u: str, v: str, patterns: list[str]): """ Check whether a given edge matches a given pattern. @@ -118,6 +101,22 @@ 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.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): + return "negative" + if Positive().apply(test_case.result): + 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 @@ -145,24 +144,37 @@ def write_dot(self, individual: CausalDAG, output_file: str): 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"] == TestResult.PASS: + print(" GREEN") individual[test["treatment"]][test["outcome"]]["color"] = "green" + individual[test["treatment"]][test["outcome"]]["fontcolor"] = "green" elif test["result"] == TestResult.INESTIMABLE: + print(" ORANGE") individual[test["treatment"]][test["outcome"]]["color"] = "orange" + individual[test["treatment"]][test["outcome"]]["fontcolor"] = "orange" elif test["result"] == TestResult.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"] == TestResult.PASS: individual[test["treatment"]][test["outcome"]]["style"] = "invis" individual[test["treatment"]][test["outcome"]]["constraint"] = False elif test["result"] == TestResult.INESTIMABLE: individual[test["treatment"]][test["outcome"]]["color"] = "orange" + individual[test["treatment"]][test["outcome"]]["fontcolor"] = "orange" elif test["result"] == TestResult.FAIL: individual[test["treatment"]][test["outcome"]]["color"] = "red" + individual[test["treatment"]][test["outcome"]]["fontcolor"] = "red" else: raise ValueError(f"Invalid test outcome {test['result']}") @@ -188,38 +200,36 @@ def evaluate_tests(self, causal_dag: CausalDAG) -> pd.DataFrame: (result, expected effect, treatment, outcome, effect direction). """ - ctf = CausalTestingFramework(None) - ctf.dag = causal_dag - ctf.data = self.df + ctf = CausalTestingFramework(dag=causal_dag, df=self.df) ctf.create_variables() ctf.scenario = Scenario(list(ctf.variables["inputs"].values()) + list(ctf.variables["outputs"].values())) - ctf.test_cases = ctf.create_test_cases( - { - "tests": [ - relation.to_json_stub( - alpha=self.alpha, - **self._json_stub_params(relation.base_test_case.outcome_variable), - ) - for relation in generate_metamorphic_relations(causal_dag) - ] - } - ) + ctf.test_cases = [ + ctf.create_causal_test( + relation.to_json_stub( + alpha=self.alpha, + **self._json_stub_params(relation.base_test_case.outcome_variable), + ) + ) + for relation in generate_metamorphic_relations(causal_dag) + ] results = [] - for test_case, result in zip(ctf.test_cases, ctf.test_cases): + for test_case in ctf.test_cases: try: - result = test_case.execute_test() + test_case.execute_test(self.df) results.append( { "result": ( - TestResult.PASS if test_case.expected_causal_effect.apply(result) else TestResult.FAIL + TestResult.PASS + if test_case.expected_causal_effect.apply(test_case.result) + else TestResult.FAIL ), "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, - "effect": effect_direction(result), + "effect": self.effect_direction(test_case), } ) except np.linalg.LinAlgError: @@ -233,4 +243,6 @@ 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/estimation/abstract_estimator.py b/causal_testing/estimation/abstract_estimator.py index b799ee01..0ffdadf5 100644 --- a/causal_testing/estimation/abstract_estimator.py +++ b/causal_testing/estimation/abstract_estimator.py @@ -4,8 +4,6 @@ from abc import ABC, abstractmethod from typing import Any -import pandas as pd - from causal_testing.testing.base_test_case import BaseTestCase logger = logging.getLogger(__name__) @@ -24,9 +22,9 @@ class Estimator(ABC): maintains a list of modelling assumptions (as strings). If a user wishes to implement their own estimator, they must implement this method and add all assumptions to the list of modelling assumptions. - 2) estimate_ate: All estimators must be capable of returning the average treatment effect as a minimum. That is, the - average effect of the intervention (changing treatment from control to treated value) on the outcome of interest - adjusted for all confounders. + 2) estimate_*: The Causal Testing Framework expects causal effects to be calculated by methods that start with + `estimate_`, followed by the name of the causal effect measure being estimated, for example `ate` or `risk_ratio`. + Naming methods this way enables estimators to hook nicely into the endpoints further up the chain. """ def __init__( @@ -37,25 +35,20 @@ def __init__( treatment_value: float, control_value: float, adjustment_set: set, - df: pd.DataFrame = None, effect_modifiers: dict[str, Any] = None, alpha: float = 0.05, - query: str = "", ): self.base_test_case = base_test_case self.treatment_value = treatment_value self.control_value = control_value self.adjustment_set = adjustment_set self.alpha = alpha - self.df = df.query(query) if query else df if effect_modifiers is None: self.effect_modifiers = {} else: self.effect_modifiers = effect_modifiers self.modelling_assumptions = [] - if query: - self.modelling_assumptions.append(query) self.add_modelling_assumptions() logger.debug("Effect Modifiers: %s", self.effect_modifiers) @@ -65,10 +58,3 @@ def add_modelling_assumptions(self): Add modelling assumptions to the estimator. This is a list of strings which list the modelling assumptions that must hold if the resulting causal inference is to be considered valid. """ - - def compute_confidence_intervals(self) -> list[float, float]: - """ - Estimate the 95% Wald confidence intervals for the effect of changing the treatment from control values to - treatment values on the outcome. - :return: 95% Wald confidence intervals. - """ diff --git a/causal_testing/estimation/abstract_regression_estimator.py b/causal_testing/estimation/abstract_regression_estimator.py index 76d1f83a..cebfe727 100644 --- a/causal_testing/estimation/abstract_regression_estimator.py +++ b/causal_testing/estimation/abstract_regression_estimator.py @@ -24,14 +24,13 @@ def __init__( # pylint: disable=too-many-arguments self, base_test_case: BaseTestCase, - treatment_value: float, - control_value: float, - adjustment_set: set, - df: pd.DataFrame = None, + treatment_value: float = None, + control_value: float = None, + adjustment_set: set = None, effect_modifiers: dict[Variable, Any] = None, + adjustment_config: dict[Variable, Any] = None, formula: str = None, alpha: float = 0.05, - query: str = "", ): # pylint: disable=R0801 super().__init__( @@ -39,14 +38,13 @@ def __init__( treatment_value=treatment_value, control_value=control_value, adjustment_set=adjustment_set, - df=df, effect_modifiers=effect_modifiers, alpha=alpha, - query=query, ) if effect_modifiers is None: - effect_modifiers = [] + effect_modifiers = {} + self.adjustment_config = {} if adjustment_config is None else adjustment_config if adjustment_set is None: adjustment_set = [] if formula is not None: @@ -56,22 +54,22 @@ def __init__( [base_test_case.treatment_variable.name] + sorted(list(adjustment_set)) + sorted(list(effect_modifiers)) ) self.formula = f"{base_test_case.outcome_variable.name} ~ {'+'.join(terms)}" - self.setup_covariates() - def setup_covariates(self): + for term in list(self.effect_modifiers) + list(self.adjustment_config): + self.adjustment_set.add(term) + + def _setup_covariates(self, df: pd.DataFrame) -> pd.Series: """ Parse the formula and set up the covariates from the design matrix so we can use them in the statsmodels array API. This allows us to only parse the formula once, rather than using the formula API, which parses it every time the regression model is fit, which can be a lot if using causal test adequacy. + :param df: The data to use. + :returns: The data and the covariate columns. """ - if self.df is not None: - _, covariate_data = dmatrices(self.formula, self.df, return_type="dataframe") - self.covariates = covariate_data.columns - self.df = pd.concat( - [self.df, covariate_data[[col for col in covariate_data.columns if col not in self.df]]], axis=1 - ) - else: - self.covariates = None + _, covariate_data = dmatrices(self.formula, df, return_type="dataframe") + df = pd.concat([df, covariate_data[[col for col in covariate_data.columns if col not in df]]], axis=1) + covariates = covariate_data.columns.tolist() + return covariates, df.dropna(subset=covariates) @property @abstractmethod @@ -94,31 +92,25 @@ def add_modelling_assumptions(self): "do not need to be linear." ) - def fit_model(self, data=None) -> RegressionResultsWrapper: + def fit_model(self, df: pd.DataFrame) -> RegressionResultsWrapper: """Run logistic regression of the treatment and adjustment set against the outcome and return the model. + :param df: The data to use. :return: The model after fitting to data. """ - if data is None: - data = self.df - if self.covariates is None: - _, covariate_data = dmatrices(self.formula, data, return_type="dataframe") - covariates = covariate_data.columns - data = pd.concat( - [data, covariate_data[[col for col in covariate_data.columns if col not in data]]], axis=1 - ).dropna() - model = self.regressor(data[self.base_test_case.outcome_variable.name], data[covariates]).fit(disp=0) - else: - data = data.dropna(subset=self.covariates) - model = self.regressor(data[self.base_test_case.outcome_variable.name], data[self.covariates]).fit(disp=0) + covariates, df = self._setup_covariates(df) + model = self.regressor(df[self.base_test_case.outcome_variable.name], df[covariates]).fit(disp=0) return model - def treatment_columns(self, model: RegressionResultsWrapper): + def treatment_columns(self, model: RegressionResultsWrapper) -> list[str]: """ Get the names of the treatment columns from the model. This is a workaround for statsmodels mangling the names of categorical variables to include the values. :param model: The fitted model from which to extract the variable names. + :returns: A list of the feature names in the model that represent the treatment. Normally this will just be + [treatment_name], but for categorical treatments, you'll have + [treatment_name[value_1], treatment_name[value_2]]. """ return [ param @@ -127,25 +119,22 @@ def treatment_columns(self, model: RegressionResultsWrapper): or param.startswith(self.base_test_case.treatment_variable.name + "[") ] - def _predict(self, data=None, adjustment_config: dict = None) -> pd.DataFrame: + def _predict(self, df) -> pd.DataFrame: """Estimate the outcomes under control and treatment. - :param data: The data to use, defaults to `self.df`. Controllable for boostrap sampling. + :param df: The data to use. :param: adjustment_config: The values of the adjustment variables to use. :return: The estimated outcome under control and treatment, with confidence intervals in the form of a dataframe with columns "predicted", "se", "ci_lower", and "ci_upper". """ - if adjustment_config is None: - adjustment_config = {} - - model = self.fit_model(data) + model = self.fit_model(df) - x = pd.DataFrame(columns=self.df.columns) + x = pd.DataFrame(columns=df.columns) x["Intercept"] = 1 # self.intercept x[self.base_test_case.treatment_variable.name] = [self.treatment_value, self.control_value] - for k, v in adjustment_config.items(): + for k, v in self.adjustment_config.items(): x[k] = v for k, v in self.effect_modifiers.items(): x[k] = v diff --git a/causal_testing/estimation/cubic_spline_estimator.py b/causal_testing/estimation/cubic_spline_estimator.py index fa4cedbe..a5a53705 100644 --- a/causal_testing/estimation/cubic_spline_estimator.py +++ b/causal_testing/estimation/cubic_spline_estimator.py @@ -3,10 +3,10 @@ import logging from typing import Any -import statsmodels.formula.api as smf -from statsmodels.regression.linear_model import RegressionResultsWrapper import pandas as pd +import statsmodels.formula.api as smf +from statsmodels.regression.linear_model import RegressionResultsWrapper from causal_testing.estimation.effect_estimate import EffectEstimate from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator @@ -31,20 +31,27 @@ def __init__( control_value: float, adjustment_set: set, basis: int, - df: pd.DataFrame = None, effect_modifiers: dict[Variable, Any] = None, formula: str = None, alpha: float = 0.05, expected_relationship=None, + adjustment_config: dict[Variable, Any] = None, ): super().__init__( - base_test_case, treatment_value, control_value, adjustment_set, df, effect_modifiers, formula, alpha + base_test_case=base_test_case, + treatment_value=treatment_value, + control_value=control_value, + adjustment_set=adjustment_set, + effect_modifiers=effect_modifiers, + formula=formula, + alpha=alpha, ) self.expected_relationship = expected_relationship + self.adjustment_config = adjustment_config if effect_modifiers is None: - effect_modifiers = [] + effect_modifiers = {} if formula is None: terms = ( @@ -52,33 +59,33 @@ def __init__( ) self.formula = f"{base_test_case.outcome_variable.name} ~ cr({'+'.join(terms)}, df={basis})" - def fit_model(self, data=None) -> RegressionResultsWrapper: + def fit_model(self, df: pd.DataFrame) -> RegressionResultsWrapper: """Run linear regression of the treatment and adjustment set against the outcome and return the model. + :param df: The data to use. :return: The model after fitting to data. """ - if data is None: - data = self.df - model = self.regressor(formula=self.formula, data=data).fit(disp=0) + model = self.regressor(formula=self.formula, data=df).fit(disp=0) return model - def estimate_ate_calculated(self, adjustment_config: dict = None) -> EffectEstimate: + def estimate_ate_calculated( + self, + df: pd.DataFrame, + ) -> EffectEstimate: """Estimate the ate effect of the treatment on the outcome. That is, the change in outcome caused by changing the treatment variable from the control value to the treatment value. Here, we actually calculate the expected outcomes under control and treatment and divide one by the other. This allows for custom terms to be put in such as squares, inverses, products, etc. - :param: adjustment_config: The configuration of the adjustment set as a dict mapping variable names to - their values. N.B. Every variable in the adjustment set MUST have a value in - order to estimate the outcome under control and treatment. + :param df: The data to use. :return: The average treatment effect. """ - model = self.fit_model() + model = self.fit_model(df) x = pd.DataFrame({"Intercept": [1], self.base_test_case.treatment_variable.name: [self.treatment_value]}) - if adjustment_config is not None: - for k, v in adjustment_config.items(): + if self.adjustment_config is not None: + for k, v in self.adjustment_config.items(): x[k] = v if self.effect_modifiers is not None: for k, v in self.effect_modifiers.items(): diff --git a/causal_testing/estimation/instrumental_variable_estimator.py b/causal_testing/estimation/instrumental_variable_estimator.py index c7775363..8ff237ee 100644 --- a/causal_testing/estimation/instrumental_variable_estimator.py +++ b/causal_testing/estimation/instrumental_variable_estimator.py @@ -29,22 +29,20 @@ def __init__( control_value: float, adjustment_set: set, instrument: str, - df: pd.DataFrame = None, alpha: float = 0.05, - query: str = "", + bootstrap_size=100, ): super().__init__( base_test_case=base_test_case, treatment_value=treatment_value, control_value=control_value, adjustment_set=adjustment_set, - df=df, effect_modifiers=None, alpha=alpha, - query=query, ) self.instrument = instrument + self.bootstrap_size = bootstrap_size def add_modelling_assumptions(self): """ @@ -77,16 +75,15 @@ def iv_coefficient(self, df) -> float: # Estimate the coefficient of I on X by cancelling return ab / a - def estimate_coefficient(self, bootstrap_size=100) -> EffectEstimate: + def estimate_coefficient(self, df: pd.DataFrame) -> EffectEstimate: """ - Estimate the unit ate (i.e. coefficient) of the treatment on the - outcome. + Estimate the unit ate (i.e. coefficient) of the treatment on the outcome. + + :param df: The data to use. """ - bootstraps = sorted( - [self.iv_coefficient(self.df.sample(len(self.df), replace=True)) for _ in range(bootstrap_size)] - ) - bound = ceil((bootstrap_size * self.alpha) / 2) + bootstraps = sorted([self.iv_coefficient(df.sample(len(df), replace=True)) for _ in range(self.bootstrap_size)]) + bound = ceil((self.bootstrap_size * self.alpha) / 2) ci_low = pd.Series(bootstraps[bound]) - ci_high = pd.Series(bootstraps[bootstrap_size - bound]) + ci_high = pd.Series(bootstraps[self.bootstrap_size - bound]) - return EffectEstimate("coefficient", pd.Series(self.iv_coefficient(self.df)), ci_low, ci_high) + return EffectEstimate("coefficient", pd.Series(self.iv_coefficient(df)), ci_low, ci_high) diff --git a/causal_testing/estimation/ipcw_estimator.py b/causal_testing/estimation/ipcw_estimator.py index a282a415..ea96c5f9 100644 --- a/causal_testing/estimation/ipcw_estimator.py +++ b/causal_testing/estimation/ipcw_estimator.py @@ -42,7 +42,6 @@ class IPCWEstimator(Estimator): # pylint: disable=too-many-instance-attributes def __init__( self, - df: pd.DataFrame, timesteps_per_observation: int, control_strategy: list[tuple[int, str, Any]], treatment_strategy: list[tuple[int, str, Any]], @@ -59,10 +58,8 @@ def __init__( treatment_value=[val for _, _, val in treatment_strategy], control_value=[val for _, _, val in control_strategy], adjustment_set=None, - df=df, effect_modifiers=None, alpha=alpha, - query="", ) self.timesteps_per_observation = timesteps_per_observation self.control_strategy = control_strategy @@ -71,7 +68,6 @@ def __init__( self.fit_bl_switch_formula = fit_bl_switch_formula self.fit_bltd_switch_formula = fit_bltd_switch_formula self.eligibility = eligibility - self.df = df.sort_values(["id", "time"]) self.len_control_group = None self.len_treatment_group = None @@ -80,7 +76,6 @@ def __init__( max(len(self.control_strategy), len(self.treatment_strategy)) + 1 ) * self.timesteps_per_observation self.total_time = total_time - self.preprocess_data() def add_modelling_assumptions(self): self.modelling_assumptions.append("The variables in the data vary over time.") @@ -162,6 +157,8 @@ def setup_fault_time(self, individual: pd.DataFrame, perturbation: float = -0.00 Return the time at which the event of interest (i.e. a fault) occurred. """ fault = individual[~individual[self.status_column]] + if (individual[self.status_column]).all(): + raise ValueError("No recorded faults") fault_time = ( individual["time"].loc[fault.index[0]] if not fault.empty @@ -173,35 +170,36 @@ def setup_fault_time(self, individual: pd.DataFrame, perturbation: float = -0.00 } ) - def preprocess_data(self): + def preprocess_data(self, df: pd.DataFrame) -> pd.DataFrame: """ Set up the treatment-specific columns in the data that are needed to estimate the hazard ratio. + + :param df: The data to use. + :returns: The preprocessed DataFrame. """ - self.df["trtrand"] = None # treatment/control arm - self.df["xo_t_do"] = None # did the individual deviate from the treatment of interest here? - self.df["eligible"] = self.df.eval(self.eligibility) if self.eligibility is not None else True + df = df.sort_values(["id", "time"]) + + df["trtrand"] = None # treatment/control arm + df["xo_t_do"] = None # did the individual deviate from the treatment of interest here? + df["eligible"] = df.eval(self.eligibility) if self.eligibility is not None else True # when did a fault occur? - fault_time_df = self.df.groupby("id", sort=False)[[self.status_column, "time", "id"]].apply( - self.setup_fault_time - ) + fault_time_df = df.groupby("id", sort=False)[[self.status_column, "time", "id"]].apply(self.setup_fault_time) - assert len(fault_time_df) == len(self.df), "Fault times error" - self.df["fault_time"] = fault_time_df["fault_time"].values + assert len(fault_time_df) == len(df), "Fault times error" + df["fault_time"] = fault_time_df["fault_time"].values assert ( - self.df.groupby("id", sort=False)["fault_time"].apply(lambda x: len(set(x)) == 1).all() + df.groupby("id", sort=False)["fault_time"].apply(lambda x: len(set(x)) == 1).all() ), "Each individual must have a unique fault time." - fault_t_do_df = self.df.groupby("id", sort=False)[["id", "time", self.status_column]].apply( - self.setup_fault_t_do - ) - assert len(fault_t_do_df) == len(self.df), "Fault t_do error" - self.df["fault_t_do"] = fault_t_do_df["fault_t_do"].values + fault_t_do_df = df.groupby("id", sort=False)[["id", "time", self.status_column]].apply(self.setup_fault_t_do) + assert len(fault_t_do_df) == len(df), "Fault t_do error" + df["fault_t_do"] = fault_t_do_df["fault_t_do"].values - living_runs = self.df.query("fault_time > 0").loc[ - (self.df["time"] % self.timesteps_per_observation == 0) & (self.df["time"] <= self.total_time) + living_runs = df.query("fault_time > 0").loc[ + (df["time"] % self.timesteps_per_observation == 0) & (df["time"] <= self.total_time) ] logging.debug(" Preprocessing groups") @@ -278,37 +276,28 @@ def preprocess_data(self): ) ] - self.df = individuals.loc[ + df = individuals.loc[ individuals["time"] < np.ceil(individuals["fault_time"] / self.timesteps_per_observation) * self.timesteps_per_observation ].reset_index() logger.debug(f"{len(individuals.groupby('id'))} individuals") + return df - def estimate_hazard_ratio(self) -> EffectEstimate: + def estimate_hazard_ratio(self, df: pd.DataFrame) -> EffectEstimate: """ Estimate the hazard ratio. + :param df: The data to use. """ - if self.df["fault_t_do"].sum() == 0: - raise ValueError("No recorded faults") - - preprocessed_data = self.df.copy() + df = self.preprocess_data(df) # Use logistic regression to predict switching given baseline covariates logger.debug("Use logistic regression to predict switching given baseline covariates") - fit_bl_switch_c = smf.logit(self.fit_bl_switch_formula, data=self.df.loc[self.df.trtrand == 0]).fit( - method="bfgs" - ) - fit_bl_switch_t = smf.logit(self.fit_bl_switch_formula, data=self.df.loc[self.df.trtrand == 1]).fit( - method="bfgs" - ) + fit_bl_switch_c = smf.logit(self.fit_bl_switch_formula, data=df.loc[df.trtrand == 0]).fit(method="bfgs") + fit_bl_switch_t = smf.logit(self.fit_bl_switch_formula, data=df.loc[df.trtrand == 1]).fit(method="bfgs") - preprocessed_data.loc[preprocessed_data["trtrand"] == 0, "pxo1"] = fit_bl_switch_c.predict( - self.df.loc[self.df.trtrand == 0] - ) - preprocessed_data.loc[preprocessed_data["trtrand"] == 1, "pxo1"] = fit_bl_switch_t.predict( - self.df.loc[self.df.trtrand == 1] - ) + df.loc[df["trtrand"] == 0, "pxo1"] = fit_bl_switch_c.predict(df.loc[df.trtrand == 0]) + df.loc[df["trtrand"] == 1, "pxo1"] = fit_bl_switch_t.predict(df.loc[df.trtrand == 1]) # Use logistic regression to predict switching given baseline and time-updated covariates (model S12) logger.debug( @@ -316,20 +305,16 @@ def estimate_hazard_ratio(self) -> EffectEstimate: ) fit_bltd_switch_c = smf.logit( self.fit_bltd_switch_formula, - data=self.df.loc[self.df.trtrand == 0], + data=df.loc[df.trtrand == 0], ).fit(method="bfgs") fit_bltd_switch_t = smf.logit( self.fit_bltd_switch_formula, - data=self.df.loc[self.df.trtrand == 1], + data=df.loc[df.trtrand == 1], ).fit(method="bfgs") - preprocessed_data.loc[preprocessed_data["trtrand"] == 0, "pxo2"] = fit_bltd_switch_c.predict( - self.df.loc[self.df.trtrand == 0] - ) - preprocessed_data.loc[preprocessed_data["trtrand"] == 1, "pxo2"] = fit_bltd_switch_t.predict( - self.df.loc[self.df.trtrand == 1] - ) - if (preprocessed_data["pxo2"] == 1).any(): + df.loc[df["trtrand"] == 0, "pxo2"] = fit_bltd_switch_c.predict(df.loc[df.trtrand == 0]) + df.loc[df["trtrand"] == 1, "pxo2"] = fit_bltd_switch_t.predict(df.loc[df.trtrand == 1]) + if (df["pxo2"] == 1).any(): raise ValueError( "Probability of switching given baseline and time-varying confounders (pxo2) cannot be one." ) @@ -338,36 +323,30 @@ def estimate_hazard_ratio(self) -> EffectEstimate: # Estimate the probabilities of remaining 'un-switched' and hence the weights logger.debug("Estimate the probabilities of remaining 'un-switched' and hence the weights") - preprocessed_data["num"] = 1 - preprocessed_data["pxo1"] - preprocessed_data["denom"] = 1 - preprocessed_data["pxo2"] - preprocessed_data[["num", "denom"]] = ( - preprocessed_data.sort_values(["id", "time"]).groupby("id")[["num", "denom"]].cumprod() - ) + df["num"] = 1 - df["pxo1"] + df["denom"] = 1 - df["pxo2"] + df[["num", "denom"]] = df.sort_values(["id", "time"]).groupby("id")[["num", "denom"]].cumprod() - assert ( - not preprocessed_data["num"].isnull().any() - ), f"{len(preprocessed_data['num'].isnull())} null numerator values" - assert ( - not preprocessed_data["denom"].isnull().any() - ), f"{len(preprocessed_data['denom'].isnull())} null denom values" + assert not df["num"].isnull().any(), f"{len(df['num'].isnull())} null numerator values" + assert not df["denom"].isnull().any(), f"{len(df['denom'].isnull())} null denom values" - preprocessed_data["weight"] = 1 / preprocessed_data["denom"] - preprocessed_data["sweight"] = preprocessed_data["num"] / preprocessed_data["denom"] + df["weight"] = 1 / df["denom"] + df["sweight"] = df["num"] / df["denom"] - preprocessed_data["tin"] = preprocessed_data["time"] - preprocessed_data["tout"] = pd.concat( - [(preprocessed_data["time"] + self.timesteps_per_observation), preprocessed_data["fault_time"]], + df["tin"] = df["time"] + df["tout"] = pd.concat( + [(df["time"] + self.timesteps_per_observation), df["fault_time"]], axis=1, ).min(axis=1) - assert (preprocessed_data["tin"] <= preprocessed_data["tout"]).all(), "Individuals left before joining." + assert (df["tin"] <= df["tout"]).all(), "Individuals left before joining." # IPCW step 4: Use these weights in a weighted analysis of the outcome model # Estimate the KM graph and IPCW hazard ratio using Cox regression. logger.debug("Estimate the KM graph and IPCW hazard ratio using Cox regression.") cox_ph = CoxPHFitter(penalizer=0.2, alpha=self.alpha) cox_ph.fit( - df=preprocessed_data, + df=df, duration_col="tout", event_col="fault_t_do", weights_col="weight", diff --git a/causal_testing/estimation/linear_regression_estimator.py b/causal_testing/estimation/linear_regression_estimator.py index 1f418358..daa8a300 100644 --- a/causal_testing/estimation/linear_regression_estimator.py +++ b/causal_testing/estimation/linear_regression_estimator.py @@ -1,7 +1,6 @@ """This module contains the LinearRegressionEstimator for estimating continuous outcomes.""" import logging -from typing import Any import pandas as pd import statsmodels.api as sm @@ -10,8 +9,6 @@ from causal_testing.estimation.abstract_regression_estimator import RegressionEstimator from causal_testing.estimation.effect_estimate import EffectEstimate from causal_testing.estimation.genetic_programming_regression_fitter import GP -from causal_testing.specification.variable import Variable -from causal_testing.testing.base_test_case import BaseTestCase logger = logging.getLogger(__name__) @@ -23,37 +20,9 @@ class LinearRegressionEstimator(RegressionEstimator): regressor = sm.OLS - def __init__( - # pylint: disable=too-many-arguments - self, - base_test_case: BaseTestCase, - treatment_value: float, - control_value: float, - adjustment_set: set, - df: pd.DataFrame = None, - effect_modifiers: dict[Variable, Any] = None, - formula: str = None, - alpha: float = 0.05, - query: str = "", - ): - # pylint: disable=too-many-arguments - # pylint: disable=R0801 - super().__init__( - base_test_case=base_test_case, - treatment_value=treatment_value, - control_value=control_value, - adjustment_set=adjustment_set, - df=df, - effect_modifiers=effect_modifiers, - alpha=alpha, - query=query, - formula=formula, - ) - for term in self.effect_modifiers: - self.adjustment_set.add(term) - def gp_formula( self, + df: pd.DataFrame, ngen: int = 100, pop_size: int = 20, num_offspring: int = 10, @@ -67,6 +36,7 @@ def gp_formula( """ Use Genetic Programming (GP) to infer the regression equation from the data. + :param df: The data to use. :param ngen: The maximum number of GP generations to run for. :param pop_size: The GP population size. :param num_offspring: The number of offspring per generation. @@ -81,7 +51,7 @@ def gp_formula( :param seed: Random seed for the GP. """ gp = GP( - df=self.df, + df=df, features=sorted(list(self.adjustment_set.union([self.base_test_case.treatment_variable.name]))), outcome=self.base_test_case.outcome_variable.name, extra_operators=extra_operators, @@ -92,27 +62,28 @@ def gp_formula( formula = gp.run_gp(ngen=ngen, pop_size=pop_size, num_offspring=num_offspring, seeds=seeds) formula = gp.simplify(formula) self.formula = f"{self.base_test_case.outcome_variable.name} ~ I({formula}) - 1" - self.setup_covariates() + self._setup_covariates(df) - def estimate_coefficient(self) -> EffectEstimate: + def estimate_coefficient(self, df: pd.DataFrame) -> EffectEstimate: """Estimate the unit average treatment effect of the treatment on the outcome. That is, the change in outcome caused by a unit change in treatment. + :param df: The data to use. :return: The unit average treatment effect and the Wald confidence intervals. """ - model = self.fit_model() + model = self.fit_model(df) newline = "\n" patsy_md = ModelDesc.from_formula(self.base_test_case.treatment_variable.name) if any( ( - not pd.api.types.is_numeric_dtype(self.df.dtypes[factor.name()]) + not pd.api.types.is_numeric_dtype(df.dtypes[factor.name()]) for factor in patsy_md.rhs_termlist[1].factors # We want to remove this long term as it prevents us from discovering categoricals within I(...) blocks - if factor.name() in self.df.dtypes + if factor.name() in df.dtypes ) ): - design_info = dmatrix(self.formula.split("~")[1], self.df).design_info + design_info = dmatrix(self.formula.split("~")[1], df).design_info treatment = design_info.column_names[ design_info.term_name_slices[self.base_test_case.treatment_variable.name] ] @@ -123,15 +94,20 @@ def estimate_coefficient(self) -> EffectEstimate: ), f"{treatment} not in\n{' ' + str(model.params.index).replace(newline, newline + ' ')}" unit_effect = model.params[treatment] # Unit effect is the coefficient of the treatment [ci_low, ci_high] = self._get_confidence_intervals(model, treatment) + + if len(unit_effect) == 0: + unit_effect = pd.Series({self.base_test_case.treatment_variable.name: None}) + return EffectEstimate("coefficient", unit_effect, ci_low, ci_high) - def estimate_ate(self) -> EffectEstimate: + def estimate_ate(self, df: pd.DataFrame) -> EffectEstimate: """Estimate the average treatment effect of the treatment on the outcome. That is, the change in outcome caused by changing the treatment variable from the control value to the treatment value. + :param df: The data to use. :return: The average treatment effect and the Wald confidence intervals. """ - model = self.fit_model() + model = self.fit_model(df) # Create an empty individual for the control and treated individuals = pd.DataFrame(1, index=["control", "treated"], columns=model.params.index) @@ -151,13 +127,15 @@ def estimate_ate(self) -> EffectEstimate: ci_low, ci_high = [pd.Series(interval) for interval in confidence_intervals] return EffectEstimate("ate", ate, ci_low, ci_high) - def estimate_risk_ratio(self, adjustment_config: dict = None) -> EffectEstimate: + def estimate_risk_ratio(self, df: pd.DataFrame) -> EffectEstimate: """Estimate the risk_ratio effect of the treatment on the outcome. That is, the change in outcome caused by changing the treatment variable from the control value to the treatment value. + :param df: The data to use. + :return: The average treatment effect and the Wald confidence intervals. """ - prediction = self._predict(adjustment_config=adjustment_config) + prediction = self._predict(df=df) control_outcome, treatment_outcome = prediction.iloc[1], prediction.iloc[0] ci_low = pd.Series(treatment_outcome["mean_ci_lower"] / control_outcome["mean_ci_upper"]) ci_high = pd.Series(treatment_outcome["mean_ci_upper"] / control_outcome["mean_ci_lower"]) @@ -165,19 +143,17 @@ def estimate_risk_ratio(self, adjustment_config: dict = None) -> EffectEstimate: "risk_ratio", pd.Series(treatment_outcome["mean"] / control_outcome["mean"]), ci_low, ci_high ) - def estimate_ate_calculated(self, adjustment_config: dict = None) -> EffectEstimate: + def estimate_ate_calculated(self, df: pd.DataFrame) -> EffectEstimate: """Estimate the ATE of the treatment on the outcome. That is, the change in outcome caused by changing the treatment variable from the control value to the treatment value. Here, we actually calculate the expected outcomes under control and treatment and divide one by the other. This allows for custom terms to be put in such as squares, inverses, products, etc. - :param: adjustment_config: The configuration of the adjustment set as a dict mapping variable names to - their values. N.B. Every variable in the adjustment set MUST have a value in - order to estimate the outcome under control and treatment. + :param df: The data to use. :return: The average treatment effect and the Wald confidence intervals. """ - prediction = self._predict(adjustment_config=adjustment_config) + prediction = self._predict(df=df) control_outcome, treatment_outcome = prediction.iloc[1], prediction.iloc[0] ci_low = pd.Series(treatment_outcome["mean_ci_lower"] - control_outcome["mean_ci_upper"]) ci_high = pd.Series(treatment_outcome["mean_ci_upper"] - control_outcome["mean_ci_lower"]) diff --git a/causal_testing/estimation/logistic_regression_estimator.py b/causal_testing/estimation/logistic_regression_estimator.py index ebdae2aa..f480c3c2 100644 --- a/causal_testing/estimation/logistic_regression_estimator.py +++ b/causal_testing/estimation/logistic_regression_estimator.py @@ -33,13 +33,14 @@ def add_modelling_assumptions(self): self.modelling_assumptions.append("The outcome must be binary.") self.modelling_assumptions.append("Independently and identically distributed errors.") - def estimate_unit_odds_ratio(self) -> EffectEstimate: + def estimate_unit_odds_ratio(self, df: pd.DataFrame) -> EffectEstimate: """Estimate the odds ratio of increasing the treatment by one. In logistic regression, this corresponds to the coefficient of the treatment of interest. + :param df: The data to use. :return: The odds ratio with confidence intervals. """ - model = self.fit_model(self.df) + model = self.fit_model(df) treatment_columns = self.treatment_columns(model) confidence_intervals = np.exp(model.conf_int(self.alpha).loc[treatment_columns]) diff --git a/causal_testing/estimation/multinomial_regression_estimator.py b/causal_testing/estimation/multinomial_regression_estimator.py index 3c1eb0d6..8bd5e7d8 100644 --- a/causal_testing/estimation/multinomial_regression_estimator.py +++ b/causal_testing/estimation/multinomial_regression_estimator.py @@ -28,13 +28,14 @@ def add_modelling_assumptions(self): super().add_modelling_assumptions() self.modelling_assumptions.append("Outcome is categorical.") - def estimate_unit_odds_ratio(self) -> EffectEstimate: + def estimate_unit_odds_ratio(self, df: pd.DataFrame) -> EffectEstimate: """Estimate the odds ratio of increasing the treatment by one. In logistic regression, this corresponds to the coefficient of the treatment of interest. + :param df: The data to use. :return: The odds ratio with confidence intervals. """ - model = self.fit_model(self.df) + model = self.fit_model(df) conf_int = model.conf_int(self.alpha) levels_of_interest = [ diff --git a/causal_testing/main.py b/causal_testing/main.py deleted file mode 100644 index 9d5ed77e..00000000 --- a/causal_testing/main.py +++ /dev/null @@ -1,639 +0,0 @@ -"""This module contains the classes that executes the Causal Testing Framework.""" - -import argparse -import json -import logging -from dataclasses import dataclass -from enum import Enum -from importlib.metadata import entry_points -from pathlib import Path -from typing import Any, Dict, List, Optional, Sequence, Union - -import numpy as np -import pandas as pd -from tqdm import tqdm - -from causal_testing.specification.causal_dag import CausalDAG -from causal_testing.specification.scenario import Scenario -from causal_testing.specification.variable import Input, Output -from causal_testing.testing.base_test_case import BaseTestCase -from causal_testing.testing.causal_test_adequacy import DataAdequacy -from causal_testing.testing.causal_test_case import CausalTestCase -from causal_testing.testing.causal_test_result import CausalTestResult - -logger = logging.getLogger(__name__) - - -class Command(Enum): - """ - Enum for supported CTF commands. - """ - - TEST = "test" - GENERATE = "generate" - DISCOVER = "discover" - - -@dataclass -class CausalTestingPaths: - """ - Class for managing paths for causal testing inputs and outputs. - - :param dag_path: Path to the DAG definition file - :param data_paths: List of paths to input data files - :param test_config_path: Path to the test configuration file - :param output_path: Path where test results will be written - """ - - dag_path: Path - data_paths: List[Path] - test_config_path: Path - output_path: Path - include_edges_path: Optional[Path] = None - exclude_edges_path: Optional[Path] = None - - def __init__( - self, - dag_path: Union[str, Path], - data_paths: List[Union[str, Path]], - test_config_path: Union[str, Path], - output_path: Union[str, Path], - include_edges_path: Optional[Union[str, Path]] = None, - exclude_edges_path: Optional[Union[str, Path]] = None, - ): - self.dag_path = Path(dag_path) - self.data_paths = [Path(p) for p in data_paths] - self.test_config_path = Path(test_config_path) - self.output_path = Path(output_path) - self.include_edges_path = Path(include_edges_path) if include_edges_path else None - self.exclude_edges_path = Path(exclude_edges_path) if exclude_edges_path else None - - def validate_paths(self) -> None: - """ - Validate existence of all input paths and writability of output path. - - :raises: FileNotFoundError if any required input file is missing. - """ - - if not self.dag_path.exists(): - raise FileNotFoundError(f"DAG file not found: {self.dag_path}") - - for data_path in self.data_paths: - if not data_path.exists(): - raise FileNotFoundError(f"Data file not found: {data_path}") - - if not self.test_config_path.exists(): - raise FileNotFoundError(f"Test configuration file not found: {self.test_config_path}") - - if not self.output_path.parent.exists(): - self.output_path.parent.mkdir(parents=True) - - if self.include_edges_path and not self.include_edges_path.exists(): - raise FileNotFoundError(f"Data file not found: {self.include_edges_path}") - - if self.exclude_edges_path and not self.exclude_edges_path.exists(): - raise FileNotFoundError(f"Data file not found: {self.exclude_edges_path}") - - -class CausalTestingFramework: - # pylint: disable=too-many-instance-attributes - """ - Main class for running causal tests. - - :param paths: CausalTestingPaths object containing required file paths - :param ignore_cycles: Flag to ignore cycles in the DAG - :param query: Optional query string to filter the input dataframe - - """ - - def __init__(self, paths: CausalTestingPaths, ignore_cycles: bool = False, query: Optional[str] = None): - self.paths = paths - self.ignore_cycles = ignore_cycles - self.query = query - - # These will be populated during setup - self.dag: Optional[CausalDAG] = None - self.data: Optional[pd.DataFrame] = None - self.variables: Dict[str, Any] = {"inputs": {}, "outputs": {}, "metas": {}} - self.scenario: Optional[Scenario] = None - self.test_cases: Optional[List[CausalTestCase]] = None - - def setup(self) -> None: - """ - Set up the framework by loading DAG, runtime csv data, creating the scenario and causal specification. - - :raises: FileNotFoundError if required files are missing - """ - - logger.info("Setting up Causal Testing Framework...") - - # Load and validate all paths - self.paths.validate_paths() - - # Load DAG - self.dag = self.load_dag() - - # Load data - self.data = self.load_data(self.query) - - # Create variables from DAG - self.create_variables() - - # Create scenario - self.scenario = Scenario( - list(self.variables["inputs"].values()) + list(self.variables["outputs"].values()), - {self.query} if self.query else None, - ) - - logger.info("Setup completed successfully") - - def load_dag(self) -> CausalDAG: - """ - Load the causal DAG from the specified file path. - """ - logger.info(f"Loading DAG from {self.paths.dag_path}") - dag = CausalDAG(str(self.paths.dag_path), ignore_cycles=self.ignore_cycles) - logger.info(f"DAG loaded with {len(dag.nodes)} nodes and {len(dag.edges)} edges") - return dag - - def _read_dataframe(self, data_path): - if str(data_path).endswith(".csv"): - return pd.read_csv(data_path) - if str(data_path).endswith(".pqt"): - return pd.read_parquet(data_path) - raise ValueError(f"Invalid file type {data_path}. Can only read CSV (.csv) or parquet (.pqt) files.") - - def load_data(self, query: Optional[str] = None) -> pd.DataFrame: - """Load and combine all data sources with optional filtering. - - :param query: Optional pandas query string to filter the loaded data - :return: Combined pandas DataFrame containing all loaded and filtered data - """ - logger.info(f"Loading data from {len(self.paths.data_paths)} source(s)") - - dfs = [self._read_dataframe(data_path) for data_path in self.paths.data_paths] - data = pd.concat(dfs, axis=0, ignore_index=True) - logger.info(f"Initial data shape: {data.shape}") - - if query: - logger.info(f"Attempting to apply query: '{query}'") - data = data.query(query) - - return data - - def create_variables(self) -> None: - """ - Create variable objects from DAG nodes based on their connectivity. - """ - for node_name, node_data in self.dag.nodes(data=True): - if node_name not in self.data.columns and not node_data.get("hidden", False): - raise ValueError(f"Node {node_name} missing from data. Should it be marked as hidden?") - - dtype = self.data.dtypes.get(node_name) - - # If node has no incoming edges, it's an input - if self.dag.in_degree(node_name) == 0: - self.variables["inputs"][node_name] = Input(name=node_name, datatype=dtype) - - # Otherwise it's an output - if self.dag.in_degree(node_name) > 0: - self.variables["outputs"][node_name] = Output(name=node_name, datatype=dtype) - - def load_tests(self) -> None: - """ - Load and prepare test configurations from file. - """ - logger.info(f"Loading test configurations from {self.paths.test_config_path}") - - with open(self.paths.test_config_path, "r", encoding="utf-8") as f: - test_configs = json.load(f) - - self.test_cases = self.create_test_cases(test_configs) - - def create_base_test(self, test: dict) -> BaseTestCase: - """ - Create base test case from test configuration. - - :param test: Dictionary containing test configuration parameters - - :return: BaseTestCase object - :raises: KeyError if required variables are not found in inputs or outputs - """ - treatment_name = test["treatment_variable"] - outcome_name = next(iter(test["expected_effect"].keys())) - - # Look for treatment variable in both inputs and outputs - treatment_var = self.variables["inputs"].get(treatment_name) or self.variables["outputs"].get(treatment_name) - if not treatment_var: - raise KeyError(f"Treatment variable '{treatment_name}' not found in inputs or outputs") - - # Look for outcome variable in both inputs and outputs - outcome_var = self.variables["inputs"].get(outcome_name) or self.variables["outputs"].get(outcome_name) - if not outcome_var: - raise KeyError(f"Outcome variable '{outcome_name}' not found in inputs or outputs") - - return BaseTestCase( - treatment_variable=treatment_var, outcome_variable=outcome_var, effect=test.get("effect", "total") - ) - - def create_test_cases(self, test_configs: dict) -> List[CausalTestCase]: - """Create test case objects from configuration dictionary. - - :param test_configs: Dictionary containing test configurations - - :return: List of CausalTestCase objects containing the initialised test cases - :raises: KeyError if required variables are not found - :raises: ValueError if invalid test configuration is provided - """ - test_cases = [] - - for test in test_configs.get("tests", []): - if test.get("skip", False): - continue - - # Create base test case - base_test = self.create_base_test(test) - - # Create causal test case - causal_test = self.create_causal_test(test, base_test) - test_cases.append(causal_test) - - return test_cases - - def create_causal_test(self, test: dict, base_test: BaseTestCase) -> CausalTestCase: - """ - Create causal test case from test configuration and base test. - - :param test: Dictionary containing test configuration parameters - :param base_test: BaseTestCase object - - :return: CausalTestCase object - :raises: ValueError if invalid estimator or configuration is provided - """ - 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: - raise ValueError( - f"Unsupported estimator {test['estimator']}. 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." - ) - - # Handle global queries - # Test-specific queries are handled by the estimator as not all estimators support them - filtered_df = self.data - if self.query: - filtered_df = self.data.query(self.query) - - # Create the estimator with correct parameters - estimator_class = estimator_map.get(test["estimator"]).load() - estimator_kwargs = test.get("estimator_kwargs", {}) - if "query" in test: - estimator_kwargs["query"] = test["query"] - estimator = estimator_class( - base_test_case=base_test, - treatment_value=test.get("treatment_value"), - control_value=test.get("control_value"), - adjustment_set=test.get( - "adjustment_set", - self.dag.identification(base_test, self.scenario.hidden_variables()), - ), - df=filtered_df, - alpha=test.get("alpha", 0.05), - **estimator_kwargs, - ) - - # Get effect type and create expected effect - effect_type = test["expected_effect"][base_test.outcome_variable.name] - if effect_type not in effect_map: - raise ValueError( - f"Unsupported causal effect {effect_type}. 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()(**test.get("effect_kwargs", {})) - - return CausalTestCase( - base_test_case=base_test, - expected_causal_effect=expected_effect, - estimate_type=test.get("estimate_type", "ate"), - estimate_params=test.get("estimate_params"), - estimator=estimator, - ) - - def run_tests_in_batches( - self, batch_size: int = 100, silent: bool = False, adequacy: bool = False, bootstrap_size: int = 100 - ) -> List[CausalTestResult]: - """ - Run tests in batches to reduce memory usage. - - :param batch_size: Number of tests to run in each batch - :param silent: Whether to suppress errors - :param adequacy: Whether to calculate causal test adequacy (defaults to False) - :param bootstrap_size: The number of bootstrap samples to use when calculating causal test adequacy - (defaults to 100) - :return: List of all test results - :raises: ValueError if no tests are loaded - """ - logger.info("Running causal tests in batches...") - - if not self.test_cases: - raise ValueError("No tests loaded. Call load_tests() first.") - - num_tests = len(self.test_cases) - num_batches = int(np.ceil(num_tests / batch_size)) - - logger.info(f"Processing {num_tests} tests in {num_batches} batches of up to {batch_size} tests each") - with tqdm(total=num_tests, desc="Overall progress", mininterval=0.1) as progress: - # Process each batch - for batch_idx in range(num_batches): - start_idx = batch_idx * batch_size - end_idx = min(start_idx + batch_size, num_tests) - - logger.info(f"Processing batch {batch_idx + 1} of {num_batches} (tests {start_idx} to {end_idx - 1})") - - # Get current batch of tests - current_batch = self.test_cases[start_idx:end_idx] - - # Process the current batch - batch_results = [] - for test_case in current_batch: - try: - result = test_case.execute_test() - if adequacy: - result.adequacy = DataAdequacy(test_case=test_case, bootstrap_size=bootstrap_size) - result.adequacy.measure_adequacy() - - batch_results.append(result) - # pylint: disable=broad-exception-caught - except Exception as e: - if not silent: - logger.error(f"Type or attribute error in test: {str(e)}") - raise - batch_results.append( - CausalTestResult(effect_estimate=None, estimator=test_case.estimator, error_message=str(e)) - ) - - progress.update(1) - - yield batch_results - logger.info(f"Completed processing in {num_batches} batches") - - def run_tests( - self, silent: bool = False, adequacy: bool = False, bootstrap_size: int = 100 - ) -> List[CausalTestResult]: - """ - Run all test cases and return their results. - - :param silent: Whether to suppress errors - :param adequacy: Whether to calculate causal test adequacy (defaults to False) - :param bootstrap_size: The number of bootstrap samples to use when calculating causal test adequacy - (defaults to 100) - - :return: List of CausalTestResult objects - :raises: ValueError if no tests are loaded - :raises: Exception if test execution fails - """ - logger.info("Running causal tests...") - - if not self.test_cases: - raise ValueError("No tests loaded. Call load_tests() first.") - - results = [] - for test_case in tqdm(self.test_cases): - try: - result = test_case.execute_test() - if adequacy: - result.adequacy = DataAdequacy(test_case=test_case, bootstrap_size=bootstrap_size) - result.adequacy.measure_adequacy() - results.append(result) - # pylint: disable=broad-exception-caught - except Exception as e: - if not silent: - logger.error(f"Error running test {test_case}: {str(e)}") - raise - result = CausalTestResult(estimator=test_case.estimator, effect_estimate=None, error_message=str(e)) - results.append(result) - logger.info(f"Test errored: {test_case}") - - return results - - def save_results(self, results: List[CausalTestResult], output_path: str = None) -> list: - """Save test results to JSON file in the expected format.""" - if output_path is None: - output_path = self.paths.output_path - logger.info(f"Saving results to {output_path}") - - # Create parent directory if it doesn't exist - Path(output_path).parent.mkdir(parents=True, exist_ok=True) - - # Load original test configs to preserve test metadata - with open(self.paths.test_config_path, "r", encoding="utf-8") as f: - test_configs = json.load(f) - - json_results = [] - result_index = 0 - - for test_config in test_configs["tests"]: - # Create a base output first of common entries - base_output = { - "name": test_config["name"], - "estimate_type": test_config["estimate_type"], - "effect": test_config.get("effect", "direct"), - "treatment_variable": test_config["treatment_variable"], - "expected_effect": test_config["expected_effect"], - "alpha": test_config.get("alpha", 0.05), - } - if test_config.get("skip", False): - # Include those skipped test entry without execution results - output = { - **base_output, - "formula": test_config.get("formula"), - "skip": True, - "passed": None, - "result": { - "status": "skipped", - "reason": "Test marked as skip:true in the causal test config file.", - }, - } - else: - # Add executed test with actual results - test_case = self.test_cases[result_index] - result = results[result_index] - 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": result.estimator.formula if hasattr(result.estimator, "formula") else None, - "skip": False, - "passed": test_passed, - "result": ( - { - "treatment": result.estimator.base_test_case.treatment_variable.name, - "outcome": result.estimator.base_test_case.outcome_variable.name, - "adjustment_set": list(result.adjustment_set) if result.adjustment_set else [], - } - | result.effect_estimate.to_dict() - | (result.adequacy.to_dict() if result.adequacy else {}) - if result.effect_estimate - else {"status": "error", "reason": result.error_message} - ), - } - - json_results.append(output) - - # Save to file - with open(output_path, "w", encoding="utf-8") as f: - json.dump(json_results, f, indent=2) - - logger.info("Results saved successfully") - return json_results - - -def setup_logging(level: str) -> None: - """Set up logging configuration.""" - logging.basicConfig( - level=level, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S" - ) - - -def parse_args(args: Optional[Sequence[str]] = None) -> argparse.Namespace: - """Parse command line arguments.""" - main_parser = argparse.ArgumentParser( - add_help=True, - description="Causal Testing Framework - " - "A causal inference-driven framework for functional black-box testing of complex software.", - ) - - main_parser.add_argument( - "-l", - "--log_level", - default="WARNING", - type=str.upper, - choices=["NONE", "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], - help="Set the logging level (default: WARNING).", - ) - - subparsers = main_parser.add_subparsers( - help="The action you want to run - call `causal_testing {action} -h` for further details", dest="command" - ) - - # Generation - parser_generate = subparsers.add_parser(Command.GENERATE.value, help="Generate causal tests from a DAG") - parser_generate.add_argument("-D", "--dag-path", help="Path to the DAG file (.dot)", required=True) - parser_generate.add_argument("-o", "--output", help="Path for output file (.json)", required=True) - parser_generate.add_argument( - "-e", - "--estimator", - help="The name of the estimator class to use when evaluating tests (defaults to LinearRegressionEstimator)", - default="LinearRegressionEstimator", - ) - parser_generate.add_argument( - "-T", - "--effect-type", - help="The effect type to estimate {direct, total}", - default="direct", - ) - parser_generate.add_argument( - "-E", - "--estimate-type", - help="The estimate type to use when evaluating tests (defaults to coefficient)", - default="coefficient", - ) - parser_generate.add_argument( - "-i", "--ignore-cycles", help="Ignore cycles in DAG", action="store_true", default=False - ) - parser_generate.add_argument( - "--threads", "-t", type=int, help="The number of parallel threads to use.", required=False, default=0 - ) - - # Testing - parser_test = subparsers.add_parser(Command.TEST.value, help="Run causal tests") - parser_test.add_argument("-D", "--dag-path", help="Path to the DAG file (.dot)", required=True) - parser_test.add_argument("-o", "--output", help="Path for output file (.json)", required=True) - parser_test.add_argument("-i", "--ignore-cycles", help="Ignore cycles in DAG", action="store_true", default=False) - parser_test.add_argument("-d", "--data-paths", help="Paths to data files (.csv)", nargs="+", required=True) - parser_test.add_argument("-t", "--test-config", help="Path to test configuration file (.json)", required=True) - parser_test.add_argument("-q", "--query", help="Query string to filter data (e.g. 'age > 18')", type=str) - parser_test.add_argument( - "-a", "--adequacy", help="Calculate causal test adequacy for each test case", action="store_true", default=False - ) - parser_test.add_argument( - "-b", - "--adequacy-bootstrap-size", - dest="bootstrap_size", - help="Number of bootstrap samples for causal test adequacy. Defaults to 100", - type=int, - ) - parser_test.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, - ) - parser_test.add_argument( - "--batch-size", - type=int, - default=0, - help="Run tests in batches of the specified size (default: 0, which means no batching)", - ) - - # 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", - help="The name of the technique to use. Currently supported are 'HillClimberDiscovery' and 'NSGADiscovery'", - required=True, - ) - parser_discover.add_argument( - "-V", - "--variables", - help="The subset of variables from the data to consider. Defaults to all.", - nargs="*", - default=[], - ) - parser_discover.add_argument("-o", "--output", help="Path for output DAG file (.dot)", required=True) - parser_discover.add_argument( - "-i", "--include-edges", help="Path to file containing edges to include", required=False - ) - parser_discover.add_argument( - "-e", "--exclude-edges", help="Path to file containing edges to exclude", required=False - ) - parser_discover.add_argument( - "--technique-kwargs", - help="Keywords for the discovery technique. These should be specified as `arg1=value1 arg2=value2...`.", - nargs="*", - default=[], - ) - - args = main_parser.parse_args(args) - - # Assume the user wants test adequacy if they're setting bootstrap_size - if getattr(args, "bootstrap_size", None) is not None: - args.adequacy = True - if getattr(args, "adequacy", False) and getattr(args, "bootstrap_size", None) is None: - # Need this here rather than a default value because otherwise the above always sets adequacy to True - args.bootstrap_size = 100 - - args.command = Command(args.command) - return args diff --git a/causal_testing/surrogate/causal_surrogate_assisted.py b/causal_testing/surrogate/causal_surrogate_assisted.py index ab308c66..4cfb6aba 100644 --- a/causal_testing/surrogate/causal_surrogate_assisted.py +++ b/causal_testing/surrogate/causal_surrogate_assisted.py @@ -34,11 +34,15 @@ class SearchAlgorithm(ABC): # pylint: disable=too-few-public-methods space to be searched""" @abstractmethod - def search(self, surrogate_models: list[CubicSplineRegressionEstimator], scenario: Scenario) -> list: + def search( + self, surrogate_models: list[CubicSplineRegressionEstimator], scenario: Scenario, df: pd.DataFrame + ) -> list: """Function which implements a search routine which searches for the optimal fitness value for the specified scenario :param surrogate_models: The surrogate models to be searched - :param scenario: The modelling scenario""" + :param scenario: The modelling scenario + :param df: The data to use + """ class Simulator(ABC): @@ -91,8 +95,8 @@ def execute( :return: tuple containing SimulationResult or str, execution number and dataframe""" for i in range(max_executions): - surrogate_models = self.generate_surrogates(df) - candidate_test_case, _, surrogate_model = self.search_algorithm.search(surrogate_models, self.scenario) + surrogate_models = self.generate_surrogates() + candidate_test_case, _, surrogate_model = self.search_algorithm.search(surrogate_models, self.scenario, df) self.simulator.startup() test_result = self.simulator.run_with_config(candidate_test_case) @@ -120,9 +124,8 @@ def execute( logger.info("No fault found") return "No fault found", i + 1, df - def generate_surrogates(self, df: pd.DataFrame) -> list[CubicSplineRegressionEstimator]: + def generate_surrogates(self) -> list[CubicSplineRegressionEstimator]: """Generate a surrogate model for each edge of the DAG that specifies it is included in the DAG metadata. - :param df: An dataframe which contains data relevant to the specified scenario :return: A list of surrogate models """ surrogate_models = [] @@ -144,7 +147,6 @@ def generate_surrogates(self, df: pd.DataFrame) -> list[CubicSplineRegressionEst 0, minimal_adjustment_set, 4, - df=df, expected_relationship=edge_metadata["expected"], ) surrogate_models.append(surrogate) diff --git a/causal_testing/surrogate/surrogate_search_algorithms.py b/causal_testing/surrogate/surrogate_search_algorithms.py index 19d05e35..36371155 100644 --- a/causal_testing/surrogate/surrogate_search_algorithms.py +++ b/causal_testing/surrogate/surrogate_search_algorithms.py @@ -4,6 +4,7 @@ from operator import itemgetter +import pandas as pd from pygad import GA from causal_testing.estimation.cubic_spline_estimator import CubicSplineRegressionEstimator @@ -27,7 +28,9 @@ def __init__(self, delta=0.05, config: dict = None) -> None: } # pylint: disable=too-many-locals - def search(self, surrogate_models: list[CubicSplineRegressionEstimator], scenario: Scenario) -> list: + def search( + self, surrogate_models: list[CubicSplineRegressionEstimator], scenario: Scenario, df: pd.DataFrame + ) -> list: solutions = [] for surrogate_model in surrogate_models: @@ -44,7 +47,8 @@ def fitness_function(ga, solution, idx): # pylint: disable=unused-argument for i, adjustment in enumerate(surrogate_model.adjustment_set): adjustment_dict[adjustment] = solution[i + 1] - ate = surrogate_model.estimate_ate_calculated(adjustment_dict).value + surrogate_model.adjustment_config = adjustment_dict + ate = surrogate_model.estimate_ate_calculated(df).value if len(ate) > 1: raise ValueError( "Multiple ate values provided but currently only single values supported in this method" @@ -96,7 +100,6 @@ def create_gene_types(surrogate_model: CubicSplineRegressionEstimator, scenario: var_space[adj] = {} for relationship in list(scenario.constraints): - print(relationship) rel_split = str(relationship).split(" ") if rel_split[0] in var_space: diff --git a/causal_testing/testing/causal_effect.py b/causal_testing/testing/causal_effect.py index 8eb3f26c..efa34732 100644 --- a/causal_testing/testing/causal_effect.py +++ b/causal_testing/testing/causal_effect.py @@ -126,9 +126,13 @@ def apply(self, res: CausalTestResult) -> bool: if len(res.effect_estimate.value) > 1: raise ValueError("Positive Effects are currently only supported on single float datatypes") if res.effect_estimate.type in {"ate", "coefficient"}: - return bool(res.effect_estimate.value[0] > 0) + return any( + 0 < ci_low < ci_high for ci_low, ci_high in zip(res.effect_estimate.ci_low, res.effect_estimate.ci_high) + ) if res.effect_estimate.type in ["risk_ratio", "unit_odds_ratio"]: - return bool(res.effect_estimate.value[0] > 1) + return any( + 1 < ci_low < ci_high for ci_low, ci_high in zip(res.effect_estimate.ci_low, res.effect_estimate.ci_high) + ) raise ValueError(f"Test Value type {res.effect_estimate.type} is not valid for this CausalEffect") @@ -140,8 +144,12 @@ def apply(self, res: CausalTestResult) -> bool: if len(res.effect_estimate.value) > 1: raise ValueError("Negative Effects are currently only supported on single float datatypes") if res.effect_estimate.type in {"ate", "coefficient"}: - return bool(res.effect_estimate.value[0] < 0) + return any( + ci_low < ci_high < 0 for ci_low, ci_high in zip(res.effect_estimate.ci_low, res.effect_estimate.ci_high) + ) if res.effect_estimate.type in ["risk_ratio", "unit_odds_ratio"]: - return bool(res.effect_estimate.value[0] < 1) + return any( + ci_low < ci_high < 1 for ci_low, ci_high in zip(res.effect_estimate.ci_low, res.effect_estimate.ci_high) + ) # Dead code but necessary for pylint raise ValueError(f"Test Value type {res.effect_estimate.type} is not valid for this CausalEffect") diff --git a/causal_testing/testing/causal_test_adequacy.py b/causal_testing/testing/causal_test_adequacy.py deleted file mode 100644 index 6dc449c4..00000000 --- a/causal_testing/testing/causal_test_adequacy.py +++ /dev/null @@ -1,125 +0,0 @@ -""" -This module contains code to measure various aspects of causal test adequacy. -""" - -import logging -from copy import deepcopy -from itertools import combinations - -import pandas as pd - -from causal_testing.specification.causal_dag import CausalDAG -from causal_testing.testing.causal_test_case import CausalTestCase - -logger = logging.getLogger(__name__) - - -class DAGAdequacy: - """ - Measures the adequacy of a given DAG by hos many edges and independences are tested. - """ - - def __init__( - self, - causal_dag: CausalDAG, - test_suite: list[CausalTestCase], - ): - self.causal_dag = causal_dag - self.test_suite = test_suite - self.tested_pairs = None - self.pairs_to_test = None - self.untested_pairs = None - self.dag_adequacy = None - - def measure_adequacy(self): - """ - Calculate the adequacy measurement, and populate the `dag_adequacy` field. - """ - self.pairs_to_test = set(combinations(self.causal_dag.nodes, 2)) - self.tested_pairs = set() - - for n1, n2 in self.pairs_to_test: - if (n1, n2) in self.causal_dag.edges(): - if any((t.treatment_variable, t.outcome_variable) == (n1, n2) for t in self.test_suite): - self.tested_pairs.add((n1, n2)) - else: - # Causal independences are not order dependent - if any((t.treatment_variable, t.outcome_variable) in {(n1, n2), (n2, n1)} for t in self.test_suite): - self.tested_pairs.add((n1, n2)) - - self.untested_pairs = self.pairs_to_test.difference(self.tested_pairs) - self.dag_adequacy = len(self.tested_pairs) / len(self.pairs_to_test) - - def to_dict(self): - """Returns the adequacy object as a dictionary.""" - return { - "causal_dag": self.causal_dag, - "test_suite": self.test_suite, - "tested_pairs": self.tested_pairs, - "pairs_to_test": self.pairs_to_test, - "untested_pairs": self.untested_pairs, - "dag_adequacy": self.dag_adequacy, - } - - -class DataAdequacy: - """ - Measures the adequacy of a given test according to the Fisher kurtosis of the bootstrapped result. - - * Positive kurtoses indicate the model doesn't have enough data so is unstable. - * Negative kurtoses indicate the model doesn't have enough data, but is too stable, - indicating that the spread of inputs is insufficient. - * Zero kurtosis is optimal. - """ - - # pylint: disable=too-many-instance-attributes - def __init__( - self, - test_case: CausalTestCase, - bootstrap_size: int = 100, - group_by=None, - ): - self.test_case = test_case - self.kurtosis = None - self.passing = None - self.results = None - self.successful = None - self.bootstrap_size = bootstrap_size - self.group_by = group_by - - def measure_adequacy(self): - """ - Calculate the adequacy measurement, and populate the data_adequacy field. - """ - results = [] - outcomes = [] - for i in range(self.bootstrap_size): - estimator = deepcopy(self.test_case.estimator) - - if self.group_by is not None: - ids = pd.Series(estimator.df[self.group_by].unique()) - ids = ids.sample(len(ids), replace=True, random_state=i) - estimator.df = estimator.df[estimator.df[self.group_by].isin(ids)] - else: - estimator.df = estimator.df.sample(len(estimator.df), replace=True, random_state=i) - result = self.test_case.execute_test(estimator) - outcomes.append(self.test_case.expected_causal_effect.apply(result)) - results.append(result.effect_estimate.to_df()) - results = pd.concat(results) - results["var"] = results.index - results["passed"] = outcomes - - self.results = results - self.kurtosis = results.groupby("var")["effect_estimate"].apply(lambda x: x.kurtosis()) - self.passing = sum(filter(lambda x: x is not None, outcomes)) - self.successful = sum(x is not None for x in outcomes) - - def to_dict(self): - """Returns the adequacy object as a dictionary.""" - return { - "kurtosis": self.kurtosis.to_dict(), - "bootstrap_size": self.bootstrap_size, - "passing": self.passing, - "successful": self.successful, - "results": self.results.reset_index(drop=True).to_dict(), - } diff --git a/causal_testing/testing/causal_test_case.py b/causal_testing/testing/causal_test_case.py index 5049f739..8b32df32 100644 --- a/causal_testing/testing/causal_test_case.py +++ b/causal_testing/testing/causal_test_case.py @@ -2,10 +2,14 @@ import logging +import numpy as np +import pandas as pd + from causal_testing.estimation.abstract_estimator import Estimator from causal_testing.testing.base_test_case import BaseTestCase from causal_testing.testing.causal_effect import CausalEffect from causal_testing.testing.causal_test_result import CausalTestResult +from causal_testing.testing.data_adequacy import DataAdequacy logger = logging.getLogger(__name__) @@ -29,8 +33,10 @@ def __init__( base_test_case: BaseTestCase, expected_causal_effect: CausalEffect, estimate_type: str = "ate", - estimate_params: dict = None, estimator: type(Estimator) = None, + name: str = None, + query: str = None, + skip: bool = False, ): self.base_test_case = base_test_case self.expected_causal_effect = expected_causal_effect @@ -38,33 +44,111 @@ def __init__( self.treatment_variable = base_test_case.treatment_variable self.estimate_type = estimate_type self.estimator = estimator - if estimate_params is None: - self.estimate_params = {} - else: - self.estimate_params = estimate_params self.effect = base_test_case.effect + self.result = None + self.name = name + self.query = query + self.skip = skip - def execute_test(self, estimator: type(Estimator) = None) -> CausalTestResult: + def measure_adequacy( + self, + df: pd.DataFrame, + bootstrap_size: int = 100, + group_by: str = None, + ) -> DataAdequacy: """ - Execute a causal test case and return the causal test result. - :param estimator: An alternative estimator. Defaults to `self.estimator`. This parameter is useful when you want - to execute a test with different data or a different equational form, but don't want to redefine the whole test - case. + Calculate the adequacy measurement, and populate the data_adequacy field. + :param df: The original dataset to use. + :param bootstrap_size: The number of bootstrap samples to use. (Defaults to 100) + :param group_by: For IPCWEstimator - the "id" column to ensure that entire individuals are sampled rather than + random rows. + """ + results = [] + outcomes = [] + for i in range(bootstrap_size): + if group_by is not None: + ids = pd.Series(df[group_by].unique()) + ids = ids.sample(len(ids), replace=True, random_state=i) + df = df[df[group_by].isin(ids)] + else: + df = df.sample(len(df), replace=True, random_state=i) + try: + result = self.estimate_effect(df) + outcomes.append(self.expected_causal_effect.apply(result)) + results.append(result.effect_estimate.to_df()) + # Could get a variety of exceptions here due to insufficient/badly formed data + # We don't want these to stop execution + except Exception: # pylint: disable=W0718 + pass + + results = pd.concat(results) + + results["var"] = results.index + results["passed"] = outcomes + + return DataAdequacy( + results=results, + kurtosis=results.groupby("var")["effect_estimate"].apply(lambda x: x.kurtosis()), + passing=sum(filter(lambda x: x is not None, outcomes)), + successful=sum(x is not None for x in outcomes), + ) + + def execute_test( + self, + df: pd.DataFrame, + estimate_params: dict[str, any] = None, + adequacy: bool = False, + suppress_estimation_errors: bool = False, + bootstrap_size: int = 100, + group_by: str = None, + ): + """ + Execute a causal test case. + :param df: The data to use. + :param estimate_params: Extra parameters for the estimate calculation. + :param adequacy: Set to True to calculate the causal test adequacy associated with the effect estimate. + :param suppress_estimation_errors: Set to True to suppress estimation errors. (Defaults to False) + :param bootstrap_size: The number of bootstrap samples to use. (Defaults to 100) + :param group_by: For IPCWEstimator - the "id" column to ensure that entire individuals are sampled rather than + random rows. :return causal_test_result: A CausalTestResult for the executed causal test case. """ + if not self.skip: + self.result = self.estimate_effect( + df=df, estimate_params=estimate_params, suppress_estimation_errors=suppress_estimation_errors + ) + if adequacy: + self.result.adequacy = self.measure_adequacy(df=df, bootstrap_size=bootstrap_size, group_by=group_by) - if estimator is None: - estimator = self.estimator + def estimate_effect( + self, + df: pd.DataFrame, + estimate_params: dict[str, any] = None, + suppress_estimation_errors: bool = False, + ) -> CausalTestResult: + """ + Execute a causal test case and return the causal test result. - if not hasattr(estimator, f"estimate_{self.estimate_type}"): - raise AttributeError(f"{estimator.__class__} has no {self.estimate_type} method.") - estimate_effect = getattr(estimator, f"estimate_{self.estimate_type}") - effect_estimate = estimate_effect(**self.estimate_params) - return CausalTestResult( - estimator=estimator, - effect_estimate=effect_estimate, - ) + :param df: The data to use. + :param estimate_params: Extra parameters for the estimate calculation. + :param suppress_estimation_errors: Set to True to suppress estimation errors. (Defaults to False) + :return causal_test_result: A CausalTestResult for the executed causal test case. + """ + if self.query: + df = df.query(self.query) + if not hasattr(self.estimator, f"estimate_{self.estimate_type}"): + raise AttributeError(f"{self.estimator.__class__} has no {self.estimate_type} method.") + estimate_effect = getattr(self.estimator, f"estimate_{self.estimate_type}") + try: + effect_estimate = estimate_effect(df, **(estimate_params if estimate_params is not None else {})) + return CausalTestResult( + effect_estimate=effect_estimate, + ) + except (np.linalg.LinAlgError, ValueError) as e: + if not suppress_estimation_errors: + raise e + return CausalTestResult(effect_estimate=None, error_message=str(e)) def __str__(self): treatment_config = {self.treatment_variable.name: self.estimator.treatment_value} diff --git a/causal_testing/testing/causal_test_result.py b/causal_testing/testing/causal_test_result.py index 8ba7afff..79d17506 100644 --- a/causal_testing/testing/causal_test_result.py +++ b/causal_testing/testing/causal_test_result.py @@ -1,9 +1,11 @@ """This module contains the CausalTestResult class, which is a container for the results of a causal test.""" -from causal_testing.estimation.abstract_estimator import Estimator +from dataclasses import dataclass + from causal_testing.estimation.effect_estimate import EffectEstimate +@dataclass class CausalTestResult: """A container to hold the results of a causal test case. Every causal test case provides a point estimate of the ATE, given a particular treatment, outcome, and adjustment set. Some but not all estimators can provide @@ -11,61 +13,10 @@ class CausalTestResult: def __init__( self, - estimator: Estimator, effect_estimate: EffectEstimate, adequacy=None, error_message: str = None, ): - self.estimator = estimator self.adequacy = adequacy - if estimator.adjustment_set: - self.adjustment_set = estimator.adjustment_set - else: - self.adjustment_set = set() self.effect_estimate = effect_estimate self.error_message = error_message - - def __str__(self): - result_str = str(self.effect_estimate.value.to_dict()) - treatment = self.estimator.base_test_case.treatment_variable.name - base_str = ( - f"Causal Test Result\n==============\n" - f"Treatment: {treatment}\n" - f"Control value: {self.estimator.control_value}\n" - f"Treatment value: {self.estimator.treatment_value}\n" - f"Outcome: {self.estimator.base_test_case.outcome_variable.name}\n" - f"Adjustment set: {self.adjustment_set}\n" - ) - if hasattr(self.estimator, "formula"): - base_str += f"Formula: {self.estimator.formula}\n" - base_str += f"{self.effect_estimate.type}: {result_str}\n" - confidence_str = "" - if self.effect_estimate.ci_valid(): - ci_str = f"CI low: {self.effect_estimate.ci_low.to_dict}\n" - ci_str += f"CI high: {self.effect_estimate.ci_high.to_dict}\n" - confidence_str += f"Confidence intervals:{ci_str}\n" - confidence_str += f"Alpha:{self.estimator.alpha}\n" - adequacy_str = "" - if self.adequacy: - adequacy_str = str(self.adequacy) - return base_str + confidence_str + adequacy_str - - def to_dict(self, json=False): - """Return result contents as a dictionary - :return: Dictionary containing contents of causal_test_result - """ - base_dict = { - "treatment": ( - self.estimator.base_test_case.treatment_variable.name - if self.estimator.base_test_case.treatment_variable is not None - else None - ), - "control_value": self.estimator.control_value, - "treatment_value": self.estimator.treatment_value, - "outcome": self.estimator.base_test_case.outcome_variable.name, - "adjustment_set": list(self.adjustment_set) if json else self.adjustment_set, - "effect_measure": self.effect_estimate.type, - } | self.effect_estimate.to_dict() - if self.adequacy: - base_dict["adequacy"] = self.adequacy.to_dict() - return base_dict diff --git a/causal_testing/testing/dag_adequacy.py b/causal_testing/testing/dag_adequacy.py new file mode 100644 index 00000000..80d5c720 --- /dev/null +++ b/causal_testing/testing/dag_adequacy.py @@ -0,0 +1,59 @@ +""" +This module contains code to measure various aspects of causal test adequacy. +""" + +import logging +from itertools import combinations + +from causal_testing.specification.causal_dag import CausalDAG +from causal_testing.testing.causal_test_case import CausalTestCase + +logger = logging.getLogger(__name__) + + +class DAGAdequacy: + """ + Measures the adequacy of a given DAG by hos many edges and independences are tested. + """ + + def __init__( + self, + causal_dag: CausalDAG, + test_suite: list[CausalTestCase], + ): + self.causal_dag = causal_dag + self.test_suite = test_suite + self.tested_pairs = None + self.pairs_to_test = None + self.untested_pairs = None + self.dag_adequacy = None + + def measure_adequacy(self): + """ + Calculate the adequacy measurement, and populate the `dag_adequacy` field. + """ + self.pairs_to_test = set(combinations(self.causal_dag.nodes, 2)) + self.tested_pairs = set() + + for n1, n2 in self.pairs_to_test: + if (n1, n2) in self.causal_dag.edges(): + if any((t.treatment_variable, t.outcome_variable) == (n1, n2) for t in self.test_suite): + self.tested_pairs.add((n1, n2)) + else: + # Causal independences are not order dependent + if any((t.treatment_variable, t.outcome_variable) in {(n1, n2), (n2, n1)} for t in self.test_suite): + self.tested_pairs.add((n1, n2)) + + self.untested_pairs = self.pairs_to_test.difference(self.tested_pairs) + self.dag_adequacy = len(self.tested_pairs) / len(self.pairs_to_test) + + def to_dict(self): + """Returns the adequacy object as a dictionary.""" + return { + "causal_dag": self.causal_dag, + "test_suite": self.test_suite, + "tested_pairs": self.tested_pairs, + "pairs_to_test": self.pairs_to_test, + "untested_pairs": self.untested_pairs, + "dag_adequacy": self.dag_adequacy, + } diff --git a/causal_testing/testing/data_adequacy.py b/causal_testing/testing/data_adequacy.py new file mode 100644 index 00000000..0f9c5105 --- /dev/null +++ b/causal_testing/testing/data_adequacy.py @@ -0,0 +1,40 @@ +""" +This module contains code to measure various aspects of causal test adequacy. +""" + +import logging + +logger = logging.getLogger(__name__) + + +class DataAdequacy: + """ + Measures the adequacy of a given test according to the Fisher kurtosis of the bootstrapped result. + + * Positive kurtoses indicate the model doesn't have enough data so is unstable. + * Negative kurtoses indicate the model doesn't have enough data, but is too stable, + indicating that the spread of inputs is insufficient. + * Zero kurtosis is optimal. + """ + + # pylint: disable=too-many-instance-attributes + def __init__( + self, + kurtosis=None, + passing=None, + results=None, + successful=None, + ): + self.kurtosis = kurtosis + self.passing = passing + self.results = results + self.successful = successful + + def to_dict(self): + """Returns the adequacy object as a dictionary.""" + return { + "kurtosis": self.kurtosis.to_dict(), + "passing": self.passing, + "successful": self.successful, + "results": self.results.reset_index(drop=True).to_dict(), + } diff --git a/docs/source/tutorials/poisson_line_process/poisson_line_process_tutorial.ipynb b/docs/source/tutorials/poisson_line_process/poisson_line_process_tutorial.ipynb index 353622a7..4d9b4cc9 100644 --- a/docs/source/tutorials/poisson_line_process/poisson_line_process_tutorial.ipynb +++ b/docs/source/tutorials/poisson_line_process/poisson_line_process_tutorial.ipynb @@ -412,7 +412,6 @@ "for (control, treatment) in zip(control_values, treatment_values): # Simultaneously loop over control and treatment\n", " \n", " estimator=LinearRegressionEstimator(\n", - " df=df, # Pass in the dataframe\n", " base_test_case=base_test_case, # Base test case we created above\n", " treatment_value=treatment, # Doubled intensity values\n", " control_value=control, # Baseline intensity values\n", @@ -427,16 +426,16 @@ " estimator = estimator) # Pass in the estimator we created above\n", "\n", "\n", - " test_results = causal_test_case.execute_test() # Execute the tests\n", + " causal_test_case.execute_test(df) # Execute the tests\n", "\n", " # Parse the test result values we need:\n", " intensity_results += [\n", " {\n", - " \"width\": test_results.estimator.control_value,\n", - " \"height\": test_results.estimator.treatment_value,\n", - " \"control\": test_results.estimator.control_value,\n", - " \"treatment\": test_results.estimator.treatment_value,\n", - " \"risk_ratio\": test_results.effect_estimate.value[0],\n", + " \"width\": estimator.control_value,\n", + " \"height\": estimator.treatment_value,\n", + " \"control\": estimator.control_value,\n", + " \"treatment\": estimator.treatment_value,\n", + " \"risk_ratio\": causal_test_case.result.effect_estimate.value[0],\n", " }]" ] }, @@ -576,7 +575,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 8, "id": "67bf5061-720f-4b3a-a371-3ff3092e81e1", "metadata": {}, "outputs": [], @@ -594,7 +593,6 @@ " for width_value in control_values:\n", " \n", " estimator = LinearRegressionEstimator(\n", - " df=df, # Pass in the dataframe\n", " base_test_case = base_test_case, # Base test case we created above\n", " treatment_value = width_value + 1.0, # Changing the width\n", " control_value=float(width_value), # Baseline width values\n", @@ -609,23 +607,23 @@ " estimate_type = \"ate_calculated\", # Calls the ate_calculated method in the linear regression estimator\n", " estimator=estimator) # Pass in the estimator we created above\n", " \n", - " test_results = causal_test_case.execute_test() # Execute the tests\n", + " causal_test_case.execute_test(df) # Execute the tests\n", "\n", " # Parse the test result values we need:\n", " width_results += [\n", " {\n", - " \"control\": test_results.estimator.control_value,\n", - " \"treatment\": test_results.estimator.treatment_value,\n", + " \"control\": estimator.control_value,\n", + " \"treatment\": estimator.treatment_value,\n", " \"intensity\": estimator.effect_modifiers[\"intensity\"],\n", - " \"ate\": test_results.effect_estimate.value[0],\n", - " \"ci_low\": test_results.effect_estimate.ci_low[0],\n", - " \"ci_high\": test_results.effect_estimate.ci_high[0],\n", + " \"ate\": causal_test_case.result.effect_estimate.value[0],\n", + " \"ci_low\": causal_test_case.result.effect_estimate.ci_low[0],\n", + " \"ci_high\": causal_test_case.result.effect_estimate.ci_high[0],\n", " }]\n" ] }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 9, "id": "6c54392c-4e6b-42b3-b39a-e0d1d0ab25b7", "metadata": {}, "outputs": [ @@ -767,7 +765,7 @@ "9 1.0 2.0 2 -7.378642 -16.381136 1.623851" ] }, - "execution_count": 8, + "execution_count": 9, "metadata": {}, "output_type": "execute_result" } diff --git a/docs/source/tutorials/vaccinating_elderly/causal_test_results.json b/docs/source/tutorials/vaccinating_elderly/causal_test_results.json index 7e83311c..06608b56 100644 --- a/docs/source/tutorials/vaccinating_elderly/causal_test_results.json +++ b/docs/source/tutorials/vaccinating_elderly/causal_test_results.json @@ -4,9 +4,7 @@ "estimate_type": "coefficient", "effect": "direct", "treatment_variable": "max_doses", - "expected_effect": { - "cum_vaccinations": "NoEffect" - }, + "expected_effect": "NoEffect", "alpha": 0.05, "formula": "cum_vaccinations ~ max_doses", "skip": false, @@ -32,9 +30,7 @@ "estimate_type": "coefficient", "effect": "direct", "treatment_variable": "max_doses", - "expected_effect": { - "cum_vaccinated": "NoEffect" - }, + "expected_effect": "NoEffect", "alpha": 0.05, "formula": "cum_vaccinated ~ max_doses", "skip": false, @@ -60,9 +56,7 @@ "estimate_type": "coefficient", "effect": "direct", "treatment_variable": "max_doses", - "expected_effect": { - "cum_infections": "NoEffect" - }, + "expected_effect": "NoEffect", "alpha": 0.05, "formula": "cum_infections ~ max_doses", "skip": false, @@ -88,9 +82,7 @@ "estimate_type": "coefficient", "effect": "direct", "treatment_variable": "vaccine", - "expected_effect": { - "cum_vaccinations": "SomeEffect" - }, + "expected_effect": "SomeEffect", "alpha": 0.05, "formula": "cum_vaccinations ~ vaccine", "skip": false, @@ -116,9 +108,7 @@ "estimate_type": "coefficient", "effect": "direct", "treatment_variable": "vaccine", - "expected_effect": { - "cum_vaccinated": "SomeEffect" - }, + "expected_effect": "SomeEffect", "alpha": 0.05, "formula": "cum_vaccinated ~ vaccine", "skip": false, @@ -144,9 +134,7 @@ "estimate_type": "coefficient", "effect": "direct", "treatment_variable": "vaccine", - "expected_effect": { - "cum_infections": "SomeEffect" - }, + "expected_effect": "SomeEffect", "alpha": 0.05, "formula": "cum_infections ~ vaccine", "skip": false, @@ -172,9 +160,7 @@ "estimate_type": "coefficient", "effect": "direct", "treatment_variable": "cum_vaccinations", - "expected_effect": { - "cum_vaccinated": "NoEffect" - }, + "expected_effect": "NoEffect", "alpha": 0.05, "formula": "cum_vaccinated ~ cum_vaccinations+vaccine", "skip": false, @@ -202,9 +188,7 @@ "estimate_type": "coefficient", "effect": "direct", "treatment_variable": "cum_vaccinations", - "expected_effect": { - "cum_infections": "NoEffect" - }, + "expected_effect": "NoEffect", "alpha": 0.05, "formula": "cum_infections ~ cum_vaccinations+vaccine", "skip": false, @@ -232,9 +216,7 @@ "estimate_type": "coefficient", "effect": "direct", "treatment_variable": "cum_vaccinated", - "expected_effect": { - "cum_infections": "NoEffect" - }, + "expected_effect": "NoEffect", "alpha": 0.05, "formula": "cum_infections ~ cum_vaccinated+vaccine", "skip": false, diff --git a/examples/covasim_/doubling_beta/example_beta.py b/examples/covasim_/doubling_beta/example_beta.py index 23842930..32cb2665 100644 --- a/examples/covasim_/doubling_beta/example_beta.py +++ b/examples/covasim_/doubling_beta/example_beta.py @@ -49,17 +49,16 @@ def doubling_beta_CATE_on_csv( base_test_case=base_test_case, expected_causal_effect=Positive, estimator=LinearRegressionEstimator( - base_test_case, - 0.032, - 0.016, - {"avg_age", "contacts"}, # We use custom adjustment set - df=past_execution_df, + base_test_case=base_test_case, + treatment_value=0.032, + control_value=0.016, + adjustment_set={"avg_age", "contacts"}, # We use custom adjustment set formula="cum_infections ~ beta + I(beta ** 2) + avg_age + contacts", ), ) # Add squared terms for beta, since it has a quadratic relationship with cumulative infections - causal_test_result = causal_test_case.execute_test() + causal_test_case.execute_test(past_execution_df) # Repeat for association estimate (no adjustment) causal_test_case = CausalTestCase( @@ -70,43 +69,42 @@ def doubling_beta_CATE_on_csv( treatment_value=0.032, control_value=0.016, adjustment_set=set(), - df=past_execution_df, formula="cum_infections ~ beta + I(beta ** 2)", ), ) - association_test_result = causal_test_case.execute_test() + causal_test_case.execute_test(past_execution_df) # Store results for plotting results_dict["association"] = { - "ate": association_test_result.effect_estimate.value, - "cis": [association_test_result.effect_estimate.ci_low, association_test_result.effect_estimate.ci_high], + "ate": causal_test_case.result.effect_estimate.value, + "cis": [causal_test_case.result.effect_estimate.ci_low, causal_test_case.result.effect_estimate.ci_high], "df": past_execution_df, } results_dict["causation"] = { - "ate": causal_test_result.effect_estimate.value, - "cis": [causal_test_result.effect_estimate.ci_low, causal_test_result.effect_estimate.ci_high], + "ate": causal_test_case.result.effect_estimate.value, + "cis": [causal_test_case.result.effect_estimate.ci_low, causal_test_case.result.effect_estimate.ci_high], "df": past_execution_df, } if verbose: - print(f"Association:\n{association_test_result}") - print(f"Causation:\n{causal_test_result}") + print(f"Association:\n{causal_test_case.result}") + print(f"Causation:\n{causal_test_case.result}") # Repeat causal inference after deleting all rows with treatment value to obtain counterfactual inferences if simulate_counterfactuals: counterfactual_past_execution_df = past_execution_df[past_execution_df["beta"] != 0.032] - counterfactual_causal_test_result = causal_test_case.execute_test() + causal_test_case.execute_test(past_execution_df) results_dict["counterfactual"] = { - "ate": counterfactual_causal_test_result.effect_estimate.value, + "ate": causal_test_case.result.effect_estimate.value, "cis": [ - counterfactual_causal_test_result.effect_estimate.ci_low, - counterfactual_causal_test_result.effect_estimate.ci_high, + causal_test_case.result.effect_estimate.ci_low, + causal_test_case.result.effect_estimate.ci_high, ], "df": counterfactual_past_execution_df, } if verbose: - print(f"Counterfactual:\n{counterfactual_causal_test_result}") + print(f"Counterfactual:\n{causal_test_case.result}") return results_dict diff --git a/examples/covasim_/vaccinating_elderly/example_vaccine.py b/examples/covasim_/vaccinating_elderly/example_vaccine.py index 5cb90aab..b877b075 100644 --- a/examples/covasim_/vaccinating_elderly/example_vaccine.py +++ b/examples/covasim_/vaccinating_elderly/example_vaccine.py @@ -49,34 +49,28 @@ def run_test_case(verbose: bool = False): causal_test_case = CausalTestCase( base_test_case=base_test_case, expected_causal_effect=expected_effect, + estimator=LinearRegressionEstimator( + base_test_case=base_test_case, + treatment_value=1, + control_value=0, + adjustment_set=causal_dag.identification(base_test_case), + ), ) - # 7. Obtain the minimal adjustment set for the causal test case from the causal DAG - minimal_adjustment_set = causal_dag.identification(base_test_case) - # 8. Build statistical model using the Linear Regression estimator - linear_regression_estimator = LinearRegressionEstimator( - base_test_case=base_test_case, - treatment_value=1, - control_value=0, - adjustment_set=minimal_adjustment_set, - df=obs_df, - ) - - # 9. Execute test and save results in dict - causal_test_result = causal_test_case.execute_test(linear_regression_estimator) + causal_test_case.execute_test(obs_df) if verbose: - logging.info("Causation:\n%s", causal_test_result) + logging.info("Causation:\n%s", causal_test_case.result) - results_dict[outcome_variable.name]["ate"] = causal_test_result.effect_estimate.value + results_dict[outcome_variable.name]["ate"] = causal_test_case.result.effect_estimate.value results_dict[outcome_variable.name]["cis"] = [ - causal_test_result.effect_estimate.ci_low, - causal_test_result.effect_estimate.ci_high, + causal_test_case.result.effect_estimate.ci_low, + causal_test_case.result.effect_estimate.ci_high, ] results_dict[outcome_variable.name]["test_passes"] = causal_test_case.expected_causal_effect.apply( - causal_test_result + causal_test_case.result ) return results_dict diff --git a/examples/lr91/example_max_conductances.py b/examples/lr91/example_max_conductances.py index 092fc004..ad9d4801 100644 --- a/examples/lr91/example_max_conductances.py +++ b/examples/lr91/example_max_conductances.py @@ -132,16 +132,15 @@ def effects_on_APD90(observational_data_path, treatment_var, control_val, treatm treatment_value=treatment_val, control_value=control_val, adjustment_set=causal_dag.identification(base_test_case), - df=pd.read_csv(observational_data_path), ), ) # 9. Run the causal test and print results - causal_test_result = causal_test_case.execute_test() - logger.info("%s", causal_test_result) - return causal_test_result.effect_estimate.value, ( - causal_test_result.effect_estimate.ci_low, - causal_test_result.effect_estimate.ci_high, + causal_test_case.execute_test(pd.read_csv(observational_data_path)) + logger.info("%s", causal_test_case.result) + return causal_test_case.result.effect_estimate.value, ( + causal_test_case.result.effect_estimate.ci_low, + causal_test_case.result.effect_estimate.ci_high, ) diff --git a/examples/poisson-line-process/example_pure_python.py b/examples/poisson-line-process/example_pure_python.py index 544ae8c3..0db5aadc 100644 --- a/examples/poisson-line-process/example_pure_python.py +++ b/examples/poisson-line-process/example_pure_python.py @@ -31,17 +31,16 @@ def add_modelling_assumptions(self): """ self.modelling_assumptions += "The data must contain runs with the exact configuration of interest." - def estimate_risk_ratio(self) -> EffectEstimate: + def estimate_risk_ratio(self, df: pd.DataFrame) -> EffectEstimate: """Estimate the outcomes under control and treatment. + :param df: The data to use. :return: The empirical average treatment effect. """ treatment_variable = self.base_test_case.treatment_variable.name outcome_variable = self.base_test_case.outcome_variable.name - control_results = self.df.where(self.df[treatment_variable] == self.control_value)[outcome_variable].dropna() - treatment_results = self.df.where(self.df[treatment_variable] == self.treatment_value)[ - outcome_variable - ].dropna() + control_results = df.where(df[treatment_variable] == self.control_value)[outcome_variable].dropna() + treatment_results = df.where(df[treatment_variable] == self.treatment_value)[outcome_variable].dropna() def risk_ratio(sample1, sample2): return sample1.mean() / sample2.mean() @@ -100,12 +99,11 @@ def test_poisson_intensity_num_shapes(save=False): treatment_value=treatment_value, control_value=control_value, adjustment_set=causal_dag.identification(base_test_case), - df=pd.read_csv(f"{ROOT}/data/smt_100/data_smt_wh{wh}_100.csv", index_col=0).astype(float), effect_modifiers=None, alpha=0.05, - query="", ), ), + f"{ROOT}/data/smt_100/data_smt_wh{wh}_100.csv", CausalTestCase( base_test_case=base_test_case, expected_causal_effect=ExactValue(4, atol=0.5), @@ -115,11 +113,9 @@ def test_poisson_intensity_num_shapes(save=False): treatment_value=treatment_value, control_value=control_value, adjustment_set=causal_dag.identification(base_test_case), - df=observational_df, effect_modifiers=None, formula="num_shapes_unit ~ I(intensity ** 2) + intensity - 1", alpha=0.05, - query="", ), ), ) @@ -127,18 +123,20 @@ def test_poisson_intensity_num_shapes(save=False): for wh in range(1, 11) ] - test_results = [(smt.execute_test(), observational.execute_test()) for smt, observational in causal_test_cases] + for smt_causal_test, datapath, obs_causal_test in causal_test_cases: + smt_causal_test.execute_test(df=pd.read_csv(datapath, index_col=0).astype(float)) + obs_causal_test.execute_test(observational_df) intensity_num_shapes_results += [ { - "width": obs_causal_test_result.estimator.control_value, - "height": obs_causal_test_result.estimator.treatment_value, - "control": obs_causal_test_result.estimator.control_value, - "treatment": obs_causal_test_result.estimator.treatment_value, - "smt_risk_ratio": smt_causal_test_result.effect_estimate.value, - "obs_risk_ratio": obs_causal_test_result.effect_estimate.value[0], + "width": obs_causal_test.estimator.control_value, + "height": obs_causal_test.estimator.treatment_value, + "control": obs_causal_test.estimator.control_value, + "treatment": obs_causal_test.estimator.treatment_value, + "smt_risk_ratio": smt_causal_test.result.effect_estimate.value, + "obs_risk_ratio": obs_causal_test.result.effect_estimate.value[0], } - for smt_causal_test_result, obs_causal_test_result in test_results + for smt_causal_test, _, obs_causal_test in causal_test_cases ] intensity_num_shapes_results = pd.DataFrame(intensity_num_shapes_results) if save: @@ -159,7 +157,6 @@ def test_poisson_width_num_shapes(save=False): treatment_value=w + 1.0, control_value=float(w), adjustment_set=causal_dag.identification(base_test_case), - df=df, effect_modifiers={"intensity": i}, formula="num_shapes_unit ~ width + I(intensity ** 2)+I(width ** -1)+intensity-1", alpha=0.05, @@ -168,17 +165,18 @@ def test_poisson_width_num_shapes(save=False): for i in range(1, 17) for w in range(1, 10) ] - test_results = [test.execute_test() for test in causal_test_cases] + for test in causal_test_cases: + test.execute_test(df) width_num_shapes_results = [ { - "control": causal_test_result.estimator.control_value, - "treatment": causal_test_result.estimator.treatment_value, - "intensity": causal_test_result.estimator.effect_modifiers["intensity"], - "ate": causal_test_result.effect_estimate.value[0], - "ci_low": causal_test_result.effect_estimate.ci_low, - "ci_high": causal_test_result.effect_estimate.ci_high, + "control": causal_test.estimator.control_value, + "treatment": causal_test.estimator.treatment_value, + "intensity": causal_test.estimator.effect_modifiers["intensity"], + "ate": causal_test.result.effect_estimate.value[0], + "ci_low": causal_test.result.effect_estimate.ci_low, + "ci_high": causal_test.result.effect_estimate.ci_high, } - for causal_test_result in test_results + for causal_test in causal_test_cases ] width_num_shapes_results = pd.DataFrame(width_num_shapes_results) if save: diff --git a/tests/discovery_tests/test_abstract_discovery.py b/tests/discovery_tests/test_abstract_discovery.py index a8104288..1f9d9e34 100644 --- a/tests/discovery_tests/test_abstract_discovery.py +++ b/tests/discovery_tests/test_abstract_discovery.py @@ -6,10 +6,12 @@ import pandas as pd from tempfile import TemporaryDirectory import os +from numpy import nan -from causal_testing.discovery.abstract_discovery import TestResult, Discovery, simple_cycle, effect_direction +from causal_testing.discovery.abstract_discovery import TestResult, Discovery, simple_cycle from causal_testing.specification.causal_dag import CausalDAG from causal_testing.testing.causal_test_result import CausalTestResult +from causal_testing.testing.causal_test_case import CausalTestCase from causal_testing.estimation.effect_estimate import EffectEstimate from causal_testing.estimation.linear_regression_estimator import LinearRegressionEstimator from causal_testing.testing.base_test_case import BaseTestCase @@ -30,10 +32,13 @@ def discover(self): class TestAbstractHillClimber(unittest.TestCase): def setUp(self) -> None: - base_test_case = BaseTestCase(Input("A", float), Output("B", float)) + self.base_test_case = BaseTestCase(Input("A", float), Output("B", float)) + self.df = pd.DataFrame({"A": [1, 2], "B": [4, 5]}) + self.abstract_discovery = AbstractDiscovery( + df=self.df, + ) self.estimator = LinearRegressionEstimator( - df=pd.DataFrame({"A": [1, 2], "B": [4, 5]}), - base_test_case=base_test_case, + base_test_case=self.base_test_case, treatment_value=1, control_value=0, adjustment_set={}, @@ -50,25 +55,29 @@ def test_simple_cycle_no_cycles(self): self.assertEqual(simple_cycle(dag), []) def test_effect_direction_positive(self): - ctr = CausalTestResult( - estimator=self.estimator, - effect_estimate=EffectEstimate(type="ate", value=pd.Series(5.05)), + causal_test_case = CausalTestCase(base_test_case=self.base_test_case, expected_causal_effect=None) + causal_test_case.result = CausalTestResult( + effect_estimate=EffectEstimate( + type="ate", value=pd.Series(5.05), ci_low=pd.Series(5), ci_high=pd.Series(6) + ), ) - self.assertEqual(effect_direction(ctr), "positive") + self.assertEqual(self.abstract_discovery.effect_direction(causal_test_case), "positive") def test_effect_direction_negative(self): - ctr = CausalTestResult( - estimator=self.estimator, - effect_estimate=EffectEstimate(type="ate", value=pd.Series(-5.05)), + causal_test_case = CausalTestCase(base_test_case=self.base_test_case, expected_causal_effect=None) + causal_test_case.result = CausalTestResult( + effect_estimate=EffectEstimate( + type="ate", value=pd.Series(-5.05), ci_low=pd.Series(-6), ci_high=pd.Series(-5) + ), ) - self.assertEqual(effect_direction(ctr), "negative") + self.assertEqual(self.abstract_discovery.effect_direction(causal_test_case), "negative") def test_effect_direction_none(self): - ctr = CausalTestResult( - estimator=self.estimator, - effect_estimate=EffectEstimate(type="ate", value=pd.Series(0)), + causal_test_case = CausalTestCase(base_test_case=self.base_test_case, expected_causal_effect=None) + causal_test_case.result = CausalTestResult( + effect_estimate=EffectEstimate(type="ate", value=pd.Series(0), ci_low=pd.Series(-1), ci_high=pd.Series(1)), ) - self.assertEqual(effect_direction(ctr), None) + self.assertEqual(self.abstract_discovery.effect_direction(causal_test_case), None) def test_include_edge_wildcard(self): abstract_discovery = AbstractDiscovery( @@ -137,13 +146,13 @@ def test_write_dot(self): dag.add_edges_from([("A", "B"), ("C", "D"), ("E", "F")]) dag.test_results = pd.DataFrame( [ # Edges - {"treatment": "A", "outcome": "B", "result": TestResult.PASS}, - {"treatment": "C", "outcome": "D", "result": TestResult.FAIL}, - {"treatment": "E", "outcome": "F", "result": TestResult.INESTIMABLE}, + {"treatment": "A", "outcome": "B", "effect": "positive", "result": TestResult.PASS}, + {"treatment": "C", "outcome": "D", "effect": "positive", "result": TestResult.FAIL}, + {"treatment": "E", "outcome": "F", "effect": "None", "result": TestResult.INESTIMABLE}, # Independences - {"treatment": "A", "outcome": "C", "result": TestResult.PASS}, - {"treatment": "A", "outcome": "D", "result": TestResult.FAIL}, - {"treatment": "A", "outcome": "E", "result": TestResult.INESTIMABLE}, + {"treatment": "A", "outcome": "C", "effect": None, "result": TestResult.PASS}, + {"treatment": "A", "outcome": "D", "effect": "negative", "result": TestResult.FAIL}, + {"treatment": "A", "outcome": "E", "effect": None, "result": TestResult.INESTIMABLE}, ] ) abstract_discovery = AbstractDiscovery(pd.DataFrame()) @@ -157,7 +166,7 @@ def test_write_dot_invalid_edge_outcome(self): dag.add_edges_from([("A", "B"), ("C", "D"), ("E", "F")]) dag.test_results = pd.DataFrame( [ # Edges - {"treatment": "A", "outcome": "B", "result": None}, + {"treatment": "A", "outcome": "B", "effect": None, "result": None}, ] ) abstract_discovery = AbstractDiscovery(pd.DataFrame()) @@ -169,7 +178,7 @@ def test_write_dot_invalid_independence_outcome(self): dag.add_edges_from([("A", "B"), ("C", "D"), ("E", "F")]) dag.test_results = pd.DataFrame( [ # Edges - {"treatment": "A", "outcome": "C", "result": None}, + {"treatment": "A", "outcome": "C", "effect": None, "result": None}, ] ) abstract_discovery = AbstractDiscovery(pd.DataFrame()) @@ -206,73 +215,64 @@ def test_evaluate_tests_inestimable(self): "expected_effect": "NoEffect", "treatment": "length_in", "outcome": "large_gauge", - "effect": "negative", }, { "result": TestResult.PASS, "expected_effect": "NoEffect", "treatment": "large_gauge", "outcome": "length_in", - "effect": "negative", }, { "result": TestResult.PASS, "expected_effect": "NoEffect", "treatment": "length_in", "outcome": "color", - "effect": None, }, { "result": TestResult.PASS, "expected_effect": "NoEffect", "treatment": "color", "outcome": "length_in", - "effect": None, }, { "result": TestResult.FAIL, "expected_effect": "SomeEffect", "treatment": "length_in", "outcome": "completed", - "effect": "positive", }, { "result": TestResult.PASS, "expected_effect": "NoEffect", "treatment": "large_gauge", "outcome": "color", - "effect": None, }, { "result": TestResult.PASS, "expected_effect": "NoEffect", "treatment": "color", "outcome": "large_gauge", - "effect": None, }, { "result": TestResult.FAIL, "expected_effect": "SomeEffect", "treatment": "large_gauge", "outcome": "completed", - "effect": "negative", }, { "result": TestResult.INESTIMABLE, "expected_effect": "NoEffect", "treatment": "color", "outcome": "completed", - "effect": None, }, { "result": TestResult.INESTIMABLE, "expected_effect": "NoEffect", "treatment": "completed", "outcome": "color", - "effect": None, }, ] ) + expected_results["effect"] = nan pd.testing.assert_frame_equal(test_results, expected_results) def test_evaluate_tests(self): @@ -291,71 +291,62 @@ def test_evaluate_tests(self): "expected_effect": "NoEffect", "treatment": "length_in", "outcome": "large_gauge", - "effect": "positive", }, { "result": TestResult.PASS, "expected_effect": "NoEffect", "treatment": "large_gauge", "outcome": "length_in", - "effect": "positive", }, { "result": TestResult.PASS, "expected_effect": "NoEffect", "treatment": "length_in", "outcome": "color", - "effect": None, }, { "result": TestResult.PASS, "expected_effect": "NoEffect", "treatment": "color", "outcome": "length_in", - "effect": None, }, { "result": TestResult.FAIL, "expected_effect": "SomeEffect", "treatment": "length_in", "outcome": "completed", - "effect": "negative", }, { "result": TestResult.PASS, "expected_effect": "NoEffect", "treatment": "large_gauge", "outcome": "color", - "effect": None, }, { "result": TestResult.PASS, "expected_effect": "NoEffect", "treatment": "color", "outcome": "large_gauge", - "effect": None, }, { "result": TestResult.FAIL, "expected_effect": "SomeEffect", "treatment": "large_gauge", "outcome": "completed", - "effect": "positive", }, { "result": TestResult.PASS, "expected_effect": "NoEffect", "treatment": "color", "outcome": "completed", - "effect": None, }, { "result": TestResult.PASS, "expected_effect": "NoEffect", "treatment": "completed", "outcome": "color", - "effect": None, }, ] ) + expected_results["effect"] = None 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 90fe4617..9338e42a 100644 --- a/tests/discovery_tests/test_hill_climber_discovery.py +++ b/tests/discovery_tests/test_hill_climber_discovery.py @@ -6,7 +6,7 @@ import pandas as pd from causal_testing.discovery.hill_climber_discovery import HillClimberDiscovery from causal_testing.specification.causal_dag import CausalDAG -from causal_testing.discovery.abstract_discovery import TestResult, Discovery, simple_cycle, effect_direction +from causal_testing.discovery.abstract_discovery import TestResult, Discovery, simple_cycle class TestHillClimber(unittest.TestCase): diff --git a/tests/estimation_tests/test_cubic_spline_estimator.py b/tests/estimation_tests/test_cubic_spline_estimator.py index 0298a893..396b190f 100644 --- a/tests/estimation_tests/test_cubic_spline_estimator.py +++ b/tests/estimation_tests/test_cubic_spline_estimator.py @@ -22,12 +22,14 @@ def test_program_11_3_cublic_spline(self): base_test_case = BaseTestCase(Input("treatments", float), Output("outcomes", float)) - cublic_spline_estimator = CubicSplineRegressionEstimator(base_test_case, 1, 0, set(), 3, df) + cublic_spline_estimator = CubicSplineRegressionEstimator( + base_test_case=base_test_case, treatment_value=1, control_value=0, adjustment_set=set(), basis=3 + ) - ate_1 = cublic_spline_estimator.estimate_ate_calculated().value + ate_1 = cublic_spline_estimator.estimate_ate_calculated(df).value cublic_spline_estimator.treatment_value = 2 - ate_2 = cublic_spline_estimator.estimate_ate_calculated().value + ate_2 = cublic_spline_estimator.estimate_ate_calculated(df).value # Doubling the treatemebnt value should roughly but not exactly double the ATE self.assertNotEqual(ate_1[0] * 2, ate_2[0]) diff --git a/tests/estimation_tests/test_instrumental_variable_estimator.py b/tests/estimation_tests/test_instrumental_variable_estimator.py index 0c831775..5c86b0ad 100644 --- a/tests/estimation_tests/test_instrumental_variable_estimator.py +++ b/tests/estimation_tests/test_instrumental_variable_estimator.py @@ -24,14 +24,13 @@ def test_estimate_coefficient(self): Test we get the correct coefficient. """ iv_estimator = InstrumentalVariableEstimator( - df=self.df, base_test_case=BaseTestCase(Input("X", float), Output("Y", float)), treatment_value=None, control_value=None, adjustment_set=set(), instrument="Z", ) - effect_estimate = iv_estimator.estimate_coefficient() + effect_estimate = iv_estimator.estimate_coefficient(self.df) self.assertEqual(effect_estimate.value[0], 2) self.assertEqual(effect_estimate.ci_low[0], 2) self.assertEqual(effect_estimate.ci_high[0], 2) diff --git a/tests/estimation_tests/test_ipcw_estimator.py b/tests/estimation_tests/test_ipcw_estimator.py index 9bdddb14..bac1d2c3 100644 --- a/tests/estimation_tests/test_ipcw_estimator.py +++ b/tests/estimation_tests/test_ipcw_estimator.py @@ -22,7 +22,6 @@ def setUp(self) -> None: def test_estimate_hazard_ratio(self): estimation_model = IPCWEstimator( - self.df, self.timesteps_per_intervention, self.control_strategy, self.treatment_strategy, @@ -32,26 +31,25 @@ def test_estimate_hazard_ratio(self): fit_bltd_switch_formula=self.fit_bl_switch_formula, eligibility=None, ) - estimate = estimation_model.estimate_hazard_ratio() + estimate = estimation_model.estimate_hazard_ratio(self.df) self.assertEqual(round(estimate.value["trtrand"], 3), 1.351) def test_invalid_treatment_strategies(self): + estimation_model = IPCWEstimator( + self.timesteps_per_intervention, + self.control_strategy, + self.treatment_strategy, + self.outcome, + self.status_column, + fit_bl_switch_formula=self.fit_bl_switch_formula, + fit_bltd_switch_formula=self.fit_bl_switch_formula, + eligibility=None, + ) with self.assertRaises(ValueError): - IPCWEstimator( - self.df.assign(t=(["1", "0"] * len(self.df))[: len(self.df)]), - self.timesteps_per_intervention, - self.control_strategy, - self.treatment_strategy, - self.outcome, - self.status_column, - fit_bl_switch_formula=self.fit_bl_switch_formula, - fit_bltd_switch_formula=self.fit_bl_switch_formula, - eligibility=None, - ) + estimation_model.preprocess_data(self.df.assign(t=(["1", "0"] * len(self.df))[: len(self.df)])) def test_invalid_fault_t_do(self): estimation_model = IPCWEstimator( - self.df.assign(outcome=1), self.timesteps_per_intervention, self.control_strategy, self.treatment_strategy, @@ -61,34 +59,47 @@ def test_invalid_fault_t_do(self): fit_bltd_switch_formula=self.fit_bl_switch_formula, eligibility=None, ) - estimation_model.df["fault_t_do"] = 0 with self.assertRaises(ValueError): - estimation_model.estimate_hazard_ratio() + estimation_model.estimate_hazard_ratio(self.df.loc[~self.df["ok"]]) def test_no_individual_began_control_strategy(self): + estimation_model = IPCWEstimator( + self.timesteps_per_intervention, + self.control_strategy, + self.treatment_strategy, + self.outcome, + self.status_column, + fit_bl_switch_formula=self.fit_bl_switch_formula, + fit_bltd_switch_formula=self.fit_bl_switch_formula, + eligibility=None, + ) with self.assertRaises(ValueError): - IPCWEstimator( - self.df.assign(t=1), - self.timesteps_per_intervention, - self.control_strategy, - self.treatment_strategy, - self.outcome, - self.status_column, - fit_bl_switch_formula=self.fit_bl_switch_formula, - fit_bltd_switch_formula=self.fit_bl_switch_formula, - eligibility=None, - ) + estimation_model.preprocess_data(self.df.assign(t=1)) def test_no_individual_began_treatment_strategy(self): + estimation_model = IPCWEstimator( + self.timesteps_per_intervention, + self.control_strategy, + self.treatment_strategy, + self.outcome, + self.status_column, + fit_bl_switch_formula=self.fit_bl_switch_formula, + fit_bltd_switch_formula=self.fit_bl_switch_formula, + eligibility=None, + ) + with self.assertRaises(ValueError): + estimation_model.preprocess_data(self.df.assign(t=0)) + + def test_preprocess_data_no_faults(self): + estimation_model = IPCWEstimator( + self.timesteps_per_intervention, + self.control_strategy, + self.treatment_strategy, + self.outcome, + self.status_column, + fit_bl_switch_formula=self.fit_bl_switch_formula, + fit_bltd_switch_formula=self.fit_bl_switch_formula, + eligibility=None, + ) with self.assertRaises(ValueError): - IPCWEstimator( - self.df.assign(t=0), - self.timesteps_per_intervention, - self.control_strategy, - self.treatment_strategy, - self.outcome, - self.status_column, - fit_bl_switch_formula=self.fit_bl_switch_formula, - fit_bltd_switch_formula=self.fit_bl_switch_formula, - eligibility=None, - ) + estimation_model.preprocess_data(self.df.assign(ok=True)) diff --git a/tests/estimation_tests/test_linear_regression_estimator.py b/tests/estimation_tests/test_linear_regression_estimator.py index b1457785..128fcb55 100644 --- a/tests/estimation_tests/test_linear_regression_estimator.py +++ b/tests/estimation_tests/test_linear_regression_estimator.py @@ -65,27 +65,20 @@ def setUpClass(cls) -> None: cls.base_test_case = BaseTestCase(Input("treatments", float), Output("outcomes", float)) cls.program_15_base_test_case = BaseTestCase(Input("qsmk", float), Output("wt82_71", float)) - def test_query(self): - df = self.nhefs_df - linear_regression_estimator = LinearRegressionEstimator( - self.program_15_base_test_case, None, None, set(), df, query="sex==1" - ) - self.assertTrue(linear_regression_estimator.df.sex.all()) - def test_linear_regression_categorical_ate(self): df = self.scarf_df.copy() base_test_case = BaseTestCase(Input("color", float), Output("completed", float)) - linear_regression_estimator = LinearRegressionEstimator(base_test_case, None, None, set(), df) - effect_estimate = linear_regression_estimator.estimate_coefficient() + linear_regression_estimator = LinearRegressionEstimator(base_test_case=base_test_case, adjustment_set=set()) + effect_estimate = linear_regression_estimator.estimate_coefficient(df) self.assertTrue( - all([ci_low < 0 < ci_high for ci_low, ci_high in zip(effect_estimate.ci_low, effect_estimate.ci_high)]) + all(ci_low < 0 < ci_high for ci_low, ci_high in zip(effect_estimate.ci_low, effect_estimate.ci_high)) ) def test_program_11_2(self): """Test whether our linear regression implementation produces the same results as program 11.2 (p. 141).""" df = self.chapter_11_df - linear_regression_estimator = LinearRegressionEstimator(self.base_test_case, None, None, set(), df) - effect_estimate = linear_regression_estimator.estimate_coefficient() + linear_regression_estimator = LinearRegressionEstimator(self.base_test_case, adjustment_set=set()) + effect_estimate = linear_regression_estimator.estimate_coefficient(df) # Increasing treatments from 90 to 100 should be the same as 10 times the unit ATE self.assertTrue( @@ -99,9 +92,11 @@ def test_program_11_3(self): """Test whether our linear regression implementation produces the same results as program 11.3 (p. 144).""" df = self.chapter_11_df.copy() linear_regression_estimator = LinearRegressionEstimator( - self.base_test_case, None, None, set(), df, formula="outcomes ~ treatments + I(treatments ** 2)" + base_test_case=self.base_test_case, + adjustment_set=set(), + formula="outcomes ~ treatments + I(treatments ** 2)", ) - effect_estimate = linear_regression_estimator.estimate_coefficient() + effect_estimate = linear_regression_estimator.estimate_coefficient(df) # Increasing treatments from 90 to 100 should be the same as 10 times the unit ATE self.assertTrue( all( @@ -130,11 +125,10 @@ def test_program_15_1A(self): "smokeyrs", } linear_regression_estimator = LinearRegressionEstimator( - self.program_15_base_test_case, - 1, - 0, - covariates, - df, + base_test_case=self.program_15_base_test_case, + treatment_value=1, + control_value=0, + adjustment_set=covariates, formula=f"""wt82_71 ~ qsmk + {'+'.join(sorted(list(covariates)))} + I(age ** 2) + @@ -144,7 +138,7 @@ def test_program_15_1A(self): (qsmk * smokeintensity)""", ) - effect_estimate = linear_regression_estimator.estimate_coefficient() + effect_estimate = linear_regression_estimator.estimate_coefficient(df) self.assertEqual(round(effect_estimate.value["qsmk"], 1), 2.6) def test_program_15_no_interaction(self): @@ -168,16 +162,15 @@ def test_program_15_no_interaction(self): "smokeyrs", } linear_regression_estimator = LinearRegressionEstimator( - self.program_15_base_test_case, - 1, - 0, - covariates, - df, + base_test_case=self.program_15_base_test_case, + treatment_value=1, + control_value=0, + adjustment_set=covariates, formula="wt82_71 ~ qsmk + age + I(age ** 2) + wt71 + I(wt71 ** 2) + smokeintensity + I(smokeintensity ** 2) + smokeyrs + I(smokeyrs ** 2)", ) # terms_to_square = ["age", "wt71", "smokeintensity", "smokeyrs"] # for term_to_square in terms_to_square: - effect_estimate = linear_regression_estimator.estimate_coefficient() + effect_estimate = linear_regression_estimator.estimate_coefficient(df) self.assertEqual(round(effect_estimate.value.iloc[0], 1), 3.5) self.assertEqual(round(effect_estimate.ci_low.iloc[0], 1), 2.6) @@ -204,23 +197,21 @@ def test_program_15_no_interaction_ate(self): "smokeyrs", } linear_regression_estimator = LinearRegressionEstimator( - self.program_15_base_test_case, - 1, - 0, - covariates, - df, + base_test_case=self.program_15_base_test_case, + treatment_value=1, + control_value=0, + adjustment_set=covariates, formula="wt82_71 ~ qsmk + age + I(age ** 2) + wt71 + I(wt71 ** 2) + smokeintensity + I(smokeintensity ** 2) + smokeyrs + I(smokeyrs ** 2)", ) # terms_to_square = ["age", "wt71", "smokeintensity", "smokeyrs"] # for term_to_square in terms_to_square: - effect_estimate = linear_regression_estimator.estimate_ate() + effect_estimate = linear_regression_estimator.estimate_ate(df) self.assertEqual(round(effect_estimate.value[0], 1), 3.5) self.assertEqual([round(effect_estimate.ci_low[0], 1), round(effect_estimate.ci_high[0], 1)], [2.6, 4.3]) def test_program_15_no_interaction_ate_calculated(self): """Test whether our linear regression implementation produces the same results as program 15.1 (p. 163, 184) without product parameter.""" - df = self.nhefs_df covariates = { "sex", "race", @@ -238,41 +229,52 @@ def test_program_15_no_interaction_ate_calculated(self): "smokeyrs", } linear_regression_estimator = LinearRegressionEstimator( - self.program_15_base_test_case, - 1, - 0, - covariates, - df, + base_test_case=self.program_15_base_test_case, + treatment_value=1, + control_value=0, + adjustment_set=covariates, formula="wt82_71 ~ qsmk + age + I(age ** 2) + wt71 + I(wt71 ** 2) + smokeintensity + I(smokeintensity ** 2) + smokeyrs + I(smokeyrs ** 2)", + adjustment_config={k: self.nhefs_df.mean()[k] for k in covariates}, ) # terms_to_square = ["age", "wt71", "smokeintensity", "smokeyrs"] # for term_to_square in terms_to_square: - effect_estimate = linear_regression_estimator.estimate_ate_calculated( - adjustment_config={k: self.nhefs_df.mean()[k] for k in covariates} - ) + effect_estimate = linear_regression_estimator.estimate_ate_calculated(df=self.nhefs_df) self.assertEqual(round(effect_estimate.value[0], 1), 3.5) self.assertEqual([round(effect_estimate.ci_low[0], 1), round(effect_estimate.ci_high[0], 1)], [1.9, 5]) def test_program_11_2_with_robustness_validation(self): """Test whether our linear regression estimator, as used in test_program_11_2 can correctly estimate robustness.""" df = self.chapter_11_df.copy() - linear_regression_estimator = LinearRegressionEstimator(self.base_test_case, 100, 90, set(), df) + linear_regression_estimator = LinearRegressionEstimator( + base_test_case=self.base_test_case, + treatment_value=100, + control_value=90, + adjustment_set=set(), + ) cv = CausalValidator() self.assertEqual( - round(cv.estimate_robustness(linear_regression_estimator.fit_model())["treatments"], 4), 0.7353 + round(cv.estimate_robustness(linear_regression_estimator.fit_model(df))["treatments"], 4), 0.7353 ) def test_gp(self): df = pd.DataFrame() - df["X"] = np.arange(10) + df["X"] = np.arange(10).astype(float) df["Y"] = 1 / (df["X"] + 1) + print(df) base_test_case = BaseTestCase(Input("X", float), Output("Y", float)) - linear_regression_estimator = LinearRegressionEstimator(base_test_case, 0, 1, set(), df.astype(float)) - linear_regression_estimator.gp_formula(seeds=["reciprocal(add(X, 1))"], extra_operators=[(reciprocal, 1)]) + linear_regression_estimator = LinearRegressionEstimator( + base_test_case=base_test_case, + treatment_value=0, + control_value=1, + adjustment_set=set(), + ) + linear_regression_estimator.gp_formula( + df=df, seeds=["reciprocal(add(X, 1))"], extra_operators=[(reciprocal, 1)] + ) self.assertEqual(linear_regression_estimator.formula, "Y ~ I(1/(X + 1)) - 1") - effect_estimate = linear_regression_estimator.estimate_ate_calculated() + effect_estimate = linear_regression_estimator.estimate_ate_calculated(df) self.assertEqual(round(effect_estimate.value[0], 2), 0.50) self.assertEqual(round(effect_estimate.ci_low[0], 2), 0.50) self.assertEqual(round(effect_estimate.ci_high[0], 2), 0.50) @@ -282,13 +284,20 @@ def test_gp_power(self): base_test_case = BaseTestCase(Input("X", float), Output("Y", float)) df["X"] = np.arange(10) df["Y"] = 2 * (df["X"] ** 2) - linear_regression_estimator = LinearRegressionEstimator(base_test_case, 0, 1, set(), df.astype(float)) - linear_regression_estimator.gp_formula(seed=1, max_order=2, seeds=["mul(2, power_2(X))"], extra_operators=[]) + linear_regression_estimator = LinearRegressionEstimator( + base_test_case=base_test_case, + treatment_value=0, + control_value=1, + adjustment_set=set(), + ) + linear_regression_estimator.gp_formula( + df=df, seed=1, max_order=2, seeds=["mul(2, power_2(X))"], extra_operators=[] + ) self.assertEqual( linear_regression_estimator.formula, "Y ~ I(2*X**2) - 1", ) - effect_estimate = linear_regression_estimator.estimate_ate_calculated() + effect_estimate = linear_regression_estimator.estimate_ate_calculated(df) self.assertEqual(round(effect_estimate.value[0], 2), -2.00) self.assertEqual(round(effect_estimate.ci_low[0], 2), -2.00) self.assertEqual(round(effect_estimate.ci_high[0], 2), -2.00) @@ -309,9 +318,14 @@ def test_X1_effect(self): """When we fix the value of X2 to 0, the effect of X1 on Y should become ~2 (because X2 terms are cancelled).""" base_test_case = BaseTestCase(Input("X1", float), Output("Y", float)) lr_model = LinearRegressionEstimator( - base_test_case, 1, 0, {"X2"}, effect_modifiers={"x2": 0}, formula="Y ~ X1 + X2 + (X1 * X2)", df=self.df + base_test_case=base_test_case, + treatment_value=1, + control_value=0, + adjustment_set={"X2"}, + effect_modifiers={"x2": 0}, + formula="Y ~ X1 + X2 + (X1 * X2)", ) - effect_estimate = lr_model.estimate_ate() + effect_estimate = lr_model.estimate_ate(self.df) self.assertAlmostEqual(effect_estimate.value[0], 2.0) def test_categorical_confidence_intervals(self): @@ -321,9 +335,8 @@ def test_categorical_confidence_intervals(self): control_value=None, treatment_value=None, adjustment_set={}, - df=self.scarf_df, ) - effect_estimate = lr_model.estimate_coefficient() + effect_estimate = lr_model.estimate_coefficient(self.scarf_df) # The precise values don't really matter. This test is primarily intended to make sure the return type is correct. self.assertTrue( diff --git a/tests/estimation_tests/test_logistic_regression_estimator.py b/tests/estimation_tests/test_logistic_regression_estimator.py index 15a0ed4b..1bd77ddf 100644 --- a/tests/estimation_tests/test_logistic_regression_estimator.py +++ b/tests/estimation_tests/test_logistic_regression_estimator.py @@ -15,18 +15,8 @@ def setUpClass(cls) -> None: cls.scarf_df = pd.read_csv("tests/resources/data/scarf_data.csv") def test_odds_ratio(self): - df = self.scarf_df.copy() - logistic_regression_estimator = LogisticRegressionEstimator( - BaseTestCase(Input("length_in", float), Output("completed", bool)), 65, 55, set(), df - ) - effect_estimate = logistic_regression_estimator.estimate_unit_odds_ratio() - self.assertEqual(round(effect_estimate.value.iloc[0], 4), 0.8948) - - def test_odds_ratio_data(self): - df = self.scarf_df.copy() logistic_regression_estimator = LogisticRegressionEstimator( BaseTestCase(Input("length_in", float), Output("completed", bool)), 65, 55, set() ) - logistic_regression_estimator.df = df - effect_estimate = logistic_regression_estimator.estimate_unit_odds_ratio() + effect_estimate = logistic_regression_estimator.estimate_unit_odds_ratio(self.scarf_df) self.assertEqual(round(effect_estimate.value.iloc[0], 4), 0.8948) diff --git a/tests/estimation_tests/test_multinomial_regression_estimator.py b/tests/estimation_tests/test_multinomial_regression_estimator.py index 48ebe711..fd69ac8a 100644 --- a/tests/estimation_tests/test_multinomial_regression_estimator.py +++ b/tests/estimation_tests/test_multinomial_regression_estimator.py @@ -16,27 +16,22 @@ def setUpClass(cls) -> None: cls.scarf_df = pd.read_csv("tests/resources/data/scarf_data.csv") def test_odds_ratio(self): - df = self.scarf_df.copy() multinomial_regression_estimator = MultinomialRegressionEstimator( - BaseTestCase(Input("length_in", float), Output("completed", bool)), 65, 55, set(), df + BaseTestCase(Input("length_in", float), Output("completed", bool)), 65, 55, set() ) - effect_estimate = multinomial_regression_estimator.estimate_unit_odds_ratio() + effect_estimate = multinomial_regression_estimator.estimate_unit_odds_ratio(self.scarf_df) self.assertEqual(round(effect_estimate.value.iloc[0], 4), 0.8948) def test_odds_ratio_category(self): - df = self.scarf_df.copy() multinomial_regression_estimator = MultinomialRegressionEstimator( - BaseTestCase(Input("length_in", float), Output("color", bool)), 65, 55, set(), df + BaseTestCase(Input("length_in", float), Output("color", bool)), 65, 55, set() ) - effect_estimate = multinomial_regression_estimator.estimate_unit_odds_ratio() - print(effect_estimate.value) + effect_estimate = multinomial_regression_estimator.estimate_unit_odds_ratio(self.scarf_df) self.assertTrue(effect_estimate.value.round(4).equals, pd.Series({"grey": 1.0072, "orange": 0.9668})) def test_odds_ratio_data(self): - df = self.scarf_df.copy() multinomial_regression_estimator = MultinomialRegressionEstimator( BaseTestCase(Input("length_in", float), Output("completed", bool)), 65, 55, set() ) - multinomial_regression_estimator.df = df - effect_estimate = multinomial_regression_estimator.estimate_unit_odds_ratio() + effect_estimate = multinomial_regression_estimator.estimate_unit_odds_ratio(self.scarf_df) self.assertEqual(round(effect_estimate.value.iloc[0], 4), 0.8948) diff --git a/tests/main_tests/test_main.py b/tests/main_tests/test_main.py index 86dd286a..27b94fb8 100644 --- a/tests/main_tests/test_main.py +++ b/tests/main_tests/test_main.py @@ -7,128 +7,60 @@ import json import pandas as pd -from causal_testing.main import CausalTestingPaths, CausalTestingFramework +from causal_testing.causal_testing_framework import CausalTestingFramework from causal_testing.__main__ import main -class TestCausalTestingPaths(unittest.TestCase): - - def setUp(self): - self.dag_path = "tests/resources/data/dag.dot" - self.data_paths = ["tests/resources/data/data.csv"] - self.test_config_path = "tests/resources/data/tests.json" - self.output_path = Path("results/results.json") - self.include_edges_path = Path("tests/resources/data/include_edges.dot") - self.exclude_edges_path = Path("tests/resources/data/exclude_edges.dot") - - def test_missing_dag(self): - with self.assertRaises(FileNotFoundError) as e: - CausalTestingPaths("missing.dot", self.data_paths, self.test_config_path, self.output_path).validate_paths() - self.assertEqual("DAG file not found: missing.dot", str(e.exception)) - - def test_missing_data(self): - with self.assertRaises(FileNotFoundError) as e: - CausalTestingPaths(self.dag_path, ["missing.csv"], self.test_config_path, self.output_path).validate_paths() - self.assertEqual("Data file not found: missing.csv", str(e.exception)) - - def test_missing_tests(self): - with self.assertRaises(FileNotFoundError) as e: - CausalTestingPaths(self.dag_path, self.data_paths, "missing.json", self.output_path).validate_paths() - self.assertEqual("Test configuration file not found: missing.json", str(e.exception)) - - def test_output_file_created(self): - self.assertFalse(self.output_path.parent.exists()) - CausalTestingPaths(self.dag_path, self.data_paths, self.test_config_path, self.output_path).validate_paths() - self.assertTrue(self.output_path.parent.exists()) - - def test_missing_include_edges(self): - with self.assertRaises(FileNotFoundError) as e: - CausalTestingPaths( - self.dag_path, - self.data_paths, - self.test_config_path, - self.output_path, - "missing.dot", - self.exclude_edges_path, - ).validate_paths() - self.assertEqual("Data file not found: missing.dot", str(e.exception)) - - def test_missing_exclude_edges(self): - with self.assertRaises(FileNotFoundError) as e: - CausalTestingPaths( - self.dag_path, - self.data_paths, - self.test_config_path, - self.output_path, - self.include_edges_path, - "missing.dot", - ).validate_paths() - self.assertEqual("Data file not found: missing.dot", str(e.exception)) - - def tearDown(self): - if self.output_path.parent.exists(): - shutil.rmtree(self.output_path.parent) - - class TestCausalTestingFramework(unittest.TestCase): def setUp(self): self.dag_path = "tests/resources/data/dag.dot" self.data_paths = ["tests/resources/data/data.csv"] - self.test_config_path = "tests/resources/data/tests.json" + self.test_cases_path = "tests/resources/data/tests.json" self.output_path = Path("results/results.json") self.include_edges_path = "tests/resources/data/include_edges.dot" self.exclude_edges_path = "tests/resources/data/exclude_edges.dot" - self.paths = CausalTestingPaths( - dag_path=self.dag_path, - data_paths=self.data_paths, - test_config_path=self.test_config_path, - output_path=self.output_path, - include_edges_path=self.include_edges_path, - exclude_edges_path=self.exclude_edges_path, - ) + self.paths = { + "dag_path": self.dag_path, + "data_paths": self.data_paths, + "test_cases_path": self.test_cases_path, + } def test_load_data(self): - csv_framework = CausalTestingFramework(self.paths) - csv_df = csv_framework.load_data() - - pqt_framework = CausalTestingFramework( - CausalTestingPaths( - dag_path=self.dag_path, - data_paths=["tests/resources/data/data.pqt"], - test_config_path=self.test_config_path, - output_path=self.output_path, - ) - ) - pqt_df = pqt_framework.load_data() - pd.testing.assert_frame_equal(csv_df, pqt_df) - - def test_load_data_invalid(self): - framework = CausalTestingFramework( - CausalTestingPaths( - dag_path=self.dag_path, - data_paths=[self.dag_path], - test_config_path=self.test_config_path, - output_path=self.output_path, - ) - ) - with self.assertRaises(ValueError): - framework.load_data() + csv_framework = CausalTestingFramework() + csv_framework.load_data(self.data_paths) + + pqt_framework = CausalTestingFramework() + pqt_framework.load_data([path.replace(".csv", ".pqt") for path in self.data_paths]) + pd.testing.assert_frame_equal(csv_framework.df, pqt_framework.df) def test_load_data_query(self): - framework = CausalTestingFramework(self.paths) - self.assertFalse((framework.load_data()["test_input"] > 4).all()) - self.assertTrue((framework.load_data("test_input > 4")["test_input"] > 4).all()) + framework = CausalTestingFramework() + framework.load_data(data_paths=self.data_paths) + self.assertFalse((framework.df["test_input"] > 4).all()) + + framework.load_data(data_paths=self.data_paths, query="test_input > 4") + self.assertTrue((framework.df["test_input"] > 4).all()) + + def test_load_data_invalid_extension(self): + framework = CausalTestingFramework() + with self.assertRaises(ValueError): + framework.load_data("data.invalid") def test_load_dag_missing_node(self): - framework = CausalTestingFramework(self.paths) - framework.setup() + framework = CausalTestingFramework() + framework.setup(**self.paths) framework.dag.add_node("missing") with self.assertRaises(ValueError): framework.create_variables() + def test_load_tests_before_dag(self): + framework = CausalTestingFramework() + with self.assertRaises(ValueError): + framework.load_test_cases_from_json(self.test_cases_path) + def test_create_base_test_case_missing_treatment(self): - framework = CausalTestingFramework(self.paths) - framework.setup() + framework = CausalTestingFramework() + framework.setup(**self.paths) with self.assertRaises(KeyError) as e: framework.create_base_test( {"treatment_variable": "missing", "expected_effect": {"test_outcome": "NoEffect"}} @@ -136,17 +68,25 @@ def test_create_base_test_case_missing_treatment(self): self.assertEqual("\"Treatment variable 'missing' not found in inputs or outputs\"", str(e.exception)) def test_create_base_test_case_missing_estimator(self): - framework = CausalTestingFramework(self.paths) - framework.setup() + framework = CausalTestingFramework() + framework.setup(**self.paths) with self.assertRaises(ValueError) as e: - framework.create_causal_test({}, None) + framework.create_causal_test( + {"treatment_variable": "test_input", "expected_effect": {"test_output": "NoEffect"}} + ) self.assertEqual("Test configuration must specify an estimator", str(e.exception)) def test_create_test_case_invalid_estimator(self): - framework = CausalTestingFramework(self.paths) - framework.setup() + framework = CausalTestingFramework() + framework.setup(**self.paths) with self.assertRaises(ValueError) as e: - framework.create_causal_test({"estimator": "InvalidEstimator"}, None) + framework.create_causal_test( + { + "treatment_variable": "test_input", + "expected_effect": {"test_output": "NoEffect"}, + "estimator": "InvalidEstimator", + } + ) self.assertEqual( f"Unsupported estimator InvalidEstimator. Supported: ['CubicSplineEstimator', 'IPCWEstimator', 'InstrumentalVariableEstimator', 'LinearRegressionEstimator', 'LogisticRegressionEstimator', 'MultinomialRegressionEstimator']. " "If you have implemented a custom estimator, you will need to add this to your entrypoints via your " @@ -155,8 +95,8 @@ def test_create_test_case_invalid_estimator(self): ) def test_create_test_case_invalid_effect(self): - framework = CausalTestingFramework(self.paths) - framework.setup() + framework = CausalTestingFramework() + framework.setup(**self.paths) test = { "name": "test1", "treatment_variable": "test_input", @@ -166,7 +106,7 @@ def test_create_test_case_invalid_effect(self): } base_test_case = framework.create_base_test(test) with self.assertRaises(ValueError) as e: - framework.create_causal_test(test, base_test_case) + framework.create_causal_test(test) self.assertEqual( f"Unsupported causal effect InvalidEffect. Supported: ['ExactValue', 'Negative', 'NoEffect', 'Positive', 'SomeEffect']. " "If you have implemented a custom causal effect, you will need to add this to your entrypoints via your " @@ -175,8 +115,8 @@ def test_create_test_case_invalid_effect(self): ) def test_create_test_case_effect_kwargs(self): - framework = CausalTestingFramework(self.paths) - framework.setup() + framework = CausalTestingFramework() + framework.setup(**self.paths) test = { "name": "test1", "treatment_variable": "test_input", @@ -186,12 +126,12 @@ def test_create_test_case_effect_kwargs(self): "effect_kwargs": {"value": 4}, } base_test_case = framework.create_base_test(test) - test_case = framework.create_causal_test(test, base_test_case) + test_case = framework.create_causal_test(test) self.assertEqual(test_case.expected_causal_effect.value, 4) def test_create_test_case_estimator_kwargs(self): - framework = CausalTestingFramework(self.paths) - framework.setup() + framework = CausalTestingFramework() + framework.setup(**self.paths) test = { "name": "test1", "treatment_variable": "test_input", @@ -201,123 +141,67 @@ def test_create_test_case_estimator_kwargs(self): "estimator_kwargs": {"instrument": "instrumental_variable"}, } base_test_case = framework.create_base_test(test) - test_case = framework.create_causal_test(test, base_test_case) + test_case = framework.create_causal_test(test) self.assertEqual(test_case.estimator.instrument, "instrumental_variable") def test_create_base_test_case_missing_outcome(self): - framework = CausalTestingFramework(self.paths) - framework.setup() + framework = CausalTestingFramework() + framework.setup(**self.paths) with self.assertRaises(KeyError) as e: framework.create_base_test({"treatment_variable": "test_input", "expected_effect": {"missing": "NoEffect"}}) self.assertEqual("\"Outcome variable 'missing' not found in inputs or outputs\"", str(e.exception)) def test_unloaded_tests(self): - framework = CausalTestingFramework(self.paths) + framework = CausalTestingFramework() with self.assertRaises(ValueError) as e: framework.run_tests() - self.assertEqual("No tests loaded. Call load_tests() first.", str(e.exception)) - - def test_unloaded_tests_batches(self): - framework = CausalTestingFramework(self.paths) - with self.assertRaises(ValueError) as e: - next(framework.run_tests_in_batches()) - self.assertEqual("No tests loaded. Call load_tests() first.", str(e.exception)) + self.assertEqual("No tests to run.", str(e.exception)) def test_ctf(self): - framework = CausalTestingFramework(self.paths) - framework.setup() + framework = CausalTestingFramework() + framework.setup(**self.paths) + framework.run_tests() + json_results = framework.save_results(self.output_path) - framework.load_tests() - results = framework.run_tests() - json_results = framework.save_results(results) - - with open(self.test_config_path, "r", encoding="utf-8") as f: + with open(self.test_cases_path, "r", encoding="utf-8") as f: test_configs = json.load(f) - self.assertEqual(len(json_results), len(test_configs["tests"])) - - result_index = 0 - for i, test_config in enumerate(test_configs["tests"]): - result = json_results[i] - - if test_config.get("skip", False): - self.assertEqual(result["skip"], True) - self.assertEqual(result["passed"], None) - self.assertEqual(result["result"]["status"], "skipped") - else: - test_case = framework.test_cases[result_index] - framework_result = results[result_index] - result_index += 1 - - test_passed = ( - test_case.expected_causal_effect.apply(framework_result) - if framework_result.effect_estimate is not None - else False - ) - self.assertEqual(result["passed"], test_passed) - - def test_ctf_batches(self): - framework = CausalTestingFramework(self.paths) - framework.setup() - - framework.load_tests() + self.assertEqual(len(json_results), len(test_configs["tests"])) - output_files = [] - with tempfile.TemporaryDirectory() as tmpdir: - for i, results in enumerate(framework.run_tests_in_batches()): - temp_file_path = os.path.join(tmpdir, f"output_{i}.json") - framework.save_results(results, temp_file_path) - output_files.append(temp_file_path) - del results + result_index = 0 + for i, test_config in enumerate(test_configs["tests"]): + result = json_results[i] - all_results = [] - for file_path in output_files: - with open(file_path, "r", encoding="utf-8") as f: - all_results.extend(json.load(f)) + if test_config.get("skip", False): + self.assertEqual(result["skip"], True) + self.assertEqual(result["passed"], None) + self.assertEqual(result["result"]["status"], "skipped") + else: + test_case = framework.test_cases[result_index] + result_index += 1 - executed_results = [result for result in all_results if not result.get("skip", False)] - self.assertEqual([result["passed"] for result in executed_results], [True]) + test_passed = ( + test_case.expected_causal_effect.apply(test_case.result) + if test_case.result.effect_estimate is not None + else False + ) + self.assertEqual(result["passed"], test_passed) def test_ctf_exception(self): - framework = CausalTestingFramework(self.paths, query="test_input < 0") - framework.setup() + framework = CausalTestingFramework(self.paths) + framework.setup(**self.paths, query="test_input < 0") - framework.load_tests() with self.assertRaises(ValueError): framework.run_tests() - def test_ctf_batches_exception_silent(self): - framework = CausalTestingFramework(self.paths, query="test_input < 0") - framework.setup() - - framework.load_tests() - - output_files = [] - with tempfile.TemporaryDirectory() as tmpdir: - for i, results in enumerate(framework.run_tests_in_batches(silent=True)): - temp_file_path = os.path.join(tmpdir, f"output_{i}.json") - framework.save_results(results, temp_file_path) - output_files.append(temp_file_path) - del results - - all_results = [] - for file_path in output_files: - with open(file_path, "r", encoding="utf-8") as f: - all_results.extend(json.load(f)) - - executed_results = [result for result in all_results if not result.get("skip", False)] - self.assertEqual([result["passed"] for result in executed_results], [False]) - self.assertIsNotNone([result.get("error") for result in executed_results]) - def test_ctf_exception_silent(self): - framework = CausalTestingFramework(self.paths, query="test_input < 0") - framework.setup() + framework = CausalTestingFramework(self.paths) + framework.setup(**self.paths, query="test_input < 0") - framework.load_tests() - results = framework.run_tests(silent=True) - json_results = framework.save_results(results) + framework.run_tests(silent=True) + json_results = framework.save_results(self.output_path) - with open(self.test_config_path, "r", encoding="utf-8") as f: + with open(self.test_cases_path, "r", encoding="utf-8") as f: test_configs = json.load(f) non_skipped_configs = [t for t in test_configs["tests"] if not t.get("skip", False)] @@ -328,105 +212,6 @@ def test_ctf_exception_silent(self): for result in non_skipped_results: self.assertEqual(result["passed"], False) - def test_ctf_batches_exception(self): - framework = CausalTestingFramework(self.paths, query="test_input < 0") - framework.setup() - - framework.load_tests() - with self.assertRaises(ValueError): - next(framework.run_tests_in_batches()) - - def test_ctf_batches_matches_run_tests(self): - framework = CausalTestingFramework(self.paths) - framework.setup() - framework.load_tests() - normal_results = framework.run_tests() - - output_files = [] - with tempfile.TemporaryDirectory() as tmpdir: - for i, results in enumerate(framework.run_tests_in_batches()): - temp_file_path = os.path.join(tmpdir, f"output_{i}.json") - framework.save_results(results, temp_file_path) - output_files.append(temp_file_path) - del results - - all_results = [] - for file_path in output_files: - with open(file_path, "r", encoding="utf-8") as f: - all_results.extend(json.load(f)) - - with tempfile.TemporaryDirectory() as tmpdir: - normal_output = os.path.join(tmpdir, "normal.json") - framework.save_results(normal_results, normal_output) - with open(normal_output) as f: - normal_json = json.load(f) - - batch_output = os.path.join(tmpdir, "batch.json") - with open(batch_output, "w") as f: - json.dump(all_results, f) - with open(batch_output) as f: - batch_json = json.load(f) - - self.assertEqual(normal_json, batch_json) - - def test_global_query(self): - framework = CausalTestingFramework(self.paths) - framework.setup() - - query_framework = CausalTestingFramework(self.paths, query="test_input > 0") - query_framework.setup() - - self.assertTrue(len(query_framework.data) > 0) - self.assertTrue((query_framework.data["test_input"] > 0).all()) - - with open(self.test_config_path, "r", encoding="utf-8") as f: - test_configs = json.load(f) - - test_config = test_configs["tests"][0].copy() - if "query" in test_config: - del test_config["query"] - - base_test = query_framework.create_base_test(test_config) - causal_test = query_framework.create_causal_test(test_config, base_test) - - self.assertTrue((causal_test.estimator.df["test_input"] > 0).all()) - - query_framework.create_variables() - self.assertIsNotNone(query_framework.scenario) - - def test_test_specific_query(self): - framework = CausalTestingFramework(self.paths) - framework.setup() - - with open(self.test_config_path, "r", encoding="utf-8") as f: - test_configs = json.load(f) - - test_config = test_configs["tests"][0].copy() - test_config["query"] = "test_input > 0" - - base_test = framework.create_base_test(test_config) - causal_test = framework.create_causal_test(test_config, base_test) - - self.assertTrue(len(causal_test.estimator.df) > 0) - self.assertTrue((causal_test.estimator.df["test_input"] > 0).all()) - - def test_combined_queries(self): - global_framework = CausalTestingFramework(self.paths, query="test_input > 0") - global_framework.setup() - - with open(self.test_config_path, "r", encoding="utf-8") as f: - test_configs = json.load(f) - - test_config = test_configs["tests"][0].copy() - test_config["query"] = "test_output > 0" - - base_test = global_framework.create_base_test(test_config) - causal_test = global_framework.create_causal_test(test_config, base_test) - - self.assertTrue(len(causal_test.estimator.df) > 0) - self.assertTrue((causal_test.estimator.df["test_input"] > 0).all()) - self.assertTrue((causal_test.estimator.df["test_output"] > 0).all()) - def test_parse_args(self): with patch( "sys.argv", @@ -438,7 +223,7 @@ def test_parse_args(self): "--data-paths", str(self.data_paths[0]), "--test-config", - str(self.test_config_path), + str(self.test_cases_path), "--output", str(self.output_path.parent / "main.json"), ], @@ -457,35 +242,10 @@ def test_parse_args_adequacy(self): "--data-paths", str(self.data_paths[0]), "--test-config", - str(self.test_config_path), - "--output", - str(self.output_path.parent / "main.json"), - "-a", - ], - ): - main() - with open(self.output_path.parent / "main.json") as f: - log = json.load(f) - executed_tests = [test for test in log if not test.get("skip", False)] - assert all(test["result"].get("bootstrap_size", 100) == 100 for test in executed_tests) - - def test_parse_args_adequacy_batches(self): - with patch( - "sys.argv", - [ - "causal_testing", - "test", - "--dag-path", - str(self.dag_path), - "--data-paths", - str(self.data_paths[0]), - "--test-config", - str(self.test_config_path), + str(self.test_cases_path), "--output", str(self.output_path.parent / "main.json"), "-a", - "--batch-size", - "5", ], ): main() @@ -505,7 +265,7 @@ def test_parse_args_bootstrap_size(self): "--data-paths", str(self.data_paths[0]), "--test-config", - str(self.test_config_path), + str(self.test_cases_path), "--output", str(self.output_path.parent / "main.json"), "-b", @@ -529,7 +289,7 @@ def test_parse_args_bootstrap_size_explicit_adequacy(self): "--data-paths", str(self.data_paths[0]), "--test-config", - str(self.test_config_path), + str(self.test_cases_path), "--output", str(self.output_path.parent / "main.json"), "-a", @@ -543,27 +303,6 @@ def test_parse_args_bootstrap_size_explicit_adequacy(self): executed_tests = [test for test in log if not test.get("skip", False)] assert all(test["result"].get("bootstrap_size", 50) == 50 for test in executed_tests) - def test_parse_args_batches(self): - with patch( - "sys.argv", - [ - "causal_testing", - "test", - "--dag-path", - str(self.dag_path), - "--data-paths", - str(self.data_paths[0]), - "--test-config", - str(self.test_config_path), - "--output", - str(self.output_path.parent / "main_batch.json"), - "--batch-size", - "5", - ], - ): - main() - self.assertTrue((self.output_path.parent / "main_batch.json").exists()) - def test_parse_args_generation(self): with tempfile.TemporaryDirectory() as tmp: with patch( diff --git a/tests/resources/data/tests.json b/tests/resources/data/tests.json index 8354a71a..54500c43 100644 --- a/tests/resources/data/tests.json +++ b/tests/resources/data/tests.json @@ -6,8 +6,8 @@ "estimate_type": "coefficient", "expected_effect": {"test_output": "NoEffect"}, "skip": false, + "query": "test_input > 0", "estimator_kwargs": { - "query": "test_input > 0", "effect_modifiers": [] } }, @@ -18,9 +18,9 @@ "estimate_type": "coefficient", "expected_effect": {"test_output": "NoEffect"}, "skip": true, + "query": "test_input <= 5", "estimator_kwargs": { - "effect_modifiers": [], - "query": "test_input <= 5" + "effect_modifiers": [] } }] } diff --git a/tests/surrogate_tests/test_causal_surrogate_assisted.py b/tests/surrogate_tests/test_causal_surrogate_assisted.py index 37f8d926..31a4bc04 100644 --- a/tests/surrogate_tests/test_causal_surrogate_assisted.py +++ b/tests/surrogate_tests/test_causal_surrogate_assisted.py @@ -19,11 +19,9 @@ class TestSimulationResult(unittest.TestCase): def setUp(self): - self.data = {"key": "value"} def test_inputs(self): - fault_values = [True, False] relationship_values = ["positive", "negative", None] @@ -55,9 +53,6 @@ def setUp(self): f.write(dag_dot) def test_surrogate_model_generation(self): - - df = self.class_df.copy() - causal_dag = CausalDAG(self.dag_dot_path) z = Input("Z", int) x = Input("X", float) @@ -66,7 +61,7 @@ def test_surrogate_model_generation(self): scenario = Scenario(variables={z, x, m, y}) c_s_a_test_case = CausalSurrogateAssistedTestCase(scenario, causal_dag, None, None) - surrogate_models = c_s_a_test_case.generate_surrogates(df) + surrogate_models = c_s_a_test_case.generate_surrogates() self.assertEqual(len(surrogate_models), 2) for surrogate_model in surrogate_models: diff --git a/tests/testing_tests/test_causal_effect.py b/tests/testing_tests/test_causal_effect.py index 29eeb61d..c0293c00 100644 --- a/tests/testing_tests/test_causal_effect.py +++ b/tests/testing_tests/test_causal_effect.py @@ -23,74 +23,40 @@ def setUp(self) -> None: def test_None_ci(self): ctr = CausalTestResult( - estimator=self.estimator, effect_estimate=EffectEstimate(type="ate", value=pd.Series(0)), ) self.assertIsNone(ctr.effect_estimate.ci_low) self.assertIsNone(ctr.effect_estimate.ci_high) - self.assertEqual( - ctr.to_dict(), - { - "treatment": "A", - "control_value": 0, - "treatment_value": 1, - "outcome": "A", - "adjustment_set": set(), - "effect_estimate": {0: 0}, - "effect_measure": "ate", - }, - ) def test_empty_adjustment_set(self): ctr = CausalTestResult( - estimator=self.estimator, effect_estimate=EffectEstimate(type="ate", value=pd.Series(0)), ) self.assertIsNone(ctr.effect_estimate.ci_low) self.assertIsNone(ctr.effect_estimate.ci_high) - self.assertEqual( - str(ctr), - ( - "Causal Test Result\n==============\n" - "Treatment: A\n" - "Control value: 0\n" - "Treatment value: 1\n" - "Outcome: A\n" - "Adjustment set: set()\n" - "Formula: A ~ A\n" - "ate: {0: 0}\n" - ), - ) def test_Positive_ate_pass(self): ctr = CausalTestResult( - estimator=self.estimator, - effect_estimate=EffectEstimate(type="ate", value=pd.Series(5.05)), + effect_estimate=EffectEstimate( + type="ate", value=pd.Series(5.05), ci_low=pd.Series(5), ci_high=pd.Series(6) + ), ) ev = Positive() self.assertTrue(ev.apply(ctr)) def test_Positive_risk_ratio_pass(self): ctr = CausalTestResult( - estimator=self.estimator, - effect_estimate=EffectEstimate(type="risk_ratio", value=pd.Series(5.05)), + effect_estimate=EffectEstimate( + type="risk_ratio", value=pd.Series(5.05), ci_low=pd.Series(5), ci_high=pd.Series(6) + ), ) ev = Positive() self.assertTrue(ev.apply(ctr)) def test_Positive_fail(self): ctr = CausalTestResult( - estimator=self.estimator, - effect_estimate=EffectEstimate(type="ate", value=pd.Series(0)), - ) - ev = Positive() - self.assertFalse(ev.apply(ctr)) - - def test_Positive_fail_ci(self): - ctr = CausalTestResult( - estimator=self.estimator, effect_estimate=EffectEstimate(type="ate", value=pd.Series(0), ci_low=pd.Series(-1), ci_high=pd.Series(1)), ) ev = Positive() @@ -98,31 +64,24 @@ def test_Positive_fail_ci(self): def test_Negative_ate_pass(self): ctr = CausalTestResult( - estimator=self.estimator, - effect_estimate=EffectEstimate(type="ate", value=pd.Series(-5.05)), + effect_estimate=EffectEstimate( + type="ate", value=pd.Series(-5.05), ci_low=pd.Series(-6), ci_high=pd.Series(-5) + ), ) ev = Negative() self.assertTrue(ev.apply(ctr)) def test_Negative_risk_ratio_pass(self): ctr = CausalTestResult( - estimator=self.estimator, - effect_estimate=EffectEstimate(type="risk_ratio", value=pd.Series(0.2)), + effect_estimate=EffectEstimate( + type="risk_ratio", value=pd.Series(0.2), ci_low=pd.Series(0.1), ci_high=pd.Series(0.5) + ), ) ev = Negative() self.assertTrue(ev.apply(ctr)) def test_Negative_fail(self): ctr = CausalTestResult( - estimator=self.estimator, - effect_estimate=EffectEstimate(type="ate", value=pd.Series(0)), - ) - ev = Negative() - self.assertFalse(ev.apply(ctr)) - - def test_Negative_fail_ci(self): - ctr = CausalTestResult( - estimator=self.estimator, effect_estimate=EffectEstimate(type="ate", value=pd.Series(0), ci_low=pd.Series(-1), ci_high=pd.Series(1)), ) ev = Negative() @@ -130,7 +89,6 @@ def test_Negative_fail_ci(self): def test_exactValue_pass(self): ctr = CausalTestResult( - estimator=self.estimator, effect_estimate=EffectEstimate(type="ate", value=pd.Series(5.05)), ) ev = ExactValue(5, 0.1) @@ -138,7 +96,6 @@ def test_exactValue_pass(self): def test_exactValue_pass_ci(self): ctr = CausalTestResult( - estimator=self.estimator, effect_estimate=EffectEstimate( type="ate", value=pd.Series(5.05), ci_low=pd.Series(4), ci_high=pd.Series(6) ), @@ -148,7 +105,6 @@ def test_exactValue_pass_ci(self): def test_exactValue_ci_pass_ci(self): ctr = CausalTestResult( - estimator=self.estimator, effect_estimate=EffectEstimate( type="ate", value=pd.Series(5.05), ci_low=pd.Series(4.1), ci_high=pd.Series(5.9) ), @@ -158,7 +114,6 @@ def test_exactValue_ci_pass_ci(self): def test_exactValue_ci_fail_ci(self): ctr = CausalTestResult( - estimator=self.estimator, effect_estimate=EffectEstimate( type="ate", value=pd.Series(5.05), ci_low=pd.Series(3.9), ci_high=pd.Series(6.1) ), @@ -168,7 +123,6 @@ def test_exactValue_ci_fail_ci(self): def test_exactValue_fail(self): ctr = CausalTestResult( - estimator=self.estimator, effect_estimate=EffectEstimate(type="ate", value=pd.Series(0)), ) ev = ExactValue(5, 0.1) @@ -196,7 +150,6 @@ def test_invalid_ci_atol(self): def test_invalid(self): ctr = CausalTestResult( - estimator=self.estimator, effect_estimate=EffectEstimate( type="invalid", value=pd.Series(5.05), ci_low=pd.Series(4.8), ci_high=pd.Series(6.7) ), @@ -212,7 +165,6 @@ def test_invalid(self): def test_someEffect_pass_coefficient(self): ctr = CausalTestResult( - estimator=self.estimator, effect_estimate=EffectEstimate( type="coefficient", value=pd.Series(5.05), ci_low=pd.Series(4.8), ci_high=pd.Series(6.7) ), @@ -222,7 +174,6 @@ def test_someEffect_pass_coefficient(self): def test_someEffect_pass_ate(self): ctr = CausalTestResult( - estimator=self.estimator, effect_estimate=EffectEstimate( type="coefficient", value=pd.Series(5.05), ci_low=pd.Series(4.8), ci_high=pd.Series(6.7) ), @@ -232,7 +183,6 @@ def test_someEffect_pass_ate(self): def test_someEffect_pass_rr(self): ctr = CausalTestResult( - estimator=self.estimator, effect_estimate=EffectEstimate( type="coefficient", value=pd.Series(5.05), ci_low=pd.Series(4.8), ci_high=pd.Series(6.7) ), @@ -242,7 +192,6 @@ def test_someEffect_pass_rr(self): def test_someEffect_fail(self): ctr = CausalTestResult( - estimator=self.estimator, effect_estimate=EffectEstimate( type="ate", value=pd.Series(0), ci_low=pd.Series(-0.1), ci_high=pd.Series(0.2) ), @@ -252,33 +201,10 @@ def test_someEffect_fail(self): def test_someEffect_None(self): ctr = CausalTestResult( - estimator=self.estimator, effect_estimate=EffectEstimate(type="ate", value=pd.Series(0)), ) self.assertEqual(SomeEffect().apply(ctr), None) - def test_someEffect_dict(self): - ctr = CausalTestResult( - estimator=self.estimator, - effect_estimate=EffectEstimate( - type="ate", value=pd.Series(0), ci_low=pd.Series(-0.1), ci_high=pd.Series(0.2) - ), - ) - self.assertEqual( - ctr.to_dict(), - { - "treatment": "A", - "control_value": 0, - "treatment_value": 1, - "outcome": "A", - "adjustment_set": set(), - "effect_estimate": {0: 0}, - "effect_measure": "ate", - "ci_low": {0: -0.1}, - "ci_high": {0: 0.2}, - }, - ) - def test_positive_risk_ratio_e_value(self): cv = CausalValidator() e_value = cv.estimate_e_value(1.5) @@ -301,7 +227,6 @@ def test_negative_risk_ratio_e_value_using_ci(self): def test_multiple_value_exception_caught(self): ctr = CausalTestResult( - estimator=self.estimator, effect_estimate=EffectEstimate(type="ate", value=pd.Series([0, 1])), ) with self.assertRaises(ValueError): diff --git a/tests/testing_tests/test_causal_test_adequacy.py b/tests/testing_tests/test_causal_test_adequacy.py index 9316d74e..36ed256a 100644 --- a/tests/testing_tests/test_causal_test_adequacy.py +++ b/tests/testing_tests/test_causal_test_adequacy.py @@ -7,10 +7,10 @@ from causal_testing.estimation.ipcw_estimator import IPCWEstimator from causal_testing.testing.base_test_case import BaseTestCase from causal_testing.testing.causal_test_case import CausalTestCase -from causal_testing.testing.causal_test_adequacy import DAGAdequacy +from causal_testing.testing.dag_adequacy import DAGAdequacy from causal_testing.testing.causal_effect import NoEffect, SomeEffect from causal_testing.specification.scenario import Scenario -from causal_testing.testing.causal_test_adequacy import DataAdequacy +from causal_testing.testing.data_adequacy import DataAdequacy from causal_testing.specification.variable import Input, Output from causal_testing.specification.causal_dag import CausalDAG @@ -37,7 +37,7 @@ def test_data_adequacy_numeric(self): Input("test_input", float, self.example_distribution), Output("test_output", float) ) estimator = LinearRegressionEstimator( - base_test_case=base_test_case, treatment_value=None, control_value=None, adjustment_set={}, df=self.df + base_test_case=base_test_case, treatment_value=None, control_value=None, adjustment_set={} ) causal_test_case = CausalTestCase( base_test_case=base_test_case, @@ -45,36 +45,29 @@ def test_data_adequacy_numeric(self): estimate_type="coefficient", estimator=estimator, ) - adequacy_metric = DataAdequacy(causal_test_case) - adequacy_metric.measure_adequacy() + adequacy_metric = causal_test_case.measure_adequacy(self.df) self.assertAlmostEqual( adequacy_metric.kurtosis["test_input"], 0, delta=1.0, msg=f"Expected kurtosis near 0, got {adequacy_metric.kurtosis['test_input']}", - ) # This adds a numerical tolerance for Pandas - self.assertEqual( - adequacy_metric.bootstrap_size, 100, f"Expected bootstrap size 100 not {adequacy_metric.bootstrap_size}" - ) - self.assertEqual(adequacy_metric.passing, 100, f"Expected passing 32 not {adequacy_metric.passing}") + ) # This adds a numerical tolerance for Pandas + self.assertEqual(adequacy_metric.passing, 19, f"Expected passing 19 not {adequacy_metric.passing}") self.assertEqual(adequacy_metric.successful, 100, f"Expected successful 100 not {adequacy_metric.successful}") def test_data_adequacy_categorical(self): base_test_case = BaseTestCase( Input("test_input_no_dist", float, self.example_distribution), Output("test_output", float) ) - estimator = LinearRegressionEstimator( - base_test_case=base_test_case, treatment_value=None, control_value=None, adjustment_set={}, df=self.df - ) + estimator = LinearRegressionEstimator(base_test_case=base_test_case, adjustment_set={}) causal_test_case = CausalTestCase( base_test_case=base_test_case, expected_causal_effect=NoEffect(), estimate_type="coefficient", estimator=estimator, ) - adequacy_metric = DataAdequacy(causal_test_case) - adequacy_metric.measure_adequacy() + adequacy_metric = causal_test_case.measure_adequacy(self.df) self.assertAlmostEqual( adequacy_metric.kurtosis["test_input_no_dist[T.b]"], @@ -82,9 +75,6 @@ def test_data_adequacy_categorical(self): delta=1.0, msg=f"Expected kurtosis near 0, got {adequacy_metric.kurtosis['test_input_no_dist[T.b]']}", ) - self.assertEqual( - adequacy_metric.bootstrap_size, 100, f"Expected bootstrap size 100 not {adequacy_metric.bootstrap_size}" - ) self.assertEqual(adequacy_metric.passing, 100, f"Expected passing 100 not {adequacy_metric.passing}") self.assertEqual(adequacy_metric.successful, 100, f"Expected successful 100 not {adequacy_metric.successful}") @@ -97,12 +87,11 @@ def test_data_adequacy_group_by(self): df = pd.read_csv("tests/resources/data/temporal_data.csv") df["ok"] = df["outcome"] == 1 estimation_model = IPCWEstimator( - df, - timesteps_per_intervention, - control_strategy, - treatment_strategy, - outcome, - "ok", + timesteps_per_observation=timesteps_per_intervention, + control_strategy=control_strategy, + treatment_strategy=treatment_strategy, + outcome=outcome, + status_column="ok", fit_bl_switch_formula=fit_bl_switch_formula, fit_bltd_switch_formula=fit_bl_switch_formula, eligibility=None, @@ -115,19 +104,15 @@ def test_data_adequacy_group_by(self): estimate_type="hazard_ratio", estimator=estimation_model, ) - adequacy_metric = DataAdequacy(causal_test_case, group_by="id") - adequacy_metric.measure_adequacy() + adequacy_metric = causal_test_case.measure_adequacy(df, group_by="id") self.assertEqual( round(adequacy_metric.kurtosis["trtrand"], 3), - -0.857, + -2.739, f"Expected kurtosis not {round(adequacy_metric.kurtosis['trtrand'], 3)}", ) - self.assertEqual( - adequacy_metric.bootstrap_size, 100, f"Expected bootstrap size 100 not {adequacy_metric.bootstrap_size}" - ) - self.assertEqual(adequacy_metric.passing, 32, f"Expected passing 32 not {adequacy_metric.passing}") - self.assertEqual(adequacy_metric.successful, 100, f"Expected successful 100 not {adequacy_metric.successful}") + self.assertEqual(adequacy_metric.passing, 1, f"Expected passing 1 not {adequacy_metric.passing}") + self.assertEqual(adequacy_metric.successful, 5, f"Expected successful 5 not {adequacy_metric.successful}") def test_dag_adequacy_dependent(self): base_test_case = BaseTestCase( diff --git a/tests/testing_tests/test_causal_test_case.py b/tests/testing_tests/test_causal_test_case.py index 8996ab3f..53bf0b4e 100644 --- a/tests/testing_tests/test_causal_test_case.py +++ b/tests/testing_tests/test_causal_test_case.py @@ -118,13 +118,17 @@ def test_execute_test_observational_linear_regression_estimator(self): """Check that executing the causal test case returns the correct results for dummy data using a linear regression estimator.""" estimation_model = LinearRegressionEstimator( - self.base_test_case_A_C, - self.treatment_value, - self.control_value, - self.minimal_adjustment_set, - self.df, + base_test_case=self.base_test_case_A_C, + treatment_value=self.treatment_value, + control_value=self.control_value, + adjustment_set=self.minimal_adjustment_set, + ) + causal_test_case = CausalTestCase( + base_test_case=self.base_test_case_A_C, + expected_causal_effect=self.expected_causal_effect, + estimator=estimation_model, ) - causal_test_result = self.causal_test_case.execute_test(estimation_model) + causal_test_result = causal_test_case.estimate_effect(self.df) pd.testing.assert_series_equal(causal_test_result.effect_estimate.value, pd.Series(4.0), atol=1e-10) def test_execute_test_observational_linear_regression_estimator_direct_effect(self): @@ -132,11 +136,10 @@ def test_execute_test_observational_linear_regression_estimator_direct_effect(se regression estimator.""" base_test_case = BaseTestCase(treatment_variable=self.A, outcome_variable=self.C, effect="direct") estimation_model = LinearRegressionEstimator( - self.base_test_case_A_C, - self.treatment_value, - self.control_value, - self.causal_dag.identification(base_test_case), - self.df, + base_test_case=self.base_test_case_A_C, + treatment_value=self.treatment_value, + control_value=self.control_value, + adjustment_set=self.causal_dag.identification(base_test_case), ) causal_test_case = CausalTestCase( @@ -148,103 +151,102 @@ def test_execute_test_observational_linear_regression_estimator_direct_effect(se # 6. Easier to access treatment and outcome values self.treatment_value = 1 self.control_value = 0 - causal_test_result = causal_test_case.execute_test() + causal_test_result = causal_test_case.estimate_effect(self.df) pd.testing.assert_series_equal(causal_test_result.effect_estimate.value, pd.Series(4.0), atol=1e-10) def test_execute_test_observational_linear_regression_estimator_coefficient(self): """Check that executing the causal test case returns the correct results for dummy data using a linear regression estimator.""" estimation_model = LinearRegressionEstimator( - self.base_test_case_D_A, - self.treatment_value, - self.control_value, - self.minimal_adjustment_set, - self.df, - ) - self.causal_test_case.estimate_type = "coefficient" - causal_test_result = self.causal_test_case.execute_test(estimation_model) + base_test_case=self.base_test_case_D_A, + treatment_value=self.treatment_value, + control_value=self.control_value, + adjustment_set=self.minimal_adjustment_set, + ) + causal_test_case = CausalTestCase( + base_test_case=self.base_test_case_A_C, + expected_causal_effect=self.expected_causal_effect, + estimator=estimation_model, + estimate_type="coefficient", + ) + causal_test_result = causal_test_case.estimate_effect(self.df) pd.testing.assert_series_equal(causal_test_result.effect_estimate.value, pd.Series({"D": 0.0}), atol=1e-1) def test_execute_test_observational_linear_regression_estimator_risk_ratio(self): """Check that executing the causal test case returns the correct results for dummy data using a linear regression estimator.""" estimation_model = LinearRegressionEstimator( - self.base_test_case_D_A, - self.treatment_value, - self.control_value, - self.minimal_adjustment_set, - self.df, - ) - self.causal_test_case.estimate_type = "risk_ratio" - causal_test_result = self.causal_test_case.execute_test(estimation_model) + base_test_case=self.base_test_case_D_A, + treatment_value=self.treatment_value, + control_value=self.control_value, + adjustment_set=self.minimal_adjustment_set, + ) + causal_test_case = CausalTestCase( + base_test_case=self.base_test_case_A_C, + expected_causal_effect=self.expected_causal_effect, + estimator=estimation_model, + estimate_type="risk_ratio", + ) + causal_test_result = causal_test_case.estimate_effect(self.df) pd.testing.assert_series_equal(causal_test_result.effect_estimate.value, pd.Series(0.0), atol=1) def test_invalid_estimate_type(self): """Check that executing the causal test case returns the correct results for dummy data using a linear regression estimator.""" estimation_model = LinearRegressionEstimator( - self.base_test_case_D_A, - self.treatment_value, - self.control_value, - self.minimal_adjustment_set, - self.df, + base_test_case=self.base_test_case_D_A, + treatment_value=self.treatment_value, + control_value=self.control_value, + adjustment_set=self.minimal_adjustment_set, + ) + causal_test_case = CausalTestCase( + base_test_case=self.base_test_case_A_C, + expected_causal_effect=self.expected_causal_effect, + estimator=estimation_model, + estimate_type="invalid", ) - self.causal_test_case.estimate_type = "invalid" with self.assertRaises(AttributeError): - self.causal_test_case.execute_test(estimation_model) + causal_test_case.execute_test(self.df) def test_execute_test_observational_linear_regression_estimator_squared_term(self): """Check that executing the causal test case returns the correct results for dummy data with a squared term using a linear regression estimator. C ~ 4*(A+2) + D + D^2""" estimation_model = LinearRegressionEstimator( - self.base_test_case_A_C, - self.treatment_value, - self.control_value, - self.minimal_adjustment_set, - self.df, - formula=f"C ~ A + {'+'.join(self.minimal_adjustment_set)} + (D ** 2)", - ) - causal_test_result = self.causal_test_case.execute_test(estimation_model) - pd.testing.assert_series_equal(causal_test_result.effect_estimate.value, pd.Series(4.0), atol=1) - - def test_estimate_params_none(self): - """Check that estimate_params defaults to empty dict when None is passed into the estimator object""" - estimator = LinearRegressionEstimator( base_test_case=self.base_test_case_A_C, - adjustment_set=set(), - control_value=0, - treatment_value=1, - formula="C ~ A + D", - df=self.df, + treatment_value=self.treatment_value, + control_value=self.control_value, + adjustment_set=self.minimal_adjustment_set, + formula=f"C ~ A + {'+'.join(self.minimal_adjustment_set)} + (D ** 2)", ) causal_test_case = CausalTestCase( base_test_case=self.base_test_case_A_C, expected_causal_effect=self.expected_causal_effect, - estimate_params=None, - estimator=estimator, - estimate_type="risk_ratio", + estimator=estimation_model, ) - self.assertEqual(causal_test_case.estimate_params, {}) - with self.assertRaises(ValueError): - causal_test_case.execute_test() + causal_test_result = causal_test_case.estimate_effect(self.df) + pd.testing.assert_series_equal(causal_test_result.effect_estimate.value, pd.Series(4.0), atol=1) def test_estimate_params_with_formula(self): """Ensure estimate params is handled correctly when a formula is passed into the estimator object""" - estimate_params = {"adjustment_config": {"D": 1}} + estimator = LinearRegressionEstimator( base_test_case=self.base_test_case_A_C, adjustment_set=set(), control_value=0, treatment_value=1, formula="C ~ A + D", - df=self.df, + adjustment_config={"D": 1}, ) causal_test_case = CausalTestCase( base_test_case=self.base_test_case_A_C, expected_causal_effect=self.expected_causal_effect, - estimate_params=estimate_params, estimate_type="risk_ratio", estimator=estimator, ) - self.assertEqual(causal_test_case.estimate_params, estimate_params) - self.assertEqual(round(causal_test_case.execute_test().effect_estimate.value[0], 3), 1.444) + self.assertEqual( + round( + causal_test_case.estimate_effect(self.df).effect_estimate.value[0], + 3, + ), + 1.444, + )