pylint

Author: christopher-wild

causal_testing/data_collection/data_collector.py

@@ -18,6 +18,7 @@ class DataCollector(ABC):
 
     def __init__(self, scenario: Scenario):
         self.scenario = scenario
+
     @abstractmethod
     def collect_data(self, **kwargs) -> pd.DataFrame:
         """

causal_testing/json_front/json_class.py

@@ -5,7 +5,7 @@
 import json
 import logging
 
-from collections.abc import Iterable, Mapping
+from collections.abc import Mapping
 from dataclasses import dataclass
 from pathlib import Path
 from statistics import StatisticsError
@@ -86,11 +86,9 @@ def setup(self, scenario: Scenario):
                 "No data found, either provide a path to a file containing data or manually populate the .data "
                 "attribute with a dataframe before calling .setup()"
             )
-        self.data_collector = ObservationalDataCollector(
-            self.scenario, data)
+        self.data_collector = ObservationalDataCollector(self.scenario, data)
         self._populate_metas()
 
-
     def _create_abstract_test_case(self, test, mutates, effects):
         assert len(test["mutations"]) == 1
         treatment_var = next(self.scenario.variables[v] for v in test["mutations"])
@@ -156,11 +154,11 @@ def run_json_tests(self, effects: dict, estimators: dict, f_flag: bool = False,
                 failed, _ = self._execute_test_case(causal_test_case=causal_test_case, test=test, f_flag=f_flag)
 
                 msg = (
-                        f"Executing concrete test: {test['name']} \n"
-                        + f"treatment variable: {test['treatment_variable']} \n"
-                        + f"outcome_variable = {outcome_variable} \n"
-                        + f"control value = {test['control_value']}, treatment value = {test['treatment_value']} \n"
-                        + f"Result: {'FAILED' if failed else 'Passed'}"
+                    f"Executing concrete test: {test['name']} \n"
+                    + f"treatment variable: {test['treatment_variable']} \n"
+                    + f"outcome_variable = {outcome_variable} \n"
+                    + f"control value = {test['control_value']}, treatment value = {test['treatment_value']} \n"
+                    + f"Result: {'FAILED' if failed else 'Passed'}"
                 )
                 print(msg)
                 self._append_to_file(msg, logging.INFO)
@@ -187,12 +185,12 @@ def _run_coefficient_test(self, test: dict, f_flag: bool, effects: dict):
         )
         result = self._execute_test_case(causal_test_case=causal_test_case, test=test, f_flag=f_flag)
         msg = (
-                f"Executing test: {test['name']} \n"
-                + f"  {causal_test_case} \n"
-                + "  "
-                + ("\n  ").join(str(result[1]).split("\n"))
-                + "==============\n"
-                + f"  Result: {'FAILED' if result[0] else 'Passed'}"
+            f"Executing test: {test['name']} \n"
+            + f"  {causal_test_case} \n"
+            + "  "
+            + ("\n  ").join(str(result[1]).split("\n"))
+            + "==============\n"
+            + f"  Result: {'FAILED' if result[0] else 'Passed'}"
         )
         return msg
 
@@ -220,13 +218,13 @@ def _run_ate_test(self, test: dict, f_flag: bool, effects: dict, mutates: dict):
         failures, _ = self._execute_tests(concrete_tests, test, f_flag)
 
         msg = (
-                f"Executing test: {test['name']} \n"
-                + "  abstract_test \n"
-                + f"  {abstract_test} \n"
-                + f"  {abstract_test.treatment_variable.name},"
-                + f"  {abstract_test.treatment_variable.distribution} \n"
-                + f"  Number of concrete tests for test case: {str(len(concrete_tests))} \n"
-                + f"  {failures}/{len(concrete_tests)} failed for {test['name']}"
+            f"Executing test: {test['name']} \n"
+            + "  abstract_test \n"
+            + f"  {abstract_test} \n"
+            + f"  {abstract_test.treatment_variable.name},"
+            + f"  {abstract_test.treatment_variable.distribution} \n"
+            + f"  Number of concrete tests for test case: {str(len(concrete_tests))} \n"
+            + f"  {failures}/{len(concrete_tests)} failed for {test['name']}"
         )
         return msg
 
@@ -251,8 +249,7 @@ def _populate_metas(self):
             meta.populate(self.data_collector.data)
 
     def _execute_test_case(
-            self, causal_test_case: CausalTestCase, test: Mapping,
-            f_flag: bool
+        self, causal_test_case: CausalTestCase, test: Mapping, f_flag: bool
     ) -> (bool, CausalTestResult):
         """Executes a singular test case, prints the results and returns the test case result
         :param causal_test_case: The concrete test case to be executed
@@ -265,8 +262,9 @@ def _execute_test_case(
         failed = False
 
         estimation_model = self._setup_test(causal_test_case=causal_test_case, test=test)
-        causal_test_result = causal_test_case.execute_test(estimator=estimation_model,
-                                                           data_collector=self.data_collector)
+        causal_test_result = causal_test_case.execute_test(
+            estimator=estimation_model, data_collector=self.data_collector
+        )
 
         test_passes = causal_test_case.expected_causal_effect.apply(causal_test_result)
 
@@ -288,9 +286,7 @@ def _execute_test_case(
             logger.warning("   FAILED- expected %s, got %s", causal_test_case.expected_causal_effect, result_string)
         return failed, causal_test_result
 
-    def _setup_test(
-            self, causal_test_case: CausalTestCase, test: Mapping
-    ) -> Estimator:
+    def _setup_test(self, causal_test_case: CausalTestCase, test: Mapping) -> Estimator:
         """Create the necessary inputs for a single test case
         :param causal_test_case: The concrete test case to be executed
         :param test: Single JSON test definition stored in a mapping (dict)
@@ -368,7 +364,7 @@ def get_args(test_args=None) -> argparse.Namespace:
         parser.add_argument(
             "-w",
             help="Specify to overwrite any existing output files. This can lead to the loss of existing outputs if not "
-                 "careful",
+            "careful",
             action="store_true",
         )
         parser.add_argument(

causal_testing/specification/causal_specification.py

@@ -1,6 +1,5 @@
 """This module holds the abstract CausalSpecification data class, which holds a Scenario and CausalDag"""
 
-from abc import ABC
 from dataclasses import dataclass
 from typing import Union
 

causal_testing/testing/causal_test_case.py

@@ -2,18 +2,14 @@
 import logging
 from typing import Any
 
-import pandas as pd
-
 from causal_testing.specification.variable import Variable
 from causal_testing.testing.causal_test_outcome import CausalTestOutcome
 from causal_testing.testing.base_test_case import BaseTestCase
 from causal_testing.testing.estimators import Estimator
 from causal_testing.testing.causal_test_result import CausalTestResult, TestValue
 from causal_testing.data_collection.data_collector import ObservationalDataCollector
-from causal_testing.specification.causal_dag import CausalDAG
-from causal_testing.specification.scenario import Scenario
 
-from causal_testing.specification.causal_specification import CausalSpecification
+
 logger = logging.getLogger(__name__)
 
 
@@ -31,15 +27,15 @@ class CausalTestCase:
     """
 
     def __init__(
-            # pylint: disable=too-many-arguments
-            self,
-            base_test_case: BaseTestCase,
-            expected_causal_effect: CausalTestOutcome,
-            control_value: Any = None,
-            treatment_value: Any = None,
-            estimate_type: str = "ate",
-            estimate_params: dict = None,
-            effect_modifier_configuration: dict[Variable:Any] = None,
+        # pylint: disable=too-many-arguments
+        self,
+        base_test_case: BaseTestCase,
+        expected_causal_effect: CausalTestOutcome,
+        control_value: Any = None,
+        treatment_value: Any = None,
+        estimate_type: str = "ate",
+        estimate_params: dict = None,
+        effect_modifier_configuration: dict[Variable:Any] = None,
     ):
         """
         :param base_test_case: A BaseTestCase object consisting of a treatment variable, outcome variable and effect
@@ -85,7 +81,7 @@ def execute_test(self, estimator: type(Estimator), data_collector: Observational
         """Execute a causal test case and return the causal test result.
 
         :param estimator: A reference to an Estimator class.
-        :param causal_test_case: The CausalTestCase object to be tested
+        :param data_collector: The data collector to be used which provides a dataframe for the Estimator
         :return causal_test_result: A CausalTestResult for the executed causal test case.
         """
         if estimator.df is None:
@@ -100,11 +96,10 @@ def execute_test(self, estimator: type(Estimator), data_collector: Observational
         causal_test_result = self._return_causal_test_results(estimator)
         return causal_test_result
 
-    def _return_causal_test_results(self, estimator):
+    def _return_causal_test_results(self, estimator) -> CausalTestResult:
         """Depending on the estimator used, calculate the 95% confidence intervals and return in a causal_test_result
 
         :param estimator: An Estimator class object
-        :param causal_test_case: The concrete test case to be executed
         :return: a CausalTestResult object containing the confidence intervals
         """
         if not hasattr(estimator, f"estimate_{self.estimate_type}"):
@@ -128,4 +123,3 @@ def __str__(self):
             f"Running {treatment_config} instead of {control_config} should cause the following "
             f"changes to {outcome_variable}: {self.expected_causal_effect}."
         )
-

causal_testing/testing/causal_test_suite.py

@@ -7,9 +7,10 @@
 from causal_testing.testing.base_test_case import BaseTestCase
 from causal_testing.testing.causal_test_case import CausalTestCase
 from causal_testing.testing.estimators import Estimator
-from causal_testing.testing.causal_test_result import CausalTestResult, TestValue
+from causal_testing.testing.causal_test_result import CausalTestResult
 from causal_testing.data_collection.data_collector import ObservationalDataCollector
 from causal_testing.specification.causal_specification import CausalSpecification
+
 logger = logging.getLogger(__name__)
 
 
@@ -26,11 +27,11 @@ class CausalTestSuite(UserDict):
     """
 
     def add_test_object(
-            self,
-            base_test_case: BaseTestCase,
-            causal_test_case_list: Iterable[CausalTestCase],
-            estimators_classes: Iterable[Type[Estimator]],
-            estimate_type: str = "ate",
+        self,
+        base_test_case: BaseTestCase,
+        causal_test_case_list: Iterable[CausalTestCase],
+        estimators_classes: Iterable[Type[Estimator]],
+        estimate_type: str = "ate",
     ):
         """
         A setter object to allow for the easy construction of the dictionary test suite structure
@@ -44,16 +45,18 @@ def add_test_object(
         test_object = {"tests": causal_test_case_list, "estimators": estimators_classes, "estimate_type": estimate_type}
         self.data[base_test_case] = test_object
 
-    def execute_test_suite(self, data_collector: ObservationalDataCollector,
-                           causal_specification: CausalSpecification) -> list[CausalTestResult]:
+    def execute_test_suite(
+        self, data_collector: ObservationalDataCollector, causal_specification: CausalSpecification
+    ) -> dict[str, CausalTestResult]:
         """Execute a suite of causal tests and return the results in a list
-        :param test_suite: CasualTestSuite object
+        :param data_collector: The data collector to be used for the test_suite. Can be observational, experimental or
+         custom
+        :param causal_specification:
         :return: A dictionary where each key is the name of the estimators specified and the values are lists of
                 causal_test_result objects
         """
         if data_collector.data.empty:
             raise ValueError("No data has been loaded. Please call load_data prior to executing a causal test case.")
-        data_collector.collect_data()
         test_suite_results = {}
         for edge in self:
             logger.info("treatment: %s", edge.treatment_variable)
@@ -76,12 +79,9 @@ def execute_test_suite(self, data_collector: ObservationalDataCollector,
                         minimal_adjustment_set,
                         test.outcome_variable.name,
                     )
-                    if estimator.df is None:
-                        estimator.df = data_collector.collect_data()
-                    causal_test_result = test._return_causal_test_results(estimator)
+                    causal_test_result = test.execute_test(estimator, data_collector)
                     causal_test_results.append(causal_test_result)
 
                 results[estimator_class.__name__] = causal_test_results
             test_suite_results[edge] = results
         return test_suite_results
-